dre_core/
values.rs

1//! Arrow values as Jinja values, for `run_query()` rows.
2
3use arrow::array::{Array, AsArray, RecordBatch};
4use arrow::datatypes::{
5    DataType, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type,
6    UInt32Type, UInt64Type,
7};
8use arrow::util::display::{ArrayFormatter, FormatOptions};
9use minijinja::Value;
10
11/// Every row of `batch`, as a list of values per row.
12pub fn batch_rows(batch: &RecordBatch) -> Vec<Vec<Value>> {
13    let opts = FormatOptions::default();
14    let cols: Vec<Box<dyn Fn(usize) -> Value + '_>> = batch
15        .columns()
16        .iter()
17        .map(|c| column(c.as_ref(), &opts))
18        .collect();
19    (0..batch.num_rows())
20        .map(|i| cols.iter().map(|f| f(i)).collect())
21        .collect()
22}
23
24fn column<'a>(a: &'a dyn Array, opts: &'a FormatOptions<'a>) -> Box<dyn Fn(usize) -> Value + 'a> {
25    macro_rules! prim {
26        ($t:ty) => {{
27            let arr = a.as_primitive::<$t>();
28            Box::new(move |i| {
29                if arr.is_null(i) {
30                    Value::from(())
31                } else {
32                    Value::from(arr.value(i))
33                }
34            })
35        }};
36    }
37    match a.data_type() {
38        DataType::Int8 => prim!(Int8Type),
39        DataType::Int16 => prim!(Int16Type),
40        DataType::Int32 => prim!(Int32Type),
41        DataType::Int64 => prim!(Int64Type),
42        DataType::UInt8 => prim!(UInt8Type),
43        DataType::UInt16 => prim!(UInt16Type),
44        DataType::UInt32 => prim!(UInt32Type),
45        DataType::UInt64 => prim!(UInt64Type),
46        DataType::Float32 => prim!(Float32Type),
47        DataType::Float64 => prim!(Float64Type),
48        DataType::Boolean => {
49            let arr = a.as_boolean();
50            Box::new(move |i| {
51                if arr.is_null(i) {
52                    Value::from(())
53                } else {
54                    Value::from(arr.value(i))
55                }
56            })
57        }
58        _ => match ArrayFormatter::try_new(a, opts) {
59            Ok(f) => Box::new(move |i| {
60                if a.is_null(i) {
61                    Value::from(())
62                } else {
63                    Value::from(f.value(i).to_string())
64                }
65            }),
66            Err(_) => Box::new(|_| Value::from(())),
67        },
68    }
69}
dre_core/values.rs

dre_core/
values.rs
