kevy_index/
agg.rs

1//! [`AggSegment`] — one shard's slice of one aggregate index (RFC
2//! v3.1, KIND agg): per-group count / sum / min / max maintained
3//! synchronously with writes. min/max stay EXACT under deletion via a
4//! per-group value multiset (BTreeMap value → multiplicity); a
5//! row → (group, value) reverse map supports O(log) updates — the
6//! same derived-by-construction discipline as [`crate::Segment`].
7
8use std::collections::{BTreeMap, HashMap};
9
10use crate::IndexValue;
11
12/// One group's live statistics.
13#[derive(Debug, Clone, PartialEq)]
14pub struct GroupStats {
15    /// Rows in the group.
16    pub count: u64,
17    /// Sum of the aggregated field (f64 accumulation — the i64
18    /// overflow guard; precision bounds documented).
19    pub sum: f64,
20    /// Exact minimum (None only when count == 0).
21    pub min: Option<IndexValue>,
22    /// Exact maximum.
23    pub max: Option<IndexValue>,
24}
25
26impl GroupStats {
27    /// Derived average.
28    pub fn avg(&self) -> Option<f64> {
29        (self.count > 0).then(|| self.sum / self.count as f64)
30    }
31}
32
33struct Group {
34    count: u64,
35    sum: f64,
36    /// value → multiplicity; min/max = first/last key.
37    values: BTreeMap<IndexValue, u32>,
38}
39
40/// Ranking metric for [`AggSegment::top_groups`].
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
42pub enum AggBy {
43    /// By row count (default).
44    #[default]
45    Count,
46    /// By sum.
47    Sum,
48    /// By minimum (ascending — smallest mins first).
49    Min,
50    /// By maximum (descending — largest maxes first).
51    Max,
52}
53
54impl AggBy {
55    /// Wire tag.
56    pub fn parse(raw: &[u8]) -> Option<AggBy> {
57        if raw.eq_ignore_ascii_case(b"count") {
58            Some(AggBy::Count)
59        } else if raw.eq_ignore_ascii_case(b"sum") {
60            Some(AggBy::Sum)
61        } else if raw.eq_ignore_ascii_case(b"min") {
62            Some(AggBy::Min)
63        } else if raw.eq_ignore_ascii_case(b"max") {
64            Some(AggBy::Max)
65        } else {
66            None
67        }
68    }
69}
70
71/// Sizing counters (memory formula / IDX.LIST).
72#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
73pub struct AggStats {
74    /// Live groups.
75    pub groups: u64,
76    /// Rows participating.
77    pub rows: u64,
78    /// Rows excluded (coerce failure / missing group field).
79    pub excluded: u64,
80    /// Approximate heap bytes (RFC D4 formula's measured side).
81    pub approx_bytes: u64,
82}
83
84/// One shard's aggregate segment.
85#[derive(Default)]
86pub struct AggSegment {
87    groups: HashMap<Vec<u8>, Group>,
88    /// row key → (group, value) for O(log) update/remove.
89    rows: HashMap<Vec<u8>, (Vec<u8>, IndexValue)>,
90    excluded: u64,
91}
92
93impl AggSegment {
94    /// Empty segment.
95    pub fn new() -> Self {
96        Self::default()
97    }
98
99    /// (Re-)register one row: `Some((group, value))` = row
100    /// participates; `None` = removed or excluded. `excluded_row`
101    /// marks the None case as a coercion/missing-field exclusion
102    /// (counted) rather than a plain delete.
103    // missing_panics_doc: the only panic is the "group of live row" expect —
104    // an internal rows↔groups invariant, never reachable from caller input.
105    #[allow(clippy::missing_panics_doc)]
106    pub fn apply(&mut self, key: &[u8], entry: Option<(Vec<u8>, IndexValue)>, excluded_row: bool) {
107        // Fast path: same-group value update (the dominant serving
108        // write shape — measured 16.8% write tax on a Zipf corpus
109        // with the retract+register path; hot groups make every
110        // full retract pay two map round-trips and two clones).
111        if let Some((group, val)) = &entry
112            && let Some((old_group, old_val)) = self.rows.get_mut(key)
113            && old_group == group
114        {
115            if old_val == val {
116                return; // nothing changed
117            }
118            let g = self.groups.get_mut(group).expect("group of live row");
119            g.sum += val.as_f64() - old_val.as_f64();
120            match g.values.get_mut(old_val) {
121                Some(m) if *m > 1 => *m -= 1,
122                _ => {
123                    g.values.remove(old_val);
124                }
125            }
126            *g.values.entry(val.clone()).or_insert(0) += 1;
127            *old_val = val.clone();
128            return;
129        }
130        self.retract_row(key);
131        match entry {
132            Some((group, val)) => {
133                let g = self.groups.entry(group.clone()).or_insert(Group {
134                    count: 0,
135                    sum: 0.0,
136                    values: BTreeMap::new(),
137                });
138                g.count += 1;
139                g.sum += val.as_f64();
140                *g.values.entry(val.clone()).or_insert(0) += 1;
141                self.rows.insert(key.to_vec(), (group, val));
142            }
143            None if excluded_row => self.excluded += 1,
144            None => {}
145        }
146    }
147
148    /// Retract one row's current contribution (no-op for an unknown
149    /// key); drops the group when its last row leaves.
150    fn retract_row(&mut self, key: &[u8]) {
151        if let Some((old_group, old_val)) = self.rows.remove(key) {
152            let empty = {
153                let g = self.groups.get_mut(&old_group).expect("group of live row");
154                g.count -= 1;
155                g.sum -= old_val.as_f64();
156                match g.values.get_mut(&old_val) {
157                    Some(m) if *m > 1 => *m -= 1,
158                    _ => {
159                        g.values.remove(&old_val);
160                    }
161                }
162                g.count == 0
163            };
164            if empty {
165                self.groups.remove(&old_group);
166            }
167        }
168    }
169
170    /// One group's stats (`count == 0` shape for an unknown group).
171    pub fn group(&self, group: &[u8]) -> GroupStats {
172        match self.groups.get(group) {
173            Some(g) => GroupStats {
174                count: g.count,
175                sum: g.sum,
176                min: g.values.keys().next().cloned(),
177                max: g.values.keys().next_back().cloned(),
178            },
179            None => GroupStats { count: 0, sum: 0.0, min: None, max: None },
180        }
181    }
182
183    /// Top `limit` groups ranked by `by` (count/sum/max descending,
184    /// min ascending), ties broken by group key ascending.
185    ///
186    /// Bounded selection over BORROWED keys — the first cut cloned
187    /// and sorted every group per query (measured as the dominant
188    /// per-shard cost at 10k groups); only the winners materialize.
189    pub fn top_groups(&self, by: AggBy, limit: usize) -> Vec<(Vec<u8>, GroupStats)> {
190        let score_of = |g: &Group| -> f64 {
191            match by {
192                AggBy::Count => g.count as f64,
193                AggBy::Sum => g.sum,
194                AggBy::Max => g.values.keys().next_back().map_or(f64::NEG_INFINITY, IndexValue::as_f64),
195                AggBy::Min => g.values.keys().next().map_or(f64::NEG_INFINITY, |v| -v.as_f64()),
196            }
197        };
198        // float_cmp: exact equality is the tiebreak trigger — an epsilon would
199        // make the top-K selection non-deterministic for equal scores.
200        #[allow(clippy::float_cmp)]
201        let better = |a: (f64, &[u8]), b: (f64, &[u8])| a.0 > b.0 || (a.0 == b.0 && a.1 < b.1);
202        let mut top: Vec<(f64, &Vec<u8>)> = Vec::with_capacity(limit.min(1024) + 1);
203        for (k, g) in &self.groups {
204            let cand = (score_of(g), k);
205            if top.len() < limit {
206                top.push(cand);
207                if top.len() == limit {
208                    top.sort_by(|a, b| b.0.total_cmp(&a.0).then_with(|| a.1.cmp(b.1)));
209                }
210            } else if let Some(last) = top.last()
211                && better((cand.0, cand.1), (last.0, last.1))
212            {
213                let pos = top.partition_point(|e| better((e.0, e.1), (cand.0, cand.1)));
214                top.insert(pos, cand);
215                top.pop();
216            }
217        }
218        if top.len() < limit {
219            top.sort_by(|a, b| b.0.total_cmp(&a.0).then_with(|| a.1.cmp(b.1)));
220        }
221        top.into_iter().map(|(_, k)| (k.clone(), self.group(k))).collect()
222    }
223
224    /// Every group, UNRANKED — the fan-out chunk shape (ranking
225    /// happens once, at the reduce, after cross-shard merge; sorting
226    /// per shard would be wasted work).
227    pub fn all_groups(&self) -> Vec<(Vec<u8>, GroupStats)> {
228        self.groups.keys().map(|k| (k.clone(), self.group(k))).collect()
229    }
230
231    /// Membership probe (verify hook).
232    pub fn contains(&self, key: &[u8]) -> bool {
233        self.rows.contains_key(key)
234    }
235
236    /// Live counters. Byte constants calibrated against measured RSS
237    /// growth at 1M rows / 10k Zipf groups (the first-cut 40/24
238    /// constants overestimated 2× — BTreeMap packs ~11 entries per
239    /// node and the reverse map's vecs are small-alloc pooled).
240    pub fn stats(&self) -> AggStats {
241        let distinct: u64 = self.groups.values().map(|g| g.values.len() as u64).sum();
242        let gkey: u64 = self.groups.keys().map(|k| k.len() as u64).sum();
243        let rowbytes: u64 = self.rows.keys().map(|k| (k.len() + 10) as u64).sum();
244        AggStats {
245            groups: self.groups.len() as u64,
246            rows: self.rows.len() as u64,
247            excluded: self.excluded,
248            approx_bytes: gkey + self.groups.len() as u64 * 64 + distinct * 18 + rowbytes,
249        }
250    }
251}
252
253/// Shared ranking order (per-shard AND at the reduce after merging
254/// shard partials — one definition, no drift).
255pub fn sort_groups(all: &mut [(Vec<u8>, GroupStats)], by: AggBy) {
256    match by {
257        AggBy::Count => all.sort_by(|a, b| b.1.count.cmp(&a.1.count).then_with(|| a.0.cmp(&b.0))),
258        AggBy::Sum => all.sort_by(|a, b| b.1.sum.total_cmp(&a.1.sum).then_with(|| a.0.cmp(&b.0))),
259        AggBy::Min => all.sort_by(|a, b| {
260            match (&a.1.min, &b.1.min) {
261                (Some(x), Some(y)) => x.cmp(y),
262                (Some(_), None) => std::cmp::Ordering::Less,
263                (None, Some(_)) => std::cmp::Ordering::Greater,
264                (None, None) => std::cmp::Ordering::Equal,
265            }
266            .then_with(|| a.0.cmp(&b.0))
267        }),
268        AggBy::Max => all.sort_by(|a, b| {
269            match (&b.1.max, &a.1.max) {
270                (Some(x), Some(y)) => x.cmp(y),
271                (Some(_), None) => std::cmp::Ordering::Less,
272                (None, Some(_)) => std::cmp::Ordering::Greater,
273                (None, None) => std::cmp::Ordering::Equal,
274            }
275            .then_with(|| a.0.cmp(&b.0))
276        }),
277    }
278}
279
280/// Merge shard partials for one group (reduce side): counts/sums add,
281/// min/max take extremes.
282pub fn merge_group(into: &mut GroupStats, part: &GroupStats) {
283    into.count += part.count;
284    into.sum += part.sum;
285    into.min = match (into.min.take(), part.min.clone()) {
286        (Some(a), Some(b)) => Some(if b < a { b } else { a }),
287        (a, b) => a.or(b),
288    };
289    into.max = match (into.max.take(), part.max.clone()) {
290        (Some(a), Some(b)) => Some(if b > a { b } else { a }),
291        (a, b) => a.or(b),
292    };
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    fn seg() -> AggSegment {
300        let mut s = AggSegment::new();
301        // orders: group = status, value = amount
302        for (k, g, v) in [
303            ("o1", "paid", 100),
304            ("o2", "paid", 250),
305            ("o3", "open", 40),
306            ("o4", "paid", 100),
307            ("o5", "open", 999),
308        ] {
309            s.apply(k.as_bytes(), Some((g.as_bytes().to_vec(), IndexValue::I64(v))), false);
310        }
311        s
312    }
313
314    #[test]
315    fn group_stats_exact() {
316        let s = seg();
317        let g = s.group(b"paid");
318        assert_eq!((g.count, g.sum), (3, 450.0));
319        assert_eq!(g.min, Some(IndexValue::I64(100)));
320        assert_eq!(g.max, Some(IndexValue::I64(250)));
321        assert_eq!(g.avg(), Some(150.0));
322        let none = s.group(b"nope");
323        assert_eq!(none.count, 0);
324        assert!(none.min.is_none() && none.avg().is_none());
325    }
326
327    #[test]
328    fn min_max_exact_under_delete_and_update() {
329        let mut s = seg();
330        // delete the paid max (o2=250): max must fall back to 100
331        s.apply(b"o2", None, false);
332        let g = s.group(b"paid");
333        assert_eq!((g.count, g.max.clone()), (2, Some(IndexValue::I64(100))));
334        // duplicate values: removing ONE 100 keeps the other
335        s.apply(b"o1", None, false);
336        let g = s.group(b"paid");
337        assert_eq!((g.count, g.min.clone()), (1, Some(IndexValue::I64(100))));
338        // update moves a row across groups
339        s.apply(b"o3", Some((b"paid".to_vec(), IndexValue::I64(40))), false);
340        assert_eq!(s.group(b"paid").count, 2);
341        assert_eq!(s.group(b"open").count, 1);
342        assert_eq!(s.group(b"paid").min, Some(IndexValue::I64(40)));
343        // last row leaving a group drops the group
344        s.apply(b"o5", None, false);
345        assert_eq!(s.group(b"open").count, 0);
346        assert_eq!(s.stats().groups, 1);
347    }
348
349    #[test]
350    fn top_groups_all_metrics() {
351        let s = seg();
352        let top = s.top_groups(AggBy::Count, 10);
353        assert_eq!(top[0].0, b"paid".to_vec());
354        let top = s.top_groups(AggBy::Sum, 10);
355        assert_eq!(top[0].0, b"open".to_vec(), "open sum 1039 > paid 450");
356        let top = s.top_groups(AggBy::Min, 10);
357        assert_eq!(top[0].0, b"open".to_vec(), "min ascending: 40 first");
358        let top = s.top_groups(AggBy::Max, 1);
359        assert_eq!(top.len(), 1);
360        assert_eq!(top[0].0, b"open".to_vec(), "max 999");
361    }
362
363    #[test]
364    fn excluded_counted_and_merge() {
365        let mut s = seg();
366        s.apply(b"bad1", None, true);
367        s.apply(b"bad2", None, true);
368        assert_eq!(s.stats().excluded, 2);
369        assert!(s.contains(b"o1") && !s.contains(b"bad1"));
370        // cross-shard merge semantics
371        let mut a = s.group(b"paid");
372        let b = seg().group(b"paid");
373        merge_group(&mut a, &b);
374        assert_eq!((a.count, a.sum), (6, 900.0));
375        assert_eq!(a.min, Some(IndexValue::I64(100)));
376        assert_eq!(a.max, Some(IndexValue::I64(250)));
377        // merge with an empty partial keeps extremes
378        let mut e = GroupStats { count: 0, sum: 0.0, min: None, max: None };
379        merge_group(&mut e, &a);
380        assert_eq!(e.max, Some(IndexValue::I64(250)));
381    }
382
383    #[test]
384    fn stats_bytes_nonzero() {
385        let s = seg();
386        let st = s.stats();
387        assert_eq!((st.groups, st.rows), (2, 5));
388        assert!(st.approx_bytes > 0);
389    }
390}
kevy_index/agg.rs

kevy_index/
agg.rs
