harmos_storage/query/transform/
aggregate.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::aggregate::{Aggregate, Measure};
17
18/* === Imports === */
19
20use crate::storage::query::transform::FieldPathExecution;
21// Standard Library Imports
22use std::cmp::Ordering;
23use std::sync::Arc;
24// Third Party Imports
25use arrow_array::RecordBatch;
26use arrow_schema::{DataType, Field, Schema};
27// Local Imports
28use super::{exact_type, invalid, numeric_type, scalar_type, values_array};
29use crate::storage::query::accumulate::{Accumulator, Scratch};
30use crate::{MemorySize, Result};
31use harmos_domain::model::storage::query::scalar::{Agg, ScalarValue};
32
33/* === Constants === */
34
35/* === Macros === */
36
37/* === Type Aliases === */
38
39/* === Structs === */
40
41pub struct Groups {
42    plan: Aggregate,
43    schema: Arc<Schema>,
44    groups: Vec<Group>,
45    capture: u64,
46    retained: usize,
47}
48
49struct Group {
50    keys: Vec<ScalarValue>,
51    accumulators: Vec<Accumulator>,
52}
53
54/* === Enums === */
55
56/* === Traits === */
57
58pub trait AggregateExecution {
59    fn output_schema(&self, input: &Schema) -> Result<Schema>;
60
61    fn start(&self, input: &Schema) -> Result<Groups>;
62
63    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch>;
64}
65
66/* === Implementations === */
67
68impl AggregateExecution for Aggregate {
69    fn output_schema(&self, input: &Schema) -> Result<Schema> {
70        if self.groups.len() > 64 || self.measures.is_empty() || self.measures.len() > 64 {
71            return Err(invalid("aggregate requires at most 64 groups and 1..64 measures"));
72        }
73        let mut fields = Vec::new();
74        for group in &self.groups {
75            let field = group.field.field(input)?;
76            if !scalar_type(field.data_type()) {
77                return Err(invalid("aggregate grouping requires scalar keys"));
78            }
79            fields.push(field.with_name(&group.name).with_nullable(true));
80        }
81        for measure in &self.measures {
82            let field = measure.field.field(input)?;
83            if !scalar_type(field.data_type()) {
84                return Err(invalid("aggregate requires scalar values"));
85            }
86            let statistic = measure.statistic;
87            if !matches!(statistic, crate::Aggregate::Count | crate::Aggregate::Min | crate::Aggregate::Max)
88                && !numeric_type(field.data_type())
89            {
90                return Err(invalid("numeric aggregate requires numeric field"));
91            }
92            if let crate::Aggregate::Percentile(rank) = statistic {
93                // The shared vocabulary validates the rank through its normal codec.
94                if rank > 100 {
95                    return Err(invalid("invalid percentile rank"));
96                }
97            }
98            let output = match statistic {
99                crate::Aggregate::Count => DataType::UInt64,
100                crate::Aggregate::Min | crate::Aggregate::Max => field.data_type().clone(),
101                crate::Aggregate::Sum | crate::Aggregate::Mean | crate::Aggregate::Median
102                    if !matches!(field.data_type(), DataType::Float32 | DataType::Float64) =>
103                {
104                    exact_type()
105                }
106                _ => DataType::Float64,
107            };
108            let mut field = Field::new(&measure.name, output, !matches!(statistic, crate::Aggregate::Count));
109            if field.data_type() == &exact_type() {
110                field = field
111                    .with_metadata(std::collections::HashMap::from([("harmos.logical_type".to_owned(), "exact-rational-v1".to_owned())]));
112            }
113            fields.push(field);
114        }
115        for field in &fields {
116            if field.name().is_empty() || fields.iter().filter(|other| other.name() == field.name()).count() != 1 {
117                return Err(invalid("aggregate output names must be nonempty and unique"));
118            }
119        }
120        Ok(Schema::new(fields))
121    }
122
123    fn start(&self, input: &Schema) -> Result<Groups> {
124        let schema = Arc::new(self.output_schema(input)?);
125        let mut groups = Groups { plan: self.clone(), schema, groups: Vec::new(), capture: 0, retained: 0 };
126        // An ungrouped reduction always yields its identity row, even for no input.
127        if self.groups.is_empty() {
128            groups.retained = group_bound(self, &[])?;
129            groups.groups.push(new_group(self, Vec::new()));
130        }
131        Ok(groups)
132    }
133
134    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch> {
135        let mut groups = self.start(input.schema().as_ref())?;
136        groups.push(input, check)?;
137        groups.finish(check)
138    }
139}
140
141impl Groups {
142    pub fn retained_bytes(&self) -> usize {
143        self.retained
144    }
145
146    pub fn finish_bytes(&self) -> Result<usize> {
147        let text = self
148            .groups
149            .iter()
150            .flat_map(|group| &group.keys)
151            .map(MemorySize::memory_size)
152            .try_fold(0usize, |total, bytes| total.checked_add(bytes))
153            .ok_or_else(|| invalid("aggregate output estimate overflow"))?;
154        self.groups
155            .len()
156            .max(1)
157            .checked_mul(self.schema.fields().len())
158            .and_then(|cells| cells.checked_mul(512))
159            .and_then(|bytes| bytes.checked_add(text.saturating_mul(4)))
160            .and_then(|bytes| bytes.checked_add(65536))
161            .ok_or_else(|| invalid("aggregate finish estimate overflow"))
162    }
163
164    pub fn push(&mut self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<()> {
165        // Extract Keys And Values Before Any Group Changes.
166        let mut rows = Vec::with_capacity(input.num_rows());
167        for row in 0..input.num_rows() {
168            check()?;
169            let keys = self.plan.groups.iter().map(|group| group.field.value(input, row).map(normalized)).collect::<Result<Vec<_>>>()?;
170            let values =
171                self.plan.measures.iter().map(|measure| measure.field.value(input, row).map(normalized)).collect::<Result<Vec<_>>>()?;
172            rows.push((keys, values));
173        }
174
175        // Fold Rows Into Sorted Groups.
176        for (keys, values) in rows {
177            check()?;
178            let index = match self.find(&keys)? {
179                Ok(index) => index,
180                Err(index) => {
181                    let bytes = group_bound(&self.plan, &keys)?;
182                    self.groups.insert(index, new_group(&self.plan, keys));
183                    self.retained = self.retained.saturating_add(bytes);
184                    index
185                }
186            };
187            let group = &mut self.groups[index];
188            for ((measure, accumulator), value) in self.plan.measures.iter().zip(&mut group.accumulators).zip(values) {
189                // Retained order statistics grow only by the selected scalars.
190                if Agg::Of(measure.statistic).retains_values() || matches!(value, ScalarValue::Text(_)) {
191                    let bytes = value.memory_size().saturating_mul(4).saturating_add(128);
192                    self.retained = self.retained.saturating_add(bytes);
193                }
194                accumulator.push(self.capture, &value, &mut Scratch::new(usize::MAX))?;
195            }
196            self.capture = self.capture.checked_add(1).ok_or_else(|| invalid("query capture count overflow"))?;
197        }
198        Ok(())
199    }
200
201    pub fn finish(self, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch> {
202        check()?;
203        let width = self.schema.fields().len();
204        let mut values = vec![Vec::with_capacity(self.groups.len()); width];
205        for group in self.groups {
206            check()?;
207            for (column, key) in group.keys.into_iter().enumerate() {
208                values[column].push(key);
209            }
210            for (column, accumulator) in group.accumulators.into_iter().enumerate() {
211                values[self.plan.groups.len() + column].push(accumulator.finish(check)?);
212            }
213        }
214        let columns = self
215            .schema
216            .fields()
217            .iter()
218            .zip(values)
219            .map(|(field, values)| values_array(field.data_type(), &values))
220            .collect::<Result<Vec<_>>>()?;
221        RecordBatch::try_new(self.schema, columns).map_err(|_| invalid("aggregate output batch is invalid"))
222    }
223
224    // Binary Search Over Groups Kept In Key Order.
225    fn find(&self, keys: &[ScalarValue]) -> Result<std::result::Result<usize, usize>> {
226        let (mut low, mut high) = (0usize, self.groups.len());
227        while low < high {
228            let middle = low + (high - low) / 2;
229            match compare_keys(&self.groups[middle].keys, keys)? {
230                Ordering::Less => low = middle + 1,
231                Ordering::Greater => high = middle,
232                Ordering::Equal => return Ok(Ok(middle)),
233            }
234        }
235        Ok(Err(low))
236    }
237}
238
239/* === Public Functions === */
240
241/* === Private Functions === */
242
243fn group_bound(plan: &Aggregate, keys: &[ScalarValue]) -> Result<usize> {
244    keys.iter()
245        .map(MemorySize::memory_size)
246        .try_fold(0usize, |total, bytes| total.checked_add(bytes))
247        .and_then(|bytes| {
248            plan.measures.len().checked_mul(std::mem::size_of::<Accumulator>() * 4).and_then(|accumulators| bytes.checked_add(accumulators))
249        })
250        .and_then(|bytes| bytes.checked_mul(4))
251        .and_then(|bytes| bytes.checked_add(1024))
252        .ok_or_else(|| invalid("aggregate group memory overflow"))
253}
254
255fn new_group(plan: &Aggregate, keys: Vec<ScalarValue>) -> Group {
256    Group { keys, accumulators: plan.measures.iter().map(|measure| Accumulator::new(Agg::Of(measure.statistic))).collect() }
257}
258
259fn compare_keys(left: &[ScalarValue], right: &[ScalarValue]) -> Result<Ordering> {
260    for (left, right) in left.iter().zip(right) {
261        let comparison = left.compare(right)?;
262        if !comparison.is_eq() {
263            return Ok(comparison);
264        }
265    }
266    Ok(left.len().cmp(&right.len()))
267}
268
269fn normalized(value: ScalarValue) -> ScalarValue {
270    match value {
271        ScalarValue::Float(number) if !number.is_finite() => ScalarValue::Null,
272        value => value,
273    }
274}
275
276/* === Tests === */
harmos_storage/query/transform/aggregate.rs

harmos_storage/query/transform/
aggregate.rs
