diff --git a/packages/cubejs-backend-native/js/index.ts b/packages/cubejs-backend-native/js/index.ts index 9a04a5bf6ccfa..3ffc19e6f4f57 100644 --- a/packages/cubejs-backend-native/js/index.ts +++ b/packages/cubejs-backend-native/js/index.ts @@ -290,6 +290,14 @@ function wrapNativeFunctionWithStream( objectMode: true, highWaterMark: chunkLength, write(row: any, encoding: BufferEncoding, callback: (error?: (Error | null)) => void) { + // Columnar batch produced by a driver (base-driver `ColumnarResponse`, + // marker `COLUMNAR_RESPONSE_TYPE`): already `{ members, columns }`, so send it + // straight through without row accumulation or the row→columnar transpose. + if (row && row.$type === 'ColumnarResponse') { + const toSend = Buffer.from(JSON.stringify({ members: row.members, columns: row.columns })); + writerOrChannel.chunk(toSend, callback); + return; + } chunkBuffer.push(row); if (chunkBuffer.length < chunkLength) { callback(null); diff --git a/packages/cubejs-backend-native/src/orchestrator.rs b/packages/cubejs-backend-native/src/orchestrator.rs index 4524a7e502b62..6b0670dfe6219 100644 --- a/packages/cubejs-backend-native/src/orchestrator.rs +++ b/packages/cubejs-backend-native/src/orchestrator.rs @@ -1,8 +1,9 @@ use crate::node_obj_deserializer::JsValueDeserializer; use crate::transport::MapCubeErrExt; +use cubeorchestrator::arrow_column::DBResponseValueRef; use cubeorchestrator::query_message_parser::QueryResult; use cubeorchestrator::query_result_transform::{ - DBResponsePrimitive, RequestResultData, RequestResultDataMulti, TransformedData, + RequestResultData, RequestResultDataMulti, TransformedData, }; use cubeorchestrator::transport::{JsRawColumnarData, TransformDataRequest}; use cubesql::compile::engine::df::scan::{ColumnarValueObject, FieldValue}; @@ -128,18 +129,22 @@ impl ResultWrapper { } } -fn db_primitive_to_field_value(value: &DBResponsePrimitive) -> FieldValue<'_> { +// String cells borrow their bytes from the backing column (a materialized +// primitive or an Arrow buffer) — no per-cell allocation on the way to the +// wire encoder. +fn db_value_to_field_value(value: DBResponseValueRef<'_>) -> FieldValue<'_> { match value { - DBResponsePrimitive::String(s) => FieldValue::String(Cow::Borrowed(s)), - DBResponsePrimitive::Int64(n) => FieldValue::Number(*n as f64), - DBResponsePrimitive::UInt64(n) => FieldValue::Number(*n as f64), - DBResponsePrimitive::Float64(n) => FieldValue::Number(*n), - DBResponsePrimitive::Boolean(b) => FieldValue::Bool(*b), - DBResponsePrimitive::Timestamp(_) => FieldValue::String(Cow::Owned(value.to_string())), - DBResponsePrimitive::Uncommon(v) => FieldValue::String(Cow::Owned( - serde_json::to_string(&v).unwrap_or_else(|_| v.to_string()), + DBResponseValueRef::Str(s) => FieldValue::String(Cow::Borrowed(s)), + DBResponseValueRef::StringOwned(s) => FieldValue::String(Cow::Owned(s)), + DBResponseValueRef::Int64(n) => FieldValue::Number(n as f64), + DBResponseValueRef::UInt64(n) => FieldValue::Number(n as f64), + DBResponseValueRef::Float64(n) => FieldValue::Number(n), + DBResponseValueRef::Boolean(b) => FieldValue::Bool(b), + DBResponseValueRef::Timestamp(_) => FieldValue::String(Cow::Owned(value.to_string())), + DBResponseValueRef::Uncommon(v) => FieldValue::String(Cow::Owned( + serde_json::to_string(v).unwrap_or_else(|_| v.to_string()), )), - DBResponsePrimitive::Null => FieldValue::Null, + DBResponseValueRef::Null => FieldValue::Null, } } @@ -189,7 +194,7 @@ impl ColumnarValueObject for ResultWrapper { }; Ok(Box::new( - column.iter().map(|v| Ok(db_primitive_to_field_value(v))), + column.iter_values().map(|v| Ok(db_value_to_field_value(v))), )) } } @@ -257,8 +262,11 @@ pub fn get_cubestore_result(mut cx: FunctionContext) -> JsResult { result.members().iter().map(|k| cx.string(k)).collect(); let row_count = result.row_count(); - let columns: Vec<_> = (0..js_keys.len()) - .map(|i| result.column(i)) + // Collect borrowed per-cell views up front: the row-major loop below + // needs random access, and value refs are cheap (string cells borrow + // from the backing column instead of copying). + let columns: Vec>> = (0..js_keys.len()) + .map(|i| result.column(i).map(|col| col.iter_values().collect())) .collect::>() .or_else(|err| cx.throw_error(err.to_string()))?; let js_array = JsArray::new(&mut cx, row_count); @@ -270,8 +278,9 @@ pub fn get_cubestore_result(mut cx: FunctionContext) -> JsResult { for (col_idx, js_key) in js_keys.iter().enumerate() { let value = &columns[col_idx][row_idx]; let js_value: Handle<'_, JsValue> = match value { - DBResponsePrimitive::Null => cx.null().upcast(), + DBResponseValueRef::Null => cx.null().upcast(), // For compatibility, we convert all primitives to strings + DBResponseValueRef::Str(s) => cx.string(s).upcast(), other => cx.string(other.to_string()).upcast(), }; diff --git a/packages/cubejs-base-driver/src/ColumnarResponse.ts b/packages/cubejs-base-driver/src/ColumnarResponse.ts new file mode 100644 index 0000000000000..1f2ffd9c6b449 --- /dev/null +++ b/packages/cubejs-base-driver/src/ColumnarResponse.ts @@ -0,0 +1,75 @@ +import * as stream from 'stream'; + +/** + * Marker set on every batch object emitted by `ColumnarResponse`. Downstream consumers + * (the query orchestrator's `QueryStream` and `cubejs-backend-native`) key on this to tell + * a columnar batch apart from a plain row object. + */ +export const COLUMNAR_RESPONSE_TYPE = 'ColumnarResponse'; + +/** + * A single columnar batch as it flows through the object-mode stream. + * `members[j]` names `columns[j]`; every `columns[j]` has the same length (the batch's + * row count). Matches the `{ members, columns }` shape decoded on the Rust side + * (`JsonColumnarValueObject`). + */ +export interface ColumnarBatch { + $type: typeof COLUMNAR_RESPONSE_TYPE; + members: string[]; + columns: any[][]; +} + +/** + * Object-mode Transform that pivots position-indexed array rows (e.g. pg `rowMode: 'array'`) + * into columnar batches. Buffering by column avoids allocating a per-row object per row and + * moves the row→columnar pivot out of `cubejs-backend-native`, so the SQL-API streaming path + * can hand columnar batches straight to Rust. + * + * Input: `unknown[]` rows whose cells are ordered to match `members`. + * Output: `ColumnarBatch` objects, one every `batchSize` rows plus a final partial batch. + */ +export class ColumnarResponse extends stream.Transform { + private readonly members: string[]; + + private readonly batchSize: number; + + private columns: any[][]; + + private count = 0; + + public constructor({ members, batchSize }: { members: string[]; batchSize: number }) { + super({ objectMode: true }); + this.members = members; + this.batchSize = batchSize; + this.columns = members.map(() => []); + } + + public _transform(row: unknown[], _encoding: BufferEncoding, callback: stream.TransformCallback) { + for (let j = 0; j < this.members.length; j++) { + this.columns[j].push(row[j]); + } + this.count++; + if (this.count >= this.batchSize) { + this.flushBatch(); + } + callback(); + } + + public _flush(callback: stream.TransformCallback) { + if (this.count > 0) { + this.flushBatch(); + } + callback(); + } + + private flushBatch() { + const batch: ColumnarBatch = { + $type: COLUMNAR_RESPONSE_TYPE, + members: this.members, + columns: this.columns, + }; + this.push(batch); + this.columns = this.members.map(() => []); + this.count = 0; + } +} diff --git a/packages/cubejs-base-driver/src/driver.interface.ts b/packages/cubejs-base-driver/src/driver.interface.ts index 5fa35e06b3e7d..7eac595417a81 100644 --- a/packages/cubejs-base-driver/src/driver.interface.ts +++ b/packages/cubejs-base-driver/src/driver.interface.ts @@ -131,6 +131,12 @@ export interface DriverCapabilities extends ExternalDriverCompatibilities { export type StreamOptions = { highWaterMark: number; + /** + * When set, the driver is asked to emit columnar batches (see `ColumnarResponse`) + * instead of row objects. It is a hint: drivers that don't support it keep emitting + * rows. Currently set to `true` only for the cubesql SQL-API streaming path. + */ + columnar?: boolean; }; export type StreamingSourceOptions = { diff --git a/packages/cubejs-base-driver/src/index.ts b/packages/cubejs-base-driver/src/index.ts index 84e8d26c5e63b..af1f2acfde197 100644 --- a/packages/cubejs-base-driver/src/index.ts +++ b/packages/cubejs-base-driver/src/index.ts @@ -1,4 +1,5 @@ export * from './BaseDriver'; +export * from './ColumnarResponse'; export * from './utils'; export * from './driver.interface'; export * from './queue-driver.interface'; diff --git a/packages/cubejs-base-driver/test/unit/ColumnarResponse.test.ts b/packages/cubejs-base-driver/test/unit/ColumnarResponse.test.ts new file mode 100644 index 0000000000000..a9ab54181a12c --- /dev/null +++ b/packages/cubejs-base-driver/test/unit/ColumnarResponse.test.ts @@ -0,0 +1,77 @@ +import { Readable } from 'stream'; +import { ColumnarResponse, COLUMNAR_RESPONSE_TYPE } from '../../src'; + +async function collect(members: string[], batchSize: number, rows: unknown[][]) { + const source = Readable.from(rows, { objectMode: true }); + const columnar = new ColumnarResponse({ members, batchSize }); + const batches: any[] = []; + source.pipe(columnar); + for await (const batch of columnar) { + batches.push(batch); + } + return batches; +} + +describe('ColumnarResponse', () => { + const members = ['a', 'b']; + + test('emits a batch at every batchSize boundary and a final partial batch', async () => { + const rows = [ + [1, 'x'], + [2, 'y'], + [3, 'z'], + [4, 'w'], + [5, 'v'], + ]; + + const batches = await collect(members, 2, rows); + + expect(batches).toHaveLength(3); // 2 + 2 + 1 + batches.forEach((batch) => { + expect(batch.$type).toBe(COLUMNAR_RESPONSE_TYPE); + expect(batch.members).toEqual(members); + }); + + expect(batches[0].columns).toEqual([[1, 2], ['x', 'y']]); + expect(batches[1].columns).toEqual([[3, 4], ['z', 'w']]); + expect(batches[2].columns).toEqual([[5], ['v']]); + }); + + test('transposes positional array rows into columns', async () => { + const rows = [ + [1, 'x'], + [2, 'y'], + ]; + + const batches = await collect(members, 100, rows); + + expect(batches).toHaveLength(1); + expect(batches[0]).toEqual({ + $type: COLUMNAR_RESPONSE_TYPE, + members, + columns: [[1, 2], ['x', 'y']], + }); + }); + + test('does not emit anything for an empty stream', async () => { + const batches = await collect(members, 2, []); + expect(batches).toHaveLength(0); + }); + + test('reuses fresh column arrays per batch (no shared references)', async () => { + const rows = [ + [1, 'x'], + [2, 'y'], + [3, 'z'], + [4, 'w'], + ]; + + const batches = await collect(members, 2, rows); + + expect(batches).toHaveLength(2); + // First batch's columns must not be mutated by the second batch. + expect(batches[0].columns).toEqual([[1, 2], ['x', 'y']]); + expect(batches[1].columns).toEqual([[3, 4], ['z', 'w']]); + expect(batches[0].columns).not.toBe(batches[1].columns); + }); +}); diff --git a/packages/cubejs-postgres-driver/src/PostgresDriver.ts b/packages/cubejs-postgres-driver/src/PostgresDriver.ts index e083b5d0104a4..89623f68b425a 100644 --- a/packages/cubejs-postgres-driver/src/PostgresDriver.ts +++ b/packages/cubejs-postgres-driver/src/PostgresDriver.ts @@ -10,6 +10,7 @@ import { types, FieldDef } from 'pg'; import { TypeId, TypeFormat } from 'pg-types'; import { BaseDriver, + ColumnarResponse, DownloadQueryResultsOptions, DownloadTableMemoryData, DriverInterface, GenericDataBaseType, IndexesSQL, TableStructure, StreamOptions, StreamTableDataWithTypes, QueryOptions, DownloadQueryResultsResult, DriverCapabilities, TableColumn, createPoolName, @@ -353,7 +354,7 @@ export class PostgresDriver { PostgresDriver.checkValuesLimit(values); @@ -366,17 +367,38 @@ export class PostgresDriver { + await this.pool.release(conn); + }; + + if (columnar) { + const members = fields.map((f) => f.name); + const columnarStream = new ColumnarResponse({ members, batchSize: highWaterMark }); + // `pipe()` does not forward errors — propagate manually so the consumer sees the + // failure and `release()` still runs during cleanup. + rowStream.once('error', (e) => columnarStream.destroy(e)); + rowStream.pipe(columnarStream); + + return { + rowStream: columnarStream, + types: mappedTypes, + release, + }; + } return { rowStream, - types: this.mapFields(fields), - release: async () => { - await this.pool.release(conn); - } + types: mappedTypes, + release, }; } catch (e) { await this.pool.release(conn); diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts index d4d6e85cecb53..75e625314a46a 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryCache.ts @@ -663,7 +663,10 @@ export class QueryCache { let logged = false; Promise .all([clientFactory()]) - .then(([client]) => (client).stream(req.query, req.values, { highWaterMark: getEnv('dbQueryStreamHighWaterMark') })) + // This handler is the cubesql SQL-API streaming path, which consumes columnar + // batches natively — ask the driver to emit them (drivers that don't support it + // fall back to row objects). + .then(([client]) => (client).stream(req.query, req.values, { highWaterMark: getEnv('dbQueryStreamHighWaterMark'), columnar: true })) .then((source) => { const cleanup = async (error) => { if (source.release) { diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryStream.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryStream.ts index 593ab7e4bfbfe..c8ec6149d0d8d 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryStream.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryStream.ts @@ -1,5 +1,6 @@ import * as stream from 'stream'; import { getEnv } from '@cubejs-backend/shared'; +import { COLUMNAR_RESPONSE_TYPE } from '@cubejs-backend/base-driver'; export class QueryStream extends stream.Transform { private timeout = 5 * 60 * 1000; @@ -43,6 +44,24 @@ export class QueryStream extends stream.Transform { if (this.streams.has(this.queryKey)) { this.streams.delete(this.queryKey); } + + // Columnar batch (see `ColumnarResponse` in base-driver): remap the `members` array + // (aliases → member names) once and forward the batch untouched. + if (chunk && chunk.$type === COLUMNAR_RESPONSE_TYPE) { + const members = this.aliasNameToMember + ? chunk.members.map((alias) => this.aliasNameToMember[alias]) + : chunk.members; + const rowCount = chunk.columns[0] ? chunk.columns[0].length : 0; + if (this.counter + rowCount < this.readableHighWaterMark) { + this.counter += rowCount; + } else { + this.counter = 0; + this.debounce(); + } + callback(null, { $type: COLUMNAR_RESPONSE_TYPE, members, columns: chunk.columns }); + return; + } + let row = {}; if (this.aliasNameToMember) { Object.keys(chunk).forEach((alias) => { diff --git a/rust/cube/cubeorchestrator/src/arrow_column.rs b/rust/cube/cubeorchestrator/src/arrow_column.rs new file mode 100644 index 0000000000000..69477d5d411b4 --- /dev/null +++ b/rust/cube/cubeorchestrator/src/arrow_column.rs @@ -0,0 +1,404 @@ +use crate::query_message_parser::ParseError; +use crate::query_result_transform::{DBResponsePrimitive, TIMESTAMP_ITEMS}; +use arrow::array::{ + Array, ArrayRef, 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 chrono::NaiveDateTime; +use serde::Serialize; +use serde_json::Value; +use std::fmt::Display; + +/// Borrowed view of one result cell. Serializes to exactly the same JSON as the +/// corresponding [`DBResponsePrimitive`], but string cells borrow their bytes +/// (from an Arrow buffer or a materialized primitive) instead of owning them, +/// so a serialize pass over an Arrow column does zero per-cell allocations. +#[derive(Debug, Clone)] +pub enum DBResponseValueRef<'a> { + Null, + Boolean(bool), + Int64(i64), + UInt64(u64), + Float64(f64), + Str(&'a str), + /// Value rendered at read time (Arrow decimals). Owning is unavoidable here: + /// the text doesn't exist in the source buffer. + StringOwned(String), + Timestamp(NaiveDateTime), + Uncommon(&'a Value), +} + +impl<'a> From<&'a DBResponsePrimitive> for DBResponseValueRef<'a> { + fn from(value: &'a DBResponsePrimitive) -> Self { + match value { + DBResponsePrimitive::Null => DBResponseValueRef::Null, + DBResponsePrimitive::Boolean(b) => DBResponseValueRef::Boolean(*b), + DBResponsePrimitive::Int64(n) => DBResponseValueRef::Int64(*n), + DBResponsePrimitive::UInt64(n) => DBResponseValueRef::UInt64(*n), + DBResponsePrimitive::Float64(n) => DBResponseValueRef::Float64(*n), + DBResponsePrimitive::String(s) => DBResponseValueRef::Str(s), + DBResponsePrimitive::Timestamp(dt) => DBResponseValueRef::Timestamp(*dt), + DBResponsePrimitive::Uncommon(v) => DBResponseValueRef::Uncommon(v), + } + } +} + +impl DBResponseValueRef<'_> { + pub fn into_primitive(self) -> DBResponsePrimitive { + match self { + DBResponseValueRef::Null => DBResponsePrimitive::Null, + DBResponseValueRef::Boolean(b) => DBResponsePrimitive::Boolean(b), + DBResponseValueRef::Int64(n) => DBResponsePrimitive::Int64(n), + DBResponseValueRef::UInt64(n) => DBResponsePrimitive::UInt64(n), + DBResponseValueRef::Float64(n) => DBResponsePrimitive::Float64(n), + DBResponseValueRef::Str(s) => DBResponsePrimitive::String(s.to_owned()), + DBResponseValueRef::StringOwned(s) => DBResponsePrimitive::String(s), + DBResponseValueRef::Timestamp(dt) => DBResponsePrimitive::Timestamp(dt), + DBResponseValueRef::Uncommon(v) => DBResponsePrimitive::Uncommon(v.clone()), + } + } +} + +// Must stay in lockstep with `Serialize for DBResponsePrimitive`: numeric +// variants render as JSON strings, timestamps use the fixed ISO format. +impl Serialize for DBResponseValueRef<'_> { + fn serialize(&self, serializer: S) -> Result { + match self { + DBResponseValueRef::Null => serializer.serialize_none(), + DBResponseValueRef::Boolean(b) => serializer.serialize_bool(*b), + DBResponseValueRef::Int64(n) => serializer.collect_str(n), + DBResponseValueRef::UInt64(n) => serializer.collect_str(n), + DBResponseValueRef::Float64(n) => serializer.collect_str(n), + DBResponseValueRef::Str(s) => serializer.serialize_str(s), + DBResponseValueRef::StringOwned(s) => serializer.serialize_str(s), + DBResponseValueRef::Timestamp(dt) => { + serializer.collect_str(&dt.format_with_items(TIMESTAMP_ITEMS.iter())) + } + DBResponseValueRef::Uncommon(v) => v.serialize(serializer), + } + } +} + +impl Display for DBResponseValueRef<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DBResponseValueRef::Null => f.write_str("null"), + DBResponseValueRef::Boolean(b) => write!(f, "{}", b), + DBResponseValueRef::Int64(n) => write!(f, "{}", n), + DBResponseValueRef::UInt64(n) => write!(f, "{}", n), + DBResponseValueRef::Float64(n) => write!(f, "{}", n), + DBResponseValueRef::Str(s) => f.write_str(s), + DBResponseValueRef::StringOwned(s) => f.write_str(s), + DBResponseValueRef::Timestamp(dt) => { + write!(f, "{}", dt.format_with_items(TIMESTAMP_ITEMS.iter())) + } + DBResponseValueRef::Uncommon(v) => { + let s = serde_json::to_string(v).unwrap_or_else(|_| v.to_string()); + f.write_str(&s) + } + } + } +} + +// Mirrors `PartialEq for DBResponsePrimitive` by value: `Str` and `StringOwned` +// compare as the same string kind, everything else is variant-strict. +impl PartialEq for DBResponseValueRef<'_> { + fn eq(&self, other: &Self) -> bool { + use DBResponseValueRef::*; + match (self, other) { + (Null, Null) => true, + (Boolean(a), Boolean(b)) => a == b, + (Int64(a), Int64(b)) => a == b, + (UInt64(a), UInt64(b)) => a == b, + (Float64(a), Float64(b)) => a == b, + (Str(a), Str(b)) => a == b, + (Str(a), StringOwned(b)) | (StringOwned(b), Str(a)) => *a == b.as_str(), + (StringOwned(a), StringOwned(b)) => a == b, + (Timestamp(a), Timestamp(b)) => a == b, + (Uncommon(a), Uncommon(b)) => a == b, + _ => false, + } + } +} + +/// Format a decimal `mantissa` with `scale` fractional digits, stripping trailing +/// fractional zeros. Generic over the mantissa's `Display`, so it renders any Arrow +/// decimal width (`i32`/`i64`/`i128`/`i256`) directly — Decimal256 needs no fallback +/// to Arrow's own string conversion. +/// +/// e.g. `(25987600, 5) -> "259.876"`, `(6199200000, 5) -> "61992"`, +/// `(-250, 3) -> "-0.25"`, `(25, 5) -> "0.00025"`. +pub 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}") + } +} + +/// One logical result column backed by the Arrow arrays it arrived in (one +/// chunk per IPC record batch). Values are read straight out of the Arrow +/// buffers on demand — parsing and the columnar transform never materialize +/// per-cell primitives, and cloning the column is an `Arc` bump per chunk. +#[derive(Debug, Clone)] +pub struct ArrowColumn { + chunks: Vec, + len: usize, +} + +impl ArrowColumn { + /// Validates every chunk's data type up front, so downstream readers can + /// assume the whole column is convertible. + pub fn try_new(chunks: Vec) -> Result { + for chunk in &chunks { + TypedChunk::try_new(chunk.as_ref())?; + } + let len = chunks.iter().map(|c| c.len()).sum(); + Ok(Self { chunks, len }) + } + + #[inline] + pub fn len(&self) -> usize { + self.len + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// True if any chunk holds string values. `transform_value` only rewrites + /// `String` cells (member type `"time"`), so a column without string chunks + /// passes the transform untouched and can be shared as-is. + pub fn has_string_values(&self) -> bool { + self.chunks.iter().any(|c| { + matches!( + c.data_type(), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) + }) + } + + pub fn iter_values(&self) -> impl Iterator> + '_ { + self.chunks.iter().flat_map(|chunk| { + let typed = TypedChunk::try_new(chunk.as_ref()) + .expect("chunk types validated in ArrowColumn::try_new"); + (0..chunk.len()).map(move |i| typed.value(i)) + }) + } +} + +/// Arrow array downcast once per chunk, so per-cell reads are a direct typed +/// access instead of a `data_type()` dispatch + `downcast_ref` per value. +#[derive(Clone, Copy)] +enum TypedChunk<'a> { + Null, + 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), + Decimal128 { + array: &'a Decimal128Array, + scale: u32, + }, + Decimal256 { + array: &'a Decimal256Array, + scale: u32, + }, +} + +macro_rules! downcast_chunk { + ($array:expr, $ty:ty) => { + $array.as_any().downcast_ref::<$ty>().ok_or_else(|| { + ParseError::ArrowError(format!( + "Failed to downcast Arrow array to {}", + stringify!($ty) + )) + })? + }; +} + +impl<'a> TypedChunk<'a> { + fn try_new(array: &'a dyn Array) -> Result { + Ok(match array.data_type() { + DataType::Null => TypedChunk::Null, + DataType::Boolean => TypedChunk::Boolean(downcast_chunk!(array, BooleanArray)), + DataType::Int8 => TypedChunk::Int8(downcast_chunk!(array, Int8Array)), + DataType::Int16 => TypedChunk::Int16(downcast_chunk!(array, Int16Array)), + DataType::Int32 => TypedChunk::Int32(downcast_chunk!(array, Int32Array)), + DataType::Int64 => TypedChunk::Int64(downcast_chunk!(array, Int64Array)), + DataType::UInt8 => TypedChunk::UInt8(downcast_chunk!(array, UInt8Array)), + DataType::UInt16 => TypedChunk::UInt16(downcast_chunk!(array, UInt16Array)), + DataType::UInt32 => TypedChunk::UInt32(downcast_chunk!(array, UInt32Array)), + DataType::UInt64 => TypedChunk::UInt64(downcast_chunk!(array, UInt64Array)), + DataType::Float16 => TypedChunk::Float16(downcast_chunk!(array, Float16Array)), + DataType::Float32 => TypedChunk::Float32(downcast_chunk!(array, Float32Array)), + DataType::Float64 => TypedChunk::Float64(downcast_chunk!(array, Float64Array)), + DataType::Utf8 => TypedChunk::Utf8(downcast_chunk!(array, StringArray)), + DataType::LargeUtf8 => TypedChunk::LargeUtf8(downcast_chunk!(array, LargeStringArray)), + DataType::Utf8View => TypedChunk::Utf8View(downcast_chunk!(array, StringViewArray)), + DataType::Date32 => TypedChunk::Date32(downcast_chunk!(array, Date32Array)), + DataType::Date64 => TypedChunk::Date64(downcast_chunk!(array, Date64Array)), + DataType::Timestamp(TimeUnit::Second, _) => { + TypedChunk::TimestampSecond(downcast_chunk!(array, TimestampSecondArray)) + } + DataType::Timestamp(TimeUnit::Millisecond, _) => { + TypedChunk::TimestampMillisecond(downcast_chunk!(array, TimestampMillisecondArray)) + } + DataType::Timestamp(TimeUnit::Microsecond, _) => { + TypedChunk::TimestampMicrosecond(downcast_chunk!(array, TimestampMicrosecondArray)) + } + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + TypedChunk::TimestampNanosecond(downcast_chunk!(array, TimestampNanosecondArray)) + } + DataType::Decimal128(_, scale) => TypedChunk::Decimal128 { + array: downcast_chunk!(array, Decimal128Array), + scale: (*scale).max(0) as u32, + }, + DataType::Decimal256(_, scale) => TypedChunk::Decimal256 { + array: downcast_chunk!(array, Decimal256Array), + scale: (*scale).max(0) as u32, + }, + other => return Err(ParseError::UnsupportedArrowType(format!("{other:?}"))), + }) + } + + fn value(&self, i: usize) -> DBResponseValueRef<'a> { + macro_rules! read { + ($a:expr, $variant:ident, $conv:expr) => {{ + if $a.is_null(i) { + DBResponseValueRef::Null + } else { + DBResponseValueRef::$variant($conv($a.value(i))) + } + }}; + } + macro_rules! read_datetime { + ($a:expr) => {{ + if $a.is_null(i) { + DBResponseValueRef::Null + } else { + match $a.value_as_datetime(i) { + Some(dt) => DBResponseValueRef::Timestamp(dt), + None => DBResponseValueRef::Null, + } + } + }}; + } + + match self { + TypedChunk::Null => DBResponseValueRef::Null, + TypedChunk::Boolean(a) => read!(a, Boolean, |v| v), + TypedChunk::Int8(a) => read!(a, Int64, |v| v as i64), + TypedChunk::Int16(a) => read!(a, Int64, |v| v as i64), + TypedChunk::Int32(a) => read!(a, Int64, |v| v as i64), + TypedChunk::Int64(a) => read!(a, Int64, |v| v), + TypedChunk::UInt8(a) => read!(a, UInt64, |v| v as u64), + TypedChunk::UInt16(a) => read!(a, UInt64, |v| v as u64), + TypedChunk::UInt32(a) => read!(a, UInt64, |v| v as u64), + TypedChunk::UInt64(a) => read!(a, UInt64, |v| v), + TypedChunk::Float16(a) => { + if a.is_null(i) { + DBResponseValueRef::Null + } else { + DBResponseValueRef::Float64(a.value(i).to_f64()) + } + } + TypedChunk::Float32(a) => read!(a, Float64, |v| v as f64), + TypedChunk::Float64(a) => read!(a, Float64, |v| v), + TypedChunk::Utf8(a) => read!(a, Str, |v| v), + TypedChunk::LargeUtf8(a) => read!(a, Str, |v| v), + TypedChunk::Utf8View(a) => read!(a, Str, |v| v), + TypedChunk::Date32(a) => read_datetime!(a), + TypedChunk::Date64(a) => read_datetime!(a), + TypedChunk::TimestampSecond(a) => read_datetime!(a), + TypedChunk::TimestampMillisecond(a) => read_datetime!(a), + TypedChunk::TimestampMicrosecond(a) => read_datetime!(a), + TypedChunk::TimestampNanosecond(a) => read_datetime!(a), + TypedChunk::Decimal128 { array, scale } => { + if array.is_null(i) { + DBResponseValueRef::Null + } else { + DBResponseValueRef::StringOwned(decimal_to_string(array.value(i), *scale)) + } + } + TypedChunk::Decimal256 { array, scale } => { + if array.is_null(i) { + DBResponseValueRef::Null + } else { + DBResponseValueRef::StringOwned(decimal_to_string(array.value(i), *scale)) + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_decimal_to_string() { + for (mantissa, scale, expected) in [ + (6199200000i128, 5u32, "61992"), + (25987600, 5, "259.876"), + (1500, 3, "1.5"), + (-250, 3, "-0.25"), + (0, 5, "0"), + (21098000, 5, "210.98"), + (100, 0, "100"), + (0, 0, "0"), + (-5, 0, "-5"), + (25, 5, "0.00025"), + (-1, 0, "-1"), + (i128::MAX, 0, "170141183460469231731687303715884105727"), + ] { + assert_eq!( + decimal_to_string(mantissa, scale), + expected, + "mantissa={mantissa} scale={scale}" + ); + } + } +} diff --git a/rust/cube/cubeorchestrator/src/lib.rs b/rust/cube/cubeorchestrator/src/lib.rs index 03f0453c16db7..5f9f8d95620af 100644 --- a/rust/cube/cubeorchestrator/src/lib.rs +++ b/rust/cube/cubeorchestrator/src/lib.rs @@ -1,3 +1,4 @@ +pub mod arrow_column; pub mod query_message_parser; 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..e0beafde2a78c 100644 --- a/rust/cube/cubeorchestrator/src/query_message_parser.rs +++ b/rust/cube/cubeorchestrator/src/query_message_parser.rs @@ -1,15 +1,9 @@ use crate::{ + arrow_column::ArrowColumn, 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::ArrayRef; use arrow::ipc::reader::StreamReader; use cubeshared::codegen::{ root_as_http_message_with_opts, HttpCommand, HttpQueryResultData, HttpResultSet, @@ -243,11 +237,10 @@ impl QueryResult { }; let n_cols = members.len(); - let data: Vec = if let Some(result_set_rows) = result_set.rows() { + let data: Vec> = if let Some(result_set_rows) = result_set.rows() { let row_count = result_set_rows.len(); - let mut data: Vec = (0..n_cols) - .map(|_| ColumnarArray::with_capacity(row_count)) - .collect(); + let mut data: Vec> = + (0..n_cols).map(|_| Vec::with_capacity(row_count)).collect(); for row in result_set_rows.iter() { let values = row.values().ok_or(ParseError::NullRow)?; @@ -270,12 +263,17 @@ impl QueryResult { data } else { - (0..n_cols).map(|_| ColumnarArray::new()).collect() + (0..n_cols).map(|_| Vec::new()).collect() }; - QueryResult::try_new(members, data) + QueryResult::try_new(members, data.into_iter().map(ColumnarArray::from).collect()) } + /// Wraps the Arrow arrays as-is (one chunk per record batch) instead of + /// materializing per-cell primitives: values are read straight from the + /// Arrow buffers when the result is transformed/serialized. Column data + /// types are still validated eagerly, so unsupported types fail here, at + /// parse time. 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 +282,31 @@ 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 columns: 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, chunks) in columns.iter_mut().enumerate() { + chunks.push(batch.column(idx).clone()); } } - let data: Vec = columns.into_iter().map(ColumnarArray::from).collect(); + let data: Vec = columns + .into_iter() + .map(|chunks| ArrowColumn::try_new(chunks).map(ColumnarArray::Arrow)) + .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), - } - } - } - }}; - } - - // 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); - } - } - 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(()) -} - #[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::{ @@ -717,7 +526,7 @@ mod tests { assert_eq!(result.data.len(), 3); assert_eq!( - result.data[0].as_slice(), + result.data[0].to_cells().as_ref(), &[ DBResponsePrimitive::String("Berlin".to_string()), DBResponsePrimitive::Null, @@ -725,7 +534,7 @@ mod tests { ] ); assert_eq!( - result.data[1].as_slice(), + result.data[1].to_cells().as_ref(), &[ DBResponsePrimitive::Float64(1.5), DBResponsePrimitive::Float64(2.0), @@ -734,18 +543,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(&result.data[1]).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 &result.data[2].to_cells()[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!(result.data[2].to_cells()[1], DBResponsePrimitive::Null); + let json = serde_json::to_value(&result.data[2]).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 +583,7 @@ mod tests { let result = QueryResult::from_arrow(&bytes)?; assert_eq!( - result.data[0].as_slice(), + result.data[0].to_cells().as_ref(), &[ DBResponsePrimitive::String("9999999999999999999999999999999999999.99".to_string()), DBResponsePrimitive::Null, @@ -807,7 +616,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(&result.data[0]).unwrap(); assert_eq!(json[0], serde_json::Value::Null); assert_eq!(json[1], "2399.96"); assert_eq!(json[2], "2249.91"); @@ -817,30 +626,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( @@ -909,14 +694,14 @@ mod tests { assert_eq!(result.members, vec!["city", "amount"]); assert_eq!(result.row_count, 2); assert_eq!( - result.data[0].as_slice(), + result.data[0].to_cells().as_ref(), &[ DBResponsePrimitive::String("Berlin".to_string()), DBResponsePrimitive::String("Lisbon".to_string()), ] ); assert_eq!( - result.data[1].as_slice(), + result.data[1].to_cells().as_ref(), &[ DBResponsePrimitive::Float64(1.5), DBResponsePrimitive::Float64(2.0), diff --git a/rust/cube/cubeorchestrator/src/query_result_transform.rs b/rust/cube/cubeorchestrator/src/query_result_transform.rs index 8be457a5c2830..fe2aa222adf34 100644 --- a/rust/cube/cubeorchestrator/src/query_result_transform.rs +++ b/rust/cube/cubeorchestrator/src/query_result_transform.rs @@ -1,4 +1,5 @@ use crate::{ + arrow_column::{ArrowColumn, DBResponseValueRef}, query_message_parser::QueryResult, transport::{ AnnotatedConfigItem, ConfigItem, MemberOrMemberExpression, MembersMap, NormalizedQuery, @@ -9,13 +10,15 @@ use anyhow::{bail, Context, Result}; use chrono::format::{Fixed, Item, Numeric, Pad}; use chrono::{DateTime, NaiveDateTime, TimeZone, Utc}; use indexmap::{Equivalent, IndexMap}; -use itertools::multizip; +use itertools::{multizip, Either}; use serde::{ de::{self, MapAccess, SeqAccess, Visitor}, + ser::SerializeSeq, Deserialize, Deserializer, Serialize, }; use serde_json::Value; use std::{ + borrow::Cow, collections::{hash_map::DefaultHasher, HashMap, HashSet}, fmt::Display, hash::{BuildHasher, Hash, Hasher}, @@ -457,12 +460,12 @@ pub fn get_members( /// eliminating the per-cell `db_data.data.get(col).and_then(...)` double /// lookup the row-major loop would otherwise do on every cell. 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`]. `column` is the + /// materialized view of the corresponding [`ColumnarArray`]: borrowed for + /// `Cells`-backed columns, converted once at plan build time for + /// Arrow-backed ones (this row-major output needs random access). Cell { - column: &'a [DBResponsePrimitive], + column: Cow<'a, [DBResponsePrimitive]>, member_type: &'a str, }, /// Constant value replicated across every row (the @@ -489,7 +492,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.data[column_index].to_cells(), member_type: annotation_item.member_type.as_deref().unwrap_or(""), }); } @@ -513,7 +516,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.data[column_index].to_cells(); entries.push(CompactPlanEntry::Cell { column, member_type, @@ -556,10 +559,10 @@ 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], + /// Materialized view of the corresponding [`ColumnarArray`]: borrowed for + /// `Cells`-backed columns, converted once at plan build time for + /// Arrow-backed ones (this row-major output needs random access). + column: Cow<'a, [DBResponsePrimitive]>, /// Interned IndexMap key for this column with a pre-computed hash. /// Cloned via [`Arc::clone`] per row (atomic refcount inc). key: Arc, @@ -634,7 +637,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.data[index].to_cells(); columns.push(VanillaColumnPlan { column, @@ -804,29 +807,43 @@ fn build_columnar_columns( db_data: &QueryResult, ) -> Vec { let row_count = db_data.row_count; - let mut columns: Vec = plan - .iter() - .map(|_| ColumnarArray::with_capacity(row_count)) - .collect(); - for (col_idx, plan_entry) in plan.iter().enumerate() { - let out = &mut columns[col_idx]; - match &plan_entry.source { + plan.iter() + .map(|plan_entry| 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)); + let src = &db_data.data[*index]; + match src { + // `transform_value` only rewrites `String` cells of "time" + // members, so an Arrow column without string chunks (or for + // a non-"time" member) passes through untouched — share it + // as-is: the clone is an `Arc` bump per chunk, not a copy. + ColumnarArray::Arrow(col) => { + if plan_entry.member_type == "time" && col.has_string_values() { + ColumnarArray::Cells( + col.iter_values() + .map(|v| { + transform_value(v.into_primitive(), plan_entry.member_type) + }) + .collect(), + ) + } else { + src.clone() + } + } + ColumnarArray::Cells(cells) => ColumnarArray::Cells( + cells + .iter() + .map(|cell| transform_value(cell.clone(), plan_entry.member_type)) + .collect(), + ), } } - ColumnarColumnSource::Constant(v) => { - out.resize(row_count, v.clone()); - } + ColumnarColumnSource::Constant(v) => ColumnarArray::Cells(vec![v.clone(); row_count]), ColumnarColumnSource::NullFilled => { - out.resize(row_count, DBResponsePrimitive::Null); + ColumnarArray::Cells(vec![DBResponsePrimitive::Null; row_count]) } - } - } - - columns + }) + .collect() } /// Convert DB response object to the vanilla output format. Keys are @@ -1166,7 +1183,7 @@ pub enum DBResponsePrimitive { } /// `%Y-%m-%dT%H:%M:%S%.3f` -const TIMESTAMP_ITEMS: &[Item<'static>] = &[ +pub(crate) const TIMESTAMP_ITEMS: &[Item<'static>] = &[ Item::Numeric(Numeric::Year, Pad::Zero), Item::Literal("-"), Item::Numeric(Numeric::Month, Pad::Zero), @@ -1325,53 +1342,114 @@ impl Display for DBResponsePrimitive { } } -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] -#[serde(transparent)] -pub struct ColumnarArray(pub Vec); +/// One logical result column. `Cells` is the materialized form (JS raw input, +/// legacy CubeStore result sets, synthetic transform columns). `Arrow` wraps +/// the Arrow arrays a CubeStore result arrived in: values are read straight +/// from the Arrow buffers at serialize time, and cloning the column — e.g. to +/// carry it into [`TransformedData::Columnar`] — is an `Arc` bump instead of a +/// per-cell copy. Both variants serialize to identical JSON. +#[derive(Debug, Clone)] +pub enum ColumnarArray { + Cells(Vec), + Arrow(ArrowColumn), +} impl ColumnarArray { #[inline] pub fn new() -> Self { - Self(Vec::new()) + Self::Cells(Vec::new()) } #[inline] pub fn with_capacity(cap: usize) -> Self { - Self(Vec::with_capacity(cap)) + Self::Cells(Vec::with_capacity(cap)) + } + + #[inline] + pub fn len(&self) -> usize { + match self { + Self::Cells(v) => v.len(), + Self::Arrow(a) => a.len(), + } + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Append a cell to a materialized column. Construction-time helper; + /// Arrow-backed columns are immutable. + #[inline] + pub fn push(&mut self, cell: DBResponsePrimitive) { + match self { + Self::Cells(v) => v.push(cell), + Self::Arrow(_) => panic!("cannot push into an Arrow-backed ColumnarArray"), + } + } + + /// Borrowed per-cell views in row order. String cells borrow their bytes + /// (no allocation) whichever variant backs the column. + pub fn iter_values(&self) -> impl Iterator> + '_ { + match self { + Self::Cells(v) => Either::Left(v.iter().map(DBResponseValueRef::from)), + Self::Arrow(a) => Either::Right(a.iter_values()), + } } + /// Materialized cells: borrowed as-is for `Cells`, converted (allocating + /// per string cell) for `Arrow`. Row-major consumers (compact/vanilla + /// plans) use this to get random access. + pub fn to_cells(&self) -> Cow<'_, [DBResponsePrimitive]> { + match self { + Self::Cells(v) => Cow::Borrowed(v.as_slice()), + Self::Arrow(a) => Cow::Owned(a.iter_values().map(|v| v.into_primitive()).collect()), + } + } +} + +impl Default for ColumnarArray { #[inline] - pub fn as_slice(&self) -> &[DBResponsePrimitive] { - &self.0 + fn default() -> Self { + Self::new() } } impl From> for ColumnarArray { #[inline] fn from(v: Vec) -> Self { - Self(v) + Self::Cells(v) } } -impl From for Vec { - #[inline] - fn from(c: ColumnarArray) -> Self { - c.0 +// Compares by value across backing variants, so an Arrow-backed column equals +// its materialized twin. +impl PartialEq for ColumnarArray { + fn eq(&self, other: &Self) -> bool { + self.len() == other.len() && self.iter_values().eq(other.iter_values()) } } -impl std::ops::Deref for ColumnarArray { - type Target = Vec; - #[inline] - fn deref(&self) -> &Self::Target { - &self.0 +impl Serialize for ColumnarArray { + fn serialize(&self, serializer: S) -> Result { + match self { + Self::Cells(v) => v.serialize(serializer), + Self::Arrow(a) => { + let mut seq = serializer.serialize_seq(Some(a.len()))?; + for value in a.iter_values() { + seq.serialize_element(&value)?; + } + seq.end() + } + } } } -impl std::ops::DerefMut for ColumnarArray { - #[inline] - fn deref_mut(&mut self) -> &mut Self::Target { - &mut self.0 +// Only the materialized form ever arrives over the wire (JS raw columnar +// JSON, cached transform fixtures); Arrow columns are created in-process. +impl<'de> Deserialize<'de> for ColumnarArray { + fn deserialize>(deserializer: D) -> Result { + Vec::::deserialize(deserializer).map(Self::Cells) } } @@ -2289,8 +2367,11 @@ mod tests { right_column.len() ))); } - for (row, left_value) in left_column.iter().enumerate() { - let right_value = &right_column[row]; + for (row, (left_value, right_value)) in left_column + .iter_values() + .zip(right_column.iter_values()) + .enumerate() + { if left_value != right_value { return Err(TestError(format!( "Columnar value at row {} for member '{}' differs: {} != {}", @@ -3538,9 +3619,10 @@ mod tests { "column {} length must equal row count", col_idx ); + let cells = column.to_cells(); for (row_idx, expected_row) in compact_dataset.iter().enumerate() { assert_eq!( - &column[row_idx], &expected_row[col_idx], + &cells[row_idx], &expected_row[col_idx], "value at column {} row {} must match compact dataset", col_idx, row_idx ); @@ -3712,4 +3794,133 @@ mod tests { ); Ok(()) } + + /// Columnar output over an Arrow-backed result: every column whose + /// transform is an identity (non-"time" members, and "time" members + /// without string chunks) must be shared as-is — an `Arc` bump, not a + /// per-cell copy — while "time" string columns still materialize through + /// `transform_value`. The serialized JSON must be identical to what the + /// materialized path produces. + #[test] + fn test_columnar_transform_shares_arrow_columns() -> Result<()> { + use arrow::array::{ArrayRef, Int64Array, StringArray, TimestampMillisecondArray}; + use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema, TimeUnit}; + use arrow::ipc::writer::StreamWriter; + use arrow::record_batch::RecordBatch; + use serde_json::json; + + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("orders__city", ArrowDataType::Utf8, true), + Field::new("orders__count", ArrowDataType::Int64, true), + Field::new( + "orders__created_at", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + true, + ), + Field::new("orders__updated_at", ArrowDataType::Utf8, true), + ])); + + // Two batches, so the shared columns are chunked. + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec![Some("Berlin"), Some("Lisbon")])) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(1i64), None])), + Arc::new(TimestampMillisecondArray::from(vec![ + Some(0i64), + Some(1_000), + ])), + Arc::new(StringArray::from(vec![Some("2024-01-02 03:04:05"), None])), + ], + ) + .unwrap(); + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(vec![None::<&str>])) as ArrayRef, + Arc::new(Int64Array::from(vec![Some(3i64)])), + Arc::new(TimestampMillisecondArray::from(vec![None::])), + Arc::new(StringArray::from(vec![Some("2024-03-04T05:06:07")])), + ], + ) + .unwrap(); + + let mut buf = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut buf, schema.as_ref()).unwrap(); + writer.write(&batch1).unwrap(); + writer.write(&batch2).unwrap(); + writer.finish().unwrap(); + } + let raw_data = QueryResult::from_arrow(&buf)?; + + let request: TransformDataRequest = serde_json::from_value(json!({ + "aliasToMemberNameMap": { + "orders__city": "Orders.city", + "orders__count": "Orders.count", + "orders__created_at": "Orders.createdAt", + "orders__updated_at": "Orders.updatedAt" + }, + "annotation": { + "Orders.city": { "type": "string" }, + "Orders.count": { "type": "number" }, + "Orders.createdAt": { "type": "time" }, + "Orders.updatedAt": { "type": "time" } + }, + "query": { + "dimensions": ["Orders.city", "Orders.createdAt", "Orders.updatedAt"], + "measures": ["Orders.count"] + }, + "queryType": "regularQuery", + "resType": "columnar" + }))?; + + let transformed = TransformedData::transform(&request, &raw_data)?; + let TransformedData::Columnar { members, columns } = &transformed else { + panic!("expected Columnar"); + }; + + assert_eq!( + members, + &[ + "Orders.city".to_string(), + "Orders.count".to_string(), + "Orders.createdAt".to_string(), + "Orders.updatedAt".to_string(), + ] + ); + + // Identity transforms keep the Arrow backing; the "time" string + // column is the only one that materializes. + assert!(matches!(columns[0], ColumnarArray::Arrow(_)), "city"); + assert!(matches!(columns[1], ColumnarArray::Arrow(_)), "count"); + assert!( + matches!(columns[2], ColumnarArray::Arrow(_)), + "createdAt (timestamp) has no string chunks" + ); + assert!( + matches!(columns[3], ColumnarArray::Cells(_)), + "updatedAt ('time' strings) must be materialized" + ); + + assert_eq!( + serde_json::to_value(&transformed)?, + json!({ + "members": [ + "Orders.city", + "Orders.count", + "Orders.createdAt", + "Orders.updatedAt" + ], + "columns": [ + ["Berlin", "Lisbon", null], + ["1", null, "3"], + ["1970-01-01T00:00:00.000", "1970-01-01T00:00:01.000", null], + ["2024-01-02T03:04:05.000", null, "2024-03-04T05:06:07.000"] + ] + }) + ); + + Ok(()) + } }