Fix Arrow batch row counts for CloudFetch results

Respect server-declared row counts when decoding Arrow IPC result streams so CloudFetch payload padding is not exposed as extra rows.

Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>

Author: msrathore-db

internal/rows/arrowbased/arrowRecordIterator.go

@@ -38,29 +38,67 @@ type arrowRecordIterator struct {
 	isFinished         bool
 	arrowSchemaBytes   []byte
 	arrowSchema        *arrow.Schema
+	nextRowNumber      int64
+	hasNextRowNumber   bool
 }
 
 var _ rows.ArrowBatchIterator = (*arrowRecordIterator)(nil)
 
 // Retrieve the next arrow record
 func (ri *arrowRecordIterator) Next() (arrow.Record, error) {
-	if !ri.HasNext() {
-		// returning EOF indicates that there are no more records to iterate
-		return nil, io.EOF
+	for {
+		if !ri.HasNext() {
+			// returning EOF indicates that there are no more records to iterate
+			return nil, io.EOF
+		}
+
+		// make sure we have the current batch
+		err := ri.getCurrentBatch()
+		if err != nil {
+			return nil, err
+		}
+
+		// return next record in current batch
+		r, err := ri.currentBatch.Next()
+		if err != nil {
+			ri.checkFinished()
+			return nil, err
+		}
+
+		r2 := ri.skipReturnedRows(r)
+		ri.checkFinished()
+		if r2 == nil {
+			continue
+		}
+
+		return r2, nil
 	}
+}
 
-	// make sure we have the current batch
-	err := ri.getCurrentBatch()
-	if err != nil {
-		return nil, err
+func (ri *arrowRecordIterator) skipReturnedRows(r SparkArrowRecord) arrow.Record {
+	if !ri.hasNextRowNumber {
+		ri.nextRowNumber = r.Start()
+		ri.hasNextRowNumber = true
 	}
 
-	// return next record in current batch
-	r, err := ri.currentBatch.Next()
+	if r.End() < ri.nextRowNumber {
+		r.Release()
+		return nil
+	}
 
-	ri.checkFinished()
+	if r.Start() < ri.nextRowNumber {
+		start := ri.nextRowNumber - r.Start()
+		sliced := r.NewSlice(start, r.NumRows())
+		r.Release()
+		if sliced == nil {
+			return nil
+		}
+		ri.nextRowNumber += sliced.NumRows()
+		return sliced
+	}
 
-	return r, err
+	ri.nextRowNumber = r.End() + 1
+	return r
 }
 
 // Indicate whether there are any more records available

internal/rows/arrowbased/arrowRecordIterator_test.go

@@ -8,6 +8,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/apache/arrow/go/v12/arrow"
 	"github.com/databricks/databricks-sql-go/driverctx"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/client"
@@ -193,6 +194,114 @@ func TestArrowRecordIterator(t *testing.T) {
 	})
 }
 
