1use super::{
4 aggregate_budget, attach_collects_with_budget, json_value_cmp, normalize, reconstruct,
5 related_specs_from_where, working_set_budget, AggregateRequest, Arc, Budget, Cap, GroupBucket,
6 Json, Lmpg, Map, OutputExpr, PointFields, QueryError, QuerySnapshot, SearchRequest,
7};
8use crate::aggregate_catalog::dictionary_grouped_count;
9use crate::aggregate_rows::{row_bytes, Partial, RowCtx, RowPlan};
10use crate::aggregate_scan::{self, Mode, Rows, Scan};
11use crate::page::serialized_bytes;
12
13pub fn aggregate(lmpg: &Lmpg, req: &AggregateRequest) -> Result<Vec<GroupBucket>, QueryError> {
22 aggregate_with_cancel(lmpg, req, &|| false).map(|(buckets, _)| buckets)
23}
24
25pub const MAX_AGGREGATE_BUCKETS: usize = 1_000;
28
29pub fn aggregate_with_cancel(
46 lmpg: &Lmpg,
47 req: &AggregateRequest,
48 is_cancelled: &dyn Fn() -> bool,
49) -> Result<(Vec<GroupBucket>, u64), QueryError> {
50 let snap = Arc::new(lmpg.store.read_view()?);
51 aggregate_with_cancel_in_view(lmpg, req, is_cancelled, snap)
52}
53
54pub fn aggregate_in_view(
61 lmpg: &Lmpg,
62 req: &AggregateRequest,
63 snap: QuerySnapshot,
64) -> Result<(Vec<GroupBucket>, u64), QueryError> {
65 aggregate_with_cancel_in_view(lmpg, req, &|| false, snap)
66}
67
68pub fn aggregate_with_cancel_in_view(
81 lmpg: &Lmpg,
82 req: &AggregateRequest,
83 is_cancelled: &dyn Fn() -> bool,
84 snap: QuerySnapshot,
85) -> Result<(Vec<GroupBucket>, u64), QueryError> {
86 let is_cancelled: &dyn Fn() -> bool = &crate::cancel_probe(lmpg, is_cancelled);
87 if !snap.belongs_to(&lmpg.store) {
88 return Err(rushdb_storage::StorageError::InvalidRead(
89 "aggregate view belongs to another Store".into(),
90 )
91 .into());
92 }
93 if is_cancelled() {
94 return Err(QueryError::Cancelled);
95 }
96 let budget = working_set_budget(lmpg);
97 Budget::WorkingSet.check(budget, std::mem::size_of::<GroupBucket>())?;
98 aggregate_budgeted(lmpg, req, is_cancelled, snap, budget)
99 .map_err(aggregate_budget::resource_error)
100}
101
102#[expect(
103 clippy::too_many_lines,
104 clippy::cast_precision_loss,
105 clippy::needless_pass_by_value,
106 reason = "aggregate accumulation keeps one live budget and pinned snapshot"
107)]
108fn aggregate_budgeted(
109 lmpg: &Lmpg,
110 req: &AggregateRequest,
111 is_cancelled: &dyn Fn() -> bool,
112 snap: QuerySnapshot,
113 budget: Option<usize>,
114) -> Result<(Vec<GroupBucket>, u64), QueryError> {
115 let page_req = SearchRequest {
116 labels: req.labels.clone(),
117 where_: req.where_.clone(),
118 related: req.related.clone(),
119 ..SearchRequest::default()
120 };
121 if is_count_only(req) {
127 if let Some(total) = catalog_count_total(lmpg, &page_req, &snap)? {
128 let mut buckets = count_only_bucket(total);
129 apply_aggregate_pagination(&mut buckets, req);
130 return Ok((buckets, total));
131 }
132 }
133 if let Some(buckets) = dictionary_grouped_count(lmpg, req, &snap, is_cancelled)? {
136 Budget::WorkingSet.check(
137 budget,
138 buckets.iter().fold(0usize, |sum, b| {
139 sum.saturating_add(aggregate_budget::bucket_bytes(b))
140 }),
141 )?;
142 let total = buckets.iter().map(|b| b.count_roots).sum();
143 return Ok((order_and_page(buckets, req), total));
144 }
145 let plan = RowPlan::new(req, related_specs_from_where(req.where_.as_ref()));
150 let points = PointFields::resolve_in(
151 &snap,
152 lmpg.tenant,
153 &req.labels,
154 plan.names.iter().map(String::as_str),
155 )?;
156 let selected = Rows::selection(&plan);
157 let rows = Rows {
158 points: points.as_ref(),
159 selected: &selected,
160 };
161 let scan = Scan {
162 lmpg,
163 search: &page_req,
164 snap: &snap,
165 budget,
166 plan: &plan,
167 };
168 let mode = if is_count_only(req) {
169 Mode::Count
170 } else if plan.per_row() {
171 Mode::Roots
172 } else {
173 Mode::Groups
174 };
175 let pass = aggregate_scan::run(&scan, &rows, mode, is_cancelled)?;
176 let total_matched = pass.total;
177 let candidate_bytes = pass.candidate_bytes;
178 let mut buckets = match mode {
179 Mode::Count => {
180 let mut buckets = count_only_bucket(total_matched);
181 apply_aggregate_pagination(&mut buckets, req);
182 return Ok((buckets, total_matched));
183 }
184 Mode::Roots => per_root_rows(&scan, &rows, &pass, req, is_cancelled)?,
185 Mode::Groups => pass.groups.finish(&plan),
186 };
187 let retained_bytes = buckets.iter().fold(candidate_bytes, |sum, b| {
188 sum.saturating_add(aggregate_budget::bucket_bytes(b))
189 });
190 Budget::WorkingSet.check(budget, retained_bytes)?;
191
192 if !req.outputs.is_empty() || !req.avg_terms.is_empty() || !req.expr_avgs.is_empty() {
197 let output_bytes = req
198 .outputs
199 .iter()
200 .map(|(name, _)| name)
201 .chain(req.avg_terms.iter().map(|(name, _)| name))
202 .chain(req.expr_avgs.iter().map(|(name, _)| name))
203 .fold(0usize, |sum, name| {
204 sum.saturating_add(name.len()).saturating_add(256)
205 })
206 .saturating_mul(buckets.len());
207 Budget::WorkingSet.check(
208 budget,
209 retained_bytes
210 .saturating_sub(candidate_bytes)
211 .saturating_add(output_bytes),
212 )?;
213 for b in &mut buckets {
214 if is_cancelled() {
215 return Err(QueryError::Cancelled);
216 }
217 let mut resolved: std::collections::BTreeMap<String, f64> = b
218 .sums
219 .clone()
220 .into_iter()
221 .filter_map(|(k, v)| v.as_f64().map(|x| (k, x)))
222 .collect();
223 resolved.insert("count_roots".into(), b.count_roots as f64);
224 for (out, field) in &req.counts {
229 if req.hidden_outputs.contains(out) {
230 continue;
231 }
232 let value = if field.as_str() == "*" {
233 b.count_roots as f64
234 } else {
235 b.counts
236 .get(out)
237 .and_then(serde_json::Value::as_u64)
238 .map_or(0.0, |n| n as f64)
239 };
240 resolved.insert(out.clone(), value);
241 }
242 for name in req
244 .avg_terms
245 .iter()
246 .map(|(n, _)| n)
247 .chain(req.expr_avgs.iter().map(|(n, _)| n))
248 {
249 let sum = resolved.get(name).copied().unwrap_or(0.0);
250 let n = b
251 .counts
252 .get(&format!("__avg_n_{name}"))
253 .and_then(serde_json::Value::as_u64)
254 .unwrap_or(0);
255 let avg = if n > 0 { sum / n as f64 } else { 0.0 };
256 let avg = match req.avg_precision.get(name) {
257 Some(&0) => avg.trunc(),
259 Some(&p) => {
260 let p = p.min(crate::MAX_AVG_PRECISION);
263 let m = 10f64.powi(i32::try_from(p).unwrap_or(0));
264 (avg * m).round() / m
265 }
266 None => avg,
267 };
268 resolved.insert(name.clone(), avg);
269 b.sums.insert(name.clone(), serde_json::json!(avg));
270 }
271 for (name, expr) in &req.outputs {
272 eval_output(&normalize(name), expr, &mut resolved)?;
273 }
274 for (k, v) in &resolved {
275 b.derived_outputs.insert(k.clone(), serde_json::json!(v));
276 b.sums
278 .entry(k.clone())
279 .or_insert_with(|| serde_json::json!(v));
280 }
281 }
282 }
283
284 Ok((order_and_page(buckets, req), total_matched))
285}
286
287fn per_root_rows(
292 scan: &Scan<'_>,
293 rows: &Rows<'_>,
294 pass: &aggregate_scan::Pass,
295 req: &AggregateRequest,
296 is_cancelled: &dyn Fn() -> bool,
297) -> Result<Vec<GroupBucket>, QueryError> {
298 let roots = if req.order_by.is_empty() {
299 let page_end = req.skip.saturating_add(req.limit.unwrap_or(u64::MAX));
300 &pass.roots[..usize::try_from(page_end)
301 .unwrap_or(usize::MAX)
302 .min(pass.roots.len())]
303 } else {
304 &pass.roots[..]
305 };
306 let (snap, tenant, budget) = (scan.snap, scan.lmpg.tenant, scan.budget);
307 let mut session = reconstruct::ReconstructSession::new(snap, tenant);
308 let mut partial = Partial::default();
309 let mut row = Vec::new();
310 let mut retained = pass
311 .candidate_bytes
312 .saturating_add(pass.roots.len().saturating_mul(8));
313 for &ord in roots {
314 if is_cancelled() {
315 return Err(QueryError::Cancelled);
316 }
317 if !rows.read(scan.plan, snap, tenant, &mut session, ord, &mut row)? {
318 continue;
319 }
320 let held = retained.saturating_add(row_bytes(&row));
321 Budget::WorkingSet.check(budget, held)?;
322 let ctx = RowCtx {
323 lmpg: scan.lmpg,
324 snap,
325 ord,
326 alias_budget: budget.map(|limit| {
327 limit.saturating_sub(held) / scan.plan.metric_count().saturating_add(1)
328 }),
329 };
330 retained = retained.saturating_add(partial.accumulate(scan.plan, &row, &ctx)?);
331 Budget::WorkingSet.check(budget, retained)?;
332 }
333 Ok(partial.finish(scan.plan))
334}
335
336fn order_and_page(mut buckets: Vec<GroupBucket>, req: &AggregateRequest) -> Vec<GroupBucket> {
339 for (name, desc) in req.order_by.iter().rev() {
340 buckets.sort_by(|a, b| {
341 let va = bucket_metric(a, name, req);
342 let vb = bucket_metric(b, name, req);
343 let ord = compare_json_metric(va.as_ref(), vb.as_ref());
344 if *desc {
345 ord.reverse()
346 } else {
347 ord
348 }
349 });
350 }
351 if buckets.is_empty() && is_global_rollup(req) {
352 buckets.push(GroupBucket::default());
354 }
355 apply_aggregate_pagination(&mut buckets, req);
356 buckets
357}
358
359fn fill_absent_metrics(row: &mut Map<String, Json>, req: &AggregateRequest) {
362 if !is_global_rollup(req) {
363 return;
364 }
365 let metrics = req
366 .sums
367 .iter()
368 .chain(&req.avg_terms)
369 .map(|(out, _)| out)
370 .chain(req.expr_avgs.iter().map(|(out, _)| out))
371 .chain(req.mins.iter().map(|(out, _)| out))
372 .chain(req.maxs.iter().map(|(out, _)| out))
373 .chain(req.outputs.iter().map(|(out, _)| out));
374 for out in metrics {
375 if !req.hidden_outputs.contains(out) {
376 row.entry(out.clone()).or_insert(Json::Null);
377 }
378 }
379}
380
381fn is_global_rollup(req: &AggregateRequest) -> bool {
385 let metrics = !req.counts.is_empty()
386 || !req.sums.is_empty()
387 || !req.mins.is_empty()
388 || !req.maxs.is_empty()
389 || !req.outputs.is_empty()
390 || !req.expr_sums.is_empty()
391 || !req.expr_avgs.is_empty()
392 || !req.avg_terms.is_empty();
393 metrics
394 && req.group_by.is_empty()
395 && req.time_bucket.is_none()
396 && req.plain_fields.is_empty()
397 && req.collects.is_empty()
398}
399
400fn is_count_only(req: &AggregateRequest) -> bool {
404 req.group_by.is_empty()
405 && req.plain_fields.is_empty()
406 && req.collects.is_empty()
407 && req.sums.is_empty()
408 && req.mins.is_empty()
409 && req.maxs.is_empty()
410 && req.outputs.is_empty()
411 && req.expr_sums.is_empty()
412 && req.expr_avgs.is_empty()
413 && req.avg_terms.is_empty()
414 && req.time_bucket.is_none()
415 && !req.counts.is_empty()
416 && req.counts.iter().all(|(_, field)| field == "*")
417}
418
419fn count_only_bucket(total: u64) -> Vec<GroupBucket> {
422 vec![GroupBucket {
423 count_roots: total,
424 ..GroupBucket::default()
425 }]
426}
427
428fn apply_aggregate_pagination(buckets: &mut Vec<GroupBucket>, req: &AggregateRequest) {
429 if req.skip > 0 || req.limit.is_some() {
430 let start = usize::try_from(req.skip)
431 .unwrap_or(usize::MAX)
432 .min(buckets.len());
433 *buckets = buckets[start..]
434 .iter()
435 .take(usize::try_from(req.limit.unwrap_or(u64::MAX)).unwrap_or(usize::MAX))
436 .cloned()
437 .collect();
438 }
439}
440
441fn catalog_count_total(
445 lmpg: &Lmpg,
446 req: &SearchRequest,
447 snap: &QuerySnapshot,
448) -> Result<Option<u64>, QueryError> {
449 if req.where_.is_some() || !req.related.is_empty() || !req.order_by.is_empty() {
450 return Ok(None);
451 }
452 if req.labels.len() != 1 {
453 return Ok(None);
454 }
455 Ok(Some(rushdb_lmpg::derived::label_count(
457 snap,
458 lmpg.tenant,
459 &req.labels[0],
460 )?))
461}
462
463pub fn aggregate_to_rows(
472 lmpg: &Lmpg,
473 buckets: &[GroupBucket],
474 req: &AggregateRequest,
475) -> Result<Vec<Json>, QueryError> {
476 let snap: QuerySnapshot = Arc::new(lmpg.store.read_view()?);
477 aggregate_to_rows_in_view(lmpg, &snap, buckets, req)
478}
479
480pub fn aggregate_to_rows_in_view(
489 lmpg: &Lmpg,
490 snap: &QuerySnapshot,
491 buckets: &[GroupBucket],
492 req: &AggregateRequest,
493) -> Result<Vec<Json>, QueryError> {
494 aggregate_to_rows_in_view_with_budget(lmpg, snap, buckets, req, None)
495}
496
497pub fn aggregate_to_rows_in_view_with_budget(
508 lmpg: &Lmpg,
509 snap: &QuerySnapshot,
510 buckets: &[GroupBucket],
511 req: &AggregateRequest,
512 budget: Option<usize>,
513) -> Result<Vec<Json>, QueryError> {
514 let cap = Cap::for_response(budget, lmpg.limits().query_memory_max_bytes());
515 let snap = cap.map_or_else(
516 || Arc::clone(snap),
517 |cap| snap.with_materialization_limit(cap.limit() / 16),
518 );
519 let mut charged = 0usize;
520 let mut rows = Vec::with_capacity(buckets.len());
521 for b in buckets {
522 let mut row = b.keys.clone();
523 for (out, field) in &req.counts {
524 if field == "*" {
525 row.insert(out.clone(), Json::Number((b.count_roots).into()));
526 } else {
527 let n = map_metric(&b.counts, out).and_then(serde_json::Value::as_u64);
528 row.insert(out.clone(), Json::Number(n.unwrap_or(0).into()));
529 }
530 }
531 for (out, _field) in &req.sums {
532 if req.hidden_outputs.contains(out) {
533 continue;
534 }
535 if let Some(v) = b.sums.get(out) {
536 row.insert(out.clone(), v.clone());
537 }
538 }
539 for out in req
540 .avg_terms
541 .iter()
542 .map(|(out, _)| out)
543 .chain(req.expr_avgs.iter().map(|(out, _)| out))
544 {
545 if !req.hidden_outputs.contains(out) {
546 if let Some(v) = b.sums.get(out) {
547 row.insert(out.clone(), v.clone());
548 }
549 }
550 }
551 for (out, _) in &req.mins {
552 if !req.hidden_outputs.contains(out) {
553 if let Some(v) = b.mins.get(out) {
554 row.insert(out.clone(), v.clone());
555 }
556 }
557 }
558 for (out, _) in &req.maxs {
559 if !req.hidden_outputs.contains(out) {
560 if let Some(v) = b.maxs.get(out) {
561 row.insert(out.clone(), v.clone());
562 }
563 }
564 }
565 for (out, _) in &req.outputs {
566 if let Some(v) = b.derived_outputs.get(&normalize(out)) {
567 row.insert(out.clone(), v.clone());
568 }
569 }
570 fill_absent_metrics(&mut row, req);
571 let mut row_json = Json::Object(row);
574 if let Some(ord) = b.root.filter(|_| !req.collects.is_empty()) {
575 row_json = attach_collects_with_budget(
576 lmpg,
577 &snap,
578 ord,
579 row_json,
580 &req.collects,
581 cap.map(|cap| cap.after(charged)),
582 &|| false,
583 )
584 .map_err(aggregate_budget::resource_error)?;
585 }
586 if cap.is_some() {
587 charged = charged.saturating_add(serialized_bytes(&row_json));
588 Cap::check_opt(cap, charged)?;
589 }
590 rows.push(row_json);
591 }
592 Ok(rows)
593}
594
595fn map_metric<'a>(map: &'a Map<String, Json>, name: &str) -> Option<&'a Json> {
596 let normalized = normalize(name);
597 map.iter()
598 .find(|(key, _)| normalize(key) == normalized)
599 .map(|(_, value)| value)
600}
601
602fn bucket_metric(b: &GroupBucket, name: &str, req: &AggregateRequest) -> Option<Json> {
603 if normalize(name) == "count_roots" {
604 return Some(Json::Number(b.count_roots.into()));
605 }
606 if let Some((output, field)) = req
607 .counts
608 .iter()
609 .find(|(output, _)| normalize(output) == normalize(name))
610 {
611 if field == "*" {
612 return Some(Json::Number(b.count_roots.into()));
613 }
614 return map_metric(&b.counts, output)
615 .map(|values| Json::Number(values.as_u64().unwrap_or(0).into()));
616 }
617 map_metric(&b.sums, name)
618 .or_else(|| map_metric(&b.mins, name))
619 .or_else(|| map_metric(&b.maxs, name))
620 .or_else(|| map_metric(&b.keys, name))
621 .cloned()
622}
623
624fn compare_json_metric(a: Option<&Json>, b: Option<&Json>) -> std::cmp::Ordering {
625 match (a, b) {
626 (None, None) => std::cmp::Ordering::Equal,
627 (None, _) => std::cmp::Ordering::Less,
628 (_, None) => std::cmp::Ordering::Greater,
629 (Some(x), Some(y)) => json_value_cmp(x, y),
630 }
631}
632
633#[expect(
636 clippy::cast_precision_loss,
637 clippy::cast_possible_truncation,
638 reason = "the integer round-trip check rejects lossy widening"
639)]
640pub(crate) fn exact_f64(value: &Json) -> Option<f64> {
641 let number = value.as_number()?;
642 if let Some(integer) = number.as_i64() {
643 let float = integer as f64;
644 return ((float as i128) == i128::from(integer)).then_some(float);
645 }
646 if let Some(integer) = number.as_u64() {
647 let float = integer as f64;
648 return ((float as i128) == i128::from(integer)).then_some(float);
649 }
650 number.as_f64()
651}
652
653fn eval_output(
654 name: &str,
655 expr: &OutputExpr,
656 results: &mut std::collections::BTreeMap<String, f64>,
657) -> Result<f64, QueryError> {
658 let v = match expr {
659 OutputExpr::Literal(l) => *l,
660 OutputExpr::Field(_) => {
661 return Ok(0.0);
664 }
665 OutputExpr::Ref(other) => results
666 .get(other)
667 .copied()
668 .ok_or_else(|| QueryError::InvalidPredicate(format!("unresolved ref `{other}`")))?,
669 OutputExpr::Add(a, b) => eval_output(name, a, results)? + eval_output(name, b, results)?,
670 OutputExpr::Subtract(a, b) => {
671 eval_output(name, a, results)? - eval_output(name, b, results)?
672 }
673 OutputExpr::Multiply(a, b) => {
674 eval_output(name, a, results)? * eval_output(name, b, results)?
675 }
676 OutputExpr::Divide {
677 dividend, divisor, ..
678 } => {
679 let d = eval_output(name, dividend, results)?;
680 let q = eval_output(name, divisor, results)?;
681 if q == 0.0 {
682 return Err(QueryError::InvalidPredicate("divide by zero".into()));
683 }
684 d / q
685 }
686 };
687 results.insert(name.to_string(), v);
688 Ok(v)
689}