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

kevy_index/
agg.rs
