kevy_index/
agg.rs

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

kevy_index/
agg.rs
