Skip to main content

perspective_client/virtual_server/
data.rs

1// ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
2// ┃ ██████ ██████ ██████       █      █      █      █      █ █▄  ▀███ █       ┃
3// ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█  ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄  ▀█ █ ▀▀▀▀▀ ┃
4// ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄   █ ▄▄▄▄▄ ┃
5// ┃ █      ██████ █  ▀█▄       █ ██████      █      ███▌▐███ ███████▄ █       ┃
6// ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫
7// ┃ Copyright (c) 2017, the Perspective Authors.                              ┃
8// ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃
9// ┃ This file is part of the Perspective library, distributed under the terms ┃
10// ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃
11// ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
12
13use std::error::Error;
14use std::sync::Arc;
15
16use arrow_array::builder::{
17    BooleanBuilder, Float64Builder, Int32Builder, StringDictionaryBuilder,
18    TimestampMillisecondBuilder,
19};
20use arrow_array::cast::AsArray;
21use arrow_array::types::Int32Type;
22use arrow_array::{
23    Array, ArrayAccessor, ArrayRef, BooleanArray, Date32Array, Date64Array, Decimal128Array,
24    Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, LargeStringArray,
25    RecordBatch, StringArray, Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray,
26    Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray,
27    TimestampNanosecondArray, TimestampSecondArray, UInt8Array, UInt16Array, UInt32Array,
28    UInt64Array,
29};
30use arrow_ipc::reader::{FileReader, StreamReader};
31use arrow_ipc::writer::StreamWriter;
32use arrow_schema::{DataType, Field, Schema, TimeUnit};
33use indexmap::IndexMap;
34use serde::Serialize;
35
36use crate::config::{GroupRollupMode, Scalar, ViewConfig};
37
38/// An Arrow column builder, used during the population phase of
39/// [`VirtualDataSlice`].
40pub enum ColumnBuilder {
41    Boolean(BooleanBuilder),
42    String(StringDictionaryBuilder<Int32Type>),
43    Float(Float64Builder),
44    Integer(Int32Builder),
45    Datetime(TimestampMillisecondBuilder),
46}
47
48fn dict_data_type() -> DataType {
49    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
50}
51
52/// A single cell value in a row-oriented data representation.
53///
54/// Used when converting [`VirtualDataSlice`] to row format for JSON
55/// serialization.
56#[derive(Debug, Serialize)]
57#[serde(untagged)]
58pub enum VirtualDataCell {
59    Boolean(Option<bool>),
60    String(Option<String>),
61    Float(Option<f64>),
62    Integer(Option<i32>),
63    Datetime(Option<i64>),
64    RowPath(Vec<Scalar>),
65}
66
67#[derive(Copy, Clone, Debug, PartialEq, Eq)]
68pub enum RowPathStyle {
69    /// Legacy: emit a single `__ROW_PATH__` sidecar (per-row nested
70    /// array in `render_to_rows`, array-of-arrays in
71    /// `render_to_columns_json`). `__ROW_PATH_N__` per-level columns
72    /// are filtered out. Matches the native engine's `to_json` /
73    /// `to_columns` shape.
74    Sidecar,
75
76    /// Native: emit per-level `__ROW_PATH_0__`, `__ROW_PATH_1__`, …
77    /// columns directly. No `__ROW_PATH__` sidecar. Matches the native
78    /// engine's Arrow IPC, CSV, and NDJSON shapes.
79    PerLevel,
80}
81
82/// Trait for types that can be written to a [`ColumnBuilder`] which
83/// enforces sequential construction.
84///
85/// This trait enables type-safe insertion of values into virtual data columns,
86/// ensuring that values are written to columns of the correct type.
87pub trait SetVirtualDataColumn {
88    /// Writes this value (sequentially) to the given column builder.
89    ///
90    /// Returns an error if the column type does not match the value type.
91    fn write_to(self, col: &mut ColumnBuilder) -> Result<(), &'static str>;
92
93    /// Creates a new empty column builder of the appropriate type for this
94    /// value.
95    fn new_builder() -> ColumnBuilder;
96
97    /// Converts this value to a [`Scalar`] representation.
98    fn to_scalar(self) -> Scalar;
99}
100
101impl SetVirtualDataColumn for Option<String> {
102    fn write_to(self, col: &mut ColumnBuilder) -> Result<(), &'static str> {
103        if let ColumnBuilder::String(builder) = col {
104            match self {
105                Some(s) => builder.append_value(&s),
106                None => builder.append_null(),
107            }
108            Ok(())
109        } else {
110            Err("Bad type")
111        }
112    }
113
114    fn new_builder() -> ColumnBuilder {
115        ColumnBuilder::String(StringDictionaryBuilder::new())
116    }
117
118    fn to_scalar(self) -> Scalar {
119        if let Some(x) = self {
120            Scalar::String(x)
121        } else {
122            Scalar::Null
123        }
124    }
125}
126
127impl SetVirtualDataColumn for Option<f64> {
128    fn write_to(self, col: &mut ColumnBuilder) -> Result<(), &'static str> {
129        if let ColumnBuilder::Float(builder) = col {
130            match self {
131                Some(v) => builder.append_value(v),
132                None => builder.append_null(),
133            }
134            Ok(())
135        } else {
136            Err("Bad type")
137        }
138    }
139
140    fn new_builder() -> ColumnBuilder {
141        ColumnBuilder::Float(Float64Builder::new())
142    }
143
144    fn to_scalar(self) -> Scalar {
145        if let Some(x) = self {
146            Scalar::Float(x)
147        } else {
148            Scalar::Null
149        }
150    }
151}
152
153impl SetVirtualDataColumn for Option<i32> {
154    fn write_to(self, col: &mut ColumnBuilder) -> Result<(), &'static str> {
155        if let ColumnBuilder::Integer(builder) = col {
156            match self {
157                Some(v) => builder.append_value(v),
158                None => builder.append_null(),
159            }
160            Ok(())
161        } else {
162            Err("Bad type")
163        }
164    }
165
166    fn new_builder() -> ColumnBuilder {
167        ColumnBuilder::Integer(Int32Builder::new())
168    }
169
170    fn to_scalar(self) -> Scalar {
171        if let Some(x) = self {
172            Scalar::Float(x as f64)
173        } else {
174            Scalar::Null
175        }
176    }
177}
178
179impl SetVirtualDataColumn for Option<i64> {
180    fn write_to(self, col: &mut ColumnBuilder) -> Result<(), &'static str> {
181        if let ColumnBuilder::Datetime(builder) = col {
182            match self {
183                Some(v) => builder.append_value(v),
184                None => builder.append_null(),
185            }
186            Ok(())
187        } else {
188            Err("Bad type")
189        }
190    }
191
192    fn new_builder() -> ColumnBuilder {
193        ColumnBuilder::Datetime(TimestampMillisecondBuilder::new())
194    }
195
196    fn to_scalar(self) -> Scalar {
197        if let Some(x) = self {
198            Scalar::Float(x as f64)
199        } else {
200            Scalar::Null
201        }
202    }
203}
204
205impl SetVirtualDataColumn for Option<bool> {
206    fn write_to(self, col: &mut ColumnBuilder) -> Result<(), &'static str> {
207        if let ColumnBuilder::Boolean(builder) = col {
208            match self {
209                Some(v) => builder.append_value(v),
210                None => builder.append_null(),
211            }
212            Ok(())
213        } else {
214            Err("Bad type")
215        }
216    }
217
218    fn new_builder() -> ColumnBuilder {
219        ColumnBuilder::Boolean(BooleanBuilder::new())
220    }
221
222    fn to_scalar(self) -> Scalar {
223        if let Some(x) = self {
224            Scalar::Bool(x)
225        } else {
226            Scalar::Null
227        }
228    }
229}
230
231/// A columnar data slice returned from a virtual server view query.
232///
233/// This struct represents a rectangular slice of data from a view, stored
234/// internally as Arrow builders during population and frozen into a
235/// `RecordBatch` on first consumption.
236#[derive(Debug)]
237pub struct VirtualDataSlice {
238    config: ViewConfig,
239    builders: IndexMap<String, ColumnBuilder>,
240    row_path: Option<Vec<Vec<Scalar>>>,
241    frozen: Option<RecordBatch>,
242}
243
244impl std::fmt::Debug for ColumnBuilder {
245    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
246        match self {
247            ColumnBuilder::Boolean(_) => write!(f, "ColumnBuilder::Boolean(..)"),
248            ColumnBuilder::String(_) => write!(f, "ColumnBuilder::String(..)"),
249            ColumnBuilder::Float(_) => write!(f, "ColumnBuilder::Float(..)"),
250            ColumnBuilder::Integer(_) => write!(f, "ColumnBuilder::Integer(..)"),
251            ColumnBuilder::Datetime(_) => write!(f, "ColumnBuilder::Datetime(..)"),
252        }
253    }
254}
255
256/// Extracts grouping ID values from an Arrow array as `i64`.
257fn cast_to_int64(array: &ArrayRef) -> Result<Vec<i64>, Box<dyn Error>> {
258    let num_rows = array.len();
259    let mut result = Vec::with_capacity(num_rows);
260    match array.data_type() {
261        DataType::Int32 => {
262            let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
263            for i in 0..num_rows {
264                result.push(if arr.is_null(i) {
265                    0
266                } else {
267                    arr.value(i) as i64
268                });
269            }
270        },
271        DataType::Int64 => {
272            let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
273            for i in 0..num_rows {
274                result.push(if arr.is_null(i) { 0 } else { arr.value(i) });
275            }
276        },
277        DataType::Float64 => {
278            let arr = array.as_any().downcast_ref::<Float64Array>().unwrap();
279            for i in 0..num_rows {
280                result.push(if arr.is_null(i) {
281                    0
282                } else {
283                    arr.value(i) as i64
284                });
285            }
286        },
287        dt => return Err(format!("Cannot cast {} to Int64", dt).into()),
288    }
289    Ok(result)
290}
291
292/// Extracts a single cell from an Arrow array as a [`Scalar`].
293fn extract_scalar(array: &ArrayRef, row_idx: usize) -> Scalar {
294    if array.is_null(row_idx) {
295        return Scalar::Null;
296    }
297    match array.data_type() {
298        DataType::Utf8 => {
299            let arr = array.as_any().downcast_ref::<StringArray>().unwrap();
300            Scalar::String(arr.value(row_idx).to_string())
301        },
302        DataType::Dictionary(..) => {
303            let dict = array.as_dictionary::<Int32Type>();
304            let values = dict.downcast_dict::<StringArray>().unwrap();
305            Scalar::String(values.value(row_idx).to_string())
306        },
307        DataType::Float64 => {
308            let arr = array.as_any().downcast_ref::<Float64Array>().unwrap();
309            Scalar::Float(arr.value(row_idx))
310        },
311        DataType::Int32 => {
312            let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
313            Scalar::Float(arr.value(row_idx) as f64)
314        },
315        DataType::Int64 => {
316            let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
317            Scalar::Float(arr.value(row_idx) as f64)
318        },
319        DataType::Boolean => {
320            let arr = array.as_any().downcast_ref::<BooleanArray>().unwrap();
321            Scalar::Bool(arr.value(row_idx))
322        },
323        DataType::Timestamp(TimeUnit::Millisecond, _) => {
324            let arr = array
325                .as_any()
326                .downcast_ref::<TimestampMillisecondArray>()
327                .unwrap();
328            Scalar::Float(arr.value(row_idx) as f64)
329        },
330        DataType::Date32 => {
331            let arr = array.as_any().downcast_ref::<Date32Array>().unwrap();
332            Scalar::Float(arr.value(row_idx) as f64 * 86_400_000.0)
333        },
334        _ => {
335            let scalar_arr = array.slice(row_idx, 1);
336            Scalar::String(format!("{:?}", scalar_arr))
337        },
338    }
339}
340
341/// Coerces an Arrow column to Perspective-compatible types, optionally
342/// renaming.
343/// Manually converts a timestamp array of any unit to milliseconds.
344fn timestamp_to_millis(array: &ArrayRef, unit: &TimeUnit) -> ArrayRef {
345    let millis: TimestampMillisecondArray = match unit {
346        TimeUnit::Second => {
347            let arr = array
348                .as_any()
349                .downcast_ref::<TimestampSecondArray>()
350                .unwrap();
351            arr.iter().map(|v| v.map(|v| v * 1_000)).collect()
352        },
353        TimeUnit::Microsecond => {
354            let arr = array
355                .as_any()
356                .downcast_ref::<TimestampMicrosecondArray>()
357                .unwrap();
358            arr.iter().map(|v| v.map(|v| v / 1_000)).collect()
359        },
360        TimeUnit::Nanosecond => {
361            let arr = array
362                .as_any()
363                .downcast_ref::<TimestampNanosecondArray>()
364                .unwrap();
365            arr.iter().map(|v| v.map(|v| v / 1_000_000)).collect()
366        },
367        TimeUnit::Millisecond => {
368            return array.clone();
369        },
370    };
371    Arc::new(millis) as ArrayRef
372}
373
374fn coerce_column(
375    name: &str,
376    field: &Field,
377    array: &ArrayRef,
378) -> Result<(Field, ArrayRef), Box<dyn Error>> {
379    match field.data_type() {
380        DataType::Boolean | DataType::Float64 | DataType::Int32 | DataType::Date32 => Ok((
381            Field::new(name, field.data_type().clone(), true),
382            array.clone(),
383        )),
384        DataType::Dictionary(..) => Ok((Field::new(name, dict_data_type(), true), array.clone())),
385        DataType::Utf8 => {
386            let arr = array.as_any().downcast_ref::<StringArray>().unwrap();
387            let mut builder = StringDictionaryBuilder::<Int32Type>::new();
388            for i in 0..arr.len() {
389                if arr.is_null(i) {
390                    builder.append_null();
391                } else {
392                    builder.append_value(arr.value(i));
393                }
394            }
395            Ok((
396                Field::new(name, dict_data_type(), true),
397                Arc::new(builder.finish()) as ArrayRef,
398            ))
399        },
400        DataType::Timestamp(TimeUnit::Millisecond, _) => Ok((
401            Field::new(name, DataType::Timestamp(TimeUnit::Millisecond, None), true),
402            array.clone(),
403        )),
404        DataType::Int8 => {
405            let arr = array.as_any().downcast_ref::<Int8Array>().unwrap();
406            let result: Int32Array = arr.iter().map(|v| v.map(|v| v as i32)).collect();
407            Ok((
408                Field::new(name, DataType::Int32, true),
409                Arc::new(result) as ArrayRef,
410            ))
411        },
412        DataType::Int16 => {
413            let arr = array.as_any().downcast_ref::<Int16Array>().unwrap();
414            let result: Int32Array = arr.iter().map(|v| v.map(|v| v as i32)).collect();
415            Ok((
416                Field::new(name, DataType::Int32, true),
417                Arc::new(result) as ArrayRef,
418            ))
419        },
420        DataType::UInt8 => {
421            let arr = array.as_any().downcast_ref::<UInt8Array>().unwrap();
422            let result: Int32Array = arr.iter().map(|v| v.map(|v| v as i32)).collect();
423            Ok((
424                Field::new(name, DataType::Int32, true),
425                Arc::new(result) as ArrayRef,
426            ))
427        },
428        DataType::UInt16 => {
429            let arr = array.as_any().downcast_ref::<UInt16Array>().unwrap();
430            let result: Int32Array = arr.iter().map(|v| v.map(|v| v as i32)).collect();
431            Ok((
432                Field::new(name, DataType::Int32, true),
433                Arc::new(result) as ArrayRef,
434            ))
435        },
436        DataType::UInt32 => {
437            let arr = array.as_any().downcast_ref::<UInt32Array>().unwrap();
438            let result: Int64Array = arr.iter().map(|v| v.map(|v| v as i64)).collect();
439            let result: Float64Array = result.iter().map(|v| v.map(|v| v as f64)).collect();
440            Ok((
441                Field::new(name, DataType::Float64, true),
442                Arc::new(result) as ArrayRef,
443            ))
444        },
445        DataType::Int64 => {
446            let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
447            let result: Float64Array = arr.iter().map(|v| v.map(|v| v as f64)).collect();
448            Ok((
449                Field::new(name, DataType::Float64, true),
450                Arc::new(result) as ArrayRef,
451            ))
452        },
453        DataType::UInt64 => {
454            let arr = array.as_any().downcast_ref::<UInt64Array>().unwrap();
455            let result: Float64Array = arr.iter().map(|v| v.map(|v| v as f64)).collect();
456            Ok((
457                Field::new(name, DataType::Float64, true),
458                Arc::new(result) as ArrayRef,
459            ))
460        },
461        DataType::Float32 => {
462            let arr = array.as_any().downcast_ref::<Float32Array>().unwrap();
463            let result: Float64Array = arr.iter().map(|v| v.map(|v| v as f64)).collect();
464            Ok((
465                Field::new(name, DataType::Float64, true),
466                Arc::new(result) as ArrayRef,
467            ))
468        },
469        DataType::Decimal128(_, scale) => {
470            let scale = *scale;
471            let arr = array.as_any().downcast_ref::<Decimal128Array>().unwrap();
472            let divisor = 10_f64.powi(scale as i32);
473            let result: Float64Array = arr.iter().map(|v| v.map(|v| v as f64 / divisor)).collect();
474            Ok((
475                Field::new(name, DataType::Float64, true),
476                Arc::new(result) as ArrayRef,
477            ))
478        },
479        DataType::Date64 => {
480            let arr = array.as_any().downcast_ref::<Date64Array>().unwrap();
481            let result: Date32Array = arr
482                .iter()
483                .map(|v| v.map(|v| (v / 86_400_000) as i32))
484                .collect();
485            Ok((
486                Field::new(name, DataType::Date32, true),
487                Arc::new(result) as ArrayRef,
488            ))
489        },
490        DataType::Timestamp(unit, _) => {
491            let casted = timestamp_to_millis(array, unit);
492            Ok((
493                Field::new(name, DataType::Timestamp(TimeUnit::Millisecond, None), true),
494                casted,
495            ))
496        },
497        DataType::Time32(TimeUnit::Second) => {
498            let arr = array.as_any().downcast_ref::<Time32SecondArray>().unwrap();
499            let result: TimestampMillisecondArray =
500                arr.iter().map(|v| v.map(|v| v as i64 * 1_000)).collect();
501            Ok((
502                Field::new(name, DataType::Timestamp(TimeUnit::Millisecond, None), true),
503                Arc::new(result) as ArrayRef,
504            ))
505        },
506        DataType::Time32(TimeUnit::Millisecond) => {
507            let arr = array
508                .as_any()
509                .downcast_ref::<Time32MillisecondArray>()
510                .unwrap();
511            let result: TimestampMillisecondArray =
512                arr.iter().map(|v| v.map(|v| v as i64)).collect();
513            Ok((
514                Field::new(name, DataType::Timestamp(TimeUnit::Millisecond, None), true),
515                Arc::new(result) as ArrayRef,
516            ))
517        },
518        DataType::Time64(TimeUnit::Microsecond) => {
519            let arr = array
520                .as_any()
521                .downcast_ref::<Time64MicrosecondArray>()
522                .unwrap();
523            let result: TimestampMillisecondArray =
524                arr.iter().map(|v| v.map(|v| v / 1_000)).collect();
525            Ok((
526                Field::new(name, DataType::Timestamp(TimeUnit::Millisecond, None), true),
527                Arc::new(result) as ArrayRef,
528            ))
529        },
530        DataType::Time64(TimeUnit::Nanosecond) => {
531            let arr = array
532                .as_any()
533                .downcast_ref::<Time64NanosecondArray>()
534                .unwrap();
535            let result: TimestampMillisecondArray =
536                arr.iter().map(|v| v.map(|v| v / 1_000_000)).collect();
537            Ok((
538                Field::new(name, DataType::Timestamp(TimeUnit::Millisecond, None), true),
539                Arc::new(result) as ArrayRef,
540            ))
541        },
542        DataType::LargeUtf8 => {
543            let arr = array.as_any().downcast_ref::<LargeStringArray>().unwrap();
544            let mut builder = StringDictionaryBuilder::<Int32Type>::new();
545            for i in 0..arr.len() {
546                if arr.is_null(i) {
547                    builder.append_null();
548                } else {
549                    builder.append_value(arr.value(i));
550                }
551            }
552            Ok((
553                Field::new(name, dict_data_type(), true),
554                Arc::new(builder.finish()) as ArrayRef,
555            ))
556        },
557        dt => {
558            tracing::warn!(
559                "Coercing unknown Arrow type {} to Dictionary for column '{}'",
560                dt,
561                name
562            );
563            let num_rows = array.len();
564            let mut builder = StringDictionaryBuilder::<Int32Type>::new();
565            for i in 0..num_rows {
566                if array.is_null(i) {
567                    builder.append_null();
568                } else {
569                    let scalar_arr = array.slice(i, 1);
570                    builder.append_value(format!("{:?}", scalar_arr));
571                }
572            }
573            Ok((
574                Field::new(name, dict_data_type(), true),
575                Arc::new(builder.finish()) as ArrayRef,
576            ))
577        },
578    }
579}
580
581impl VirtualDataSlice {
582    pub fn new(config: ViewConfig) -> Self {
583        VirtualDataSlice {
584            config,
585            builders: IndexMap::default(),
586            row_path: None,
587            frozen: None,
588        }
589    }
590
591    /// Loads data from Arrow IPC file format bytes, with automatic
592    /// post-processing based on the view configuration.
593    ///
594    /// When `group_by` is active, extracts `__GROUPING_ID__` and
595    /// `__ROW_PATH_N__` columns to build `self.row_path`, then removes
596    /// `__GROUPING_ID__` from the output `RecordBatch`. The
597    /// `__ROW_PATH_N__` columns are *kept* in the frozen batch so
598    /// downstream Arrow IPC consumers (`with_typed_arrays`, used by
599    /// viewer-charts to drive its categorical/numeric axis resolvers
600    /// and tree-hierarchy walkers) see them inline — matching the
601    /// native `perspective-server`'s `to_arrow` output when
602    /// `emit_legacy_row_path_names: false`.
603    ///
604    /// Also coerces non-standard Arrow types (e.g. `Decimal128`, `Int64`)
605    /// to Perspective-compatible types. Data column names are passed
606    /// through verbatim — pivoted views already name columns with
607    /// Perspective's column-path separator.
608    pub fn from_arrow_ipc(&mut self, ipc: &[u8]) -> Result<(), Box<dyn Error>> {
609        let cursor = std::io::Cursor::new(ipc);
610        let batches: Vec<RecordBatch> = if &ipc[0..6] == "ARROW1".as_bytes() {
611            FileReader::try_new(cursor, None)?.collect::<Result<Vec<_>, _>>()?
612        } else {
613            StreamReader::try_new(cursor, None)?.collect::<Result<Vec<_>, _>>()?
614        };
615
616        let batch = match batches.len() {
617            0 => return Err("Arrow IPC stream contained no record batches".into()),
618            1 => batches.into_iter().next().unwrap(),
619            _ => arrow_select::concat::concat_batches(&batches[0].schema(), &batches)?,
620        };
621
622        let has_group_by = !self.config.group_by.is_empty();
623        let has_split_by = !self.config.split_by.is_empty();
624        let is_total = self.config.group_rollup_mode == GroupRollupMode::Total;
625
626        if !has_group_by && !has_split_by && !is_total {
627            self.frozen = Some(batch);
628            return Ok(());
629        }
630
631        let num_rows = batch.num_rows();
632        let schema = batch.schema();
633
634        // Phase A: Extract row_path from __GROUPING_ID__ and __ROW_PATH_N__
635        if has_group_by {
636            let group_by_len = self.config.group_by.len();
637            let is_flat = self.config.group_rollup_mode == GroupRollupMode::Flat;
638            let grouping_ids = if is_flat {
639                None
640            } else {
641                let grouping_id_idx = schema
642                    .index_of("__GROUPING_ID__")
643                    .map_err(|_| "Missing __GROUPING_ID__ column")?;
644                Some(cast_to_int64(batch.column(grouping_id_idx))?)
645            };
646
647            let mut row_paths: Vec<Vec<Scalar>> = (0..num_rows).map(|_| Vec::new()).collect();
648            for gidx in 0..group_by_len {
649                let col_name = format!("__ROW_PATH_{}__", gidx);
650                let col_idx = schema
651                    .index_of(&col_name)
652                    .map_err(|_| format!("Missing {} column", col_name))?;
653
654                let col = batch.column(col_idx);
655
656                // In flat mode, all rows are leaf rows
657                if is_flat {
658                    // TODO I may be dumb but I'm not exactly sure what Clippy
659                    // wants here. This could be an `enumerate` but how is this
660                    // better?
661                    #[allow(clippy::needless_range_loop)]
662                    for row_idx in 0..num_rows {
663                        row_paths[row_idx].push(extract_scalar(col, row_idx));
664                    }
665                } else {
666                    let gids = grouping_ids.as_ref().unwrap();
667                    let max_grouping_id = 2_i64.pow(group_by_len as u32 - gidx as u32) - 1;
668                    for row_idx in 0..num_rows {
669                        if gids[row_idx] < max_grouping_id {
670                            row_paths[row_idx].push(extract_scalar(col, row_idx));
671                        }
672                    }
673                }
674            }
675
676            self.row_path = Some(row_paths);
677        }
678
679        // Phase B: Rebuild RecordBatch without metadata columns, with
680        // column renames and type coercion.
681        let mut new_fields = Vec::new();
682        let mut new_arrays: Vec<ArrayRef> = Vec::new();
683        for (col_idx, field) in schema.fields().iter().enumerate() {
684            let name = field.name();
685            // `__GROUPING_ID__` is an internal SQL-rollup discriminator
686            // (used in Phase A above to decide which row-path levels
687            // belong to each row). No JS consumer reads it, so it's
688            // dropped from the frozen batch.
689            //
690            // `__ROW_PATH_N__` columns are kept. Phase A copied their
691            // values into `self.row_path` for the JSON sidecar paths
692            // (`render_to_columns_json`, `render_to_rows`), but
693            // viewer-charts' `with_typed_arrays` callback needs the
694            // per-level columns inline in the Arrow stream — its
695            // categorical-axis resolver, numeric-position lookup, and
696            // tree hierarchy walker all do `columns.get(\`__ROW_PATH_${n}__\`)`.
697            // Keeping the columns here lets `render_to_arrow_ipc`
698            // serialize them naturally, matching native
699            // `perspective-server`'s `to_arrow` output.
700            if name == "__GROUPING_ID__" {
701                continue;
702            }
703
704            let (coerced_field, coerced_array) = coerce_column(name, field, batch.column(col_idx))?;
705            new_fields.push(coerced_field);
706            new_arrays.push(coerced_array);
707        }
708
709        let new_schema = Arc::new(Schema::new(new_fields));
710        self.frozen = if new_arrays.is_empty() {
711            Some(RecordBatch::new_empty(new_schema))
712        } else {
713            Some(RecordBatch::try_new(new_schema, new_arrays)?)
714        };
715        Ok(())
716    }
717
718    /// Freezes the builders into a `RecordBatch`. Idempotent — subsequent
719    /// calls return the cached batch.
720    pub(crate) fn freeze(&mut self) -> &RecordBatch {
721        if self.frozen.is_none() {
722            let mut fields = Vec::new();
723            let mut arrays: Vec<ArrayRef> = Vec::new();
724
725            for (name, builder) in &mut self.builders {
726                let (field, array): (Field, ArrayRef) = match builder {
727                    ColumnBuilder::Boolean(b) => (
728                        Field::new(name, DataType::Boolean, true),
729                        Arc::new(b.finish()),
730                    ),
731                    ColumnBuilder::String(b) => (
732                        Field::new(name, dict_data_type(), true),
733                        Arc::new(b.finish()),
734                    ),
735                    ColumnBuilder::Float(b) => (
736                        Field::new(name, DataType::Float64, true),
737                        Arc::new(b.finish()),
738                    ),
739                    ColumnBuilder::Integer(b) => (
740                        Field::new(name, DataType::Int32, true),
741                        Arc::new(b.finish()),
742                    ),
743                    ColumnBuilder::Datetime(b) => (
744                        Field::new(name, DataType::Timestamp(TimeUnit::Millisecond, None), true),
745                        Arc::new(b.finish()),
746                    ),
747                };
748                fields.push(field);
749                arrays.push(array);
750            }
751
752            let schema = Arc::new(Schema::new(fields));
753            self.frozen = Some(
754                RecordBatch::try_new(schema, arrays)
755                    .expect("RecordBatch construction should not fail for well-formed builders"),
756            );
757        }
758        self.frozen.as_ref().unwrap()
759    }
760
761    /// Serializes the data to Arrow IPC streaming format.
762    pub(crate) fn render_to_arrow_ipc(&mut self) -> Result<Vec<u8>, Box<dyn Error>> {
763        let batch = self.freeze().clone();
764        let schema = batch.schema();
765        let mut buf = Vec::new();
766        {
767            let mut writer = StreamWriter::try_new(&mut buf, &schema)?;
768            writer.write(&batch)?;
769            writer.finish()?;
770        }
771        Ok(buf)
772    }
773
774    /// Converts the columnar data to a row-oriented representation for JSON
775    /// serialization.
776    ///
777    /// `style` selects between the legacy `__ROW_PATH__` sidecar
778    /// (`Sidecar`, used by `to_json`) and the native per-level
779    /// `__ROW_PATH_N__` columns (`PerLevel`, used by `to_csv` /
780    /// `to_ndjson`). See [`RowPathStyle`] for the deprecation plan.
781    pub(crate) fn render_to_rows(
782        &mut self,
783        style: RowPathStyle,
784    ) -> Vec<IndexMap<String, VirtualDataCell>> {
785        let batch = self.freeze().clone();
786        let num_rows = batch.num_rows();
787        let schema = batch.schema();
788
789        (0..num_rows)
790            .map(|row_idx| {
791                let mut row = IndexMap::new();
792                if style == RowPathStyle::Sidecar
793                    && let Some(ref rp) = self.row_path
794                    && row_idx < rp.len()
795                {
796                    row.insert(
797                        "__ROW_PATH__".to_string(),
798                        VirtualDataCell::RowPath(rp[row_idx].clone()),
799                    );
800                }
801
802                for (col_idx, field) in schema.fields().iter().enumerate() {
803                    if style == RowPathStyle::Sidecar && field.name().starts_with("__ROW_PATH_") {
804                        continue;
805                    }
806
807                    let col = batch.column(col_idx);
808                    let cell = if col.is_null(row_idx) {
809                        match field.data_type() {
810                            DataType::Boolean => VirtualDataCell::Boolean(None),
811                            DataType::Utf8 | DataType::Dictionary(..) => {
812                                VirtualDataCell::String(None)
813                            },
814                            DataType::Float64 => VirtualDataCell::Float(None),
815                            DataType::Int32 => VirtualDataCell::Integer(None),
816                            DataType::Timestamp(TimeUnit::Millisecond, _) => {
817                                VirtualDataCell::Datetime(None)
818                            },
819                            _ => continue,
820                        }
821                    } else {
822                        match field.data_type() {
823                            DataType::Boolean => {
824                                let arr = col.as_any().downcast_ref::<BooleanArray>().unwrap();
825                                VirtualDataCell::Boolean(Some(arr.value(row_idx)))
826                            },
827                            DataType::Utf8 => {
828                                let arr = col.as_any().downcast_ref::<StringArray>().unwrap();
829                                VirtualDataCell::String(Some(arr.value(row_idx).to_string()))
830                            },
831                            DataType::Dictionary(..) => {
832                                let dict = col.as_dictionary::<Int32Type>();
833                                let values = dict.downcast_dict::<StringArray>().unwrap();
834                                VirtualDataCell::String(Some(values.value(row_idx).to_string()))
835                            },
836                            DataType::Float64 => {
837                                let arr = col.as_any().downcast_ref::<Float64Array>().unwrap();
838                                VirtualDataCell::Float(Some(arr.value(row_idx)))
839                            },
840                            DataType::Int32 => {
841                                let arr = col.as_any().downcast_ref::<Int32Array>().unwrap();
842                                VirtualDataCell::Integer(Some(arr.value(row_idx)))
843                            },
844                            DataType::Int64 => {
845                                // TODO ????
846                                let arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
847                                VirtualDataCell::Float(Some(arr.value(row_idx) as f64))
848                            },
849                            DataType::Time64(TimeUnit::Microsecond) => {
850                                let arr = col
851                                    .as_any()
852                                    .downcast_ref::<Time64MicrosecondArray>()
853                                    .unwrap();
854                                VirtualDataCell::Float(Some(arr.value(row_idx) as f64))
855                            },
856                            DataType::Timestamp(TimeUnit::Microsecond, _) => {
857                                let arr = col
858                                    .as_any()
859                                    .downcast_ref::<Time64MicrosecondArray>()
860                                    .unwrap();
861                                VirtualDataCell::Datetime(Some(arr.value(row_idx) * 1000))
862                            },
863                            DataType::Timestamp(TimeUnit::Millisecond, _) => {
864                                let arr = col
865                                    .as_any()
866                                    .downcast_ref::<TimestampMillisecondArray>()
867                                    .unwrap();
868                                VirtualDataCell::Datetime(Some(arr.value(row_idx)))
869                            },
870                            DataType::Date32 => {
871                                let arr = col.as_any().downcast_ref::<Date32Array>().unwrap();
872                                VirtualDataCell::Datetime(Some(
873                                    arr.value(row_idx) as i64 * 86_400_000,
874                                ))
875                            },
876                            x => {
877                                tracing::error!("Unknown Arrow IPC type {}", x);
878                                continue;
879                            },
880                        }
881                    };
882                    row.insert(field.name().clone(), cell);
883                }
884
885                row
886            })
887            .collect()
888    }
889
890    /// Serializes the data to a column-oriented JSON string.
891    ///
892    /// `style` selects between the legacy `__ROW_PATH__` sidecar
893    /// (`Sidecar`, used by `to_columns`) and the native per-level
894    /// `__ROW_PATH_N__` columns (`PerLevel`, currently unused — reserved
895    /// for the future deprecation of `__ROW_PATH__`). See
896    /// [`RowPathStyle`] for context.
897    pub fn render_to_columns_json(
898        &mut self,
899        style: RowPathStyle,
900    ) -> Result<String, Box<dyn Error>> {
901        let batch = self.freeze().clone();
902        let schema = batch.schema();
903        let mut map = serde_json::Map::new();
904
905        if style == RowPathStyle::Sidecar
906            && let Some(ref rp) = self.row_path
907        {
908            map.insert("__ROW_PATH__".to_string(), serde_json::to_value(rp)?);
909        }
910
911        for (col_idx, field) in schema.fields().iter().enumerate() {
912            if style == RowPathStyle::Sidecar && field.name().starts_with("__ROW_PATH_") {
913                continue;
914            }
915
916            let col = batch.column(col_idx);
917            let num_rows = col.len();
918            let values: serde_json::Value = match field.data_type() {
919                DataType::Boolean => {
920                    let arr = col.as_any().downcast_ref::<BooleanArray>().unwrap();
921                    serde_json::to_value(
922                        (0..num_rows)
923                            .map(|i| {
924                                if arr.is_null(i) {
925                                    None
926                                } else {
927                                    Some(arr.value(i))
928                                }
929                            })
930                            .collect::<Vec<_>>(),
931                    )?
932                },
933                DataType::Utf8 => {
934                    let arr = col.as_any().downcast_ref::<StringArray>().unwrap();
935                    serde_json::to_value(
936                        (0..num_rows)
937                            .map(|i| {
938                                if arr.is_null(i) {
939                                    None
940                                } else {
941                                    Some(arr.value(i))
942                                }
943                            })
944                            .collect::<Vec<_>>(),
945                    )?
946                },
947                DataType::Dictionary(..) => {
948                    let dict = col.as_dictionary::<Int32Type>();
949                    let values = dict.downcast_dict::<StringArray>().unwrap();
950                    serde_json::to_value(
951                        (0..num_rows)
952                            .map(|i| {
953                                if col.is_null(i) {
954                                    None
955                                } else {
956                                    Some(values.value(i))
957                                }
958                            })
959                            .collect::<Vec<_>>(),
960                    )?
961                },
962                DataType::Float64 => {
963                    let arr = col.as_any().downcast_ref::<Float64Array>().unwrap();
964                    serde_json::to_value(
965                        (0..num_rows)
966                            .map(|i| {
967                                if arr.is_null(i) {
968                                    None
969                                } else {
970                                    Some(arr.value(i))
971                                }
972                            })
973                            .collect::<Vec<_>>(),
974                    )?
975                },
976                DataType::Int32 => {
977                    let arr = col.as_any().downcast_ref::<Int32Array>().unwrap();
978                    serde_json::to_value(
979                        (0..num_rows)
980                            .map(|i| {
981                                if arr.is_null(i) {
982                                    None
983                                } else {
984                                    Some(arr.value(i))
985                                }
986                            })
987                            .collect::<Vec<_>>(),
988                    )?
989                },
990                DataType::Int64 => {
991                    let arr = col.as_any().downcast_ref::<Int64Array>().unwrap();
992                    serde_json::to_value(
993                        (0..num_rows)
994                            .map(|i| {
995                                if arr.is_null(i) {
996                                    None
997                                } else {
998                                    Some(arr.value(i) as f64)
999                                }
1000                            })
1001                            .collect::<Vec<_>>(),
1002                    )?
1003                },
1004                DataType::Timestamp(TimeUnit::Millisecond, _) => {
1005                    let arr = col
1006                        .as_any()
1007                        .downcast_ref::<TimestampMillisecondArray>()
1008                        .unwrap();
1009                    serde_json::to_value(
1010                        (0..num_rows)
1011                            .map(|i| {
1012                                if arr.is_null(i) {
1013                                    None
1014                                } else {
1015                                    Some(arr.value(i))
1016                                }
1017                            })
1018                            .collect::<Vec<_>>(),
1019                    )?
1020                },
1021                DataType::Time64(TimeUnit::Microsecond) => {
1022                    let arr = col
1023                        .as_any()
1024                        .downcast_ref::<Time64MicrosecondArray>()
1025                        .unwrap();
1026                    serde_json::to_value(
1027                        (0..num_rows)
1028                            .map(|i| {
1029                                if arr.is_null(i) {
1030                                    None
1031                                } else {
1032                                    Some(arr.value(i) as f64)
1033                                }
1034                            })
1035                            .collect::<Vec<_>>(),
1036                    )?
1037                },
1038                DataType::Date32 => {
1039                    let arr = col.as_any().downcast_ref::<Date32Array>().unwrap();
1040                    serde_json::to_value(
1041                        (0..num_rows)
1042                            .map(|i| {
1043                                if arr.is_null(i) {
1044                                    None
1045                                } else {
1046                                    Some(arr.value(i) as i64 * 86_400_000)
1047                                }
1048                            })
1049                            .collect::<Vec<_>>(),
1050                    )?
1051                },
1052                x => {
1053                    tracing::error!("Unknown Arrow IPC type {}", x);
1054                    continue;
1055                },
1056            };
1057            map.insert(field.name().clone(), values);
1058        }
1059
1060        Ok(serde_json::to_string(&map)?)
1061    }
1062
1063    /// Sets a value in a column at the specified row index.
1064    ///
1065    /// If `group_by_index` is `Some`, the value is added to the `__ROW_PATH__`
1066    /// column as part of the row's group-by path. Otherwise, the value is
1067    /// inserted into the named column.
1068    ///
1069    /// Creates the column if it does not already exist.
1070    pub fn set_col<T: SetVirtualDataColumn>(
1071        &mut self,
1072        name: &str,
1073        grouping_id: Option<usize>,
1074        index: usize,
1075        value: T,
1076    ) -> Result<(), Box<dyn Error>> {
1077        if name == "__GROUPING_ID__" {
1078            return Ok(());
1079        }
1080
1081        if name.starts_with("__ROW_PATH_") {
1082            let group_by_index: u32 = name[11..name.len() - 2].parse()?;
1083            let max_grouping_id =
1084                2_i32.pow((self.config.group_by.len() as u32) - group_by_index) - 1;
1085
1086            if grouping_id.map(|x| x as i32).unwrap_or(i32::MAX) < max_grouping_id {
1087                let col = self.row_path.get_or_insert_with(Vec::new);
1088                if let Some(row) = col.get_mut(index) {
1089                    let scalar = value.to_scalar();
1090                    row.push(scalar);
1091                } else {
1092                    while col.len() < index {
1093                        col.push(vec![])
1094                    }
1095
1096                    let scalar = value.to_scalar();
1097                    col.push(vec![scalar]);
1098                }
1099            }
1100
1101            Ok(())
1102        } else {
1103            if !self.builders.contains_key(name) {
1104                self.builders.insert(name.to_owned(), T::new_builder());
1105            }
1106
1107            let col = self
1108                .builders
1109                .get_mut(name)
1110                .ok_or_else(|| format!("Column '{}' not found after insertion", name))?;
1111
1112            Ok(value.write_to(col)?)
1113        }
1114    }
1115}