nodedb_lite/engine/array/ops/
aggregate.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! `ArrayOp::Aggregate` handler for NodeDB-Lite.
4//!
5//! Scans all segments and the memtable, applies `aggregate_attr` (or
6//! `group_by_dim`) per tile, then merges partials across tiles via
7//! `AggregateResult::merge`. Bitemporal system-time ceiling is honoured:
8//! tiles whose `system_from_ms` exceeds the cutoff are skipped, mirroring
9//! the same filter applied in `engine.rs::slice`.
10
11use std::collections::HashMap;
12use std::sync::Arc;
13
14use nodedb_array::query::aggregate::{
15    AggregateResult, GroupAggregate, Reducer, aggregate_attr, group_by_dim,
16};
17use nodedb_array::tile::sparse_tile::SparseTile;
18use nodedb_array::{SegmentReader, TilePayload};
19use nodedb_physical::physical_plan::ArrayReducer;
20use nodedb_types::result::QueryResult;
21use nodedb_types::value::Value;
22
23use crate::engine::array::engine::ArrayEngineState;
24use crate::error::LiteError;
25use crate::runtime::now_millis_i64;
26use crate::storage::engine::StorageEngine;
27
28fn map_reducer(r: ArrayReducer) -> Reducer {
29    match r {
30        ArrayReducer::Sum => Reducer::Sum,
31        ArrayReducer::Count => Reducer::Count,
32        ArrayReducer::Min => Reducer::Min,
33        ArrayReducer::Max => Reducer::Max,
34        ArrayReducer::Mean => Reducer::Mean,
35    }
36}
37
38fn result_to_value(r: AggregateResult) -> Value {
39    match r.finalize() {
40        Some(v) => Value::Float(v),
41        None => Value::Null,
42    }
43}
44
45fn coord_value_to_value(cv: nodedb_array::types::coord::value::CoordValue) -> Value {
46    use nodedb_array::types::coord::value::CoordValue;
47    match cv {
48        CoordValue::Int64(i) => Value::Integer(i),
49        CoordValue::TimestampMs(i) => Value::Integer(i),
50        CoordValue::Float64(f) => Value::Float(f),
51        CoordValue::String(s) => Value::String(s),
52    }
53}
54
55/// Emit one row per group.  Key column name comes from the dim at `dim_idx`.
56fn emit_groups(groups: Vec<GroupAggregate>, rows: &mut Vec<Vec<Value>>) {
57    for g in groups {
58        rows.push(vec![coord_value_to_value(g.key), result_to_value(g.result)]);
59    }
60}
61
62/// Merge a group map `dst` with groups from `src`.
63fn merge_groups(
64    dst: &mut HashMap<Vec<u8>, AggregateResult>,
65    src: Vec<GroupAggregate>,
66) -> Result<(), LiteError> {
67    for g in src {
68        let key_bytes = zerompk::to_msgpack_vec(&g.key).map_err(|e| LiteError::Serialization {
69            detail: format!("encode group key: {e}"),
70        })?;
71        let entry = dst
72            .entry(key_bytes)
73            .or_insert(AggregateResult::Empty(match g.result {
74                AggregateResult::Sum { .. } => Reducer::Sum,
75                AggregateResult::Count { .. } => Reducer::Count,
76                AggregateResult::Min { .. } => Reducer::Min,
77                AggregateResult::Max { .. } => Reducer::Max,
78                AggregateResult::Mean { .. } => Reducer::Mean,
79                AggregateResult::Empty(r) => r,
80            }));
81        *entry = entry.merge(g.result);
82    }
83    Ok(())
84}
85
86/// Accumulate a single tile into the running scalar partial or group map.
87fn accumulate_tile(
88    tile: &SparseTile,
89    attr_idx: usize,
90    reducer: Reducer,
91    group_dim: Option<usize>,
92    scalar: &mut AggregateResult,
93    groups: &mut HashMap<Vec<u8>, AggregateResult>,
94) -> Result<(), LiteError> {
95    match group_dim {
96        None => {
97            let partial = aggregate_attr(tile, attr_idx, reducer);
98            *scalar = scalar.merge(partial);
99        }
100        Some(dim_idx) => {
101            let tile_groups = group_by_dim(tile, dim_idx, attr_idx, reducer);
102            merge_groups(groups, tile_groups)?;
103        }
104    }
105    Ok(())
106}
107
108/// Execute `ArrayOp::Aggregate` for the Lite engine.
109pub async fn aggregate<S: StorageEngine>(
110    array_state: &Arc<tokio::sync::Mutex<ArrayEngineState>>,
111    storage: &Arc<S>,
112    name: &str,
113    attr_idx: u32,
114    reducer: ArrayReducer,
115    group_by_dim_idx: i32,
116) -> Result<QueryResult, LiteError> {
117    let system_as_of = now_millis_i64();
118    let reducer_inner = map_reducer(reducer);
119    let group_dim: Option<usize> = if group_by_dim_idx >= 0 {
120        Some(group_by_dim_idx as usize)
121    } else {
122        None
123    };
124
125    let (seg_ids, schema, attr_count, dim_count) = {
126        let state = array_state.lock().await;
127        let arr = state
128            .arrays
129            .get(name)
130            .ok_or_else(|| LiteError::BadRequest {
131                detail: format!("array '{name}' not found"),
132            })?;
133        let seg_ids: Vec<u64> = arr.manifest.segments.iter().map(|s| s.id).collect();
134        let attr_count = arr.schema.attrs.len();
135        let dim_count = arr.schema.dims.len();
136        (seg_ids, arr.schema.clone(), attr_count, dim_count)
137    };
138
139    let attr_idx_usize = attr_idx as usize;
140    if attr_idx_usize >= attr_count {
141        return Err(LiteError::BadRequest {
142            detail: format!(
143                "aggregate attr index {attr_idx} out of range (array '{name}' has {attr_count} attrs)"
144            ),
145        });
146    }
147    if let Some(d) = group_dim
148        && d >= dim_count
149    {
150        return Err(LiteError::BadRequest {
151            detail: format!(
152                "group_by dim index {d} out of range (array '{name}' has {dim_count} dims)"
153            ),
154        });
155    }
156
157    let mut scalar = AggregateResult::Empty(reducer_inner);
158    // Key: msgpack-encoded CoordValue, Value: running partial.
159    let mut groups: HashMap<Vec<u8>, AggregateResult> = HashMap::new();
160
161    // Segments.
162    for seg_id in &seg_ids {
163        let bytes = crate::engine::array::segments::load_segment(storage, name, *seg_id).await?;
164        let reader = SegmentReader::open(&bytes).map_err(|e| LiteError::Storage {
165            detail: format!("open segment {seg_id}: {e}"),
166        })?;
167        for idx in 0..reader.tile_count() {
168            let entry_tile_id = reader.tiles()[idx].tile_id;
169            if entry_tile_id.system_from_ms > system_as_of {
170                continue;
171            }
172            let payload = reader.read_tile(idx).map_err(|e| LiteError::Storage {
173                detail: format!("read_tile seg {seg_id} idx {idx}: {e}"),
174            })?;
175            let TilePayload::Sparse(tile) = payload else {
176                continue;
177            };
178            accumulate_tile(
179                &tile,
180                attr_idx_usize,
181                reducer_inner,
182                group_dim,
183                &mut scalar,
184                &mut groups,
185            )?;
186        }
187    }
188
189    // Memtable.
190    {
191        let mut state = array_state.lock().await;
192        let arr = state
193            .arrays
194            .get_mut(name)
195            .ok_or_else(|| LiteError::BadRequest {
196                detail: format!("array '{name}' not found"),
197            })?;
198        let mem_tiles = arr
199            .memtable
200            .drain_all_tiles_read_only(system_as_of, &schema)
201            .map_err(|e| LiteError::Storage {
202                detail: format!("memtable drain: {e}"),
203            })?;
204        for (_tile_id, tile) in &mem_tiles {
205            accumulate_tile(
206                tile,
207                attr_idx_usize,
208                reducer_inner,
209                group_dim,
210                &mut scalar,
211                &mut groups,
212            )?;
213        }
214    }
215
216    // Build result.
217    match group_dim {
218        None => {
219            let columns = vec!["value".to_string()];
220            let row = vec![result_to_value(scalar)];
221            Ok(QueryResult {
222                columns,
223                rows: vec![row],
224                rows_affected: 0,
225                command: None,
226            })
227        }
228        Some(_) => {
229            let columns = vec!["key".to_string(), "value".to_string()];
230            let mut rows: Vec<Vec<Value>> = Vec::new();
231            // Decode group keys back to CoordValue for display.
232            for (key_bytes, result) in groups {
233                let coord_val: nodedb_array::types::coord::value::CoordValue =
234                    zerompk::from_msgpack(&key_bytes).map_err(|e| LiteError::Serialization {
235                        detail: format!("decode group key: {e}"),
236                    })?;
237                emit_groups(
238                    vec![GroupAggregate {
239                        key: coord_val,
240                        result,
241                    }],
242                    &mut rows,
243                );
244            }
245            Ok(QueryResult {
246                columns,
247                rows,
248                rows_affected: 0,
249                command: None,
250            })
251        }
252    }
253}
nodedb_lite/engine/array/ops/aggregate.rs

nodedb_lite/engine/array/ops/
aggregate.rs
