1use 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
38pub 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
52fn 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#[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 Sidecar,
101
102 PerLevel,
106}
107
108pub trait SetVirtualDataColumn {
114 fn write_to(self, col: &mut ColumnBuilder) -> Result<(), &'static str>;
118
119 fn new_builder() -> ColumnBuilder;
122
123 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#[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
282fn 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
318fn 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
386fn 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 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 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 if is_flat {
726 #[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 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 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 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 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 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 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 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}