Skip to main content

kevy_index/
agg_stats.rs

1//! Per-group statistics and the one ranking order every face shares.
2
3use crate::IndexValue;
4
5/// One group's live statistics.
6///
7/// ```
8/// use kevy_index::{GroupStats, IndexValue};
9/// let mut total = GroupStats::default();
10/// let mut part = GroupStats::default();
11/// (part.count, part.sum, part.min, part.max) = (2, 10.0, Some(IndexValue::I64(3)), Some(IndexValue::I64(7)));
12/// total.merge(&part);
13/// assert_eq!(total.avg(), Some(5.0));
14/// ```
15#[derive(Debug, Clone, Default, PartialEq)]
16#[non_exhaustive]
17pub struct GroupStats {
18    /// Rows in the group.
19    ///
20    /// ```
21    /// # use kevy_index::{AggRow, AggSegment, IndexValue};
22    /// # let mut s = AggSegment::new();
23    /// # let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
24    /// s.apply(b"o:1", row(b"eu", 30));
25    /// s.apply(b"o:2", row(b"eu", 12));
26    /// assert_eq!(s.group(b"eu").count, 2);
27    /// s.apply(b"o:2", AggRow::Removed);
28    /// assert_eq!(s.group(b"eu").count, 1);
29    /// ```
30    pub count: u64,
31    /// Sum of the aggregated field (f64 accumulation — the i64
32    /// overflow guard; precision bounds documented).
33    ///
34    /// ```
35    /// # use kevy_index::{AggRow, AggSegment, IndexValue};
36    /// # let mut s = AggSegment::new();
37    /// # let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
38    /// s.apply(b"o:1", row(b"eu", 30));
39    /// s.apply(b"o:2", row(b"eu", 12));
40    /// assert_eq!(s.group(b"eu").sum, 42.0);
41    /// ```
42    pub sum: f64,
43    /// Exact minimum (None only when count == 0).
44    ///
45    /// ```
46    /// # use kevy_index::{AggRow, AggSegment, IndexValue};
47    /// # let mut s = AggSegment::new();
48    /// # let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
49    /// s.apply(b"o:1", row(b"eu", 30));
50    /// s.apply(b"o:2", row(b"eu", 12));
51    /// assert_eq!(s.group(b"eu").min, Some(IndexValue::I64(12)));
52    /// // exact under deletion: the next-smallest value takes over
53    /// s.apply(b"o:2", AggRow::Removed);
54    /// assert_eq!(s.group(b"eu").min, Some(IndexValue::I64(30)));
55    /// assert_eq!(s.group(b"us").min, None, "an unknown group has no rows");
56    /// ```
57    pub min: Option<IndexValue>,
58    /// Exact maximum.
59    ///
60    /// ```
61    /// # use kevy_index::{AggRow, AggSegment, IndexValue};
62    /// # let mut s = AggSegment::new();
63    /// # let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
64    /// s.apply(b"o:1", row(b"eu", 30));
65    /// s.apply(b"o:2", row(b"eu", 12));
66    /// assert_eq!(s.group(b"eu").max, Some(IndexValue::I64(30)));
67    /// s.apply(b"o:1", row(b"eu", 5));
68    /// assert_eq!(s.group(b"eu").max, Some(IndexValue::I64(12)));
69    /// ```
70    pub max: Option<IndexValue>,
71}
72
73impl GroupStats {
74    /// Derived average.
75    ///
76    /// ```
77    /// assert_eq!(kevy_index::GroupStats::default().avg(), None);
78    /// ```
79    pub fn avg(&self) -> Option<f64> {
80        (self.count > 0).then(|| self.sum / self.count as f64)
81    }
82
83    /// This group's standing under `by`, oriented so a larger score ranks
84    /// higher for every metric — the order [`sort_groups`] puts groups in
85    /// (a smaller minimum ranks higher; a group without one ranks last).
86    ///
87    /// ```
88    /// use kevy_index::{AggBy, GroupStats, IndexValue};
89    /// let mut g = GroupStats::default();
90    /// (g.count, g.min) = (3, Some(IndexValue::I64(2)));
91    /// assert_eq!((g.rank_score(AggBy::Count), g.rank_score(AggBy::Min)), (3.0, -2.0));
92    /// assert_eq!(g.rank_score(AggBy::Max), f64::NEG_INFINITY);
93    /// ```
94    pub fn rank_score(&self, by: AggBy) -> f64 {
95        match by {
96            AggBy::Count => self.count as f64,
97            AggBy::Sum => self.sum,
98            AggBy::Max => self.max.as_ref().map_or(f64::NEG_INFINITY, IndexValue::as_f64),
99            AggBy::Min => self.min.as_ref().map_or(f64::NEG_INFINITY, |v| -v.as_f64()),
100        }
101    }
102
103    /// Fold one shard's partial for the same group into this one (reduce
104    /// side): counts and sums add, min and max take the extremes.
105    ///
106    /// ```
107    /// use kevy_index::{GroupStats, IndexValue};
108    /// let mut a = GroupStats::default();
109    /// a.count = 1;
110    /// a.min = Some(IndexValue::I64(5));
111    /// let mut b = GroupStats::default();
112    /// b.count = 1;
113    /// b.min = Some(IndexValue::I64(2));
114    /// a.merge(&b);
115    /// assert_eq!((a.count, a.min), (2, Some(IndexValue::I64(2))));
116    /// ```
117    pub fn merge(&mut self, part: &GroupStats) {
118        self.count += part.count;
119        self.sum += part.sum;
120        self.min = match (self.min.take(), part.min.clone()) {
121            (Some(a), Some(b)) => Some(if b < a { b } else { a }),
122            (a, b) => a.or(b),
123        };
124        self.max = match (self.max.take(), part.max.clone()) {
125            (Some(a), Some(b)) => Some(if b > a { b } else { a }),
126            (a, b) => a.or(b),
127        };
128    }
129}
130
131/// Ranking metric for [`AggSegment::top_groups`](crate::AggSegment::top_groups).
132///
133/// ```
134/// use kevy_index::AggBy;
135/// assert_eq!(AggBy::default(), AggBy::Count);
136/// assert_eq!(AggBy::parse(b"MAX"), Some(AggBy::Max), "tags are case-insensitive");
137/// assert_eq!(AggBy::parse(b"avg"), None);
138/// ```
139#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
140#[non_exhaustive]
141pub enum AggBy {
142    /// By row count (default).
143    ///
144    /// ```
145    /// # use kevy_index::{AggRow, AggSegment, IndexValue};
146    /// # let mut s = AggSegment::new();
147    /// # let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
148    /// s.apply(b"o:1", row(b"eu", 30));
149    /// s.apply(b"o:2", row(b"eu", 12));
150    /// s.apply(b"o:3", row(b"us", 99));
151    /// let top = s.top_groups(kevy_index::AggBy::Count, 1);
152    /// assert_eq!((top[0].0.as_slice(), top[0].1.count), (&b"eu"[..], 2));
153    /// ```
154    #[default]
155    Count,
156    /// By sum.
157    ///
158    /// ```
159    /// # use kevy_index::{AggRow, AggSegment, IndexValue};
160    /// # let mut s = AggSegment::new();
161    /// # let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
162    /// s.apply(b"o:1", row(b"eu", 30));
163    /// s.apply(b"o:2", row(b"eu", 12));
164    /// s.apply(b"o:3", row(b"us", 99));
165    /// let top = s.top_groups(kevy_index::AggBy::Sum, 1);
166    /// assert_eq!((top[0].0.as_slice(), top[0].1.sum), (&b"us"[..], 99.0));
167    /// ```
168    Sum,
169    /// By minimum (ascending — smallest mins first).
170    ///
171    /// ```
172    /// # use kevy_index::{AggRow, AggSegment, IndexValue};
173    /// # let mut s = AggSegment::new();
174    /// # let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
175    /// s.apply(b"o:1", row(b"eu", 30));
176    /// s.apply(b"o:2", row(b"eu", 12));
177    /// s.apply(b"o:3", row(b"us", 99));
178    /// let top = s.top_groups(kevy_index::AggBy::Min, 2);
179    /// assert_eq!(top[0].0, b"eu", "12 < 99");
180    /// ```
181    Min,
182    /// By maximum (descending — largest maxes first).
183    ///
184    /// ```
185    /// # use kevy_index::{AggRow, AggSegment, IndexValue};
186    /// # let mut s = AggSegment::new();
187    /// # let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
188    /// s.apply(b"o:1", row(b"eu", 30));
189    /// s.apply(b"o:2", row(b"eu", 12));
190    /// s.apply(b"o:3", row(b"us", 99));
191    /// let top = s.top_groups(kevy_index::AggBy::Max, 2);
192    /// assert_eq!(top[0].0, b"us", "99 > 30");
193    /// ```
194    Max,
195}
196
197impl AggBy {
198    /// The wire tag, as [`AggBy::parse`] reads it.
199    ///
200    /// ```
201    /// use kevy_index::AggBy;
202    /// assert_eq!(AggBy::parse(AggBy::Sum.tag().as_bytes()), Some(AggBy::Sum));
203    /// ```
204    pub fn tag(self) -> &'static str {
205        match self {
206            AggBy::Count => "count",
207            AggBy::Sum => "sum",
208            AggBy::Min => "min",
209            AggBy::Max => "max",
210        }
211    }
212
213    /// Wire tag.
214    pub fn parse(raw: &[u8]) -> Option<AggBy> {
215        if raw.eq_ignore_ascii_case(b"count") {
216            Some(AggBy::Count)
217        } else if raw.eq_ignore_ascii_case(b"sum") {
218            Some(AggBy::Sum)
219        } else if raw.eq_ignore_ascii_case(b"min") {
220            Some(AggBy::Min)
221        } else if raw.eq_ignore_ascii_case(b"max") {
222            Some(AggBy::Max)
223        } else {
224            None
225        }
226    }
227}
228
229/// Shared ranking order (per-shard AND at the reduce after merging
230/// shard partials — one definition, no drift).
231///
232/// ```
233/// use kevy_index::{AggBy, AggRow, AggSegment, IndexValue, sort_groups};
234/// let row = |g: &[u8], v| AggRow::Member { group: g.to_vec(), value: IndexValue::I64(v) };
235/// let (mut a, mut b) = (AggSegment::new(), AggSegment::new());
236/// a.apply(b"o:1", row(b"eu", 5));
237/// b.apply(b"o:2", row(b"us", 9));
238/// b.apply(b"o:3", row(b"us", 1));
239/// // reduce: concatenate the shards' unranked chunks, then rank once
240/// let mut all = a.all_groups();
241/// all.extend(b.all_groups());
242/// sort_groups(&mut all, AggBy::Sum);
243/// assert_eq!(all[0].0, b"us");
244/// sort_groups(&mut all, AggBy::Min);
245/// assert_eq!(all[0].0, b"us", "1 is the smallest minimum");
246/// ```
247pub fn sort_groups(all: &mut [(Vec<u8>, GroupStats)], by: AggBy) {
248    match by {
249        AggBy::Count => all.sort_by(|a, b| b.1.count.cmp(&a.1.count).then_with(|| a.0.cmp(&b.0))),
250        AggBy::Sum => all.sort_by(|a, b| b.1.sum.total_cmp(&a.1.sum).then_with(|| a.0.cmp(&b.0))),
251        AggBy::Min => all.sort_by(|a, b| {
252            match (&a.1.min, &b.1.min) {
253                (Some(x), Some(y)) => x.cmp(y),
254                (Some(_), None) => std::cmp::Ordering::Less,
255                (None, Some(_)) => std::cmp::Ordering::Greater,
256                (None, None) => std::cmp::Ordering::Equal,
257            }
258            .then_with(|| a.0.cmp(&b.0))
259        }),
260        AggBy::Max => all.sort_by(|a, b| {
261            match (&a.1.max, &b.1.max) {
262                (Some(x), Some(y)) => y.cmp(x),
263                (Some(_), None) => std::cmp::Ordering::Less,
264                (None, Some(_)) => std::cmp::Ordering::Greater,
265                (None, None) => std::cmp::Ordering::Equal,
266            }
267            .then_with(|| a.0.cmp(&b.0))
268        }),
269    }
270}