harmos_storage/query/transform/
mod.rs

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

harmos_storage/query/transform/
mod.rs
