gofmt

Author: jadewang-db

internal/rows/arrowbased/ipc_stream_iterator.go

@@ -4,7 +4,7 @@ import (
 	"bytes"
 	"context"
 	"io"
-	
+
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/config"
 	"github.com/databricks/databricks-sql-go/internal/rows/rowscanner"
@@ -14,13 +14,13 @@ import (
 
 // ipcStreamIterator provides access to raw Arrow IPC streams without deserialization
 type ipcStreamIterator struct {
-	ctx              context.Context
+	ctx                context.Context
 	resultPageIterator rowscanner.ResultPageIterator
-	currentBatches   []*cli_service.TSparkArrowBatch
-	currentIndex     int
-	arrowSchemaBytes []byte
-	useLz4           bool
-	hasMorePages     bool
+	currentBatches     []*cli_service.TSparkArrowBatch
+	currentIndex       int
+	arrowSchemaBytes   []byte
+	useLz4             bool
+	hasMorePages       bool
 }
 
 // NewIPCStreamIterator creates an iterator that returns raw IPC streams
@@ -35,12 +35,12 @@ func NewIPCStreamIterator(
 	if cfg != nil {
 		useLz4 = cfg.UseLz4Compression
 	}
-	
+
 	var batches []*cli_service.TSparkArrowBatch
 	if initialRowSet != nil {
 		batches = initialRowSet.ArrowBatches
 	}
-	
+
 	return &ipcStreamIterator{
 		ctx:                ctx,
 		resultPageIterator: resultPageIterator,
@@ -59,38 +59,38 @@ func (it *ipcStreamIterator) NextIPCStream() (io.Reader, error) {
 		if !it.hasMorePages || it.resultPageIterator == nil {
 			return nil, io.EOF
 		}
-		
+
 		// Fetch next page
 		fetchResult, err := it.resultPageIterator.Next()
 		if err != nil {
 			return nil, err
 		}
-		
+
 		if fetchResult == nil || fetchResult.Results == nil || fetchResult.Results.ArrowBatches == nil {
 			return nil, io.EOF
 		}
-		
+
 		it.currentBatches = fetchResult.Results.ArrowBatches
 		it.currentIndex = 0
 		it.hasMorePages = it.resultPageIterator.HasNext()
-		
+
 		// If no batches in this page, recurse to try next page
 		if len(it.currentBatches) == 0 {
 			return it.NextIPCStream()
 		}
 	}
-	
+
 	batch := it.currentBatches[it.currentIndex]
 	it.currentIndex++
-	
+
 	// Create reader for the batch data
 	var batchReader io.Reader = bytes.NewReader(batch.Batch)
-	
+
 	// Handle LZ4 decompression if needed
 	if it.useLz4 {
 		batchReader = lz4.NewReader(batchReader)
 	}
-	
+
 	// Combine schema and batch data into a complete IPC stream
 	// Arrow IPC format expects: [Schema][Batch1][Batch2]...
 	return io.MultiReader(
@@ -112,4 +112,5 @@ func (it *ipcStreamIterator) Close() {
 // GetSchemaBytes returns the Arrow schema in IPC format
 func (it *ipcStreamIterator) GetSchemaBytes() ([]byte, error) {
 	return it.arrowSchemaBytes, nil
-}
+}
+

rows/ipc_stream.go

@@ -9,13 +9,13 @@ type IPCStreamIterator interface {
 	// GetNextIPCStream returns the next Arrow batch as an IPC stream reader
 	// Returns io.EOF when no more batches are available
 	NextIPCStream() (io.Reader, error)
-	
+
 	// HasNext returns true if there are more batches
 	HasNext() bool
-	
+
 	// Close releases any resources
 	Close()
-	
+
 	// GetSchemaBytes returns the Arrow schema in IPC format
 	GetSchemaBytes() ([]byte, error)
 }
@@ -25,4 +25,5 @@ type RowsWithIPCStream interface {
 	Rows
 	// GetIPCStreams returns an iterator for raw Arrow IPC streams
 	GetIPCStreams() (IPCStreamIterator, error)
-}
+}
+