Change the columnar buffer interface

Author: vitalyisaev2

app/server/paging/columnar_buffer_arrow_ipc_streaming_default.go

@@ -20,6 +20,8 @@ type columnarBufferArrowIPCStreamingDefault[T Acceptor] struct {
 	builders       []array.Builder
 	schema         *arrow.Schema
 	logger         *zap.Logger
+	arrowRecord    arrow.Record // Store the Arrow Record directly
+	rowsAdded      bool         // Track if rows were added via addRow
 }
 
 // AddRow saves a row obtained from the datasource into the buffer
@@ -30,23 +32,62 @@ func (cb *columnarBufferArrowIPCStreamingDefault[T]) addRow(transformer RowTrans
 		return fmt.Errorf("append values to arrow builders: %w", err)
 	}
 
+	cb.rowsAdded = true
+
 	return nil
 }
 
-// ToResponse returns all the accumulated data and clears buffer
-func (cb *columnarBufferArrowIPCStreamingDefault[T]) ToResponse() (*api_service_protos.TReadSplitsResponse, error) {
-	// chunk consists of columns
-	chunk := make([]arrow.Array, 0, len(cb.builders))
+// addArrowRecord saves an Arrow Block obtained from the datasource into the columnar buffer
+func (cb *columnarBufferArrowIPCStreamingDefault[T]) addArrowRecord(record arrow.Record) error {
+	// Create a new record with the same schema as the buffer
+	if !cb.schema.Equal(record.Schema()) {
+		return fmt.Errorf("record schema does not match buffer schema")
+	}
 
-	// prepare arrow record
-	for _, builder := range cb.builders {
-		chunk = append(chunk, builder.NewArray())
+	// Store the record directly
+	if cb.arrowRecord != nil {
+		// Release the previous record if it exists
+		cb.arrowRecord.Release()
 	}
 
-	record := array.NewRecord(cb.schema, chunk, -1)
+	// Retain the record to prevent it from being garbage collected
+	record.Retain()
+	cb.arrowRecord = record
 
-	for _, col := range chunk {
-		col.Release()
+	return nil
+}
+
+// ToResponse returns all the accumulated data and clears buffer
+func (cb *columnarBufferArrowIPCStreamingDefault[T]) ToResponse() (*api_service_protos.TReadSplitsResponse, error) {
+	var record arrow.Record
+	var releaseRecord bool
+
+	// If we have a stored Arrow Record, use it directly
+	if cb.arrowRecord != nil {
+		record = cb.arrowRecord
+		// We'll release our reference to the record at the end
+		releaseRecord = false
+	} else if cb.rowsAdded {
+		// If rows were added, create a new record from the builders
+		chunk := make([]arrow.Array, 0, len(cb.builders))
+
+		// prepare arrow record
+		for _, builder := range cb.builders {
+			chunk = append(chunk, builder.NewArray())
+		}
+
+		record = array.NewRecord(cb.schema, chunk, -1)
+
+		// We need to release the arrays after creating the record
+		for _, col := range chunk {
+			col.Release()
+		}
+
+		// We'll need to release this record after writing it
+		releaseRecord = true
+	} else {
+		// No data to return
+		return &api_service_protos.TReadSplitsResponse{}, nil
 	}
 
 	// prepare arrow writer
@@ -55,13 +96,24 @@ func (cb *columnarBufferArrowIPCStreamingDefault[T]) ToResponse() (*api_service_
 	writer := ipc.NewWriter(&buf, ipc.WithSchema(cb.schema), ipc.WithAllocator(cb.arrowAllocator))
 
 	if err := writer.Write(record); err != nil {
+		if releaseRecord {
+			record.Release()
+		}
 		return nil, fmt.Errorf("write record: %w", err)
 	}
 
 	if err := writer.Close(); err != nil {
+		if releaseRecord {
+			record.Release()
+		}
 		return nil, fmt.Errorf("close arrow writer: %w", err)
 	}
 
+	// Release the record if we created it
+	if releaseRecord {
+		record.Release()
+	}
+
 	out := &api_service_protos.TReadSplitsResponse{
 		Payload: &api_service_protos.TReadSplitsResponse_ArrowIpcStreaming{
 			ArrowIpcStreaming: buf.Bytes(),
@@ -71,12 +123,26 @@ func (cb *columnarBufferArrowIPCStreamingDefault[T]) ToResponse() (*api_service_
 	return out, nil
 }
 
-func (cb *columnarBufferArrowIPCStreamingDefault[T]) TotalRows() int { return cb.builders[0].Len() }
+func (cb *columnarBufferArrowIPCStreamingDefault[T]) TotalRows() int {
+	if cb.arrowRecord != nil {
+		return int(cb.arrowRecord.NumRows())
+	}
+	if len(cb.builders) > 0 {
+		return cb.builders[0].Len()
+	}
+	return 0
+}
 
 // Frees resources if buffer is no longer used
 func (cb *columnarBufferArrowIPCStreamingDefault[T]) Release() {
 	// cleanup builders
 	for _, b := range cb.builders {
 		b.Release()
 	}
+
+	// Release the stored Arrow Record if it exists
+	if cb.arrowRecord != nil {
+		cb.arrowRecord.Release()
+		cb.arrowRecord = nil
+	}
 }

app/server/paging/columnar_buffer_arrow_ipc_streaming_empty_columns.go

@@ -34,6 +34,13 @@ func (cb *columnarBufferArrowIPCStreamingEmptyColumns[T]) addRow(transformer Row
 	return nil
 }
 
+// addArrowRecord saves an Arrow Block obtained from the datasource into the columnar buffer
+func (cb *columnarBufferArrowIPCStreamingEmptyColumns[T]) addArrowRecord(record arrow.Record) error {
+	// For empty columns buffer, we just need to count the number of rows
+	cb.rowsAdded += int(record.NumRows())
+	return nil
+}
+
 // ToResponse returns all the accumulated data and clears buffer
 func (cb *columnarBufferArrowIPCStreamingEmptyColumns[T]) ToResponse() (*api_service_protos.TReadSplitsResponse, error) {
 	columns := make([]arrow.Array, 0)

app/server/paging/interface.go

@@ -24,6 +24,8 @@ type RowTransformer[T Acceptor] interface {
 type ColumnarBuffer[T Acceptor] interface {
 	// addRow saves a row obtained from the datasource into the columnar buffer
 	addRow(rowTransformer RowTransformer[T]) error
+	// addArrowRecord saves an Arrow Block obtained from the datasource into the columnar buffer
+	addArrowRecord(record arrow.Record) error
 	// ToResponse returns all the accumulated data and clears buffer
 	ToResponse() (*api_service_protos.TReadSplitsResponse, error)
 	// Release frees resources if buffer is no longer used
@@ -54,6 +56,9 @@ type Sink[T Acceptor] interface {
 	// AddRow saves the row obtained from a stream incoming from an external data source.
 	AddRow(rowTransformer RowTransformer[T]) error
 
+	// AddArrowRecord saves the Arrow block obtained from a stream incoming from an external data source.
+	// AddArrowRecord(record arrow.Record) error
+
 	// Finish reports the successful (!) completion of data stream reading.
 	// Never call this method if the request has failed.
 	// This method can be called only once.

app/server/paging/mock.go

@@ -1,6 +1,7 @@
 package paging
 
 import (
+	"github.com/apache/arrow/go/v13/arrow"
 	"github.com/stretchr/testify/mock"
 	"go.uber.org/zap"
 
@@ -66,6 +67,10 @@ func (*ColumnarBufferMock) addRow(_ RowTransformer[any]) error {
 	panic("not implemented") // TODO: Implement
 }
 
+func (*ColumnarBufferMock) addArrowRecord(_ arrow.Record) error {
+	panic("not implemented") // TODO: Implement
+}
+
 func (m *ColumnarBufferMock) ToResponse() (*api_service_protos.TReadSplitsResponse, error) {
 	args := m.Called()