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..37da441f17f34 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]; @@ -79,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; @@ -99,19 +126,40 @@ 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() { - 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)); + 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( + alias.clone(), + DataType::Decimal128(38, 2), + false, + )); + // Same magnitudes as the Float64 arm, as a scale-2 mantissa. + columns.push(Arc::new( + Decimal128Array::from_iter_values( + (0..row_count).map(|i| ((i * (j + 1)) as i128) * 50), + ) + .with_precision_and_scale(38, 2) + .expect("decimal precision"), + )); + } + } } for (j, td) in time_dims.iter().enumerate() { fields.push(Field::new( @@ -120,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)); @@ -137,3 +184,56 @@ 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], + measure_kind: MeasureKind, +) -> QueryResult { + 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 7358bd56ed6de..bcb150053e5ac 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, MeasureKind, 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"); @@ -202,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(); @@ -214,8 +182,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..4089b55014fda 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, MeasureKind, + TimeColumn, COLUMN_COUNTS, ROW_COUNTS, }; /// Total columns and row count used by `bench_transform_time_scenarios`. @@ -152,32 +154,49 @@ 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, + &[], + MeasureKind::Float64, + ), + ), + ]; // 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 +221,28 @@ 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, + MeasureKind::Float64, + ), + ), + ]; // Throughput in cells/sec; total cells = row_count * total_cols, where // total_cols == SCENARIO_COL_COUNT regardless of scenario. @@ -216,30 +250,157 @@ 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"), + ), + ( + 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, + ), + ), + ]; + + 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..47303989d36bc --- /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) + })?; + } + 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..257c97520c1be 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,13 @@ mod tests { use cubeshared::flatbuffers::FlatBufferBuilder; use std::sync::Arc; + /// 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 fn create_test_message(num_rows: usize, num_columns: usize) -> Vec { let mut builder = FlatBufferBuilder::new(); @@ -717,7 +563,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 +571,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 +580,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 +620,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 +653,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 +663,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 +680,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 +843,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..a5df57e71c399 --- /dev/null +++ b/rust/cube/cubeorchestrator/src/query_result_column.rs @@ -0,0 +1,954 @@ +//! 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::{i256, ArrowTemporalType, DataType, TimeUnit}; +use serde::{Serialize, Serializer}; +use std::{ + borrow::Cow, + convert::Infallible, + fmt::{self, Write as _}, +}; + +/// 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()?)), + } + } +} + +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() + } + + /// 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) + } +} + +/// A cell on its way out of a column. +/// +/// [`DBResponsePrimitive`] owns its `String`, so building one from an Arrow `Utf8` +/// 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), + /// 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), +} + +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::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` 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), + } + } +} + +/// 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)) + } + }}; +} + +/// 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> { + 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`, 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(CellRef<'_>) -> R, + ) -> R { + 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.cell_without_alloc(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. + /// 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(CellRef<'_>) -> Result<(), E>, + ) -> Result<(), E> { + match self { + ColumnReader::Primitives(cells) => { + if is_identity_transform(member_type) { + for cell in cells.iter() { + visit(CellRef::Primitive(cell))?; + } + } else { + for cell in cells.iter() { + visit(CellRef::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 + ))) + } + } + } + + /// 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 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, + } + } + + /// 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(CellRef<'_>) -> Result<(), E>, + ) -> Result<(), E> { + let len = self.len(); + + macro_rules! fill { + ($read:expr) => {{ + for row in 0..len { + visit(CellRef::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)) + }}; + } + + /// 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() + )) + ) + } + }}; + } + + /// 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(_) => { + for _ in 0..len { + visit(CellRef::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_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)), + 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_decimal!(a, Decimal128, *scale, |v| DBResponsePrimitive::String( + decimal_to_string(v, *scale) + )) + } + ArrowCellReader::Decimal256(a, scale) => { + fill_decimal!(a, Decimal256, *scale, |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, + } +} + +/// 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) 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::*; + + #[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}" + ); + } + } + + /// 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 8be457a5c2830..7914fe938075b 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,8 +1422,9 @@ impl std::ops::DerefMut for ColumnarArray { } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::*; + use crate::query_result_column::QueryResultColumn; use crate::transport::JsRawColumnarData; use anyhow::Result; use serde_json::from_str; @@ -1392,13 +1439,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 +2219,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 +3731,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 +3761,242 @@ 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)]; + + /// 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 + /// 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, + ), + // 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( + 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, + ])), + Arc::new(StringArray::from(vec![ + Some("2024-06-15 00:00:00.000"), + None, + Some("2024-06-16T00:00:00"), + ])), + ], + ) + .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(), materialize_columns(&arrow)?)?; + + 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"), + ("cube__shipped_at_day", "Cube.shippedAt.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())) + // `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 would flake. + .chain([MemberOrMemberExpression::Member( + "Cube.createdAt".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(); + 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"), + "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, + } + } }