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