nodedb_query/msgpack_scan/
aggregate_helpers.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Public helpers for streaming aggregate accumulators.
4//!
5//! These thin wrappers expose field-extraction primitives used by the
6//! streaming aggregate accumulators in the `nodedb` crate.
7//! Each function operates on a single raw MessagePack document byte slice and
8//! returns only the scalar value needed by the calling accumulator — no
9//! document bytes are retained after the call returns.
10
11use nodedb_types::Value;
12
13use crate::expr::{EvalError, SqlExpr};
14use crate::msgpack_scan::field::extract_field;
15use crate::msgpack_scan::reader::{read_f64, read_numeric, read_str, read_value};
16use crate::value_ops;
17
18// ── Expression evaluator ───────────────────────────────────────────────────
19
20/// Evaluate `expr` against a decoded document.
21///
22/// The `Result<Option<Value>, EvalError>` shape distinguishes the two ways a
23/// per-row aggregate argument can decline to contribute a value:
24/// - `Ok(None)` — the row is *skipped* (document could not be decoded, or the
25///   expression legitimately yields no value). Matches SQL aggregate semantics
26///   where a NULL/absent argument is excluded from the accumulation.
27/// - `Err(EvalError::DivisionByZero)` — the expression divided or took a
28///   modulus by zero. This is a statement failure that
29///   surfaces as SQLSTATE `22012`, exactly like a WHERE/projection expression,
30///   rather than silently dropping the row from the aggregate.
31#[inline]
32fn eval_expr(doc: &[u8], expr: &SqlExpr) -> Result<Option<Value>, EvalError> {
33    let Ok(doc_val) = nodedb_types::json_msgpack::value_from_msgpack(doc) else {
34        return Ok(None);
35    };
36    Ok(Some(expr.eval(&doc_val)?))
37}
38
39// ── Public extraction helpers ──────────────────────────────────────────────
40
41/// Extract a numeric (f64) value from `field`, or evaluate `expr` if provided.
42/// Returns `Ok(None)` when the field is absent or cannot be converted to f64,
43/// and `Err(EvalError::DivisionByZero)` when `expr` divides/mods by zero.
44#[inline]
45pub fn extract_f64(
46    doc: &[u8],
47    field: &str,
48    expr: Option<&SqlExpr>,
49) -> Result<Option<f64>, EvalError> {
50    if let Some(expr) = expr {
51        return Ok(eval_expr(doc, expr)?.and_then(|v| value_ops::value_to_f64(&v, false)));
52    }
53    let Some((start, _end)) = extract_field(doc, 0, field) else {
54        return Ok(None);
55    };
56    Ok(read_f64(doc, start))
57}
58
59/// Extract the SUM / AVG input of `field`, or of `expr` if provided, as the
60/// number it contributes (`Integer`, `Float`, or `Decimal`). A raw field
61/// and an expression result contribute by the same rule,
62/// [`crate::numeric_sum::sum_input`]: a msgpack number read exactly, or a
63/// numeric string read per [`crate::numeric_sum::numeric_text`] (a
64/// `DECIMAL` cell is stored as its text).
65/// Returns `Ok(None)` when nothing contributes, and
66/// `Err(EvalError::DivisionByZero)` when `expr` divides/mods by zero.
67#[inline]
68pub fn extract_sum_value(
69    doc: &[u8],
70    field: &str,
71    expr: Option<&SqlExpr>,
72) -> Result<Option<Value>, EvalError> {
73    if let Some(expr) = expr {
74        return Ok(eval_expr(doc, expr)?.and_then(|v| crate::numeric_sum::sum_input(&v)));
75    }
76    let Some((start, _end)) = extract_field(doc, 0, field) else {
77        return Ok(None);
78    };
79    Ok(match read_numeric(doc, start) {
80        Some(n) => Some(crate::numeric_sum::numeric_to_value(n)),
81        None => read_str(doc, start).and_then(crate::numeric_sum::numeric_text),
82    })
83}
84
85/// Extract a display string from `field`, or evaluate `expr` if provided.
86/// Returns `Ok(None)` when the field is absent.
87pub fn extract_str(
88    doc: &[u8],
89    field: &str,
90    expr: Option<&SqlExpr>,
91) -> Result<Option<String>, EvalError> {
92    if let Some(expr) = expr {
93        return Ok(eval_expr(doc, expr)?.map(|v| value_ops::value_to_display_string(&v)));
94    }
95    let Some((start, _end)) = extract_field(doc, 0, field) else {
96        return Ok(None);
97    };
98    Ok(read_str(doc, start).map(|s| s.to_string()))
99}
100
101/// Extract a field as `Value`.  Uses direct msgpack→Value for scalars;
102/// falls back to full document decode only for complex types.
103pub fn extract_value(
104    doc: &[u8],
105    field: &str,
106    expr: Option<&SqlExpr>,
107) -> Result<Option<Value>, EvalError> {
108    if let Some(expr) = expr {
109        return eval_expr(doc, expr);
110    }
111    let Some((start, end)) = extract_field(doc, 0, field) else {
112        return Ok(None);
113    };
114    if let Some(v) = read_value(doc, start) {
115        return Ok(Some(v));
116    }
117    let field_bytes = &doc[start..end];
118    Ok(nodedb_types::json_msgpack::value_from_msgpack(field_bytes).ok())
119}
120
121/// Extract a field or expression result as raw msgpack bytes.
122/// Used by `count_distinct`, `approx_count_distinct`, `approx_topk`, etc.
123pub fn extract_bytes(
124    doc: &[u8],
125    field: &str,
126    expr: Option<&SqlExpr>,
127) -> Result<Option<Vec<u8>>, EvalError> {
128    if let Some(expr) = expr {
129        let Some(val) = eval_expr(doc, expr)? else {
130            return Ok(None);
131        };
132        return Ok(nodedb_types::json_msgpack::value_to_msgpack(&val).ok());
133    }
134    let Some((start, end)) = extract_field(doc, 0, field) else {
135        return Ok(None);
136    };
137    Ok(Some(doc[start..end].to_vec()))
138}
139
140/// Returns `Ok(Some(()))` when the field is present and non-null.
141/// Used by `count(field)` accumulator to count non-null values.
142#[inline]
143pub fn extract_non_null(
144    doc: &[u8],
145    field: &str,
146    expr: Option<&SqlExpr>,
147) -> Result<Option<()>, EvalError> {
148    let Some(v) = extract_value(doc, field, expr)? else {
149        return Ok(None);
150    };
151    Ok(if v.is_null() { None } else { Some(()) })
152}
nodedb_query/msgpack_scan/aggregate_helpers.rs

nodedb_query/msgpack_scan/
aggregate_helpers.rs
