perf(cubeorchestrator): Arrow-backed CubeStore results + columnar driver batches

packages/cubejs-backend-native/js/index.ts

@@ -290,6 +290,14 @@ function wrapNativeFunctionWithStream(
           objectMode: true,
           highWaterMark: chunkLength,
           write(row: any, encoding: BufferEncoding, callback: (error?: (Error | null)) => void) {
+            // Columnar batch produced by a driver (base-driver `ColumnarResponse`,
+            // marker `COLUMNAR_RESPONSE_TYPE`): already `{ members, columns }`, so send it
+            // straight through without row accumulation or the row→columnar transpose.
+            if (row && row.$type === 'ColumnarResponse') {
+              const toSend = Buffer.from(JSON.stringify({ members: row.members, columns: row.columns }));
+              writerOrChannel.chunk(toSend, callback);
+              return;
+            }
             chunkBuffer.push(row);
             if (chunkBuffer.length < chunkLength) {
               callback(null);

packages/cubejs-backend-native/src/orchestrator.rs

@@ -1,8 +1,9 @@
 use crate::node_obj_deserializer::JsValueDeserializer;
 use crate::transport::MapCubeErrExt;
+use cubeorchestrator::arrow_column::DBResponseValueRef;
 use cubeorchestrator::query_message_parser::QueryResult;
 use cubeorchestrator::query_result_transform::{
-    DBResponsePrimitive, RequestResultData, RequestResultDataMulti, TransformedData,
+    RequestResultData, RequestResultDataMulti, TransformedData,
 };
 use cubeorchestrator::transport::{JsRawColumnarData, TransformDataRequest};
 use cubesql::compile::engine::df::scan::{ColumnarValueObject, FieldValue};
@@ -128,18 +129,22 @@ impl ResultWrapper {
     }
 }
 
-fn db_primitive_to_field_value(value: &DBResponsePrimitive) -> FieldValue<'_> {
+// String cells borrow their bytes from the backing column (a materialized
+// primitive or an Arrow buffer) — no per-cell allocation on the way to the
+// wire encoder.
+fn db_value_to_field_value(value: DBResponseValueRef<'_>) -> FieldValue<'_> {
     match value {
-        DBResponsePrimitive::String(s) => FieldValue::String(Cow::Borrowed(s)),
-        DBResponsePrimitive::Int64(n) => FieldValue::Number(*n as f64),
-        DBResponsePrimitive::UInt64(n) => FieldValue::Number(*n as f64),
-        DBResponsePrimitive::Float64(n) => FieldValue::Number(*n),
-        DBResponsePrimitive::Boolean(b) => FieldValue::Bool(*b),
-        DBResponsePrimitive::Timestamp(_) => FieldValue::String(Cow::Owned(value.to_string())),
-        DBResponsePrimitive::Uncommon(v) => FieldValue::String(Cow::Owned(
-            serde_json::to_string(&v).unwrap_or_else(|_| v.to_string()),
+        DBResponseValueRef::Str(s) => FieldValue::String(Cow::Borrowed(s)),
+        DBResponseValueRef::StringOwned(s) => FieldValue::String(Cow::Owned(s)),
+        DBResponseValueRef::Int64(n) => FieldValue::Number(n as f64),
+        DBResponseValueRef::UInt64(n) => FieldValue::Number(n as f64),
+        DBResponseValueRef::Float64(n) => FieldValue::Number(n),
+        DBResponseValueRef::Boolean(b) => FieldValue::Bool(b),
+        DBResponseValueRef::Timestamp(_) => FieldValue::String(Cow::Owned(value.to_string())),
+        DBResponseValueRef::Uncommon(v) => FieldValue::String(Cow::Owned(
+            serde_json::to_string(v).unwrap_or_else(|_| v.to_string()),
         )),
-        DBResponsePrimitive::Null => FieldValue::Null,
+        DBResponseValueRef::Null => FieldValue::Null,
     }
 }
 
@@ -189,7 +194,7 @@ impl ColumnarValueObject for ResultWrapper {
         };
 
         Ok(Box::new(
-            column.iter().map(|v| Ok(db_primitive_to_field_value(v))),
+            column.iter_values().map(|v| Ok(db_value_to_field_value(v))),
         ))
     }
 }
@@ -257,8 +262,11 @@ pub fn get_cubestore_result(mut cx: FunctionContext) -> JsResult<JsValue> {
             result.members().iter().map(|k| cx.string(k)).collect();
 
         let row_count = result.row_count();
-        let columns: Vec<_> = (0..js_keys.len())
-            .map(|i| result.column(i))
+        // Collect borrowed per-cell views up front: the row-major loop below
+        // needs random access, and value refs are cheap (string cells borrow
+        // from the backing column instead of copying).
+        let columns: Vec<Vec<DBResponseValueRef<'_>>> = (0..js_keys.len())
+            .map(|i| result.column(i).map(|col| col.iter_values().collect()))
             .collect::<Result<_, _>>()
             .or_else(|err| cx.throw_error(err.to_string()))?;
         let js_array = JsArray::new(&mut cx, row_count);
@@ -270,8 +278,9 @@ pub fn get_cubestore_result(mut cx: FunctionContext) -> JsResult<JsValue> {
                 for (col_idx, js_key) in js_keys.iter().enumerate() {
                     let value = &columns[col_idx][row_idx];
                     let js_value: Handle<'_, JsValue> = match value {
-                        DBResponsePrimitive::Null => cx.null().upcast(),
+                        DBResponseValueRef::Null => cx.null().upcast(),
                         // For compatibility, we convert all primitives to strings
+                        DBResponseValueRef::Str(s) => cx.string(s).upcast(),
                         other => cx.string(other.to_string()).upcast(),
                     };
 

packages/cubejs-base-driver/src/ColumnarResponse.ts

@@ -0,0 +1,75 @@
+import * as stream from 'stream';
+
+/**
+ * Marker set on every batch object emitted by `ColumnarResponse`. Downstream consumers
+ * (the query orchestrator's `QueryStream` and `cubejs-backend-native`) key on this to tell
+ * a columnar batch apart from a plain row object.
+ */
+export const COLUMNAR_RESPONSE_TYPE = 'ColumnarResponse';
+
+/**
+ * A single columnar batch as it flows through the object-mode stream.
+ * `members[j]` names `columns[j]`; every `columns[j]` has the same length (the batch's
+ * row count). Matches the `{ members, columns }` shape decoded on the Rust side
+ * (`JsonColumnarValueObject`).
+ */
+export interface ColumnarBatch {
+  $type: typeof COLUMNAR_RESPONSE_TYPE;
+  members: string[];
+  columns: any[][];
+}
+
+/**
+ * Object-mode Transform that pivots position-indexed array rows (e.g. pg `rowMode: 'array'`)
+ * into columnar batches. Buffering by column avoids allocating a per-row object per row and
+ * moves the row→columnar pivot out of `cubejs-backend-native`, so the SQL-API streaming path
+ * can hand columnar batches straight to Rust.
+ *
+ * Input: `unknown[]` rows whose cells are ordered to match `members`.
+ * Output: `ColumnarBatch` objects, one every `batchSize` rows plus a final partial batch.
+ */
+export class ColumnarResponse extends stream.Transform {
+  private readonly members: string[];
+
+  private readonly batchSize: number;
+
+  private columns: any[][];
+
+  private count = 0;
+
+  public constructor({ members, batchSize }: { members: string[]; batchSize: number }) {
+    super({ objectMode: true });
+    this.members = members;
+    this.batchSize = batchSize;
+    this.columns = members.map(() => []);
+  }
+
+  public _transform(row: unknown[], _encoding: BufferEncoding, callback: stream.TransformCallback) {
+    for (let j = 0; j < this.members.length; j++) {
+      this.columns[j].push(row[j]);
+    }
+    this.count++;
+    if (this.count >= this.batchSize) {
+      this.flushBatch();
+    }
+    callback();
+  }
+
+  public _flush(callback: stream.TransformCallback) {
+    if (this.count > 0) {
+      this.flushBatch();
+    }
+    callback();
+  }
+
+  private flushBatch() {
+    const batch: ColumnarBatch = {
+      $type: COLUMNAR_RESPONSE_TYPE,
+      members: this.members,
+      columns: this.columns,
+    };
+    this.push(batch);
+    this.columns = this.members.map(() => []);
+    this.count = 0;
+  }
+}

packages/cubejs-base-driver/src/driver.interface.ts

@@ -131,6 +131,12 @@ export interface DriverCapabilities extends ExternalDriverCompatibilities {
 
 export type StreamOptions = {
   highWaterMark: number;
+  /**
+   * When set, the driver is asked to emit columnar batches (see `ColumnarResponse`)
+   * instead of row objects. It is a hint: drivers that don't support it keep emitting
+   * rows. Currently set to `true` only for the cubesql SQL-API streaming path.
+   */
+  columnar?: boolean;
 };
 
 export type StreamingSourceOptions = {

packages/cubejs-base-driver/src/index.ts

@@ -1,4 +1,5 @@
 export * from './BaseDriver';
+export * from './ColumnarResponse';
 export * from './utils';
 export * from './driver.interface';
 export * from './queue-driver.interface';

packages/cubejs-base-driver/test/unit/ColumnarResponse.test.ts

@@ -0,0 +1,77 @@
+import { Readable } from 'stream';
+import { ColumnarResponse, COLUMNAR_RESPONSE_TYPE } from '../../src';
+
+async function collect(members: string[], batchSize: number, rows: unknown[][]) {
+  const source = Readable.from(rows, { objectMode: true });
+  const columnar = new ColumnarResponse({ members, batchSize });
+  const batches: any[] = [];
+  source.pipe(columnar);
+  for await (const batch of columnar) {
+    batches.push(batch);
+  }
+  return batches;
+}
+
+describe('ColumnarResponse', () => {
+  const members = ['a', 'b'];
+
+  test('emits a batch at every batchSize boundary and a final partial batch', async () => {
+    const rows = [
+      [1, 'x'],
+      [2, 'y'],
+      [3, 'z'],
+      [4, 'w'],
+      [5, 'v'],
+    ];
+
+    const batches = await collect(members, 2, rows);
+
+    expect(batches).toHaveLength(3); // 2 + 2 + 1
+    batches.forEach((batch) => {
+      expect(batch.$type).toBe(COLUMNAR_RESPONSE_TYPE);
+      expect(batch.members).toEqual(members);
+    });
+
+    expect(batches[0].columns).toEqual([[1, 2], ['x', 'y']]);
+    expect(batches[1].columns).toEqual([[3, 4], ['z', 'w']]);
+    expect(batches[2].columns).toEqual([[5], ['v']]);
+  });
+
+  test('transposes positional array rows into columns', async () => {
+    const rows = [
+      [1, 'x'],
+      [2, 'y'],
+    ];
+
+    const batches = await collect(members, 100, rows);
+
+    expect(batches).toHaveLength(1);
+    expect(batches[0]).toEqual({
+      $type: COLUMNAR_RESPONSE_TYPE,
+      members,
+      columns: [[1, 2], ['x', 'y']],
+    });
+  });
+
+  test('does not emit anything for an empty stream', async () => {
+    const batches = await collect(members, 2, []);
+    expect(batches).toHaveLength(0);
+  });
+
+  test('reuses fresh column arrays per batch (no shared references)', async () => {
+    const rows = [
+      [1, 'x'],
+      [2, 'y'],
+      [3, 'z'],
+      [4, 'w'],
+    ];
+
+    const batches = await collect(members, 2, rows);
+
+    expect(batches).toHaveLength(2);
+    // First batch's columns must not be mutated by the second batch.
+    expect(batches[0].columns).toEqual([[1, 2], ['x', 'y']]);
+    expect(batches[1].columns).toEqual([[3, 4], ['z', 'w']]);
+    expect(batches[0].columns).not.toBe(batches[1].columns);
+  });
+});

packages/cubejs-postgres-driver/src/PostgresDriver.ts

@@ -10,6 +10,7 @@ import { types, FieldDef } from 'pg';
 import { TypeId, TypeFormat } from 'pg-types';
 import {
   BaseDriver,
+  ColumnarResponse,
   DownloadQueryResultsOptions, DownloadTableMemoryData, DriverInterface,
   GenericDataBaseType, IndexesSQL, TableStructure, StreamOptions,
   StreamTableDataWithTypes, QueryOptions, DownloadQueryResultsResult, DriverCapabilities, TableColumn, createPoolName,
@@ -353,7 +354,7 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
   public async stream(
     query: string,
     values: unknown[],
-    { highWaterMark }: StreamOptions
+    { highWaterMark, columnar }: StreamOptions
   ): Promise<StreamTableDataWithTypes> {
     PostgresDriver.checkValuesLimit(values);
 
@@ -366,17 +367,38 @@ export class PostgresDriver<Config extends PostgresDriverConfiguration = Postgre
         types: {
           getTypeParser: this.getTypeParser,
         },
-        highWaterMark
-      });
+        highWaterMark,
+        // In columnar mode read rows as positional arrays instead of per-row objects, so
+        // `ColumnarResponse` can bucket cells into columns without any object allocation.
+        // @types/pg-query-stream (1.0.3) does not declare `rowMode`, hence the cast.
+        ...(columnar ? { rowMode: 'array' } : {}),
+      } as any);
       const rowStream: QueryStream = await conn.query(queryStream);
       const fields = await rowStream.fields();
+      const mappedTypes = this.mapFields(fields);
+      const release = async () => {
+        await this.pool.release(conn);
+      };
+
+      if (columnar) {
+        const members = fields.map((f) => f.name);
+        const columnarStream = new ColumnarResponse({ members, batchSize: highWaterMark });
+        // `pipe()` does not forward errors — propagate manually so the consumer sees the
+        // failure and `release()` still runs during cleanup.
+        rowStream.once('error', (e) => columnarStream.destroy(e));
+        rowStream.pipe(columnarStream);
+
+        return {
+          rowStream: columnarStream,
+          types: mappedTypes,
+          release,
+        };
+      }
 
       return {
         rowStream,
-        types: this.mapFields(fields),
-        release: async () => {
-          await this.pool.release(conn);
-        }
+        types: mappedTypes,
+        release,
       };
     } catch (e) {
       await this.pool.release(conn);

packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts

@@ -663,7 +663,10 @@ export class QueryCache {
           let logged = false;
           Promise
             .all([clientFactory()])
-            .then(([client]) => (<DriverInterface>client).stream(req.query, req.values, { highWaterMark: getEnv('dbQueryStreamHighWaterMark') }))
+            // This handler is the cubesql SQL-API streaming path, which consumes columnar
+            // batches natively — ask the driver to emit them (drivers that don't support it
+            // fall back to row objects).
+            .then(([client]) => (<DriverInterface>client).stream(req.query, req.values, { highWaterMark: getEnv('dbQueryStreamHighWaterMark'), columnar: true }))
             .then((source) => {
               const cleanup = async (error) => {
                 if (source.release) {

packages/cubejs-query-orchestrator/src/orchestrator/QueryStream.ts

@@ -1,5 +1,6 @@
 import * as stream from 'stream';
 import { getEnv } from '@cubejs-backend/shared';
+import { COLUMNAR_RESPONSE_TYPE } from '@cubejs-backend/base-driver';
 
 export class QueryStream extends stream.Transform {
   private timeout = 5 * 60 * 1000;
@@ -43,6 +44,24 @@ export class QueryStream extends stream.Transform {
     if (this.streams.has(this.queryKey)) {
       this.streams.delete(this.queryKey);
     }
+
+    // Columnar batch (see `ColumnarResponse` in base-driver): remap the `members` array
+    // (aliases → member names) once and forward the batch untouched.
+    if (chunk && chunk.$type === COLUMNAR_RESPONSE_TYPE) {
+      const members = this.aliasNameToMember
+        ? chunk.members.map((alias) => this.aliasNameToMember[alias])
+        : chunk.members;
+      const rowCount = chunk.columns[0] ? chunk.columns[0].length : 0;
+      if (this.counter + rowCount < this.readableHighWaterMark) {
+        this.counter += rowCount;
+      } else {
+        this.counter = 0;
+        this.debounce();
+      }
+      callback(null, { $type: COLUMNAR_RESPONSE_TYPE, members, columns: chunk.columns });
+      return;
+    }
+
     let row = {};
     if (this.aliasNameToMember) {
       Object.keys(chunk).forEach((alias) => {