PECO-1054 Expose Arrow batches to users

Step one of exposing arrow batches directly to users.
Moved the logic for iterating of the pages in a result set into ResultPageEditor.  Rows now composes in ResultPageIterator.
Introduced Delimeter type. Delimeter tracks a start/end point and provides functions for determining if a point is within the delimiter range and the direction of the point if it is outside the delimeter range.
Updated sparArrowBatch, arrowRowScanner, columnRows, rows to use Delimiter.
Updated the Fetch logic for cloudURL and localBatch so that the concurrentFetcher doesn't need to hold or pass through a Config instance.

Signed-off-by: Raymond Cypher <raymond.cypher@databricks.com>

Author: rcypher-databricks

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())
 			}
 		}