harmos_storage/query/transform/
histogram.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::histogram::Histogram;
17
18/* === Imports === */
19
20use crate::storage::query::transform::FieldPathExecution;
21// Standard Library Imports
22use std::sync::Arc;
23// Third Party Imports
24use arrow_array::{Float64Array, RecordBatch, UInt64Array};
25use arrow_schema::{DataType, Field, Schema};
26// Local Imports
27use super::downsample::number;
28use super::{invalid, numeric_type};
29use crate::Result;
30use harmos_domain::model::storage::query::scalar::ScalarValue;
31
32/* === Constants === */
33
34/* === Macros === */
35
36/* === Type Aliases === */
37
38/* === Structs === */
39
40pub struct Bins {
41    plan: Histogram,
42    schema: Arc<Schema>,
43    counts: Vec<u64>,
44}
45
46/* === Enums === */
47
48/* === Traits === */
49
50pub trait HistogramExecution {
51    fn output_schema(&self, input: &Schema) -> Result<Schema>;
52
53    fn start(&self, input: &Schema) -> Result<Bins>;
54
55    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch>;
56}
57
58/* === Implementations === */
59
60impl HistogramExecution for Histogram {
61    fn output_schema(&self, input: &Schema) -> Result<Schema> {
62        if self.edges.len() < 2
63            || self.edges.len() > 65536
64            || self.edges.iter().any(|edge| !edge.is_finite())
65            || self.edges.windows(2).any(|pair| pair[0] >= pair[1])
66        {
67            return Err(invalid("histogram requires 2..65536 finite strictly increasing edges"));
68        }
69        if !numeric_type(self.field.field(input)?.data_type()) {
70            return Err(invalid("histogram requires numeric field"));
71        }
72        Ok(Schema::new(vec![
73            Field::new("lower", DataType::Float64, true),
74            Field::new("upper", DataType::Float64, true),
75            Field::new("count", DataType::UInt64, false),
76        ]))
77    }
78
79    fn start(&self, input: &Schema) -> Result<Bins> {
80        Ok(Bins { schema: Arc::new(self.output_schema(input)?), plan: self.clone(), counts: vec![0; self.edges.len() + 1] })
81    }
82
83    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch> {
84        let mut bins = self.start(input.schema().as_ref())?;
85        bins.push(input, check)?;
86        bins.finish()
87    }
88}
89
90impl Bins {
91    pub fn retained_bytes(&self) -> usize {
92        self.counts.len().saturating_mul(16).saturating_add(1024)
93    }
94
95    pub fn finish_bytes(&self) -> usize {
96        self.counts.len().saturating_mul(96).saturating_add(4096)
97    }
98
99    pub fn push(&mut self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<()> {
100        let edges = &self.plan.edges;
101        let mut counts = self.counts.clone();
102        for row in 0..input.num_rows() {
103            check()?;
104            let value = self.plan.field.value(input, row)?;
105            if matches!(value, ScalarValue::Null | ScalarValue::Missing) {
106                continue;
107            }
108            let value = number(&value)?;
109            if !value.is_finite() {
110                continue;
111            }
112            let bin = if value < edges[0] {
113                0
114            } else if value > edges[edges.len() - 1] {
115                edges.len()
116            } else if value == edges[edges.len() - 1] {
117                edges.len() - 1
118            } else {
119                edges.partition_point(|edge| *edge <= value)
120            };
121            counts[bin] = counts[bin].checked_add(1).ok_or_else(|| invalid("histogram count overflow"))?;
122        }
123        self.counts = counts;
124        Ok(())
125    }
126
127    pub fn finish(self) -> Result<RecordBatch> {
128        let lower = std::iter::once(None).chain(self.plan.edges.iter().copied().map(Some)).collect::<Vec<_>>();
129        let upper = self.plan.edges.iter().copied().map(Some).chain(std::iter::once(None)).collect::<Vec<_>>();
130        RecordBatch::try_new(
131            self.schema,
132            vec![Arc::new(Float64Array::from(lower)), Arc::new(Float64Array::from(upper)), Arc::new(UInt64Array::from(self.counts))],
133        )
134        .map_err(|_| invalid("histogram output batch is invalid"))
135    }
136}
137
138/* === Public Functions === */
139
140/* === Private Functions === */
141
142/* === Tests === */
harmos_storage/query/transform/histogram.rs

harmos_storage/query/transform/
histogram.rs
