From 8b535c5687172c344f7fd6c88f1ad4bbcd7c42de Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Wed, 29 Jul 2026 17:20:44 +0200 Subject: [PATCH 1/5] perf(cubeorchestrator): keep CubeStore Arrow results in Arrow memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CubeStore answers with Arrow IPC, but `QueryResult::from_arrow` threw that away: it walked every cell and materialized a `Vec` per column (32 B/cell plus a heap `String` per text/decimal cell), then the transform cloned every cell again into an output dataset that was itself fully built before serialization began. A 100k x 32 result was materialized twice. Introduce `QueryResultColumn { Columnar(ColumnarArray) | Arrow(ArrowArray) }`. Arrow columns stay as Arrow buffers; `ColumnReader` resolves the concrete array once per column so the per-cell path is a jump table plus an index, never a `DataType` match and a downcast. Multi-batch streams are concatenated per column at parse, so a single-batch stream — everything outside CubeStore's streaming path — is copied not at all. On top of that, `DirectData` serializes the `compact` and `columnar` response formats straight from the source columns, so the output dataset is never materialized either. Cells that need no transform are handed to the serializer by reference, which is where most of the JSON-side win comes from. `TransformedData::transform` is unchanged: the SQL API needs the in-memory `Columnar` form, and it doubles as the test oracle. The vanilla format still materializes — streaming it needs a duplicate-key guard first, since today's row `IndexMap` silently dedupes the deprecated-granularity and blending keys. 16 columns x 100k rows, before -> after: | stage | before | after | | ------------------------- | -------- | ------- | | Arrow parse | 18.34 ms | 1.85 ms | | final JSON, compact | 66.8 ms | 57.1 ms | | final JSON, columnar | 53.8 ms | 51.3 ms | | **end-to-end, compact** | 85.1 ms | 59.0 ms | | **end-to-end, columnar** | 72.1 ms | 53.2 ms | Peak memory drops by the whole intermediate primitive buffer (~3.2M cells at that size) for compact and columnar. Behaviour is held byte-identical by a new oracle test: for every fixture, in every response format, with both Arrow- and primitive-backed columns, the streamed JSON must equal the materialized JSON exactly. Two deliberate differences: plan-time failures on the streaming path surface as a serde error wrapping the message instead of an `anyhow` error raised before serialization, and a zero-row column of an unsupported Arrow type still parses (no cell is ever read), matching the pre-existing behaviour of the per-cell loop. Co-Authored-By: Claude Opus 5 (1M context) --- .../cubejs-backend-native/src/orchestrator.rs | 34 +- .../cubeorchestrator/benches/common/mod.rs | 58 ++ rust/cube/cubeorchestrator/benches/parser.rs | 47 +- .../cubeorchestrator/benches/transform.rs | 231 +++++-- .../cubeorchestrator/src/direct_result.rs | 366 ++++++++++ rust/cube/cubeorchestrator/src/lib.rs | 2 + .../src/query_message_parser.rs | 425 +++++------- .../src/query_result_column.rs | 648 ++++++++++++++++++ .../src/query_result_transform.rs | 402 ++++++++--- 9 files changed, 1789 insertions(+), 424 deletions(-) create mode 100644 rust/cube/cubeorchestrator/src/direct_result.rs create mode 100644 rust/cube/cubeorchestrator/src/query_result_column.rs diff --git a/packages/cubejs-backend-native/src/orchestrator.rs b/packages/cubejs-backend-native/src/orchestrator.rs index 4524a7e502b62..7734510610964 100644 --- a/packages/cubejs-backend-native/src/orchestrator.rs +++ b/packages/cubejs-backend-native/src/orchestrator.rs @@ -1,5 +1,6 @@ use crate::node_obj_deserializer::JsValueDeserializer; use crate::transport::MapCubeErrExt; +use cubeorchestrator::direct_result::DirectData; use cubeorchestrator::query_message_parser::QueryResult; use cubeorchestrator::query_result_transform::{ DBResponsePrimitive, RequestResultData, RequestResultDataMulti, TransformedData, @@ -257,8 +258,10 @@ pub fn get_cubestore_result(mut cx: FunctionContext) -> JsResult { result.members().iter().map(|k| cx.string(k)).collect(); let row_count = result.row_count(); + // One reader per column, so an Arrow-backed column resolves its type once + // instead of on every cell. let columns: Vec<_> = (0..js_keys.len()) - .map(|i| result.column(i)) + .map(|i| result.reader(i)) .collect::>() .or_else(|err| cx.throw_error(err.to_string()))?; let js_array = JsArray::new(&mut cx, row_count); @@ -268,11 +271,11 @@ pub fn get_cubestore_result(mut cx: FunctionContext) -> JsResult { let js_row = JsObject::new(&mut cx); 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(), - // For compatibility, we convert all primitives to strings - other => cx.string(other.to_string()).upcast(), + // For compatibility, we convert all primitives to strings + let js_value: Handle<'_, JsValue> = match columns[col_idx].value_as_str(row_idx) + { + None => cx.null().upcast(), + Some(text) => cx.string(text).upcast(), }; js_row.set(&mut cx, *js_key, js_value)?; @@ -307,14 +310,17 @@ pub fn final_query_result(mut cx: FunctionContext) -> JsResult { let result_data_js_object = cx.argument::(2)?; let deserializer = JsValueDeserializer::new(&mut cx, result_data_js_object); - let mut result_data: RequestResultData = match Deserialize::deserialize(deserializer) { + let result_data: RequestResultData = match Deserialize::deserialize(deserializer) { Ok(data) => data, Err(err) => return cx.throw_error(err.to_string()), }; let promise = cx .task(move || { - result_data.prepare_results(&transform_request_data, &cube_store_result)?; + // The result is serialized and dropped, so render `data` straight + // from the source columns instead of materializing it first. + let result_data = + result_data.with_data(DirectData::new(&transform_request_data, &cube_store_result)); match serde_json::to_string(&result_data) { Ok(json) => Ok(json), @@ -353,7 +359,17 @@ pub fn final_query_result_multi(mut cx: FunctionContext) -> JsResult let promise = cx .task(move || { - result_data.prepare_results(&transform_requests, &cube_store_results)?; + result_data.prepare_pivot_query()?; + + // As in `final_query_result`: render each `data` member while + // serializing, straight from its source columns. + let result_data = result_data.with_data( + transform_requests + .iter() + .zip(cube_store_results.iter()) + .map(|(request, source)| DirectData::new(request, source)) + .collect(), + )?; match serde_json::to_string(&result_data) { Ok(json) => Ok(json), diff --git a/rust/cube/cubeorchestrator/benches/common/mod.rs b/rust/cube/cubeorchestrator/benches/common/mod.rs index 6b2aeb0200a70..8a57247404a50 100644 --- a/rust/cube/cubeorchestrator/benches/common/mod.rs +++ b/rust/cube/cubeorchestrator/benches/common/mod.rs @@ -1,7 +1,13 @@ #![allow(dead_code)] +use cubeorchestrator::query_message_parser::QueryResult; use cubeorchestrator::query_result_transform::{ColumnarArray, DBResponsePrimitive}; use cubeorchestrator::transport::JsRawColumnarData; +use cubeshared::codegen::{ + HttpCommand, HttpMessage, HttpMessageArgs, HttpQueryResult, HttpQueryResultArgs, + HttpQueryResultArrow, HttpQueryResultArrowArgs, HttpQueryResultData, +}; +use cubeshared::flatbuffers::FlatBufferBuilder; pub const ROW_COUNTS: &[usize] = &[1_000, 10_000, 50_000, 100_000]; pub const COLUMN_COUNTS: &[usize] = &[8, 16, 32, 64]; @@ -137,3 +143,55 @@ pub fn build_arrow_ipc( } buf } + +/// Wrap raw Arrow IPC bytes in an `HttpMessage` FlatBuffer carrying +/// `HttpQueryResultArrow`, exactly as CubeStore sends it. +pub fn build_cubestore_fb_arrow_message(arrow_ipc: &[u8]) -> Vec { + let mut builder = FlatBufferBuilder::new(); + let data_vec = builder.create_vector(arrow_ipc); + let arrow = HttpQueryResultArrow::create( + &mut builder, + &HttpQueryResultArrowArgs { + data: Some(data_vec), + is_last: true, + }, + ); + let query_result = HttpQueryResult::create( + &mut builder, + &HttpQueryResultArgs { + data_type: HttpQueryResultData::HttpQueryResultArrow, + data: Some(arrow.as_union_value()), + }, + ); + let connection_id = builder.create_string("bench_connection"); + let message = HttpMessage::create( + &mut builder, + &HttpMessageArgs { + message_id: 1, + command_type: HttpCommand::HttpQueryResult, + command: Some(query_result.as_union_value()), + connection_id: Some(connection_id), + }, + ); + builder.finish(message, None); + builder.finished_data().to_vec() +} + +/// An Arrow-backed `QueryResult` with the same logical shape as +/// [`build_dataset`], so transform throughput can be compared column storage +/// against column storage. +/// +/// Note the time dimensions differ in kind, not just encoding: the Arrow fixture +/// carries `Timestamp(Millisecond)` cells, which the `time` member type passes +/// straight through, while [`build_dataset`] carries ISO strings that +/// `transform_value` re-parses and re-formats per cell. +pub fn build_arrow_query_result( + row_count: usize, + dimensions: &[(String, String)], + measures: &[(String, String)], + time_dims: &[TimeColumn], +) -> QueryResult { + let ipc = build_arrow_ipc(row_count, dimensions, measures, time_dims); + let payload = build_cubestore_fb_arrow_message(&ipc); + QueryResult::from_cubestore_fb(&payload).expect("arrow query result") +} diff --git a/rust/cube/cubeorchestrator/benches/parser.rs b/rust/cube/cubeorchestrator/benches/parser.rs index 7358bd56ed6de..2f80fc10902f8 100644 --- a/rust/cube/cubeorchestrator/benches/parser.rs +++ b/rust/cube/cubeorchestrator/benches/parser.rs @@ -4,17 +4,16 @@ use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Through use cubeorchestrator::query_message_parser::QueryResult; use cubeorchestrator::transport::JsRawColumnarData; use cubeshared::codegen::{ - HttpColumnValue, HttpColumnValueArgs, HttpCommand, HttpMessage, HttpMessageArgs, - HttpQueryResult, HttpQueryResultArgs, HttpQueryResultArrow, HttpQueryResultArrowArgs, - HttpQueryResultData, HttpResultSet, HttpResultSetArgs, HttpRow, HttpRowArgs, + HttpColumnValue, HttpColumnValueArgs, HttpCommand, HttpMessage, HttpMessageArgs, HttpResultSet, + HttpResultSetArgs, HttpRow, HttpRowArgs, }; use cubeshared::flatbuffers::FlatBufferBuilder; #[path = "common/mod.rs"] mod common; use common::{ - build_arrow_ipc, build_dataset, make_member_aliases, split_dim_measure, COLUMN_COUNTS, - ROW_COUNTS, + build_arrow_ipc, build_cubestore_fb_arrow_message, build_dataset, make_member_aliases, + split_dim_measure, COLUMN_COUNTS, ROW_COUNTS, }; /// Build a FlatBuffer `HttpMessage` payload mirroring CubeStore's wire format @@ -153,38 +152,6 @@ fn bench_from_js_raw_data(c: &mut Criterion) { group.finish(); } -/// Wrap raw Arrow IPC bytes in an `HttpMessage` FlatBuffer carrying -fn build_cubestore_fb_arrow_message(arrow_ipc: &[u8]) -> Vec { - let mut builder = FlatBufferBuilder::new(); - let data_vec = builder.create_vector(arrow_ipc); - let arrow = HttpQueryResultArrow::create( - &mut builder, - &HttpQueryResultArrowArgs { - data: Some(data_vec), - is_last: true, - }, - ); - let query_result = HttpQueryResult::create( - &mut builder, - &HttpQueryResultArgs { - data_type: HttpQueryResultData::HttpQueryResultArrow, - data: Some(arrow.as_union_value()), - }, - ); - let connection_id = builder.create_string("bench_connection"); - let message = HttpMessage::create( - &mut builder, - &HttpMessageArgs { - message_id: 1, - command_type: HttpCommand::HttpQueryResult, - command: Some(query_result.as_union_value()), - connection_id: Some(connection_id), - }, - ); - builder.finish(message, None); - builder.finished_data().to_vec() -} - fn bench_from_cubestore_fb_arrow(c: &mut Criterion) { let mut group = c.benchmark_group("QueryResult::from_cubestore_fb_arrow"); @@ -214,8 +181,10 @@ fn bench_from_cubestore_fb_arrow(c: &mut Criterion) { group.throughput(Throughput::Elements((row_count * col_count) as u64)); let id = format!("c{:02}_r{}", col_count, row_count); - // Arrow IPC parse always materializes the QueryResult, so this measures - // the equivalent of from_js_raw_data's `parse_plus_build`. + // Arrow columns are kept in Arrow memory, so this measures the IPC decode + // and the column wiring — not a per-cell conversion. The conversion cost + // moved to `TransformedData::transform`; see the `arrow` axis in + // benches/transform.rs. group.bench_with_input(BenchmarkId::from_parameter(id), &(), |b, _| { b.iter(|| { let built = QueryResult::from_cubestore_fb(black_box(&payload)) diff --git a/rust/cube/cubeorchestrator/benches/transform.rs b/rust/cube/cubeorchestrator/benches/transform.rs index 57562aeaca7b4..2dd70a047a7c8 100644 --- a/rust/cube/cubeorchestrator/benches/transform.rs +++ b/rust/cube/cubeorchestrator/benches/transform.rs @@ -2,8 +2,9 @@ use std::collections::HashMap; use std::hint::black_box; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use cubeorchestrator::direct_result::DirectData; use cubeorchestrator::query_message_parser::QueryResult; -use cubeorchestrator::query_result_transform::TransformedData; +use cubeorchestrator::query_result_transform::{RequestResultData, TransformedData}; use cubeorchestrator::transport::{ ConfigItem, MemberOrMemberExpression, NormalizedQuery, QueryType, ResultType, TransformDataRequest, @@ -12,7 +13,8 @@ use cubeorchestrator::transport::{ #[path = "common/mod.rs"] mod common; use common::{ - build_dataset, make_member_aliases, split_dim_measure, TimeColumn, COLUMN_COUNTS, ROW_COUNTS, + build_arrow_query_result, build_dataset, make_member_aliases, split_dim_measure, TimeColumn, + COLUMN_COUNTS, ROW_COUNTS, }; /// Total columns and row count used by `bench_transform_time_scenarios`. @@ -152,32 +154,43 @@ fn bench_transform(c: &mut Criterion) { let measures = make_member_aliases("measure", measure_count); for &row_count in ROW_COUNTS { - let raw = QueryResult::from_js_raw_data(build_dataset( - row_count, - &dimensions, - &measures, - &[], - )) - .expect("from_js_raw_data"); + let sources = [ + ( + "js_raw", + QueryResult::from_js_raw_data(build_dataset( + row_count, + &dimensions, + &measures, + &[], + )) + .expect("from_js_raw_data"), + ), + ( + "arrow", + build_arrow_query_result(row_count, &dimensions, &measures, &[]), + ), + ]; // Throughput in cells/sec so numbers are comparable across widths. group.throughput(Throughput::Elements((row_count * col_count) as u64)); - for (label, res_type) in [ - ("compact", Some(ResultType::Compact)), - ("columnar", Some(ResultType::Columnar)), - ("vanilla", None), - ] { - let request = build_request(res_type, &dimensions, &measures, &[]); - let id_param = format!("c{:02}_r{}", col_count, row_count); - group.bench_with_input(BenchmarkId::new(label, id_param), &(), |b, _| { - b.iter(|| { - let result = - TransformedData::transform(black_box(&request), black_box(&raw)) - .expect("transform"); - black_box(result); + for (source, raw) in &sources { + for (label, res_type) in [ + ("compact", Some(ResultType::Compact)), + ("columnar", Some(ResultType::Columnar)), + ("vanilla", None), + ] { + let request = build_request(res_type, &dimensions, &measures, &[]); + let id_param = format!("{}/c{:02}_r{}", source, col_count, row_count); + group.bench_with_input(BenchmarkId::new(label, id_param), &(), |b, _| { + b.iter(|| { + let result = + TransformedData::transform(black_box(&request), black_box(raw)) + .expect("transform"); + black_box(result); + }); }); - }); + } } } } @@ -202,13 +215,22 @@ fn bench_transform_time_scenarios(c: &mut Criterion) { let dimensions = make_member_aliases("dim", dim_count); let measures = make_member_aliases("measure", measure_count); - let raw = QueryResult::from_js_raw_data(build_dataset( - SCENARIO_ROW_COUNT, - &dimensions, - &measures, - &time_dims, - )) - .expect("from_js_raw_data"); + let sources = [ + ( + "js_raw", + QueryResult::from_js_raw_data(build_dataset( + SCENARIO_ROW_COUNT, + &dimensions, + &measures, + &time_dims, + )) + .expect("from_js_raw_data"), + ), + ( + "arrow", + build_arrow_query_result(SCENARIO_ROW_COUNT, &dimensions, &measures, &time_dims), + ), + ]; // Throughput in cells/sec; total cells = row_count * total_cols, where // total_cols == SCENARIO_COL_COUNT regardless of scenario. @@ -216,30 +238,141 @@ fn bench_transform_time_scenarios(c: &mut Criterion) { (SCENARIO_ROW_COUNT * SCENARIO_COL_COUNT) as u64, )); - for (label, res_type) in [ - ("compact", Some(ResultType::Compact)), - ("columnar", Some(ResultType::Columnar)), - ("vanilla", None), - ] { - let request = build_request(res_type, &dimensions, &measures, &time_dims); - let id_param = format!( - "{}/c{:02}_r{}", - scenario.label(), - SCENARIO_COL_COUNT, - SCENARIO_ROW_COUNT - ); - group.bench_with_input(BenchmarkId::new(label, id_param), &(), |b, _| { - b.iter(|| { - let result = TransformedData::transform(black_box(&request), black_box(&raw)) - .expect("transform"); - black_box(result); + for (source, raw) in &sources { + for (label, res_type) in [ + ("compact", Some(ResultType::Compact)), + ("columnar", Some(ResultType::Columnar)), + ("vanilla", None), + ] { + let request = build_request(res_type, &dimensions, &measures, &time_dims); + let id_param = format!( + "{}/{}/c{:02}_r{}", + source, + scenario.label(), + SCENARIO_COL_COUNT, + SCENARIO_ROW_COUNT + ); + group.bench_with_input(BenchmarkId::new(label, id_param), &(), |b, _| { + b.iter(|| { + let result = + TransformedData::transform(black_box(&request), black_box(raw)) + .expect("transform"); + black_box(result); + }); }); - }); + } } } group.finish(); } -criterion_group!(benches, bench_transform, bench_transform_time_scenarios); +/// The whole job the `getFinalQueryResult` bridge does: turn a source result +/// into the response JSON. `materialized` builds a `TransformedData` first, +/// `direct` renders the `data` member while serializing. +fn bench_final_json(c: &mut Criterion) { + let mut group = c.benchmark_group("final_json"); + + let col_count = 16usize; + let (dim_count, measure_count) = split_dim_measure(col_count); + let dimensions = make_member_aliases("dim", dim_count); + let measures = make_member_aliases("measure", measure_count); + + for &row_count in &[10_000usize, 100_000] { + let sources = [ + ( + "js_raw", + QueryResult::from_js_raw_data(build_dataset( + row_count, + &dimensions, + &measures, + &[], + )) + .expect("from_js_raw_data"), + ), + ( + "arrow", + build_arrow_query_result(row_count, &dimensions, &measures, &[]), + ), + ]; + + group.throughput(Throughput::Elements((row_count * col_count) as u64)); + + for (source, raw) in &sources { + for (label, res_type) in [ + ("compact", Some(ResultType::Compact)), + ("columnar", Some(ResultType::Columnar)), + ("vanilla", None), + ] { + let request = build_request(res_type, &dimensions, &measures, &[]); + let head = result_head(); + let id_param = format!("{}/{}/c{:02}_r{}", label, source, col_count, row_count); + + group.bench_with_input(BenchmarkId::new("materialized", &id_param), &(), |b, _| { + b.iter(|| { + let mut result = head.clone(); + result + .prepare_results(black_box(&request), black_box(raw)) + .expect("prepare_results"); + black_box(serde_json::to_string(&result).expect("to_string")); + }); + }); + + group.bench_with_input(BenchmarkId::new("direct", &id_param), &(), |b, _| { + b.iter(|| { + let result = head + .clone() + .with_data(DirectData::new(black_box(&request), black_box(raw))); + black_box(serde_json::to_string(&result).expect("to_string")); + }); + }); + } + } + } + + group.finish(); +} + +/// Response envelope the neon bridge deserializes from JS, with `data` still empty. +fn result_head() -> RequestResultData { + RequestResultData { + query: NormalizedQuery { + measures: None, + dimensions: None, + time_dimensions: None, + segments: None, + limit: None, + offset: None, + total: None, + total_query: None, + timezone: Some("UTC".to_string()), + ungrouped: None, + response_format: None, + filters: None, + row_limit: None, + order: None, + query_type: Some(QueryType::RegularQuery), + }, + last_refresh_time: None, + refresh_key_values: None, + used_pre_aggregations: None, + transformed_query: None, + request_id: Some("bench".to_string()), + annotation: HashMap::new(), + data_source: Some("default".to_string()), + db_type: Some("postgres".to_string()), + ext_db_type: None, + external: Some(false), + slow_query: false, + total: None, + data: None, + } +} + +criterion_group!( + benches, + bench_transform, + bench_transform_time_scenarios, + bench_final_json +); criterion_main!(benches); diff --git a/rust/cube/cubeorchestrator/src/direct_result.rs b/rust/cube/cubeorchestrator/src/direct_result.rs new file mode 100644 index 0000000000000..878d14351fb52 --- /dev/null +++ b/rust/cube/cubeorchestrator/src/direct_result.rs @@ -0,0 +1,366 @@ +//! Serializing a response straight from the source columns. +//! +//! [`crate::query_result_transform::TransformedData::transform`] builds the whole +//! output in memory before anything is written. For the callers that only ever +//! serialize the result and drop it (`getFinalQueryResult` and its multi +//! variant), [`DirectData`] renders the `data` member while serializing instead: +//! one cell is materialized at a time, read from the source column — Arrow memory +//! included — so neither the intermediate primitives nor the output dataset are +//! ever fully materialized. +//! +//! Values are still rendered by `Serialize for DBResponsePrimitive` and +//! [`transform_value`], so the JSON is byte-for-byte what the materializing path +//! produces. `test_direct_matches_transformed_*` asserts exactly that. + +use crate::{ + query_message_parser::QueryResult, + query_result_transform::{ + build_columnar_plan, build_compact_plan, get_members, ColumnarColumnPlan, + ColumnarColumnSource, CompactPlan, CompactPlanEntry, DBResponsePrimitive, TransformedData, + }, + transport::{QueryType, ResultType, TransformDataRequest}, +}; +use serde::{ + ser::{Error as SerError, SerializeSeq, SerializeStruct}, + Serialize, Serializer, +}; + +/// The `data` member of a response, rendered from `source` as it is serialized. +pub struct DirectData<'a> { + request: &'a TransformDataRequest, + source: &'a QueryResult, +} + +impl<'a> DirectData<'a> { + pub fn new(request: &'a TransformDataRequest, source: &'a QueryResult) -> Self { + Self { request, source } + } + + fn query_type(&self) -> QueryType { + self.request.query_type.clone().unwrap_or_default() + } + + fn serialize_compact(&self, serializer: S) -> Result { + let query_type = self.query_type(); + let request = self.request; + + let (members_to_alias_map, members) = get_members( + &query_type, + &request.query, + self.source, + &request.alias_to_member_name_map, + &request.annotation, + ) + .map_err(S::Error::custom)?; + + let plan = build_compact_plan( + &members, + &members_to_alias_map, + &request.annotation, + self.source, + &query_type, + request.query.time_dimensions.as_ref(), + ) + .map_err(S::Error::custom)?; + + let mut out = serializer.serialize_struct("TransformedData", 2)?; + out.serialize_field("members", &members)?; + out.serialize_field( + "dataset", + &CompactDataset { + plan: &plan, + row_count: self.source.row_count(), + }, + )?; + out.end() + } + + fn serialize_columnar(&self, serializer: S) -> Result { + let query_type = self.query_type(); + let request = self.request; + + let (members_to_alias_map, members) = get_members( + &query_type, + &request.query, + self.source, + &request.alias_to_member_name_map, + &request.annotation, + ) + .map_err(S::Error::custom)?; + + let plan = build_columnar_plan( + &members, + &members_to_alias_map, + &request.annotation, + &self.source.columns_pos, + &query_type, + request.query.time_dimensions.as_ref(), + ) + .map_err(S::Error::custom)?; + + let mut out = serializer.serialize_struct("TransformedData", 2)?; + out.serialize_field("members", &members)?; + out.serialize_field( + "columns", + &ColumnarColumns { + plan: &plan, + source: self.source, + }, + )?; + out.end() + } +} + +impl Serialize for DirectData<'_> { + fn serialize(&self, serializer: S) -> Result { + match self.request.res_type { + Some(ResultType::Compact) => self.serialize_compact(serializer), + Some(ResultType::Columnar) => self.serialize_columnar(serializer), + // The vanilla format keeps materializing: its rows are maps, and the + // deprecated-granularity and blending keys rely on the row map + // deduplicating them. + _ => TransformedData::transform(self.request, self.source) + .map_err(S::Error::custom)? + .serialize(serializer), + } + } +} + +/// `dataset` of a compact result: one row at a time, no row ever kept. +struct CompactDataset<'a> { + plan: &'a CompactPlan<'a>, + row_count: usize, +} + +impl Serialize for CompactDataset<'_> { + fn serialize(&self, serializer: S) -> Result { + let mut seq = serializer.serialize_seq(Some(self.row_count))?; + for row_idx in 0..self.row_count { + seq.serialize_element(&CompactRow { + plan: self.plan, + row_idx, + })?; + } + seq.end() + } +} + +struct CompactRow<'a> { + plan: &'a CompactPlan<'a>, + row_idx: usize, +} + +impl Serialize for CompactRow<'_> { + fn serialize(&self, serializer: S) -> Result { + let entries = &self.plan.entries; + let mut seq = serializer.serialize_seq(Some(entries.len()))?; + + for entry in entries { + match entry { + CompactPlanEntry::Cell { + column, + member_type, + } => column.with_transformed(self.row_idx, member_type, |value| { + seq.serialize_element(value) + })?, + CompactPlanEntry::Constant(value) => seq.serialize_element(value)?, + } + } + + seq.end() + } +} + +/// `columns` of a columnar result: one column at a time, read column-major. +struct ColumnarColumns<'a> { + plan: &'a [ColumnarColumnPlan<'a>], + source: &'a QueryResult, +} + +impl Serialize for ColumnarColumns<'_> { + fn serialize(&self, serializer: S) -> Result { + let mut seq = serializer.serialize_seq(Some(self.plan.len()))?; + for entry in self.plan { + seq.serialize_element(&ColumnarColumn { + entry, + source: self.source, + })?; + } + seq.end() + } +} + +struct ColumnarColumn<'a> { + entry: &'a ColumnarColumnPlan<'a>, + source: &'a QueryResult, +} + +impl Serialize for ColumnarColumn<'_> { + fn serialize(&self, serializer: S) -> Result { + let row_count = self.source.row_count(); + let mut seq = serializer.serialize_seq(Some(row_count))?; + + match &self.entry.source { + ColumnarColumnSource::DbColumn { index } => { + // Column-major, so the reader keeps its type dispatch outside the + // row loop — the same tight fill the materializing path uses. + let reader = self.source.reader(*index).map_err(S::Error::custom)?; + reader.for_each_transformed(self.entry.member_type, |value| { + seq.serialize_element(value.as_ref()) + })?; + } + ColumnarColumnSource::Constant(value) => { + for _ in 0..row_count { + seq.serialize_element(value)?; + } + } + ColumnarColumnSource::NullFilled => { + for _ in 0..row_count { + seq.serialize_element(&DBResponsePrimitive::Null)?; + } + } + } + + seq.end() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::query_result_transform::tests::{ + make_result_head, StorageFixture, ALL_RES_TYPES, TEST_SUITE_DATA, + }; + use crate::query_result_transform::RequestResultDataMulti; + use anyhow::Result; + use std::sync::Arc; + + /// The streamed response must be byte-identical to the materialized one. + fn assert_direct_matches_transformed( + request: &TransformDataRequest, + source: &QueryResult, + context: &str, + ) -> Result<()> { + let head = make_result_head(request.query.clone()); + + let materialized = head + .clone() + .with_data(TransformedData::transform(request, source)?); + let direct = head.with_data(DirectData::new(request, source)); + + assert_eq!( + serde_json::to_string(&materialized)?, + serde_json::to_string(&direct)?, + "{context}: streamed JSON must match the materialized JSON" + ); + + Ok(()) + } + + /// Covers the regular, compare-date-range and blending query shapes — the + /// three that build different plans — in every response format. + #[test] + fn test_direct_matches_transformed_for_all_fixtures() -> Result<()> { + for (name, test_data) in TEST_SUITE_DATA.iter() { + let source = QueryResult::from_js_raw_data(test_data.query_result.clone())?; + + for res_type in ALL_RES_TYPES { + let mut request = test_data.request.clone(); + request.res_type = res_type.clone(); + + assert_direct_matches_transformed( + &request, + &source, + &format!("{name} / {res_type:?}"), + )?; + } + } + + Ok(()) + } + + /// Same check with Arrow-backed columns, which decode cells while serializing. + #[test] + fn test_direct_matches_transformed_for_arrow_source() -> Result<()> { + let fixture = StorageFixture::new()?; + + for res_type in ALL_RES_TYPES { + let request = fixture.request(res_type.clone()); + assert_direct_matches_transformed(&request, &fixture.arrow, &format!("{res_type:?}"))?; + assert_direct_matches_transformed( + &request, + &fixture.columnar, + &format!("{res_type:?}"), + )?; + } + + Ok(()) + } + + /// The multi envelope must match too, `pivotQuery` included. + #[test] + fn test_direct_matches_transformed_multi() -> Result<()> { + let fixture = StorageFixture::new()?; + let requests = [ + fixture.request(Some(ResultType::Compact)), + fixture.request(Some(ResultType::Columnar)), + ]; + let sources = [&fixture.arrow, &fixture.columnar]; + + let envelope = RequestResultDataMulti { + query_type: QueryType::RegularQuery, + results: requests + .iter() + .map(|request| make_result_head(request.query.clone())) + .collect(), + pivot_query: None, + slow_query: false, + }; + + let mut materialized = envelope.clone(); + let owned_sources: Vec<_> = sources.iter().map(|s| Arc::new((*s).clone())).collect(); + materialized.prepare_results(&requests, &owned_sources)?; + + let mut direct = envelope; + direct.prepare_pivot_query()?; + let direct = direct.with_data( + requests + .iter() + .zip(sources) + .map(|(request, source)| DirectData::new(request, source)) + .collect(), + )?; + + assert_eq!( + serde_json::to_string(&materialized)?, + serde_json::to_string(&direct)?, + "streamed multi JSON must match the materialized multi JSON" + ); + + Ok(()) + } + + /// A plan that cannot be built still fails on the streaming path — it just + /// surfaces as a serializer error instead of before serialization. + #[test] + fn test_direct_reports_plan_errors() -> Result<()> { + let fixture = StorageFixture::new()?; + let mut request = fixture.request(Some(ResultType::Compact)); + request.alias_to_member_name_map.clear(); + + assert!( + TransformedData::transform(&request, &fixture.arrow).is_err(), + "materializing path must reject an unmapped alias" + ); + + let head = make_result_head(request.query.clone()); + let err = serde_json::to_string(&head.with_data(DirectData::new(&request, &fixture.arrow))) + .expect_err("streaming path must reject an unmapped alias"); + assert!( + err.to_string().contains("Member name not found for alias"), + "unexpected error: {err}" + ); + + Ok(()) + } +} diff --git a/rust/cube/cubeorchestrator/src/lib.rs b/rust/cube/cubeorchestrator/src/lib.rs index 03f0453c16db7..a5928a07b77cb 100644 --- a/rust/cube/cubeorchestrator/src/lib.rs +++ b/rust/cube/cubeorchestrator/src/lib.rs @@ -1,3 +1,5 @@ +pub mod direct_result; pub mod query_message_parser; +pub mod query_result_column; pub mod query_result_transform; pub mod transport; diff --git a/rust/cube/cubeorchestrator/src/query_message_parser.rs b/rust/cube/cubeorchestrator/src/query_message_parser.rs index 7b4d1bcd151bc..dba70654f8a1f 100644 --- a/rust/cube/cubeorchestrator/src/query_message_parser.rs +++ b/rust/cube/cubeorchestrator/src/query_message_parser.rs @@ -1,15 +1,10 @@ use crate::{ + query_result_column::{ArrowArray, ColumnReader, QueryResultColumn}, query_result_transform::{ColumnarArray, DBResponsePrimitive}, transport::JsRawColumnarData, }; -use arrow::array::{ - Array, BooleanArray, Date32Array, Date64Array, Decimal128Array, Decimal256Array, Float16Array, - Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, LargeStringArray, - StringArray, StringViewArray, TimestampMicrosecondArray, TimestampMillisecondArray, - TimestampNanosecondArray, TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, - UInt8Array, -}; -use arrow::datatypes::{DataType, TimeUnit}; +use arrow::array::{new_empty_array, Array, ArrayRef}; +use arrow::compute::concat; use arrow::ipc::reader::StreamReader; use cubeshared::codegen::{ root_as_http_message_with_opts, HttpCommand, HttpQueryResultData, HttpResultSet, @@ -100,7 +95,7 @@ pub struct QueryResult { pub(crate) members: Vec, pub(crate) columns_pos: IndexMap, pub(crate) row_count: usize, - pub(crate) data: Vec, + pub(crate) data: Vec, } impl Finalize for QueryResult {} @@ -115,7 +110,7 @@ impl QueryResult { } } - pub fn try_new(members: Vec, data: Vec) -> Result { + pub fn try_new(members: Vec, data: Vec) -> Result { if members.len() != data.len() { return Err(ParseError::MembersColumnsMismatch { members_len: members.len(), @@ -161,16 +156,30 @@ impl QueryResult { } #[inline] - pub fn column(&self, idx: usize) -> Result<&ColumnarArray, ParseError> { + pub fn column(&self, idx: usize) -> Result<&QueryResultColumn, ParseError> { self.data.get(idx).ok_or(ParseError::ColumnIndexOutOfRange { idx, data_len: self.data.len(), }) } + /// Per-cell accessor for column `idx`. Resolve it once per column and read + /// rows off it — for Arrow-backed columns that keeps the type dispatch and + /// downcast out of the per-cell path. + #[inline] + pub fn reader(&self, idx: usize) -> Result, ParseError> { + self.column(idx)?.reader() + } + pub fn from_js_raw_data(js_raw_data: JsRawColumnarData) -> Result { let JsRawColumnarData { members, columns } = js_raw_data; - QueryResult::try_new(members, columns) + QueryResult::try_new( + members, + columns + .into_iter() + .map(QueryResultColumn::Columnar) + .collect(), + ) } pub fn from_cubestore_fb(msg_data: &[u8]) -> Result { @@ -273,9 +282,14 @@ impl QueryResult { (0..n_cols).map(|_| ColumnarArray::new()).collect() }; - QueryResult::try_new(members, data) + QueryResult::try_new( + members, + data.into_iter().map(QueryResultColumn::Columnar).collect(), + ) } + /// Build a result straight from an Arrow IPC stream, keeping every column in + /// Arrow memory — cells are decoded on read, not here. pub(crate) fn from_arrow(bytes: &[u8]) -> Result { let reader = StreamReader::try_new(Cursor::new(bytes), None) .map_err(|err| ParseError::ArrowError(err.to_string()))?; @@ -284,220 +298,45 @@ impl QueryResult { let members: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); let n_cols = members.len(); - let mut columns: Vec> = (0..n_cols).map(|_| Vec::new()).collect(); + let mut chunks: Vec> = (0..n_cols).map(|_| Vec::new()).collect(); for batch in reader { let batch = batch.map_err(|err| ParseError::ArrowError(err.to_string()))?; - for (idx, col) in columns.iter_mut().enumerate() { - append_arrow_array(col, batch.column(idx).as_ref())?; + for (idx, col_chunks) in chunks.iter_mut().enumerate() { + col_chunks.push(batch.column(idx).clone()); } } - let data: Vec = columns.into_iter().map(ColumnarArray::from).collect(); - QueryResult::try_new(members, data) - } -} - -/// Format a decimal `mantissa` with `scale` fractional digits, stripping trailing -/// fractional zeros. Generic over the mantissa's `Display`, so it renders any Arrow -/// decimal width (`i32`/`i64`/`i128`/`i256`) directly — Decimal256 needs no fallback -/// to Arrow's own string conversion. -/// -/// e.g. `(25987600, 5) -> "259.876"`, `(6199200000, 5) -> "61992"`, -/// `(-250, 3) -> "-0.25"`, `(25, 5) -> "0.00025"`. -fn decimal_to_string(mantissa: T, scale: u32) -> String { - let raw = mantissa.to_string(); - if scale == 0 { - return raw; - } - - let scale = scale as usize; - let (sign, digits) = match raw.strip_prefix('-') { - Some(rest) => ("-", rest), - None => ("", raw.as_str()), - }; - - let (int_part, frac) = if digits.len() > scale { - let (int_part, frac) = digits.split_at(digits.len() - scale); - (int_part, frac.to_string()) - } else { - let pad = "0".repeat(scale - digits.len()); - ("0", format!("{pad}{digits}")) - }; - - let frac = frac.trim_end_matches('0'); - if frac.is_empty() { - format!("{sign}{int_part}") - } else { - format!("{sign}{int_part}.{frac}") - } -} - -/// Append every element of an Arrow `array` to a column accumulator, converting -/// each value to [`DBResponsePrimitive`]. -fn append_arrow_array( - col: &mut Vec, - array: &dyn Array, -) -> Result<(), ParseError> { - let len = array.len(); - col.reserve(len); - - macro_rules! downcast_array_ref { - ($ty:ty) => { - array.as_any().downcast_ref::<$ty>().ok_or_else(|| { - ParseError::ArrowError(format!( - "Failed to downcast Arrow array to {}", - stringify!($ty) - )) - })? - }; - } - - macro_rules! push_int { - ($ty:ty) => {{ - let a = downcast_array_ref!($ty); - for i in 0..len { - if a.is_null(i) { - col.push(DBResponsePrimitive::Null); - } else { - col.push(DBResponsePrimitive::Int64(a.value(i) as i64)); - } - } - }}; - } - - macro_rules! push_uint { - ($ty:ty) => {{ - let a = downcast_array_ref!($ty); - for i in 0..len { - if a.is_null(i) { - col.push(DBResponsePrimitive::Null); - } else { - col.push(DBResponsePrimitive::UInt64(a.value(i) as u64)); - } - } - }}; - } - - macro_rules! push_float { - ($ty:ty) => {{ - let a = downcast_array_ref!($ty); - for i in 0..len { - if a.is_null(i) { - col.push(DBResponsePrimitive::Null); - } else { - col.push(DBResponsePrimitive::Float64(a.value(i) as f64)); - } - } - }}; - } - - macro_rules! push_str { - ($ty:ty) => {{ - let a = downcast_array_ref!($ty); - for i in 0..len { - if a.is_null(i) { - col.push(DBResponsePrimitive::Null); - } else { - col.push(DBResponsePrimitive::String(a.value(i).to_owned())); - } - } - }}; - } - - macro_rules! push_datetime { - ($ty:ty) => {{ - let a = downcast_array_ref!($ty); - for i in 0..len { - if a.is_null(i) { - col.push(DBResponsePrimitive::Null); - } else { - match a.value_as_datetime(i) { - Some(dt) => col.push(DBResponsePrimitive::Timestamp(dt)), - None => col.push(DBResponsePrimitive::Null), - } + let mut data = Vec::with_capacity(n_cols); + for (idx, col_chunks) in chunks.into_iter().enumerate() { + // One array per column, so the read path downcasts once for the whole + // column. A stream with a single batch — everything CubeStore sends + // outside its streaming path — needs no copy at all. + let array: ArrayRef = match col_chunks.len() { + 0 => new_empty_array(schema.field(idx).data_type()), + 1 => col_chunks.into_iter().next().unwrap(), + _ => { + let refs: Vec<&dyn Array> = + col_chunks.iter().map(|chunk| chunk.as_ref()).collect(); + concat(&refs).map_err(|err| ParseError::ArrowError(err.to_string()))? } - } - }}; - } + }; - // Format decimal arrays from their mantissa and scale. `decimal_to_string` is - // generic over the mantissa width, so one path handles Decimal128/256. - macro_rules! push_decimal { - ($ty:ty) => {{ - let a = downcast_array_ref!($ty); - let scale = a.scale().max(0) as u32; - for i in 0..len { - if a.is_null(i) { - col.push(DBResponsePrimitive::Null); - } else { - col.push(DBResponsePrimitive::String(decimal_to_string( - a.value(i), - scale, - ))); - } - } - }}; - } - - match array.data_type() { - DataType::Null => { - for _ in 0..len { - col.push(DBResponsePrimitive::Null); - } + data.push(QueryResultColumn::Arrow(ArrowArray::try_new(array)?)); } - DataType::Boolean => { - let a = downcast_array_ref!(BooleanArray); - for i in 0..len { - if a.is_null(i) { - col.push(DBResponsePrimitive::Null); - } else { - col.push(DBResponsePrimitive::Boolean(a.value(i))); - } - } - } - DataType::Int8 => push_int!(Int8Array), - DataType::Int16 => push_int!(Int16Array), - DataType::Int32 => push_int!(Int32Array), - DataType::Int64 => push_int!(Int64Array), - DataType::UInt8 => push_uint!(UInt8Array), - DataType::UInt16 => push_uint!(UInt16Array), - DataType::UInt32 => push_uint!(UInt32Array), - DataType::UInt64 => push_uint!(UInt64Array), - DataType::Float32 => push_float!(Float32Array), - DataType::Float64 => push_float!(Float64Array), - DataType::Float16 => { - let a = downcast_array_ref!(Float16Array); - for i in 0..len { - if a.is_null(i) { - col.push(DBResponsePrimitive::Null); - } else { - col.push(DBResponsePrimitive::Float64(a.value(i).to_f64())); - } - } - } - DataType::Utf8 => push_str!(StringArray), - DataType::LargeUtf8 => push_str!(LargeStringArray), - DataType::Utf8View => push_str!(StringViewArray), - DataType::Date32 => push_datetime!(Date32Array), - DataType::Date64 => push_datetime!(Date64Array), - DataType::Timestamp(TimeUnit::Second, _) => push_datetime!(TimestampSecondArray), - DataType::Timestamp(TimeUnit::Millisecond, _) => push_datetime!(TimestampMillisecondArray), - DataType::Timestamp(TimeUnit::Microsecond, _) => push_datetime!(TimestampMicrosecondArray), - DataType::Timestamp(TimeUnit::Nanosecond, _) => push_datetime!(TimestampNanosecondArray), - DataType::Decimal128(_, _) => push_decimal!(Decimal128Array), - DataType::Decimal256(_, _) => push_decimal!(Decimal256Array), - other => return Err(ParseError::UnsupportedArrowType(format!("{other:?}"))), - } - Ok(()) + QueryResult::try_new(members, data) + } } #[cfg(test)] mod tests { use super::*; - use arrow::array::BinaryArray; - use arrow::datatypes::{Field, Schema}; + use arrow::array::{ + BinaryArray, Decimal128Array, Decimal256Array, Float64Array, StringArray, + TimestampMillisecondArray, + }; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use arrow::ipc::writer::StreamWriter; use arrow::record_batch::RecordBatch; use cubeshared::codegen::{ @@ -509,6 +348,12 @@ mod tests { use cubeshared::flatbuffers::FlatBufferBuilder; use std::sync::Arc; + /// Materialize one column so assertions can compare against a plain slice, + /// whatever the column's backing storage is. + fn column(result: &QueryResult, idx: usize) -> ColumnarArray { + result.column(idx).unwrap().to_columnar().unwrap() + } + /// Helper function to create a test HttpMessage with a given number of rows and columns fn create_test_message(num_rows: usize, num_columns: usize) -> Vec { let mut builder = FlatBufferBuilder::new(); @@ -717,7 +562,7 @@ mod tests { assert_eq!(result.data.len(), 3); assert_eq!( - result.data[0].as_slice(), + column(&result, 0).as_slice(), &[ DBResponsePrimitive::String("Berlin".to_string()), DBResponsePrimitive::Null, @@ -725,7 +570,7 @@ mod tests { ] ); assert_eq!( - result.data[1].as_slice(), + column(&result, 1).as_slice(), &[ DBResponsePrimitive::Float64(1.5), DBResponsePrimitive::Float64(2.0), @@ -734,18 +579,18 @@ mod tests { ); // Numeric values serialize as JSON strings, matching the legacy result set. - let amounts_json = serde_json::to_value(result.data[1].as_slice()).unwrap(); + let amounts_json = serde_json::to_value(column(&result, 1).as_slice()).unwrap(); assert_eq!(amounts_json[0], "1.5"); assert_eq!(amounts_json[1], "2"); assert_eq!(amounts_json[2], serde_json::Value::Null); // Timestamps land in the dedicated variant and serialize to the ISO format. - match &result.data[2].as_slice()[0] { + match &column(&result, 2).as_slice()[0] { DBResponsePrimitive::Timestamp(_) => {} other => panic!("expected Timestamp, got {other:?}"), } - assert_eq!(result.data[2].as_slice()[1], DBResponsePrimitive::Null); - let json = serde_json::to_value(result.data[2].as_slice()).unwrap(); + assert_eq!(column(&result, 2).as_slice()[1], DBResponsePrimitive::Null); + let json = serde_json::to_value(column(&result, 2).as_slice()).unwrap(); assert_eq!(json[0], "1970-01-01T00:00:00.000"); assert_eq!(json[1], serde_json::Value::Null); assert_eq!(json[2], "1970-01-01T00:00:01.000"); @@ -774,7 +619,7 @@ mod tests { let result = QueryResult::from_arrow(&bytes)?; assert_eq!( - result.data[0].as_slice(), + column(&result, 0).as_slice(), &[ DBResponsePrimitive::String("9999999999999999999999999999999999999.99".to_string()), DBResponsePrimitive::Null, @@ -807,7 +652,7 @@ mod tests { let result = QueryResult::from_arrow(&bytes)?; assert_eq!(result.row_count, 5); - let json = serde_json::to_value(result.data[0].as_slice()).unwrap(); + let json = serde_json::to_value(column(&result, 0).as_slice()).unwrap(); assert_eq!(json[0], serde_json::Value::Null); assert_eq!(json[1], "2399.96"); assert_eq!(json[2], "2249.91"); @@ -817,30 +662,6 @@ mod tests { Ok(()) } - #[test] - fn test_decimal_to_string() { - for (mantissa, scale, expected) in [ - (6199200000i128, 5u32, "61992"), - (25987600, 5, "259.876"), - (1500, 3, "1.5"), - (-250, 3, "-0.25"), - (0, 5, "0"), - (21098000, 5, "210.98"), - (100, 0, "100"), - (0, 0, "0"), - (-5, 0, "-5"), - (25, 5, "0.00025"), - (-1, 0, "-1"), - (i128::MAX, 0, "170141183460469231731687303715884105727"), - ] { - assert_eq!( - decimal_to_string(mantissa, scale), - expected, - "mantissa={mantissa} scale={scale}" - ); - } - } - #[test] fn test_from_arrow_unsupported_type() -> Result<(), ParseError> { let schema = Arc::new(Schema::new(vec![Field::new( @@ -858,6 +679,118 @@ mod tests { Ok(()) } + /// An unsupported type with no rows still parses: no cell is ever read, and + /// the pre-Arrow-column code reached the same conclusion by never entering + /// its per-cell loop. + #[test] + fn test_from_arrow_unsupported_type_zero_rows() -> Result<(), ParseError> { + let schema = Arc::new(Schema::new(vec![Field::new( + "blob", + DataType::Binary, + false, + )])); + let blobs: Vec<&[u8]> = vec![]; + let batch = RecordBatch::try_new(schema, vec![Arc::new(BinaryArray::from_vec(blobs))]) + .expect("empty batch"); + + let bytes = arrow_ipc_bytes(&batch); + let result = QueryResult::from_arrow(&bytes)?; + assert_eq!(result.members, vec!["blob"]); + assert_eq!(result.row_count, 0); + + Ok(()) + } + + /// A stream with more than one batch is concatenated per column, so values + /// stay in stream order across the chunk boundary. + #[test] + fn test_from_arrow_multiple_batches() -> Result<(), ParseError> { + let schema = Arc::new(Schema::new(vec![ + Field::new("city", DataType::Utf8, true), + Field::new("amount", DataType::Float64, true), + ])); + + let batches = [ + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![Some("Berlin"), None])), + Arc::new(Float64Array::from(vec![Some(1.5), Some(2.0)])), + ], + ) + .unwrap(), + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![Some("Lisbon")])), + Arc::new(Float64Array::from(vec![None])), + ], + ) + .unwrap(), + ]; + + let mut bytes = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut bytes, schema.as_ref()).unwrap(); + for batch in &batches { + writer.write(batch).unwrap(); + } + writer.finish().unwrap(); + } + + let result = QueryResult::from_arrow(&bytes)?; + assert_eq!(result.row_count, 3); + assert_eq!( + column(&result, 0).as_slice(), + &[ + DBResponsePrimitive::String("Berlin".to_string()), + DBResponsePrimitive::Null, + DBResponsePrimitive::String("Lisbon".to_string()), + ] + ); + assert_eq!( + column(&result, 1).as_slice(), + &[ + DBResponsePrimitive::Float64(1.5), + DBResponsePrimitive::Float64(2.0), + DBResponsePrimitive::Null, + ] + ); + + // Reading through the reader must agree with the materialized column. + let reader = result.reader(0)?; + assert_eq!( + reader.value(2), + DBResponsePrimitive::String("Lisbon".to_string()) + ); + assert_eq!(reader.value(1), DBResponsePrimitive::Null); + + Ok(()) + } + + /// A schema with no batches at all yields empty columns, not a missing column. + #[test] + fn test_from_arrow_no_batches() -> Result<(), ParseError> { + let schema = Arc::new(Schema::new(vec![ + Field::new("city", DataType::Utf8, true), + Field::new("amount", DataType::Float64, true), + ])); + + let mut bytes = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut bytes, schema.as_ref()).unwrap(); + writer.finish().unwrap(); + } + + let result = QueryResult::from_arrow(&bytes)?; + assert_eq!(result.members, vec!["city", "amount"]); + assert_eq!(result.row_count, 0); + assert_eq!(result.data.len(), 2); + assert!(result.data.iter().all(|c| c.is_empty())); + + Ok(()) + } + #[test] fn test_from_cubestore_fb_arrow_query_result() -> Result<(), ParseError> { // Arrow IPC stream payload, as CubeStore would emit it. @@ -909,14 +842,14 @@ mod tests { assert_eq!(result.members, vec!["city", "amount"]); assert_eq!(result.row_count, 2); assert_eq!( - result.data[0].as_slice(), + column(&result, 0).as_slice(), &[ DBResponsePrimitive::String("Berlin".to_string()), DBResponsePrimitive::String("Lisbon".to_string()), ] ); assert_eq!( - result.data[1].as_slice(), + column(&result, 1).as_slice(), &[ DBResponsePrimitive::Float64(1.5), DBResponsePrimitive::Float64(2.0), diff --git a/rust/cube/cubeorchestrator/src/query_result_column.rs b/rust/cube/cubeorchestrator/src/query_result_column.rs new file mode 100644 index 0000000000000..130b09d98fb85 --- /dev/null +++ b/rust/cube/cubeorchestrator/src/query_result_column.rs @@ -0,0 +1,648 @@ +//! Column storage for [`crate::query_message_parser::QueryResult`]. +//! +//! A column arrives either already materialized as [`DBResponsePrimitive`] cells +//! (the legacy CubeStore `HttpResultSet` and the JS→Rust columnar transport) or as +//! Arrow memory straight off the wire (`HttpQueryResultArrow`). Arrow columns are +//! kept as-is and decoded per cell on read, so a result set is never materialized +//! just to be re-cloned into the transform output. + +use crate::{ + query_message_parser::ParseError, + query_result_transform::{ + is_identity_transform, transform_value, ColumnarArray, DBResponsePrimitive, + }, +}; +use arrow::array::{ + Array, ArrayRef, BooleanArray, Date32Array, Date64Array, Decimal128Array, Decimal256Array, + Float16Array, Float32Array, Float64Array, Int16Array, Int32Array, Int64Array, Int8Array, + LargeStringArray, PrimitiveArray, StringArray, StringViewArray, TimestampMicrosecondArray, + TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt16Array, + UInt32Array, UInt64Array, UInt8Array, +}; +use arrow::datatypes::{ArrowTemporalType, DataType, TimeUnit}; +use std::{borrow::Cow, convert::Infallible}; + +/// One logical column of a query result. +#[derive(Debug, Clone)] +pub enum QueryResultColumn { + /// Cells already materialized as primitives: legacy `HttpResultSet` rows and + /// `JsRawColumnarData` coming from the JS drivers. + Columnar(ColumnarArray), + /// Cells still held in Arrow memory, decoded on read. + Arrow(ArrowArray), +} + +impl QueryResultColumn { + #[inline] + pub fn len(&self) -> usize { + match self { + QueryResultColumn::Columnar(c) => c.len(), + QueryResultColumn::Arrow(a) => a.len(), + } + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Per-cell accessor for this column. Resolving it once per column keeps the + /// Arrow type dispatch out of the per-cell path. + pub fn reader(&self) -> Result, ParseError> { + match self { + QueryResultColumn::Columnar(c) => Ok(ColumnReader::Primitives(c.as_slice())), + QueryResultColumn::Arrow(a) => Ok(ColumnReader::Arrow(a.cell_reader()?)), + } + } + + /// Materialize the column as primitives. Only for callers that genuinely need + /// an owned slice — the read paths use [`QueryResultColumn::reader`] instead. + pub fn to_columnar(&self) -> Result { + match self { + QueryResultColumn::Columnar(c) => Ok(c.clone()), + QueryResultColumn::Arrow(a) => a.to_columnar(), + } + } +} + +impl From for QueryResultColumn { + #[inline] + fn from(c: ColumnarArray) -> Self { + QueryResultColumn::Columnar(c) + } +} + +impl From> for QueryResultColumn { + #[inline] + fn from(v: Vec) -> Self { + QueryResultColumn::Columnar(ColumnarArray::from(v)) + } +} + +impl From for QueryResultColumn { + #[inline] + fn from(a: ArrowArray) -> Self { + QueryResultColumn::Arrow(a) + } +} + +/// A single logical column backed by Arrow memory. +#[derive(Debug, Clone)] +pub struct ArrowArray(ArrayRef); + +impl ArrowArray { + /// Wrap an Arrow array, rejecting types the cell reader cannot decode. + /// + /// Validation goes through [`ArrowArray::cell_reader`] so the set of supported + /// types has a single definition, and an unsupported type is reported while + /// parsing rather than halfway through a transform. Empty arrays are accepted + /// whatever their type: no cell is ever read, matching the previous behaviour + /// where a zero-row column of an unsupported type parsed fine because the + /// per-cell loop never ran. + pub fn try_new(array: ArrayRef) -> Result { + let this = Self(array); + this.cell_reader()?; + Ok(this) + } + + #[inline] + pub fn len(&self) -> usize { + self.0.len() + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + #[inline] + pub fn data_type(&self) -> &DataType { + self.0.data_type() + } + + #[inline] + pub fn array(&self) -> &ArrayRef { + &self.0 + } + + pub fn to_columnar(&self) -> Result { + let mut out = ColumnarArray::with_capacity(self.len()); + // `member_type` is empty: materializing must not apply the `time` reformat + // that a transform would. + let Ok(()) = self + .cell_reader()? + .for_each_transformed::("", |value| { + out.push(value.into_owned()); + Ok(()) + }); + Ok(out) + } + + /// Downcast to the concrete Arrow array once, for the whole column. + fn cell_reader(&self) -> Result, ParseError> { + let array = self.0.as_ref(); + + if array.is_empty() { + return Ok(ArrowCellReader::Empty); + } + + macro_rules! downcast { + ($variant:ident, $ty:ty) => {{ + let a = array.as_any().downcast_ref::<$ty>().ok_or_else(|| { + ParseError::ArrowError(format!( + "Failed to downcast Arrow array to {}", + stringify!($ty) + )) + })?; + ArrowCellReader::$variant(a) + }}; + } + + macro_rules! downcast_decimal { + ($variant:ident, $ty:ty) => {{ + let a = array.as_any().downcast_ref::<$ty>().ok_or_else(|| { + ParseError::ArrowError(format!( + "Failed to downcast Arrow array to {}", + stringify!($ty) + )) + })?; + ArrowCellReader::$variant(a, a.scale().max(0) as u32) + }}; + } + + let reader = match array.data_type() { + DataType::Null => ArrowCellReader::Null(array.len()), + DataType::Boolean => downcast!(Boolean, BooleanArray), + DataType::Int8 => downcast!(Int8, Int8Array), + DataType::Int16 => downcast!(Int16, Int16Array), + DataType::Int32 => downcast!(Int32, Int32Array), + DataType::Int64 => downcast!(Int64, Int64Array), + DataType::UInt8 => downcast!(UInt8, UInt8Array), + DataType::UInt16 => downcast!(UInt16, UInt16Array), + DataType::UInt32 => downcast!(UInt32, UInt32Array), + DataType::UInt64 => downcast!(UInt64, UInt64Array), + DataType::Float16 => downcast!(Float16, Float16Array), + DataType::Float32 => downcast!(Float32, Float32Array), + DataType::Float64 => downcast!(Float64, Float64Array), + DataType::Utf8 => downcast!(Utf8, StringArray), + DataType::LargeUtf8 => downcast!(LargeUtf8, LargeStringArray), + DataType::Utf8View => downcast!(Utf8View, StringViewArray), + DataType::Date32 => downcast!(Date32, Date32Array), + DataType::Date64 => downcast!(Date64, Date64Array), + DataType::Timestamp(TimeUnit::Second, _) => { + downcast!(TimestampSecond, TimestampSecondArray) + } + DataType::Timestamp(TimeUnit::Millisecond, _) => { + downcast!(TimestampMillisecond, TimestampMillisecondArray) + } + DataType::Timestamp(TimeUnit::Microsecond, _) => { + downcast!(TimestampMicrosecond, TimestampMicrosecondArray) + } + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + downcast!(TimestampNanosecond, TimestampNanosecondArray) + } + DataType::Decimal128(_, _) => downcast_decimal!(Decimal128, Decimal128Array), + DataType::Decimal256(_, _) => downcast_decimal!(Decimal256, Decimal256Array), + other => return Err(ParseError::UnsupportedArrowType(format!("{other:?}"))), + }; + + Ok(reader) + } +} + +/// Per-cell accessor for one column, resolved once per column so the per-cell path +/// is a jump table plus an index — never a `DataType` match and a downcast. +pub enum ColumnReader<'a> { + Primitives(&'a [DBResponsePrimitive]), + Arrow(ArrowCellReader<'a>), +} + +impl ColumnReader<'_> { + #[inline] + pub fn len(&self) -> usize { + match self { + ColumnReader::Primitives(s) => s.len(), + ColumnReader::Arrow(a) => a.len(), + } + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Value at `row`. Panics when out of range, like the slice indexing it + /// replaces — `QueryResult::try_new` guarantees every column holds `row_count` + /// cells. + #[inline] + pub fn value(&self, row: usize) -> DBResponsePrimitive { + match self { + ColumnReader::Primitives(s) => s[row].clone(), + ColumnReader::Arrow(a) => a.value(row), + } + } + + /// Value at `row` rendered as text, or `None` when the cell is null. Borrows + /// when the underlying storage already holds a string, so string columns cross + /// into JS without an intermediate allocation. + pub fn value_as_str(&self, row: usize) -> Option> { + match self { + ColumnReader::Primitives(s) => match &s[row] { + DBResponsePrimitive::Null => None, + DBResponsePrimitive::String(s) => Some(Cow::Borrowed(s)), + other => Some(Cow::Owned(other.to_string())), + }, + ColumnReader::Arrow(a) => a.value_as_str(row), + } + } + + /// The cell at `row`, transformed for `member_type`, passed to `f` by + /// reference. A materialized column whose member type needs no transform is + /// handed over borrowed, so a reader that only forwards the value — the + /// row-major serializers — never clones it. + #[inline] + pub fn with_transformed( + &self, + row: usize, + member_type: &str, + f: impl FnOnce(&DBResponsePrimitive) -> R, + ) -> R { + match self { + ColumnReader::Primitives(cells) if is_identity_transform(member_type) => f(&cells[row]), + _ => f(&transform_value(self.value(row), member_type)), + } + } + + /// Hand every cell, transformed for `member_type`, to `visit` in row order. + /// The Arrow type dispatch happens once, outside the row loop — so prefer this + /// over [`ColumnReader::value`] whenever a whole column is being consumed. + /// Cells that need no transform arrive borrowed, the rest owned. + pub fn for_each_transformed( + &self, + member_type: &str, + mut visit: impl FnMut(Cow<'_, DBResponsePrimitive>) -> Result<(), E>, + ) -> Result<(), E> { + match self { + ColumnReader::Primitives(cells) => { + if is_identity_transform(member_type) { + for cell in cells.iter() { + visit(Cow::Borrowed(cell))?; + } + } else { + for cell in cells.iter() { + visit(Cow::Owned(transform_value(cell.clone(), member_type)))?; + } + } + Ok(()) + } + ColumnReader::Arrow(a) => a.for_each_transformed(member_type, visit), + } + } + + /// Append every cell, transformed for `member_type`, to `out`. + pub fn append_transformed(&self, out: &mut ColumnarArray, member_type: &str) { + out.reserve(self.len()); + + let Ok(()) = self.for_each_transformed::(member_type, |value| { + out.push(value.into_owned()); + Ok(()) + }); + } +} + +/// The concrete Arrow array behind one column. +pub enum ArrowCellReader<'a> { + /// Zero-length column of any type — no cell is ever read. + Empty, + /// `DataType::Null` column, which carries nothing but its row count. + Null(usize), + Boolean(&'a BooleanArray), + Int8(&'a Int8Array), + Int16(&'a Int16Array), + Int32(&'a Int32Array), + Int64(&'a Int64Array), + UInt8(&'a UInt8Array), + UInt16(&'a UInt16Array), + UInt32(&'a UInt32Array), + UInt64(&'a UInt64Array), + Float16(&'a Float16Array), + Float32(&'a Float32Array), + Float64(&'a Float64Array), + Utf8(&'a StringArray), + LargeUtf8(&'a LargeStringArray), + Utf8View(&'a StringViewArray), + Date32(&'a Date32Array), + Date64(&'a Date64Array), + TimestampSecond(&'a TimestampSecondArray), + TimestampMillisecond(&'a TimestampMillisecondArray), + TimestampMicrosecond(&'a TimestampMicrosecondArray), + TimestampNanosecond(&'a TimestampNanosecondArray), + /// Array plus its scale, used to render the mantissa as a decimal string. + Decimal128(&'a Decimal128Array, u32), + Decimal256(&'a Decimal256Array, u32), +} + +/// Read `$array[$row]` as `$make(value)`, or `Null` when the cell is null. +macro_rules! read_cell { + ($array:expr, $row:expr, $make:expr) => {{ + let a = $array; + if a.is_null($row) { + DBResponsePrimitive::Null + } else { + #[allow(clippy::redundant_closure_call)] + ($make)(a.value($row)) + } + }}; +} + +impl ArrowCellReader<'_> { + pub fn len(&self) -> usize { + match self { + ArrowCellReader::Empty => 0, + ArrowCellReader::Null(len) => *len, + ArrowCellReader::Boolean(a) => a.len(), + ArrowCellReader::Int8(a) => a.len(), + ArrowCellReader::Int16(a) => a.len(), + ArrowCellReader::Int32(a) => a.len(), + ArrowCellReader::Int64(a) => a.len(), + ArrowCellReader::UInt8(a) => a.len(), + ArrowCellReader::UInt16(a) => a.len(), + ArrowCellReader::UInt32(a) => a.len(), + ArrowCellReader::UInt64(a) => a.len(), + ArrowCellReader::Float16(a) => a.len(), + ArrowCellReader::Float32(a) => a.len(), + ArrowCellReader::Float64(a) => a.len(), + ArrowCellReader::Utf8(a) => a.len(), + ArrowCellReader::LargeUtf8(a) => a.len(), + ArrowCellReader::Utf8View(a) => a.len(), + ArrowCellReader::Date32(a) => a.len(), + ArrowCellReader::Date64(a) => a.len(), + ArrowCellReader::TimestampSecond(a) => a.len(), + ArrowCellReader::TimestampMillisecond(a) => a.len(), + ArrowCellReader::TimestampMicrosecond(a) => a.len(), + ArrowCellReader::TimestampNanosecond(a) => a.len(), + ArrowCellReader::Decimal128(a, _) => a.len(), + ArrowCellReader::Decimal256(a, _) => a.len(), + } + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + #[inline] + pub fn value(&self, row: usize) -> DBResponsePrimitive { + match self { + // `Empty` columns hold no cells and `Null` columns hold nothing but + // nulls, so neither needs a bounds check to answer. + ArrowCellReader::Empty | ArrowCellReader::Null(_) => DBResponsePrimitive::Null, + ArrowCellReader::Boolean(a) => read_cell!(a, row, DBResponsePrimitive::Boolean), + ArrowCellReader::Int8(a) => { + read_cell!(a, row, |v| DBResponsePrimitive::Int64(v as i64)) + } + ArrowCellReader::Int16(a) => { + read_cell!(a, row, |v| DBResponsePrimitive::Int64(v as i64)) + } + ArrowCellReader::Int32(a) => { + read_cell!(a, row, |v| DBResponsePrimitive::Int64(v as i64)) + } + ArrowCellReader::Int64(a) => read_cell!(a, row, DBResponsePrimitive::Int64), + ArrowCellReader::UInt8(a) => { + read_cell!(a, row, |v| DBResponsePrimitive::UInt64(v as u64)) + } + ArrowCellReader::UInt16(a) => { + read_cell!(a, row, |v| DBResponsePrimitive::UInt64(v as u64)) + } + ArrowCellReader::UInt32(a) => { + read_cell!(a, row, |v| DBResponsePrimitive::UInt64(v as u64)) + } + ArrowCellReader::UInt64(a) => read_cell!(a, row, DBResponsePrimitive::UInt64), + // `half` is not a direct dependency, so the closure form of + // `read_cell!` can't name `f16` for inference — spell the read out. + ArrowCellReader::Float16(a) => { + if a.is_null(row) { + DBResponsePrimitive::Null + } else { + DBResponsePrimitive::Float64(a.value(row).to_f64()) + } + } + ArrowCellReader::Float32(a) => { + read_cell!(a, row, |v| DBResponsePrimitive::Float64(v as f64)) + } + ArrowCellReader::Float64(a) => read_cell!(a, row, DBResponsePrimitive::Float64), + ArrowCellReader::Utf8(a) => { + read_cell!(a, row, |v: &str| DBResponsePrimitive::String(v.to_owned())) + } + ArrowCellReader::LargeUtf8(a) => { + read_cell!(a, row, |v: &str| DBResponsePrimitive::String(v.to_owned())) + } + ArrowCellReader::Utf8View(a) => { + read_cell!(a, row, |v: &str| DBResponsePrimitive::String(v.to_owned())) + } + ArrowCellReader::Date32(a) => datetime_cell(*a, row), + ArrowCellReader::Date64(a) => datetime_cell(*a, row), + ArrowCellReader::TimestampSecond(a) => datetime_cell(*a, row), + ArrowCellReader::TimestampMillisecond(a) => datetime_cell(*a, row), + ArrowCellReader::TimestampMicrosecond(a) => datetime_cell(*a, row), + ArrowCellReader::TimestampNanosecond(a) => datetime_cell(*a, row), + ArrowCellReader::Decimal128(a, scale) => { + read_cell!(a, row, |v| DBResponsePrimitive::String(decimal_to_string( + v, *scale + ))) + } + ArrowCellReader::Decimal256(a, scale) => { + read_cell!(a, row, |v| DBResponsePrimitive::String(decimal_to_string( + v, *scale + ))) + } + } + } + + /// Text rendering of `row`, borrowing from the Arrow buffer for string types. + fn value_as_str(&self, row: usize) -> Option> { + macro_rules! borrowed_str { + ($array:expr) => {{ + let a = $array; + if a.is_null(row) { + None + } else { + Some(Cow::Borrowed(a.value(row))) + } + }}; + } + + match self { + ArrowCellReader::Utf8(a) => borrowed_str!(a), + ArrowCellReader::LargeUtf8(a) => borrowed_str!(a), + ArrowCellReader::Utf8View(a) => borrowed_str!(a), + other => match other.value(row) { + DBResponsePrimitive::Null => None, + // A freshly decoded primitive, so `String` can be unwrapped + // instead of re-rendered through `Display`. + DBResponsePrimitive::String(s) => Some(Cow::Owned(s)), + value => Some(Cow::Owned(value.to_string())), + }, + } + } + + /// Column-major read with the type match hoisted out of the row loop. + fn for_each_transformed( + &self, + member_type: &str, + mut visit: impl FnMut(Cow<'_, DBResponsePrimitive>) -> Result<(), E>, + ) -> Result<(), E> { + let len = self.len(); + + macro_rules! fill { + ($read:expr) => {{ + for row in 0..len { + visit(Cow::Owned(transform_value($read(row), member_type)))?; + } + }}; + } + + macro_rules! fill_with { + ($array:expr, $make:expr) => {{ + let a = $array; + fill!(|row| read_cell!(a, row, $make)) + }}; + } + + match self { + ArrowCellReader::Empty => {} + ArrowCellReader::Null(_) => { + for _ in 0..len { + visit(Cow::Owned(DBResponsePrimitive::Null))?; + } + } + ArrowCellReader::Boolean(a) => fill_with!(a, DBResponsePrimitive::Boolean), + ArrowCellReader::Int8(a) => fill_with!(a, |v| DBResponsePrimitive::Int64(v as i64)), + ArrowCellReader::Int16(a) => fill_with!(a, |v| DBResponsePrimitive::Int64(v as i64)), + ArrowCellReader::Int32(a) => fill_with!(a, |v| DBResponsePrimitive::Int64(v as i64)), + ArrowCellReader::Int64(a) => fill_with!(a, DBResponsePrimitive::Int64), + ArrowCellReader::UInt8(a) => fill_with!(a, |v| DBResponsePrimitive::UInt64(v as u64)), + ArrowCellReader::UInt16(a) => fill_with!(a, |v| DBResponsePrimitive::UInt64(v as u64)), + ArrowCellReader::UInt32(a) => fill_with!(a, |v| DBResponsePrimitive::UInt64(v as u64)), + ArrowCellReader::UInt64(a) => fill_with!(a, DBResponsePrimitive::UInt64), + ArrowCellReader::Float16(a) => fill!(|row| if a.is_null(row) { + DBResponsePrimitive::Null + } else { + DBResponsePrimitive::Float64(a.value(row).to_f64()) + }), + ArrowCellReader::Float32(a) => { + fill_with!(a, |v| DBResponsePrimitive::Float64(v as f64)) + } + ArrowCellReader::Float64(a) => fill_with!(a, DBResponsePrimitive::Float64), + ArrowCellReader::Utf8(a) => { + fill_with!(a, |v: &str| DBResponsePrimitive::String(v.to_owned())) + } + ArrowCellReader::LargeUtf8(a) => { + fill_with!(a, |v: &str| DBResponsePrimitive::String(v.to_owned())) + } + ArrowCellReader::Utf8View(a) => { + fill_with!(a, |v: &str| DBResponsePrimitive::String(v.to_owned())) + } + ArrowCellReader::Date32(a) => fill!(|row| datetime_cell(*a, row)), + ArrowCellReader::Date64(a) => fill!(|row| datetime_cell(*a, row)), + ArrowCellReader::TimestampSecond(a) => fill!(|row| datetime_cell(*a, row)), + ArrowCellReader::TimestampMillisecond(a) => fill!(|row| datetime_cell(*a, row)), + ArrowCellReader::TimestampMicrosecond(a) => fill!(|row| datetime_cell(*a, row)), + ArrowCellReader::TimestampNanosecond(a) => fill!(|row| datetime_cell(*a, row)), + ArrowCellReader::Decimal128(a, scale) => { + fill_with!(a, |v| DBResponsePrimitive::String(decimal_to_string( + v, *scale + ))) + } + ArrowCellReader::Decimal256(a, scale) => { + fill_with!(a, |v| DBResponsePrimitive::String(decimal_to_string( + v, *scale + ))) + } + } + + Ok(()) + } +} + +/// Read a date/timestamp cell as [`DBResponsePrimitive::Timestamp`]. The timezone +/// an Arrow timestamp field may carry is ignored, as it was before this refactor. +#[inline] +fn datetime_cell(array: &PrimitiveArray, row: usize) -> DBResponsePrimitive +where + T: ArrowTemporalType, + i64: From, +{ + if array.is_null(row) { + return DBResponsePrimitive::Null; + } + + match array.value_as_datetime(row) { + Some(dt) => DBResponsePrimitive::Timestamp(dt), + None => DBResponsePrimitive::Null, + } +} + +/// Format a decimal `mantissa` with `scale` fractional digits, stripping trailing +/// fractional zeros. Generic over the mantissa's `Display`, so it renders any Arrow +/// decimal width (`i32`/`i64`/`i128`/`i256`) directly — Decimal256 needs no fallback +/// to Arrow's own string conversion. +/// +/// e.g. `(25987600, 5) -> "259.876"`, `(6199200000, 5) -> "61992"`, +/// `(-250, 3) -> "-0.25"`, `(25, 5) -> "0.00025"`. +pub(crate) fn decimal_to_string(mantissa: T, scale: u32) -> String { + let raw = mantissa.to_string(); + if scale == 0 { + return raw; + } + + let scale = scale as usize; + let (sign, digits) = match raw.strip_prefix('-') { + Some(rest) => ("-", rest), + None => ("", raw.as_str()), + }; + + let (int_part, frac) = if digits.len() > scale { + let (int_part, frac) = digits.split_at(digits.len() - scale); + (int_part, frac.to_string()) + } else { + let pad = "0".repeat(scale - digits.len()); + ("0", format!("{pad}{digits}")) + }; + + let frac = frac.trim_end_matches('0'); + if frac.is_empty() { + format!("{sign}{int_part}") + } else { + format!("{sign}{int_part}.{frac}") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_decimal_to_string() { + for (mantissa, scale, expected) in [ + (6199200000i128, 5u32, "61992"), + (25987600, 5, "259.876"), + (1500, 3, "1.5"), + (-250, 3, "-0.25"), + (0, 5, "0"), + (21098000, 5, "210.98"), + (100, 0, "100"), + (0, 0, "0"), + (-5, 0, "-5"), + (25, 5, "0.00025"), + (-1, 0, "-1"), + (i128::MAX, 0, "170141183460469231731687303715884105727"), + ] { + assert_eq!( + decimal_to_string(mantissa, scale), + expected, + "mantissa={mantissa} scale={scale}" + ); + } + } +} diff --git a/rust/cube/cubeorchestrator/src/query_result_transform.rs b/rust/cube/cubeorchestrator/src/query_result_transform.rs index 8be457a5c2830..88ea1cad73663 100644 --- a/rust/cube/cubeorchestrator/src/query_result_transform.rs +++ b/rust/cube/cubeorchestrator/src/query_result_transform.rs @@ -1,5 +1,6 @@ use crate::{ query_message_parser::QueryResult, + query_result_column::ColumnReader, transport::{ AnnotatedConfigItem, ConfigItem, MemberOrMemberExpression, MembersMap, NormalizedQuery, QueryTimeDimension, QueryType, ResultType, TransformDataRequest, @@ -174,6 +175,15 @@ pub fn empty_vanilla_row(capacity: usize) -> VanillaRow { IndexMap::with_capacity_and_hasher(capacity, PrehashedBuildHasher) } +/// True when [`transform_value`] would return `type_`'s values unchanged, so a +/// caller that already holds the value can skip the call — and the clone it needs. +/// Keep in step with [`transform_value`] below: it rewrites nothing but `String` +/// cells of `time` members. +#[inline] +pub fn is_identity_transform(type_: &str) -> bool { + type_ != "time" +} + /// Transform specified `value` with specified `type` to the network protocol type. pub fn transform_value(value: DBResponsePrimitive, type_: &str) -> DBResponsePrimitive { match value { @@ -452,17 +462,15 @@ pub fn get_members( /// One output cell in a compact row. Built once per request by /// [`build_compact_plan`] so the per-row materializer ([`get_compact_row`]) -/// only does a single bounds check (`column.get(row_idx)`) and the -/// [`transform_value`] call. The plan borrows the column slice directly, -/// eliminating the per-cell `db_data.data.get(col).and_then(...)` double -/// lookup the row-major loop would otherwise do on every cell. +/// only does a single column read and the [`transform_value`] call. The plan +/// holds the column's reader directly, eliminating the per-cell +/// `db_data.data.get(col).and_then(...)` double lookup the row-major loop would +/// otherwise do on every cell — and, for Arrow-backed columns, the type +/// dispatch and downcast too. pub(crate) enum CompactPlanEntry<'a> { - /// Read `column[row_idx]` and run [`transform_value`]. `column` is a slice - /// of the corresponding [`ColumnarArray`]; the fat pointer inlines - /// `(ptr, len)` so the per-cell access avoids the extra Vec metadata - /// indirection. + /// Read `column[row_idx]` and run [`transform_value`]. Cell { - column: &'a [DBResponsePrimitive], + column: ColumnReader<'a>, member_type: &'a str, }, /// Constant value replicated across every row (the @@ -471,7 +479,7 @@ pub(crate) enum CompactPlanEntry<'a> { } pub struct CompactPlan<'a> { - entries: Vec>, + pub(crate) entries: Vec>, } pub(crate) fn build_compact_plan<'a>( @@ -489,7 +497,7 @@ pub(crate) fn build_compact_plan<'a>( if let Some(alias) = members_to_alias_map.get(m) { if let Some(&column_index) = cube_store_result.columns_pos.get(alias) { entries.push(CompactPlanEntry::Cell { - column: cube_store_result.data[column_index].as_slice(), + column: cube_store_result.reader(column_index)?, member_type: annotation_item.member_type.as_deref().unwrap_or(""), }); } @@ -513,7 +521,7 @@ pub(crate) fn build_compact_plan<'a>( let member_type = annotation .get(alias) .map_or("", |a| a.member_type.as_deref().unwrap_or("")); - let column = cube_store_result.data[column_index].as_slice(); + let column = cube_store_result.reader(column_index)?; entries.push(CompactPlanEntry::Cell { column, member_type, @@ -528,8 +536,8 @@ pub(crate) fn build_compact_plan<'a>( } /// Convert DB response row to the compact output. The plan carries the -/// per-cell column slice directly, so this loop only does one bounds check -/// (`column.get(row_idx)`) per cell — no `db_data.data.get(col)` indirection. +/// per-cell column reader directly, so this loop only does one column read per +/// cell — no `db_data.data.get(col)` indirection. pub fn get_compact_row(plan: &CompactPlan<'_>, row_idx: usize) -> Vec { let mut row: Vec = Vec::with_capacity(plan.entries.len()); @@ -539,7 +547,7 @@ pub fn get_compact_row(plan: &CompactPlan<'_>, row_idx: usize) -> Vec { - row.push(transform_value(column[row_idx].clone(), member_type)); + row.push(transform_value(column.value(row_idx), member_type)); } CompactPlanEntry::Constant(v) => { row.push(v.clone()); @@ -552,14 +560,13 @@ pub fn get_compact_row(plan: &CompactPlan<'_>, row_idx: usize) -> Vec { - /// Slice of the corresponding [`ColumnarArray`]. Fat pointer inlines - /// `(ptr, len)`, so the per-cell access avoids the extra Vec metadata - /// indirection. - column: &'a [DBResponsePrimitive], + /// Reader for the corresponding column, resolved once so Arrow-backed + /// columns skip the type dispatch and downcast on every cell. + column: ColumnReader<'a>, /// Interned IndexMap key for this column with a pre-computed hash. /// Cloned via [`Arc::clone`] per row (atomic refcount inc). key: Arc, @@ -634,7 +641,7 @@ pub fn build_vanilla_plan<'a>( .push((track.level, Arc::clone(&key))); } - let column = cube_store_result.data[index].as_slice(); + let column = cube_store_result.reader(index)?; columns.push(VanillaColumnPlan { column, @@ -721,11 +728,11 @@ pub(crate) enum ColumnarColumnSource { } pub(crate) struct ColumnarColumnPlan<'a> { - member_type: &'a str, - source: ColumnarColumnSource, + pub(crate) member_type: &'a str, + pub(crate) source: ColumnarColumnSource, } -fn build_columnar_plan<'a>( +pub(crate) fn build_columnar_plan<'a>( members: &[String], members_to_alias_map: &IndexMap, annotation: &'a HashMap, @@ -802,7 +809,7 @@ fn build_columnar_plan<'a>( fn build_columnar_columns( plan: &[ColumnarColumnPlan<'_>], db_data: &QueryResult, -) -> Vec { +) -> Result> { let row_count = db_data.row_count; let mut columns: Vec = plan .iter() @@ -813,9 +820,11 @@ fn build_columnar_columns( let out = &mut columns[col_idx]; match &plan_entry.source { ColumnarColumnSource::DbColumn { index } => { - for cell in db_data.data[*index].iter() { - out.push(transform_value(cell.clone(), plan_entry.member_type)); - } + // Column-major, so the reader resolves its type once for the + // whole column and the row loop stays a tight typed fill. + db_data + .reader(*index)? + .append_transformed(out, plan_entry.member_type); } ColumnarColumnSource::Constant(v) => { out.resize(row_count, v.clone()); @@ -826,15 +835,15 @@ fn build_columnar_columns( } } - columns + Ok(columns) } /// Convert DB response object to the vanilla output format. Keys are /// pre-hashed [`InternedKey`] values shared via [`Arc::clone`] from the plan, /// turning per-cell hashing/key allocation into an atomic refcount inc. The -/// plan also carries the column slice directly, so the per-row loop does one -/// bounds check (`column.column.get(row_idx)`) per cell instead of the -/// `db_data.data.get(col).and_then(...)` double lookup. +/// plan also carries the column reader directly, so the per-row loop does one +/// column read per cell instead of the `db_data.data.get(col).and_then(...)` +/// double lookup. pub fn get_vanilla_row(plan: &VanillaPlan<'_>, row_idx: usize) -> Result { // +1 to cover the optional tail entry (compareDateRange / blending key). let mut row = IndexMap::with_capacity_and_hasher( @@ -843,7 +852,7 @@ pub fn get_vanilla_row(plan: &VanillaPlan<'_>, row_idx: usize) -> Result], - result_data: &mut [RequestResultData], -) -> Result<()> { - for (transform_data, cube_store_result, result) in multizip(( - transform_requests.iter(), - cube_store_results.iter(), - result_data.iter_mut(), - )) { - result.prepare_results(transform_data, cube_store_result)?; - } - - Ok(()) -} - #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(untagged)] pub enum TransformedData { @@ -1045,7 +1038,7 @@ impl TransformedData { query_type, query.time_dimensions.as_ref(), )?; - let columns = build_columnar_columns(&plan, cube_store_result); + let columns = build_columnar_columns(&plan, cube_store_result)?; Ok(TransformedData::Columnar { members, columns }) } _ => { @@ -1066,16 +1059,60 @@ impl TransformedData { } } +/// The `data` member is generic so a response can carry either a materialized +/// [`TransformedData`] or something that renders it while serializing (see +/// [`crate::direct_result::DirectData`]). Defaulting the parameter keeps the +/// derived `Serialize`/`Deserialize` — and with them the field names, order and +/// `skip_serializing_if` rules — as the single definition of the wire shape. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct RequestResultDataMulti { +pub struct RequestResultDataMulti { pub query_type: QueryType, - pub results: Vec, + pub results: Vec>, #[serde(skip_serializing_if = "Option::is_none")] pub pivot_query: Option, pub slow_query: bool, } +impl RequestResultDataMulti { + /// Computes the pivot query from the per-result queries. Independent of the + /// result data, so it can run before serialization on the streaming path. + pub fn prepare_pivot_query(&mut self) -> Result<()> { + let normalized_queries = self + .results + .iter() + .map(|result| &result.query) + .collect::>(); + + self.pivot_query = Some(get_pivot_query(&self.query_type, &normalized_queries)?); + + Ok(()) + } + + /// Attach one data payload per result, keeping every other field as it is. + pub fn with_data(self, data: Vec) -> Result> { + if self.results.len() != data.len() { + bail!( + "Expected {} result data entries, got {}", + self.results.len(), + data.len() + ); + } + + Ok(RequestResultDataMulti { + query_type: self.query_type, + results: self + .results + .into_iter() + .zip(data) + .map(|(result, data)| result.with_data(data)) + .collect(), + pivot_query: self.pivot_query, + slow_query: self.slow_query, + }) + } +} + impl RequestResultDataMulti { /// Processes multiple results and populates the final `RequestResultDataMulti` structure /// which is sent to the client. @@ -1092,21 +1129,13 @@ impl RequestResultDataMulti { result.prepare_results(transform_data, cube_store_result)?; } - let normalized_queries = self - .results - .iter() - .map(|result| &result.query) - .collect::>(); - - self.pivot_query = Some(get_pivot_query(&self.query_type, &normalized_queries)?); - - Ok(()) + self.prepare_pivot_query() } } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct RequestResultData { +pub struct RequestResultData { pub query: NormalizedQuery, #[serde(skip_serializing_if = "Option::is_none")] pub last_refresh_time: Option, @@ -1131,7 +1160,31 @@ pub struct RequestResultData { #[serde(skip_serializing_if = "Option::is_none")] pub total: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub data: Option, + pub data: Option, +} + +impl RequestResultData { + /// Replace the data payload, keeping every other field as it is. Written out + /// field by field on purpose: a field added later fails to compile here + /// instead of being silently dropped from the response. + pub fn with_data(self, data: T) -> RequestResultData { + RequestResultData { + query: self.query, + last_refresh_time: self.last_refresh_time, + refresh_key_values: self.refresh_key_values, + used_pre_aggregations: self.used_pre_aggregations, + transformed_query: self.transformed_query, + request_id: self.request_id, + annotation: self.annotation, + data_source: self.data_source, + db_type: self.db_type, + ext_db_type: self.ext_db_type, + external: self.external, + slow_query: self.slow_query, + total: self.total, + data: Some(data), + } + } } impl RequestResultData { @@ -1353,13 +1406,6 @@ impl From> for ColumnarArray { } } -impl From for Vec { - #[inline] - fn from(c: ColumnarArray) -> Self { - c.0 - } -} - impl std::ops::Deref for ColumnarArray { type Target = Vec; #[inline] @@ -1376,7 +1422,7 @@ impl std::ops::DerefMut for ColumnarArray { } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::*; use crate::transport::JsRawColumnarData; use anyhow::Result; @@ -1392,13 +1438,13 @@ mod tests { assert_eq!(std::mem::size_of::(), 32); } - type TestSuiteData = HashMap; + pub(crate) type TestSuiteData = HashMap; #[derive(Clone, Deserialize)] #[serde(rename_all = "camelCase")] - struct TestData { - request: TransformDataRequest, - query_result: JsRawColumnarData, + pub(crate) struct TestData { + pub(crate) request: TransformDataRequest, + pub(crate) query_result: JsRawColumnarData, final_result_default: Option, final_result_compact: Option, } @@ -2172,7 +2218,7 @@ mod tests { } "#; - static TEST_SUITE_DATA: LazyLock = + pub(crate) static TEST_SUITE_DATA: LazyLock = LazyLock::new(|| from_str(TEST_SUITE_JSON).unwrap()); #[derive(Debug)] @@ -3684,12 +3730,14 @@ mod tests { let raw_data = QueryResult::try_new( vec!["t_day".to_string(), "t_month".to_string()], vec![ - ColumnarArray::from(vec![DBResponsePrimitive::String( + vec![DBResponsePrimitive::String( "2024-06-15T00:00:00.000".to_string(), - )]), - ColumnarArray::from(vec![DBResponsePrimitive::String( + )] + .into(), + vec![DBResponsePrimitive::String( "2024-06-01T00:00:00.000".to_string(), - )]), + )] + .into(), ], )?; let plan = build_vanilla_plan( @@ -3712,4 +3760,196 @@ mod tests { ); Ok(()) } + + /// Every response format, so a test can assert across all of them. `None` is + /// the vanilla (default) format. + pub(crate) const ALL_RES_TYPES: [Option; 3] = + [None, Some(ResultType::Compact), Some(ResultType::Columnar)]; + + /// The same three-row result twice: once backed by Arrow memory (as CubeStore + /// sends it) and once by materialized primitives (as the legacy and JS-driver + /// paths hand it over). Types are mixed on purpose — string, float, decimal + /// and timestamp cells all take different routes through the cell reader. + pub(crate) struct StorageFixture { + pub arrow: QueryResult, + pub columnar: QueryResult, + alias_to_member_name_map: HashMap, + annotation: HashMap, + query: NormalizedQuery, + } + + impl StorageFixture { + pub fn new() -> Result { + use arrow::array::{ + Decimal128Array, Float64Array, StringArray, TimestampMillisecondArray, + }; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use arrow::ipc::writer::StreamWriter; + use arrow::record_batch::RecordBatch; + + let schema = Arc::new(Schema::new(vec![ + Field::new("cube__city", DataType::Utf8, true), + Field::new("cube__amount", DataType::Float64, true), + Field::new("cube__total", DataType::Decimal128(38, 2), true), + Field::new( + "cube__created_at_day", + DataType::Timestamp(TimeUnit::Millisecond, None), + true, + ), + ])); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![ + Some("Berlin"), + None, + Some("Lisbon"), + ])), + Arc::new(Float64Array::from(vec![Some(1.5), Some(2.0), None])), + Arc::new( + Decimal128Array::from(vec![Some(239996i128), None, Some(215490)]) + .with_precision_and_scale(38, 2) + .unwrap(), + ), + Arc::new(TimestampMillisecondArray::from(vec![ + Some(0i64), + Some(1_000), + None, + ])), + ], + ) + .unwrap(); + + let mut ipc = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut ipc, schema.as_ref()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + let arrow = QueryResult::from_arrow(&ipc)?; + let columnar = QueryResult::try_new( + arrow.members().to_vec(), + (0..arrow.members().len()) + .map(|idx| arrow.column(idx)?.to_columnar().map(Into::into)) + .collect::, _>>()?, + )?; + + let members = [ + ("cube__city", "Cube.city", "string"), + ("cube__amount", "Cube.amount", "number"), + ("cube__total", "Cube.total", "number"), + ("cube__created_at_day", "Cube.createdAt.day", "time"), + ]; + + let mut alias_to_member_name_map: HashMap = HashMap::new(); + let mut annotation: HashMap = HashMap::new(); + for (alias, member, member_type) in members { + alias_to_member_name_map.insert(alias.to_string(), member.to_string()); + annotation.insert(member.to_string(), make_config_item(member_type)); + } + + let query = make_query_with_dims(Some( + members + .iter() + .map(|(_, member, _)| MemberOrMemberExpression::Member(member.to_string())) + .collect(), + )); + + Ok(Self { + arrow, + columnar, + alias_to_member_name_map, + annotation, + query, + }) + } + + pub fn request(&self, res_type: Option) -> TransformDataRequest { + TransformDataRequest { + alias_to_member_name_map: self.alias_to_member_name_map.clone(), + annotation: self.annotation.clone(), + query: self.query.clone(), + query_type: Some(QueryType::RegularQuery), + res_type, + } + } + } + + /// Arrow-backed and primitive-backed columns must transform identically in + /// every response format — the Arrow path reads cells lazily through + /// `ColumnReader`, the primitive path off a slice. + #[test] + fn test_transform_matches_across_column_storage() -> Result<()> { + let fixture = StorageFixture::new()?; + + for res_type in ALL_RES_TYPES { + let request = fixture.request(res_type.clone()); + + let from_arrow = TransformedData::transform(&request, &fixture.arrow)?; + let from_columnar = TransformedData::transform(&request, &fixture.columnar)?; + assert_eq!( + from_arrow, from_columnar, + "res_type {res_type:?} must not depend on column storage" + ); + + // Serialized shape too, since that is what reaches the client. + assert_eq!( + serde_json::to_string(&from_arrow)?, + serde_json::to_string(&from_columnar)?, + "res_type {res_type:?} JSON must not depend on column storage" + ); + } + + // Spot-check the rendering rules the Arrow reader is responsible for. + let TransformedData::Compact { members, dataset } = TransformedData::transform( + &fixture.request(Some(ResultType::Compact)), + &fixture.arrow, + )? + else { + panic!("expected Compact"); + }; + let total_idx = members.iter().position(|m| m == "Cube.total").unwrap(); + let created_idx = members + .iter() + .position(|m| m == "Cube.createdAt.day") + .unwrap(); + assert_eq!( + dataset[0][total_idx], + DBResponsePrimitive::String("2399.96".to_string()), + "decimals render from mantissa and scale" + ); + assert_eq!( + serde_json::to_value(&dataset[1][created_idx])?, + serde_json::json!("1970-01-01T00:00:01.000"), + "timestamps serialize in the legacy ISO shape" + ); + assert_eq!(dataset[2][created_idx], DBResponsePrimitive::Null); + + Ok(()) + } + + /// Response envelope with every kind of optional field represented, so the + /// direct-serialization tests also cover field order and `skip_serializing_if`. + pub(crate) fn make_result_head(query: NormalizedQuery) -> RequestResultData { + RequestResultData { + query, + last_refresh_time: Some("2024-06-15T00:00:00.000".to_string()), + refresh_key_values: None, + used_pre_aggregations: None, + transformed_query: None, + request_id: Some("test-request".to_string()), + // Left empty on purpose: `HashMap` iteration order is not stable + // across instances, and these tests compare serialized bytes. + annotation: HashMap::new(), + data_source: Some("default".to_string()), + db_type: Some("postgres".to_string()), + ext_db_type: None, + external: Some(false), + slow_query: false, + total: Some(3), + data: None, + } + } } From 7905950ae489ccf26760b2ab236579fb5f6ce2cf Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Wed, 29 Jul 2026 17:36:59 +0200 Subject: [PATCH 2/5] refactor(cubeorchestrator): drop the materializing column escape hatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `QueryResultColumn::to_columnar` and `ArrowArray::to_columnar` had no production callers — only two test sites — and `ArrowArray::data_type`/`::array` had none at all. A materializing helper on the type whose point is *not* to materialize is an invitation to reintroduce the copy this branch removed, so drop all four. The two tests now read columns through `ColumnReader`, which is what the transform paths use, so they exercise the production accessor instead of a test-only one. Co-Authored-By: Claude Opus 5 (1M context) --- .../src/query_message_parser.rs | 9 +++--- .../src/query_result_column.rs | 32 ------------------- .../src/query_result_transform.rs | 24 ++++++++++---- 3 files changed, 23 insertions(+), 42 deletions(-) diff --git a/rust/cube/cubeorchestrator/src/query_message_parser.rs b/rust/cube/cubeorchestrator/src/query_message_parser.rs index dba70654f8a1f..257c97520c1be 100644 --- a/rust/cube/cubeorchestrator/src/query_message_parser.rs +++ b/rust/cube/cubeorchestrator/src/query_message_parser.rs @@ -348,10 +348,11 @@ mod tests { use cubeshared::flatbuffers::FlatBufferBuilder; use std::sync::Arc; - /// Materialize one column so assertions can compare against a plain slice, - /// whatever the column's backing storage is. - fn column(result: &QueryResult, idx: usize) -> ColumnarArray { - result.column(idx).unwrap().to_columnar().unwrap() + /// Read one column through the production accessor, so assertions can compare + /// against a plain slice whatever the column's backing storage is. + fn column(result: &QueryResult, idx: usize) -> Vec { + let reader = result.reader(idx).unwrap(); + (0..reader.len()).map(|row| reader.value(row)).collect() } /// Helper function to create a test HttpMessage with a given number of rows and columns diff --git a/rust/cube/cubeorchestrator/src/query_result_column.rs b/rust/cube/cubeorchestrator/src/query_result_column.rs index 130b09d98fb85..8205b72c69dc3 100644 --- a/rust/cube/cubeorchestrator/src/query_result_column.rs +++ b/rust/cube/cubeorchestrator/src/query_result_column.rs @@ -54,15 +54,6 @@ impl QueryResultColumn { QueryResultColumn::Arrow(a) => Ok(ColumnReader::Arrow(a.cell_reader()?)), } } - - /// Materialize the column as primitives. Only for callers that genuinely need - /// an owned slice — the read paths use [`QueryResultColumn::reader`] instead. - pub fn to_columnar(&self) -> Result { - match self { - QueryResultColumn::Columnar(c) => Ok(c.clone()), - QueryResultColumn::Arrow(a) => a.to_columnar(), - } - } } impl From for QueryResultColumn { @@ -115,29 +106,6 @@ impl ArrowArray { self.0.is_empty() } - #[inline] - pub fn data_type(&self) -> &DataType { - self.0.data_type() - } - - #[inline] - pub fn array(&self) -> &ArrayRef { - &self.0 - } - - pub fn to_columnar(&self) -> Result { - let mut out = ColumnarArray::with_capacity(self.len()); - // `member_type` is empty: materializing must not apply the `time` reformat - // that a transform would. - let Ok(()) = self - .cell_reader()? - .for_each_transformed::("", |value| { - out.push(value.into_owned()); - Ok(()) - }); - Ok(out) - } - /// Downcast to the concrete Arrow array once, for the whole column. fn cell_reader(&self) -> Result, ParseError> { let array = self.0.as_ref(); diff --git a/rust/cube/cubeorchestrator/src/query_result_transform.rs b/rust/cube/cubeorchestrator/src/query_result_transform.rs index 88ea1cad73663..8b668e6bc212b 100644 --- a/rust/cube/cubeorchestrator/src/query_result_transform.rs +++ b/rust/cube/cubeorchestrator/src/query_result_transform.rs @@ -1424,6 +1424,7 @@ impl std::ops::DerefMut for ColumnarArray { #[cfg(test)] pub(crate) mod tests { use super::*; + use crate::query_result_column::QueryResultColumn; use crate::transport::JsRawColumnarData; use anyhow::Result; use serde_json::from_str; @@ -3766,6 +3767,21 @@ pub(crate) mod tests { pub(crate) const ALL_RES_TYPES: [Option; 3] = [None, Some(ResultType::Compact), Some(ResultType::Columnar)]; + /// Read every column of `source` into primitives, i.e. the shape the legacy + /// and JS-driver paths hand over. + fn materialize_columns(source: &QueryResult) -> Result> { + (0..source.members().len()) + .map(|idx| { + let reader = source.reader(idx)?; + Ok(QueryResultColumn::from( + (0..reader.len()) + .map(|row| reader.value(row)) + .collect::>(), + )) + }) + .collect() + } + /// The same three-row result twice: once backed by Arrow memory (as CubeStore /// sends it) and once by materialized primitives (as the legacy and JS-driver /// paths hand it over). Types are mixed on purpose — string, float, decimal @@ -3829,12 +3845,8 @@ pub(crate) mod tests { } let arrow = QueryResult::from_arrow(&ipc)?; - let columnar = QueryResult::try_new( - arrow.members().to_vec(), - (0..arrow.members().len()) - .map(|idx| arrow.column(idx)?.to_columnar().map(Into::into)) - .collect::, _>>()?, - )?; + let columnar = + QueryResult::try_new(arrow.members().to_vec(), materialize_columns(&arrow)?)?; let members = [ ("cube__city", "Cube.city", "string"), From 04a5202e434ceea535cf6e62074ca498c9b8c6c6 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Thu, 30 Jul 2026 14:43:03 +0200 Subject: [PATCH 3/5] perf(cubeorchestrator): borrow Arrow text instead of allocating per cell MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `read_cell!` had to produce a `DBResponsePrimitive`, whose `String` variant owns its bytes — so every cell of an Arrow `Utf8`/`LargeUtf8`/`Utf8View` column cost an allocation even when the caller only wrote the cell out and dropped it. Add `CellRef`: a cell on its way out of a column, which either points at a materialized primitive, borrows text straight from the Arrow buffer, or owns a value decoded on read. The response serializers take that, so text columns now reach the JSON writer with no allocation at all. `Serialize for CellRef` delegates to `Serialize for DBResponsePrimitive` for everything except `Str`, which renders exactly as that impl's `String` arm — one line, so the two cannot drift apart. Borrowing is gated on `is_identity_transform`: a text column annotated as a `time` member is rewritten by `transform_value` and still needs an owned value. `StorageFixture` grew a `Utf8` column typed `time` to cover precisely that, and asserts the reformatted output rather than the raw Arrow text. 16 columns x 100k rows, Arrow source, response JSON: | format | materialized | direct before | direct now | | -------- | ------------ | ------------- | ---------- | | compact | 67.9 ms | 57.1 ms | 48.1 ms | | columnar | 55.2 ms | 51.3 ms | 38.6 ms | The fixture also had to pin down member ordering: `get_members` walks a `HashMap` when appending deprecated-style time members, so a query with two of them yields either order from one call to the next. The fixture now requests one base dimension explicitly, leaving a single appended member — the nondeterminism itself is pre-existing and left alone. Co-Authored-By: Claude Opus 5 (1M context) --- .../cubeorchestrator/src/direct_result.rs | 4 +- .../src/query_result_column.rs | 136 +++++++++++++++--- .../src/query_result_transform.rs | 35 +++++ 3 files changed, 150 insertions(+), 25 deletions(-) diff --git a/rust/cube/cubeorchestrator/src/direct_result.rs b/rust/cube/cubeorchestrator/src/direct_result.rs index 878d14351fb52..47303989d36bc 100644 --- a/rust/cube/cubeorchestrator/src/direct_result.rs +++ b/rust/cube/cubeorchestrator/src/direct_result.rs @@ -161,7 +161,7 @@ impl Serialize for CompactRow<'_> { column, member_type, } => column.with_transformed(self.row_idx, member_type, |value| { - seq.serialize_element(value) + seq.serialize_element(&value) })?, CompactPlanEntry::Constant(value) => seq.serialize_element(value)?, } @@ -206,7 +206,7 @@ impl Serialize for ColumnarColumn<'_> { // row loop — the same tight fill the materializing path uses. let reader = self.source.reader(*index).map_err(S::Error::custom)?; reader.for_each_transformed(self.entry.member_type, |value| { - seq.serialize_element(value.as_ref()) + seq.serialize_element(&value) })?; } ColumnarColumnSource::Constant(value) => { diff --git a/rust/cube/cubeorchestrator/src/query_result_column.rs b/rust/cube/cubeorchestrator/src/query_result_column.rs index 8205b72c69dc3..b384f9fe064e6 100644 --- a/rust/cube/cubeorchestrator/src/query_result_column.rs +++ b/rust/cube/cubeorchestrator/src/query_result_column.rs @@ -20,6 +20,7 @@ use arrow::array::{ UInt32Array, UInt64Array, UInt8Array, }; use arrow::datatypes::{ArrowTemporalType, DataType, TimeUnit}; +use serde::{Serialize, Serializer}; use std::{borrow::Cow, convert::Infallible}; /// One logical column of a query result. @@ -178,6 +179,58 @@ impl ArrowArray { } } +/// A cell on its way out of a column. +/// +/// [`DBResponsePrimitive`] owns its `String`, so building one from an Arrow `Utf8` +/// column costs an allocation per cell. Callers that only render a cell and drop +/// it — the response serializers — take this instead and borrow the text where it +/// already lives. +pub enum CellRef<'a> { + /// A cell already materialized in the column. + Primitive(&'a DBResponsePrimitive), + /// Text borrowed straight from the column's buffer. + Str(&'a str), + /// Decoded on read: booleans, numbers, timestamps, decimals. + Owned(DBResponsePrimitive), +} + +impl CellRef<'_> { + /// The same cell as an owned primitive, for callers that keep it. + #[inline] + pub fn into_owned(self) -> DBResponsePrimitive { + match self { + CellRef::Primitive(value) => value.clone(), + CellRef::Str(text) => DBResponsePrimitive::String(text.to_owned()), + CellRef::Owned(value) => value, + } + } +} + +/// Mirrors `Serialize for DBResponsePrimitive`: `Str` renders exactly as that +/// impl's `String` arm, and every other cell delegates to it outright. +impl Serialize for CellRef<'_> { + #[inline] + fn serialize(&self, serializer: S) -> Result { + match self { + CellRef::Primitive(value) => value.serialize(serializer), + CellRef::Str(text) => serializer.serialize_str(text), + CellRef::Owned(value) => value.serialize(serializer), + } + } +} + +/// Read `$array[$row]` as borrowed text, or a null cell. +macro_rules! borrowed_str_cell { + ($array:expr, $row:expr) => {{ + let a = $array; + if a.is_null($row) { + CellRef::Owned(DBResponsePrimitive::Null) + } else { + CellRef::Str(a.value($row)) + } + }}; +} + /// Per-cell accessor for one column, resolved once per column so the per-cell path /// is a jump table plus an index — never a `DataType` match and a downcast. pub enum ColumnReader<'a> { @@ -224,21 +277,32 @@ impl ColumnReader<'_> { } } - /// The cell at `row`, transformed for `member_type`, passed to `f` by - /// reference. A materialized column whose member type needs no transform is - /// handed over borrowed, so a reader that only forwards the value — the - /// row-major serializers — never clones it. + /// The cell at `row`, transformed for `member_type`, handed to `f`. Cells that + /// need no transform are passed borrowed — from the column's own storage or + /// straight out of the Arrow buffer — so a caller that only renders the value + /// never allocates for it. #[inline] pub fn with_transformed( &self, row: usize, member_type: &str, - f: impl FnOnce(&DBResponsePrimitive) -> R, + f: impl FnOnce(CellRef<'_>) -> R, ) -> R { - match self { - ColumnReader::Primitives(cells) if is_identity_transform(member_type) => f(&cells[row]), - _ => f(&transform_value(self.value(row), member_type)), + if is_identity_transform(member_type) { + match self { + ColumnReader::Primitives(cells) => return f(CellRef::Primitive(&cells[row])), + ColumnReader::Arrow(a) => { + if let Some(cell) = a.borrowed_cell(row) { + return f(cell); + } + } + } } + + f(CellRef::Owned(transform_value( + self.value(row), + member_type, + ))) } /// Hand every cell, transformed for `member_type`, to `visit` in row order. @@ -248,17 +312,17 @@ impl ColumnReader<'_> { pub fn for_each_transformed( &self, member_type: &str, - mut visit: impl FnMut(Cow<'_, DBResponsePrimitive>) -> Result<(), E>, + mut visit: impl FnMut(CellRef<'_>) -> Result<(), E>, ) -> Result<(), E> { match self { ColumnReader::Primitives(cells) => { if is_identity_transform(member_type) { for cell in cells.iter() { - visit(Cow::Borrowed(cell))?; + visit(CellRef::Primitive(cell))?; } } else { for cell in cells.iter() { - visit(Cow::Owned(transform_value(cell.clone(), member_type)))?; + visit(CellRef::Owned(transform_value(cell.clone(), member_type)))?; } } Ok(()) @@ -427,6 +491,19 @@ impl ArrowCellReader<'_> { } } + /// The cell at `row` borrowed from the Arrow buffer, for the types that store + /// their values as text. `None` for every other type, so the caller falls back + /// to decoding an owned value. + #[inline] + fn borrowed_cell(&self, row: usize) -> Option> { + match self { + ArrowCellReader::Utf8(a) => Some(borrowed_str_cell!(a, row)), + ArrowCellReader::LargeUtf8(a) => Some(borrowed_str_cell!(a, row)), + ArrowCellReader::Utf8View(a) => Some(borrowed_str_cell!(a, row)), + _ => None, + } + } + /// Text rendering of `row`, borrowing from the Arrow buffer for string types. fn value_as_str(&self, row: usize) -> Option> { macro_rules! borrowed_str { @@ -458,14 +535,14 @@ impl ArrowCellReader<'_> { fn for_each_transformed( &self, member_type: &str, - mut visit: impl FnMut(Cow<'_, DBResponsePrimitive>) -> Result<(), E>, + mut visit: impl FnMut(CellRef<'_>) -> Result<(), E>, ) -> Result<(), E> { let len = self.len(); macro_rules! fill { ($read:expr) => {{ for row in 0..len { - visit(Cow::Owned(transform_value($read(row), member_type)))?; + visit(CellRef::Owned(transform_value($read(row), member_type)))?; } }}; } @@ -477,11 +554,30 @@ impl ArrowCellReader<'_> { }}; } + /// Text columns are handed over borrowed unless the member type asks for + /// a transform, which needs an owned `String` to rewrite. + macro_rules! fill_str { + ($array:expr) => {{ + let a = $array; + if is_identity_transform(member_type) { + for row in 0..len { + visit(borrowed_str_cell!(a, row))?; + } + } else { + fill!( + |row| read_cell!(a, row, |v: &str| DBResponsePrimitive::String( + v.to_owned() + )) + ) + } + }}; + } + match self { ArrowCellReader::Empty => {} ArrowCellReader::Null(_) => { for _ in 0..len { - visit(Cow::Owned(DBResponsePrimitive::Null))?; + visit(CellRef::Owned(DBResponsePrimitive::Null))?; } } ArrowCellReader::Boolean(a) => fill_with!(a, DBResponsePrimitive::Boolean), @@ -502,15 +598,9 @@ impl ArrowCellReader<'_> { fill_with!(a, |v| DBResponsePrimitive::Float64(v as f64)) } ArrowCellReader::Float64(a) => fill_with!(a, DBResponsePrimitive::Float64), - ArrowCellReader::Utf8(a) => { - fill_with!(a, |v: &str| DBResponsePrimitive::String(v.to_owned())) - } - ArrowCellReader::LargeUtf8(a) => { - fill_with!(a, |v: &str| DBResponsePrimitive::String(v.to_owned())) - } - ArrowCellReader::Utf8View(a) => { - fill_with!(a, |v: &str| DBResponsePrimitive::String(v.to_owned())) - } + ArrowCellReader::Utf8(a) => fill_str!(a), + ArrowCellReader::LargeUtf8(a) => fill_str!(a), + ArrowCellReader::Utf8View(a) => fill_str!(a), ArrowCellReader::Date32(a) => fill!(|row| datetime_cell(*a, row)), ArrowCellReader::Date64(a) => fill!(|row| datetime_cell(*a, row)), ArrowCellReader::TimestampSecond(a) => fill!(|row| datetime_cell(*a, row)), diff --git a/rust/cube/cubeorchestrator/src/query_result_transform.rs b/rust/cube/cubeorchestrator/src/query_result_transform.rs index 8b668e6bc212b..f1fb1bdab853f 100644 --- a/rust/cube/cubeorchestrator/src/query_result_transform.rs +++ b/rust/cube/cubeorchestrator/src/query_result_transform.rs @@ -3812,6 +3812,10 @@ pub(crate) mod tests { DataType::Timestamp(TimeUnit::Millisecond, None), true, ), + // A text column annotated as `time`: the one shape where a cell + // cannot be handed over borrowed, because `transform_value` + // rewrites it. + Field::new("cube__shipped_at_day", DataType::Utf8, true), ])); let batch = RecordBatch::try_new( @@ -3833,6 +3837,11 @@ pub(crate) mod tests { Some(1_000), None, ])), + Arc::new(StringArray::from(vec![ + Some("2024-06-15 00:00:00.000"), + None, + Some("2024-06-16T00:00:00"), + ])), ], ) .unwrap(); @@ -3853,6 +3862,7 @@ pub(crate) mod tests { ("cube__amount", "Cube.amount", "number"), ("cube__total", "Cube.total", "number"), ("cube__created_at_day", "Cube.createdAt.day", "time"), + ("cube__shipped_at_day", "Cube.shippedAt.day", "time"), ]; let mut alias_to_member_name_map: HashMap = HashMap::new(); @@ -3866,6 +3876,15 @@ pub(crate) mod tests { members .iter() .map(|(_, member, _)| MemberOrMemberExpression::Member(member.to_string())) + // `Cube.createdAt` is requested without a granularity, so + // `get_members` appends a deprecated-style member only for + // `Cube.shippedAt`. That keeps the appended-member order + // deterministic across calls: `get_members` walks a `HashMap` + // there, so two of them would come out in either order and the + // byte-equality assertions below would flake. + .chain([MemberOrMemberExpression::Member( + "Cube.createdAt".to_string(), + )]) .collect(), )); @@ -3927,11 +3946,27 @@ pub(crate) mod tests { .iter() .position(|m| m == "Cube.createdAt.day") .unwrap(); + let shipped_idx = members + .iter() + .position(|m| m == "Cube.shippedAt.day") + .unwrap(); assert_eq!( dataset[0][total_idx], DBResponsePrimitive::String("2399.96".to_string()), "decimals render from mantissa and scale" ); + // A `time` member's text is reformatted, so it must not be passed through + // as the raw Arrow value. + assert_eq!( + dataset[0][shipped_idx], + DBResponsePrimitive::String("2024-06-15T00:00:00.000".to_string()), + "text of a time member is reformatted, not borrowed as-is" + ); + assert_eq!( + dataset[2][shipped_idx], + DBResponsePrimitive::String("2024-06-16T00:00:00.000".to_string()), + ); + assert_eq!(dataset[1][shipped_idx], DBResponsePrimitive::Null); assert_eq!( serde_json::to_value(&dataset[1][created_idx])?, serde_json::json!("1970-01-01T00:00:01.000"), From 264824af1b183dc007fce819120b5ae7e9a623a7 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Thu, 30 Jul 2026 14:53:12 +0200 Subject: [PATCH 4/5] perf(cubeorchestrator): render decimals without allocating MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CubeStore answers `SUM` with `Decimal128`, so most measure cells are decimals — and `decimal_to_string` allocated twice for each one: a `String` for the mantissa's digits, then a `format!` for the result. Nothing measured that, because the Arrow bench fixture typed its measures as `Float64`; with `Decimal128` measures the response JSON cost 66% more (compact) and 90% more (columnar) than the numbers this branch had been reporting. Render through `fmt::Write` and a stack digit buffer instead. `DecimalText` holds the mantissa and scale and writes the text on demand, so: - the serializers hand `CellRef::Decimal128`/`Decimal256` to `collect_str`, which streams the digits into the response with no allocation at all; - `decimal_to_string` is that same renderer collected into one `String`, which also speeds up the materializing path the SQL API uses. There is one definition of the format, so the two paths cannot drift. Since this replaces hand-written numeric formatting, `test_decimal_to_string_matches_reference` keeps the previous String-building algorithm as an oracle and checks both agree across `i128` extremes, `i256` extremes and eleven scales, and `test_decimal_cell_serializes_like_owned_string` pins the serialized bytes to what the owned `String` cell produces. 16 columns x 100k rows, Decimal128 measures, response JSON: | path | format | before | after | | ------------ | -------- | -------- | ------- | | direct | compact | 79.9 ms | 45.1 ms | | direct | columnar | 73.4 ms | 40.2 ms | | materialized | compact | 103.5 ms | 78.8 ms | | materialized | columnar | 87.8 ms | 64.2 ms | The bench fixture gained a `MeasureKind` axis so the decimal shape stays measured from here on. Co-Authored-By: Claude Opus 5 (1M context) --- .../cubeorchestrator/benches/common/mod.rs | 62 +++- rust/cube/cubeorchestrator/benches/parser.rs | 5 +- .../cubeorchestrator/benches/transform.rs | 40 +- .../src/query_result_column.rs | 344 +++++++++++++++--- .../src/query_result_transform.rs | 2 +- 5 files changed, 386 insertions(+), 67 deletions(-) diff --git a/rust/cube/cubeorchestrator/benches/common/mod.rs b/rust/cube/cubeorchestrator/benches/common/mod.rs index 8a57247404a50..3df2b004bf658 100644 --- a/rust/cube/cubeorchestrator/benches/common/mod.rs +++ b/rust/cube/cubeorchestrator/benches/common/mod.rs @@ -85,17 +85,38 @@ pub fn build_dataset( JsRawColumnarData { members, columns } } +/// How measure columns are typed in an Arrow fixture. CubeStore answers `SUM` +/// with `Decimal128`, so that is the shape most measure cells really have; the +/// `Float64` variant is kept because it is the cheaper baseline to compare against. +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum MeasureKind { + Float64, + Decimal128, +} + +impl MeasureKind { + pub fn label(self) -> &'static str { + match self { + MeasureKind::Float64 => "arrow", + MeasureKind::Decimal128 => "arrow_dec", + } + } +} + /// Build an Arrow IPC **stream** payload with the same logical data shape as -/// [`build_dataset`]: dimensions as Utf8, measures as Float64, time dimensions -/// as Timestamp(Millisecond). Used to compare Arrow parse throughput against the -/// JSON path. +/// [`build_dataset`]: dimensions as Utf8, measures per `measure_kind`, time +/// dimensions as Timestamp(Millisecond). Used to compare Arrow parse throughput +/// against the JSON path. pub fn build_arrow_ipc( row_count: usize, dimensions: &[(String, String)], measures: &[(String, String)], time_dims: &[TimeColumn], + measure_kind: MeasureKind, ) -> Vec { - use arrow::array::{ArrayRef, Float64Array, StringArray, TimestampMillisecondArray}; + use arrow::array::{ + ArrayRef, Decimal128Array, Float64Array, StringArray, TimestampMillisecondArray, + }; use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use arrow::ipc::writer::StreamWriter; use arrow::record_batch::RecordBatch; @@ -113,11 +134,31 @@ pub fn build_arrow_ipc( columns.push(Arc::new(StringArray::from(values))); } for (j, (_, alias)) in measures.iter().enumerate() { - fields.push(Field::new(alias.clone(), DataType::Float64, false)); - let values: Vec = (0..row_count) - .map(|i| ((i * (j + 1)) as f64) * 0.5) - .collect(); - columns.push(Arc::new(Float64Array::from(values))); + match measure_kind { + MeasureKind::Float64 => { + fields.push(Field::new(alias.clone(), DataType::Float64, false)); + let values: Vec = (0..row_count) + .map(|i| ((i * (j + 1)) as f64) * 0.5) + .collect(); + columns.push(Arc::new(Float64Array::from(values))); + } + MeasureKind::Decimal128 => { + fields.push(Field::new( + alias.clone(), + DataType::Decimal128(38, 2), + false, + )); + // Same magnitudes as the Float64 arm, as a scale-2 mantissa. + let values: Vec = (0..row_count) + .map(|i| ((i * (j + 1)) as i128) * 50) + .collect(); + columns.push(Arc::new( + Decimal128Array::from(values) + .with_precision_and_scale(38, 2) + .expect("decimal precision"), + )); + } + } } for (j, td) in time_dims.iter().enumerate() { fields.push(Field::new( @@ -190,8 +231,9 @@ pub fn build_arrow_query_result( dimensions: &[(String, String)], measures: &[(String, String)], time_dims: &[TimeColumn], + measure_kind: MeasureKind, ) -> QueryResult { - let ipc = build_arrow_ipc(row_count, dimensions, measures, time_dims); + let ipc = build_arrow_ipc(row_count, dimensions, measures, time_dims, measure_kind); let payload = build_cubestore_fb_arrow_message(&ipc); QueryResult::from_cubestore_fb(&payload).expect("arrow query result") } diff --git a/rust/cube/cubeorchestrator/benches/parser.rs b/rust/cube/cubeorchestrator/benches/parser.rs index 2f80fc10902f8..bcb150053e5ac 100644 --- a/rust/cube/cubeorchestrator/benches/parser.rs +++ b/rust/cube/cubeorchestrator/benches/parser.rs @@ -13,7 +13,7 @@ use cubeshared::flatbuffers::FlatBufferBuilder; mod common; use common::{ build_arrow_ipc, build_cubestore_fb_arrow_message, build_dataset, make_member_aliases, - split_dim_measure, COLUMN_COUNTS, ROW_COUNTS, + split_dim_measure, MeasureKind, COLUMN_COUNTS, ROW_COUNTS, }; /// Build a FlatBuffer `HttpMessage` payload mirroring CubeStore's wire format @@ -169,7 +169,8 @@ fn bench_from_cubestore_fb_arrow(c: &mut Criterion) { let dimensions = make_member_aliases("dim", dim_count); let measures = make_member_aliases("measure", measure_count); - let arrow_ipc = build_arrow_ipc(row_count, &dimensions, &measures, &[]); + let arrow_ipc = + build_arrow_ipc(row_count, &dimensions, &measures, &[], MeasureKind::Float64); let payload = build_cubestore_fb_arrow_message(&arrow_ipc); let payload_len = payload.len(); diff --git a/rust/cube/cubeorchestrator/benches/transform.rs b/rust/cube/cubeorchestrator/benches/transform.rs index 2dd70a047a7c8..4089b55014fda 100644 --- a/rust/cube/cubeorchestrator/benches/transform.rs +++ b/rust/cube/cubeorchestrator/benches/transform.rs @@ -13,8 +13,8 @@ use cubeorchestrator::transport::{ #[path = "common/mod.rs"] mod common; use common::{ - build_arrow_query_result, build_dataset, make_member_aliases, split_dim_measure, TimeColumn, - COLUMN_COUNTS, ROW_COUNTS, + build_arrow_query_result, build_dataset, make_member_aliases, split_dim_measure, MeasureKind, + TimeColumn, COLUMN_COUNTS, ROW_COUNTS, }; /// Total columns and row count used by `bench_transform_time_scenarios`. @@ -167,7 +167,13 @@ fn bench_transform(c: &mut Criterion) { ), ( "arrow", - build_arrow_query_result(row_count, &dimensions, &measures, &[]), + build_arrow_query_result( + row_count, + &dimensions, + &measures, + &[], + MeasureKind::Float64, + ), ), ]; @@ -228,7 +234,13 @@ fn bench_transform_time_scenarios(c: &mut Criterion) { ), ( "arrow", - build_arrow_query_result(SCENARIO_ROW_COUNT, &dimensions, &measures, &time_dims), + build_arrow_query_result( + SCENARIO_ROW_COUNT, + &dimensions, + &measures, + &time_dims, + MeasureKind::Float64, + ), ), ]; @@ -291,8 +303,24 @@ fn bench_final_json(c: &mut Criterion) { .expect("from_js_raw_data"), ), ( - "arrow", - build_arrow_query_result(row_count, &dimensions, &measures, &[]), + MeasureKind::Float64.label(), + build_arrow_query_result( + row_count, + &dimensions, + &measures, + &[], + MeasureKind::Float64, + ), + ), + ( + MeasureKind::Decimal128.label(), + build_arrow_query_result( + row_count, + &dimensions, + &measures, + &[], + MeasureKind::Decimal128, + ), ), ]; diff --git a/rust/cube/cubeorchestrator/src/query_result_column.rs b/rust/cube/cubeorchestrator/src/query_result_column.rs index b384f9fe064e6..a5df57e71c399 100644 --- a/rust/cube/cubeorchestrator/src/query_result_column.rs +++ b/rust/cube/cubeorchestrator/src/query_result_column.rs @@ -19,9 +19,13 @@ use arrow::array::{ TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array, }; -use arrow::datatypes::{ArrowTemporalType, DataType, TimeUnit}; +use arrow::datatypes::{i256, ArrowTemporalType, DataType, TimeUnit}; use serde::{Serialize, Serializer}; -use std::{borrow::Cow, convert::Infallible}; +use std::{ + borrow::Cow, + convert::Infallible, + fmt::{self, Write as _}, +}; /// One logical column of a query result. #[derive(Debug, Clone)] @@ -182,15 +186,24 @@ impl ArrowArray { /// A cell on its way out of a column. /// /// [`DBResponsePrimitive`] owns its `String`, so building one from an Arrow `Utf8` -/// column costs an allocation per cell. Callers that only render a cell and drop -/// it — the response serializers — take this instead and borrow the text where it -/// already lives. +/// or decimal column costs an allocation per cell. Callers that only render a cell +/// and drop it — the response serializers — take this instead, which borrows text +/// where it already lives and renders a decimal straight into the output. pub enum CellRef<'a> { /// A cell already materialized in the column. Primitive(&'a DBResponsePrimitive), /// Text borrowed straight from the column's buffer. Str(&'a str), - /// Decoded on read: booleans, numbers, timestamps, decimals. + /// A decimal still in its Arrow form, rendered on the way out. + Decimal128 { + mantissa: i128, + scale: u32, + }, + Decimal256 { + mantissa: i256, + scale: u32, + }, + /// Decoded on read: booleans, numbers, timestamps. Owned(DBResponsePrimitive), } @@ -201,19 +214,35 @@ impl CellRef<'_> { match self { CellRef::Primitive(value) => value.clone(), CellRef::Str(text) => DBResponsePrimitive::String(text.to_owned()), + CellRef::Decimal128 { mantissa, scale } => { + DBResponsePrimitive::String(decimal_to_string(mantissa, scale)) + } + CellRef::Decimal256 { mantissa, scale } => { + DBResponsePrimitive::String(decimal_to_string(mantissa, scale)) + } CellRef::Owned(value) => value, } } } -/// Mirrors `Serialize for DBResponsePrimitive`: `Str` renders exactly as that -/// impl's `String` arm, and every other cell delegates to it outright. +/// Mirrors `Serialize for DBResponsePrimitive`: `Str` and the decimals render +/// exactly as that impl's `String` arm would — a JSON string, from the same +/// [`DecimalText`] the owned path formats with — and every other cell delegates to +/// it outright. impl Serialize for CellRef<'_> { #[inline] fn serialize(&self, serializer: S) -> Result { match self { CellRef::Primitive(value) => value.serialize(serializer), CellRef::Str(text) => serializer.serialize_str(text), + CellRef::Decimal128 { mantissa, scale } => serializer.collect_str(&DecimalText { + mantissa: *mantissa, + scale: *scale, + }), + CellRef::Decimal256 { mantissa, scale } => serializer.collect_str(&DecimalText { + mantissa: *mantissa, + scale: *scale, + }), CellRef::Owned(value) => value.serialize(serializer), } } @@ -231,6 +260,21 @@ macro_rules! borrowed_str_cell { }}; } +/// Read `$array[$row]` as an unrendered decimal, or a null cell. +macro_rules! decimal_cell { + ($array:expr, $row:expr, $variant:ident, $scale:expr) => {{ + let a = $array; + if a.is_null($row) { + CellRef::Owned(DBResponsePrimitive::Null) + } else { + CellRef::$variant { + mantissa: a.value($row), + scale: $scale, + } + } + }}; +} + /// Per-cell accessor for one column, resolved once per column so the per-cell path /// is a jump table plus an index — never a `DataType` match and a downcast. pub enum ColumnReader<'a> { @@ -292,7 +336,7 @@ impl ColumnReader<'_> { match self { ColumnReader::Primitives(cells) => return f(CellRef::Primitive(&cells[row])), ColumnReader::Arrow(a) => { - if let Some(cell) = a.borrowed_cell(row) { + if let Some(cell) = a.cell_without_alloc(row) { return f(cell); } } @@ -491,15 +535,22 @@ impl ArrowCellReader<'_> { } } - /// The cell at `row` borrowed from the Arrow buffer, for the types that store - /// their values as text. `None` for every other type, so the caller falls back - /// to decoding an owned value. + /// The cell at `row` for the types whose [`ArrowCellReader::value`] would + /// allocate — text, which is borrowed instead, and decimals, which are left + /// unrendered. `None` for every other type, whose owned form allocates nothing + /// anyway, so the caller just decodes it. #[inline] - fn borrowed_cell(&self, row: usize) -> Option> { + fn cell_without_alloc(&self, row: usize) -> Option> { match self { ArrowCellReader::Utf8(a) => Some(borrowed_str_cell!(a, row)), ArrowCellReader::LargeUtf8(a) => Some(borrowed_str_cell!(a, row)), ArrowCellReader::Utf8View(a) => Some(borrowed_str_cell!(a, row)), + ArrowCellReader::Decimal128(a, scale) => { + Some(decimal_cell!(a, row, Decimal128, *scale)) + } + ArrowCellReader::Decimal256(a, scale) => { + Some(decimal_cell!(a, row, Decimal256, *scale)) + } _ => None, } } @@ -573,6 +624,21 @@ impl ArrowCellReader<'_> { }}; } + /// Decimals likewise: handed over unrendered when nothing has to rewrite + /// them, so their digits go straight into the output. + macro_rules! fill_decimal { + ($array:expr, $variant:ident, $scale:expr, $make:expr) => {{ + let a = $array; + if is_identity_transform(member_type) { + for row in 0..len { + visit(decimal_cell!(a, row, $variant, $scale))?; + } + } else { + fill!(|row| read_cell!(a, row, $make)) + } + }}; + } + match self { ArrowCellReader::Empty => {} ArrowCellReader::Null(_) => { @@ -608,14 +674,14 @@ impl ArrowCellReader<'_> { ArrowCellReader::TimestampMicrosecond(a) => fill!(|row| datetime_cell(*a, row)), ArrowCellReader::TimestampNanosecond(a) => fill!(|row| datetime_cell(*a, row)), ArrowCellReader::Decimal128(a, scale) => { - fill_with!(a, |v| DBResponsePrimitive::String(decimal_to_string( - v, *scale - ))) + fill_decimal!(a, Decimal128, *scale, |v| DBResponsePrimitive::String( + decimal_to_string(v, *scale) + )) } ArrowCellReader::Decimal256(a, scale) => { - fill_with!(a, |v| DBResponsePrimitive::String(decimal_to_string( - v, *scale - ))) + fill_decimal!(a, Decimal256, *scale, |v| DBResponsePrimitive::String( + decimal_to_string(v, *scale) + )) } } @@ -641,41 +707,109 @@ where } } -/// Format a decimal `mantissa` with `scale` fractional digits, stripping trailing -/// fractional zeros. Generic over the mantissa's `Display`, so it renders any Arrow -/// decimal width (`i32`/`i64`/`i128`/`i256`) directly — Decimal256 needs no fallback -/// to Arrow's own string conversion. +/// Renders a decimal from its `mantissa` and `scale`: the mantissa's digits with a +/// point inserted `scale` places from the right and trailing fractional zeros +/// stripped. Generic over the mantissa's `Display`, so it covers every Arrow +/// decimal width (`i32`/`i64`/`i128`/`i256`) — Decimal256 needs no fallback to +/// Arrow's own string conversion. +/// +/// Rendering through `fmt::Write` and a stack buffer means no allocation at all: +/// `collect_str` streams this straight into the response, and +/// [`decimal_to_string`] is the same text collected into a `String`. /// /// e.g. `(25987600, 5) -> "259.876"`, `(6199200000, 5) -> "61992"`, /// `(-250, 3) -> "-0.25"`, `(25, 5) -> "0.00025"`. -pub(crate) fn decimal_to_string(mantissa: T, scale: u32) -> String { - let raw = mantissa.to_string(); - if scale == 0 { - return raw; - } - - let scale = scale as usize; - let (sign, digits) = match raw.strip_prefix('-') { - Some(rest) => ("-", rest), - None => ("", raw.as_str()), - }; - - let (int_part, frac) = if digits.len() > scale { - let (int_part, frac) = digits.split_at(digits.len() - scale); - (int_part, frac.to_string()) - } else { - let pad = "0".repeat(scale - digits.len()); - ("0", format!("{pad}{digits}")) - }; - - let frac = frac.trim_end_matches('0'); - if frac.is_empty() { - format!("{sign}{int_part}") - } else { - format!("{sign}{int_part}.{frac}") +pub(crate) struct DecimalText { + pub mantissa: T, + pub scale: u32, +} + +impl fmt::Display for DecimalText { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if self.scale == 0 { + return write!(f, "{}", self.mantissa); + } + + let mut rendered = DigitBuf::default(); + write!(&mut rendered, "{}", self.mantissa)?; + let raw = rendered.as_str()?; + + let scale = self.scale as usize; + let (sign, digits) = match raw.strip_prefix('-') { + Some(rest) => ("-", rest), + None => ("", raw), + }; + + if digits.len() > scale { + let (int_part, frac) = digits.split_at(digits.len() - scale); + let frac = frac.trim_end_matches('0'); + f.write_str(sign)?; + f.write_str(int_part)?; + if !frac.is_empty() { + f.write_char('.')?; + f.write_str(frac)?; + } + return Ok(()); + } + + // Fewer digits than the scale, so the value is `0.` followed by the digits + // padded out to `scale`. Trailing zeros are stripped from the digits, which + // is where any of them can be. + let frac = digits.trim_end_matches('0'); + f.write_str(sign)?; + if frac.is_empty() { + return f.write_str("0"); + } + + f.write_str("0.")?; + for _ in 0..scale - digits.len() { + f.write_char('0')?; + } + f.write_str(frac) } } +/// Fixed-size sink for a mantissa's digits. `i256::MIN` is 78 digits plus a sign, +/// so this covers every Arrow decimal width with room to spare; a mantissa that +/// somehow overflowed it would surface as a formatting error rather than bad text. +struct DigitBuf { + bytes: [u8; 96], + len: usize, +} + +impl Default for DigitBuf { + fn default() -> Self { + Self { + bytes: [0; 96], + len: 0, + } + } +} + +impl DigitBuf { + fn as_str(&self) -> Result<&str, fmt::Error> { + std::str::from_utf8(&self.bytes[..self.len]).map_err(|_| fmt::Error) + } +} + +impl fmt::Write for DigitBuf { + fn write_str(&mut self, s: &str) -> fmt::Result { + let end = self.len + s.len(); + if end > self.bytes.len() { + return Err(fmt::Error); + } + + self.bytes[self.len..end].copy_from_slice(s.as_bytes()); + self.len = end; + Ok(()) + } +} + +/// [`DecimalText`] collected into a `String`, for the paths that keep the value. +pub(crate) fn decimal_to_string(mantissa: T, scale: u32) -> String { + DecimalText { mantissa, scale }.to_string() +} + #[cfg(test)] mod tests { use super::*; @@ -703,4 +837,118 @@ mod tests { ); } } + + /// The allocation-free renderer must agree with the straightforward + /// String-building version everywhere, not just on the cases above. This is + /// that version, kept only as the oracle below. + fn decimal_to_string_reference(raw: String, scale: u32) -> String { + if scale == 0 { + return raw; + } + + let scale = scale as usize; + let (sign, digits) = match raw.strip_prefix('-') { + Some(rest) => ("-", rest), + None => ("", raw.as_str()), + }; + + let (int_part, frac) = if digits.len() > scale { + let (int_part, frac) = digits.split_at(digits.len() - scale); + (int_part, frac.to_string()) + } else { + let pad = "0".repeat(scale - digits.len()); + ("0", format!("{pad}{digits}")) + }; + + let frac = frac.trim_end_matches('0'); + if frac.is_empty() { + format!("{sign}{int_part}") + } else { + format!("{sign}{int_part}.{frac}") + } + } + + #[test] + fn test_decimal_to_string_matches_reference() { + let mantissas: Vec = [ + 0i128, + 1, + 5, + 9, + 10, + 99, + 100, + 101, + 1_000, + 1_005, + 10_000_000, + 123_456_789, + 999_999_999_999, + 1_000_000_000_000, + i128::MAX, + i128::MIN, + ] + .into_iter() + // `-i128::MIN` overflows, so only negate what can be negated. + .flat_map(|m| [Some(m), m.checked_neg()]) + .flatten() + .collect(); + + for mantissa in mantissas { + for scale in [0u32, 1, 2, 3, 5, 9, 12, 20, 38, 39, 40] { + assert_eq!( + decimal_to_string(mantissa, scale), + decimal_to_string_reference(mantissa.to_string(), scale), + "mantissa={mantissa} scale={scale}" + ); + } + } + } + + /// The widest mantissa Arrow can hand over must still render, i.e. the stack + /// buffer has to be large enough for it. + #[test] + fn test_decimal_to_string_i256_extremes() { + for mantissa in [i256::MAX, i256::MIN, i256::from_i128(-1), i256::ZERO] { + for scale in [0u32, 2, 38, 76] { + assert_eq!( + decimal_to_string(mantissa, scale), + decimal_to_string_reference(mantissa.to_string(), scale), + "mantissa={mantissa} scale={scale}" + ); + } + } + } + + /// Serializing an unrendered decimal must produce exactly the JSON the owned + /// `String` cell produces. + #[test] + fn test_decimal_cell_serializes_like_owned_string() { + for (mantissa, scale) in [ + (239996i128, 2u32), + (-250, 3), + (0, 5), + (i128::MIN, 38), + (6199200000, 5), + ] { + let unrendered = CellRef::Decimal128 { mantissa, scale }; + let owned = DBResponsePrimitive::String(decimal_to_string(mantissa, scale)); + + assert_eq!( + serde_json::to_string(&unrendered).unwrap(), + serde_json::to_string(&owned).unwrap(), + "mantissa={mantissa} scale={scale}" + ); + } + + let big = i256::MAX; + assert_eq!( + serde_json::to_string(&CellRef::Decimal256 { + mantissa: big, + scale: 4 + }) + .unwrap(), + serde_json::to_string(&DBResponsePrimitive::String(decimal_to_string(big, 4))).unwrap(), + ); + } } diff --git a/rust/cube/cubeorchestrator/src/query_result_transform.rs b/rust/cube/cubeorchestrator/src/query_result_transform.rs index f1fb1bdab853f..7914fe938075b 100644 --- a/rust/cube/cubeorchestrator/src/query_result_transform.rs +++ b/rust/cube/cubeorchestrator/src/query_result_transform.rs @@ -3881,7 +3881,7 @@ pub(crate) mod tests { // `Cube.shippedAt`. That keeps the appended-member order // deterministic across calls: `get_members` walks a `HashMap` // there, so two of them would come out in either order and the - // byte-equality assertions below would flake. + // byte-equality assertions would flake. .chain([MemberOrMemberExpression::Member( "Cube.createdAt".to_string(), )]) From 3fc4a84851be66a88b6a6a359bcd92b7b7faaf12 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Thu, 30 Jul 2026 14:58:18 +0200 Subject: [PATCH 5/5] refactor(cubeorchestrator): build bench columns straight from the iterator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each Arrow column in the bench fixture was collected into a `Vec` and handed to `Array::from`, which copies it into the Arrow buffer — so every column existed twice during setup, and the string case held `row_count` live `String` allocations at once (100k of them at the top fixture size). `from_iter_values` fills the buffer from the iterator directly. The emitted IPC payloads are byte-identical, so recorded measurements stay comparable. Co-Authored-By: Claude Opus 5 (1M context) --- .../cubeorchestrator/benches/common/mod.rs | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/rust/cube/cubeorchestrator/benches/common/mod.rs b/rust/cube/cubeorchestrator/benches/common/mod.rs index 3df2b004bf658..37da441f17f34 100644 --- a/rust/cube/cubeorchestrator/benches/common/mod.rs +++ b/rust/cube/cubeorchestrator/benches/common/mod.rs @@ -126,21 +126,23 @@ pub fn build_arrow_ipc( let mut fields = Vec::with_capacity(total_cols); let mut columns: Vec = Vec::with_capacity(total_cols); + // Every column is built with `from_iter_values`, which fills the Arrow buffer + // straight from the iterator. Collecting into a `Vec` first would hold a + // second copy of the whole column — and for the string case, `row_count` live + // `String` allocations — just to hand it over. for (j, (_, alias)) in dimensions.iter().enumerate() { fields.push(Field::new(alias.clone(), DataType::Utf8, false)); - let values: Vec = (0..row_count) - .map(|i| format!("dim_{}_{}", j, i % 1000)) - .collect(); - columns.push(Arc::new(StringArray::from(values))); + columns.push(Arc::new(StringArray::from_iter_values( + (0..row_count).map(|i| format!("dim_{}_{}", j, i % 1000)), + ))); } for (j, (_, alias)) in measures.iter().enumerate() { match measure_kind { MeasureKind::Float64 => { fields.push(Field::new(alias.clone(), DataType::Float64, false)); - let values: Vec = (0..row_count) - .map(|i| ((i * (j + 1)) as f64) * 0.5) - .collect(); - columns.push(Arc::new(Float64Array::from(values))); + columns.push(Arc::new(Float64Array::from_iter_values( + (0..row_count).map(|i| ((i * (j + 1)) as f64) * 0.5), + ))); } MeasureKind::Decimal128 => { fields.push(Field::new( @@ -149,13 +151,12 @@ pub fn build_arrow_ipc( false, )); // Same magnitudes as the Float64 arm, as a scale-2 mantissa. - let values: Vec = (0..row_count) - .map(|i| ((i * (j + 1)) as i128) * 50) - .collect(); columns.push(Arc::new( - Decimal128Array::from(values) - .with_precision_and_scale(38, 2) - .expect("decimal precision"), + Decimal128Array::from_iter_values( + (0..row_count).map(|i| ((i * (j + 1)) as i128) * 50), + ) + .with_precision_and_scale(38, 2) + .expect("decimal precision"), )); } } @@ -167,10 +168,9 @@ pub fn build_arrow_ipc( false, )); // One day apart, offset per column — arbitrary but realistic spread. - let values: Vec = (0..row_count) - .map(|i| ((i + j) as i64) * 86_400_000) - .collect(); - columns.push(Arc::new(TimestampMillisecondArray::from(values))); + columns.push(Arc::new(TimestampMillisecondArray::from_iter_values( + (0..row_count).map(|i| ((i + j) as i64) * 86_400_000), + ))); } let schema = Arc::new(Schema::new(fields));