drive_proof_verifier/proof/document_split_average.rs
1//! Verified per-entry average result.
2//!
3//! Average-side analog of [`super::document_split_sum::DocumentSplitSums`].
4//! Holds one verified `(in_key, key, count, sum)` 4-tuple per matched
5//! group — the client divides each `sum / count` to compute per-group
6//! averages. Returned by `select=AVG, group_by=[...]` queries;
7//! aggregate averages use [`super::document_average::DocumentAverage`]
8//! instead.
9//!
10//! The generic `FromProof<Q>` impl below intentionally rejects
11//! calls (matching [`super::document_split_count::DocumentSplitCounts`]'s
12//! pattern). Real dispatch lives in the
13//! `FromProof<DocumentQuery>` impl in
14//! `rs-sdk/src/platform/documents/document_split_averages.rs`,
15//! which picks the right per-shape verifier (PCPS carrier-aggregate
16//! / primary-key direct read) based on the resolved
17//! `DocumentAverageMode`.
18
19/// A single verified `(in_key?, key, count, sum)` entry from an
20/// average query with `group_by`. Mirrors sum's `SplitSumEntry`
21/// shape with both metrics carried alongside.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct SplitAverageEntry {
24 /// Outer In-prefix value for compound `(In, range)` queries;
25 /// `None` for flat queries.
26 pub in_key: Option<Vec<u8>>,
27 /// The terminator key value.
28 pub key: Vec<u8>,
29 /// Matched-document count at that key. `Some(n)` for matched
30 /// keys; `None` for keys proven absent.
31 pub count: Option<u64>,
32 /// Aggregated sum at that key. `Some(n)` for matched keys;
33 /// `None` for keys proven absent.
34 pub sum: Option<i64>,
35}
36
37impl SplitAverageEntry {
38 /// Convenience: compute the average for this entry as `f64`, or
39 /// `None` if the entry was proven absent (`count`/`sum` is `None`)
40 /// or has zero count.
41 pub fn as_f64(&self) -> Option<f64> {
42 match (self.count, self.sum) {
43 (Some(c), Some(s)) if c > 0 => Some(s as f64 / c as f64),
44 _ => None,
45 }
46 }
47}
48
49/// The full per-entry average result, verified from proof.
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub struct DocumentSplitAverages(pub Vec<SplitAverageEntry>);
52
53impl DocumentSplitAverages {
54 /// Convenience: collapse compound `(in_key, key)` entries into a
55 /// flat `BTreeMap<key, (summed_count, summed_sum)>` by combining
56 /// each In-fork's contribution at the same terminator key.
57 /// Mirrors [`super::document_split_sum::DocumentSplitSums::try_into_flat_map`].
58 ///
59 /// Uses [`u64::checked_add`] on the count axis and
60 /// [`i64::checked_add`] on the sum axis; overflow on either
61 /// surfaces as [`Error::RequestError`](crate::Error::RequestError)
62 /// rather than panicking (debug) or wrapping (release).
63 /// Mirrors the SDK's aggregate-side
64 /// `DocumentAverage::fold_average_entries`, which hardens the
65 /// same two axes for the single-aggregate response shape.
66 ///
67 /// On overflow, drop back to iterating `DocumentSplitAverages.0`
68 /// directly — per-branch entries preserve the verified
69 /// `(u64, i64)` pair, and the caller can fold with its own
70 /// arithmetic.
71 pub fn try_into_flat_map(
72 self,
73 ) -> Result<std::collections::BTreeMap<Vec<u8>, (u64, i64)>, crate::Error> {
74 let mut out: std::collections::BTreeMap<Vec<u8>, (u64, i64)> =
75 std::collections::BTreeMap::new();
76 for entry in self.0 {
77 if let (Some(c), Some(s)) = (entry.count, entry.sum) {
78 let acc = out.entry(entry.key).or_insert((0u64, 0i64));
79 acc.0 = acc
80 .0
81 .checked_add(c)
82 .ok_or_else(|| crate::Error::RequestError {
83 error: "DocumentSplitAverages::try_into_flat_map: u64 overflow merging \
84 per-In-fork counts at the same terminator key. Iterate \
85 DocumentSplitAverages.0 directly to access per-branch counts and \
86 fold with your own arithmetic."
87 .to_string(),
88 })?;
89 acc.1 = acc
90 .1
91 .checked_add(s)
92 .ok_or_else(|| crate::Error::RequestError {
93 error: "DocumentSplitAverages::try_into_flat_map: i64 overflow merging \
94 per-In-fork sums at the same terminator key. Iterate \
95 DocumentSplitAverages.0 directly to access per-branch sums and fold \
96 with your own arithmetic (e.g. i128)."
97 .to_string(),
98 })?;
99 }
100 }
101 Ok(out)
102 }
103}
104
105#[cfg(test)]
106mod tests {
107 use super::*;
108
109 /// Single-fork pass-through: count + sum preserved unchanged.
110 #[test]
111 fn try_into_flat_map_passes_through_flat_entries() {
112 let splits = DocumentSplitAverages(vec![SplitAverageEntry {
113 in_key: None,
114 key: b"key-a".to_vec(),
115 count: Some(3),
116 sum: Some(42),
117 }]);
118 let flat = splits.try_into_flat_map().expect("non-overflowing fold");
119 assert_eq!(flat.get(b"key-a".as_slice()), Some(&(3u64, 42i64)));
120 }
121
122 /// Two In-forks at the same terminator key fold both axes.
123 #[test]
124 fn try_into_flat_map_sums_across_in_forks() {
125 let splits = DocumentSplitAverages(vec![
126 SplitAverageEntry {
127 in_key: Some(b"alice".to_vec()),
128 key: b"red".to_vec(),
129 count: Some(2),
130 sum: Some(10),
131 },
132 SplitAverageEntry {
133 in_key: Some(b"bob".to_vec()),
134 key: b"red".to_vec(),
135 count: Some(5),
136 sum: Some(32),
137 },
138 ]);
139 let flat = splits.try_into_flat_map().expect("non-overflowing fold");
140 assert_eq!(flat.get(b"red".as_slice()), Some(&(7u64, 42i64)));
141 }
142
143 /// `u64::MAX + 1` on count surfaces a `RequestError` (not
144 /// panic / wrap). Independent of the sum axis.
145 #[test]
146 fn try_into_flat_map_count_overflow_surfaces_request_error() {
147 let splits = DocumentSplitAverages(vec![
148 SplitAverageEntry {
149 in_key: Some(b"alice".to_vec()),
150 key: b"shared".to_vec(),
151 count: Some(u64::MAX),
152 sum: Some(0),
153 },
154 SplitAverageEntry {
155 in_key: Some(b"bob".to_vec()),
156 key: b"shared".to_vec(),
157 count: Some(1),
158 sum: Some(0),
159 },
160 ]);
161 let err = splits
162 .try_into_flat_map()
163 .expect_err("u64::MAX + 1 on the count axis must surface as overflow error");
164 let msg = format!("{err:?}");
165 assert!(
166 msg.contains("u64 overflow") && msg.contains("DocumentSplitAverages"),
167 "error must name the count-axis overflow and the helper: got {msg}"
168 );
169 }
170
171 /// `i64::MAX + 1` on sum surfaces a `RequestError` independently
172 /// of count.
173 #[test]
174 fn try_into_flat_map_sum_overflow_surfaces_request_error() {
175 let splits = DocumentSplitAverages(vec![
176 SplitAverageEntry {
177 in_key: Some(b"alice".to_vec()),
178 key: b"shared".to_vec(),
179 count: Some(0),
180 sum: Some(i64::MAX),
181 },
182 SplitAverageEntry {
183 in_key: Some(b"bob".to_vec()),
184 key: b"shared".to_vec(),
185 count: Some(0),
186 sum: Some(1),
187 },
188 ]);
189 let err = splits
190 .try_into_flat_map()
191 .expect_err("i64::MAX + 1 on the sum axis must surface as overflow error");
192 let msg = format!("{err:?}");
193 assert!(
194 msg.contains("i64 overflow") && msg.contains("DocumentSplitAverages"),
195 "error must name the sum-axis overflow and the helper: got {msg}"
196 );
197 }
198
199 /// An entry where EITHER half is `None` (proven absent) is
200 /// skipped entirely — must not poison the accumulator.
201 #[test]
202 fn try_into_flat_map_skips_partial_or_absent_entries() {
203 let splits = DocumentSplitAverages(vec![
204 SplitAverageEntry {
205 in_key: None,
206 key: b"present".to_vec(),
207 count: Some(2),
208 sum: Some(5),
209 },
210 SplitAverageEntry {
211 in_key: None,
212 key: b"absent".to_vec(),
213 count: None,
214 sum: Some(99),
215 },
216 SplitAverageEntry {
217 in_key: None,
218 key: b"absent".to_vec(),
219 count: Some(99),
220 sum: None,
221 },
222 ]);
223 let flat = splits.try_into_flat_map().expect("partial entries skipped");
224 assert_eq!(flat.get(b"present".as_slice()), Some(&(2u64, 5i64)));
225 assert!(!flat.contains_key(b"absent".as_slice()));
226 }
227}
228
229// No generic `FromProof<Q>` impl — callers reach this through
230// `FromProof<DocumentQuery> for DocumentSplitAverages` in
231// `rs-sdk/src/platform/documents/document_split_averages.rs`.