Avoid putting Arrow Records into the buckets

Author: vitalyisaev2

app/server/data_source_collection.go

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

app/server/datasource/rdbms/clickhouse/connection_http.go

@@ -73,6 +73,7 @@ func (c *connectionHTTP) Query(params *rdbms_utils.QueryParams) (*rdbms_utils.Qu
 	}
 
 	rows := &rows{Rows: out}
+
 	return &rdbms_utils.QueryResult{
 		Rows: rows,
 	}, nil

app/server/datasource/rdbms/clickhouse/connection_native.go

@@ -73,6 +73,7 @@ func (c *connectionNative) Query(params *rdbms_utils.QueryParams) (*rdbms_utils.
 	}
 
 	rows := &rowsNative{Rows: out}
+
 	return &rdbms_utils.QueryResult{
 		Rows: rows,
 	}, nil

app/server/datasource/rdbms/datasource.go

@@ -219,6 +219,7 @@ func (ds *dataSourceImpl) doReadSplitSingleConn(
 
 	rows := queryResult.Rows
 	transformer, err := rows.MakeTransformer(query.YdbColumns, ds.converterCollection)
+
 	if err != nil {
 		return 0, fmt.Errorf("make transformer: %w", err)
 	}

app/server/datasource/rdbms/utils/interface.go

@@ -3,9 +3,9 @@ package utils
 import (
 	"context"
 
+	"github.com/apache/arrow/go/v13/arrow"
 	"go.uber.org/zap"
 
-	"github.com/apache/arrow/go/v13/arrow"
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 
 	api_common "github.com/ydb-platform/fq-connector-go/api/common"

app/server/datasource/rdbms/ydb/connection_database_sql.go

@@ -74,6 +74,7 @@ func (c *connectionDatabaseSQL) Query(params *rdbms_utils.QueryParams) (*rdbms_u
 	}
 
 	rows := rowsDatabaseSQL{Rows: out}
+
 	return &rdbms_utils.QueryResult{
 		Rows: rows,
 	}, nil

app/server/datasource/rdbms/ydb/connection_native.go

@@ -8,15 +8,14 @@ import (
 	"io"
 	"time"
 
+	"github.com/apache/arrow/go/v13/arrow"
+	"github.com/apache/arrow/go/v13/arrow/ipc"
 	"go.uber.org/zap"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
 	ydb_sdk "github.com/ydb-platform/ydb-go-sdk/v3"
 	ydb_sdk_query "github.com/ydb-platform/ydb-go-sdk/v3/query"
 
-	"github.com/apache/arrow/go/v13/arrow"
-	"github.com/apache/arrow/go/v13/arrow/ipc"
-
 	api_common "github.com/ydb-platform/fq-connector-go/api/common"
 	"github.com/ydb-platform/fq-connector-go/app/config"
 	"github.com/ydb-platform/fq-connector-go/app/server/conversion"
@@ -126,6 +125,7 @@ func (c *columnsNative) Close() error {
 	if err := c.arrowResult.Close(c.ctx); err != nil {
 		return fmt.Errorf("arrow result close: %w", err)
 	}
+
 	return nil
 }
 
@@ -142,6 +142,7 @@ func (c *columnsNative) Next() bool {
 
 	// Try to get the next part
 	var part io.Reader
+
 	var err error
 
 	for p, e := range c.arrowResult.Parts(c.ctx) {
@@ -167,10 +168,12 @@ func (c *columnsNative) Next() bool {
 	// Create a new reader for this part
 	c.currentPart = part
 	reader, err := ipc.NewReader(part)
+
 	if err != nil {
 		c.err = fmt.Errorf("create arrow reader: %w", err)
 		return false
 	}
+
 	c.reader = reader
 
 	// Get the first record from this part
@@ -180,6 +183,7 @@ func (c *columnsNative) Next() bool {
 	}
 
 	c.record = c.reader.Record()
+
 	return true
 }
 

app/server/paging/columnar_buffer_arrow_ipc_streaming_default.go

@@ -60,6 +60,7 @@ func (cb *columnarBufferArrowIPCStreamingDefault[T]) addArrowRecord(record arrow
 // 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
@@ -99,13 +100,15 @@ func (cb *columnarBufferArrowIPCStreamingDefault[T]) ToResponse() (*api_service_
 		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)
 	}
 
@@ -127,9 +130,11 @@ 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
 }