Split columnar buffer implementations

Author: vitalyisaev2

app/server/data_source_collection.go

@@ -318,7 +318,8 @@ func doReadSplit[T paging.Acceptor](
 		logger,
 		memoryAllocator,
 		request.Format,
-		split.Select.What)
+		split.Select.What,
+		false) // Default to row-based approach for now
 	if err != nil {
 		return fmt.Errorf("new columnar buffer factory: %w", err)
 	}

app/server/paging/columnar_buffer_arrow_ipc_streaming_records.go

@@ -0,0 +1,96 @@
+package paging
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/apache/arrow/go/v13/arrow"
+	"github.com/apache/arrow/go/v13/arrow/ipc"
+	"github.com/apache/arrow/go/v13/arrow/memory"
+	"go.uber.org/zap"
+
+	api_service_protos "github.com/ydb-platform/fq-connector-go/api/service/protos"
+)
+
+var _ ColumnarBuffer[any] = (*columnarBufferArrowIPCStreamingRecords[any])(nil)
+
+// columnarBufferArrowIPCStreamingRecords is a specialized implementation for Arrow-based data
+// that stores Arrow records directly without deconstructing them
+type columnarBufferArrowIPCStreamingRecords[T Acceptor] struct {
+	arrowAllocator memory.Allocator
+	schema         *arrow.Schema
+	logger         *zap.Logger
+	arrowRecord    arrow.Record // Store the Arrow Record directly
+}
+
+// addRow is not the primary method for this implementation
+// It returns an error since this implementation is optimized for Arrow records
+func (cb *columnarBufferArrowIPCStreamingRecords[T]) addRow(transformer RowTransformer[T]) error {
+	return fmt.Errorf("this implementation is optimized for Arrow records, use columnarBufferArrowIPCStreamingRows for row-based data")
+}
+
+// addArrowRecord saves an Arrow Block obtained from the datasource into the columnar buffer
+func (cb *columnarBufferArrowIPCStreamingRecords[T]) addArrowRecord(record arrow.Record) error {
+	// Verify schema compatibility
+	if !cb.schema.Equal(record.Schema()) {
+		return fmt.Errorf("record schema does not match buffer schema")
+	}
+
+	// Store the record directly
+	if cb.arrowRecord != nil {
+		// Release the previous record if it exists
+		cb.arrowRecord.Release()
+	}
+
+	// Retain the record to prevent it from being garbage collected
+	record.Retain()
+	cb.arrowRecord = record
+
+	return nil
+}
+
+// ToResponse returns all the accumulated data and clears buffer
+func (cb *columnarBufferArrowIPCStreamingRecords[T]) ToResponse() (*api_service_protos.TReadSplitsResponse, error) {
+	// If no record was added, return an empty response
+	if cb.arrowRecord == nil {
+		return &api_service_protos.TReadSplitsResponse{}, nil
+	}
+
+	// prepare arrow writer
+	var buf bytes.Buffer
+
+	writer := ipc.NewWriter(&buf, ipc.WithSchema(cb.schema), ipc.WithAllocator(cb.arrowAllocator))
+
+	if err := writer.Write(cb.arrowRecord); err != nil {
+		return nil, fmt.Errorf("write record: %w", err)
+	}
+
+	if err := writer.Close(); err != nil {
+		return nil, fmt.Errorf("close arrow writer: %w", err)
+	}
+
+	out := &api_service_protos.TReadSplitsResponse{
+		Payload: &api_service_protos.TReadSplitsResponse_ArrowIpcStreaming{
+			ArrowIpcStreaming: buf.Bytes(),
+		},
+	}
+
+	return out, nil
+}
+
+// TotalRows returns the number of rows in the buffer
+func (cb *columnarBufferArrowIPCStreamingRecords[T]) TotalRows() int {
+	if cb.arrowRecord == nil {
+		return 0
+	}
+	return int(cb.arrowRecord.NumRows())
+}
+
+// Release frees resources if buffer is no longer used
+func (cb *columnarBufferArrowIPCStreamingRecords[T]) 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_rows.go

