Skip to main content

harmos_storage/query/transform/
downsample.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
14/* === Re-Exports === */
15
16pub use harmos_domain::model::storage::query::transform::downsample::{Algorithm, Downsample};
17
18/* === Imports === */
19
20// Standard Library Imports
21
22// Third Party Imports
23use arrow_array::RecordBatch;
24use arrow_schema::{DataType, Schema, TimeUnit};
25use harmos_domain::model::storage::query::scalar::ScalarValue;
26
27// Local Imports
28use super::{invalid, numeric_type};
29use crate::Result;
30use crate::storage::query::transform::FieldPathExecution;
31
32/* === Constants === */
33
34/* === Macros === */
35
36/* === Type Aliases === */
37
38/* === Structs === */
39
40/* === Enums === */
41
42/* === Traits === */
43
44pub trait DownsampleExecution {
45    fn output_schema(&self, input: &Schema) -> Result<Schema>;
46
47    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch>;
48}
49
50/* === Implementations === */
51
52impl DownsampleExecution for Downsample {
53    fn output_schema(&self, input: &Schema) -> Result<Schema> {
54        if self.max_rows == 0 {
55            return Err(invalid("downsample max_rows must be positive"));
56        }
57        if !matches!(self.algorithm, Algorithm::Decimate) {
58            let field = self.value.as_ref().ok_or_else(|| invalid("downsampling algorithm requires value field"))?.field(input)?;
59            if !numeric_type(field.data_type()) {
60                return Err(invalid("downsampling requires numeric values"));
61            }
62            let minimum = if matches!(self.algorithm, Algorithm::Extrema) { 2 } else { 3 };
63            if self.max_rows < minimum {
64                return Err(invalid("downsampling output bound is below algorithm minimum"));
65            }
66        }
67        if matches!(self.algorithm, Algorithm::LttbLog) && self.axis.is_none() {
68            return Err(invalid("logarithmic LTTB requires an explicit positive numeric axis"));
69        }
70        if let Some(axis) = &self.axis {
71            let field = axis.field(input)?;
72            if matches!(self.algorithm, Algorithm::LttbLog) && !numeric_type(field.data_type()) {
73                return Err(invalid("logarithmic LTTB requires a numeric axis"));
74            }
75            if !numeric_type(field.data_type()) && !matches!(field.data_type(), DataType::Timestamp(TimeUnit::Nanosecond, _)) {
76                return Err(invalid("downsample axis must be numeric or a timestamp"));
77            }
78        }
79        Ok(input.clone())
80    }
81
82    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch> {
83        if input.num_rows() <= self.max_rows {
84            return Ok(input.clone());
85        }
86        let mut spool = crate::query::spill::Builder::new(input.schema())?;
87        spool.push(input, check)?;
88        let output = crate::query::reduction::downsample(spool.finish()?, self, check)?;
89        let mut pages = output.pages()?;
90        let mut batches = Vec::new();
91        while let Some(batch) = pages.next()? {
92            check()?;
93            batches.push(batch);
94        }
95        // This Eager API Owns Only Its Requested Output, Never Coordinate Copies.
96        arrow_select::concat::concat_batches(&input.schema(), batches.iter()).map_err(|_| invalid("downsample output construction failed"))
97    }
98}
99
100/* === Public Functions === */
101
102/* === Private Functions === */
103
104pub(in crate::query) fn number(value: &ScalarValue) -> Result<f64> {
105    match value {
106        ScalarValue::Float(value) => Ok(*value),
107        value => {
108            let exact = harmos_domain::model::storage::query::scalar::exact(value)
109                .ok_or_else(|| invalid("numeric field has a null or nonnumeric value"))?;
110            let value = exact.numerator as f64 / exact.denominator.get() as f64;
111            Ok(if exact.negative { -value } else { value })
112        }
113    }
114}
115
116/* === Tests === */