harmos_storage/query/transform/
mod.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Storage Arrow Query Contract
11
12/* === Modules === */
13
14pub mod aggregate;
15pub mod downsample;
16pub mod filter;
17pub mod histogram;
18pub mod project;
19pub mod sort;
20
21/* === Re-Exports === */
22
23pub use aggregate::AggregateExecution;
24pub use downsample::DownsampleExecution;
25pub use filter::FilterExecution;
26pub use harmos_domain::model::storage::query::transform::{FieldPath, Transform};
27pub use histogram::HistogramExecution;
28pub use project::ProjectExecution;
29pub use sort::SortExecution;
30
31/* === Imports === */
32
33// Standard Library Imports
34use std::sync::Arc;
35
36// Third Party Imports
37use arrow_array::{
38    Array, ArrayRef, BinaryArray, BooleanArray, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, NullArray,
39    RecordBatch, StringArray, StructArray, TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array, UInt64Array,
40};
41use arrow_schema::{DataType, Field, Schema, TimeUnit};
42use harmos_domain::model::storage::query::scalar::{ExactNumber, ScalarValue};
43
44// Local Imports
45use crate::{Error, Result};
46
47/* === Constants === */
48
49/* === Macros === */
50
51/* === Type Aliases === */
52
53/* === Structs === */
54
55/* === Enums === */
56
57/* === Traits === */
58
59pub trait FieldPathExecution {
60    fn field(&self, schema: &Schema) -> Result<Field>;
61
62    fn array<'a>(&self, batch: &'a RecordBatch) -> Result<&'a ArrayRef>;
63
64    fn value(&self, batch: &RecordBatch, row: usize) -> Result<ScalarValue>;
65}
66
67pub trait TransformExecution {
68    fn output_schema(&self, input: &Schema) -> Result<Schema>;
69
70    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch>;
71}
72
73/* === Implementations === */
74
75impl FieldPathExecution for FieldPath {
76    fn field(&self, schema: &Schema) -> Result<Field> {
77        if self.fields.is_empty() || self.fields.len() > 32 {
78            return Err(invalid("field path requires 1..32 components"));
79        }
80        let mut fields = schema.fields();
81        let mut found = None;
82        let mut nullable = false;
83        for (index, name) in self.fields.iter().enumerate() {
84            let field = fields.iter().find(|field| field.name() == name).ok_or_else(|| invalid("query field does not exist"))?;
85            nullable |= field.is_nullable();
86            found = Some(field.as_ref().clone().with_nullable(nullable));
87            if index + 1 < self.fields.len() {
88                let DataType::Struct(children) = field.data_type() else {
89                    return Err(invalid("field path traverses a non-struct value"));
90                };
91                fields = children;
92            }
93        }
94        Ok(found.expect("nonempty field path"))
95    }
96
97    fn array<'a>(&self, batch: &'a RecordBatch) -> Result<&'a ArrayRef> {
98        let first = self.fields.first().ok_or_else(|| invalid("empty field path"))?;
99        let mut array = batch.column_by_name(first).ok_or_else(|| invalid("query field does not exist"))?;
100        for name in self.fields.iter().skip(1) {
101            let structure =
102                array.as_any().downcast_ref::<StructArray>().ok_or_else(|| invalid("field path traverses a non-struct value"))?;
103            array = structure.column_by_name(name).ok_or_else(|| invalid("query child field does not exist"))?;
104        }
105        Ok(array)
106    }
107
108    fn value(&self, batch: &RecordBatch, row: usize) -> Result<ScalarValue> {
109        let first = self.fields.first().ok_or_else(|| invalid("empty field path"))?;
110        let mut array = batch.column_by_name(first).ok_or_else(|| invalid("query field does not exist"))?;
111        if array.is_null(row) {
112            return Ok(ScalarValue::Null);
113        }
114        for name in self.fields.iter().skip(1) {
115            let structure =
116                array.as_any().downcast_ref::<StructArray>().ok_or_else(|| invalid("field path traverses a non-struct value"))?;
117            array = structure.column_by_name(name).ok_or_else(|| invalid("query child field does not exist"))?;
118            if array.is_null(row) {
119                return Ok(ScalarValue::Null);
120            }
121        }
122        scalar(array.as_ref(), row)
123    }
124}
125
126impl TransformExecution for Transform {
127    fn output_schema(&self, input: &Schema) -> Result<Schema> {
128        match self {
129            Self::Filter(value) => value.output_schema(input),
130            Self::Project(value) => value.output_schema(input),
131            Self::Aggregate(value) => value.output_schema(input),
132            Self::Histogram(value) => value.output_schema(input),
133            Self::Downsample(value) => value.output_schema(input),
134            Self::Sort(value) => value.output_schema(input),
135        }
136    }
137
138    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch> {
139        match self {
140            Self::Filter(value) => value.apply(input, check),
141            Self::Project(value) => value.apply(input, check),
142            Self::Aggregate(value) => value.apply(input, check),
143            Self::Histogram(value) => value.apply(input, check),
144            Self::Downsample(value) => value.apply(input, check),
145            Self::Sort(value) => value.apply(input, check),
146        }
147    }
148}
149
150/* === Public Functions === */
151
152/* === Private Functions === */
153
154pub(crate) fn invalid(reason: &'static str) -> Error {
155    Error::InvalidInput { field: "storage query transform", reason }
156}
157
158pub(super) fn take(input: &RecordBatch, indices: Vec<u64>) -> Result<RecordBatch> {
159    let indices = UInt64Array::from(indices);
160    let columns = input
161        .columns()
162        .iter()
163        .map(|column| arrow_select::take::take(column.as_ref(), &indices, None).map_err(|_| invalid("Arrow row selection failed")))
164        .collect::<Result<Vec<_>>>()?;
165    RecordBatch::try_new_with_options(input.schema(), columns, &arrow_array::RecordBatchOptions::new().with_row_count(Some(indices.len())))
166        .map_err(|_| invalid("Arrow selected batch is invalid"))
167}
168
169pub(super) fn scalar(array: &dyn Array, row: usize) -> Result<ScalarValue> {
170    if array.is_null(row) {
171        return Ok(ScalarValue::Null);
172    }
173    macro_rules! read {
174        ($ty:ty, $variant:ident, $convert:expr) => {
175            if let Some(array) = array.as_any().downcast_ref::<$ty>() {
176                return Ok(ScalarValue::$variant(($convert)(array.value(row))));
177            }
178        };
179    }
180    read!(BooleanArray, Bool, |value| value);
181    read!(Int8Array, Signed, i64::from);
182    read!(Int16Array, Signed, i64::from);
183    read!(Int32Array, Signed, i64::from);
184    read!(Int64Array, Signed, |value| value);
185    read!(UInt8Array, Unsigned, u64::from);
186    read!(UInt16Array, Unsigned, u64::from);
187    read!(UInt32Array, Unsigned, u64::from);
188    read!(UInt64Array, Unsigned, |value| value);
189    read!(Float32Array, Float, f64::from);
190    read!(Float64Array, Float, |value| value);
191    read!(StringArray, Text, str::to_owned);
192    read!(TimestampNanosecondArray, Datetime, i128::from);
193    if let Some(exact) = array.as_any().downcast_ref::<StructArray>() {
194        return exact_value(exact, row);
195    }
196    Err(invalid("field is not a supported scalar"))
197}
198
199pub(super) fn exact_type() -> DataType {
200    DataType::Struct(
201        vec![
202            Arc::new(Field::new("negative", DataType::Boolean, false)),
203            Arc::new(Field::new("numerator", DataType::Binary, false)),
204            Arc::new(Field::new("denominator", DataType::Binary, false)),
205        ]
206        .into(),
207    )
208}
209
210fn exact_value(value: &StructArray, row: usize) -> Result<ScalarValue> {
211    if value.data_type() != &exact_type() {
212        return Err(invalid("struct is not an exact number"));
213    }
214    let sign = value.column(0).as_any().downcast_ref::<BooleanArray>().expect("validated exact sign");
215    let numerator = value.column(1).as_any().downcast_ref::<BinaryArray>().expect("validated exact numerator");
216    let denominator = value.column(2).as_any().downcast_ref::<BinaryArray>().expect("validated exact denominator");
217    let numerator = u128::from_le_bytes(numerator.value(row).try_into().map_err(|_| invalid("exact numerator requires 16 bytes"))?);
218    let denominator = u128::from_le_bytes(denominator.value(row).try_into().map_err(|_| invalid("exact denominator requires 16 bytes"))?);
219    let denominator = std::num::NonZeroU128::new(denominator).ok_or_else(|| invalid("zero exact denominator"))?;
220    let exact = ExactNumber::new(sign.value(row), numerator, denominator);
221    if exact.numerator != numerator || exact.denominator != denominator || exact.negative != sign.value(row) {
222        return Err(invalid("noncanonical exact number"));
223    }
224    Ok(ScalarValue::Exact(exact))
225}
226
227pub(super) fn scalar_type(data_type: &DataType) -> bool {
228    matches!(
229        data_type,
230        DataType::Null
231            | DataType::Boolean
232            | DataType::Int8
233            | DataType::Int16
234            | DataType::Int32
235            | DataType::Int64
236            | DataType::UInt8
237            | DataType::UInt16
238            | DataType::UInt32
239            | DataType::UInt64
240            | DataType::Float32
241            | DataType::Float64
242            | DataType::Utf8
243            | DataType::Timestamp(TimeUnit::Nanosecond, _)
244    ) || data_type == &exact_type()
245}
246
247pub(super) fn numeric_type(data_type: &DataType) -> bool {
248    matches!(
249        data_type,
250        DataType::Int8
251            | DataType::Int16
252            | DataType::Int32
253            | DataType::Int64
254            | DataType::UInt8
255            | DataType::UInt16
256            | DataType::UInt32
257            | DataType::UInt64
258            | DataType::Float32
259            | DataType::Float64
260    ) || data_type == &exact_type()
261}
262
263pub(super) fn values_array(data_type: &DataType, values: &[ScalarValue]) -> Result<ArrayRef> {
264    macro_rules! number {
265        ($array:ty, $variant:ident, $native:ty) => {{
266            let values = values
267                .iter()
268                .map(|value| match value {
269                    ScalarValue::$variant(value) => {
270                        <$native>::try_from(*value).ok().map(Some).ok_or_else(|| invalid("aggregate output width overflow"))
271                    }
272                    ScalarValue::Null | ScalarValue::Missing => Ok(None),
273                    _ => Err(invalid("aggregate output type mismatch")),
274                })
275                .collect::<Result<Vec<_>>>()?;
276            Arc::new(<$array>::from(values)) as ArrayRef
277        }};
278    }
279    Ok(match data_type {
280        DataType::Null => Arc::new(NullArray::new(values.len())),
281        DataType::Int8 => number!(Int8Array, Signed, i8),
282        DataType::Int16 => number!(Int16Array, Signed, i16),
283        DataType::Int32 => number!(Int32Array, Signed, i32),
284        DataType::Int64 => number!(Int64Array, Signed, i64),
285        DataType::UInt8 => number!(UInt8Array, Unsigned, u8),
286        DataType::UInt16 => number!(UInt16Array, Unsigned, u16),
287        DataType::UInt32 => number!(UInt32Array, Unsigned, u32),
288        DataType::UInt64 => number!(UInt64Array, Unsigned, u64),
289        DataType::Boolean => Arc::new(BooleanArray::from(
290            values
291                .iter()
292                .map(|value| match value {
293                    ScalarValue::Bool(value) => Some(*value),
294                    _ => None,
295                })
296                .collect::<Vec<_>>(),
297        )),
298        DataType::Float32 => Arc::new(Float32Array::from(
299            values
300                .iter()
301                .map(|value| match value {
302                    ScalarValue::Float(value) => Some(*value as f32),
303                    _ => None,
304                })
305                .collect::<Vec<_>>(),
306        )),
307        DataType::Float64 => Arc::new(Float64Array::from(
308            values
309                .iter()
310                .map(|value| match value {
311                    ScalarValue::Float(value) => Some(*value),
312                    _ => None,
313                })
314                .collect::<Vec<_>>(),
315        )),
316        DataType::Utf8 => Arc::new(StringArray::from(
317            values
318                .iter()
319                .map(|value| match value {
320                    ScalarValue::Text(value) => Some(value.as_str()),
321                    _ => None,
322                })
323                .collect::<Vec<_>>(),
324        )),
325        DataType::Timestamp(TimeUnit::Nanosecond, zone) => {
326            let values = values
327                .iter()
328                .map(|value| match value {
329                    ScalarValue::Datetime(value) => i64::try_from(*value).map(Some).map_err(|_| invalid("timestamp output overflow")),
330                    ScalarValue::Null | ScalarValue::Missing => Ok(None),
331                    _ => Err(invalid("timestamp output type mismatch")),
332                })
333                .collect::<Result<Vec<_>>>()?;
334            Arc::new(TimestampNanosecondArray::from(values).with_timezone_opt(zone.clone()))
335        }
336        data_type if data_type == &exact_type() => {
337            let values = values
338                .iter()
339                .map(|value| match value {
340                    ScalarValue::Null | ScalarValue::Missing => Ok(None),
341                    value => harmos_domain::model::storage::query::scalar::exact(value)
342                        .map(Some)
343                        .ok_or_else(|| invalid("exact output type mismatch")),
344                })
345                .collect::<Result<Vec<_>>>()?;
346            let signs = BooleanArray::from(values.iter().map(|value| value.is_some_and(|value| value.negative)).collect::<Vec<_>>());
347            let numerators = values.iter().map(|value| value.map_or(0, |value| value.numerator).to_le_bytes()).collect::<Vec<_>>();
348            let denominators =
349                values.iter().map(|value| value.map_or(1, |value| value.denominator.get()).to_le_bytes()).collect::<Vec<_>>();
350            let numerators = BinaryArray::from_iter_values(numerators.iter().map(|value| value.as_slice()));
351            let denominators = BinaryArray::from_iter_values(denominators.iter().map(|value| value.as_slice()));
352            let DataType::Struct(fields) = exact_type() else { unreachable!() };
353            let nulls = arrow_buffer::NullBuffer::from(values.iter().map(Option::is_some).collect::<Vec<_>>());
354            Arc::new(StructArray::new(fields, vec![Arc::new(signs), Arc::new(numerators), Arc::new(denominators)], Some(nulls)))
355        }
356        _ => return Err(invalid("unsupported scalar output type")),
357    })
358}
359
360/* === Tests === */
harmos_storage/query/transform/mod.rs

harmos_storage/query/transform/
mod.rs
