dre_core/
values.rs

1//! Arrow values as Jinja values, for `run_query()` rows and message results.
2
3use arrow::array::{Array, AsArray, RecordBatch};
4use arrow::datatypes::{
5    DataType, Date32Type, Date64Type, Decimal128Type, Float32Type, Float64Type, Int8Type, Int16Type,
6    Int32Type, Int64Type, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
7    TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
8};
9use chrono::TimeZone;
10
11use crate::dates::{Calendar, Date, DateTime};
12use arrow::util::display::{ArrayFormatter, FormatOptions};
13use minijinja::Value;
14
15/// Every row of `batch`, as a list of values per row.
16pub fn batch_rows(batch: &RecordBatch) -> Vec<Vec<Value>> {
17    let opts = FormatOptions::default();
18    let cols: Vec<Box<dyn Fn(usize) -> Value + '_>> = batch
19        .columns()
20        .iter()
21        .map(|c| column(c.as_ref(), &opts))
22        .collect();
23    (0..batch.num_rows())
24        .map(|i| cols.iter().map(|f| f(i)).collect())
25        .collect()
26}
27
28/// Like [`batch_rows`], keeping more types: dates and timestamps become DRE dates and datetimes
29/// (in `cal`'s timezone), decimals become numbers.
30pub fn typed_rows(batch: &RecordBatch, cal: Calendar) -> Vec<Vec<Value>> {
31    let opts = FormatOptions::default();
32    let cols: Vec<Box<dyn Fn(usize) -> Value + '_>> = batch
33        .columns()
34        .iter()
35        .map(|c| typed_column(c.as_ref(), &opts, cal))
36        .collect();
37    (0..batch.num_rows())
38        .map(|i| cols.iter().map(|f| f(i)).collect())
39        .collect()
40}
41
42fn typed_column<'a>(
43    a: &'a dyn Array,
44    opts: &'a FormatOptions<'a>,
45    cal: Calendar,
46) -> Box<dyn Fn(usize) -> Value + 'a> {
47    let nullable = move |f: Box<dyn Fn(usize) -> Value + 'a>| -> Box<dyn Fn(usize) -> Value + 'a> {
48        Box::new(move |i| if a.is_null(i) { Value::from(()) } else { f(i) })
49    };
50    match a.data_type() {
51        DataType::Date32 | DataType::Date64 => nullable(Box::new(move |i| match date_at(a, i) {
52            Some(d) => Date::value(d, cal),
53            None => Value::from(()),
54        })),
55        DataType::Timestamp(_, tz) => {
56            let tz = tz.clone();
57            nullable(Box::new(move |i| match timestamp_at(a, i) {
58                Some(utc) => {
59                    let utc = match tz {
60                        Some(_) => utc,
61                        // A naive timestamp is wall time in the run's zone.
62                        None => match cal.tz.from_local_datetime(&utc.naive_utc()).earliest() {
63                            Some(t) => t.with_timezone(&chrono::Utc),
64                            None => utc,
65                        },
66                    };
67                    DateTime::value(utc.with_timezone(&cal.tz), cal)
68                }
69                None => Value::from(()),
70            }))
71        }
72        DataType::Decimal128(_, scale) => {
73            let arr = a.as_primitive::<Decimal128Type>();
74            let scale = *scale as i32;
75            nullable(Box::new(move |i| {
76                Value::from(arr.value(i) as f64 / 10f64.powi(scale))
77            }))
78        }
79        _ => column(a, opts),
80    }
81}
82
83fn date_at(a: &dyn Array, i: usize) -> Option<chrono::NaiveDate> {
84    match a.data_type() {
85        DataType::Date32 => a.as_primitive::<Date32Type>().value_as_date(i),
86        DataType::Date64 => a.as_primitive::<Date64Type>().value_as_date(i),
87        _ => None,
88    }
89}
90
91fn timestamp_at(a: &dyn Array, i: usize) -> Option<chrono::DateTime<chrono::Utc>> {
92    use arrow::datatypes::TimeUnit;
93    let DataType::Timestamp(unit, _) = a.data_type() else {
94        return None;
95    };
96    let naive = match unit {
97        TimeUnit::Second => a.as_primitive::<TimestampSecondType>().value_as_datetime(i),
98        TimeUnit::Millisecond => a.as_primitive::<TimestampMillisecondType>().value_as_datetime(i),
99        TimeUnit::Microsecond => a.as_primitive::<TimestampMicrosecondType>().value_as_datetime(i),
100        TimeUnit::Nanosecond => a.as_primitive::<TimestampNanosecondType>().value_as_datetime(i),
101    }?;
102    Some(naive.and_utc())
103}
104
105fn column<'a>(a: &'a dyn Array, opts: &'a FormatOptions<'a>) -> Box<dyn Fn(usize) -> Value + 'a> {
106    macro_rules! prim {
107        ($t:ty) => {{
108            let arr = a.as_primitive::<$t>();
109            Box::new(move |i| {
110                if arr.is_null(i) {
111                    Value::from(())
112                } else {
113                    Value::from(arr.value(i))
114                }
115            })
116        }};
117    }
118    match a.data_type() {
119        DataType::Int8 => prim!(Int8Type),
120        DataType::Int16 => prim!(Int16Type),
121        DataType::Int32 => prim!(Int32Type),
122        DataType::Int64 => prim!(Int64Type),
123        DataType::UInt8 => prim!(UInt8Type),
124        DataType::UInt16 => prim!(UInt16Type),
125        DataType::UInt32 => prim!(UInt32Type),
126        DataType::UInt64 => prim!(UInt64Type),
127        DataType::Float32 => prim!(Float32Type),
128        DataType::Float64 => prim!(Float64Type),
129        DataType::Boolean => {
130            let arr = a.as_boolean();
131            Box::new(move |i| {
132                if arr.is_null(i) {
133                    Value::from(())
134                } else {
135                    Value::from(arr.value(i))
136                }
137            })
138        }
139        _ => match ArrayFormatter::try_new(a, opts) {
140            Ok(f) => Box::new(move |i| {
141                if a.is_null(i) {
142                    Value::from(())
143                } else {
144                    Value::from(f.value(i).to_string())
145                }
146            }),
147            Err(_) => Box::new(|_| Value::from(())),
148        },
149    }
150}
dre_core/values.rs

dre_core/
values.rs
