From 08fc3891252be648945436d5346c7cc594896c47 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Fri, 10 Jul 2026 13:50:59 +0200 Subject: [PATCH 1/4] perf(cubesql): Migrate streaming to columnar JSON batches The SQL API streaming path shipped rows across the JS->Rust native bridge as an array of row objects and rebuilt each `RecordBatch` cell-by-cell through per-cell Neon/NAPI downcasts (`JsValueObject::get`). This mirrors the row-oriented transport the `/load` path already replaced with a columnar JSON `Buffer`. --- packages/cubejs-backend-native/js/index.ts | 6 +- packages/cubejs-backend-native/src/stream.rs | 118 +++--------------- .../cubesql/src/compile/engine/df/scan.rs | 3 +- 3 files changed, 25 insertions(+), 102 deletions(-) diff --git a/packages/cubejs-backend-native/js/index.ts b/packages/cubejs-backend-native/js/index.ts index 24715c926b0e0..32a37a5ff9ced 100644 --- a/packages/cubejs-backend-native/js/index.ts +++ b/packages/cubejs-backend-native/js/index.ts @@ -4,7 +4,7 @@ import path from 'path'; import { Writable } from 'stream'; import type { Request as ExpressRequest } from 'express'; import { CacheMode } from '@cubejs-backend/shared'; -import { NativeQueryResultRef, ResultWrapper } from './ResultWrapper'; +import { NativeQueryResultRef, ResultWrapper, rowsToColumnar } from './ResultWrapper'; export * from './ResultWrapper'; @@ -294,7 +294,7 @@ function wrapNativeFunctionWithStream( if (chunkBuffer.length < chunkLength) { callback(null); } else { - const toSend = chunkBuffer; + const toSend = Buffer.from(JSON.stringify(rowsToColumnar(chunkBuffer))); chunkBuffer = []; writerOrChannel.chunk(toSend, callback); } @@ -308,7 +308,7 @@ function wrapNativeFunctionWithStream( } }; if (chunkBuffer.length > 0) { - const toSend = chunkBuffer; + const toSend = Buffer.from(JSON.stringify(rowsToColumnar(chunkBuffer))); chunkBuffer = []; writerOrChannel.chunk(toSend, end); } else { diff --git a/packages/cubejs-backend-native/src/stream.rs b/packages/cubejs-backend-native/src/stream.rs index 9bc270ed1d01a..84b4a06c272c5 100644 --- a/packages/cubejs-backend-native/src/stream.rs +++ b/packages/cubejs-backend-native/src/stream.rs @@ -1,7 +1,6 @@ use cubesql::compile::engine::df::scan::{ - transform_response, FieldValue, MemberField, RecordBatch, SchemaRef, ValueObject, + transform_columnar_response, JsonColumnarValueObject, MemberField, RecordBatch, SchemaRef, }; -use std::borrow::Cow; use std::cell::RefCell; use std::future::Future; @@ -15,13 +14,12 @@ use crate::channel::call_js_fn; use cubesql::CubeError; use neon::prelude::*; +use neon::types::buffer::TypedArray; use tokio::sync::{oneshot, Semaphore}; #[cfg(feature = "neon-debug")] use log::trace; -use neon::types::JsDate; - use crate::utils::bind_method; use tokio::sync::mpsc::{channel as mpsc_channel, Receiver, Sender}; @@ -256,83 +254,6 @@ fn wait_for_future_and_execute_callback( }); } -pub struct JsValueObject<'a> { - pub cx: FunctionContext<'a>, - pub handle: Handle<'a, JsArray>, -} - -fn js_value_to_json_string<'a, C: Context<'a>>( - cx: &mut C, - value: Handle<'a, JsValue>, -) -> Result { - let global = cx.global_object(); - let json = global - .get::(cx, "JSON") - .map_err(|e| CubeError::internal(format!("Can't get JSON global: {}", e)))?; - let stringify = json - .get::(cx, "stringify") - .map_err(|e| CubeError::internal(format!("Can't get JSON.stringify: {}", e)))?; - let undefined = cx.undefined().upcast::(); - let result = stringify - .call(cx, undefined, [value]) - .map_err(|e| CubeError::internal(format!("JSON.stringify failed: {}", e)))?; - let s = result.downcast::(cx).map_err(|e| { - CubeError::internal(format!("JSON.stringify did not return a string: {}", e)) - })?; - Ok(s.value(cx)) -} - -impl ValueObject for JsValueObject<'_> { - fn len(&mut self) -> Result { - Ok(self.handle.len(&mut self.cx) as usize) - } - - fn get(&mut self, index: usize, field_name: &str) -> Result, CubeError> { - let value = self - .handle - .get::(&mut self.cx, index as u32) - .map_err(|e| { - CubeError::internal(format!("Can't get object at array index {}: {}", index, e)) - })? - .get::(&mut self.cx, field_name) - .map_err(|e| { - CubeError::internal(format!("Can't get '{}' field value: {}", field_name, e)) - })?; - if let Ok(s) = value.downcast::(&mut self.cx) { - Ok(FieldValue::String(Cow::Owned(s.value(&mut self.cx)))) - } else if let Ok(n) = value.downcast::(&mut self.cx) { - Ok(FieldValue::Number(n.value(&mut self.cx))) - } else if let Ok(b) = value.downcast::(&mut self.cx) { - Ok(FieldValue::Bool(b.value(&mut self.cx))) - } else if value.downcast::(&mut self.cx).is_ok() - || value.downcast::(&mut self.cx).is_ok() - { - Ok(FieldValue::Null) - } else if value.is_a::(&mut self.cx) { - Ok(FieldValue::String(Cow::Owned(js_value_to_json_string( - &mut self.cx, - value, - )?))) - } else if let Ok(b) = value.downcast::(&mut self.cx) { - // TODO: Support it? - Err(CubeError::internal(format!( - "Expected primitive value but found JsDate({:?})", - b - ))) - } else if value.is_a::(&mut self.cx) { - Ok(FieldValue::String(Cow::Owned(js_value_to_json_string( - &mut self.cx, - value, - )?))) - } else { - Err(CubeError::internal(format!( - "Expected primitive value but found: {:?}", - value - ))) - } - } -} - fn js_stream_push_chunk(mut cx: FunctionContext) -> JsResult { #[cfg(feature = "neon-debug")] trace!("JsWriteStream.push_chunk"); @@ -340,26 +261,27 @@ fn js_stream_push_chunk(mut cx: FunctionContext) -> JsResult { let this = cx .this::()? .downcast_or_throw::, _>(&mut cx)?; - let chunk_array = cx.argument::(0)?; let callback = cx.argument::(1)?.root(&mut cx); - let mut value_object = JsValueObject { - cx, - handle: chunk_array, - }; - let value = - match transform_response(&mut value_object, this.schema.clone(), &this.member_fields) { - Ok(value) => value, - Err(e) => return value_object.cx.throw_error(e.message), + + let chunk_buffer = cx.argument::(0)?; + let mut value_object = + match serde_json::from_slice::(chunk_buffer.as_slice(&cx)) { + Ok(v) => v, + Err(e) => return cx.throw_error(format!("Can't parse columnar chunk JSON: {}", e)), }; + let value = match transform_columnar_response( + &mut value_object, + this.schema.clone(), + &this.member_fields, + ) { + Ok(value) => value, + Err(e) => return cx.throw_error(e.message), + }; + let future = this.push_chunk(value); - wait_for_future_and_execute_callback( - this.tokio_handle.clone(), - value_object.cx.channel(), - callback, - future, - ); - - Ok(value_object.cx.undefined()) + wait_for_future_and_execute_callback(this.tokio_handle.clone(), cx.channel(), callback, future); + + Ok(cx.undefined()) } fn js_stream_start(mut cx: FunctionContext) -> JsResult { diff --git a/rust/cubesql/cubesql/src/compile/engine/df/scan.rs b/rust/cubesql/cubesql/src/compile/engine/df/scan.rs index 5a498b4fd220b..226b7161af990 100644 --- a/rust/cubesql/cubesql/src/compile/engine/df/scan.rs +++ b/rust/cubesql/cubesql/src/compile/engine/df/scan.rs @@ -45,7 +45,7 @@ use datafusion::{ }; use futures::Stream; use log::warn; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use serde_json::Value; use std::str::FromStr; use std::{ @@ -380,6 +380,7 @@ impl ValueObject for JsonValueObject { } } +#[derive(Deserialize)] pub struct JsonColumnarValueObject { members: Vec, columns: Vec>, From 4d765e76b4fa43d7ec6fca728572cff0eef328c0 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Fri, 10 Jul 2026 14:24:33 +0200 Subject: [PATCH 2/4] refactor(native): Extract rowsToColumnarBuffer for columnar wire encoding --- .../cubejs-backend-native/js/ResultWrapper.ts | 32 +++++++++++-------- packages/cubejs-backend-native/js/index.ts | 6 ++-- 2 files changed, 22 insertions(+), 16 deletions(-) diff --git a/packages/cubejs-backend-native/js/ResultWrapper.ts b/packages/cubejs-backend-native/js/ResultWrapper.ts index 8cfd378118b2e..b5a1b5d15d417 100644 --- a/packages/cubejs-backend-native/js/ResultWrapper.ts +++ b/packages/cubejs-backend-native/js/ResultWrapper.ts @@ -54,6 +54,24 @@ export function rowsToColumnar(rawData: any): JsRawColumnarData { return { members, columns }; } +/** + * Pivot to columnar before serializing: the row-oriented form repeats + * every column name on every row, which inflates JSON size and forces + * the Rust side to allocate a per-row map before transposing back to + * its native columnar `QueryResult` representation. + * + * Serialize to a Buffer so the Rust side can decode via + * serde_json::from_slice instead of walking a JsValue through the + * Neon bridge with JsValueDeserializer. On 5 MB of AoO rows + * (~21k rows × 8 fields) the JsValue walk costs ~80 ms locally; + * Buffer + serde_json is ~7× faster (M3 MAX) and tracks V8's JSON.parse + * (~11 ms on the same payload). On a real server it should be 3-6× slower, + * so avoiding the JsValue walk matters even more there. + */ +export function rowsToColumnarBuffer(rawData: any): Buffer { + return Buffer.from(JSON.stringify(rowsToColumnar(rawData))); +} + class BaseWrapper { public readonly isWrapper: boolean = true; } @@ -191,19 +209,7 @@ export class ResultWrapper extends BaseWrapper implements DataResult { return [this[NATIVE_REFERENCE]]; } - // Pivot to columnar before serializing: the row-oriented form repeats - // every column name on every row, which inflates JSON size and forces - // the Rust side to allocate a per-row map before transposing back to - // its native columnar `QueryResult` representation. - // - // Serialize to a Buffer so the Rust side can decode via - // serde_json::from_slice instead of walking a JsValue through the - // Neon bridge with JsValueDeserializer. On 5 MB of AoO rows - // (~21k rows × 8 fields) the JsValue walk costs ~80 ms locally; - // Buffer + serde_json is ~7× faster (M3 MAX) and tracks V8's JSON.parse - // (~11 ms on the same payload). On a real server it should be 3-6× slower, - // so avoiding the JsValue walk matters even more there. - return [Buffer.from(JSON.stringify(rowsToColumnar(this.jsResult)))]; + return [rowsToColumnarBuffer(this.jsResult)]; } public setTransformData(td: any) { diff --git a/packages/cubejs-backend-native/js/index.ts b/packages/cubejs-backend-native/js/index.ts index 32a37a5ff9ced..9a04a5bf6ccfa 100644 --- a/packages/cubejs-backend-native/js/index.ts +++ b/packages/cubejs-backend-native/js/index.ts @@ -4,7 +4,7 @@ import path from 'path'; import { Writable } from 'stream'; import type { Request as ExpressRequest } from 'express'; import { CacheMode } from '@cubejs-backend/shared'; -import { NativeQueryResultRef, ResultWrapper, rowsToColumnar } from './ResultWrapper'; +import { NativeQueryResultRef, ResultWrapper, rowsToColumnarBuffer } from './ResultWrapper'; export * from './ResultWrapper'; @@ -294,7 +294,7 @@ function wrapNativeFunctionWithStream( if (chunkBuffer.length < chunkLength) { callback(null); } else { - const toSend = Buffer.from(JSON.stringify(rowsToColumnar(chunkBuffer))); + const toSend = rowsToColumnarBuffer(chunkBuffer); chunkBuffer = []; writerOrChannel.chunk(toSend, callback); } @@ -308,7 +308,7 @@ function wrapNativeFunctionWithStream( } }; if (chunkBuffer.length > 0) { - const toSend = Buffer.from(JSON.stringify(rowsToColumnar(chunkBuffer))); + const toSend = rowsToColumnarBuffer(chunkBuffer); chunkBuffer = []; writerOrChannel.chunk(toSend, end); } else { From f53cc43db7bef86bf528e15eceb3b8eb27c5a717 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Fri, 10 Jul 2026 15:31:59 +0200 Subject: [PATCH 3/4] refactor: move to columnar --- .../cubejs-backend-native/src/orchestrator.rs | 88 +------ packages/cubejs-backend-native/src/stream.rs | 15 +- .../cubejs-backend-native/src/transport.rs | 16 +- .../cubesql/benches/transform_response.rs | 59 +---- .../cubesql/src/compile/engine/df/scan.rs | 220 ++++++------------ rust/cubesql/cubesql/src/compile/mod.rs | 28 ++- rust/cubesql/cubesql/src/compile/test/mod.rs | 8 +- rust/cubesql/cubesql/src/transport/mod.rs | 1 - rust/cubesql/cubesql/src/transport/service.rs | 4 + 9 files changed, 122 insertions(+), 317 deletions(-) diff --git a/packages/cubejs-backend-native/src/orchestrator.rs b/packages/cubejs-backend-native/src/orchestrator.rs index cce0910fed1c5..4524a7e502b62 100644 --- a/packages/cubejs-backend-native/src/orchestrator.rs +++ b/packages/cubejs-backend-native/src/orchestrator.rs @@ -2,11 +2,10 @@ use crate::node_obj_deserializer::JsValueDeserializer; use crate::transport::MapCubeErrExt; use cubeorchestrator::query_message_parser::QueryResult; use cubeorchestrator::query_result_transform::{ - DBResponsePrimitive, InternedKeyLookup, RequestResultData, RequestResultDataMulti, - TransformedData, + DBResponsePrimitive, RequestResultData, RequestResultDataMulti, TransformedData, }; use cubeorchestrator::transport::{JsRawColumnarData, TransformDataRequest}; -use cubesql::compile::engine::df::scan::{ColumnarValueObject, FieldValue, ValueObject}; +use cubesql::compile::engine::df::scan::{ColumnarValueObject, FieldValue}; use cubesql::CubeError; use neon::context::{Context, FunctionContext, ModuleContext}; use neon::handle::Handle; @@ -144,89 +143,6 @@ fn db_primitive_to_field_value(value: &DBResponsePrimitive) -> FieldValue<'_> { } } -impl ValueObject for ResultWrapper { - fn len(&mut self) -> Result { - if self.transformed_data.is_none() { - self.transform_result()?; - } - - let data = self.transformed_data.as_ref().unwrap(); - - match data { - TransformedData::Compact { - members: _members, - dataset, - } => Ok(dataset.len()), - TransformedData::Columnar { - members: _members, - columns, - } => Ok(columns.first().map(|c| c.len()).unwrap_or(0)), - TransformedData::Vanilla(dataset) => Ok(dataset.len()), - } - } - - fn get(&mut self, index: usize, field_name: &str) -> Result, CubeError> { - if self.transformed_data.is_none() { - self.transform_result()?; - } - - let data = self.transformed_data.as_ref().unwrap(); - - let value = match data { - TransformedData::Compact { members, dataset } => { - let Some(row) = dataset.get(index) else { - return Err(CubeError::internal(format!( - "Unexpected response from Cube, can't get {} row", - index - ))); - }; - - let Some(member_index) = members.iter().position(|m| m == field_name) else { - // Missing field → NULL, matching `Vanilla` semantics below. - return Ok(FieldValue::Null); - }; - - row.get(member_index).unwrap_or(&DBResponsePrimitive::Null) - } - TransformedData::Columnar { members, columns } => { - let Some(member_index) = members.iter().position(|m| m == field_name) else { - // Missing field → NULL, matching `Vanilla` semantics below. - return Ok(FieldValue::Null); - }; - - let Some(column) = columns.get(member_index) else { - return Err(CubeError::internal(format!( - "Unexpected response from Cube, missing column for '{}'", - field_name - ))); - }; - - let Some(value) = column.get(index) else { - return Err(CubeError::user(format!( - "Unexpected response from Cube, can't get {} row", - index - ))); - }; - - value - } - TransformedData::Vanilla(dataset) => { - let Some(row) = dataset.get(index) else { - return Err(CubeError::internal(format!( - "Unexpected response from Cube, can't get {} row", - index - ))); - }; - - row.get(&InternedKeyLookup::new(field_name)) - .unwrap_or(&DBResponsePrimitive::Null) - } - }; - - Ok(db_primitive_to_field_value(value)) - } -} - impl ColumnarValueObject for ResultWrapper { fn len(&mut self) -> Result { if self.transformed_data.is_none() { diff --git a/packages/cubejs-backend-native/src/stream.rs b/packages/cubejs-backend-native/src/stream.rs index 84b4a06c272c5..38b907ed84f6f 100644 --- a/packages/cubejs-backend-native/src/stream.rs +++ b/packages/cubejs-backend-native/src/stream.rs @@ -1,5 +1,5 @@ use cubesql::compile::engine::df::scan::{ - transform_columnar_response, JsonColumnarValueObject, MemberField, RecordBatch, SchemaRef, + transform_response, JsonColumnarValueObject, MemberField, RecordBatch, SchemaRef, }; use std::cell::RefCell; @@ -269,14 +269,11 @@ fn js_stream_push_chunk(mut cx: FunctionContext) -> JsResult { Ok(v) => v, Err(e) => return cx.throw_error(format!("Can't parse columnar chunk JSON: {}", e)), }; - let value = match transform_columnar_response( - &mut value_object, - this.schema.clone(), - &this.member_fields, - ) { - Ok(value) => value, - Err(e) => return cx.throw_error(e.message), - }; + let value = + match transform_response(&mut value_object, this.schema.clone(), &this.member_fields) { + Ok(value) => value, + Err(e) => return cx.throw_error(e.message), + }; let future = this.push_chunk(value); wait_for_future_and_execute_callback(this.tokio_handle.clone(), cx.channel(), callback, future); diff --git a/packages/cubejs-backend-native/src/transport.rs b/packages/cubejs-backend-native/src/transport.rs index be4f17623b66e..f6bca69f71deb 100644 --- a/packages/cubejs-backend-native/src/transport.rs +++ b/packages/cubejs-backend-native/src/transport.rs @@ -15,8 +15,8 @@ use crate::{ use async_trait::async_trait; use cubeorchestrator::query_result_transform::RequestResultData; use cubesql::compile::engine::df::scan::{ - build_response_schema, convert_transport_response_columnar, transform_columnar_response, - CacheMode, MemberField, RecordBatch, SchemaRef, + build_response_schema, convert_transport_response, transform_response, CacheMode, MemberField, + RecordBatch, SchemaRef, }; use cubesql::compile::engine::df::wrapper::SqlQuery; use cubesql::transport::{ @@ -524,11 +524,7 @@ impl TransportService for NodeBridgeTransport { } }; - break convert_transport_response_columnar( - response, - schema.clone(), - member_fields, - ); + break convert_transport_response(response, schema.clone(), member_fields); } ValueFromJs::ResultWrapper(result_wrappers) => { break result_wrappers @@ -540,11 +536,7 @@ impl TransportService for NodeBridgeTransport { wrapper.external, ); - transform_columnar_response( - &mut wrapper, - updated_schema, - &member_fields, - ) + transform_response(&mut wrapper, updated_schema, &member_fields) }) .collect::, _>>(); } diff --git a/rust/cubesql/cubesql/benches/transform_response.rs b/rust/cubesql/cubesql/benches/transform_response.rs index a03e1a01e7311..0baba390ff3dc 100644 --- a/rust/cubesql/cubesql/benches/transform_response.rs +++ b/rust/cubesql/cubesql/benches/transform_response.rs @@ -2,10 +2,9 @@ use std::sync::Arc; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; use cubesql::compile::engine::df::scan::{ - convert_transport_response, convert_transport_response_columnar, DataType, MemberField, Schema, - SchemaRef, + convert_transport_response, DataType, MemberField, Schema, SchemaRef, }; -use cubesql::transport::{TransportLoadResponse, TransportLoadResponseColumnar}; +use cubesql::transport::TransportLoadResponseColumnar; use datafusion::arrow::datatypes::{Field, TimeUnit}; use serde_json::json; @@ -113,33 +112,6 @@ fn annotation_value() -> serde_json::Value { }) } -fn build_row_json(rows: usize, kinds: &[ColKind]) -> String { - let names: Vec = kinds - .iter() - .enumerate() - .map(|(i, k)| field_name(i, *k)) - .collect(); - - let data: Vec = (0..rows) - .map(|r| { - let mut row_obj = serde_json::Map::with_capacity(kinds.len()); - for (c, kind) in kinds.iter().enumerate() { - row_obj.insert(names[c].clone(), cell_value(r, c, *kind)); - } - serde_json::Value::Object(row_obj) - }) - .collect(); - - let response = json!({ - "results": [{ - "annotation": annotation_value(), - "data": data, - }] - }); - - serde_json::to_string(&response).expect("serialize row json") -} - fn build_columnar_json(rows: usize, kinds: &[ColKind]) -> String { let names: Vec = kinds .iter() @@ -169,7 +141,6 @@ fn build_columnar_json(rows: usize, kinds: &[ColKind]) -> String { struct Inputs { schema: SchemaRef, member_fields: Vec, - row_json: String, columnar_json: String, } @@ -178,7 +149,6 @@ fn build_inputs(rows: usize, cols: usize, time_dims: usize) -> Inputs { Inputs { schema: build_schema(&kinds), member_fields: build_member_fields(&kinds), - row_json: build_row_json(rows, &kinds), columnar_json: build_columnar_json(rows, &kinds), } } @@ -212,30 +182,9 @@ fn bench_transform_response(c: &mut Criterion) { let Inputs { schema, member_fields, - row_json, columnar_json, } = &inputs; - let row_id = format!("row/rows={}/cols={}/td={}", rows, cols, td); - group.bench_with_input( - BenchmarkId::from_parameter(&row_id), - row_json.as_str(), - |b, json| { - b.iter(|| { - let value: serde_json::Value = - serde_json::from_str(json).expect("row from_str"); - let response: TransportLoadResponse = - serde_json::from_value(value).expect("row from_value"); - convert_transport_response( - response, - schema.clone(), - member_fields.clone(), - ) - .expect("convert_transport_response") - }) - }, - ); - let col_id = format!("columnar/rows={}/cols={}/td={}", rows, cols, td); group.bench_with_input( BenchmarkId::from_parameter(&col_id), @@ -246,12 +195,12 @@ fn bench_transform_response(c: &mut Criterion) { serde_json::from_str(json).expect("columnar from_str"); let response: TransportLoadResponseColumnar = serde_json::from_value(value).expect("columnar from_value"); - convert_transport_response_columnar( + convert_transport_response( response, schema.clone(), member_fields.clone(), ) - .expect("convert_transport_response_columnar") + .expect("convert_transport_response") }) }, ); diff --git a/rust/cubesql/cubesql/src/compile/engine/df/scan.rs b/rust/cubesql/cubesql/src/compile/engine/df/scan.rs index 226b7161af990..e75c1dbf11db2 100644 --- a/rust/cubesql/cubesql/src/compile/engine/df/scan.rs +++ b/rust/cubesql/cubesql/src/compile/engine/df/scan.rs @@ -310,16 +310,6 @@ pub enum FieldValue<'a> { Null, } -pub trait ValueObject { - fn len(&mut self) -> std::result::Result; - - fn get( - &mut self, - index: usize, - field_name: &str, - ) -> std::result::Result, CubeError>; -} - pub trait ColumnarValueObject { fn len(&mut self) -> std::result::Result; @@ -332,16 +322,6 @@ pub trait ColumnarValueObject { >; } -pub struct JsonValueObject { - rows: Vec, -} - -impl JsonValueObject { - pub fn new(rows: Vec) -> Self { - JsonValueObject { rows } - } -} - fn json_value_to_field_value(value: &Value) -> std::result::Result, CubeError> { Ok(match value { Value::String(s) => FieldValue::String(Cow::Borrowed(s)), @@ -358,28 +338,6 @@ fn json_value_to_field_value(value: &Value) -> std::result::Result std::result::Result { - Ok(self.rows.len()) - } - - fn get( - &mut self, - index: usize, - field_name: &str, - ) -> std::result::Result, CubeError> { - let Some(as_object) = self.rows[index].as_object() else { - return Err(CubeError::internal(format!( - "Unexpected response from Cube, row is not an object: {:?}", - self.rows[index] - ))); - }; - - let value = as_object.get(field_name).unwrap_or(&Value::Null); - json_value_to_field_value(value) - } -} - #[derive(Deserialize)] pub struct JsonColumnarValueObject { members: Vec, @@ -410,8 +368,7 @@ impl ColumnarValueObject for JsonColumnarValueObject { CubeError, > { let Some(idx) = self.members.iter().position(|m| m == field_name) else { - // Match the row-mode `JsonValueObject::get` semantics: a missing field - // is treated as a column of NULLs. + // A missing field is treated as a column of NULLs. let len = self.columns.first().map(|c| c.len()).unwrap_or(0); return Ok(Box::new((0..len).map(|_| Ok(FieldValue::Null)))); }; @@ -427,39 +384,46 @@ impl ColumnarValueObject for JsonColumnarValueObject { } } -// `$mode` is one of `row` or `columnar`. The `row` arm calls `$response.get(i, field_name)?` -// per cell; the `columnar` arm fetches the entire column slice once via -// `$response.column(field_name)?` and iterates it. -macro_rules! build_column_iter_loop { - (row, $response:expr, $len:expr, $field_name:expr, $value:ident, $body:block) => {{ - for i in 0..$len { - let $value = $response.get(i, $field_name)?; - $body - } - }}; - (columnar, $response:expr, $len:expr, $field_name:expr, $value:ident, $body:block) => {{ - for cell in $response.column($field_name)? { - let $value = cell?; - $body - } - }}; +/// A columnar value object with no data columns, representing `row_count` rows of a +/// literal-only projection (the `no_members_query` shortcut). Every schema field is a +/// `MemberField::Literal`, so `column()` is never actually invoked — only `len()` matters. +struct LiteralRowsValueObject { + row_count: usize, +} + +impl ColumnarValueObject for LiteralRowsValueObject { + fn len(&mut self) -> std::result::Result { + Ok(self.row_count) + } + + fn column<'a>( + &'a mut self, + _field_name: &str, + ) -> std::result::Result< + Box, CubeError>> + 'a>, + CubeError, + > { + let row_count = self.row_count; + Ok(Box::new((0..row_count).map(|_| Ok(FieldValue::Null)))) + } } macro_rules! build_column { - ($data_type:expr, $builder_ty:ty, $mode:tt, $response:expr, $field_name:expr, { $($builder_block:tt)* }, { $($scalar_block:tt)* }) => {{ + ($data_type:expr, $builder_ty:ty, $response:expr, $field_name:expr, { $($builder_block:tt)* }, { $($scalar_block:tt)* }) => {{ let len = $response.len()?; let mut builder = <$builder_ty>::new(len); - build_column_custom_builder!($data_type, $mode, len, builder, $response, $field_name, { $($builder_block)* }, { $($scalar_block)* }) + build_column_custom_builder!($data_type, len, builder, $response, $field_name, { $($builder_block)* }, { $($scalar_block)* }) }} } macro_rules! build_column_custom_builder { - ($data_type:expr, $mode:tt, $len:expr, $builder:expr, $response:expr, $field_name: expr, { $($builder_block:tt)* }, { $($scalar_block:tt)* }) => {{ + ($data_type:expr, $len:expr, $builder:expr, $response:expr, $field_name: expr, { $($builder_block:tt)* }, { $($scalar_block:tt)* }) => {{ match $field_name { MemberField::Member(member) => { let field_name = &member.field_name; - build_column_iter_loop!($mode, $response, $len, &field_name, value, { + for cell in $response.column(&field_name)? { + let value = cell?; match (value, &mut $builder) { (FieldValue::Null, builder) => builder.append_null()?, $($builder_block)* @@ -472,7 +436,7 @@ macro_rules! build_column_custom_builder { ))); } }; - }); + } } MemberField::Literal(value) => { for _ in 0..$len { @@ -794,13 +758,10 @@ async fn load_data( let result = if no_members_query { let limit = request.limit.unwrap_or(1); - let mut data = Vec::new(); - - for _ in 0..limit { - data.push(serde_json::Value::Null) - } - let mut response = JsonValueObject::new(data); + let mut response = LiteralRowsValueObject { + row_count: limit as usize, + }; let rec = transform_response(&mut response, schema.clone(), &member_fields) .map_err(|e| ArrowError::ExternalError(Box::new(e)))?; @@ -897,12 +858,10 @@ fn load_to_stream_sync(one_shot_stream: &mut CubeScanOneShotStream) -> Result<() Ok(()) } -// Body of `transform_response` / `transform_columnar_response`. The two functions differ -// only in how they iterate per-cell values: row-major (`$mode = row`) goes through -// `ValueObject::get` per cell, columnar (`$mode = columnar`) fetches the whole column once -// via `ColumnarValueObject::column`. Per-`DataType` coercion arms are shared. +// Body of `transform_response`: builds one Arrow column per schema field from a +// `ColumnarValueObject`, fetching each column once via `ColumnarValueObject::column`. macro_rules! transform_response_body { - ($mode:tt, $response:expr, $schema:expr, $member_fields:expr) => {{ + ($response:expr, $schema:expr, $member_fields:expr) => {{ let mut columns = vec![]; for (i, schema_field) in $schema.fields().iter().enumerate() { @@ -912,7 +871,6 @@ macro_rules! transform_response_body { build_column!( DataType::Utf8, StringBuilder, - $mode, $response, field_name, { @@ -929,7 +887,6 @@ macro_rules! transform_response_body { build_column!( DataType::Int16, Int16Builder, - $mode, $response, field_name, { @@ -955,7 +912,6 @@ macro_rules! transform_response_body { build_column!( DataType::Int32, Int32Builder, - $mode, $response, field_name, { @@ -981,7 +937,6 @@ macro_rules! transform_response_body { build_column!( DataType::Int64, Int64Builder, - $mode, $response, field_name, { @@ -1007,7 +962,6 @@ macro_rules! transform_response_body { build_column!( DataType::Float32, Float32Builder, - $mode, $response, field_name, { @@ -1033,7 +987,6 @@ macro_rules! transform_response_body { build_column!( DataType::Float64, Float64Builder, - $mode, $response, field_name, { @@ -1059,7 +1012,6 @@ macro_rules! transform_response_body { build_column!( DataType::Boolean, BooleanBuilder, - $mode, $response, field_name, { @@ -1083,7 +1035,6 @@ macro_rules! transform_response_body { build_column!( DataType::Timestamp(TimeUnit::Nanosecond, None), TimestampNanosecondBuilder, - $mode, $response, field_name, { @@ -1110,7 +1061,6 @@ macro_rules! transform_response_body { build_column!( DataType::Timestamp(TimeUnit::Millisecond, None), TimestampMillisecondBuilder, - $mode, $response, field_name, { @@ -1128,7 +1078,6 @@ macro_rules! transform_response_body { build_column!( DataType::Date32, Date32Builder, - $mode, $response, field_name, { @@ -1161,7 +1110,6 @@ macro_rules! transform_response_body { build_column_custom_builder!( DataType::Decimal(*precision, *scale), - $mode, len, builder, $response, @@ -1208,7 +1156,6 @@ macro_rules! transform_response_body { build_column!( DataType::Interval(IntervalUnit::YearMonth), IntervalYearMonthBuilder, - $mode, $response, field_name, { @@ -1223,7 +1170,6 @@ macro_rules! transform_response_body { build_column!( DataType::Interval(IntervalUnit::DayTime), IntervalDayTimeBuilder, - $mode, $response, field_name, { @@ -1238,7 +1184,6 @@ macro_rules! transform_response_body { build_column!( DataType::Interval(IntervalUnit::MonthDayNano), IntervalMonthDayNanoBuilder, - $mode, $response, field_name, { @@ -1269,20 +1214,12 @@ macro_rules! transform_response_body { }}; } -pub fn transform_response( - response: &mut V, - schema: SchemaRef, - member_fields: &Vec, -) -> std::result::Result { - transform_response_body!(row, response, schema, member_fields) -} - -pub fn transform_columnar_response( +pub fn transform_response( response: &mut C, schema: SchemaRef, member_fields: &Vec, ) -> std::result::Result { - transform_response_body!(columnar, response, schema, member_fields) + transform_response_body!(response, schema, member_fields) } /// Builds a schema with `lastRefreshTime` / `external` metadata. @@ -1317,31 +1254,6 @@ pub fn build_response_schema( } pub fn convert_transport_response( - response: V1LoadResponse, - schema: SchemaRef, - member_fields: Vec, -) -> std::result::Result, CubeError> { - response - .results - .into_iter() - .map(|result| { - let V1LoadResult { - data, - last_refresh_time, - external, - .. - } = result; - - let mut response = JsonValueObject::new(data); - let updated_schema = - build_response_schema(&schema, last_refresh_time, external.unwrap_or(false)); - - transform_response(&mut response, updated_schema, &member_fields) - }) - .collect::, CubeError>>() -} - -pub fn convert_transport_response_columnar( response: V1LoadResponse, schema: SchemaRef, member_fields: Vec, @@ -1362,7 +1274,7 @@ pub fn convert_transport_response_columnar( let updated_schema = build_response_schema(&schema, last_refresh_time, external.unwrap_or(false)); - transform_columnar_response(&mut response, updated_schema, &member_fields) + transform_response(&mut response, updated_schema, &member_fields) }) .collect::, CubeError>>() } @@ -1376,7 +1288,7 @@ mod tests { transport::{MetaContext, SqlResponse}, CubeError, }; - use cubeclient::models::V1LoadResponse; + use cubeclient::models::{V1LoadResponse, V1LoadResultDataColumnar}; use datafusion::{ arrow::{ array::{ @@ -1447,11 +1359,10 @@ mod tests { #[test] fn convert_transport_response_threads_external_flag_into_schema_metadata() { - // End-to-end coverage of the row-format `convert_transport_response` - // path: a V1LoadResponse with `external: true` and a - // `lastRefreshTime` should produce a RecordBatch whose schema - // metadata has the `external` marker set and `lastRefreshTime` - // passed through unchanged. + // End-to-end coverage of `convert_transport_response`: a columnar + // V1LoadResponse with `external: true` and a `lastRefreshTime` should + // produce a RecordBatch whose schema metadata has the `external` marker + // set and `lastRefreshTime` passed through unchanged. let raw = r#" { "results": [{ @@ -1461,13 +1372,13 @@ mod tests { "segments": [], "timeDimensions": [] }, - "data": [{"c": 1}], + "data": {"members": ["c"], "columns": [[1]]}, "lastRefreshTime": "2000-01-01T00:00:00.000Z", "external": true }] } "#; - let response: V1LoadResponse = serde_json::from_str(raw).unwrap(); + let response: V1LoadResponse = serde_json::from_str(raw).unwrap(); let schema = build_schema(); let member_fields = vec![MemberField::regular("c".to_string())]; let batches = convert_transport_response(response, schema, member_fields).unwrap(); @@ -1491,14 +1402,17 @@ mod tests { "segments": [], "timeDimensions": [] }, - "data": [ - {"c": ["a", "b"], "d": {"k": 1}}, - {"c": null, "d": "plain"} - ] + "data": { + "members": ["c", "d"], + "columns": [ + [["a", "b"], null], + [{"k": 1}, "plain"] + ] + } }] } "#; - let response: V1LoadResponse = serde_json::from_str(raw).unwrap(); + let response: V1LoadResponse = serde_json::from_str(raw).unwrap(); let schema = Arc::new(Schema::new(vec![ Field::new("c", DataType::Utf8, true), Field::new("d", DataType::Utf8, true), @@ -1579,18 +1493,30 @@ mod tests { "segments": [], "timeDimensions": [] }, - "data": [ - {"KibanaSampleDataEcommerce.count": null, "KibanaSampleDataEcommerce.maxPrice": null, "KibanaSampleDataEcommerce.isBool": null, "KibanaSampleDataEcommerce.orderTimestamp": null, "KibanaSampleDataEcommerce.orderDate": null, "KibanaSampleDataEcommerce.city": "City 1"}, - {"KibanaSampleDataEcommerce.count": 5, "KibanaSampleDataEcommerce.maxPrice": 5.05, "KibanaSampleDataEcommerce.isBool": true, "KibanaSampleDataEcommerce.orderTimestamp": "2022-01-01 00:00:00.000", "KibanaSampleDataEcommerce.orderDate": "2022-01-01", "KibanaSampleDataEcommerce.city": "City 2"}, - {"KibanaSampleDataEcommerce.count": "5", "KibanaSampleDataEcommerce.maxPrice": "5.05", "KibanaSampleDataEcommerce.isBool": false, "KibanaSampleDataEcommerce.orderTimestamp": "2023-01-01 00:00:00.000", "KibanaSampleDataEcommerce.orderDate": "2023-01-01", "KibanaSampleDataEcommerce.city": "City 3"}, - {"KibanaSampleDataEcommerce.count": null, "KibanaSampleDataEcommerce.maxPrice": null, "KibanaSampleDataEcommerce.isBool": "true", "KibanaSampleDataEcommerce.orderTimestamp": "9999-12-31 00:00:00.000", "KibanaSampleDataEcommerce.orderDate": "9999-12-31", "KibanaSampleDataEcommerce.city": "City 4"}, - {"KibanaSampleDataEcommerce.count": null, "KibanaSampleDataEcommerce.maxPrice": null, "KibanaSampleDataEcommerce.isBool": "false", "KibanaSampleDataEcommerce.orderTimestamp": null, "KibanaSampleDataEcommerce.orderDate": null, "KibanaSampleDataEcommerce.city": null} - ] + "data": { + "members": [ + "KibanaSampleDataEcommerce.count", + "KibanaSampleDataEcommerce.maxPrice", + "KibanaSampleDataEcommerce.isBool", + "KibanaSampleDataEcommerce.orderTimestamp", + "KibanaSampleDataEcommerce.orderDate", + "KibanaSampleDataEcommerce.city" + ], + "columns": [ + [null, 5, "5", null, null], + [null, 5.05, "5.05", null, null], + [null, true, false, "true", "false"], + [null, "2022-01-01 00:00:00.000", "2023-01-01 00:00:00.000", "9999-12-31 00:00:00.000", null], + [null, "2022-01-01", "2023-01-01", "9999-12-31", null], + ["City 1", "City 2", "City 3", "City 4", null] + ] + } }] } "#; - let result: V1LoadResponse = serde_json::from_str(response).unwrap(); + let result: V1LoadResponse = + serde_json::from_str(response).unwrap(); convert_transport_response(result, schema.clone(), member_fields) } diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index a2df7daa4fdda..29d9e07201966 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -46,7 +46,7 @@ mod tests { use chrono::{Datelike, Duration, Utc}; use cubeclient::models::{ V1LoadRequestQuery, V1LoadRequestQueryFilterItem, V1LoadRequestQueryTimeDimension, - V1LoadResponse, V1LoadResult, V1LoadResultAnnotation, + V1LoadResponse, V1LoadResult, V1LoadResultAnnotation, V1LoadResultDataColumnar, }; use datafusion::{arrow::datatypes::DataType, physical_plan::displayable}; use itertools::Itertools; @@ -13921,8 +13921,30 @@ ORDER BY "source"."str0" ASC V1LoadResultAnnotation::new(json!([]), json!([]), json!([]), json!([])) } - pub(crate) fn simple_load_response(data: Vec) -> V1LoadResponse { - V1LoadResponse::new(vec![V1LoadResult::new(empty_annotation(), data)]) + // Transpose row-shaped `json!({...})` fixtures into the columnar + // `{ members, columns }` wire format the transport now consumes. Keeping the + // call sites row-shaped keeps the mocks readable. + pub(crate) fn simple_load_response( + data: Vec, + ) -> V1LoadResponse { + let members: Vec = data + .first() + .and_then(|row| row.as_object()) + .map(|row| row.keys().cloned().collect()) + .unwrap_or_default(); + let columns = members + .iter() + .map(|member| { + data.iter() + .map(|row| row.get(member).cloned().unwrap_or(serde_json::Value::Null)) + .collect() + }) + .collect(); + + V1LoadResponse::new(vec![V1LoadResult::new( + empty_annotation(), + V1LoadResultDataColumnar::new(members, columns), + )]) } #[tokio::test] diff --git a/rust/cubesql/cubesql/src/compile/test/mod.rs b/rust/cubesql/cubesql/src/compile/test/mod.rs index b65f633628f0b..3a934c1da5f1a 100644 --- a/rust/cubesql/cubesql/src/compile/test/mod.rs +++ b/rust/cubesql/cubesql/src/compile/test/mod.rs @@ -14,7 +14,7 @@ use crate::{ transport::{ CubeMeta, CubeMetaDimension, CubeMetaJoin, CubeMetaMeasure, CubeMetaSegment, CubeStreamReceiver, LoadRequestMeta, MetaContext, SpanId, SqlGenerator, SqlResponse, - SqlTemplates, TransportLoadRequestQuery, TransportLoadResponse, TransportService, + SqlTemplates, TransportLoadRequestQuery, TransportLoadResponseColumnar, TransportService, }, CubeError, CubeErrorCauseType, }; @@ -917,7 +917,7 @@ pub struct TestTransportLoadCall { #[derive(Debug)] struct TestConnectionTransport { meta_context: Arc, - load_mocks: tokio::sync::Mutex>, + load_mocks: tokio::sync::Mutex>, load_calls: tokio::sync::Mutex>, } @@ -937,7 +937,7 @@ impl TestConnectionTransport { pub async fn add_cube_load_mock( &self, req: TransportLoadRequestQuery, - res: TransportLoadResponse, + res: TransportLoadResponseColumnar, ) { self.load_mocks.lock().await.push((req, res)); } @@ -1147,7 +1147,7 @@ impl TestContext { pub async fn add_cube_load_mock( &self, mut req: TransportLoadRequestQuery, - res: TransportLoadResponse, + res: TransportLoadResponseColumnar, ) { // Fill in default limit to simplify passing queries as they were in logical plan let config_limit = self.config_obj.non_streaming_query_max_row_limit(); diff --git a/rust/cubesql/cubesql/src/transport/mod.rs b/rust/cubesql/cubesql/src/transport/mod.rs index d95c1957e0707..68592e16c0489 100644 --- a/rust/cubesql/cubesql/src/transport/mod.rs +++ b/rust/cubesql/cubesql/src/transport/mod.rs @@ -25,7 +25,6 @@ pub type CubeMetaDimensionOrder = cubeclient::models::V1CubeMetaDimensionOrder; pub type CubeMetaFormat = cubeclient::models::V1CubeMetaFormat; // Request/Response -pub type TransportLoadResponse = cubeclient::models::V1LoadResponse; pub type TransportLoadResponseColumnar = cubeclient::models::V1LoadResponse; pub type TransportLoadRequestQuery = cubeclient::models::V1LoadRequestQuery; diff --git a/rust/cubesql/cubesql/src/transport/service.rs b/rust/cubesql/cubesql/src/transport/service.rs index 5367f488ea9b2..90eb60df62293 100644 --- a/rust/cubesql/cubesql/src/transport/service.rs +++ b/rust/cubesql/cubesql/src/transport/service.rs @@ -310,6 +310,10 @@ impl TransportService for HttpTransport { }, }; + let mut query = query; + query.response_format = + Some(cubeclient::models::v1_load_request_query::ResponseFormat::Columnar); + // TODO: support meta_fields for HTTP let request = TransportLoadRequest { query: Some(Box::new(query)), From a3a5467425bac01b88aa8ad132bdff48c1193adf Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Fri, 10 Jul 2026 17:42:54 +0200 Subject: [PATCH 4/4] refactor(cubesql): Tidy leftovers from columnar transform consolidation --- rust/cubesql/cubesql/benches/transform_response.rs | 2 -- rust/cubesql/cubesql/src/transport/service.rs | 3 +-- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/rust/cubesql/cubesql/benches/transform_response.rs b/rust/cubesql/cubesql/benches/transform_response.rs index 0baba390ff3dc..e6dad770d8e17 100644 --- a/rust/cubesql/cubesql/benches/transform_response.rs +++ b/rust/cubesql/cubesql/benches/transform_response.rs @@ -204,8 +204,6 @@ fn bench_transform_response(c: &mut Criterion) { }) }, ); - - drop(inputs); } } } diff --git a/rust/cubesql/cubesql/src/transport/service.rs b/rust/cubesql/cubesql/src/transport/service.rs index 90eb60df62293..509290d6223f5 100644 --- a/rust/cubesql/cubesql/src/transport/service.rs +++ b/rust/cubesql/cubesql/src/transport/service.rs @@ -282,7 +282,7 @@ impl TransportService for HttpTransport { async fn load( &self, _span_id: Option>, - query: TransportLoadRequestQuery, + mut query: TransportLoadRequestQuery, _sql_query: Option, ctx: AuthContextRef, meta: LoadRequestMeta, @@ -310,7 +310,6 @@ impl TransportService for HttpTransport { }, }; - let mut query = query; query.response_format = Some(cubeclient::models::v1_load_request_query::ResponseFormat::Columnar);