+func TestArrowRecordIterator_SkipsOverlappingReturnedRows(t *testing.T) {
+	var releasedOverlappingRecord bool
+	var slicedStart int64
+	var slicedEnd int64
+
+	overlappingRecord := fakeRecord{
+		fnNumRows: func() int64 { return 10 },
+		fnRelease: func() {
+			releasedOverlappingRecord = true
+		},
+		fnNewSlice: func(i, j int64) arrow.Record {
+			slicedStart = i
+			slicedEnd = j
+			return fakeRecord{fnNumRows: func() int64 { return j - i }}
+		},
+	}
+
+	rs := &arrowRecordIterator{
+		batchIterator: &fakeBatchIterator{
+			index: -1,
+			batches: []SparkArrowBatch{
+				&sparkArrowBatch{
+					Delimiter: rowscanner.NewDelimiter(0, 5),
+					arrowRecords: []SparkArrowRecord{
+						&sparkArrowRecord{
+							Delimiter: rowscanner.NewDelimiter(0, 5),
+							Record:    fakeRecord{fnNumRows: func() int64 { return 5 }},
+						},
+					},
+				},
+				&sparkArrowBatch{
+					Delimiter: rowscanner.NewDelimiter(0, 10),
+					arrowRecords: []SparkArrowRecord{
+						&sparkArrowRecord{
+							Delimiter: rowscanner.NewDelimiter(0, 10),
+							Record:    overlappingRecord,
+						},
+					},
+				},
+			},
+		},
+	}
+	defer rs.Close()
+
+	r1, err := rs.Next()
+	assert.NoError(t, err)
+	assert.Equal(t, int64(5), r1.NumRows())
+	r1.Release()
+
+	r2, err := rs.Next()
+	assert.NoError(t, err)
+	assert.Equal(t, int64(5), r2.NumRows())
+	r2.Release()
+
+	assert.True(t, releasedOverlappingRecord)
+	assert.Equal(t, int64(5), slicedStart)
+	assert.Equal(t, int64(10), slicedEnd)
+
+	r3, err := rs.Next()
+	assert.Nil(t, r3)
+	assert.ErrorIs(t, err, io.EOF)
+}
+
+func TestLimitArrowRecordsUsesExpectedRowCount(t *testing.T) {
+	var releasedOriginal bool
+	var releasedExtra bool
+	var slicedStart int64
+	var slicedEnd int64
+
+	records := []SparkArrowRecord{
+		&sparkArrowRecord{
+			Delimiter: rowscanner.NewDelimiter(10, 5),
+			Record: fakeRecord{
+				fnNumRows: func() int64 { return 5 },
+				fnRelease: func() {
+					releasedOriginal = true
+				},
+				fnNewSlice: func(i, j int64) arrow.Record {
+					slicedStart = i
+					slicedEnd = j
+					return fakeRecord{fnNumRows: func() int64 { return j - i }}
+				},
+			},
+		},
+		&sparkArrowRecord{
+			Delimiter: rowscanner.NewDelimiter(15, 5),
+			Record: fakeRecord{
+				fnNumRows: func() int64 { return 5 },
+				fnRelease: func() {
+					releasedExtra = true
+				},
+			},
+		},
+	}
+
+	limited := limitArrowRecords(records, 3)
+	defer limited[0].Release()
+
+	assert.Len(t, limited, 1)
+	assert.Equal(t, int64(10), limited[0].Start())
+	assert.Equal(t, int64(3), limited[0].Count())
+	assert.Equal(t, int64(3), limited[0].NumRows())
+	assert.True(t, releasedOriginal)
+	assert.True(t, releasedExtra)
+	assert.Equal(t, int64(0), slicedStart)
+	assert.Equal(t, int64(3), slicedEnd)
+}
+
 func TestArrowRecordIteratorSchema(t *testing.T) {
 	// Test with arrowSchemaBytes available
 	t.Run("schema with initial schema bytes", func(t *testing.T) {

internal/rows/arrowbased/batchloader.go

@@ -31,6 +31,10 @@ type IPCStreamIterator interface {
 	Close()
 }
 
+type positionedIPCStreamIterator interface {
+	NextWithMetadata() (io.Reader, int64, int64, error)
+}
+
 func NewCloudIPCStreamIterator(
 	ctx context.Context,
 	files []*cli_service.TSparkArrowResultLink,
@@ -136,21 +140,30 @@ type localIPCStreamIterator struct {
 var _ IPCStreamIterator = (*localIPCStreamIterator)(nil)
 
 func (bi *localIPCStreamIterator) Next() (io.Reader, error) {
+	reader, _, _, err := bi.NextWithMetadata()
+	return reader, err
+}
+
+func (bi *localIPCStreamIterator) NextWithMetadata() (io.Reader, int64, int64, error) {
 	cnt := len(bi.batches)
 	bi.index++
 	if bi.index < cnt {
 		ab := bi.batches[bi.index]
+		startRowOffset := bi.startRowOffset
+		for i := 0; i < bi.index; i++ {
+			startRowOffset += bi.batches[i].RowCount
+		}
 
 		reader := io.MultiReader(
 			bytes.NewReader(bi.arrowSchemaBytes),
 			getReader(bytes.NewReader(ab.Batch), bi.cfg.UseLz4Compression),
 		)
 
-		return reader, nil
+		return reader, startRowOffset, ab.RowCount, nil
 	}
 
 	bi.index = cnt
-	return nil, io.EOF
+	return nil, 0, 0, io.EOF
 }
 
 func (bi *localIPCStreamIterator) HasNext() bool {
@@ -176,6 +189,11 @@ type cloudIPCStreamIterator struct {
 var _ IPCStreamIterator = (*cloudIPCStreamIterator)(nil)
 
 func (bi *cloudIPCStreamIterator) Next() (io.Reader, error) {
+	reader, _, _, err := bi.NextWithMetadata()
+	return reader, err
+}
+
+func (bi *cloudIPCStreamIterator) NextWithMetadata() (io.Reader, int64, int64, error) {
 	for (bi.downloadTasks.Len() < bi.cfg.MaxDownloadThreads) && (bi.pendingLinks.Len() > 0) {
 		link := bi.pendingLinks.Dequeue()
 		logger.Debug().Msgf(
@@ -204,15 +222,15 @@ func (bi *cloudIPCStreamIterator) Next() (io.Reader, error) {
 
 	task := bi.downloadTasks.Dequeue()
 	if task == nil {
-		return nil, io.EOF
+		return nil, 0, 0, io.EOF
 	}
 
 	data, downloadMs, err := task.GetResult()
 
 	// once we've got an errored out task - cancel the remaining ones
 	if err != nil {
 		bi.Close()
-		return nil, err
+		return nil, 0, 0, err
 	}
 
 	// explicitly call cancel function on successfully completed task to avoid context leak
@@ -226,7 +244,7 @@ func (bi *cloudIPCStreamIterator) Next() (io.Reader, error) {
 		bi.onFileDownloaded(downloadMs)
 	}
 
-	return data, nil
+	return data, task.link.StartRowOffset, task.link.RowCount, nil
 }
 
 func (bi *cloudIPCStreamIterator) HasNext() bool {
@@ -558,15 +576,26 @@ func NewBatchIterator(ipcIterator IPCStreamIterator, startRowOffset int64) Batch
 }
 
 func (bi *batchIterator) Next() (SparkArrowBatch, error) {
-	reader, err := bi.ipcIterator.Next()
+	startRowOffset := bi.startRowOffset
+	expectedRows := int64(-1)
+	var reader io.Reader
+	var err error
+	if positionedIterator, ok := bi.ipcIterator.(positionedIPCStreamIterator); ok {
+		reader, startRowOffset, expectedRows, err = positionedIterator.NextWithMetadata()
+	} else {
+		reader, err = bi.ipcIterator.Next()
+	}
 	if err != nil {
 		return nil, err
 	}
 
-	records, err := getArrowRecords(reader, bi.startRowOffset)
+	records, err := getArrowRecords(reader, startRowOffset)
 	if err != nil {
 		return nil, err
 	}
+	if expectedRows >= 0 {
+		records = limitArrowRecords(records, expectedRows)
+	}
 
 	// When using CloudFetch, cached Arrow IPC files may contain stale column
 	// names from a previous query. Replace the embedded schema with the
@@ -593,14 +622,48 @@ func (bi *batchIterator) Next() (SparkArrowBatch, error) {
 	}
 
 	batch := &sparkArrowBatch{
-		Delimiter:    rowscanner.NewDelimiter(bi.startRowOffset, totalRows),
+		Delimiter:    rowscanner.NewDelimiter(startRowOffset, totalRows),
 		arrowRecords: records,
 	}
 
-	bi.startRowOffset += totalRows
+	bi.startRowOffset = startRowOffset + totalRows
 	return batch, nil
 }
 
+func limitArrowRecords(records []SparkArrowRecord, expectedRows int64) []SparkArrowRecord {
+	if expectedRows < 0 {
+		return records
+	}
+
+	remaining := expectedRows
+	limited := records[:0]
+	for _, record := range records {
+		if remaining <= 0 {
+			record.Release()
+			continue
+		}
+
+		if record.NumRows() <= remaining {
+			limited = append(limited, record)
+			remaining -= record.NumRows()
+			continue
+		}
+
+		start := record.Start()
+		sliced := record.NewSlice(0, remaining)
+		record.Release()
+		if sliced != nil {
+			limited = append(limited, &sparkArrowRecord{
+				Delimiter: rowscanner.NewDelimiter(start, sliced.NumRows()),
+				Record:    sliced,
+			})
+		}
+		remaining = 0
+	}
+
+	return limited
+}
+
 func (bi *batchIterator) HasNext() bool {
 	return bi.ipcIterator.HasNext()
 }