PECO-1054 Expose Arrow batches to users

internal/fetcher/fetcher.go

@@ -2,15 +2,14 @@ package fetcher
 
 import (
 	"context"
-	"github.com/databricks/databricks-sql-go/internal/config"
 	"sync"
 
 	"github.com/databricks/databricks-sql-go/driverctx"
 	dbsqllog "github.com/databricks/databricks-sql-go/logger"
 )
 
 type FetchableItems[OutputType any] interface {
-	Fetch(ctx context.Context, cfg *config.Config) ([]OutputType, error)
+	Fetch(ctx context.Context) ([]OutputType, error)
 }
 
 type Fetcher[OutputType any] interface {
@@ -24,7 +23,6 @@ type concurrentFetcher[I FetchableItems[O], O any] struct {
 	outChan    chan O
 	err        error
 	nWorkers   int
-	cfg        *config.Config
 	mu         sync.Mutex
 	start      sync.Once
 	ctx        context.Context
@@ -100,10 +98,17 @@ func (f *concurrentFetcher[I, O]) logger() *dbsqllog.DBSQLLogger {
 	return f.DBSQLLogger
 }
 
-func NewConcurrentFetcher[I FetchableItems[O], O any](ctx context.Context, nWorkers int, cfg *config.Config, inputChan <-chan FetchableItems[O]) (Fetcher[O], error) {
+func NewConcurrentFetcher[I FetchableItems[O], O any](ctx context.Context, nWorkers, maxItemsInMemory int, inputChan <-chan FetchableItems[O]) (Fetcher[O], error) {
+	if nWorkers < 1 {
+		nWorkers = 1
+	}
+	if maxItemsInMemory < 1 {
+		maxItemsInMemory = 1
+	}
+
 	// channel for loaded items
 	// TODO: pass buffer size
-	outputChannel := make(chan O, 100)
+	outputChannel := make(chan O, maxItemsInMemory)
 
 	// channel to signal a cancel
 	stopChannel := make(chan bool)
@@ -118,7 +123,6 @@ func NewConcurrentFetcher[I FetchableItems[O], O any](ctx context.Context, nWork
 		cancelChan: stopChannel,
 		ctx:        ctx,
 		nWorkers:   nWorkers,
-		cfg:        cfg,
 	}
 
 	return fetcher, nil
@@ -139,7 +143,7 @@ func work[I FetchableItems[O], O any](f *concurrentFetcher[I, O], workerIndex in
 		case input, ok := <-f.inputChan:
 			if ok {
 				f.logger().Debug().Msgf("concurrent fetcher worker %d loading item", workerIndex)
-				result, err := input.Fetch(f.ctx, f.cfg)
+				result, err := input.Fetch(f.ctx)
 				if err != nil {
 					f.logger().Debug().Msgf("concurrent fetcher worker %d received error", workerIndex)
 					f.setErr(err)

internal/fetcher/fetcher_test.go

@@ -2,11 +2,11 @@ package fetcher
 
 import (
 	"context"
-	"github.com/databricks/databricks-sql-go/internal/config"
-	"github.com/pkg/errors"
 	"math"
 	"testing"
 	"time"
+
+	"github.com/pkg/errors"
 )
 
 // Create a mock struct for FetchableItems
@@ -20,7 +20,7 @@ type mockOutput struct {
 }
 
 // Implement the Fetch method
-func (m *mockFetchableItem) Fetch(ctx context.Context, cfg *config.Config) ([]*mockOutput, error) {
+func (m *mockFetchableItem) Fetch(ctx context.Context) ([]*mockOutput, error) {
 	time.Sleep(m.wait)
 	outputs := make([]*mockOutput, 5)
 	for i := range outputs {
@@ -35,7 +35,7 @@ var _ FetchableItems[*mockOutput] = (*mockFetchableItem)(nil)
 func TestConcurrentFetcher(t *testing.T) {
 	t.Run("Comprehensively tests the concurrent fetcher", func(t *testing.T) {
 		ctx := context.Background()
-		cfg := &config.Config{}
+
 		inputChan := make(chan FetchableItems[*mockOutput], 10)
 		for i := 0; i < 10; i++ {
 			item := mockFetchableItem{item: i, wait: 1 * time.Second}
@@ -44,7 +44,7 @@ func TestConcurrentFetcher(t *testing.T) {
 		close(inputChan)
 
 		// Create a fetcher
-		fetcher, err := NewConcurrentFetcher[*mockFetchableItem](ctx, 3, cfg, inputChan)
+		fetcher, err := NewConcurrentFetcher[*mockFetchableItem](ctx, 3, 3, inputChan)
 		if err != nil {
 			t.Fatalf("Error creating fetcher: %v", err)
 		}
@@ -95,7 +95,7 @@ func TestConcurrentFetcher(t *testing.T) {
 		close(inputChan)
 
 		// Create a new fetcher
-		fetcher, err := NewConcurrentFetcher[*mockFetchableItem](ctx, 2, &config.Config{}, inputChan)
+		fetcher, err := NewConcurrentFetcher[*mockFetchableItem](ctx, 2, 2, inputChan)
 		if err != nil {
 			t.Fatalf("Error creating fetcher: %v", err)
 		}

internal/rows/arrowbased/arrowRows.go

@@ -27,13 +27,9 @@ type valueContainerMaker interface {
 }
 
 type sparkArrowBatch struct {
-	rowCount, startRow, endRow int64
-	arrowRecordBytes           []byte
-	hasSchema                  bool
-}
-
-func (sab *sparkArrowBatch) contains(rowIndex int64) bool {
-	return sab != nil && sab.startRow <= rowIndex && sab.endRow >= rowIndex
+	rowscanner.Delimiter
+	arrowRecordBytes []byte
+	hasSchema        bool
 }
 
 type timeStampFn func(arrow.Timestamp) time.Time
@@ -46,6 +42,7 @@ type colInfo struct {
 
 // arrowRowScanner handles extracting values from arrow records
 type arrowRowScanner struct {
+	rowscanner.Delimiter
 	recordReader
 	valueContainerMaker
 
@@ -61,9 +58,6 @@ type arrowRowScanner struct {
 	// database types for the columns
 	colInfo []colInfo
 
-	// number of rows in the current TRowSet
-	nRows int64
-
 	// a TRowSet contains multiple arrow batches
 	currentBatch *sparkArrowBatch
 
@@ -140,12 +134,12 @@ func NewArrowRowScanner(resultSetMetadata *cli_service.TGetResultSetMetadataResp
 	}
 
 	rs := &arrowRowScanner{
+		Delimiter: rowscanner.NewDelimiter(rowSet.StartRowOffset, rowscanner.CountRows(rowSet)),
 		recordReader: sparkRecordReader{
 			ctx: ctx,
 		},
 		valueContainerMaker: &arrowValueContainerMaker{},
 		ArrowConfig:         arrowConfig,
-		nRows:               countRows(rowSet),
 		arrowSchemaBytes:    schemaBytes,
 		arrowSchema:         arrowSchema,
 		toTimestampFn:       ttsf,
@@ -172,7 +166,7 @@ func (ars *arrowRowScanner) Close() {
 // NRows returns the number of rows in the current set of batches
 func (ars *arrowRowScanner) NRows() int64 {
 	if ars != nil {
-		return ars.nRows
+		return ars.Count()
 	}
 
 	return 0
@@ -203,7 +197,7 @@ func (ars *arrowRowScanner) ScanRow(
 		return err
 	}
 
-	var rowInBatchIndex int = int(rowIndex - ars.currentBatch.startRow)
+	var rowInBatchIndex int = int(rowIndex - ars.currentBatch.Start())
 
 	// if no location is provided default to UTC
 	if ars.location == nil {
@@ -248,41 +242,14 @@ func isIntervalType(typeId cli_service.TTypeId) bool {
 	return ok
 }
 
-// countRows returns the number of rows in the TRowSet
-func countRows(rowSet *cli_service.TRowSet) int64 {
-	if rowSet == nil {
-		return 0
-	}
-
-	if rowSet.ArrowBatches != nil {
-		batches := rowSet.ArrowBatches
-		var n int64
-		for i := range batches {
-			n += batches[i].RowCount
-		}
-		return n
-	}
-
-	if rowSet.ResultLinks != nil {
-		links := rowSet.ResultLinks
-		var n int64
-		for i := range links {
-			n += links[i].RowCount
-		}
-		return n
-	}
-
-	return 0
-}
-
 // loadBatchFor loads the batch containing the specified row if necessary
 func (ars *arrowRowScanner) loadBatchFor(rowIndex int64) dbsqlerr.DBError {
 
 	if ars == nil || ars.BatchLoader == nil {
 		return dbsqlerrint.NewDriverError(context.Background(), errArrowRowsNoArrowBatches, nil)
 	}
 	// if the batch already loaded we can just return
-	if ars.currentBatch != nil && ars.currentBatch.contains(rowIndex) && ars.columnValues != nil {
+	if ars.currentBatch != nil && ars.currentBatch.Contains(rowIndex) && ars.columnValues != nil {
 		return nil
 	}
 

internal/rows/arrowbased/arrowRows_test.go

@@ -663,23 +663,23 @@ func TestArrowRowScanner(t *testing.T) {
 			assert.Nil(t, err)
 			assert.NotNil(t, lastReadBatch)
 			assert.Equal(t, 1, callCount)
-			assert.Equal(t, int64(0), lastReadBatch.startRow)
+			assert.Equal(t, int64(0), lastReadBatch.Start())
 		}
 
 		for _, i := range []int64{5, 6, 7} {
 			err := ars.loadBatchFor(i)
 			assert.Nil(t, err)
 			assert.NotNil(t, lastReadBatch)
 			assert.Equal(t, 2, callCount)
-			assert.Equal(t, int64(5), lastReadBatch.startRow)
+			assert.Equal(t, int64(5), lastReadBatch.Start())
 		}
 
 		for _, i := range []int64{8, 9, 10, 11, 12, 13, 14} {
 			err := ars.loadBatchFor(i)
 			assert.Nil(t, err)
 			assert.NotNil(t, lastReadBatch)
 			assert.Equal(t, 3, callCount)
-			assert.Equal(t, int64(8), lastReadBatch.startRow)
+			assert.Equal(t, int64(8), lastReadBatch.Start())
 		}
 
 		err := ars.loadBatchFor(-1)
@@ -983,13 +983,13 @@ func TestArrowRowScanner(t *testing.T) {
 
 			if i%1000 == 0 {
 				assert.NotNil(t, ars.currentBatch)
-				assert.Equal(t, int64(i), ars.currentBatch.startRow)
+				assert.Equal(t, int64(i), ars.currentBatch.Start())
 				if i < 53000 {
-					assert.Equal(t, int64(1000), ars.currentBatch.rowCount)
+					assert.Equal(t, int64(1000), ars.currentBatch.Count())
 				} else {
-					assert.Equal(t, int64(940), ars.currentBatch.rowCount)
+					assert.Equal(t, int64(940), ars.currentBatch.Count())
 				}
-				assert.Equal(t, ars.currentBatch.startRow+ars.currentBatch.rowCount-1, ars.currentBatch.endRow)
+				assert.Equal(t, ars.currentBatch.Start()+ars.currentBatch.Count()-1, ars.currentBatch.End())
 			}
 		}