@@ -0,0 +1,105 @@
+package paging
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/apache/arrow/go/v13/arrow"
+	"github.com/apache/arrow/go/v13/arrow/array"
+	"github.com/apache/arrow/go/v13/arrow/ipc"
+	"github.com/apache/arrow/go/v13/arrow/memory"
+	"go.uber.org/zap"
+
+	api_service_protos "github.com/ydb-platform/fq-connector-go/api/service/protos"
+)
+
+var _ ColumnarBuffer[any] = (*columnarBufferArrowIPCStreamingRows[any])(nil)
+
+// columnarBufferArrowIPCStreamingRows is a specialized implementation for row-based data
+// that builds Arrow arrays from individual rows
+type columnarBufferArrowIPCStreamingRows[T Acceptor] struct {
+	arrowAllocator memory.Allocator
+	builders       []array.Builder
+	schema         *arrow.Schema
+	logger         *zap.Logger
+}
+
+// addRow saves a row obtained from the datasource into the buffer
+func (cb *columnarBufferArrowIPCStreamingRows[T]) addRow(transformer RowTransformer[T]) error {
+	if err := transformer.AppendToArrowBuilders(cb.schema, cb.builders); err != nil {
+		return fmt.Errorf("append values to arrow builders: %w", err)
+	}
+
+	return nil
+}
+
+// addArrowRecord is not the primary method for this implementation
+// It returns an error since this implementation is optimized for row-based data
+func (cb *columnarBufferArrowIPCStreamingRows[T]) addArrowRecord(record arrow.Record) error {
+	return fmt.Errorf("this implementation is optimized for row-based data, use columnarBufferArrowIPCStreamingRecords for Arrow records")
+}
+
+// ToResponse returns all the accumulated data and clears buffer
+func (cb *columnarBufferArrowIPCStreamingRows[T]) ToResponse() (*api_service_protos.TReadSplitsResponse, error) {
+	// If no rows were added, return an empty response
+	if cb.TotalRows() == 0 {
+		return &api_service_protos.TReadSplitsResponse{}, nil
+	}
+
+	// chunk consists of columns
+	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()
+	}
+
+	// prepare arrow writer
+	var buf bytes.Buffer
+
+	writer := ipc.NewWriter(&buf, ipc.WithSchema(cb.schema), ipc.WithAllocator(cb.arrowAllocator))
+
+	if err := writer.Write(record); err != nil {
+		record.Release()
+		return nil, fmt.Errorf("write record: %w", err)
+	}
+
+	if err := writer.Close(); err != nil {
+		record.Release()
+		return nil, fmt.Errorf("close arrow writer: %w", err)
+	}
+
+	// Release the record after writing
+	record.Release()
+
+	out := &api_service_protos.TReadSplitsResponse{
+		Payload: &api_service_protos.TReadSplitsResponse_ArrowIpcStreaming{
+			ArrowIpcStreaming: buf.Bytes(),
+		},
+	}
+
+	return out, nil
+}
+
+// TotalRows returns the number of rows in the buffer
+func (cb *columnarBufferArrowIPCStreamingRows[T]) TotalRows() int {
+	if len(cb.builders) == 0 {
+		return 0
+	}
+	return cb.builders[0].Len()
+}
+
+// Release frees resources if buffer is no longer used
+func (cb *columnarBufferArrowIPCStreamingRows[T]) Release() {
+	// cleanup builders
+	for _, b := range cb.builders {
+		b.Release()
+	}
+}

app/server/paging/columnar_buffer_factory.go

