[SPARK-52780] Add StreamRows and Arrow Record Streaming

.github/workflows/build.yml

@@ -32,7 +32,7 @@ on:
       - master
 
 env:
-  SPARK_VERSION: '4.0.0'
+  SPARK_VERSION: '4.0.1'
   HADOOP_VERSION: '3'
 
 permissions:

internal/tests/integration/dataframe_test.go

@@ -1196,7 +1196,6 @@ func TestDataFrame_RangeIter(t *testing.T) {
 	}
 	assert.Equal(t, 10, cnt)
 
-	// Check that errors are properly propagated
 	df, err = spark.Sql(ctx, "select if(id = 5, raise_error('handle'), false) from range(10)")
 	assert.NoError(t, err)
 	for _, err := range df.All(ctx) {
@@ -1224,3 +1223,197 @@ func TestDataFrame_SchemaTreeString(t *testing.T) {
 	assert.Contains(t, ts, "|-- second: array")
 	assert.Contains(t, ts, "|-- third: map")
 }
+
+func TestDataFrame_StreamRowsThroughChannel(t *testing.T) {
+	// Demonstrates how StreamRows can be used to pipe data through a channel for scenarios like:
+	// - Proxying Spark data through gRPC streaming or unary RPCs
+	// - Implementing producer-consumer patterns with backpressure based on Spark results
+	// - Buffering and rate-limiting data flow between systems
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select id, id * 2 as doubled, 'test_' || cast(id as string) as label from range(100)")
+	assert.NoError(t, err)
+
+	iter, err := df.StreamRows(ctx)
+	assert.NoError(t, err)
+
+	rowChan := make(chan map[string]interface{}, 10)
+	errChan := make(chan error, 1)
+
+	go func() {
+		defer close(rowChan)
+		for row, err := range iter {
+			if err != nil {
+				errChan <- err
+				return
+			}
+
+			rowData := make(map[string]interface{})
+			names := row.FieldNames()
+			for i, name := range names {
+				rowData[name] = row.At(i)
+			}
+
+			select {
+			case rowChan <- rowData:
+			case <-ctx.Done():
+				errChan <- ctx.Err()
+				return
+			}
+		}
+	}()
+
+	// In a gRPC scenario, this would be your response handler
+	receivedRows := make([]map[string]interface{}, 0)
+	consumerDone := make(chan struct{})
+
+	go func() {
+		defer close(consumerDone)
+		for rowData := range rowChan {
+			receivedRows = append(receivedRows, rowData)
+
+			id := rowData["id"].(int64)
+			doubled := rowData["doubled"].(int64)
+			assert.Equal(t, id*2, doubled)
+		}
+	}()
+
+	<-consumerDone
+
+	select {
+	case err := <-errChan:
+		assert.NoError(t, err)
+	default:
+		// continue
+	}
+
+	assert.Equal(t, 100, len(receivedRows))
+
+	assert.Equal(t, int64(0), receivedRows[0]["id"])
+	assert.Equal(t, int64(99), receivedRows[99]["id"])
+	assert.Equal(t, "test_0", receivedRows[0]["label"])
+	assert.Equal(t, "test_99", receivedRows[99]["label"])
+}
+
+func TestDataFrame_StreamRows(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select * from range(100)")
+	assert.NoError(t, err)
+
+	iter, err := df.StreamRows(ctx)
+	assert.NoError(t, err)
+	assert.NotNil(t, iter)
+
+	cnt := 0
+
+	for row, err := range iter {
+		assert.NoError(t, err)
+		assert.NotNil(t, row)
+		assert.Equal(t, 1, row.Len())
+		cnt++
+	}
+	assert.Equal(t, 100, cnt)
+}
+
+func TestDataFrame_StreamRowsWithFilter(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select * from range(100)")
+	assert.NoError(t, err)
+
+	df, err = df.Filter(ctx, functions.Col("id").Lt(functions.IntLit(10)))
+	assert.NoError(t, err)
+
+	iter, err := df.StreamRows(ctx)
+	assert.NoError(t, err)
+
+	cnt := 0
+	for row, err := range iter {
+		assert.NoError(t, err)
+		assert.NotNil(t, row)
+		assert.Equal(t, 1, row.Len())
+		assert.Less(t, row.At(0).(int64), int64(10))
+		cnt++
+	}
+	assert.Equal(t, 10, cnt)
+}
+
+func TestDataFrame_StreamRowsEmpty(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select * from range(0)")
+	assert.NoError(t, err)
+
+	iter, err := df.StreamRows(ctx)
+	assert.NoError(t, err)
+
+	cnt := 0
+	for row, err := range iter {
+		assert.NoError(t, err)
+		assert.NotNil(t, row)
+		cnt++
+	}
+	assert.Equal(t, 0, cnt)
+}
+
+func TestDataFrame_StreamRowsWithError(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select if(id = 5, raise_error('test error'), id) as id from range(10)")
+	assert.NoError(t, err)
+
+	iter, err := df.StreamRows(ctx)
+	assert.NoError(t, err)
+
+	errorEncountered := false
+	for _, err := range iter {
+		if err != nil {
+			errorEncountered = true
+			assert.Error(t, err)
+			break
+		}
+	}
+	assert.True(t, errorEncountered, "Expected to encounter an error during iteration")
+}
+
+func TestDataFrame_StreamRowsMultipleColumns(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select id, id * 2 as doubled, 'test' as name from range(50)")
+	assert.NoError(t, err)
+
+	iter, err := df.StreamRows(ctx)
+	assert.NoError(t, err)
+
+	cnt := 0
+	for row, err := range iter {
+		assert.NoError(t, err)
+		assert.NotNil(t, row)
+		assert.Equal(t, 3, row.Len())
+
+		id := row.At(0).(int64)
+		doubled := row.At(1).(int64)
+		name := row.At(2).(string)
+
+		assert.Equal(t, id*2, doubled)
+		assert.Equal(t, "test", name)
+		cnt++
+	}
+	assert.Equal(t, 50, cnt)
+}
+
+func TestDataFrame_StreamRowsLargeDataset(t *testing.T) {
+	ctx, spark := connect()
+	df, err := spark.Sql(ctx, "select * from range(10000)")
+	assert.NoError(t, err)
+
+	iter, err := df.StreamRows(ctx)
+	assert.NoError(t, err)
+
+	cnt := 0
+	lastValue := int64(-1)
+	for row, err := range iter {
+		assert.NoError(t, err)
+		assert.NotNil(t, row)
+		currentValue := row.At(0).(int64)
+		assert.Greater(t, currentValue, lastValue)
+		lastValue = currentValue
+		cnt++
+	}
+	assert.Equal(t, 10000, cnt)
+}

spark/client/base/base.go

@@ -17,6 +17,7 @@ package base
 
 import (
 	"context"
+	"iter"
 
 	"github.com/apache/spark-connect-go/spark/sql/utils"
 
@@ -47,6 +48,9 @@ type SparkConnectClient interface {
 }
 
 type ExecuteResponseStream interface {
+	// ToTable consumes all arrow.Record batches to a single arrow.Table. Useful for collecting all query results into a client DF.
 	ToTable() (*types.StructType, arrow.Table, error)
+	// ToRecordSequence lazily consumes each arrow.Record retrieved by a query. Useful for streaming query results.
+	ToRecordSequence(ctx context.Context) iter.Seq2[arrow.Record, error]
 	Properties() map[string]any
 }

spark/client/client.go

@@ -20,6 +20,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"iter"
 
 	"github.com/apache/spark-connect-go/spark/sql/utils"
 
@@ -368,6 +369,15 @@ func (c *ExecutePlanClient) ToTable() (*types.StructType, arrow.Table, error) {
 	c.done = false
 	for {
 		resp, err := c.responseStream.Recv()
+		if err != nil {
+			fmt.Printf("DEBUG: Recv error: %v, is EOF: %v\n", err, errors.Is(err, io.EOF))
+		}
+		if err == nil && resp != nil {
+			fmt.Printf("DEBUG: Received response type: %T\n", resp.ResponseType)
+			if _, ok := resp.ResponseType.(*proto.ExecutePlanResponse_ResultComplete_); ok {
+				fmt.Println("DEBUG: Got ResultComplete!")
+			}
+		}
 		// EOF is received when the last message has been processed and the stream
 		// finished normally.
 		if errors.Is(err, io.EOF) {
@@ -434,6 +444,110 @@ func (c *ExecutePlanClient) ToTable() (*types.StructType, arrow.Table, error) {
 	}
 }
 
+// ToRecordSequence returns a single Seq2 iterator
+func (c *ExecutePlanClient) ToRecordSequence(ctx context.Context) iter.Seq2[arrow.Record, error] {
+	// Return Seq2 iterator that directly yields results as they arrive, upstream callers can convert this as needed
+	iterator := func(yield func(arrow.Record, error) bool) {
+		// Explicitly needed when tracking re-attachable execution.
+		c.done = false
+
+		for {
+			// Check for context cancellation before each iteration
+			select {
+			case <-ctx.Done():
+				// Yield the context error and stop
+				yield(nil, ctx.Err())
+				return
+			default:
+			}
+
+			resp, err := c.responseStream.Recv()
+
+			// Check for context cancellation after potentially blocking operations
+			select {
+			case <-ctx.Done():
+				yield(nil, ctx.Err())
+				return
+			default:
+			}
+
+			// EOF is received when the last message has been processed (Observed on Databricks instances)
+			if errors.Is(err, io.EOF) {
+				return // Clean end of stream
+			}
+
+			// Handle other errors
+			if err != nil {
+				if se := sparkerrors.FromRPCError(err); se != nil {
+					yield(nil, sparkerrors.WithType(se, sparkerrors.ExecutionError))
+				} else {
+					yield(nil, err)
+				}
+				return // Stop on error
+			}
+
+			// Only proceed if we have a valid response
+			if resp == nil {
+				continue
+			}
+
+			// Validate session ID
+			if resp.GetSessionId() != c.sessionId {
+				yield(nil, sparkerrors.WithType(
+					&sparkerrors.InvalidServerSideSessionDetailsError{
+						OwnSessionId:      c.sessionId,
+						ReceivedSessionId: resp.GetSessionId(),
+					}, sparkerrors.InvalidServerSideSessionError))
+				return
+			}
+
+			// Process schema if present
+			if resp.Schema != nil {
+				var schemaErr error
+				c.schema, schemaErr = types.ConvertProtoDataTypeToStructType(resp.Schema)
+				if schemaErr != nil {
+					yield(nil, sparkerrors.WithType(schemaErr, sparkerrors.ExecutionError))
+					return
+				}
+			}
+
+			// Process response types
+			switch x := resp.ResponseType.(type) {
+			case *proto.ExecutePlanResponse_SqlCommandResult_:
+				if val := x.SqlCommandResult.GetRelation(); val != nil {
+					c.properties["sql_command_result"] = val
+				}
+
+			case *proto.ExecutePlanResponse_ArrowBatch_:
+				record, err := types.ReadArrowBatchToRecord(x.ArrowBatch.Data, c.schema)
+				if err != nil {
+					yield(nil, err)
+					return
+				}
+
+				// Yield the record and check if consumer wants to continue
+				if !yield(record, nil) {
+					// Consumer stopped iteration early
+					// Note: Consumer is responsible for releasing the record
+					return
+				}
+
+			case *proto.ExecutePlanResponse_ResultComplete_:
+				c.done = true
+				return
+
+			case *proto.ExecutePlanResponse_ExecutionProgress_:
+				// Progress updates - ignore for now
+
+			default:
+				// Explicitly ignore unknown message types
+			}
+		}
+	}
+
+	return iterator
+}
+
 func NewExecuteResponseStream(
 	responseClient proto.SparkConnectService_ExecutePlanClient,
 	sessionId string,