spg_engine/aggregate.rs
1//! Aggregate executor.
2//!
3//! Handles `SELECT … <aggs> … [GROUP BY …]` queries. The planning strategy
4//! is straightforward:
5//!
6//! 1. Walk the SELECT (and ORDER BY) expressions to find every aggregate
7//! function call. Dedupe by AST equality and assign each `__agg_<i>`.
8//! 2. Same for every `GROUP BY` expression: assign `__grp_<j>`.
9//! 3. Stream the WHERE-filtered rows, group by the tuple of GROUP BY
10//! values, and update per-group aggregate state.
11//! 4. Materialise a synthetic per-group row containing
12//! `[__grp_0..__grp_K, __agg_0..__agg_N]` and rewrite the user's
13//! SELECT / ORDER BY expressions to reference those synthetic columns
14//! instead of the originals.
15//! 5. Evaluate the rewritten expressions against the synthetic schema and
16//! emit results.
17//!
18//! v1.8 implements `count(*)`, `count(expr)`, `sum`, `min`, `max`, `avg`.
19//! NULL semantics follow PG: aggregates skip NULL inputs (except
20//! `count(*)`, which counts rows). `sum(int)` widens to `BigInt`;
21//! `avg(int|bigint)` returns `Float`.
22
23use alloc::borrow::Cow;
24use alloc::boxed::Box;
25use alloc::collections::BTreeSet;
26use alloc::format;
27use alloc::string::{String, ToString};
28use alloc::vec::Vec;
29
30use spg_sql::ast::{Expr, SelectItem, SelectStatement};
31use spg_storage::{ColumnSchema, DataType, Row, Value};
32
33use crate::eval::{self, EvalContext, EvalError};
34use crate::join::RowRef;
35
36/// True if this statement should go through the aggregate path.
37pub fn uses_aggregate(stmt: &SelectStatement) -> bool {
38 if stmt.group_by.is_some() || stmt.having.is_some() {
39 return true;
40 }
41 for item in &stmt.items {
42 if let SelectItem::Expr { expr, .. } = item
43 && contains_aggregate(expr)
44 {
45 return true;
46 }
47 }
48 for o in &stmt.order_by {
49 if contains_aggregate(&o.expr) {
50 return true;
51 }
52 }
53 if let Some(h) = &stmt.having
54 && contains_aggregate(h)
55 {
56 return true;
57 }
58 false
59}
60
61pub fn contains_aggregate(e: &Expr) -> bool {
62 match e {
63 Expr::FunctionCall { name, args } => {
64 is_aggregate_name(name) || args.iter().any(contains_aggregate)
65 }
66 Expr::AggregateOrdered { .. } => true,
67 Expr::Binary { lhs, rhs, .. } => contains_aggregate(lhs) || contains_aggregate(rhs),
68 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
69 contains_aggregate(expr)
70 }
71 Expr::Like { expr, pattern, .. } => contains_aggregate(expr) || contains_aggregate(pattern),
72 Expr::Extract { source, .. } => contains_aggregate(source),
73 // v4.10 subqueries + v4.12 window functions / Literal /
74 // Column — all non-aggregate leaves from the regular
75 // aggregate planner's POV. Window-bearing projections are
76 // routed to exec_select_with_window before this runs.
77 Expr::ScalarSubquery(_)
78 | Expr::Exists { .. }
79 | Expr::InSubquery { .. }
80 | Expr::WindowFunction { .. }
81 | Expr::Literal(_)
82 | Expr::Placeholder(_)
83 | Expr::Column(_) => false,
84 // v7.10.10 — recurse into array constructor / subscript /
85 // ANY/ALL children. Aggregates inside `ARRAY[SUM(x)]` are
86 // valid PG and must be detected here.
87 Expr::Array(items) => items.iter().any(contains_aggregate),
88 Expr::ArraySubscript { target, index } => {
89 contains_aggregate(target) || contains_aggregate(index)
90 }
91 Expr::AnyAll { expr, array, .. } => contains_aggregate(expr) || contains_aggregate(array),
92 Expr::InList { expr, list, .. } => {
93 contains_aggregate(expr) || list.iter().any(contains_aggregate)
94 }
95 // v7.13.0 — CASE WHEN … END. Recurse into operand,
96 // every (WHEN, THEN) pair, and the ELSE branch.
97 Expr::Case {
98 operand,
99 branches,
100 else_branch,
101 } => {
102 operand.as_deref().is_some_and(contains_aggregate)
103 || branches
104 .iter()
105 .any(|(w, t)| contains_aggregate(w) || contains_aggregate(t))
106 || else_branch.as_deref().is_some_and(contains_aggregate)
107 }
108 }
109}
110
111pub fn is_aggregate_name(name: &str) -> bool {
112 matches!(
113 name.to_ascii_lowercase().as_str(),
114 "count"
115 | "count_star"
116 | "sum"
117 | "min"
118 | "max"
119 | "avg"
120 // v7.17.0 — variadic / collection aggregates. ORM
121 // reports (Hibernate / Rails / Django) emit these in
122 // GROUP BY rollups; pre-7.17 SPG hit "unknown
123 // aggregate".
124 | "string_agg"
125 | "array_agg"
126 // v7.17.0 — boolean aggregates. `every` is SQL-standard
127 // alias for `bool_and`.
128 | "bool_and"
129 | "bool_or"
130 | "every"
131 // v7.32 (round-29) — statistical aggregates (every BI /
132 // dashboard emits these in rollups).
133 | "stddev" | "stddev_samp" | "stddev_pop"
134 | "variance" | "var_samp" | "var_pop"
135 // v7.32 (round-29) — bitwise aggregates.
136 | "bit_and" | "bit_or" | "bit_xor"
137 // v7.32 (round-29) — ordered-set aggregates (used with
138 // `WITHIN GROUP (ORDER BY …)`).
139 | "percentile_cont" | "percentile_disc" | "mode"
140 // v7.32 (round-29) — hypothetical-set aggregates (also
141 // `WITHIN GROUP`): the rank the direct args WOULD have.
142 | "rank" | "dense_rank" | "percent_rank" | "cume_dist"
143 // v7.32 (round-29) — two-argument regression family.
144 | "covar_pop" | "covar_samp" | "corr"
145 | "regr_count" | "regr_avgx" | "regr_avgy" | "regr_slope"
146 | "regr_intercept" | "regr_r2" | "regr_sxx" | "regr_syy" | "regr_sxy"
147 // v7.32 (round-29) — JSON aggregates.
148 | "json_agg" | "jsonb_agg" | "json_object_agg" | "jsonb_object_agg"
149 )
150}
151
152/// v7.32 (round-29) — two-argument regression aggregates `f(Y, X)`.
153fn is_regression_name(name: &str) -> bool {
154 matches!(
155 name,
156 "covar_pop"
157 | "covar_samp"
158 | "corr"
159 | "regr_count"
160 | "regr_avgx"
161 | "regr_avgy"
162 | "regr_slope"
163 | "regr_intercept"
164 | "regr_r2"
165 | "regr_sxx"
166 | "regr_syy"
167 | "regr_sxy"
168 )
169}
170
171/// v7.32 (round-29) — aggregates that consume a second positional
172/// argument: `string_agg(v, sep)`, the regression family `f(Y, X)`, and
173/// `json_object_agg(key, value)`.
174fn agg_uses_second_arg(name: &str) -> bool {
175 name == "string_agg"
176 || name == "json_object_agg"
177 || name == "jsonb_object_agg"
178 || is_regression_name(name)
179}
180
181/// v7.32 (round-29) — ordered-set aggregates: the value to aggregate
182/// comes from the `WITHIN GROUP (ORDER BY …)` sort spec, and any
183/// in-parens arguments are *direct* arguments (the percentile fraction).
184/// `mode()` takes no direct argument.
185pub fn is_ordered_set_name(name: &str) -> bool {
186 // v7.32 — `eq_ignore_ascii_case` instead of `to_ascii_lowercase()`:
187 // these classifiers run in the aggregate row/group loop, where the
188 // old per-call `String` allocation showed up as ~16% of the inbox's
189 // aggregate path in a sampled profile (the names are constant).
190 ["percentile_cont", "percentile_disc", "mode"]
191 .iter()
192 .any(|k| name.eq_ignore_ascii_case(k))
193}
194
195/// v7.32 (round-29) — hypothetical-set aggregates: `rank(args) WITHIN
196/// GROUP (ORDER BY …)` and friends compute the rank the hypothetical
197/// row would have. Like ordered-set, the value stream comes from the
198/// sort spec and the in-parens args are direct (the hypothetical row).
199pub fn is_hypothetical_set_name(name: &str) -> bool {
200 ["rank", "dense_rank", "percent_rank", "cume_dist"]
201 .iter()
202 .any(|k| name.eq_ignore_ascii_case(k))
203}
204
205/// v7.32 (round-29) — every aggregate that takes its value stream from
206/// a `WITHIN GROUP (ORDER BY …)` clause (ordered-set + hypothetical-set).
207pub fn is_within_group_name(name: &str) -> bool {
208 is_ordered_set_name(name) || is_hypothetical_set_name(name)
209}
210
211/// v7.37.4 (R34) — pre-computed aggregate kind. Replaces per-row
212/// string matches in `update_state` with a single `match` on a
213/// `Copy` enum (compiles to a jump table). For the mailrs prod
214/// `/api/conversations` shape (14 aggregates × 100 k rows = 1.4 M
215/// inner-loop iterations) this is the dominant per-row cost.
216///
217/// Lowered from `AggSpec::name` at spec build time via
218/// [`classify_agg_name`]; populated by the three `AggSpec`
219/// construction sites (window+ORDER, plain, `first_ordered`
220/// `array_agg`).
221#[derive(Copy, Clone, Debug, PartialEq, Eq)]
222enum AggKind {
223 CountStar,
224 Count,
225 Sum,
226 Avg,
227 Min,
228 Max,
229 StringAgg,
230 ArrayAgg,
231 BoolAnd,
232 BoolOr,
233 /// stddev / stddev_samp / stddev_pop / variance / var_samp / var_pop.
234 StddevFamily,
235 BitAnd,
236 BitOr,
237 BitXor,
238 /// ordered-set (`percentile_cont/disc`, `mode`) +
239 /// hypothetical-set (`rank`/`dense_rank`/etc.) aggregates that
240 /// share the WITHIN-GROUP collection path.
241 WithinGroup,
242 /// covar_samp / covar_pop / corr / regr_*.
243 Regression,
244 JsonAgg,
245 JsonObjectAgg,
246}
247
248/// v7.37.4 (R34) — name → kind, called once per spec at build time.
249/// Hot path (`update_state_kind`) only sees the enum; the canonical
250/// string still travels with the spec so `finalize` and errors can
251/// quote it.
252fn classify_agg_name(name: &str) -> AggKind {
253 match name {
254 "count_star" => AggKind::CountStar,
255 "count" => AggKind::Count,
256 "sum" => AggKind::Sum,
257 "avg" => AggKind::Avg,
258 "min" => AggKind::Min,
259 "max" => AggKind::Max,
260 "string_agg" => AggKind::StringAgg,
261 "array_agg" => AggKind::ArrayAgg,
262 "bool_and" => AggKind::BoolAnd,
263 "bool_or" => AggKind::BoolOr,
264 "stddev" | "stddev_samp" | "stddev_pop" | "variance" | "var_samp" | "var_pop" => {
265 AggKind::StddevFamily
266 }
267 "bit_and" => AggKind::BitAnd,
268 "bit_or" => AggKind::BitOr,
269 "bit_xor" => AggKind::BitXor,
270 "json_agg" | "jsonb_agg" => AggKind::JsonAgg,
271 "json_object_agg" | "jsonb_object_agg" => AggKind::JsonObjectAgg,
272 n if is_within_group_name(n) => AggKind::WithinGroup,
273 n if is_regression_name(n) => AggKind::Regression,
274 other => panic!("classify_agg_name: unknown aggregate {other}"),
275 }
276}
277
278/// Per-aggregate running state.
279#[derive(Debug, Default, Clone)]
280struct AggState {
281 count: i64,
282 sum_int: i64,
283 sum_float: f64,
284 extreme: Option<Value<'static>>,
285 use_float: bool,
286 /// v7.17.0 — running collection for string_agg / array_agg.
287 /// Each entry is one row's contribution (NULL preserved as
288 /// `Value::Null`; string_agg's finalize step drops them, but
289 /// array_agg keeps them). Pushing in insertion order matches
290 /// PG behaviour when no `ORDER BY` is given inside the
291 /// aggregate call.
292 items: Vec<Value<'static>>,
293 /// v7.25 (round-17) — per-group dedupe set for DISTINCT
294 /// aggregates (encoded values; NULLs never reach it because
295 /// the caller's skip runs after the per-aggregate NULL rules).
296 /// v7.37.4 measured `hashbrown::HashSet` as worse at this
297 /// shape — the per-(group × distinct-spec) hash table alloc
298 /// overhead beats the lookup-speed gain when each set is
299 /// small. Sticking with `BTreeSet`; the dispatch-side enum
300 /// fix in `update_state` is the R34 win.
301 seen: BTreeSet<String>,
302 /// v7.37.x (docker-fair DISTA attack) — fast-path BigInt seen
303 /// set. The hot DISTINCT path used `encode_key_refs_into` to
304 /// turn `Value::BigInt(n)` into a string key like `"I<n>|"` then
305 /// inserted that into the String BTreeSet — ~100 ns of pure alloc
306 /// + format churn per row × 25 k rows × 1 BigInt DISTINCT spec
307 /// (the DISTA `COUNT(DISTINCT m.id)` shape) ≈ 2.5 ms of waste.
308 /// Direct `BTreeSet<i64>` skips encode entirely; lookups stay
309 /// O(log small) on the per-group set. Lazy-allocated — only the
310 /// BigInt-DISTINCT path constructs it.
311 seen_int: Option<BTreeSet<i64>>,
312 /// v7.24 (round-16 A) — per-item ORDER BY key tuples, parallel
313 /// to `items` (pushed under the same skip/keep conditions).
314 /// Empty when the aggregate carries no internal ordering.
315 item_keys: Vec<Vec<Value<'static>>>,
316 /// v7.17.0 — captured separator for string_agg. PG accepts a
317 /// non-constant separator expression but in practice every
318 /// caller passes a literal; the engine snapshots the last
319 /// non-NULL text it sees, which matches PG's "use the latest
320 /// row's value" behaviour.
321 separator: Option<String>,
322 /// v7.17.0 — running boolean accumulator for bool_and /
323 /// bool_or / every. `None` until the first non-NULL input;
324 /// at finalize None → SQL NULL.
325 bool_acc: Option<bool>,
326 /// v7.32 (round-29) — sum of squares for the variance / stddev
327 /// family (`sum_float` carries the running sum; `count` the n).
328 sum_sq: f64,
329 /// v7.32 (round-29) — running accumulator for bit_and / bit_or /
330 /// bit_xor. `None` until the first non-NULL input → SQL NULL.
331 bit_acc: Option<i64>,
332 /// v7.32 (round-29) — two-argument regression family
333 /// (`covar_*` / `corr` / `regr_*`), PG arg order `f(Y, X)`. Only
334 /// rows where BOTH inputs are non-NULL contribute (`count` is the
335 /// paired n, independent of the single-arg `sum_*`).
336 reg_n: i64,
337 reg_sx: f64,
338 reg_sy: f64,
339 reg_sxx: f64,
340 reg_syy: f64,
341 reg_sxy: f64,
342 /// v7.32 (round-29) — second value stream for `json_object_agg`
343 /// (`items` holds the keys, `aux_items` the values).
344 aux_items: Vec<Value<'static>>,
345 /// v7.33 (array_agg argmax) — for a `first_ordered` spec
346 /// (`(array_agg(x ORDER BY y))[1]`), the running first-by-order
347 /// (sort-key tuple, value). Replaced only when a new row's key sorts
348 /// strictly before the current best (ties keep the earliest row, =
349 /// the stable-sort `[1]`). No items/item_keys array is built.
350 first_best: Option<(Vec<Value<'static>>, Value<'static>)>,
351}
352
353#[derive(Debug, Clone)]
354struct AggSpec {
355 name: String, // lowercased
356 /// First argument (value expression) for every aggregate
357 /// except `count(*)`. `None` for `count_star`.
358 arg: Option<Expr>,
359 /// v7.17.0 — second argument. Only `string_agg(value, sep)`
360 /// uses it today. `None` for every other aggregate (or for
361 /// `array_agg`, which is single-arg). Carried in the spec so
362 /// per-row evaluation can re-use the same separator
363 /// expression across calls.
364 arg2: Option<Expr>,
365 /// v7.25 (round-17) — `COUNT(DISTINCT x)` & friends: dedupe
366 /// the input stream per group before accumulation.
367 distinct: bool,
368 /// v7.24 (round-16 A) — aggregate-internal ORDER BY keys
369 /// (`array_agg(x ORDER BY y DESC NULLS LAST)`). Empty for the
370 /// plain form. Only the collection aggregates honour it;
371 /// other aggregates are order-insensitive and ignore it (PG
372 /// accepts the syntax everywhere too).
373 order_by: Vec<spg_sql::ast::OrderBy>,
374 /// v7.32 (round-29) — `FILTER (WHERE cond)`: a per-row predicate
375 /// evaluated against the source row before accumulation. A row
376 /// whose `cond` is not TRUE (false or NULL) is excluded from this
377 /// aggregate only. `None` for the unfiltered form.
378 filter: Option<Expr>,
379 /// v7.32 (round-29) — ordered-set aggregates only: the *direct*
380 /// argument (the percentile fraction for `percentile_cont/disc`).
381 /// PG requires it constant, so it is evaluated once. `None` for
382 /// `mode()` and for every non-ordered-set aggregate.
383 direct_arg: Option<Expr>,
384 /// v7.33 (array_agg argmax) — set when this spec came from
385 /// `(array_agg(x ORDER BY y))[1]`: accumulate only the first-by-order
386 /// element (a running argmax/argmin) and finalise to that scalar
387 /// value, instead of collecting + sorting + materialising the whole
388 /// per-group array just to take element 1. Returns the element type,
389 /// not the array type.
390 first_ordered: bool,
391 /// v7.37.4 (R34) — derived from `name` at spec build time so the
392 /// per-row inner loop dispatches via a `match` on `Copy` enum
393 /// instead of a string compare for every (row × aggregate)
394 /// iteration.
395 kind: AggKind,
396}
397
398/// Output of running the aggregate path. Schema describes one row per
399/// group; rows are not yet ORDER BY-sorted (caller does it).
400#[derive(Debug)]
401pub struct AggResult {
402 pub columns: Vec<ColumnSchema>,
403 pub rows: Vec<Row<'static>>,
404 /// v7.31 (perf — PG lesson #1, post-LIMIT subquery projection):
405 /// select-list items whose rewritten expr carries a subquery and
406 /// is referenced by neither ORDER BY nor HAVING. Their output
407 /// cells hold NULL placeholders; the caller truncates to
408 /// LIMIT+OFFSET first and only then evaluates these for the
409 /// surviving rows (PG runs the same shape with SubPlan loops=50
410 /// instead of loops=24000). `(output_col, rewritten_expr)`.
411 pub deferred: Vec<(usize, Expr)>,
412 /// Synthetic group rows aligned 1:1 with `rows`; populated only
413 /// when `deferred` is non-empty.
414 pub synth_rows: Vec<Row<'static>>,
415 /// Schema the deferred exprs evaluate against.
416 pub synth_schema: Vec<ColumnSchema>,
417}
418
419/// Execute aggregate logic against an already-WHERE-filtered iterator of
420/// rows. `table_alias` is the alias accepted by column resolution.
421#[allow(clippy::too_many_lines)]
422/// v7.25.2 (round-19 A) — caller-injected evaluator for synth-row
423/// expressions that still carry subquery nodes after the rewrite
424/// (correlated subqueries in the select list / HAVING / aggregate
425/// ORDER BY of a GROUP BY query). The engine passes its
426/// correlated-aware evaluator; pure-library callers pass None and
427/// surviving subqueries keep erroring loudly.
428pub type CorrelatedEval<'a> =
429 &'a dyn Fn(&Expr, &Row<'static>, &EvalContext<'_>) -> Result<Value<'static>, EvalError>;
430
431/// Output of the per-group projection stage (`project_groups`): the
432/// output schema, the projected rows, the synth rows kept alongside
433/// them for post-LIMIT deferred evaluation, the deferred subquery
434/// items, and the rewritten ORDER BY exprs (shared with the sort).
435struct Projection {
436 columns: Vec<ColumnSchema>,
437 out_rows: Vec<Row<'static>>,
438 kept_synth: Vec<Row<'static>>,
439 deferred: Vec<(usize, Expr)>,
440 order_rewritten: Vec<Expr>,
441 /// v7.37.x — when `defer_projection` is requested, `out_rows`
442 /// carries empty placeholders and the caller runs the per-item
443 /// eval pass after sort+truncate over the surviving ≤ keep_n
444 /// rows. `None` when projection was performed inline.
445 deferred_project: Option<DeferredProject>,
446}
447
448struct DeferredProject {
449 items_rewritten: Vec<Option<Expr>>,
450 items_compiled: Vec<Option<eval::CompiledExpr>>,
451}
452
453/// v7.35.0 — detect the `SELECT COUNT(*) FROM … [WHERE …]` shape
454/// (single item, no GROUP BY / HAVING / ORDER BY / DISTINCT /
455/// LIMIT WITH TIES / FILTER / window). For this shape the answer
456/// is exactly `rows.len()` as `BigInt`, no group state needed.
457/// Returns `None` for any deviation so the caller's full pipeline
458/// runs verbatim.
459///
460/// v7.35.2 — also short-circuit `COUNT(<literal>)` (e.g.
461/// `COUNT(1)`) and `COUNT(<column>)` when the column is declared
462/// NOT NULL on the input schema. PG handles both cases as
463/// `COUNT(*)` (the non-null filter is a no-op), so doing the same
464/// here keeps every `count this thing` shape on the same fast path
465/// instead of routing the literal / non-null-col variants through
466/// the four-stage aggregate pipeline.
467fn try_pure_count_star_short_circuit(
468 stmt: &SelectStatement,
469 rows: &[RowRef<'_>],
470 schema_cols: &[ColumnSchema],
471 table_alias: Option<&str>,
472) -> Option<AggResult> {
473 if stmt.distinct
474 || stmt.limit_with_ties
475 || stmt.group_by.is_some()
476 || stmt.having.is_some()
477 || !stmt.order_by.is_empty()
478 {
479 return None;
480 }
481 if stmt.items.len() != 1 {
482 return None;
483 }
484 let SelectItem::Expr { expr, alias } = &stmt.items[0] else {
485 return None;
486 };
487 let Expr::FunctionCall { name, args } = expr else {
488 return None;
489 };
490 if !name.eq_ignore_ascii_case("count") && !name.eq_ignore_ascii_case("count_star") {
491 return None;
492 }
493 let count_star_shape = match args.as_slice() {
494 // `COUNT(*)` parses to `count_star` with no args.
495 [] if name.eq_ignore_ascii_case("count_star") => true,
496 // `COUNT(<literal>)` — the per-row test is "is this literal
497 // non-null?" which is constant, so it's COUNT(*) when the
498 // literal is non-null.
499 [Expr::Literal(lit)] => !matches!(lit, spg_sql::ast::Literal::Null),
500 // `COUNT(<column>)` — same answer as COUNT(*) when the
501 // column is statically declared NOT NULL on the input
502 // schema. Resolve through the alias if one is set.
503 [Expr::Column(c)] => {
504 if let Some(q) = c.qualifier.as_deref()
505 && let Some(alias) = table_alias
506 && !q.eq_ignore_ascii_case(alias)
507 {
508 return None;
509 }
510 schema_cols
511 .iter()
512 .find(|s| s.name.eq_ignore_ascii_case(&c.name))
513 .is_some_and(|s| !s.nullable)
514 }
515 _ => return None,
516 };
517 if !count_star_shape {
518 return None;
519 }
520 let col_name = alias.clone().unwrap_or_else(|| "count".to_string());
521 let count = i64::try_from(rows.len()).unwrap_or(i64::MAX);
522 Some(AggResult {
523 columns: alloc::vec![ColumnSchema::new(col_name, DataType::BigInt, false)],
524 rows: alloc::vec![Row::new(alloc::vec![Value::BigInt(count)])],
525 deferred: Vec::new(),
526 synth_rows: Vec::new(),
527 synth_schema: Vec::new(),
528 })
529}
530
531pub(crate) fn run(
532 stmt: &SelectStatement,
533 rows: &[RowRef<'_>],
534 schema_cols: &[ColumnSchema],
535 table_alias: Option<&str>,
536 correlated_eval: Option<CorrelatedEval<'_>>,
537) -> Result<AggResult, EvalError> {
538 // v7.35.0 — pure `SELECT COUNT(*) FROM … WHERE …` short-circuit.
539 // The caller already filtered rows by WHERE (we run on the
540 // post-WHERE survivor set), so for the canonical pure-COUNT(*)
541 // shape (no GROUP BY / HAVING / ORDER BY / DISTINCT / FILTER /
542 // window) the answer is simply `rows.len()`. The four-stage
543 // aggregate pipeline below (accumulate_groups → build_synth_schema
544 // → finalize_synth_rows → project_groups) collapses to a single
545 // BigInt cell when there's a single group, but each stage still
546 // pays its own allocation tax — group state map, synth schema
547 // vec, finalize loop. `exists_in_60` (mailrs prod #4 baseline)
548 // is exactly this shape on a 25 k-row JOIN.
549 if let Some(short) = try_pure_count_star_short_circuit(stmt, rows, schema_cols, table_alias) {
550 return Ok(short);
551 }
552 let group_exprs: Vec<Expr> = stmt.group_by.clone().unwrap_or_default();
553
554 // Collect aggregate sub-expressions across items + order_by.
555 let mut agg_specs: Vec<AggSpec> = Vec::new();
556 for item in &stmt.items {
557 if let SelectItem::Expr { expr, .. } = item {
558 collect_aggregates(expr, &mut agg_specs);
559 }
560 }
561 for o in &stmt.order_by {
562 collect_aggregates(&o.expr, &mut agg_specs);
563 }
564 if let Some(h) = &stmt.having {
565 collect_aggregates(h, &mut agg_specs);
566 }
567 // v7.17.0 — arity validation. The collector tolerates an
568 // arbitrary positional-arg count; here we enforce the
569 // per-aggregate contract so a malformed call (e.g.
570 // `array_agg()` or `string_agg(x)`) surfaces as a SQL error
571 // rather than silently coercing to a degenerate aggregate.
572 validate_agg_arities(stmt, &agg_specs)?;
573 validate_within_group(&agg_specs)?;
574
575 // (1) Stream the WHERE-filtered rows into insertion-ordered group state.
576 let order = accumulate_groups(
577 rows,
578 &group_exprs,
579 &agg_specs,
580 schema_cols,
581 table_alias,
582 correlated_eval,
583 )?;
584
585 // (2) Build the synthetic per-group schema and finalise each group's row.
586 let synth_schema =
587 build_synth_schema(rows, &group_exprs, &agg_specs, schema_cols, table_alias)?;
588 let synth_rows = finalize_synth_rows(
589 &order,
590 &agg_specs,
591 &synth_schema,
592 rows,
593 schema_cols,
594 table_alias,
595 )?;
596
597 // v7.37.x (mailrs Track A 100k attack) — defer the bound
598 // per-item SELECT projection on the synth rows until AFTER
599 // sort + LIMIT truncation. On a `GROUP BY t ORDER BY agg DESC
600 // LIMIT 50` with 20 000 groups (the mailrs minimal 100k shape)
601 // pre-defer ran 20 000 × N_items compiled-VM evals + Row
602 // allocations before discarding 99.75 % at the sort truncation
603 // step. HAVING still runs inline on every group because it
604 // filters BEFORE the LIMIT; we only skip the SELECT-list eval.
605 let defer_projection = !stmt.order_by.is_empty()
606 && !stmt.distinct
607 && !stmt.limit_with_ties
608 && stmt.having.is_none()
609 && stmt.limit_literal().is_some_and(|l| {
610 let off = stmt.offset_literal().unwrap_or(0) as usize;
611 let k = (l as usize).saturating_add(off);
612 k > 0 && k < synth_rows.len()
613 });
614
615 // (3) Rewrite the user's expressions, filter groups by HAVING and project.
616 let Projection {
617 columns,
618 mut out_rows,
619 mut kept_synth,
620 deferred,
621 order_rewritten,
622 deferred_project,
623 } = project_groups(
624 synth_rows,
625 stmt,
626 &group_exprs,
627 &agg_specs,
628 &synth_schema,
629 correlated_eval,
630 defer_projection,
631 )?;
632
633 // (4) ORDER BY on the aggregated output (the caller applies LIMIT).
634 //
635 // v7.37.3 (mailrs prod /api/contacts 3.21× regression — and the
636 // general inbox-listing-shape SPG-vs-PG gap) — top-K sink for
637 // `ORDER BY <agg> [DESC] LIMIT k`. Pre-7.37.3 this stage ran a
638 // full O(N log N) sort over every surviving group, then the
639 // caller truncated to `k`. With high-cardinality GROUP BY (a
640 // sender column with hundreds-thousands of distinct values) the
641 // truncated set is a tiny fraction of `N` — keep an O(k) top-K
642 // sink and never sort the discarded majority. Matches PG /
643 // MySQL / MariaDB's standard "LIMIT k under ORDER BY agg"
644 // optimisation; SPG previously implemented it only on the
645 // streamed inner-join path (`try_streamed_inner_join_topn`)
646 // and not on the aggregate output.
647 //
648 // Gate: needs a literal LIMIT (placeholder LIMIT we can't bound
649 // statically here), no DISTINCT (would need post-dedup, can't
650 // truncate during sort), no LIMIT WITH TIES (which extends past
651 // the literal k by run-time tie-key comparison).
652 let keep_n: Option<usize> =
653 if !stmt.order_by.is_empty() && !stmt.distinct && !stmt.limit_with_ties {
654 stmt.limit_literal().map(|l| {
655 let off = stmt.offset_literal().unwrap_or(0) as usize;
656 (l as usize).saturating_add(off)
657 })
658 } else {
659 None
660 };
661 if !stmt.order_by.is_empty() {
662 let (sorted_synth, sorted_out) = sort_synth_by_order_by(
663 &synth_schema,
664 &stmt.order_by,
665 &order_rewritten,
666 kept_synth,
667 out_rows,
668 correlated_eval,
669 keep_n,
670 )?;
671 kept_synth = sorted_synth;
672 out_rows = sorted_out;
673 }
674
675 // v7.37.x — run deferred SELECT-list projection on the truncated
676 // top-K survivors. For `GROUP BY thread_id ORDER BY MAX(date) DESC
677 // LIMIT 50` against 20 000 groups, this turns ~40 000 compiled-VM
678 // evals + Row allocations into 100, saving ~2-3 ms on the mailrs
679 // minimal 100k shape.
680 if let Some(DeferredProject {
681 items_rewritten,
682 items_compiled,
683 }) = deferred_project
684 {
685 let synth_ctx = EvalContext::new(&synth_schema, None);
686 let mut stack: Vec<Value<'static>> = Vec::new();
687 for (idx, srow) in kept_synth.iter().enumerate() {
688 let mut values: Vec<Value<'static>> = Vec::with_capacity(columns.len());
689 for (i, rewritten) in items_rewritten.iter().enumerate() {
690 let Some(rewritten) = rewritten else { continue };
691 if deferred.iter().any(|(c, _)| *c == i) {
692 values.push(Value::Null);
693 continue;
694 }
695 values.push(if let Some(cc) = &items_compiled[i] {
696 eval::eval_compiled(cc, srow, &synth_ctx, &mut stack)?
697 } else {
698 match correlated_eval {
699 Some(f) if crate::expr_has_subquery(rewritten) => {
700 f(rewritten, srow, &synth_ctx)?
701 }
702 _ => eval::eval_expr(rewritten, srow, &synth_ctx)?,
703 }
704 });
705 }
706 out_rows[idx] = Row::new(values);
707 }
708 }
709
710 let (synth_rows_out, synth_schema_out) = if deferred.is_empty() {
711 (Vec::new(), Vec::new())
712 } else {
713 (kept_synth, synth_schema.clone())
714 };
715 Ok(AggResult {
716 columns,
717 rows: out_rows,
718 deferred,
719 synth_rows: synth_rows_out,
720 synth_schema: synth_schema_out,
721 })
722}
723
724/// v7.32 (round-29) — validate the structural requirements of WITHIN
725/// GROUP (ordered-set / hypothetical-set) aggregates up front, so a
726/// malformed call surfaces as a SQL error rather than a silently
727/// degenerate aggregate.
728fn validate_within_group(agg_specs: &[AggSpec]) -> Result<(), EvalError> {
729 // v7.32 (round-29) — WITHIN GROUP aggregates require the clause (PG
730 // raises a hard error otherwise rather than silently degrading), and
731 // SPG supports the single-sort-key form only.
732 for spec in agg_specs {
733 if is_within_group_name(&spec.name) {
734 if spec.order_by.is_empty() {
735 return Err(EvalError::TypeMismatch {
736 detail: format!("{}() requires WITHIN GROUP (ORDER BY …)", spec.name),
737 });
738 }
739 // mode() is the only WITHIN GROUP aggregate with no direct
740 // argument; the rest carry one (percentile fraction /
741 // hypothetical value).
742 if spec.name != "mode" && spec.direct_arg.is_none() {
743 return Err(EvalError::TypeMismatch {
744 detail: format!("{}() requires a direct argument", spec.name),
745 });
746 }
747 // Multi-key WITHIN GROUP (multiple sort keys / hypothetical
748 // args) is not supported yet — error loudly instead of
749 // silently using only the first key.
750 if spec.order_by.len() > 1 {
751 return Err(EvalError::TypeMismatch {
752 detail: format!(
753 "{}() with multiple WITHIN GROUP sort keys is not supported yet",
754 spec.name
755 ),
756 });
757 }
758 }
759 }
760 Ok(())
761}
762
763/// (1) Stream the WHERE-filtered rows, group by the GROUP BY value
764/// tuple, and update per-group aggregate state. Returns the groups in
765/// insertion order. See `run` for the bind-once fast path rationale.
766#[allow(clippy::too_many_lines, clippy::type_complexity)]
767fn accumulate_groups(
768 rows: &[RowRef<'_>],
769 group_exprs: &[Expr],
770 agg_specs: &[AggSpec],
771 schema_cols: &[ColumnSchema],
772 table_alias: Option<&str>,
773 correlated_eval: Option<CorrelatedEval<'_>>,
774) -> Result<Vec<(Vec<Value<'static>>, Vec<AggState>)>, EvalError> {
775 let ctx = EvalContext::new(schema_cols, table_alias);
776 // Map group key (vec of values, encoded as canonical string) -> group state.
777 // v7.32 (architecture v2, P2b) — insertion-ordered group state in
778 // a Vec; the hash map only maps key → index. Removes the parallel
779 // `key_order: Vec<String>` (a second per-group key clone) and the
780 // per-group re-probe `groups[k]` at finalize (24k hash lookups for
781 // the inbox shape). The map owns its key once on vacant insert.
782 let mut order: Vec<(Vec<Value<'static>>, Vec<AggState>)> = Vec::new();
783 let mut groups: hashbrown::HashMap<String, usize> = hashbrown::HashMap::new();
784 // v7.37.x (mailrs Track A perf — SPGE ≫ PG18) — single-Text GROUP
785 // BY column fast path. The canonical-string encode (`S<text>|`)
786 // + `encode_key_refs_into` reuse-buffer churn dominated the 30 k-
787 // row mailrs minimal probe (~3-4 ms / 30 k). For `GROUP BY t` on
788 // a TEXT column (the inbox-listing / conversation-grouping shape)
789 // the column text IS the canonical key — no encoder, no prefix
790 // byte, no `refs` Vec rebuild per row. The fallback `groups` map
791 // above is retained for multi-col / non-Text / collation paths;
792 // this map only fires when the schema and value structurally
793 // permit it. `null_group_idx` collects NULL group rows (SQL groups
794 // all NULLs into one bucket).
795 let mut groups_text: hashbrown::HashMap<String, usize> = hashbrown::HashMap::new();
796 let mut null_group_idx: Option<usize> = None;
797 // When there are no GROUP BY exprs *and* there is at least one aggregate,
798 // every row collapses into a single anonymous group keyed by "".
799 if rows.is_empty() && group_exprs.is_empty() {
800 // Single empty-aggregate group: count=0, sum=0, max=NULL, etc.
801 // No rows follow, so the map is never probed — seed `order` only.
802 let init: Vec<AggState> = (0..agg_specs.len()).map(|_| AggState::default()).collect();
803 order.push((Vec::new(), init));
804 }
805
806 // v7.30 (perf campaign) - hoist the per-row work that doesn't
807 // depend on the row: which group exprs need collation folding
808 // (none, for most queries - the old code cloned the whole
809 // group_vals vec per row just in case).
810 // v7.30 (perf campaign) - the no-tax row loop. When a group
811 // expr or an aggregate argument is a bare column reference
812 // (the overwhelmingly common shape), bind its position ONCE
813 // and read row cells by offset in the loop - no per-row tree
814 // walk, no owned-Value clone out of resolve_column. Anything
815 // more complex keeps the eval path.
816 let col_pos = |e: &Expr| -> Option<usize> {
817 // Qualified references only: the bare-name resolver carries
818 // alias/ambiguity logic the bind-once path must not fork.
819 if let Expr::Column(c) = e
820 && c.qualifier.is_some()
821 {
822 eval::find_column_pos(c, &ctx)
823 } else {
824 None
825 }
826 };
827 let group_pos: Vec<Option<usize>> = group_exprs.iter().map(col_pos).collect();
828 let all_groups_bound = group_pos.iter().all(Option::is_some);
829 // v7.37.x — single-col GROUP BY on a TEXT-typed column lets the
830 // hot loop key the hash map by the column text directly. Resolved
831 // once from the bound position against `schema_cols`.
832 let single_text_group_col: bool = group_pos.len() == 1
833 && group_pos[0].is_some_and(|p| {
834 schema_cols
835 .get(p)
836 .is_some_and(|c| matches!(c.ty, spg_storage::DataType::Text))
837 });
838 let arg_pos: Vec<Option<usize>> = agg_specs
839 .iter()
840 .map(|spec| spec.arg.as_ref().and_then(|e| col_pos(e)))
841 .collect();
842 // v7.37.x (mailrs Track A 100k attack) — dedicated tight loop
843 // for the "single-Text GROUP BY + single MAX(bound numeric arg)"
844 // shape. This is the mailrs `/api/conversations` minimal shape
845 // (`GROUP BY thread_id, MAX(internal_date)`) and an inbox-listing
846 // staple across the SPG customer set. Skipping the per-row spec
847 // loop, FILTER / arg2 / order_keys checks, and the union-typed
848 // `update_state` enum jump saves ~80-100 ns/row at 100 k input
849 // — the gap closing the SPGE vs PG18 ratio at this scale.
850 let dedicated_max_loop: bool = single_text_group_col
851 && agg_specs.len() == 1
852 && matches!(agg_specs[0].kind, AggKind::Max)
853 && agg_specs[0].filter.is_none()
854 && agg_specs[0].arg2.is_none()
855 && agg_specs[0].order_by.is_empty()
856 && !agg_specs[0].distinct
857 && !agg_specs[0].first_ordered
858 && arg_pos[0].is_some();
859 // v7.36 (perf — mailrs Ask 1 SUM(LENGTH(text_body)) 18ms → ?) —
860 // pre-compile every aggregate arg that's a `fully_compilable`
861 // PURE expression over bound columns. Without this, `LENGTH(col)`
862 // / `COALESCE(col, '')` / `CAST(col AS BIGINT)` etc. ALL fell
863 // through to the `(None, Some(e)) => eval_arg(e, mat, ...)` slow
864 // path that materialises a Cow<Row> per input row — for a 25k-row
865 // JOIN that's 25k full-row clones for one column read. The Step
866 // VM (`eval_compiled_ref`) reads columns by RowRef::get and runs
867 // the same `apply_function` dispatcher with zero materialisation.
868 let arg_compiled: Vec<Option<eval::CompiledExpr>> = agg_specs
869 .iter()
870 .enumerate()
871 .map(|(i, spec)| match (&arg_pos[i], &spec.arg) {
872 (Some(_), _) => None,
873 (None, Some(e)) if eval::fully_compilable(e) => Some(eval::compile_expr(e, &ctx)),
874 _ => None,
875 })
876 .collect();
877 // v7.37.4 (L1 — executor-time CSE / mailrs P0) — dedupe
878 // compiled aggregate-arg expressions across specs. mailrs's
879 // `/api/conversations` SQL has 14 aggregates whose compiled
880 // CASE/CAST arg expressions overlap heavily (`m.message_id != ''`
881 // re-appears 4×, the inner `CASE WHEN m.message_id != '' THEN
882 // m.message_id ELSE CAST(m.id AS TEXT) END` re-appears 3×). Each
883 // dup currently costs one Step-VM walk per row — 100k rows ×
884 // ~3-4 redundant evals = ~300-400k wasted Step-VM runs.
885 //
886 // Dedupe key = source `Expr` (PartialEq). `CompiledExpr` itself
887 // is not `Hash` / `Eq`, but n_specs is small (≤ ~20 in practice);
888 // O(n²) PartialEq probe cost = ~196 cmp per query, vs millions
889 // of saved per-row evals. `fully_compilable` requires PURE
890 // scalars (no NOW / RANDOM / sequence accessors), so an earlier
891 // eval has identical observable semantics to the original.
892 //
893 // `arg_slot[i] = Some(s)` means spec `i`'s compiled arg lives in
894 // slot `s` of `arg_unique_idx` (which points back into
895 // `arg_compiled` for the canonical owner). Per-row cache fills
896 // LAZILY — preserves the current FILTER semantics where an arg
897 // whose spec is filtered out is never evaluated (and never
898 // surfaces a type error). Reset to `None` at the top of each row.
899 let mut arg_unique_idx: Vec<usize> = Vec::new();
900 let mut arg_slot: Vec<Option<usize>> = Vec::with_capacity(agg_specs.len());
901 arg_slot.resize(agg_specs.len(), None);
902 for (i, spec) in agg_specs.iter().enumerate() {
903 if arg_pos[i].is_some() || arg_compiled[i].is_none() {
904 continue;
905 }
906 let src = spec.arg.as_ref().expect("arg_compiled => spec.arg is Some");
907 let pos = arg_unique_idx
908 .iter()
909 .position(|&j| agg_specs[j].arg.as_ref().is_some_and(|other| other == src));
910 arg_slot[i] = Some(match pos {
911 Some(p) => p,
912 None => {
913 arg_unique_idx.push(i);
914 arg_unique_idx.len() - 1
915 }
916 });
917 }
918 let mut row_eval_cache: Vec<Option<Value>> = Vec::with_capacity(arg_unique_idx.len());
919 row_eval_cache.resize(arg_unique_idx.len(), None);
920 // v7.33 (array_agg perf) — bound positions for each spec's internal
921 // ORDER BY keys, so an ordered aggregate (`array_agg(x ORDER BY y)`)
922 // reads the sort key by reference (RowRef::get) instead of
923 // materialising the whole combined join row per input row just to
924 // eval one bound column. Mirrors arg_pos. On the inbox shape this
925 // turned 24k full-row (~1 KB each) clones into 24k single-cell reads.
926 let order_pos: Vec<Vec<Option<usize>>> = agg_specs
927 .iter()
928 .map(|spec| spec.order_by.iter().map(|o| col_pos(&o.expr)).collect())
929 .collect();
930 // v7.37.43 (DISTA A-3) — precompute the per-spec arg2 when it is a
931 // bare literal. `string_agg(DISTINCT col, ',')` and every other
932 // call with a constant separator goes through this path; PG evaluates
933 // arg2 as a Const once at plan time. SPG was paying a Cow row
934 // materialisation per input row purely so `eval_arg(literal, &row)`
935 // could run — but a literal doesn't read the row at all. Hoist the
936 // literal value into a per-query table; per-row arg2 just clones it.
937 //
938 // Sentinel: when arg2 is present but NOT a literal, the entry stays
939 // `None` and the per-row path still falls into the eval branch
940 // (which forces `needs_mat`).
941 let arg2_literal_val: Vec<Option<Value<'static>>> = agg_specs
942 .iter()
943 .map(|s| match &s.arg2 {
944 Some(Expr::Literal(l)) => Some(eval::literal_to_value(l)),
945 _ => None,
946 })
947 .collect();
948 // Does any spec need the fully-materialised row in the bound fast
949 // path — a FILTER, a non-bound value arg, a NON-LITERAL second arg,
950 // or a non-bound ORDER key? When false (every aggregate arg/key is a
951 // bound column — the inbox shape, and the DISTA shape after A-3)
952 // the bound fast path never materialises a row.
953 let needs_mat = agg_specs.iter().enumerate().any(|(i, s)| {
954 s.filter.is_some()
955 || (s.arg.is_some() && arg_pos[i].is_none() && arg_compiled[i].is_none())
956 || (s.arg2.is_some() && arg2_literal_val[i].is_none())
957 || order_pos[i].iter().any(Option::is_none)
958 });
959 let ci_positions: Vec<usize> = group_exprs
960 .iter()
961 .enumerate()
962 .filter(|(_, g)| {
963 matches!(
964 eval::column_collation(g, &ctx),
965 Some(spg_storage::Collation::CaseInsensitive)
966 )
967 })
968 .map(|(i, _)| i)
969 .collect();
970 // v7.31 (perf 3e) — per-row scratch buffers. The fast path used
971 // to allocate a key String (and a refs Vec) for EVERY row just
972 // to probe the group map; hits — the overwhelming case — now
973 // touch the allocator zero times.
974 let mut keybuf_s = String::new();
975 // v7.36 — reused Step VM eval stack for compiled aggregate args.
976 let mut eval_stack: Vec<Value<'static>> = Vec::new();
977 let mut dkeybuf = String::new();
978 let mut refs: Vec<&Value> = Vec::with_capacity(group_pos.len());
979 // v7.32 (round-31) — an aggregate's argument / FILTER / second arg /
980 // ORDER key may itself be a *correlated* subquery, e.g.
981 // `MAX((SELECT i.v FROM inner i WHERE i.fk = o.id))`. A non-correlated
982 // subquery is pre-resolved to a literal before this loop, but a
983 // correlated one survives as a subquery node and must be evaluated per
984 // outer row through the correlated evaluator — the same hook the
985 // select-list / HAVING / ORDER finalisers already use below. Plain
986 // `eval_expr` would hit "subquery reached row eval".
987 //
988 // The `any_agg_subquery` gate is computed once here so the common case
989 // (no subquery anywhere in the aggregate args — including every hot
990 // scan/group aggregate) short-circuits before the per-row
991 // `expr_has_subquery` walk: `eval_arg` is then exactly `eval_expr`.
992 let any_agg_subquery = correlated_eval.is_some()
993 && agg_specs.iter().any(|s| {
994 s.filter
995 .as_ref()
996 .is_some_and(|e| crate::expr_has_subquery(e))
997 || s.arg.as_ref().is_some_and(|e| crate::expr_has_subquery(e))
998 || s.arg2.as_ref().is_some_and(|e| crate::expr_has_subquery(e))
999 || s.order_by.iter().any(|o| crate::expr_has_subquery(&o.expr))
1000 });
1001 let eval_arg =
1002 |e: &Expr, r: &Row<'static>, c: &EvalContext<'_>| -> Result<Value<'static>, EvalError> {
1003 match correlated_eval {
1004 Some(f) if any_agg_subquery && crate::expr_has_subquery(e) => f(e, r, c),
1005 _ => eval::eval_expr(e, r, c),
1006 }
1007 };
1008 // v7.36 (perf — mailrs Phase 1, post u64-hash) — single
1009 // anonymous group fast path. When the query has no GROUP BY
1010 // (`SELECT SUM(LENGTH(col)) FROM ...`, COUNT, AVG, etc.) the
1011 // whole input collapses into one group. The fast path below
1012 // still pays one `groups.get("")` hash probe per row plus
1013 // `entry = &mut order[0]` reindex even when the empty-key
1014 // path encodes nothing — measured ~50 ns/row across 25 k rows
1015 // = ~1.25 ms of pure bookkeeping on the user_storage_usage
1016 // baseline.
1017 //
1018 // Bypass: lift `entry` outside the loop and feed every row
1019 // straight into it. Same `update_state` machinery, zero
1020 // per-row hash work, zero per-row index lookup.
1021 let single_anon_group = group_exprs.is_empty() && !rows.is_empty();
1022 if single_anon_group {
1023 // Seed the single group at idx 0 once.
1024 let init: Vec<AggState> = (0..agg_specs.len()).map(|_| AggState::default()).collect();
1025 order.clear();
1026 order.push((Vec::new(), init));
1027 }
1028 // v7.36 (perf — mailrs Phase 1, count_messages 2.58 → ?) —
1029 // `COUNT(*)` short-circuit. For a single-anon-group `COUNT(*)`
1030 // with no FILTER / DISTINCT, every survivor counts once — the
1031 // answer IS `rows.len()`. Skips the 25 k iterations of
1032 // `update_state("count_star", …)` on the mailrs count_messages
1033 // shape; the JOIN already produced exactly the set of rows
1034 // that must be counted.
1035 if single_anon_group
1036 && agg_specs.len() == 1
1037 && agg_specs[0].name == "count_star"
1038 && agg_specs[0].filter.is_none()
1039 && agg_specs[0].arg.is_none()
1040 && agg_specs[0].arg2.is_none()
1041 && agg_specs[0].order_by.is_empty()
1042 && !agg_specs[0].distinct
1043 {
1044 let state = &mut order[0].1[0];
1045 state.count = rows.len() as i64;
1046 return Ok(order);
1047 }
1048 // v7.36 (perf — mailrs Phase 1) — `COUNT(<bound col>)` (non-`*`)
1049 // collapses to: read the cell, increment when not NULL. Skips
1050 // the per-row spec dispatch + `update_state("count", …)`.
1051 if single_anon_group
1052 && agg_specs.len() == 1
1053 && agg_specs[0].name == "count"
1054 && agg_specs[0].filter.is_none()
1055 && agg_specs[0].arg2.is_none()
1056 && agg_specs[0].order_by.is_empty()
1057 && !agg_specs[0].distinct
1058 && arg_pos[0].is_some()
1059 {
1060 let p = arg_pos[0].unwrap();
1061 let mut count: i64 = 0;
1062 for row in rows {
1063 if !matches!(row.get(p), Some(Value::Null) | None) {
1064 count += 1;
1065 }
1066 }
1067 let state = &mut order[0].1[0];
1068 state.count = count;
1069 return Ok(order);
1070 }
1071 // v7.36 (perf — mailrs Phase 1, user_storage_usage 7.5 → ?) —
1072 // single-aggregate streaming accumulator. For
1073 // `SUM(<compiled-expr>)` / `SUM(<bound col>)` with no GROUP BY,
1074 // no FILTER, no arg2, no ORDER BY, no DISTINCT, the whole
1075 // per-row work collapses to: eval the arg, match the Value
1076 // variant, accumulate. Skips the spec-dispatch loop +
1077 // `update_state` per-row name match. On a 25 k-row JOIN
1078 // (user_storage_usage `SUM(LENGTH(text_body))`) that's
1079 // ~50-100 ns/row of pure spec-dispatch overhead removed.
1080 if single_anon_group
1081 && agg_specs.len() == 1
1082 && agg_specs[0].filter.is_none()
1083 && agg_specs[0].arg2.is_none()
1084 && agg_specs[0].order_by.is_empty()
1085 && !agg_specs[0].distinct
1086 && (agg_specs[0].name == "sum" || agg_specs[0].name == "avg")
1087 && (arg_pos[0].is_some() || arg_compiled[0].is_some())
1088 {
1089 let arg_pos0 = arg_pos[0];
1090 let arg_c0 = &arg_compiled[0];
1091 let mut sum_int: i64 = 0;
1092 let mut sum_float: f64 = 0.0;
1093 let mut use_float = false;
1094 let mut count: i64 = 0;
1095 // Borrow-aware fast inner: avoid the per-row clone when arg
1096 // is a bound column position.
1097 if let Some(p) = arg_pos0 {
1098 for row in rows {
1099 let v_ref = row.get(p).unwrap_or(&Value::Null);
1100 match v_ref {
1101 Value::Null => continue,
1102 Value::SmallInt(n) => {
1103 sum_int += i64::from(*n);
1104 count += 1;
1105 }
1106 Value::Int(n) => {
1107 sum_int += i64::from(*n);
1108 count += 1;
1109 }
1110 Value::BigInt(n) => {
1111 sum_int += *n;
1112 count += 1;
1113 }
1114 Value::Float(x) => {
1115 sum_float += *x;
1116 use_float = true;
1117 count += 1;
1118 }
1119 other => {
1120 return Err(EvalError::TypeMismatch {
1121 detail: format!("sum/avg need numeric, got {:?}", other.data_type()),
1122 });
1123 }
1124 }
1125 }
1126 } else if let Some(p) = arg_c0.as_ref().and_then(|c| c.as_single_column_length()) {
1127 // v7.36 (perf — mailrs Phase 1, user_storage_usage hot
1128 // inner) — `SUM(LENGTH(<text col>))` collapses to a
1129 // straight scan: read the cell by ref, branch on the
1130 // variant, do an ASCII probe + `len()` (or
1131 // `chars().count()` on non-ASCII), accumulate. No Step
1132 // VM, no stack push/pop, no `BigInt` boxing on the way
1133 // out — pure i64 sum. The original Step VM path keeps
1134 // running for everything outside this shape (`SUM(col)`,
1135 // `SUM(expr)`, multi-step compiled args).
1136 for row in rows {
1137 let Some(v_ref) = row.get(p) else {
1138 continue;
1139 };
1140 let n = match v_ref {
1141 Value::Null => continue,
1142 Value::Text(s) => {
1143 if s.is_ascii() {
1144 s.len() as i64
1145 } else {
1146 s.chars().count() as i64
1147 }
1148 }
1149 other => {
1150 return Err(EvalError::TypeMismatch {
1151 detail: format!("length() needs text, got {:?}", other.data_type()),
1152 });
1153 }
1154 };
1155 sum_int += n;
1156 count += 1;
1157 }
1158 } else {
1159 let c = arg_c0.as_ref().unwrap();
1160 for row in rows {
1161 let v = eval::eval_compiled_ref(c, row, &ctx, &mut eval_stack)?;
1162 match v {
1163 Value::Null => continue,
1164 Value::SmallInt(n) => {
1165 sum_int += i64::from(n);
1166 count += 1;
1167 }
1168 Value::Int(n) => {
1169 sum_int += i64::from(n);
1170 count += 1;
1171 }
1172 Value::BigInt(n) => {
1173 sum_int += n;
1174 count += 1;
1175 }
1176 Value::Float(x) => {
1177 sum_float += x;
1178 use_float = true;
1179 count += 1;
1180 }
1181 other => {
1182 return Err(EvalError::TypeMismatch {
1183 detail: format!("sum/avg need numeric, got {:?}", other.data_type()),
1184 });
1185 }
1186 }
1187 }
1188 }
1189 let state = &mut order[0].1[0];
1190 state.count = count;
1191 state.sum_int = sum_int;
1192 state.sum_float = sum_float;
1193 state.use_float = use_float;
1194 return Ok(order);
1195 }
1196 // v7.37.x (mailrs Track A 100k attack) — tight inlined loop for
1197 // the "single-Text GROUP BY + single MAX(bound numeric arg)"
1198 // shape. See `dedicated_max_loop` above for the gate. Returns
1199 // straight to the caller; the rest of the function (single-anon,
1200 // bound-fast, eval-slow paths) is skipped.
1201 if dedicated_max_loop && !single_anon_group {
1202 let gpos = group_pos[0].expect("dedicated_max_loop gates on Some");
1203 let apos = arg_pos[0].expect("dedicated_max_loop gates on Some");
1204 for row in rows {
1205 let kv = row.get(gpos).unwrap_or(&Value::Null);
1206 let idx = match kv {
1207 Value::Text(s) => match groups_text.get(s.as_ref()) {
1208 Some(&i) => i,
1209 None => {
1210 let i = order.len();
1211 order.push((
1212 alloc::vec![Value::text(s.clone())],
1213 alloc::vec![AggState::default()],
1214 ));
1215 groups_text.insert(s.to_string(), i);
1216 i
1217 }
1218 },
1219 Value::Null => match null_group_idx {
1220 Some(i) => i,
1221 None => {
1222 let i = order.len();
1223 order.push((alloc::vec![Value::Null], alloc::vec![AggState::default()]));
1224 null_group_idx = Some(i);
1225 i
1226 }
1227 },
1228 _ => {
1229 // Schema said Text but value isn't — fall back to
1230 // the generic encoded path for correctness.
1231 refs.clear();
1232 refs.push(kv);
1233 encode_key_refs_into(&refs, &mut keybuf_s);
1234 match groups.get(keybuf_s.as_str()) {
1235 Some(&i) => i,
1236 None => {
1237 let i = order.len();
1238 order.push((
1239 alloc::vec![kv.clone().into_owned()],
1240 alloc::vec![AggState::default()],
1241 ));
1242 groups.insert(keybuf_s.clone(), i);
1243 i
1244 }
1245 }
1246 }
1247 };
1248 // Inline MAX accumulator — skip the union-typed
1249 // `update_state` enum jump and per-spec arg dispatch.
1250 let av = row.get(apos).unwrap_or(&Value::Null);
1251 if !matches!(av, Value::Null) {
1252 let st = &mut order[idx].1[0];
1253 let upd = match &st.extreme {
1254 None => true,
1255 Some(prev) => value_cmp(av, prev) == core::cmp::Ordering::Greater,
1256 };
1257 if upd {
1258 st.extreme = Some(av.clone().into_owned());
1259 }
1260 }
1261 }
1262 return Ok(order);
1263 }
1264
1265 for row in rows {
1266 // v7.37.4 (L1 CSE) — reset per-row cache for shared compiled
1267 // aggregate-arg evals. No-op when no dedupe (empty vec).
1268 for slot in row_eval_cache.iter_mut() {
1269 *slot = None;
1270 }
1271 if single_anon_group {
1272 let entry = &mut order[0];
1273 let mat: Option<Cow<'_, Row>> = if needs_mat { Some(row.as_row()) } else { None };
1274 for (i, spec) in agg_specs.iter().enumerate() {
1275 if let Some(f) = &spec.filter
1276 && !matches!(
1277 eval_arg(f, mat.as_deref().expect("needs_mat for FILTER"), &ctx)?,
1278 Value::Bool(true)
1279 )
1280 {
1281 continue;
1282 }
1283 let arg_owned: Value;
1284 let arg_ref: &Value = match (&arg_pos[i], arg_slot[i], &spec.arg) {
1285 (Some(p), _, _) => row.get(*p).unwrap_or(&Value::Null),
1286 (None, None, None) => {
1287 arg_owned = Value::Bool(true);
1288 &arg_owned
1289 }
1290 (None, Some(s), _) => {
1291 if row_eval_cache[s].is_none() {
1292 let c = arg_compiled[arg_unique_idx[s]]
1293 .as_ref()
1294 .expect("arg_unique_idx points at a compiled spec");
1295 let v = eval::eval_compiled_ref(c, row, &ctx, &mut eval_stack)?;
1296 row_eval_cache[s] = Some(v);
1297 }
1298 row_eval_cache[s].as_ref().expect("just filled above")
1299 }
1300 (None, None, Some(e)) => {
1301 arg_owned = eval_arg(
1302 e,
1303 mat.as_deref().expect("needs_mat for non-bound arg"),
1304 &ctx,
1305 )?;
1306 &arg_owned
1307 }
1308 };
1309 let arg2_val = match (&spec.arg2, &arg2_literal_val[i]) {
1310 (None, _) => None,
1311 // v7.37.43 (DISTA A-3) — literal arg2: clone the
1312 // precomputed value, skip per-row eval & row mat.
1313 (Some(_), Some(lit)) => Some(lit.clone()),
1314 (Some(e), None) => Some(eval_arg(
1315 e,
1316 mat.as_deref().expect("needs_mat for arg2"),
1317 &ctx,
1318 )?),
1319 };
1320 let order_keys: Option<Vec<Value<'static>>> = if spec.order_by.is_empty() {
1321 None
1322 } else {
1323 let mut keys: Vec<Value<'static>> = Vec::with_capacity(spec.order_by.len());
1324 for (k, o) in spec.order_by.iter().enumerate() {
1325 let v: Value<'static> = if let Some(p) = order_pos[i][k] {
1326 row.get(p)
1327 .cloned()
1328 .map(Value::into_owned)
1329 .unwrap_or(Value::Null)
1330 } else {
1331 eval_arg(
1332 &o.expr,
1333 mat.as_deref().expect("needs_mat for ORDER key"),
1334 &ctx,
1335 )?
1336 };
1337 keys.push(v);
1338 }
1339 Some(keys)
1340 };
1341 // v7.36 (perf — bugfix v7.36.1 candidate) — first_ordered
1342 // was missing from the single_anon_group fast path,
1343 // sending `(array_agg(x ORDER BY y))[1]` values into
1344 // `update_state(array_agg, …)` whose finalize ignored
1345 // the absent `first_best` and returned `[]`. The slow
1346 // path below has the same branch — keep them aligned.
1347 if spec.first_ordered {
1348 if let Some(keys) = order_keys {
1349 let st = &mut entry.1[i];
1350 let better = match &st.first_best {
1351 None => true,
1352 Some((bk, _)) => {
1353 cmp_order_keys(&spec.order_by, &keys, bk)
1354 == core::cmp::Ordering::Less
1355 }
1356 };
1357 if better {
1358 st.first_best = Some((keys, arg_ref.clone().into_owned()));
1359 }
1360 }
1361 continue;
1362 }
1363 if spec.distinct {
1364 // v7.37.x (mailrs Track A 100k distinct_aggs attack)
1365 // — single-Text DISTINCT fast path. Within a single
1366 // distinct spec all input values come from one
1367 // expression and share one type, so the encode-
1368 // prefix (`S<text>|`) is redundant: the column
1369 // text alone is collision-free within this spec's
1370 // `seen` set. Skips encode_one + 2-walk
1371 // contains+insert; only Text arms apply, others
1372 // ride the encoded path unchanged.
1373 //
1374 // v7.37.x (docker-fair DISTA attack) — extend the
1375 // single-family fast path to BigInt via a parallel
1376 // `seen_int: Option<BTreeSet<i64>>`. The DISTA
1377 // `COUNT(DISTINCT m.id)` shape pumps 25 k BigInt
1378 // probes; skipping `encode_key_refs_into` saves
1379 // ~100 ns of alloc + format churn per row.
1380 if let Value::Text(s) = arg_ref {
1381 if entry.1[i].seen.contains(s.as_ref()) {
1382 continue;
1383 }
1384 entry.1[i].seen.insert(s.to_string());
1385 } else if let Value::BigInt(n) = arg_ref {
1386 let set = entry.1[i].seen_int.get_or_insert_with(BTreeSet::new);
1387 if !set.insert(*n) {
1388 continue;
1389 }
1390 } else if let Value::Int(n) = arg_ref {
1391 let set = entry.1[i].seen_int.get_or_insert_with(BTreeSet::new);
1392 if !set.insert(i64::from(*n)) {
1393 continue;
1394 }
1395 } else {
1396 encode_key_refs_into(core::slice::from_ref(&arg_ref), &mut dkeybuf);
1397 if entry.1[i].seen.contains(dkeybuf.as_str()) {
1398 continue;
1399 }
1400 entry.1[i].seen.insert(dkeybuf.clone());
1401 }
1402 }
1403 // v7.37.x (mailrs Track A 100k attack) — inline the
1404 // common aggregate kinds (MAX / MIN / Count / CountStar
1405 // / BoolOr / BoolAnd) here instead of dispatching
1406 // through `update_state`'s enum jump + per-kind branch.
1407 // Skipping the function-call overhead saves ~20-30 ns
1408 // per spec per row at 100 k; the slow kinds keep the
1409 // dispatched call.
1410 match spec.kind {
1411 AggKind::Max => {
1412 if !matches!(arg_ref, Value::Null) {
1413 let st = &mut entry.1[i];
1414 let upd = match &st.extreme {
1415 None => true,
1416 Some(prev) => {
1417 value_cmp(arg_ref, prev) == core::cmp::Ordering::Greater
1418 }
1419 };
1420 if upd {
1421 st.extreme = Some(arg_ref.clone().into_owned());
1422 }
1423 }
1424 }
1425 AggKind::Min => {
1426 if !matches!(arg_ref, Value::Null) {
1427 let st = &mut entry.1[i];
1428 let upd = match &st.extreme {
1429 None => true,
1430 Some(prev) => value_cmp(arg_ref, prev) == core::cmp::Ordering::Less,
1431 };
1432 if upd {
1433 st.extreme = Some(arg_ref.clone().into_owned());
1434 }
1435 }
1436 }
1437 AggKind::CountStar => {
1438 entry.1[i].count += 1;
1439 }
1440 AggKind::Count => {
1441 if !matches!(arg_ref, Value::Null) {
1442 entry.1[i].count += 1;
1443 }
1444 }
1445 AggKind::BoolOr => match arg_ref {
1446 Value::Bool(b) => {
1447 let st = &mut entry.1[i];
1448 st.bool_acc = Some(st.bool_acc.unwrap_or(false) || *b);
1449 }
1450 Value::Null => {}
1451 _ => update_state(
1452 &mut entry.1[i],
1453 spec.kind,
1454 &spec.name,
1455 arg_ref,
1456 arg2_val.as_ref(),
1457 order_keys,
1458 )?,
1459 },
1460 AggKind::BoolAnd => match arg_ref {
1461 Value::Bool(b) => {
1462 let st = &mut entry.1[i];
1463 st.bool_acc = Some(st.bool_acc.unwrap_or(true) && *b);
1464 }
1465 Value::Null => {}
1466 _ => update_state(
1467 &mut entry.1[i],
1468 spec.kind,
1469 &spec.name,
1470 arg_ref,
1471 arg2_val.as_ref(),
1472 order_keys,
1473 )?,
1474 },
1475 _ => {
1476 update_state(
1477 &mut entry.1[i],
1478 spec.kind,
1479 &spec.name,
1480 arg_ref,
1481 arg2_val.as_ref(),
1482 order_keys,
1483 )?;
1484 }
1485 }
1486 }
1487 continue;
1488 }
1489 // Fast key: bound positions + no ci folding -> encode
1490 // straight from borrowed cells; group_vals materialise
1491 // only when the group is NEW.
1492 if all_groups_bound && ci_positions.is_empty() {
1493 // v7.37.x — single-Text fast path uses the raw text as the
1494 // map key (no encode_one's `S<text>|` prefix/suffix push,
1495 // no refs Vec rebuild). NULL values land in a dedicated
1496 // slot so SQL's "all NULLs share one group" semantics hold.
1497 let idx = if single_text_group_col {
1498 let v = row.get(group_pos[0].unwrap()).unwrap_or(&Value::Null);
1499 match v {
1500 Value::Text(s) => match groups_text.get(s.as_ref()) {
1501 Some(&i) => i,
1502 None => {
1503 let i = order.len();
1504 let init: Vec<AggState> =
1505 (0..agg_specs.len()).map(|_| AggState::default()).collect();
1506 order.push((alloc::vec![Value::text(s.clone())], init));
1507 groups_text.insert(s.to_string(), i);
1508 i
1509 }
1510 },
1511 Value::Null => match null_group_idx {
1512 Some(i) => i,
1513 None => {
1514 let i = order.len();
1515 let init: Vec<AggState> =
1516 (0..agg_specs.len()).map(|_| AggState::default()).collect();
1517 order.push((alloc::vec![Value::Null], init));
1518 null_group_idx = Some(i);
1519 i
1520 }
1521 },
1522 _ => {
1523 // Schema says Text but value is something else
1524 // (coercion edge case). Fall back to the encoded
1525 // path for correctness — same logic as the
1526 // non-single-Text branch below.
1527 refs.clear();
1528 refs.push(v);
1529 encode_key_refs_into(&refs, &mut keybuf_s);
1530 match groups.get(keybuf_s.as_str()) {
1531 Some(&i) => i,
1532 None => {
1533 let i = order.len();
1534 let init: Vec<AggState> =
1535 (0..agg_specs.len()).map(|_| AggState::default()).collect();
1536 order.push((alloc::vec![v.clone().into_owned()], init));
1537 groups.insert(keybuf_s.clone(), i);
1538 i
1539 }
1540 }
1541 }
1542 }
1543 } else {
1544 refs.clear();
1545 refs.extend(
1546 group_pos
1547 .iter()
1548 .map(|p| row.get(p.unwrap()).unwrap_or(&Value::Null)),
1549 );
1550 encode_key_refs_into(&refs, &mut keybuf_s);
1551 match groups.get(keybuf_s.as_str()) {
1552 Some(&i) => i,
1553 None => {
1554 let i = order.len();
1555 let init: Vec<AggState> =
1556 (0..agg_specs.len()).map(|_| AggState::default()).collect();
1557 let owned: Vec<Value<'static>> =
1558 refs.iter().map(|v| (*v).clone().into_owned()).collect();
1559 order.push((owned, init));
1560 groups.insert(keybuf_s.clone(), i);
1561 i
1562 }
1563 }
1564 };
1565 let entry = &mut order[idx];
1566 // v7.33 (array_agg perf) — materialise the combined row AT
1567 // MOST once per input row, and only when a spec actually
1568 // needs the eval path (FILTER / non-bound arg / arg2 / non-
1569 // bound ORDER key). Bound args and bound ORDER keys read
1570 // cells by reference below, so the inbox shape (all bound)
1571 // never materialises — killing the per-row ~1 KB clone that
1572 // dominated the ordered-aggregate cost.
1573 let mat: Option<Cow<'_, Row>> = if needs_mat { Some(row.as_row()) } else { None };
1574 for (i, spec) in agg_specs.iter().enumerate() {
1575 // v7.32 (round-29) — FILTER (WHERE cond): exclude rows
1576 // where cond is not TRUE before they reach this
1577 // aggregate's accumulator (and before DISTINCT dedup).
1578 if let Some(f) = &spec.filter
1579 && !matches!(
1580 eval_arg(f, mat.as_deref().expect("needs_mat for FILTER"), &ctx)?,
1581 Value::Bool(true)
1582 )
1583 {
1584 continue;
1585 }
1586 let arg_owned: Value;
1587 let arg_ref: &Value = match (&arg_pos[i], arg_slot[i], &spec.arg) {
1588 (Some(p), _, _) => row.get(*p).unwrap_or(&Value::Null),
1589 (None, None, None) => {
1590 arg_owned = Value::Bool(true);
1591 &arg_owned
1592 }
1593 (None, Some(s), _) => {
1594 // v7.37.4 (L1 CSE) — shared compiled-arg slot.
1595 // First spec that needs slot `s` this row pays
1596 // the Step-VM eval; siblings reading the same
1597 // slot get the cached Value for free. Preserves
1598 // FILTER semantics: a spec filtered out above
1599 // never reaches here, so its arg stays unevaled.
1600 if row_eval_cache[s].is_none() {
1601 let c = arg_compiled[arg_unique_idx[s]]
1602 .as_ref()
1603 .expect("arg_unique_idx points at a compiled spec");
1604 let v = eval::eval_compiled_ref(c, row, &ctx, &mut eval_stack)?;
1605 row_eval_cache[s] = Some(v);
1606 }
1607 row_eval_cache[s].as_ref().expect("just filled above")
1608 }
1609 (None, None, Some(e)) => {
1610 arg_owned = eval_arg(
1611 e,
1612 mat.as_deref().expect("needs_mat for non-bound arg"),
1613 &ctx,
1614 )?;
1615 &arg_owned
1616 }
1617 };
1618 let arg2_val = match (&spec.arg2, &arg2_literal_val[i]) {
1619 (None, _) => None,
1620 // v7.37.43 (DISTA A-3) — literal arg2: clone the
1621 // precomputed value, skip per-row eval & row mat.
1622 (Some(_), Some(lit)) => Some(lit.clone()),
1623 (Some(e), None) => Some(eval_arg(
1624 e,
1625 mat.as_deref().expect("needs_mat for arg2"),
1626 &ctx,
1627 )?),
1628 };
1629 let order_keys: Option<Vec<Value<'static>>> = if spec.order_by.is_empty() {
1630 None
1631 } else {
1632 let mut keys: Vec<Value<'static>> = Vec::with_capacity(spec.order_by.len());
1633 for (k, o) in spec.order_by.iter().enumerate() {
1634 // Bound ORDER key → read the cell by reference; only
1635 // a non-bound key falls to the materialised eval path.
1636 keys.push(match order_pos[i][k] {
1637 Some(p) => row
1638 .get(p)
1639 .cloned()
1640 .map(Value::into_owned)
1641 .unwrap_or(Value::Null),
1642 None => eval_arg(
1643 &o.expr,
1644 mat.as_deref().expect("needs_mat for non-bound ORDER key"),
1645 &ctx,
1646 )?,
1647 });
1648 }
1649 Some(keys)
1650 };
1651 // v7.33 (array_agg argmax) — first_ordered: keep only the
1652 // running first-by-order element (strict-less replacement
1653 // = ties keep the earliest row, matching the stable-sort
1654 // `[1]`), no array build.
1655 if spec.first_ordered {
1656 if let Some(keys) = order_keys {
1657 let st = &mut entry.1[i];
1658 let better = match &st.first_best {
1659 None => true,
1660 Some((bk, _)) => {
1661 cmp_order_keys(&spec.order_by, &keys, bk)
1662 == core::cmp::Ordering::Less
1663 }
1664 };
1665 if better {
1666 st.first_best = Some((keys, arg_ref.clone().into_owned()));
1667 }
1668 }
1669 continue;
1670 }
1671 if spec.distinct {
1672 // v7.37.x — single-Text DISTINCT fast path (see
1673 // bound fast path counterpart above). Per-spec
1674 // type invariance lets us use the column text as
1675 // the `seen` key directly, no `S<text>|` prefix.
1676 // v7.37.x (docker-fair DISTA) — BigInt parallel
1677 // path skips encode_key_refs_into entirely.
1678 if let Value::Text(s) = arg_ref {
1679 if entry.1[i].seen.contains(s.as_ref()) {
1680 continue;
1681 }
1682 entry.1[i].seen.insert(s.to_string());
1683 } else if let Value::BigInt(n) = arg_ref {
1684 let set = entry.1[i].seen_int.get_or_insert_with(BTreeSet::new);
1685 if !set.insert(*n) {
1686 continue;
1687 }
1688 } else if let Value::Int(n) = arg_ref {
1689 let set = entry.1[i].seen_int.get_or_insert_with(BTreeSet::new);
1690 if !set.insert(i64::from(*n)) {
1691 continue;
1692 }
1693 } else {
1694 encode_key_refs_into(core::slice::from_ref(&arg_ref), &mut dkeybuf);
1695 if entry.1[i].seen.contains(dkeybuf.as_str()) {
1696 continue;
1697 }
1698 entry.1[i].seen.insert(dkeybuf.clone());
1699 }
1700 }
1701 // v7.37.x (mailrs Track A 100k attack) — inline the
1702 // common aggregate kinds (MAX / MIN / Count / CountStar
1703 // / BoolOr / BoolAnd) here instead of dispatching
1704 // through `update_state`'s enum jump + per-kind branch.
1705 // Skipping the function-call overhead saves ~20-30 ns
1706 // per spec per row at 100 k; the slow kinds keep the
1707 // dispatched call.
1708 match spec.kind {
1709 AggKind::Max => {
1710 if !matches!(arg_ref, Value::Null) {
1711 let st = &mut entry.1[i];
1712 let upd = match &st.extreme {
1713 None => true,
1714 Some(prev) => {
1715 value_cmp(arg_ref, prev) == core::cmp::Ordering::Greater
1716 }
1717 };
1718 if upd {
1719 st.extreme = Some(arg_ref.clone().into_owned());
1720 }
1721 }
1722 }
1723 AggKind::Min => {
1724 if !matches!(arg_ref, Value::Null) {
1725 let st = &mut entry.1[i];
1726 let upd = match &st.extreme {
1727 None => true,
1728 Some(prev) => value_cmp(arg_ref, prev) == core::cmp::Ordering::Less,
1729 };
1730 if upd {
1731 st.extreme = Some(arg_ref.clone().into_owned());
1732 }
1733 }
1734 }
1735 AggKind::CountStar => {
1736 entry.1[i].count += 1;
1737 }
1738 AggKind::Count => {
1739 if !matches!(arg_ref, Value::Null) {
1740 entry.1[i].count += 1;
1741 }
1742 }
1743 AggKind::BoolOr => match arg_ref {
1744 Value::Bool(b) => {
1745 let st = &mut entry.1[i];
1746 st.bool_acc = Some(st.bool_acc.unwrap_or(false) || *b);
1747 }
1748 Value::Null => {}
1749 _ => update_state(
1750 &mut entry.1[i],
1751 spec.kind,
1752 &spec.name,
1753 arg_ref,
1754 arg2_val.as_ref(),
1755 order_keys,
1756 )?,
1757 },
1758 AggKind::BoolAnd => match arg_ref {
1759 Value::Bool(b) => {
1760 let st = &mut entry.1[i];
1761 st.bool_acc = Some(st.bool_acc.unwrap_or(true) && *b);
1762 }
1763 Value::Null => {}
1764 _ => update_state(
1765 &mut entry.1[i],
1766 spec.kind,
1767 &spec.name,
1768 arg_ref,
1769 arg2_val.as_ref(),
1770 order_keys,
1771 )?,
1772 },
1773 _ => {
1774 update_state(
1775 &mut entry.1[i],
1776 spec.kind,
1777 &spec.name,
1778 arg_ref,
1779 arg2_val.as_ref(),
1780 order_keys,
1781 )?;
1782 }
1783 }
1784 }
1785 continue;
1786 }
1787 // v7.32 (P4 increment 2) — eval (non-bound) path: present the
1788 // row as a borrowed Row once (Owned → zero-cost borrow; a join
1789 // tuple materialises here exactly once, never on the bound fast
1790 // path above), then the original eval loop runs unchanged.
1791 let row_materialised = row.as_row();
1792 let row: &Row<'static> = &row_materialised;
1793 let group_vals: Vec<Value<'static>> = group_exprs
1794 .iter()
1795 .map(|g| eval::eval_expr(g, row, &ctx))
1796 .collect::<Result<_, _>>()?;
1797 // v7.17.0 Phase 2.5b — case-insensitive group keying: fold
1798 // only the ci columns, and only when any exist. Display
1799 // value (`group_vals`) stays original — only the key folds.
1800 let key = if ci_positions.is_empty() {
1801 encode_key(&group_vals)
1802 } else {
1803 let mut key_vals = group_vals.clone();
1804 for &i in &ci_positions {
1805 if let Value::Text(s) = &key_vals[i] {
1806 key_vals[i] = Value::text(s.to_ascii_lowercase());
1807 }
1808 }
1809 encode_key(&key_vals)
1810 };
1811 // Probe by index; the map owns the key once on vacant insert.
1812 let idx = match groups.get(key.as_str()) {
1813 Some(&i) => i,
1814 None => {
1815 let i = order.len();
1816 let init: Vec<AggState> =
1817 (0..agg_specs.len()).map(|_| AggState::default()).collect();
1818 order.push((group_vals.clone(), init));
1819 groups.insert(key, i);
1820 i
1821 }
1822 };
1823 let entry = &mut order[idx];
1824 for (i, spec) in agg_specs.iter().enumerate() {
1825 // v7.32 (round-29) — FILTER (WHERE cond): exclude rows where
1826 // cond is not TRUE before accumulation (and before DISTINCT).
1827 if let Some(f) = &spec.filter
1828 && !matches!(eval_arg(f, row, &ctx)?, Value::Bool(true))
1829 {
1830 continue;
1831 }
1832 let arg_val = match &spec.arg {
1833 None => Value::Bool(true), // count_star: sentinel non-null
1834 Some(e) => eval_arg(e, row, &ctx)?,
1835 };
1836 // v7.17.0 — `string_agg(value, separator)` evaluates the
1837 // separator per row but PG treats it as constant; we
1838 // pass the per-row value into update_state so a future
1839 // varying-separator caller still sees correct output,
1840 // even though SPG (like PG) only uses the most recent.
1841 let arg2_val = match &spec.arg2 {
1842 None => None,
1843 Some(e) => Some(eval_arg(e, row, &ctx)?),
1844 };
1845 // v7.24 (round-16 A) — aggregate-internal ORDER BY:
1846 // evaluate the key tuple against the source row.
1847 let order_keys: Option<Vec<Value<'static>>> = if spec.order_by.is_empty() {
1848 None
1849 } else {
1850 let mut keys: Vec<Value<'static>> = Vec::with_capacity(spec.order_by.len());
1851 for o in &spec.order_by {
1852 keys.push(eval_arg(&o.expr, row, &ctx)?);
1853 }
1854 Some(keys)
1855 };
1856 // v7.33 (array_agg argmax) — first_ordered: keep the running
1857 // first-by-order element only (mirrors the bound fast path).
1858 if spec.first_ordered {
1859 if let Some(keys) = order_keys {
1860 let st = &mut entry.1[i];
1861 let better = match &st.first_best {
1862 None => true,
1863 Some((bk, _)) => {
1864 cmp_order_keys(&spec.order_by, &keys, bk) == core::cmp::Ordering::Less
1865 }
1866 };
1867 if better {
1868 st.first_best = Some((keys, arg_val.clone().into_owned()));
1869 }
1870 }
1871 continue;
1872 }
1873 // v7.25 (round-17) — DISTINCT: drop repeated inputs
1874 // before they reach the accumulator. NULLs flow through
1875 // (each aggregate's own NULL rule applies; PG also
1876 // treats NULL as a single distinct value for array_agg).
1877 // v7.37.x — single-Text fast path same shape as the
1878 // bound/slow paths above.
1879 if spec.distinct {
1880 // v7.37.x (docker-fair DISTA) — single-family fast
1881 // paths skip encode_key for Text/BigInt/Int.
1882 let inserted = match &arg_val {
1883 Value::Text(s) => entry.1[i].seen.insert(s.to_string()),
1884 Value::BigInt(n) => entry.1[i]
1885 .seen_int
1886 .get_or_insert_with(BTreeSet::new)
1887 .insert(*n),
1888 Value::Int(n) => entry.1[i]
1889 .seen_int
1890 .get_or_insert_with(BTreeSet::new)
1891 .insert(i64::from(*n)),
1892 _ => {
1893 let key = encode_key(core::slice::from_ref(&arg_val));
1894 entry.1[i].seen.insert(key)
1895 }
1896 };
1897 if !inserted {
1898 continue;
1899 }
1900 }
1901 update_state(
1902 &mut entry.1[i],
1903 spec.kind,
1904 &spec.name,
1905 &arg_val,
1906 arg2_val.as_ref(),
1907 order_keys,
1908 )?;
1909 }
1910 }
1911 Ok(order)
1912}
1913
1914/// (2a) Build the synthetic per-group schema: `__grp_0..K` then
1915/// `__agg_0..N`. Group types are probed from the first row; aggregate
1916/// types from each spec.
1917fn build_synth_schema(
1918 rows: &[RowRef<'_>],
1919 group_exprs: &[Expr],
1920 agg_specs: &[AggSpec],
1921 schema_cols: &[ColumnSchema],
1922 table_alias: Option<&str>,
1923) -> Result<Vec<ColumnSchema>, EvalError> {
1924 let ctx = EvalContext::new(schema_cols, table_alias);
1925 // Build synthetic schema: __grp_0..K then __agg_0..N.
1926 let group_types: Vec<DataType> = if rows.is_empty() {
1927 // Use Text as a safe stand-in — empty result means schema isn't
1928 // observable. Avoids needing to evaluate group exprs on no row.
1929 group_exprs.iter().map(|_| DataType::Text).collect()
1930 } else {
1931 let probe_row = rows[0].as_row();
1932 let probe: &Row<'static> = &probe_row;
1933 group_exprs
1934 .iter()
1935 .map(|g| {
1936 eval::eval_expr(g, probe, &ctx).map(|v| v.data_type().unwrap_or(DataType::Text))
1937 })
1938 .collect::<Result<_, _>>()?
1939 };
1940 let agg_types: Vec<DataType> = agg_specs
1941 .iter()
1942 .map(|spec| infer_agg_type(spec, schema_cols))
1943 .collect();
1944 let mut synth_schema: Vec<ColumnSchema> = Vec::new();
1945 for (i, ty) in group_types.iter().enumerate() {
1946 synth_schema.push(ColumnSchema::new(format!("__grp_{i}"), *ty, true));
1947 }
1948 for (i, ty) in agg_types.iter().enumerate() {
1949 synth_schema.push(ColumnSchema::new(format!("__agg_{i}"), *ty, true));
1950 }
1951 Ok(synth_schema)
1952}
1953
1954/// (2b) Materialise one synthetic row per group (insertion order):
1955/// apply each aggregate's internal ORDER BY, then finalise the running
1956/// state into the group + aggregate cells.
1957/// v7.33 — compare two aggregate-internal ORDER BY key tuples under the
1958/// per-key DESC / NULLS directives. This is the exact comparator the
1959/// finalize sort uses, factored out so the `first_ordered` argmax
1960/// accumulator's "keep first" decision is provably identical to taking
1961/// element `[1]` of the fully-sorted array.
1962fn cmp_order_keys(
1963 order_by: &[spg_sql::ast::OrderBy],
1964 a: &[Value<'static>],
1965 b: &[Value<'static>],
1966) -> core::cmp::Ordering {
1967 for (k, o) in order_by.iter().enumerate() {
1968 let cmp = crate::order_by_value_cmp(o.desc, o.nulls_first, &a[k], &b[k]);
1969 if cmp != core::cmp::Ordering::Equal {
1970 return cmp;
1971 }
1972 }
1973 core::cmp::Ordering::Equal
1974}
1975
1976fn finalize_synth_rows(
1977 order: &[(Vec<Value<'static>>, Vec<AggState>)],
1978 agg_specs: &[AggSpec],
1979 synth_schema: &[ColumnSchema],
1980 rows: &[RowRef<'_>],
1981 schema_cols: &[ColumnSchema],
1982 table_alias: Option<&str>,
1983) -> Result<Vec<Row<'static>>, EvalError> {
1984 let ctx = EvalContext::new(schema_cols, table_alias);
1985 // v7.32 (round-29) — ordered-set direct arguments (the percentile
1986 // fraction) are constant per PG, so evaluate each once up front.
1987 let direct_arg_vals: Vec<Option<Value>> = agg_specs
1988 .iter()
1989 .map(|spec| match (&spec.direct_arg, rows.first()) {
1990 (Some(e), Some(r)) => eval::eval_expr(e, &r.as_row(), &ctx).map(Some),
1991 _ => Ok(None),
1992 })
1993 .collect::<Result<_, _>>()?;
1994
1995 // Materialise synthetic rows (insertion order = `order`).
1996 let mut synth_rows: Vec<Row<'static>> = Vec::new();
1997 for (gvals, states) in order {
1998 let mut values: Vec<Value<'static>> = Vec::with_capacity(synth_schema.len());
1999 values.extend(gvals.iter().cloned());
2000 for (i, st) in states.iter().enumerate() {
2001 // v7.33 (array_agg argmax) — first_ordered: the running
2002 // first-by-order value IS the result; no array build/sort.
2003 if agg_specs[i].first_ordered {
2004 values.push(
2005 st.first_best
2006 .as_ref()
2007 .map_or(Value::Null, |(_, v)| v.clone()),
2008 );
2009 continue;
2010 }
2011 // v7.24 (round-16 A) — order the collected items per the
2012 // aggregate-internal ORDER BY before finalize consumes
2013 // them.
2014 let st_sorted;
2015 let st_final: &AggState =
2016 if !agg_specs[i].order_by.is_empty() && st.item_keys.len() == st.items.len() {
2017 let mut idx: Vec<usize> = (0..st.items.len()).collect();
2018 let ob = &agg_specs[i].order_by;
2019 idx.sort_by(|&x, &y| cmp_order_keys(ob, &st.item_keys[x], &st.item_keys[y]));
2020 let mut sorted = st.clone();
2021 sorted.items = idx.iter().map(|&j| st.items[j].clone()).collect();
2022 st_sorted = sorted;
2023 &st_sorted
2024 } else {
2025 st
2026 };
2027 // Ordered-set aggregates compute from the sorted items + the
2028 // direct fraction; everything else uses the running state.
2029 let v = if is_within_group_name(&agg_specs[i].name) {
2030 finalize_ordered_set(
2031 &agg_specs[i].name,
2032 st_final,
2033 direct_arg_vals[i].as_ref(),
2034 agg_specs[i].order_by.first(),
2035 )
2036 } else {
2037 finalize(&agg_specs[i].name, st_final)
2038 };
2039 values.push(v);
2040 }
2041 synth_rows.push(Row::new(values));
2042 }
2043 Ok(synth_rows)
2044}
2045
2046/// (3) Rewrite the user's SELECT items + HAVING to reference the
2047/// synthetic columns, filter groups by HAVING, and project each
2048/// surviving group into an output row. The synth rows ride alongside
2049/// (`kept_synth`) so post-LIMIT deferred subqueries can evaluate later.
2050#[allow(clippy::too_many_lines)]
2051fn project_groups(
2052 synth_rows: Vec<Row<'static>>,
2053 stmt: &SelectStatement,
2054 group_exprs: &[Expr],
2055 agg_specs: &[AggSpec],
2056 synth_schema: &[ColumnSchema],
2057 correlated_eval: Option<CorrelatedEval<'_>>,
2058 defer_projection: bool,
2059) -> Result<Projection, EvalError> {
2060 // Rewrite the user's SELECT items + ORDER BY to reference synthetic
2061 // columns. After rewriting, every remaining `Expr::Column` must
2062 // resolve against the synthetic schema (i.e. must have been a GROUP
2063 // BY expression).
2064 let columns: Vec<ColumnSchema> = stmt
2065 .items
2066 .iter()
2067 .map(|item| match item {
2068 SelectItem::Wildcard => Err(EvalError::TypeMismatch {
2069 detail: "SELECT * with aggregates is not supported".into(),
2070 }),
2071 SelectItem::Expr { expr, alias } => {
2072 let rewritten = rewrite_expr(expr, group_exprs, agg_specs);
2073 let name = alias.clone().unwrap_or_else(|| expr.to_string());
2074 Ok(ColumnSchema::new(
2075 name,
2076 agg_or_group_type(&rewritten, synth_schema),
2077 true,
2078 ))
2079 }
2080 })
2081 .collect::<Result<_, _>>()?;
2082
2083 // Project per synthetic row. HAVING filters out groups *before*
2084 // we keep the projected row — same semantics as PG: HAVING runs
2085 // against the aggregated row (so `HAVING count(*) > 1` works) and
2086 // sees only group-by'd columns plus aggregate values.
2087 let synth_ctx = EvalContext::new(synth_schema, None);
2088 let having_rewritten = stmt
2089 .having
2090 .as_ref()
2091 .map(|h| rewrite_expr(h, group_exprs, agg_specs));
2092 // v7.30 (phase 3e-1) - rewrite SELECT items ONCE. This ran per
2093 // GROUP (23.5k x 9 items of AST cloning = ~48% of the inbox
2094 // query in sampled stacks); the rewrite is group-independent.
2095 // Stable addresses also let the per-expression subquery plans
2096 // (v7.29 3c) hit across groups instead of rebuilding.
2097 let items_rewritten: alloc::vec::Vec<Option<Expr>> = stmt
2098 .items
2099 .iter()
2100 .map(|item| match item {
2101 SelectItem::Expr { expr, .. } => Some(rewrite_expr(expr, group_exprs, agg_specs)),
2102 SelectItem::Wildcard => None,
2103 })
2104 .collect();
2105 // v7.31 (perf — PG lesson #1): subquery-bearing select items
2106 // deferred to post-LIMIT, when no sort/filter key can observe
2107 // them. ORDER BY rewrites are hoisted here so the safety check
2108 // and the sort below share one rewrite pass.
2109 let order_rewritten: Vec<Expr> = stmt
2110 .order_by
2111 .iter()
2112 .map(|o| rewrite_expr(&o.expr, group_exprs, agg_specs))
2113 .collect();
2114 let defer_enabled = correlated_eval.is_some()
2115 && !stmt.distinct
2116 && !having_rewritten
2117 .as_ref()
2118 .is_some_and(crate::expr_has_subquery)
2119 && !order_rewritten.iter().any(crate::expr_has_subquery);
2120 let deferred: Vec<(usize, Expr)> = if defer_enabled {
2121 items_rewritten
2122 .iter()
2123 .enumerate()
2124 .filter_map(|(i, r)| {
2125 r.as_ref()
2126 .filter(|e| crate::expr_has_subquery(e))
2127 .map(|e| (i, e.clone()))
2128 })
2129 .collect()
2130 } else {
2131 Vec::new()
2132 };
2133 // v7.32 (architecture v2, P2) — compile the per-group synth-row
2134 // expressions ONCE. The projection / HAVING here run per GROUP
2135 // (24k for the inbox shape) × per item; the rewritten exprs are
2136 // mostly `Column(__agg_N)` / `Column(__grp_K)` against the synth
2137 // schema — flat step programs, no tree walk per group.
2138 let having_compiled = having_rewritten
2139 .as_ref()
2140 .filter(|h| eval::fully_compilable(h))
2141 .map(|h| eval::compile_expr(h, &synth_ctx));
2142 let items_compiled: Vec<Option<eval::CompiledExpr>> = items_rewritten
2143 .iter()
2144 .enumerate()
2145 .map(|(i, r)| {
2146 r.as_ref()
2147 .filter(|e| !deferred.iter().any(|(c, _)| *c == i) && eval::fully_compilable(e))
2148 .map(|e| eval::compile_expr(e, &synth_ctx))
2149 })
2150 .collect();
2151 let mut kept_synth: Vec<Row<'static>> = Vec::new();
2152 let mut out_rows: Vec<Row<'static>> = Vec::new();
2153 let mut stack: Vec<Value<'static>> = Vec::new();
2154 for srow in synth_rows {
2155 if let Some(hc) = &having_compiled {
2156 let cond = eval::eval_compiled(hc, &srow, &synth_ctx, &mut stack)?;
2157 if !matches!(cond, Value::Bool(true)) {
2158 continue;
2159 }
2160 } else if let Some(h) = &having_rewritten {
2161 let cond = match correlated_eval {
2162 Some(f) if crate::expr_has_subquery(h) => f(h, &srow, &synth_ctx)?,
2163 _ => eval::eval_expr(h, &srow, &synth_ctx)?,
2164 };
2165 if !matches!(cond, Value::Bool(true)) {
2166 continue;
2167 }
2168 }
2169 // v7.37.x — when caller pre-truncates via ORDER BY+LIMIT, skip
2170 // per-item projection here; the caller fills the placeholder
2171 // out_rows from the top-K survivors below.
2172 if defer_projection {
2173 kept_synth.push(srow);
2174 out_rows.push(Row::new(Vec::new()));
2175 continue;
2176 }
2177 let mut values: Vec<Value<'static>> = Vec::with_capacity(columns.len());
2178 for (i, rewritten) in items_rewritten.iter().enumerate() {
2179 let Some(rewritten) = rewritten else { continue };
2180 if deferred.iter().any(|(c, _)| *c == i) {
2181 values.push(Value::Null);
2182 continue;
2183 }
2184 values.push(if let Some(cc) = &items_compiled[i] {
2185 eval::eval_compiled(cc, &srow, &synth_ctx, &mut stack)?
2186 } else {
2187 match correlated_eval {
2188 Some(f) if crate::expr_has_subquery(rewritten) => {
2189 f(rewritten, &srow, &synth_ctx)?
2190 }
2191 _ => eval::eval_expr(rewritten, &srow, &synth_ctx)?,
2192 }
2193 });
2194 }
2195 kept_synth.push(srow);
2196 out_rows.push(Row::new(values));
2197 }
2198 let deferred_project_state = if defer_projection {
2199 Some(DeferredProject {
2200 items_rewritten,
2201 items_compiled,
2202 })
2203 } else {
2204 None
2205 };
2206 Ok(Projection {
2207 columns,
2208 out_rows,
2209 kept_synth,
2210 deferred,
2211 order_rewritten,
2212 deferred_project: deferred_project_state,
2213 })
2214}
2215
2216/// (4) Sort the projected output by the rewritten ORDER BY keys. The
2217/// synth rows ride through the sort so deferred subqueries evaluate
2218/// against the surviving groups after the caller's LIMIT truncation.
2219fn sort_synth_by_order_by(
2220 synth_schema: &[ColumnSchema],
2221 order_by: &[spg_sql::ast::OrderBy],
2222 order_rewritten: &[Expr],
2223 mut kept_synth: Vec<Row<'static>>,
2224 mut out_rows: Vec<Row<'static>>,
2225 correlated_eval: Option<CorrelatedEval<'_>>,
2226 keep_n: Option<usize>,
2227) -> Result<(Vec<Row<'static>>, Vec<Row<'static>>), EvalError> {
2228 let synth_ctx = EvalContext::new(synth_schema, None);
2229 // v6.4.0 — multi-key ORDER BY on aggregate output. Each key
2230 // gets its own rewrite + per-key DESC flag. (Rewrites hoisted
2231 // above as `order_rewritten` — shared with the deferral
2232 // safety check.)
2233 let keys_meta: Vec<(bool, Option<bool>)> =
2234 order_by.iter().map(|o| (o.desc, o.nulls_first)).collect();
2235 // P2: compile order-by keys once (per-group sort keys are
2236 // the same `__agg_N` / `__grp_K` shape as the projection).
2237 let order_compiled: Vec<Option<eval::CompiledExpr>> = order_rewritten
2238 .iter()
2239 .map(|e| {
2240 Some(e)
2241 .filter(|e| eval::fully_compilable(e))
2242 .map(|e| eval::compile_expr(e, &synth_ctx))
2243 })
2244 .collect();
2245 // The synth row rides through the sort so deferred exprs can
2246 // evaluate against the surviving groups after the caller's
2247 // LIMIT truncation.
2248 let mut keystack: Vec<Value<'static>> = Vec::new();
2249 let mut tagged: Vec<(Vec<Value<'static>>, Row, Row)> = Vec::with_capacity(kept_synth.len());
2250 for (s, o) in kept_synth.into_iter().zip(out_rows) {
2251 let mut keys = Vec::with_capacity(order_rewritten.len());
2252 for (e, oc) in order_rewritten.iter().zip(&order_compiled) {
2253 keys.push(if let Some(oc) = oc {
2254 eval::eval_compiled(oc, &s, &synth_ctx, &mut keystack)?
2255 } else {
2256 match correlated_eval {
2257 Some(f) if crate::expr_has_subquery(e) => f(e, &s, &synth_ctx)?,
2258 _ => eval::eval_expr(e, &s, &synth_ctx)?,
2259 }
2260 });
2261 }
2262 tagged.push((keys, s, o));
2263 }
2264 let cmp = |a: &(Vec<Value<'static>>, Row, Row), b: &(Vec<Value<'static>>, Row, Row)| {
2265 use core::cmp::Ordering;
2266 for (i, (ka, kb)) in a.0.iter().zip(b.0.iter()).enumerate() {
2267 let (desc, nf) = keys_meta[i];
2268 let c = crate::order_by_value_cmp(desc, nf, ka, kb);
2269 if c != Ordering::Equal {
2270 return c;
2271 }
2272 }
2273 Ordering::Equal
2274 };
2275 // v7.37.3 — top-K partial sort when `keep_n` is small enough to
2276 // matter (`Some(k)` with `k < tagged.len()` and `k > 0`).
2277 // `select_nth_unstable_by` partitions in O(N), then we sort the
2278 // surviving prefix in O(K log K). Total = O(N + K log K) vs
2279 // O(N log N) the full sort would pay — matches the inbox-listing
2280 // shape PG uses.
2281 //
2282 match keep_n {
2283 Some(k) if k < tagged.len() && k > 0 => {
2284 let pivot = k - 1;
2285 tagged.select_nth_unstable_by(pivot, cmp);
2286 tagged[..k].sort_by(cmp);
2287 tagged.truncate(k);
2288 }
2289 _ => {
2290 tagged.sort_by(cmp);
2291 }
2292 }
2293 kept_synth = Vec::with_capacity(tagged.len());
2294 out_rows = Vec::with_capacity(tagged.len());
2295 for (_, s, o) in tagged {
2296 kept_synth.push(s);
2297 out_rows.push(o);
2298 }
2299 Ok((kept_synth, out_rows))
2300}
2301
2302/// v7.17.0 — walk the statement again to validate the positional
2303/// arity of every aggregate call site. Done after AST collection
2304/// rather than inside `collect_aggregates` so the collector stays
2305/// infallible; callers in `run()` can do a single early-error
2306/// exit before any per-row work.
2307fn validate_agg_arities(stmt: &SelectStatement, _specs: &[AggSpec]) -> Result<(), EvalError> {
2308 fn walk(e: &Expr) -> Result<(), EvalError> {
2309 if let Expr::FunctionCall { name, args } = e {
2310 let lower = name.to_ascii_lowercase();
2311 let expected: Option<usize> = match lower.as_str() {
2312 "count_star" => Some(0),
2313 "count" | "sum" | "avg" | "min" | "max" | "array_agg"
2314 // v7.17.0 — boolean aggregates also take exactly
2315 // one arg. `every` is an alias normalised inside
2316 // collect_aggregates / rewrite_expr.
2317 | "bool_and" | "bool_or" | "every"
2318 // v7.32 (round-29) — statistical + bitwise aggregates
2319 // + single-arg JSON aggregate.
2320 | "stddev" | "stddev_samp" | "stddev_pop"
2321 | "variance" | "var_samp" | "var_pop"
2322 | "bit_and" | "bit_or" | "bit_xor"
2323 | "json_agg" | "jsonb_agg" => Some(1),
2324 // v7.32 (round-29) — two-argument aggregates: string_agg,
2325 // the regression family f(Y, X), and json_object_agg.
2326 "string_agg"
2327 | "covar_pop" | "covar_samp" | "corr"
2328 | "regr_count" | "regr_avgx" | "regr_avgy" | "regr_slope"
2329 | "regr_intercept" | "regr_r2" | "regr_sxx" | "regr_syy" | "regr_sxy"
2330 | "json_object_agg" | "jsonb_object_agg" => Some(2),
2331 _ => None,
2332 };
2333 if let Some(want) = expected
2334 && args.len() != want
2335 {
2336 return Err(EvalError::TypeMismatch {
2337 detail: alloc::format!("{lower}() takes {want} arg(s), got {}", args.len()),
2338 });
2339 }
2340 for a in args {
2341 walk(a)?;
2342 }
2343 } else if let Expr::Binary { lhs, rhs, .. } = e {
2344 walk(lhs)?;
2345 walk(rhs)?;
2346 } else if let Expr::Unary { expr, .. }
2347 | Expr::Cast { expr, .. }
2348 | Expr::IsNull { expr, .. } = e
2349 {
2350 walk(expr)?;
2351 }
2352 Ok(())
2353 }
2354 for item in &stmt.items {
2355 if let SelectItem::Expr { expr, .. } = item {
2356 walk(expr)?;
2357 }
2358 }
2359 for o in &stmt.order_by {
2360 walk(&o.expr)?;
2361 }
2362 if let Some(h) = &stmt.having {
2363 walk(h)?;
2364 }
2365 Ok(())
2366}
2367
2368/// v7.33 (array_agg argmax) — recognise `(array_agg(x ORDER BY y))[1]`,
2369/// the argmax/argmin idiom: a non-DISTINCT ordered `array_agg`
2370/// subscripted by the constant 1. Returns `(value_arg, order_by,
2371/// filter)` on a match. When matched, the whole per-group array build +
2372/// sort + materialise is replaced by a running first-by-order scalar
2373/// accumulator and the subscript node is consumed (replaced by the
2374/// synthetic column). collect_aggregates and rewrite_expr share this one
2375/// matcher so their `__agg_<i>` assignment stays in lockstep.
2376fn first_ordered_array_agg(e: &Expr) -> Option<(&Expr, &[spg_sql::ast::OrderBy], Option<&Expr>)> {
2377 let Expr::ArraySubscript { target, index } = e else {
2378 return None;
2379 };
2380 if !matches!(
2381 index.as_ref(),
2382 Expr::Literal(spg_sql::ast::Literal::Integer(1))
2383 ) {
2384 return None;
2385 }
2386 let Expr::AggregateOrdered {
2387 call,
2388 order_by,
2389 distinct,
2390 filter,
2391 } = target.as_ref()
2392 else {
2393 return None;
2394 };
2395 if *distinct || order_by.is_empty() {
2396 return None;
2397 }
2398 let Expr::FunctionCall { name, args } = call.as_ref() else {
2399 return None;
2400 };
2401 if !name.eq_ignore_ascii_case("array_agg") || args.len() != 1 {
2402 return None;
2403 }
2404 Some((&args[0], order_by, filter.as_deref()))
2405}
2406
2407fn collect_aggregates(e: &Expr, out: &mut Vec<AggSpec>) {
2408 match e {
2409 // v7.24 (round-16 A) — ordered aggregate: register the inner
2410 // call's spec with the ordering attached.
2411 Expr::AggregateOrdered {
2412 call,
2413 order_by,
2414 distinct,
2415 filter,
2416 } => {
2417 if let Expr::FunctionCall { name, args } = call.as_ref() {
2418 let lower = name.to_ascii_lowercase();
2419 if is_aggregate_name(&lower) {
2420 let canonical = if lower == "every" {
2421 "bool_and".to_string()
2422 } else {
2423 lower
2424 };
2425 // Ordered-set aggregates (`percentile_cont(f)
2426 // WITHIN GROUP (ORDER BY x)`) take the value to
2427 // aggregate from the sort spec and the in-parens
2428 // arg as the direct (fraction) argument.
2429 let ordered_set = is_within_group_name(&canonical);
2430 let (arg, direct_arg) = if ordered_set {
2431 (
2432 order_by.first().map(|o| o.expr.clone()),
2433 args.first().cloned(),
2434 )
2435 } else {
2436 (args.first().cloned(), None)
2437 };
2438 let spec = AggSpec {
2439 kind: classify_agg_name(&canonical),
2440 name: canonical.clone(),
2441 arg,
2442 arg2: if agg_uses_second_arg(&canonical) {
2443 args.get(1).cloned()
2444 } else {
2445 None
2446 },
2447 distinct: *distinct,
2448 order_by: order_by.clone(),
2449 filter: filter.as_deref().cloned(),
2450 direct_arg,
2451 first_ordered: false,
2452 };
2453 if !out.iter().any(|s| {
2454 s.name == spec.name
2455 && s.arg == spec.arg
2456 && s.arg2 == spec.arg2
2457 && s.distinct == spec.distinct
2458 && s.order_by == spec.order_by
2459 && s.filter == spec.filter
2460 && s.direct_arg == spec.direct_arg
2461 && s.first_ordered == spec.first_ordered
2462 }) {
2463 out.push(spec);
2464 }
2465 return;
2466 }
2467 }
2468 collect_aggregates(call, out);
2469 for o in order_by {
2470 collect_aggregates(&o.expr, out);
2471 }
2472 }
2473 Expr::FunctionCall { name, args } => {
2474 let lower = name.to_ascii_lowercase();
2475 if is_aggregate_name(&lower) {
2476 let arg = if lower == "count_star" {
2477 None
2478 } else {
2479 args.first().cloned()
2480 };
2481 // v7.17.0 — second positional arg for
2482 // `string_agg(value, separator)`; v7.32 — also the
2483 // regression family `f(Y, X)` and `json_object_agg`.
2484 let arg2 = if agg_uses_second_arg(&lower) {
2485 args.get(1).cloned()
2486 } else {
2487 None
2488 };
2489 // v7.17.0 — `every` is the SQL-standard alias for
2490 // `bool_and`; collapse at collection time so
2491 // update_state / finalize need only one arm.
2492 let canonical = if lower == "every" {
2493 "bool_and".to_string()
2494 } else {
2495 lower
2496 };
2497 let spec = AggSpec {
2498 kind: classify_agg_name(&canonical),
2499 name: canonical,
2500 arg: arg.clone(),
2501 arg2: arg2.clone(),
2502 distinct: false,
2503 order_by: Vec::new(),
2504 filter: None,
2505 direct_arg: None,
2506 first_ordered: false,
2507 };
2508 if !out.iter().any(|s| {
2509 s.name == spec.name
2510 && s.arg == spec.arg
2511 && s.arg2 == spec.arg2
2512 && !s.distinct
2513 && s.order_by == spec.order_by
2514 && s.filter.is_none()
2515 && !s.first_ordered
2516 }) {
2517 out.push(spec);
2518 }
2519 // Don't recurse into the arg — nested aggregates are
2520 // illegal in standard SQL.
2521 } else {
2522 for a in args {
2523 collect_aggregates(a, out);
2524 }
2525 }
2526 }
2527 Expr::Binary { lhs, rhs, .. } => {
2528 collect_aggregates(lhs, out);
2529 collect_aggregates(rhs, out);
2530 }
2531 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
2532 collect_aggregates(expr, out);
2533 }
2534 Expr::Like { expr, pattern, .. } => {
2535 collect_aggregates(expr, out);
2536 collect_aggregates(pattern, out);
2537 }
2538 Expr::InList { expr, list, .. } => {
2539 collect_aggregates(expr, out);
2540 for item in list {
2541 collect_aggregates(item, out);
2542 }
2543 }
2544 Expr::Extract { source, .. } => collect_aggregates(source, out),
2545 // v4.10 subquery + v4.12 window / Literal / Column —
2546 // non-recursing leaves for the aggregate collector.
2547 Expr::ScalarSubquery(_)
2548 | Expr::Exists { .. }
2549 | Expr::InSubquery { .. }
2550 | Expr::WindowFunction { .. }
2551 | Expr::Literal(_)
2552 | Expr::Placeholder(_)
2553 | Expr::Column(_) => {}
2554 // v7.10.10 — recurse into array constructor children +
2555 // subscript / ANY/ALL operands.
2556 Expr::Array(items) => {
2557 for elem in items {
2558 collect_aggregates(elem, out);
2559 }
2560 }
2561 Expr::ArraySubscript { target, index } => {
2562 // v7.33 (array_agg argmax) — `(array_agg(x ORDER BY y))[1]`
2563 // collects as a first_ordered spec; the subscript is consumed
2564 // here (do NOT recurse into the array_agg, or it would also
2565 // register a plain full-array spec).
2566 if let Some((arg, order_by, filter)) = first_ordered_array_agg(e) {
2567 let spec = AggSpec {
2568 kind: AggKind::ArrayAgg,
2569 name: "array_agg".to_string(),
2570 arg: Some(arg.clone()),
2571 arg2: None,
2572 distinct: false,
2573 order_by: order_by.to_vec(),
2574 filter: filter.cloned(),
2575 direct_arg: None,
2576 first_ordered: true,
2577 };
2578 if !out.iter().any(|s| {
2579 s.name == spec.name
2580 && s.arg == spec.arg
2581 && s.order_by == spec.order_by
2582 && s.filter == spec.filter
2583 && s.first_ordered
2584 }) {
2585 out.push(spec);
2586 }
2587 return;
2588 }
2589 collect_aggregates(target, out);
2590 collect_aggregates(index, out);
2591 }
2592 Expr::AnyAll { expr, array, .. } => {
2593 collect_aggregates(expr, out);
2594 collect_aggregates(array, out);
2595 }
2596 Expr::Case {
2597 operand,
2598 branches,
2599 else_branch,
2600 } => {
2601 if let Some(o) = operand {
2602 collect_aggregates(o, out);
2603 }
2604 for (w, t) in branches {
2605 collect_aggregates(w, out);
2606 collect_aggregates(t, out);
2607 }
2608 if let Some(e) = else_branch {
2609 collect_aggregates(e, out);
2610 }
2611 }
2612 }
2613}
2614
2615fn update_state(
2616 st: &mut AggState,
2617 kind: AggKind,
2618 name: &str,
2619 v: &Value<'_>,
2620 arg2: Option<&Value<'_>>,
2621 order_keys: Option<Vec<Value<'static>>>,
2622) -> Result<(), EvalError> {
2623 let is_null = matches!(v, Value::Null);
2624 // v7.37.4 (R34) — dispatch by pre-classified `kind` (`Copy`
2625 // enum), not by per-row string match. Hot inner loop on
2626 // multi-aggregate queries (mailrs `/api/conversations`: 14
2627 // aggregates × 100 k rows = 1.4 M dispatches) sees an enum
2628 // jump table instead of a sequence of `eq_str` checks. `name`
2629 // is still threaded through for error messages so the user-
2630 // facing wording is unchanged.
2631 match kind {
2632 AggKind::CountStar => st.count += 1,
2633 AggKind::Count => {
2634 if !is_null {
2635 st.count += 1;
2636 }
2637 }
2638 AggKind::Sum | AggKind::Avg => {
2639 if is_null {
2640 return Ok(());
2641 }
2642 st.count += 1;
2643 match v {
2644 Value::Int(n) => st.sum_int += i64::from(*n),
2645 Value::BigInt(n) => st.sum_int += *n,
2646 Value::Float(x) => {
2647 st.use_float = true;
2648 st.sum_float += *x;
2649 }
2650 other => {
2651 return Err(EvalError::TypeMismatch {
2652 detail: format!("sum/avg need numeric, got {:?}", other.data_type()),
2653 });
2654 }
2655 }
2656 }
2657 AggKind::Min => {
2658 if is_null {
2659 return Ok(());
2660 }
2661 match &st.extreme {
2662 None => st.extreme = Some(v.clone().into_owned()),
2663 Some(cur) => {
2664 if value_cmp(v, cur) == core::cmp::Ordering::Less {
2665 st.extreme = Some(v.clone().into_owned());
2666 }
2667 }
2668 }
2669 }
2670 AggKind::Max => {
2671 if is_null {
2672 return Ok(());
2673 }
2674 match &st.extreme {
2675 None => st.extreme = Some(v.clone().into_owned()),
2676 Some(cur) => {
2677 if value_cmp(v, cur) == core::cmp::Ordering::Greater {
2678 st.extreme = Some(v.clone().into_owned());
2679 }
2680 }
2681 }
2682 }
2683 // v7.17.0 — string_agg(value, separator). NULL value is
2684 // skipped (PG aggregate-skip-null). Separator captured
2685 // from the latest row that flows through; matches PG's
2686 // semantics of evaluating the separator per row but using
2687 // the last value at finalize time (in practice it's
2688 // constant). count is bumped so we can distinguish "empty
2689 // group → NULL" from "all-NULL group → NULL".
2690 AggKind::StringAgg => {
2691 if let Some(sep) = arg2
2692 && let Value::Text(s) = sep
2693 {
2694 st.separator = Some(s.to_string());
2695 }
2696 if is_null {
2697 return Ok(());
2698 }
2699 if let Value::Text(s) = v {
2700 st.items.push(Value::text(s.clone()));
2701 if let Some(k) = order_keys {
2702 st.item_keys.push(k);
2703 }
2704 st.count += 1;
2705 } else {
2706 return Err(EvalError::TypeMismatch {
2707 detail: format!("string_agg requires text value, got {:?}", v.data_type()),
2708 });
2709 }
2710 }
2711 // v7.17.0 — array_agg(value). Unlike string_agg, NULL
2712 // elements are KEPT in the array (PG behaviour); the
2713 // result is NULL only when ZERO rows fed in. Element type
2714 // is locked from the first row's value type; subsequent
2715 // rows must match (PG also rejects mixed-type array_agg).
2716 AggKind::ArrayAgg => {
2717 st.items.push(v.clone().into_owned());
2718 if let Some(k) = order_keys {
2719 st.item_keys.push(k);
2720 }
2721 st.count += 1;
2722 }
2723 // v7.17.0 — bool_and(p): TRUE iff every non-NULL input is
2724 // TRUE. NULL skipped; running accumulator stays at TRUE
2725 // until the first non-NULL FALSE.
2726 AggKind::BoolAnd => {
2727 if is_null {
2728 return Ok(());
2729 }
2730 let b = match v {
2731 Value::Bool(b) => *b,
2732 other => {
2733 return Err(EvalError::TypeMismatch {
2734 detail: format!("bool_and requires bool, got {:?}", other.data_type()),
2735 });
2736 }
2737 };
2738 st.bool_acc = Some(st.bool_acc.map_or(b, |acc| acc && b));
2739 }
2740 // v7.17.0 — bool_or(p): TRUE iff any non-NULL input is
2741 // TRUE. NULL skipped.
2742 AggKind::BoolOr => {
2743 if is_null {
2744 return Ok(());
2745 }
2746 let b = match v {
2747 Value::Bool(b) => *b,
2748 other => {
2749 return Err(EvalError::TypeMismatch {
2750 detail: format!("bool_or requires bool, got {:?}", other.data_type()),
2751 });
2752 }
2753 };
2754 st.bool_acc = Some(st.bool_acc.map_or(b, |acc| acc || b));
2755 }
2756 // v7.32 (round-29) — variance / stddev family. Accumulate the
2757 // running sum (sum_float) and sum of squares (sum_sq) over the
2758 // non-NULL numeric inputs; finalize divides by n or n-1.
2759 AggKind::StddevFamily => {
2760 if is_null {
2761 return Ok(());
2762 }
2763 let x = match v {
2764 Value::Int(n) => f64::from(*n),
2765 Value::SmallInt(n) => f64::from(*n),
2766 Value::BigInt(n) => *n as f64,
2767 Value::Float(x) => *x,
2768 other => {
2769 return Err(EvalError::TypeMismatch {
2770 detail: format!("{name} needs numeric, got {:?}", other.data_type()),
2771 });
2772 }
2773 };
2774 st.count += 1;
2775 st.sum_float += x;
2776 st.sum_sq += x * x;
2777 }
2778 // v7.32 (round-29) — bitwise aggregates over integer inputs.
2779 AggKind::BitAnd | AggKind::BitOr | AggKind::BitXor => {
2780 if is_null {
2781 return Ok(());
2782 }
2783 let n = match v {
2784 Value::Int(n) => i64::from(*n),
2785 Value::SmallInt(n) => i64::from(*n),
2786 Value::BigInt(n) => *n,
2787 other => {
2788 return Err(EvalError::TypeMismatch {
2789 detail: format!("{name} needs integer, got {:?}", other.data_type()),
2790 });
2791 }
2792 };
2793 st.bit_acc = Some(match (st.bit_acc, kind) {
2794 (None, _) => n,
2795 (Some(acc), AggKind::BitAnd) => acc & n,
2796 (Some(acc), AggKind::BitOr) => acc | n,
2797 (Some(acc), _) => acc ^ n, // BitXor
2798 });
2799 }
2800 // v7.32 (round-29) — WITHIN GROUP aggregates (ordered-set +
2801 // hypothetical-set) collect the sort value (NULLs ignored, per
2802 // PG) into `items`, sorted at finalize by the parallel
2803 // `item_keys`.
2804 AggKind::WithinGroup => {
2805 if is_null {
2806 return Ok(());
2807 }
2808 st.items.push(v.clone().into_owned());
2809 if let Some(k) = order_keys {
2810 st.item_keys.push(k);
2811 }
2812 st.count += 1;
2813 }
2814 // v7.32 (round-29) — regression family f(Y, X). Only rows with
2815 // BOTH inputs non-NULL contribute (PG semantics). `v` is Y,
2816 // `arg2` is X.
2817 AggKind::Regression => {
2818 let (Some(y), Some(x)) = (agg_value_to_f64(v), arg2.and_then(agg_value_to_f64)) else {
2819 return Ok(()); // NULL (or non-numeric) in either input
2820 };
2821 st.reg_n += 1;
2822 st.reg_sx += x;
2823 st.reg_sy += y;
2824 st.reg_sxx += x * x;
2825 st.reg_syy += y * y;
2826 st.reg_sxy += x * y;
2827 }
2828 // v7.32 (round-29) — json_agg / jsonb_agg collect every input
2829 // (NULL becomes JSON null, per PG) in row order.
2830 AggKind::JsonAgg => {
2831 st.items.push(v.clone().into_owned());
2832 st.count += 1;
2833 }
2834 // v7.32 (round-29) — json_object_agg(key, value): keys in
2835 // `items`, values in `aux_items`. A NULL key is skipped (PG
2836 // raises; we drop it rather than abort the whole query).
2837 AggKind::JsonObjectAgg => {
2838 if is_null {
2839 return Ok(());
2840 }
2841 st.items.push(v.clone().into_owned());
2842 st.aux_items
2843 .push(arg2.cloned().map(Value::into_owned).unwrap_or(Value::Null));
2844 st.count += 1;
2845 }
2846 }
2847 Ok(())
2848}
2849
2850#[allow(clippy::cast_precision_loss)]
2851fn finalize(name: &str, st: &AggState) -> Value<'static> {
2852 match name {
2853 "count" | "count_star" => Value::BigInt(st.count),
2854 "sum" => {
2855 if st.count == 0 {
2856 Value::Null
2857 } else if st.use_float {
2858 Value::Float(st.sum_float + (st.sum_int as f64))
2859 } else {
2860 Value::BigInt(st.sum_int)
2861 }
2862 }
2863 "avg" => {
2864 if st.count == 0 {
2865 Value::Null
2866 } else {
2867 let total = if st.use_float {
2868 st.sum_float + (st.sum_int as f64)
2869 } else {
2870 st.sum_int as f64
2871 };
2872 Value::Float(total / (st.count as f64))
2873 }
2874 }
2875 "min" | "max" => st.extreme.clone().unwrap_or(Value::Null),
2876 // v7.17.0 — string_agg: join all collected text items with
2877 // the captured separator. Empty / all-NULL group → NULL
2878 // (PG semantics).
2879 "string_agg" => {
2880 if st.items.is_empty() {
2881 return Value::Null;
2882 }
2883 let sep = st.separator.clone().unwrap_or_default();
2884 let mut out = String::new();
2885 for (i, item) in st.items.iter().enumerate() {
2886 if i > 0 {
2887 out.push_str(&sep);
2888 }
2889 if let Value::Text(s) = item {
2890 out.push_str(s);
2891 }
2892 }
2893 Value::text(out)
2894 }
2895 // v7.17.0 — array_agg: collect into a typed array. NULL
2896 // elements are preserved per PG. Result type is decided
2897 // by the first non-NULL element seen (or Text fallback
2898 // when the whole group is NULL — PG would surface the
2899 // declared input type, but SPG hasn't yet wired the
2900 // aggregate's static input-type from `describe`).
2901 "array_agg" => {
2902 if st.items.is_empty() {
2903 return Value::Null;
2904 }
2905 let probe = st.items.iter().find(|v| !v.is_null());
2906 match probe.and_then(spg_storage::Value::data_type) {
2907 Some(DataType::Int) | Some(DataType::SmallInt) => {
2908 let items: Vec<Option<i32>> = st
2909 .items
2910 .iter()
2911 .map(|v| match v {
2912 Value::Int(n) => Some(*n),
2913 Value::SmallInt(n) => Some(i32::from(*n)),
2914 _ => None,
2915 })
2916 .collect();
2917 Value::IntArray(items)
2918 }
2919 Some(DataType::BigInt) => {
2920 let items: Vec<Option<i64>> = st
2921 .items
2922 .iter()
2923 .map(|v| match v {
2924 Value::BigInt(n) => Some(*n),
2925 _ => None,
2926 })
2927 .collect();
2928 Value::BigIntArray(items)
2929 }
2930 _ => {
2931 let items: Vec<Option<String>> = st
2932 .items
2933 .iter()
2934 .map(|v| match v {
2935 Value::Text(s) => Some(s.to_string()),
2936 Value::Null => None,
2937 other => Some(format!("{other:?}")),
2938 })
2939 .collect();
2940 Value::TextArray(items)
2941 }
2942 }
2943 }
2944 // v7.17.0 — bool_and / bool_or finalize: lazy-init pattern
2945 // means `None` is exactly "empty group or all-NULL", which
2946 // PG surfaces as SQL NULL.
2947 "bool_and" | "bool_or" => st.bool_acc.map_or(Value::Null, Value::Bool),
2948 // v7.32 (round-29) — variance / stddev. PG: `variance` ==
2949 // `var_samp`, `stddev` == `stddev_samp`. samp needs n >= 2
2950 // (n < 2 → NULL); pop needs n >= 1 (n == 1 → 0).
2951 "variance" | "var_samp" | "var_pop" | "stddev" | "stddev_samp" | "stddev_pop" => {
2952 let n = st.count;
2953 if n == 0 {
2954 return Value::Null;
2955 }
2956 let nf = n as f64;
2957 // Sum of squared deviations from the mean.
2958 let ss = st.sum_sq - (st.sum_float * st.sum_float) / nf;
2959 let pop = name.ends_with("_pop");
2960 let denom = if pop { nf } else { nf - 1.0 };
2961 if denom <= 0.0 {
2962 // var_samp / stddev (samp) with n == 1 → NULL.
2963 return Value::Null;
2964 }
2965 let var = (ss / denom).max(0.0); // clamp fp noise below 0
2966 if name.starts_with("stddev") {
2967 Value::Float(crate::eval::f64_sqrt(var))
2968 } else {
2969 Value::Float(var)
2970 }
2971 }
2972 // v7.32 (round-29) — bitwise aggregates: None (empty / all-NULL)
2973 // → SQL NULL.
2974 "bit_and" | "bit_or" | "bit_xor" => st.bit_acc.map_or(Value::Null, Value::BigInt),
2975 // v7.32 (round-29) — regression family. `regr_count` is the
2976 // paired n; everything else is NULL over an empty set. Terms
2977 // are the mean-centred sums of squares / cross-products.
2978 "regr_count" => Value::BigInt(st.reg_n),
2979 "covar_pop" | "covar_samp" | "corr" | "regr_avgx" | "regr_avgy" | "regr_slope"
2980 | "regr_intercept" | "regr_r2" | "regr_sxx" | "regr_syy" | "regr_sxy" => {
2981 let n = st.reg_n;
2982 if n == 0 {
2983 return Value::Null;
2984 }
2985 let nf = n as f64;
2986 let sxx = st.reg_sxx - st.reg_sx * st.reg_sx / nf;
2987 let syy = st.reg_syy - st.reg_sy * st.reg_sy / nf;
2988 let sxy = st.reg_sxy - st.reg_sx * st.reg_sy / nf;
2989 let avgx = st.reg_sx / nf;
2990 let avgy = st.reg_sy / nf;
2991 let out = match name {
2992 "regr_avgx" => Some(avgx),
2993 "regr_avgy" => Some(avgy),
2994 "regr_sxx" => Some(sxx),
2995 "regr_syy" => Some(syy),
2996 "regr_sxy" => Some(sxy),
2997 "covar_pop" => Some(sxy / nf),
2998 "covar_samp" => (n >= 2).then(|| sxy / (nf - 1.0)),
2999 "regr_slope" => (sxx != 0.0).then(|| sxy / sxx),
3000 "regr_intercept" => (sxx != 0.0).then(|| avgy - (sxy / sxx) * avgx),
3001 "corr" => {
3002 let d = sxx * syy;
3003 (d > 0.0).then(|| sxy / crate::eval::f64_sqrt(d))
3004 }
3005 // PG: NULL when sxx==0; 1 when syy==0 (and sxx>0).
3006 "regr_r2" => {
3007 if sxx == 0.0 {
3008 None
3009 } else if syy == 0.0 {
3010 Some(1.0)
3011 } else {
3012 Some((sxy * sxy) / (sxx * syy))
3013 }
3014 }
3015 _ => None,
3016 };
3017 out.map_or(Value::Null, Value::Float)
3018 }
3019 // v7.32 (round-29) — json_agg / jsonb_agg: a JSON array of every
3020 // collected element in row order; empty set → SQL NULL.
3021 "json_agg" | "jsonb_agg" => {
3022 if st.items.is_empty() {
3023 return Value::Null;
3024 }
3025 let mut out = String::from("[");
3026 for (i, item) in st.items.iter().enumerate() {
3027 if i > 0 {
3028 out.push_str(", ");
3029 }
3030 out.push_str(&crate::json::value_to_json_text(item));
3031 }
3032 out.push(']');
3033 Value::json(out)
3034 }
3035 // v7.32 (round-29) — json_object_agg: a JSON object built from
3036 // the parallel key (`items`) / value (`aux_items`) streams.
3037 "json_object_agg" | "jsonb_object_agg" => {
3038 if st.items.is_empty() {
3039 return Value::Null;
3040 }
3041 let mut out = String::from("{");
3042 for (i, key) in st.items.iter().enumerate() {
3043 if i > 0 {
3044 out.push_str(", ");
3045 }
3046 // Object keys are always JSON strings (PG coerces).
3047 let key_text = match key {
3048 Value::Text(s) | Value::Json(s) => s.to_string(),
3049 other => crate::json::value_to_json_text(other),
3050 };
3051 out.push_str(&crate::json::value_to_json_text(&Value::text(key_text)));
3052 out.push_str(": ");
3053 let val = st.aux_items.get(i).unwrap_or(&Value::Null);
3054 out.push_str(&crate::json::value_to_json_text(val));
3055 }
3056 out.push('}');
3057 Value::json(out)
3058 }
3059 // Ordered-set aggregates are finalized in `run` (they need the
3060 // sorted items + the direct fraction argument), never here.
3061 _ => unreachable!(),
3062 }
3063}
3064
3065/// v7.32 (round-29) — numeric coercion for the percentile interpolation.
3066fn agg_value_to_f64(v: &Value) -> Option<f64> {
3067 match v {
3068 Value::Int(n) => Some(f64::from(*n)),
3069 Value::SmallInt(n) => Some(f64::from(*n)),
3070 Value::BigInt(n) => Some(*n as f64),
3071 Value::Float(x) => Some(*x),
3072 _ => None,
3073 }
3074}
3075
3076/// v7.32 (round-29) — finalize a WITHIN GROUP aggregate. `st.items` is
3077/// already sorted by the `WITHIN GROUP (ORDER BY …)` spec. `direct` is
3078/// the evaluated direct argument: the fraction for `percentile_*`, the
3079/// hypothetical value for the hypothetical-set family (`rank` etc.),
3080/// and unused by `mode`. `order` is the (single) sort key, needed by
3081/// the hypothetical-set family to compare in the sort direction.
3082#[allow(
3083 clippy::cast_precision_loss,
3084 clippy::cast_possible_truncation,
3085 clippy::cast_sign_loss
3086)]
3087fn finalize_ordered_set(
3088 name: &str,
3089 st: &AggState,
3090 direct: Option<&Value>,
3091 order: Option<&spg_sql::ast::OrderBy>,
3092) -> Value<'static> {
3093 let fraction = direct;
3094 let items = &st.items;
3095 if items.is_empty() {
3096 // A hypothetical row ranks first over an empty group; the
3097 // distribution functions are 0 / divide-by-(n+1).
3098 return match name {
3099 "rank" | "dense_rank" => Value::BigInt(1),
3100 "percent_rank" => Value::Float(0.0),
3101 "cume_dist" => Value::Float(1.0),
3102 _ => Value::Null,
3103 };
3104 }
3105 let n = items.len();
3106 match name {
3107 // v7.32 (round-29) — hypothetical-set: the rank the direct value
3108 // would have if inserted into the group, in the sort direction.
3109 "rank" | "dense_rank" | "percent_rank" | "cume_dist" => {
3110 let Some(h) = fraction else {
3111 return Value::Null;
3112 };
3113 let (desc, nulls_first) = order.map_or((false, None), |o| (o.desc, o.nulls_first));
3114 let mut before = 0usize; // sort strictly before h
3115 let mut before_or_eq = 0usize; // sort before-or-peer with h
3116 let mut distinct_before = 0usize;
3117 let mut last_before: Option<&Value> = None;
3118 for it in items {
3119 match crate::order_by_value_cmp(desc, nulls_first, it, h) {
3120 core::cmp::Ordering::Less => {
3121 before += 1;
3122 before_or_eq += 1;
3123 if last_before
3124 .is_none_or(|p| value_cmp(p, it) != core::cmp::Ordering::Equal)
3125 {
3126 distinct_before += 1;
3127 last_before = Some(it);
3128 }
3129 }
3130 core::cmp::Ordering::Equal => before_or_eq += 1,
3131 core::cmp::Ordering::Greater => {}
3132 }
3133 }
3134 let nn = n as f64;
3135 match name {
3136 "rank" => Value::BigInt((before + 1) as i64),
3137 "dense_rank" => Value::BigInt((distinct_before + 1) as i64),
3138 "percent_rank" => Value::Float(before as f64 / nn),
3139 "cume_dist" => Value::Float((before_or_eq as f64 + 1.0) / (nn + 1.0)),
3140 _ => unreachable!(),
3141 }
3142 }
3143 // Most frequent value; equal values are adjacent in the sorted
3144 // run, and a frequency tie resolves to the earliest run (the
3145 // smallest value under an ascending sort), matching PG.
3146 "mode" => {
3147 let (mut best_i, mut best_cnt) = (0usize, 1usize);
3148 let (mut run_i, mut run_cnt) = (0usize, 1usize);
3149 for i in 1..n {
3150 if value_cmp(&items[i], &items[run_i]) == core::cmp::Ordering::Equal {
3151 run_cnt += 1;
3152 } else {
3153 run_i = i;
3154 run_cnt = 1;
3155 }
3156 if run_cnt > best_cnt {
3157 best_cnt = run_cnt;
3158 best_i = run_i;
3159 }
3160 }
3161 items[best_i].clone()
3162 }
3163 // The first value whose cumulative fraction reaches `f`.
3164 "percentile_disc" => {
3165 let f = fraction
3166 .and_then(agg_value_to_f64)
3167 .unwrap_or(0.0)
3168 .clamp(0.0, 1.0);
3169 let idx = if f <= 0.0 {
3170 0
3171 } else {
3172 (crate::eval::f64_ceil(f * n as f64) as usize)
3173 .saturating_sub(1)
3174 .min(n - 1)
3175 };
3176 items[idx].clone()
3177 }
3178 // Linear interpolation between the two bracketing values.
3179 "percentile_cont" => {
3180 let f = fraction
3181 .and_then(agg_value_to_f64)
3182 .unwrap_or(0.0)
3183 .clamp(0.0, 1.0);
3184 let Some(nums) = items
3185 .iter()
3186 .map(agg_value_to_f64)
3187 .collect::<Option<Vec<f64>>>()
3188 else {
3189 return Value::Null; // non-numeric ordered set
3190 };
3191 if n == 1 {
3192 return Value::Float(nums[0]);
3193 }
3194 let rank = f * (n as f64 - 1.0);
3195 let lo = crate::eval::f64_floor(rank) as usize;
3196 let hi = crate::eval::f64_ceil(rank) as usize;
3197 let frac = rank - lo as f64;
3198 Value::Float(nums[lo] + (nums[hi] - nums[lo]) * frac)
3199 }
3200 _ => unreachable!(),
3201 }
3202}
3203
3204fn infer_agg_type(spec: &AggSpec, schema_cols: &[ColumnSchema]) -> DataType {
3205 // v7.26 (round-20 C) — the argument's statically-derived shape
3206 // types MIN/MAX/SUM/array_agg properly; RowDescription used to
3207 // report TEXT for these, breaking every sqlx typed decode.
3208 let arg_ty = spec
3209 .arg
3210 .as_ref()
3211 .and_then(|a| crate::describe::describe_expr(a, schema_cols))
3212 .map(|shape| shape.ty);
3213 // v7.33 (array_agg argmax) — `(array_agg(x ORDER BY y))[1]` yields the
3214 // ELEMENT type (x), not the array type.
3215 if spec.first_ordered {
3216 return arg_ty.unwrap_or(DataType::Text);
3217 }
3218 match spec.name.as_str() {
3219 "count" | "count_star" => DataType::BigInt,
3220 "sum" => match arg_ty {
3221 Some(DataType::Float) => DataType::Float,
3222 _ => DataType::BigInt,
3223 },
3224 "avg" => DataType::Float,
3225 // v7.17.0 — string_agg always returns TEXT.
3226 "string_agg" => DataType::Text,
3227 "array_agg" => match arg_ty {
3228 Some(DataType::Int | DataType::SmallInt) => DataType::IntArray,
3229 Some(DataType::BigInt) => DataType::BigIntArray,
3230 _ => DataType::TextArray,
3231 },
3232 // v7.17.0 — boolean aggregates always return BOOL (nullable
3233 // — empty / all-NULL group → NULL).
3234 "bool_and" | "bool_or" => DataType::Bool,
3235 // v7.32 (round-29) — variance / stddev are floating point;
3236 // percentile_cont interpolates to float; the regression family
3237 // (except regr_count) is floating point.
3238 "stddev" | "stddev_samp" | "stddev_pop" | "variance" | "var_samp" | "var_pop"
3239 | "percentile_cont" | "covar_pop" | "covar_samp" | "corr" | "regr_avgx" | "regr_avgy"
3240 | "regr_slope" | "regr_intercept" | "regr_r2" | "regr_sxx" | "regr_syy" | "regr_sxy" => {
3241 DataType::Float
3242 }
3243 // v7.32 (round-29) — bitwise aggregates, regr_count, and the
3244 // integer hypothetical-set ranks return an integer.
3245 "bit_and" | "bit_or" | "bit_xor" | "regr_count" | "rank" | "dense_rank" => DataType::BigInt,
3246 // v7.32 (round-29) — hypothetical-set distribution functions.
3247 "percent_rank" | "cume_dist" => DataType::Float,
3248 // v7.32 (round-29) — JSON aggregates return JSON.
3249 "json_agg" | "jsonb_agg" | "json_object_agg" | "jsonb_object_agg" => DataType::Json,
3250 // min/max, percentile_disc, mode, and anything pass-through:
3251 // the argument's shape (for ordered-set aggs `spec.arg` is the
3252 // WITHIN GROUP value expression).
3253 _ => arg_ty.unwrap_or(DataType::Text),
3254 }
3255}
3256
3257fn agg_or_group_type(e: &Expr, synth: &[ColumnSchema]) -> DataType {
3258 if let Expr::Column(c) = e
3259 && let Some(s) = synth.iter().find(|s| s.name == c.name)
3260 {
3261 return s.ty;
3262 }
3263 // v7.26 (round-20 C) — compound expressions over aggregates
3264 // (COALESCE(BOOL_OR(…), false), (array_agg(…))[1], CASE …)
3265 // derive their shape statically against the synth schema; the
3266 // old Text fallback broke sqlx typed decodes of exactly these
3267 // columns.
3268 crate::describe::describe_expr(e, synth)
3269 .map(|shape| shape.ty)
3270 .unwrap_or(DataType::Text)
3271}
3272
3273fn rewrite_expr(e: &Expr, group_exprs: &[Expr], aggs: &[AggSpec]) -> Expr {
3274 // v7.33 (array_agg argmax) — `(array_agg(x ORDER BY y))[1]` rewrites
3275 // to its first_ordered synth column, consuming the subscript. Checked
3276 // before the AggregateOrdered/recursion arms (which would otherwise
3277 // rewrite the inner array_agg and leave the subscript). Same matcher
3278 // as collect_aggregates, so the spec it finds is the one collected.
3279 if let Some((arg, order_by, filter)) = first_ordered_array_agg(e) {
3280 let arg_owned = Some(arg.clone());
3281 let filter_owned = filter.cloned();
3282 for (i, spec) in aggs.iter().enumerate() {
3283 if spec.first_ordered
3284 && spec.name == "array_agg"
3285 && spec.arg == arg_owned
3286 && spec.order_by == *order_by
3287 && spec.filter == filter_owned
3288 {
3289 return Expr::Column(spg_sql::ast::ColumnName {
3290 qualifier: None,
3291 name: format!("__agg_{i}"),
3292 });
3293 }
3294 }
3295 }
3296 // v7.24 (round-16 A) — ordered aggregate: match on the inner
3297 // call PLUS the ordering keys.
3298 if let Expr::AggregateOrdered {
3299 call,
3300 order_by,
3301 distinct,
3302 filter,
3303 } = e
3304 && let Expr::FunctionCall { name, args } = call.as_ref()
3305 {
3306 let lower = name.to_ascii_lowercase();
3307 if is_aggregate_name(&lower) {
3308 let canonical: &str = if lower == "every" { "bool_and" } else { &lower };
3309 // Mirror collect_aggregates: ordered-set aggregates take the
3310 // value from the sort spec and the in-parens arg as direct.
3311 let (arg, direct_arg) = if is_within_group_name(canonical) {
3312 (
3313 order_by.first().map(|o| o.expr.clone()),
3314 args.first().cloned(),
3315 )
3316 } else {
3317 (args.first().cloned(), None)
3318 };
3319 let arg2 = if agg_uses_second_arg(canonical) {
3320 args.get(1).cloned()
3321 } else {
3322 None
3323 };
3324 let filter_owned = filter.as_deref().cloned();
3325 for (i, spec) in aggs.iter().enumerate() {
3326 if spec.name == canonical
3327 && spec.arg == arg
3328 && spec.arg2 == arg2
3329 && spec.distinct == *distinct
3330 && spec.order_by == *order_by
3331 && spec.filter == filter_owned
3332 && spec.direct_arg == direct_arg
3333 {
3334 return Expr::Column(spg_sql::ast::ColumnName {
3335 qualifier: None,
3336 name: format!("__agg_{i}"),
3337 });
3338 }
3339 }
3340 }
3341 }
3342 // Match aggregate FunctionCalls first — they sit outside group_by.
3343 if let Expr::FunctionCall { name, args } = e {
3344 let lower = name.to_ascii_lowercase();
3345 if is_aggregate_name(&lower) {
3346 let arg = if lower == "count_star" {
3347 None
3348 } else {
3349 args.first().cloned()
3350 };
3351 // v7.17.0 — match the spec we registered for
3352 // string_agg(value, separator) on the full pair; v7.32 also
3353 // the regression family and json_object_agg.
3354 let arg2 = if agg_uses_second_arg(&lower) {
3355 args.get(1).cloned()
3356 } else {
3357 None
3358 };
3359 // v7.17.0 — `every` collapses into `bool_and` at
3360 // collection; mirror that here so the rewrite finds
3361 // the matching synth column.
3362 let canonical: &str = if lower == "every" {
3363 "bool_and"
3364 } else {
3365 lower.as_str()
3366 };
3367 for (i, spec) in aggs.iter().enumerate() {
3368 if spec.name == canonical
3369 && spec.arg == arg
3370 && spec.arg2 == arg2
3371 && !spec.distinct
3372 && spec.order_by.is_empty()
3373 {
3374 return Expr::Column(spg_sql::ast::ColumnName {
3375 qualifier: None,
3376 name: format!("__agg_{i}"),
3377 });
3378 }
3379 }
3380 }
3381 }
3382 // Match a group_by expression by AST equality.
3383 for (i, g) in group_exprs.iter().enumerate() {
3384 if g == e {
3385 return Expr::Column(spg_sql::ast::ColumnName {
3386 qualifier: None,
3387 name: format!("__grp_{i}"),
3388 });
3389 }
3390 }
3391 // Recurse into children.
3392 match e {
3393 Expr::AggregateOrdered {
3394 call,
3395 order_by,
3396 distinct,
3397 filter,
3398 } => Expr::AggregateOrdered {
3399 call: Box::new(rewrite_expr(call, group_exprs, aggs)),
3400 distinct: *distinct,
3401 order_by: order_by
3402 .iter()
3403 .map(|o| spg_sql::ast::OrderBy {
3404 expr: rewrite_expr(&o.expr, group_exprs, aggs),
3405 desc: o.desc,
3406 nulls_first: o.nulls_first,
3407 })
3408 .collect(),
3409 // The filter is evaluated against SOURCE rows during
3410 // accumulation, never against synth rows — keep it as-is.
3411 filter: filter.clone(),
3412 },
3413 Expr::Binary { lhs, op, rhs } => Expr::Binary {
3414 lhs: Box::new(rewrite_expr(lhs, group_exprs, aggs)),
3415 op: *op,
3416 rhs: Box::new(rewrite_expr(rhs, group_exprs, aggs)),
3417 },
3418 Expr::Unary { op, expr } => Expr::Unary {
3419 op: *op,
3420 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
3421 },
3422 Expr::Cast { expr, target } => Expr::Cast {
3423 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
3424 target: target.clone(),
3425 },
3426 Expr::IsNull { expr, negated } => Expr::IsNull {
3427 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
3428 negated: *negated,
3429 },
3430 Expr::FunctionCall { name, args } => Expr::FunctionCall {
3431 name: name.clone(),
3432 args: args
3433 .iter()
3434 .map(|a| rewrite_expr(a, group_exprs, aggs))
3435 .collect(),
3436 },
3437 Expr::Like {
3438 expr,
3439 pattern,
3440 negated,
3441 case_insensitive,
3442 } => Expr::Like {
3443 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
3444 pattern: Box::new(rewrite_expr(pattern, group_exprs, aggs)),
3445 negated: *negated,
3446 case_insensitive: *case_insensitive,
3447 },
3448 Expr::Extract { field, source } => Expr::Extract {
3449 field: *field,
3450 source: Box::new(rewrite_expr(source, group_exprs, aggs)),
3451 },
3452 // v7.25.2 (round-19 A) — subquery nodes: rewrite group-key
3453 // references INSIDE the body to `__grp_N` so the correlated
3454 // resolver can substitute them against the synthesised group
3455 // row (aggs are NOT matched inside the body — a COUNT in the
3456 // subquery is the subquery's own aggregate).
3457 Expr::ScalarSubquery(s) => {
3458 Expr::ScalarSubquery(Box::new(rewrite_group_keys_in_select(s, group_exprs)))
3459 }
3460 Expr::Exists { subquery, negated } => Expr::Exists {
3461 subquery: Box::new(rewrite_group_keys_in_select(subquery, group_exprs)),
3462 negated: *negated,
3463 },
3464 Expr::InSubquery {
3465 expr,
3466 subquery,
3467 negated,
3468 } => Expr::InSubquery {
3469 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
3470 subquery: Box::new(rewrite_group_keys_in_select(subquery, group_exprs)),
3471 negated: *negated,
3472 },
3473 // v4.12 window / Literal / Column — clone-pass (these don't
3474 // participate in aggregate rewrite).
3475 Expr::WindowFunction { .. } | Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => {
3476 e.clone()
3477 }
3478 // v7.10.10 — recurse children for array nodes.
3479 Expr::Array(items) => Expr::Array(
3480 items
3481 .iter()
3482 .map(|elem| rewrite_expr(elem, group_exprs, aggs))
3483 .collect(),
3484 ),
3485 Expr::ArraySubscript { target, index } => Expr::ArraySubscript {
3486 target: Box::new(rewrite_expr(target, group_exprs, aggs)),
3487 index: Box::new(rewrite_expr(index, group_exprs, aggs)),
3488 },
3489 Expr::AnyAll {
3490 expr,
3491 op,
3492 array,
3493 is_any,
3494 } => Expr::AnyAll {
3495 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
3496 op: *op,
3497 array: Box::new(rewrite_expr(array, group_exprs, aggs)),
3498 is_any: *is_any,
3499 },
3500 Expr::InList {
3501 expr,
3502 list,
3503 negated,
3504 } => Expr::InList {
3505 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
3506 list: list
3507 .iter()
3508 .map(|item| rewrite_expr(item, group_exprs, aggs))
3509 .collect(),
3510 negated: *negated,
3511 },
3512 Expr::Case {
3513 operand,
3514 branches,
3515 else_branch,
3516 } => Expr::Case {
3517 operand: operand
3518 .as_deref()
3519 .map(|o| Box::new(rewrite_expr(o, group_exprs, aggs))),
3520 branches: branches
3521 .iter()
3522 .map(|(w, t)| {
3523 (
3524 rewrite_expr(w, group_exprs, aggs),
3525 rewrite_expr(t, group_exprs, aggs),
3526 )
3527 })
3528 .collect(),
3529 else_branch: else_branch
3530 .as_deref()
3531 .map(|e| Box::new(rewrite_expr(e, group_exprs, aggs))),
3532 },
3533 }
3534}
3535
3536/// v7.25.2 (round-19 A) — rewrite group-key references inside a
3537/// subquery body to `__grp_N` synthetic columns (aggregates are
3538/// not touched: empty spec list). Runs through the canonical
3539/// Select walker so every expression slot is covered.
3540fn rewrite_group_keys_in_select(
3541 s: &spg_sql::ast::SelectStatement,
3542 group_exprs: &[Expr],
3543) -> spg_sql::ast::SelectStatement {
3544 let mut out = s.clone();
3545 let _ = crate::walk_select_exprs_mut(&mut out, &mut |e| {
3546 *e = rewrite_expr(e, group_exprs, &[]);
3547 Ok(())
3548 });
3549 out
3550}
3551
3552/// Canonical string key for a tuple of group values. Used as map key.
3553/// Per-value group-key encoding (shared by owned and borrowed paths).
3554fn encode_one(out: &mut String, v: &Value) {
3555 use core::fmt::Write;
3556 match v {
3557 Value::Null => out.push_str("N|"),
3558 // v7.36 (perf — mailrs Phase 1) — switch the integer / float
3559 // encoders to `write!`. `n.to_string()` allocates a fresh
3560 // `String` per cell just to push its bytes into the
3561 // (already-cleared) reuse buffer — for the 25 k-row JOIN
3562 // probe in `count_messages` that's 25 k heap allocs per
3563 // query. `write!(&mut String, ...)` formats straight into
3564 // the buffer; no intermediate alloc.
3565 Value::SmallInt(n) => {
3566 let _ = write!(out, "s{n}|");
3567 }
3568 Value::Int(n) => {
3569 let _ = write!(out, "I{n}|");
3570 }
3571 Value::BigInt(n) => {
3572 let _ = write!(out, "B{n}|");
3573 }
3574 Value::Float(x) => {
3575 let _ = write!(out, "F{x}|");
3576 }
3577 Value::Bool(b) => {
3578 out.push(if *b { 'T' } else { 'f' });
3579 out.push('|');
3580 }
3581 Value::Text(s) => {
3582 out.push('S');
3583 out.push_str(s);
3584 out.push('|');
3585 }
3586 Value::Vector(v) => {
3587 out.push('V');
3588 for x in v.iter() {
3589 out.push_str(&x.to_string());
3590 out.push(',');
3591 }
3592 out.push('|');
3593 }
3594 // v6.0.1: GROUP BY on a `VECTOR(N) USING SQ8` column.
3595 // Two cells with byte-identical `(min, max, bytes)`
3596 // share the same group; equivalence is byte-equality
3597 // (same as f32 grouping today — neither path tries to
3598 // normalise nan/-0).
3599 Value::Sq8Vector(q) => {
3600 out.push('Q');
3601 out.push_str(&q.min.to_string());
3602 out.push('@');
3603 out.push_str(&q.max.to_string());
3604 out.push(':');
3605 for b in &q.bytes {
3606 out.push_str(&b.to_string());
3607 out.push(',');
3608 }
3609 out.push('|');
3610 }
3611 // v6.0.3: GROUP BY on a `VECTOR(N) USING HALF` column.
3612 // Byte-equality over the raw u16 bits; matches the SQ8
3613 // path's byte-key model.
3614 Value::HalfVector(h) => {
3615 out.push('H');
3616 for b in &h.bytes {
3617 out.push_str(&b.to_string());
3618 out.push(',');
3619 }
3620 out.push('|');
3621 }
3622 Value::Numeric { scaled, scale } => {
3623 out.push('D');
3624 out.push_str(&scaled.to_string());
3625 out.push('@');
3626 out.push_str(&scale.to_string());
3627 out.push('|');
3628 }
3629 Value::Date(d) => {
3630 out.push('d');
3631 out.push_str(&d.to_string());
3632 out.push('|');
3633 }
3634 Value::Timestamp(t) => {
3635 out.push('t');
3636 out.push_str(&t.to_string());
3637 out.push('|');
3638 }
3639 Value::Interval {
3640 months,
3641 days,
3642 micros,
3643 } => {
3644 out.push('i');
3645 out.push_str(&months.to_string());
3646 out.push('m');
3647 out.push_str(&days.to_string());
3648 out.push('d');
3649 out.push_str(µs.to_string());
3650 out.push('|');
3651 }
3652 Value::Json(s) => {
3653 out.push('j');
3654 out.push_str(s);
3655 out.push('|');
3656 }
3657 // v7.5.0 — Value is #[non_exhaustive] for downstream
3658 // forward-compat. Any future variant lacking explicit
3659 // handling here will share a debug-derived group key,
3660 // which is observably wrong but won't crash.
3661 _ => {
3662 out.push('?');
3663 out.push_str(&format!("{v:?}"));
3664 out.push('|');
3665 }
3666 }
3667}
3668
3669/// v7.30 (perf campaign) - encode from borrowed cells without
3670/// materialising an owned Vec<Value<'static>> first.
3671pub(crate) fn encode_key_refs(vals: &[&Value]) -> String {
3672 let mut out = String::new();
3673 for v in vals {
3674 encode_one(&mut out, v);
3675 }
3676 out
3677}
3678
3679/// v7.31 (perf 3e) — encode into a caller-owned scratch buffer.
3680/// The per-row key paths (group hash, DISTINCT set, join build/
3681/// probe) ran 24k+ String allocations per query through the
3682/// allocator just to LOOK UP a map; the scratch form allocates
3683/// only when a map actually has to take ownership (vacant insert).
3684pub(crate) fn encode_key_refs_into(vals: &[&Value], out: &mut String) {
3685 out.clear();
3686 for v in vals {
3687 encode_one(out, v);
3688 }
3689}
3690
3691pub(crate) fn encode_key(vals: &[Value<'static>]) -> String {
3692 let mut out = String::new();
3693 for v in vals {
3694 encode_one(&mut out, v);
3695 }
3696 out
3697}
3698
3699#[allow(clippy::cast_precision_loss)]
3700fn value_cmp(a: &Value, b: &Value) -> core::cmp::Ordering {
3701 use core::cmp::Ordering::Equal;
3702 match (a, b) {
3703 (Value::Null, Value::Null) => Equal,
3704 (Value::Null, _) => core::cmp::Ordering::Greater, // NULLs last
3705 (_, Value::Null) => core::cmp::Ordering::Less,
3706 (Value::Int(x), Value::Int(y)) => x.cmp(y),
3707 (Value::BigInt(x), Value::BigInt(y)) => x.cmp(y),
3708 (Value::Int(x), Value::BigInt(y)) => i64::from(*x).cmp(y),
3709 (Value::BigInt(x), Value::Int(y)) => x.cmp(&i64::from(*y)),
3710 (Value::Float(x), Value::Float(y)) => x.partial_cmp(y).unwrap_or(Equal),
3711 (Value::Int(x), Value::Float(y)) => f64::from(*x).partial_cmp(y).unwrap_or(Equal),
3712 (Value::Float(x), Value::Int(y)) => x.partial_cmp(&f64::from(*y)).unwrap_or(Equal),
3713 (Value::BigInt(x), Value::Float(y)) => (*x as f64).partial_cmp(y).unwrap_or(Equal),
3714 (Value::Float(x), Value::BigInt(y)) => x.partial_cmp(&(*y as f64)).unwrap_or(Equal),
3715 (Value::Text(x), Value::Text(y)) => x.cmp(y),
3716 (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
3717 _ => Equal,
3718 }
3719}