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, 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
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
52#[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 Sidecar,
75
76 PerLevel,
80}
81
82pub trait SetVirtualDataColumn {
88 fn write_to(self, col: &mut ColumnBuilder) -> Result<(), &'static str>;
92
93 fn new_builder() -> ColumnBuilder;
96
97 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#[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
256fn 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
292fn 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
341fn 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 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 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 if is_flat {
658 #[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 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 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 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 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 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 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 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 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}