@@ -14,21 +14,18 @@ import (
 )
 
 type columnarBufferFactoryImpl[T Acceptor] struct {
-	arrowAllocator memory.Allocator
-	logger         *zap.Logger
-	format         api_service_protos.TReadSplitsRequest_EFormat
-	schema         *arrow.Schema
-	ydbTypes       []*Ydb.Type
+	arrowAllocator  memory.Allocator
+	logger          *zap.Logger
+	format          api_service_protos.TReadSplitsRequest_EFormat
+	schema          *arrow.Schema
+	ydbTypes        []*Ydb.Type
+	useArrowRecords bool // Whether to use Arrow records directly
 }
 
 func (cbf *columnarBufferFactoryImpl[T]) MakeBuffer() (ColumnarBuffer[T], error) {
 	switch cbf.format {
 	case api_service_protos.TReadSplitsRequest_ARROW_IPC_STREAMING:
-		builders, err := common.YdbTypesToArrowBuilders(cbf.ydbTypes, cbf.arrowAllocator)
-		if err != nil {
-			return nil, fmt.Errorf("convert Select.What to arrow.Schema: %w", err)
-		}
-
+		// Special case for empty columns
 		if len(cbf.ydbTypes) == 0 {
 			return &columnarBufferArrowIPCStreamingEmptyColumns[T]{
 				arrowAllocator: cbf.arrowAllocator,
@@ -37,12 +34,26 @@ func (cbf *columnarBufferFactoryImpl[T]) MakeBuffer() (ColumnarBuffer[T], error)
 			}, nil
 		}
 
-		return &columnarBufferArrowIPCStreamingDefault[T]{
-			arrowAllocator: cbf.arrowAllocator,
-			builders:       builders,
-			schema:         cbf.schema,
-			logger:         cbf.logger,
-		}, nil
+		// Choose implementation based on whether we're using Arrow records directly
+		if cbf.useArrowRecords {
+			return &columnarBufferArrowIPCStreamingRecords[T]{
+				arrowAllocator: cbf.arrowAllocator,
+				schema:         cbf.schema,
+				logger:         cbf.logger,
+			}, nil
+		} else {
+			builders, err := common.YdbTypesToArrowBuilders(cbf.ydbTypes, cbf.arrowAllocator)
+			if err != nil {
+				return nil, fmt.Errorf("convert Select.What to arrow.Schema: %w", err)
+			}
+
+			return &columnarBufferArrowIPCStreamingRows[T]{
+				arrowAllocator: cbf.arrowAllocator,
+				builders:       builders,
+				schema:         cbf.schema,
+				logger:         cbf.logger,
+			}, nil
+		}
 	default:
 		return nil, fmt.Errorf("unknown format: %v", cbf.format)
 	}
@@ -53,6 +64,7 @@ func NewColumnarBufferFactory[T Acceptor](
 	arrowAllocator memory.Allocator,
 	format api_service_protos.TReadSplitsRequest_EFormat,
 	selectWhat *api_service_protos.TSelect_TWhat,
+	useArrowRecords bool,
 ) (ColumnarBufferFactory[T], error) {
 	ydbTypes, err := common.SelectWhatToYDBTypes(selectWhat)
 	if err != nil {
@@ -65,11 +77,12 @@ func NewColumnarBufferFactory[T Acceptor](
 	}
 
 	cbf := &columnarBufferFactoryImpl[T]{
-		logger:         logger,
-		arrowAllocator: arrowAllocator,
-		format:         format,
-		schema:         schema,
-		ydbTypes:       ydbTypes,
+		logger:          logger,
+		arrowAllocator:  arrowAllocator,
+		format:          format,
+		schema:          schema,
+		ydbTypes:        ydbTypes,
+		useArrowRecords: useArrowRecords,
 	}
 
 	return cbf, nil

app/server/streaming/read_splits_streamer_test.go

@@ -271,7 +271,8 @@ func (tc testCaseStreaming) execute(t *testing.T) {
 		logger,
 		memory.NewGoAllocator(),
 		api_service_protos.TReadSplitsRequest_ARROW_IPC_STREAMING,
-		split.Select.What)
+		split.Select.What,
+		false) // Use row-based approach for tests
 	require.NoError(t, err)
 
 	pagingCfg := &config.TPagingConfig{RowsPerPage: uint64(tc.rowsPerPage)}