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/// Per-aggregate running state.
212#[derive(Debug, Default, Clone)]
213struct AggState {
214 count: i64,
215 sum_int: i64,
216 sum_float: f64,
217 extreme: Option<Value>,
218 use_float: bool,
219 /// v7.17.0 — running collection for string_agg / array_agg.
220 /// Each entry is one row's contribution (NULL preserved as
221 /// `Value::Null`; string_agg's finalize step drops them, but
222 /// array_agg keeps them). Pushing in insertion order matches
223 /// PG behaviour when no `ORDER BY` is given inside the
224 /// aggregate call.
225 items: Vec<Value>,
226 /// v7.25 (round-17) — per-group dedupe set for DISTINCT
227 /// aggregates (encoded values; NULLs never reach it because
228 /// the caller's skip runs after the per-aggregate NULL rules).
229 seen: BTreeSet<String>,
230 /// v7.24 (round-16 A) — per-item ORDER BY key tuples, parallel
231 /// to `items` (pushed under the same skip/keep conditions).
232 /// Empty when the aggregate carries no internal ordering.
233 item_keys: Vec<Vec<Value>>,
234 /// v7.17.0 — captured separator for string_agg. PG accepts a
235 /// non-constant separator expression but in practice every
236 /// caller passes a literal; the engine snapshots the last
237 /// non-NULL text it sees, which matches PG's "use the latest
238 /// row's value" behaviour.
239 separator: Option<String>,
240 /// v7.17.0 — running boolean accumulator for bool_and /
241 /// bool_or / every. `None` until the first non-NULL input;
242 /// at finalize None → SQL NULL.
243 bool_acc: Option<bool>,
244 /// v7.32 (round-29) — sum of squares for the variance / stddev
245 /// family (`sum_float` carries the running sum; `count` the n).
246 sum_sq: f64,
247 /// v7.32 (round-29) — running accumulator for bit_and / bit_or /
248 /// bit_xor. `None` until the first non-NULL input → SQL NULL.
249 bit_acc: Option<i64>,
250 /// v7.32 (round-29) — two-argument regression family
251 /// (`covar_*` / `corr` / `regr_*`), PG arg order `f(Y, X)`. Only
252 /// rows where BOTH inputs are non-NULL contribute (`count` is the
253 /// paired n, independent of the single-arg `sum_*`).
254 reg_n: i64,
255 reg_sx: f64,
256 reg_sy: f64,
257 reg_sxx: f64,
258 reg_syy: f64,
259 reg_sxy: f64,
260 /// v7.32 (round-29) — second value stream for `json_object_agg`
261 /// (`items` holds the keys, `aux_items` the values).
262 aux_items: Vec<Value>,
263 /// v7.33 (array_agg argmax) — for a `first_ordered` spec
264 /// (`(array_agg(x ORDER BY y))[1]`), the running first-by-order
265 /// (sort-key tuple, value). Replaced only when a new row's key sorts
266 /// strictly before the current best (ties keep the earliest row, =
267 /// the stable-sort `[1]`). No items/item_keys array is built.
268 first_best: Option<(Vec<Value>, Value)>,
269}
270
271#[derive(Debug, Clone)]
272struct AggSpec {
273 name: String, // lowercased
274 /// First argument (value expression) for every aggregate
275 /// except `count(*)`. `None` for `count_star`.
276 arg: Option<Expr>,
277 /// v7.17.0 — second argument. Only `string_agg(value, sep)`
278 /// uses it today. `None` for every other aggregate (or for
279 /// `array_agg`, which is single-arg). Carried in the spec so
280 /// per-row evaluation can re-use the same separator
281 /// expression across calls.
282 arg2: Option<Expr>,
283 /// v7.25 (round-17) — `COUNT(DISTINCT x)` & friends: dedupe
284 /// the input stream per group before accumulation.
285 distinct: bool,
286 /// v7.24 (round-16 A) — aggregate-internal ORDER BY keys
287 /// (`array_agg(x ORDER BY y DESC NULLS LAST)`). Empty for the
288 /// plain form. Only the collection aggregates honour it;
289 /// other aggregates are order-insensitive and ignore it (PG
290 /// accepts the syntax everywhere too).
291 order_by: Vec<spg_sql::ast::OrderBy>,
292 /// v7.32 (round-29) — `FILTER (WHERE cond)`: a per-row predicate
293 /// evaluated against the source row before accumulation. A row
294 /// whose `cond` is not TRUE (false or NULL) is excluded from this
295 /// aggregate only. `None` for the unfiltered form.
296 filter: Option<Expr>,
297 /// v7.32 (round-29) — ordered-set aggregates only: the *direct*
298 /// argument (the percentile fraction for `percentile_cont/disc`).
299 /// PG requires it constant, so it is evaluated once. `None` for
300 /// `mode()` and for every non-ordered-set aggregate.
301 direct_arg: Option<Expr>,
302 /// v7.33 (array_agg argmax) — set when this spec came from
303 /// `(array_agg(x ORDER BY y))[1]`: accumulate only the first-by-order
304 /// element (a running argmax/argmin) and finalise to that scalar
305 /// value, instead of collecting + sorting + materialising the whole
306 /// per-group array just to take element 1. Returns the element type,
307 /// not the array type.
308 first_ordered: bool,
309}
310
311/// Output of running the aggregate path. Schema describes one row per
312/// group; rows are not yet ORDER BY-sorted (caller does it).
313#[derive(Debug)]
314pub struct AggResult {
315 pub columns: Vec<ColumnSchema>,
316 pub rows: Vec<Row>,
317 /// v7.31 (perf — PG lesson #1, post-LIMIT subquery projection):
318 /// select-list items whose rewritten expr carries a subquery and
319 /// is referenced by neither ORDER BY nor HAVING. Their output
320 /// cells hold NULL placeholders; the caller truncates to
321 /// LIMIT+OFFSET first and only then evaluates these for the
322 /// surviving rows (PG runs the same shape with SubPlan loops=50
323 /// instead of loops=24000). `(output_col, rewritten_expr)`.
324 pub deferred: Vec<(usize, Expr)>,
325 /// Synthetic group rows aligned 1:1 with `rows`; populated only
326 /// when `deferred` is non-empty.
327 pub synth_rows: Vec<Row>,
328 /// Schema the deferred exprs evaluate against.
329 pub synth_schema: Vec<ColumnSchema>,
330}
331
332/// Execute aggregate logic against an already-WHERE-filtered iterator of
333/// rows. `table_alias` is the alias accepted by column resolution.
334#[allow(clippy::too_many_lines)]
335/// v7.25.2 (round-19 A) — caller-injected evaluator for synth-row
336/// expressions that still carry subquery nodes after the rewrite
337/// (correlated subqueries in the select list / HAVING / aggregate
338/// ORDER BY of a GROUP BY query). The engine passes its
339/// correlated-aware evaluator; pure-library callers pass None and
340/// surviving subqueries keep erroring loudly.
341pub type CorrelatedEval<'a> = &'a dyn Fn(&Expr, &Row, &EvalContext<'_>) -> Result<Value, EvalError>;
342
343/// Output of the per-group projection stage (`project_groups`): the
344/// output schema, the projected rows, the synth rows kept alongside
345/// them for post-LIMIT deferred evaluation, the deferred subquery
346/// items, and the rewritten ORDER BY exprs (shared with the sort).
347struct Projection {
348 columns: Vec<ColumnSchema>,
349 out_rows: Vec<Row>,
350 kept_synth: Vec<Row>,
351 deferred: Vec<(usize, Expr)>,
352 order_rewritten: Vec<Expr>,
353}
354
355/// v7.35.0 — detect the `SELECT COUNT(*) FROM … [WHERE …]` shape
356/// (single item, no GROUP BY / HAVING / ORDER BY / DISTINCT /
357/// LIMIT WITH TIES / FILTER / window). For this shape the answer
358/// is exactly `rows.len()` as `BigInt`, no group state needed.
359/// Returns `None` for any deviation so the caller's full pipeline
360/// runs verbatim.
361///
362/// v7.35.2 — also short-circuit `COUNT(<literal>)` (e.g.
363/// `COUNT(1)`) and `COUNT(<column>)` when the column is declared
364/// NOT NULL on the input schema. PG handles both cases as
365/// `COUNT(*)` (the non-null filter is a no-op), so doing the same
366/// here keeps every `count this thing` shape on the same fast path
367/// instead of routing the literal / non-null-col variants through
368/// the four-stage aggregate pipeline.
369fn try_pure_count_star_short_circuit(
370 stmt: &SelectStatement,
371 rows: &[RowRef<'_>],
372 schema_cols: &[ColumnSchema],
373 table_alias: Option<&str>,
374) -> Option<AggResult> {
375 if stmt.distinct
376 || stmt.limit_with_ties
377 || stmt.group_by.is_some()
378 || stmt.having.is_some()
379 || !stmt.order_by.is_empty()
380 {
381 return None;
382 }
383 if stmt.items.len() != 1 {
384 return None;
385 }
386 let SelectItem::Expr { expr, alias } = &stmt.items[0] else {
387 return None;
388 };
389 let Expr::FunctionCall { name, args } = expr else {
390 return None;
391 };
392 if !name.eq_ignore_ascii_case("count") && !name.eq_ignore_ascii_case("count_star") {
393 return None;
394 }
395 let count_star_shape = match args.as_slice() {
396 // `COUNT(*)` parses to `count_star` with no args.
397 [] if name.eq_ignore_ascii_case("count_star") => true,
398 // `COUNT(<literal>)` — the per-row test is "is this literal
399 // non-null?" which is constant, so it's COUNT(*) when the
400 // literal is non-null.
401 [Expr::Literal(lit)] => !matches!(lit, spg_sql::ast::Literal::Null),
402 // `COUNT(<column>)` — same answer as COUNT(*) when the
403 // column is statically declared NOT NULL on the input
404 // schema. Resolve through the alias if one is set.
405 [Expr::Column(c)] => {
406 if let Some(q) = c.qualifier.as_deref()
407 && let Some(alias) = table_alias
408 && !q.eq_ignore_ascii_case(alias)
409 {
410 return None;
411 }
412 schema_cols
413 .iter()
414 .find(|s| s.name.eq_ignore_ascii_case(&c.name))
415 .is_some_and(|s| !s.nullable)
416 }
417 _ => return None,
418 };
419 if !count_star_shape {
420 return None;
421 }
422 let col_name = alias.clone().unwrap_or_else(|| "count".to_string());
423 let count = i64::try_from(rows.len()).unwrap_or(i64::MAX);
424 Some(AggResult {
425 columns: alloc::vec![ColumnSchema::new(col_name, DataType::BigInt, false)],
426 rows: alloc::vec![Row::new(alloc::vec![Value::BigInt(count)])],
427 deferred: Vec::new(),
428 synth_rows: Vec::new(),
429 synth_schema: Vec::new(),
430 })
431}
432
433pub(crate) fn run(
434 stmt: &SelectStatement,
435 rows: &[RowRef<'_>],
436 schema_cols: &[ColumnSchema],
437 table_alias: Option<&str>,
438 correlated_eval: Option<CorrelatedEval<'_>>,
439) -> Result<AggResult, EvalError> {
440 // v7.35.0 — pure `SELECT COUNT(*) FROM … WHERE …` short-circuit.
441 // The caller already filtered rows by WHERE (we run on the
442 // post-WHERE survivor set), so for the canonical pure-COUNT(*)
443 // shape (no GROUP BY / HAVING / ORDER BY / DISTINCT / FILTER /
444 // window) the answer is simply `rows.len()`. The four-stage
445 // aggregate pipeline below (accumulate_groups → build_synth_schema
446 // → finalize_synth_rows → project_groups) collapses to a single
447 // BigInt cell when there's a single group, but each stage still
448 // pays its own allocation tax — group state map, synth schema
449 // vec, finalize loop. `exists_in_60` (mailrs prod #4 baseline)
450 // is exactly this shape on a 25 k-row JOIN.
451 if let Some(short) = try_pure_count_star_short_circuit(stmt, rows, schema_cols, table_alias) {
452 return Ok(short);
453 }
454 let group_exprs: Vec<Expr> = stmt.group_by.clone().unwrap_or_default();
455
456 // Collect aggregate sub-expressions across items + order_by.
457 let mut agg_specs: Vec<AggSpec> = Vec::new();
458 for item in &stmt.items {
459 if let SelectItem::Expr { expr, .. } = item {
460 collect_aggregates(expr, &mut agg_specs);
461 }
462 }
463 for o in &stmt.order_by {
464 collect_aggregates(&o.expr, &mut agg_specs);
465 }
466 if let Some(h) = &stmt.having {
467 collect_aggregates(h, &mut agg_specs);
468 }
469 // v7.17.0 — arity validation. The collector tolerates an
470 // arbitrary positional-arg count; here we enforce the
471 // per-aggregate contract so a malformed call (e.g.
472 // `array_agg()` or `string_agg(x)`) surfaces as a SQL error
473 // rather than silently coercing to a degenerate aggregate.
474 validate_agg_arities(stmt, &agg_specs)?;
475 validate_within_group(&agg_specs)?;
476
477 // (1) Stream the WHERE-filtered rows into insertion-ordered group state.
478 let order = accumulate_groups(
479 rows,
480 &group_exprs,
481 &agg_specs,
482 schema_cols,
483 table_alias,
484 correlated_eval,
485 )?;
486
487 // (2) Build the synthetic per-group schema and finalise each group's row.
488 let synth_schema =
489 build_synth_schema(rows, &group_exprs, &agg_specs, schema_cols, table_alias)?;
490 let synth_rows = finalize_synth_rows(
491 &order,
492 &agg_specs,
493 &synth_schema,
494 rows,
495 schema_cols,
496 table_alias,
497 )?;
498
499 // (3) Rewrite the user's expressions, filter groups by HAVING and project.
500 let Projection {
501 columns,
502 mut out_rows,
503 mut kept_synth,
504 deferred,
505 order_rewritten,
506 } = project_groups(
507 synth_rows,
508 stmt,
509 &group_exprs,
510 &agg_specs,
511 &synth_schema,
512 correlated_eval,
513 )?;
514
515 // (4) ORDER BY on the aggregated output (the caller applies LIMIT).
516 if !stmt.order_by.is_empty() {
517 let (sorted_synth, sorted_out) = sort_synth_by_order_by(
518 &synth_schema,
519 &stmt.order_by,
520 &order_rewritten,
521 kept_synth,
522 out_rows,
523 correlated_eval,
524 )?;
525 kept_synth = sorted_synth;
526 out_rows = sorted_out;
527 }
528
529 let (synth_rows_out, synth_schema_out) = if deferred.is_empty() {
530 (Vec::new(), Vec::new())
531 } else {
532 (kept_synth, synth_schema.clone())
533 };
534 Ok(AggResult {
535 columns,
536 rows: out_rows,
537 deferred,
538 synth_rows: synth_rows_out,
539 synth_schema: synth_schema_out,
540 })
541}
542
543/// v7.32 (round-29) — validate the structural requirements of WITHIN
544/// GROUP (ordered-set / hypothetical-set) aggregates up front, so a
545/// malformed call surfaces as a SQL error rather than a silently
546/// degenerate aggregate.
547fn validate_within_group(agg_specs: &[AggSpec]) -> Result<(), EvalError> {
548 // v7.32 (round-29) — WITHIN GROUP aggregates require the clause (PG
549 // raises a hard error otherwise rather than silently degrading), and
550 // SPG supports the single-sort-key form only.
551 for spec in agg_specs {
552 if is_within_group_name(&spec.name) {
553 if spec.order_by.is_empty() {
554 return Err(EvalError::TypeMismatch {
555 detail: format!("{}() requires WITHIN GROUP (ORDER BY …)", spec.name),
556 });
557 }
558 // mode() is the only WITHIN GROUP aggregate with no direct
559 // argument; the rest carry one (percentile fraction /
560 // hypothetical value).
561 if spec.name != "mode" && spec.direct_arg.is_none() {
562 return Err(EvalError::TypeMismatch {
563 detail: format!("{}() requires a direct argument", spec.name),
564 });
565 }
566 // Multi-key WITHIN GROUP (multiple sort keys / hypothetical
567 // args) is not supported yet — error loudly instead of
568 // silently using only the first key.
569 if spec.order_by.len() > 1 {
570 return Err(EvalError::TypeMismatch {
571 detail: format!(
572 "{}() with multiple WITHIN GROUP sort keys is not supported yet",
573 spec.name
574 ),
575 });
576 }
577 }
578 }
579 Ok(())
580}
581
582/// (1) Stream the WHERE-filtered rows, group by the GROUP BY value
583/// tuple, and update per-group aggregate state. Returns the groups in
584/// insertion order. See `run` for the bind-once fast path rationale.
585#[allow(clippy::too_many_lines, clippy::type_complexity)]
586fn accumulate_groups(
587 rows: &[RowRef<'_>],
588 group_exprs: &[Expr],
589 agg_specs: &[AggSpec],
590 schema_cols: &[ColumnSchema],
591 table_alias: Option<&str>,
592 correlated_eval: Option<CorrelatedEval<'_>>,
593) -> Result<Vec<(Vec<Value>, Vec<AggState>)>, EvalError> {
594 let ctx = EvalContext::new(schema_cols, table_alias);
595 // Map group key (vec of values, encoded as canonical string) -> group state.
596 // v7.32 (architecture v2, P2b) — insertion-ordered group state in
597 // a Vec; the hash map only maps key → index. Removes the parallel
598 // `key_order: Vec<String>` (a second per-group key clone) and the
599 // per-group re-probe `groups[k]` at finalize (24k hash lookups for
600 // the inbox shape). The map owns its key once on vacant insert.
601 let mut order: Vec<(Vec<Value>, Vec<AggState>)> = Vec::new();
602 let mut groups: hashbrown::HashMap<String, usize> = hashbrown::HashMap::new();
603 // When there are no GROUP BY exprs *and* there is at least one aggregate,
604 // every row collapses into a single anonymous group keyed by "".
605 if rows.is_empty() && group_exprs.is_empty() {
606 // Single empty-aggregate group: count=0, sum=0, max=NULL, etc.
607 // No rows follow, so the map is never probed — seed `order` only.
608 let init: Vec<AggState> = (0..agg_specs.len()).map(|_| AggState::default()).collect();
609 order.push((Vec::new(), init));
610 }
611
612 // v7.30 (perf campaign) - hoist the per-row work that doesn't
613 // depend on the row: which group exprs need collation folding
614 // (none, for most queries - the old code cloned the whole
615 // group_vals vec per row just in case).
616 // v7.30 (perf campaign) - the no-tax row loop. When a group
617 // expr or an aggregate argument is a bare column reference
618 // (the overwhelmingly common shape), bind its position ONCE
619 // and read row cells by offset in the loop - no per-row tree
620 // walk, no owned-Value clone out of resolve_column. Anything
621 // more complex keeps the eval path.
622 let col_pos = |e: &Expr| -> Option<usize> {
623 // Qualified references only: the bare-name resolver carries
624 // alias/ambiguity logic the bind-once path must not fork.
625 if let Expr::Column(c) = e
626 && c.qualifier.is_some()
627 {
628 eval::find_column_pos(c, &ctx)
629 } else {
630 None
631 }
632 };
633 let group_pos: Vec<Option<usize>> = group_exprs.iter().map(col_pos).collect();
634 let all_groups_bound = group_pos.iter().all(Option::is_some);
635 let arg_pos: Vec<Option<usize>> = agg_specs
636 .iter()
637 .map(|spec| spec.arg.as_ref().and_then(|e| col_pos(e)))
638 .collect();
639 // v7.36 (perf — mailrs Ask 1 SUM(LENGTH(text_body)) 18ms → ?) —
640 // pre-compile every aggregate arg that's a `fully_compilable`
641 // PURE expression over bound columns. Without this, `LENGTH(col)`
642 // / `COALESCE(col, '')` / `CAST(col AS BIGINT)` etc. ALL fell
643 // through to the `(None, Some(e)) => eval_arg(e, mat, ...)` slow
644 // path that materialises a Cow<Row> per input row — for a 25k-row
645 // JOIN that's 25k full-row clones for one column read. The Step
646 // VM (`eval_compiled_ref`) reads columns by RowRef::get and runs
647 // the same `apply_function` dispatcher with zero materialisation.
648 let arg_compiled: Vec<Option<eval::CompiledExpr>> = agg_specs
649 .iter()
650 .enumerate()
651 .map(|(i, spec)| match (&arg_pos[i], &spec.arg) {
652 (Some(_), _) => None,
653 (None, Some(e)) if eval::fully_compilable(e) => Some(eval::compile_expr(e, &ctx)),
654 _ => None,
655 })
656 .collect();
657 // v7.33 (array_agg perf) — bound positions for each spec's internal
658 // ORDER BY keys, so an ordered aggregate (`array_agg(x ORDER BY y)`)
659 // reads the sort key by reference (RowRef::get) instead of
660 // materialising the whole combined join row per input row just to
661 // eval one bound column. Mirrors arg_pos. On the inbox shape this
662 // turned 24k full-row (~1 KB each) clones into 24k single-cell reads.
663 let order_pos: Vec<Vec<Option<usize>>> = agg_specs
664 .iter()
665 .map(|spec| spec.order_by.iter().map(|o| col_pos(&o.expr)).collect())
666 .collect();
667 // Does any spec need the fully-materialised row in the bound fast
668 // path — a FILTER, a non-bound value arg, a second arg, or a non-bound
669 // ORDER key? When false (every aggregate arg/key is a bound column —
670 // the inbox shape) the bound fast path never materialises a row.
671 let needs_mat = agg_specs.iter().enumerate().any(|(i, s)| {
672 s.filter.is_some()
673 || (s.arg.is_some() && arg_pos[i].is_none() && arg_compiled[i].is_none())
674 || s.arg2.is_some()
675 || order_pos[i].iter().any(Option::is_none)
676 });
677 let ci_positions: Vec<usize> = group_exprs
678 .iter()
679 .enumerate()
680 .filter(|(_, g)| {
681 matches!(
682 eval::column_collation(g, &ctx),
683 Some(spg_storage::Collation::CaseInsensitive)
684 )
685 })
686 .map(|(i, _)| i)
687 .collect();
688 // v7.31 (perf 3e) — per-row scratch buffers. The fast path used
689 // to allocate a key String (and a refs Vec) for EVERY row just
690 // to probe the group map; hits — the overwhelming case — now
691 // touch the allocator zero times.
692 let mut keybuf_s = String::new();
693 // v7.36 — reused Step VM eval stack for compiled aggregate args.
694 let mut eval_stack: Vec<Value> = Vec::new();
695 let mut dkeybuf = String::new();
696 let mut refs: Vec<&Value> = Vec::with_capacity(group_pos.len());
697 // v7.32 (round-31) — an aggregate's argument / FILTER / second arg /
698 // ORDER key may itself be a *correlated* subquery, e.g.
699 // `MAX((SELECT i.v FROM inner i WHERE i.fk = o.id))`. A non-correlated
700 // subquery is pre-resolved to a literal before this loop, but a
701 // correlated one survives as a subquery node and must be evaluated per
702 // outer row through the correlated evaluator — the same hook the
703 // select-list / HAVING / ORDER finalisers already use below. Plain
704 // `eval_expr` would hit "subquery reached row eval".
705 //
706 // The `any_agg_subquery` gate is computed once here so the common case
707 // (no subquery anywhere in the aggregate args — including every hot
708 // scan/group aggregate) short-circuits before the per-row
709 // `expr_has_subquery` walk: `eval_arg` is then exactly `eval_expr`.
710 let any_agg_subquery = correlated_eval.is_some()
711 && agg_specs.iter().any(|s| {
712 s.filter
713 .as_ref()
714 .is_some_and(|e| crate::expr_has_subquery(e))
715 || s.arg.as_ref().is_some_and(|e| crate::expr_has_subquery(e))
716 || s.arg2.as_ref().is_some_and(|e| crate::expr_has_subquery(e))
717 || s.order_by.iter().any(|o| crate::expr_has_subquery(&o.expr))
718 });
719 let eval_arg = |e: &Expr, r: &Row, c: &EvalContext<'_>| -> Result<Value, EvalError> {
720 match correlated_eval {
721 Some(f) if any_agg_subquery && crate::expr_has_subquery(e) => f(e, r, c),
722 _ => eval::eval_expr(e, r, c),
723 }
724 };
725 // v7.36 (perf — mailrs Phase 1, post u64-hash) — single
726 // anonymous group fast path. When the query has no GROUP BY
727 // (`SELECT SUM(LENGTH(col)) FROM ...`, COUNT, AVG, etc.) the
728 // whole input collapses into one group. The fast path below
729 // still pays one `groups.get("")` hash probe per row plus
730 // `entry = &mut order[0]` reindex even when the empty-key
731 // path encodes nothing — measured ~50 ns/row across 25 k rows
732 // = ~1.25 ms of pure bookkeeping on the user_storage_usage
733 // baseline.
734 //
735 // Bypass: lift `entry` outside the loop and feed every row
736 // straight into it. Same `update_state` machinery, zero
737 // per-row hash work, zero per-row index lookup.
738 let single_anon_group = group_exprs.is_empty() && !rows.is_empty();
739 if single_anon_group {
740 // Seed the single group at idx 0 once.
741 let init: Vec<AggState> = (0..agg_specs.len()).map(|_| AggState::default()).collect();
742 order.clear();
743 order.push((Vec::new(), init));
744 }
745 // v7.36 (perf — mailrs Phase 1, count_messages 2.58 → ?) —
746 // `COUNT(*)` short-circuit. For a single-anon-group `COUNT(*)`
747 // with no FILTER / DISTINCT, every survivor counts once — the
748 // answer IS `rows.len()`. Skips the 25 k iterations of
749 // `update_state("count_star", …)` on the mailrs count_messages
750 // shape; the JOIN already produced exactly the set of rows
751 // that must be counted.
752 if single_anon_group
753 && agg_specs.len() == 1
754 && agg_specs[0].name == "count_star"
755 && agg_specs[0].filter.is_none()
756 && agg_specs[0].arg.is_none()
757 && agg_specs[0].arg2.is_none()
758 && agg_specs[0].order_by.is_empty()
759 && !agg_specs[0].distinct
760 {
761 let state = &mut order[0].1[0];
762 state.count = rows.len() as i64;
763 return Ok(order);
764 }
765 // v7.36 (perf — mailrs Phase 1) — `COUNT(<bound col>)` (non-`*`)
766 // collapses to: read the cell, increment when not NULL. Skips
767 // the per-row spec dispatch + `update_state("count", …)`.
768 if single_anon_group
769 && agg_specs.len() == 1
770 && agg_specs[0].name == "count"
771 && agg_specs[0].filter.is_none()
772 && agg_specs[0].arg2.is_none()
773 && agg_specs[0].order_by.is_empty()
774 && !agg_specs[0].distinct
775 && arg_pos[0].is_some()
776 {
777 let p = arg_pos[0].unwrap();
778 let mut count: i64 = 0;
779 for row in rows {
780 if !matches!(row.get(p), Some(Value::Null) | None) {
781 count += 1;
782 }
783 }
784 let state = &mut order[0].1[0];
785 state.count = count;
786 return Ok(order);
787 }
788 // v7.36 (perf — mailrs Phase 1, user_storage_usage 7.5 → ?) —
789 // single-aggregate streaming accumulator. For
790 // `SUM(<compiled-expr>)` / `SUM(<bound col>)` with no GROUP BY,
791 // no FILTER, no arg2, no ORDER BY, no DISTINCT, the whole
792 // per-row work collapses to: eval the arg, match the Value
793 // variant, accumulate. Skips the spec-dispatch loop +
794 // `update_state` per-row name match. On a 25 k-row JOIN
795 // (user_storage_usage `SUM(LENGTH(text_body))`) that's
796 // ~50-100 ns/row of pure spec-dispatch overhead removed.
797 if single_anon_group
798 && agg_specs.len() == 1
799 && agg_specs[0].filter.is_none()
800 && agg_specs[0].arg2.is_none()
801 && agg_specs[0].order_by.is_empty()
802 && !agg_specs[0].distinct
803 && (agg_specs[0].name == "sum" || agg_specs[0].name == "avg")
804 && (arg_pos[0].is_some() || arg_compiled[0].is_some())
805 {
806 let arg_pos0 = arg_pos[0];
807 let arg_c0 = &arg_compiled[0];
808 let mut sum_int: i64 = 0;
809 let mut sum_float: f64 = 0.0;
810 let mut use_float = false;
811 let mut count: i64 = 0;
812 // Borrow-aware fast inner: avoid the per-row clone when arg
813 // is a bound column position.
814 if let Some(p) = arg_pos0 {
815 for row in rows {
816 let v_ref = row.get(p).unwrap_or(&Value::Null);
817 match v_ref {
818 Value::Null => continue,
819 Value::SmallInt(n) => {
820 sum_int += i64::from(*n);
821 count += 1;
822 }
823 Value::Int(n) => {
824 sum_int += i64::from(*n);
825 count += 1;
826 }
827 Value::BigInt(n) => {
828 sum_int += *n;
829 count += 1;
830 }
831 Value::Float(x) => {
832 sum_float += *x;
833 use_float = true;
834 count += 1;
835 }
836 other => {
837 return Err(EvalError::TypeMismatch {
838 detail: format!("sum/avg need numeric, got {:?}", other.data_type()),
839 });
840 }
841 }
842 }
843 } else if let Some(p) = arg_c0.as_ref().and_then(|c| c.as_single_column_length()) {
844 // v7.36 (perf — mailrs Phase 1, user_storage_usage hot
845 // inner) — `SUM(LENGTH(<text col>))` collapses to a
846 // straight scan: read the cell by ref, branch on the
847 // variant, do an ASCII probe + `len()` (or
848 // `chars().count()` on non-ASCII), accumulate. No Step
849 // VM, no stack push/pop, no `BigInt` boxing on the way
850 // out — pure i64 sum. The original Step VM path keeps
851 // running for everything outside this shape (`SUM(col)`,
852 // `SUM(expr)`, multi-step compiled args).
853 for row in rows {
854 let Some(v_ref) = row.get(p) else {
855 continue;
856 };
857 let n = match v_ref {
858 Value::Null => continue,
859 Value::Text(s) => {
860 if s.is_ascii() {
861 s.len() as i64
862 } else {
863 s.chars().count() as i64
864 }
865 }
866 other => {
867 return Err(EvalError::TypeMismatch {
868 detail: format!("length() needs text, got {:?}", other.data_type()),
869 });
870 }
871 };
872 sum_int += n;
873 count += 1;
874 }
875 } else {
876 let c = arg_c0.as_ref().unwrap();
877 for row in rows {
878 let v = eval::eval_compiled_ref(c, row, &ctx, &mut eval_stack)?;
879 match v {
880 Value::Null => continue,
881 Value::SmallInt(n) => {
882 sum_int += i64::from(n);
883 count += 1;
884 }
885 Value::Int(n) => {
886 sum_int += i64::from(n);
887 count += 1;
888 }
889 Value::BigInt(n) => {
890 sum_int += n;
891 count += 1;
892 }
893 Value::Float(x) => {
894 sum_float += x;
895 use_float = true;
896 count += 1;
897 }
898 other => {
899 return Err(EvalError::TypeMismatch {
900 detail: format!("sum/avg need numeric, got {:?}", other.data_type()),
901 });
902 }
903 }
904 }
905 }
906 let state = &mut order[0].1[0];
907 state.count = count;
908 state.sum_int = sum_int;
909 state.sum_float = sum_float;
910 state.use_float = use_float;
911 return Ok(order);
912 }
913 for row in rows {
914 if single_anon_group {
915 let entry = &mut order[0];
916 let mat: Option<Cow<'_, Row>> = if needs_mat { Some(row.as_row()) } else { None };
917 for (i, spec) in agg_specs.iter().enumerate() {
918 if let Some(f) = &spec.filter
919 && !matches!(
920 eval_arg(f, mat.as_deref().expect("needs_mat for FILTER"), &ctx)?,
921 Value::Bool(true)
922 )
923 {
924 continue;
925 }
926 let arg_owned: Value;
927 let arg_ref: &Value = match (&arg_pos[i], &arg_compiled[i], &spec.arg) {
928 (Some(p), _, _) => row.get(*p).unwrap_or(&Value::Null),
929 (None, _, None) => {
930 arg_owned = Value::Bool(true);
931 &arg_owned
932 }
933 (None, Some(c), _) => {
934 arg_owned = eval::eval_compiled_ref(c, row, &ctx, &mut eval_stack)?;
935 &arg_owned
936 }
937 (None, None, Some(e)) => {
938 arg_owned = eval_arg(
939 e,
940 mat.as_deref().expect("needs_mat for non-bound arg"),
941 &ctx,
942 )?;
943 &arg_owned
944 }
945 };
946 let arg2_val = match &spec.arg2 {
947 None => None,
948 Some(e) => Some(eval_arg(
949 e,
950 mat.as_deref().expect("needs_mat for arg2"),
951 &ctx,
952 )?),
953 };
954 let order_keys = if spec.order_by.is_empty() {
955 None
956 } else {
957 let mut keys = Vec::with_capacity(spec.order_by.len());
958 for (k, o) in spec.order_by.iter().enumerate() {
959 let v = if let Some(p) = order_pos[i][k] {
960 row.get(p).cloned().unwrap_or(Value::Null)
961 } else {
962 eval_arg(
963 &o.expr,
964 mat.as_deref().expect("needs_mat for ORDER key"),
965 &ctx,
966 )?
967 };
968 keys.push(v);
969 }
970 Some(keys)
971 };
972 if spec.distinct {
973 encode_key_refs_into(core::slice::from_ref(&arg_ref), &mut dkeybuf);
974 if entry.1[i].seen.contains(dkeybuf.as_str()) {
975 continue;
976 }
977 entry.1[i].seen.insert(dkeybuf.clone());
978 }
979 update_state(
980 &mut entry.1[i],
981 &spec.name,
982 arg_ref,
983 arg2_val.as_ref(),
984 order_keys,
985 )?;
986 }
987 continue;
988 }
989 // Fast key: bound positions + no ci folding -> encode
990 // straight from borrowed cells; group_vals materialise
991 // only when the group is NEW.
992 if all_groups_bound && ci_positions.is_empty() {
993 refs.clear();
994 refs.extend(
995 group_pos
996 .iter()
997 .map(|p| row.get(p.unwrap()).unwrap_or(&Value::Null)),
998 );
999 encode_key_refs_into(&refs, &mut keybuf_s);
1000 let idx = match groups.get(keybuf_s.as_str()) {
1001 Some(&i) => i,
1002 None => {
1003 let i = order.len();
1004 let init: Vec<AggState> =
1005 (0..agg_specs.len()).map(|_| AggState::default()).collect();
1006 let owned: Vec<Value> = refs.iter().map(|v| (*v).clone()).collect();
1007 order.push((owned, init));
1008 groups.insert(keybuf_s.clone(), i);
1009 i
1010 }
1011 };
1012 let entry = &mut order[idx];
1013 // v7.33 (array_agg perf) — materialise the combined row AT
1014 // MOST once per input row, and only when a spec actually
1015 // needs the eval path (FILTER / non-bound arg / arg2 / non-
1016 // bound ORDER key). Bound args and bound ORDER keys read
1017 // cells by reference below, so the inbox shape (all bound)
1018 // never materialises — killing the per-row ~1 KB clone that
1019 // dominated the ordered-aggregate cost.
1020 let mat: Option<Cow<'_, Row>> = if needs_mat { Some(row.as_row()) } else { None };
1021 for (i, spec) in agg_specs.iter().enumerate() {
1022 // v7.32 (round-29) — FILTER (WHERE cond): exclude rows
1023 // where cond is not TRUE before they reach this
1024 // aggregate's accumulator (and before DISTINCT dedup).
1025 if let Some(f) = &spec.filter
1026 && !matches!(
1027 eval_arg(f, mat.as_deref().expect("needs_mat for FILTER"), &ctx)?,
1028 Value::Bool(true)
1029 )
1030 {
1031 continue;
1032 }
1033 let arg_owned: Value;
1034 let arg_ref: &Value = match (&arg_pos[i], &arg_compiled[i], &spec.arg) {
1035 (Some(p), _, _) => row.get(*p).unwrap_or(&Value::Null),
1036 (None, _, None) => {
1037 arg_owned = Value::Bool(true);
1038 &arg_owned
1039 }
1040 (None, Some(c), _) => {
1041 // v7.36 — compiled-arg fast path. `eval_stack`
1042 // is reused across rows; the Step VM never
1043 // materialises a row for column reads.
1044 arg_owned = eval::eval_compiled_ref(c, row, &ctx, &mut eval_stack)?;
1045 &arg_owned
1046 }
1047 (None, None, Some(e)) => {
1048 arg_owned = eval_arg(
1049 e,
1050 mat.as_deref().expect("needs_mat for non-bound arg"),
1051 &ctx,
1052 )?;
1053 &arg_owned
1054 }
1055 };
1056 let arg2_val = match &spec.arg2 {
1057 None => None,
1058 Some(e) => Some(eval_arg(
1059 e,
1060 mat.as_deref().expect("needs_mat for arg2"),
1061 &ctx,
1062 )?),
1063 };
1064 let order_keys = if spec.order_by.is_empty() {
1065 None
1066 } else {
1067 let mut keys = Vec::with_capacity(spec.order_by.len());
1068 for (k, o) in spec.order_by.iter().enumerate() {
1069 // Bound ORDER key → read the cell by reference; only
1070 // a non-bound key falls to the materialised eval path.
1071 keys.push(match order_pos[i][k] {
1072 Some(p) => row.get(p).cloned().unwrap_or(Value::Null),
1073 None => eval_arg(
1074 &o.expr,
1075 mat.as_deref().expect("needs_mat for non-bound ORDER key"),
1076 &ctx,
1077 )?,
1078 });
1079 }
1080 Some(keys)
1081 };
1082 // v7.33 (array_agg argmax) — first_ordered: keep only the
1083 // running first-by-order element (strict-less replacement
1084 // = ties keep the earliest row, matching the stable-sort
1085 // `[1]`), no array build.
1086 if spec.first_ordered {
1087 if let Some(keys) = order_keys {
1088 let st = &mut entry.1[i];
1089 let better = match &st.first_best {
1090 None => true,
1091 Some((bk, _)) => {
1092 cmp_order_keys(&spec.order_by, &keys, bk)
1093 == core::cmp::Ordering::Less
1094 }
1095 };
1096 if better {
1097 st.first_best = Some((keys, arg_ref.clone()));
1098 }
1099 }
1100 continue;
1101 }
1102 if spec.distinct {
1103 encode_key_refs_into(core::slice::from_ref(&arg_ref), &mut dkeybuf);
1104 if entry.1[i].seen.contains(dkeybuf.as_str()) {
1105 continue;
1106 }
1107 entry.1[i].seen.insert(dkeybuf.clone());
1108 }
1109 update_state(
1110 &mut entry.1[i],
1111 &spec.name,
1112 arg_ref,
1113 arg2_val.as_ref(),
1114 order_keys,
1115 )?;
1116 }
1117 continue;
1118 }
1119 // v7.32 (P4 increment 2) — eval (non-bound) path: present the
1120 // row as a borrowed Row once (Owned → zero-cost borrow; a join
1121 // tuple materialises here exactly once, never on the bound fast
1122 // path above), then the original eval loop runs unchanged.
1123 let row_materialised = row.as_row();
1124 let row: &Row = &row_materialised;
1125 let group_vals: Vec<Value> = group_exprs
1126 .iter()
1127 .map(|g| eval::eval_expr(g, row, &ctx))
1128 .collect::<Result<_, _>>()?;
1129 // v7.17.0 Phase 2.5b — case-insensitive group keying: fold
1130 // only the ci columns, and only when any exist. Display
1131 // value (`group_vals`) stays original — only the key folds.
1132 let key = if ci_positions.is_empty() {
1133 encode_key(&group_vals)
1134 } else {
1135 let mut key_vals = group_vals.clone();
1136 for &i in &ci_positions {
1137 if let Value::Text(s) = &key_vals[i] {
1138 key_vals[i] = Value::Text(s.to_ascii_lowercase());
1139 }
1140 }
1141 encode_key(&key_vals)
1142 };
1143 // Probe by index; the map owns the key once on vacant insert.
1144 let idx = match groups.get(key.as_str()) {
1145 Some(&i) => i,
1146 None => {
1147 let i = order.len();
1148 let init: Vec<AggState> =
1149 (0..agg_specs.len()).map(|_| AggState::default()).collect();
1150 order.push((group_vals.clone(), init));
1151 groups.insert(key, i);
1152 i
1153 }
1154 };
1155 let entry = &mut order[idx];
1156 for (i, spec) in agg_specs.iter().enumerate() {
1157 // v7.32 (round-29) — FILTER (WHERE cond): exclude rows where
1158 // cond is not TRUE before accumulation (and before DISTINCT).
1159 if let Some(f) = &spec.filter
1160 && !matches!(eval_arg(f, row, &ctx)?, Value::Bool(true))
1161 {
1162 continue;
1163 }
1164 let arg_val = match &spec.arg {
1165 None => Value::Bool(true), // count_star: sentinel non-null
1166 Some(e) => eval_arg(e, row, &ctx)?,
1167 };
1168 // v7.17.0 — `string_agg(value, separator)` evaluates the
1169 // separator per row but PG treats it as constant; we
1170 // pass the per-row value into update_state so a future
1171 // varying-separator caller still sees correct output,
1172 // even though SPG (like PG) only uses the most recent.
1173 let arg2_val = match &spec.arg2 {
1174 None => None,
1175 Some(e) => Some(eval_arg(e, row, &ctx)?),
1176 };
1177 // v7.24 (round-16 A) — aggregate-internal ORDER BY:
1178 // evaluate the key tuple against the source row.
1179 let order_keys = if spec.order_by.is_empty() {
1180 None
1181 } else {
1182 let mut keys = Vec::with_capacity(spec.order_by.len());
1183 for o in &spec.order_by {
1184 keys.push(eval_arg(&o.expr, row, &ctx)?);
1185 }
1186 Some(keys)
1187 };
1188 // v7.33 (array_agg argmax) — first_ordered: keep the running
1189 // first-by-order element only (mirrors the bound fast path).
1190 if spec.first_ordered {
1191 if let Some(keys) = order_keys {
1192 let st = &mut entry.1[i];
1193 let better = match &st.first_best {
1194 None => true,
1195 Some((bk, _)) => {
1196 cmp_order_keys(&spec.order_by, &keys, bk) == core::cmp::Ordering::Less
1197 }
1198 };
1199 if better {
1200 st.first_best = Some((keys, arg_val.clone()));
1201 }
1202 }
1203 continue;
1204 }
1205 // v7.25 (round-17) — DISTINCT: drop repeated inputs
1206 // before they reach the accumulator. NULLs flow through
1207 // (each aggregate's own NULL rule applies; PG also
1208 // treats NULL as a single distinct value for array_agg).
1209 if spec.distinct {
1210 let key = encode_key(core::slice::from_ref(&arg_val));
1211 if !entry.1[i].seen.insert(key) {
1212 continue;
1213 }
1214 }
1215 update_state(
1216 &mut entry.1[i],
1217 &spec.name,
1218 &arg_val,
1219 arg2_val.as_ref(),
1220 order_keys,
1221 )?;
1222 }
1223 }
1224 Ok(order)
1225}
1226
1227/// (2a) Build the synthetic per-group schema: `__grp_0..K` then
1228/// `__agg_0..N`. Group types are probed from the first row; aggregate
1229/// types from each spec.
1230fn build_synth_schema(
1231 rows: &[RowRef<'_>],
1232 group_exprs: &[Expr],
1233 agg_specs: &[AggSpec],
1234 schema_cols: &[ColumnSchema],
1235 table_alias: Option<&str>,
1236) -> Result<Vec<ColumnSchema>, EvalError> {
1237 let ctx = EvalContext::new(schema_cols, table_alias);
1238 // Build synthetic schema: __grp_0..K then __agg_0..N.
1239 let group_types: Vec<DataType> = if rows.is_empty() {
1240 // Use Text as a safe stand-in — empty result means schema isn't
1241 // observable. Avoids needing to evaluate group exprs on no row.
1242 group_exprs.iter().map(|_| DataType::Text).collect()
1243 } else {
1244 let probe_row = rows[0].as_row();
1245 let probe: &Row = &probe_row;
1246 group_exprs
1247 .iter()
1248 .map(|g| {
1249 eval::eval_expr(g, probe, &ctx).map(|v| v.data_type().unwrap_or(DataType::Text))
1250 })
1251 .collect::<Result<_, _>>()?
1252 };
1253 let agg_types: Vec<DataType> = agg_specs
1254 .iter()
1255 .map(|spec| infer_agg_type(spec, schema_cols))
1256 .collect();
1257 let mut synth_schema: Vec<ColumnSchema> = Vec::new();
1258 for (i, ty) in group_types.iter().enumerate() {
1259 synth_schema.push(ColumnSchema::new(format!("__grp_{i}"), *ty, true));
1260 }
1261 for (i, ty) in agg_types.iter().enumerate() {
1262 synth_schema.push(ColumnSchema::new(format!("__agg_{i}"), *ty, true));
1263 }
1264 Ok(synth_schema)
1265}
1266
1267/// (2b) Materialise one synthetic row per group (insertion order):
1268/// apply each aggregate's internal ORDER BY, then finalise the running
1269/// state into the group + aggregate cells.
1270/// v7.33 — compare two aggregate-internal ORDER BY key tuples under the
1271/// per-key DESC / NULLS directives. This is the exact comparator the
1272/// finalize sort uses, factored out so the `first_ordered` argmax
1273/// accumulator's "keep first" decision is provably identical to taking
1274/// element `[1]` of the fully-sorted array.
1275fn cmp_order_keys(
1276 order_by: &[spg_sql::ast::OrderBy],
1277 a: &[Value],
1278 b: &[Value],
1279) -> core::cmp::Ordering {
1280 for (k, o) in order_by.iter().enumerate() {
1281 let cmp = crate::order_by_value_cmp(o.desc, o.nulls_first, &a[k], &b[k]);
1282 if cmp != core::cmp::Ordering::Equal {
1283 return cmp;
1284 }
1285 }
1286 core::cmp::Ordering::Equal
1287}
1288
1289fn finalize_synth_rows(
1290 order: &[(Vec<Value>, Vec<AggState>)],
1291 agg_specs: &[AggSpec],
1292 synth_schema: &[ColumnSchema],
1293 rows: &[RowRef<'_>],
1294 schema_cols: &[ColumnSchema],
1295 table_alias: Option<&str>,
1296) -> Result<Vec<Row>, EvalError> {
1297 let ctx = EvalContext::new(schema_cols, table_alias);
1298 // v7.32 (round-29) — ordered-set direct arguments (the percentile
1299 // fraction) are constant per PG, so evaluate each once up front.
1300 let direct_arg_vals: Vec<Option<Value>> = agg_specs
1301 .iter()
1302 .map(|spec| match (&spec.direct_arg, rows.first()) {
1303 (Some(e), Some(r)) => eval::eval_expr(e, &r.as_row(), &ctx).map(Some),
1304 _ => Ok(None),
1305 })
1306 .collect::<Result<_, _>>()?;
1307
1308 // Materialise synthetic rows (insertion order = `order`).
1309 let mut synth_rows: Vec<Row> = Vec::new();
1310 for (gvals, states) in order {
1311 let mut values: Vec<Value> = Vec::with_capacity(synth_schema.len());
1312 values.extend(gvals.iter().cloned());
1313 for (i, st) in states.iter().enumerate() {
1314 // v7.33 (array_agg argmax) — first_ordered: the running
1315 // first-by-order value IS the result; no array build/sort.
1316 if agg_specs[i].first_ordered {
1317 values.push(
1318 st.first_best
1319 .as_ref()
1320 .map_or(Value::Null, |(_, v)| v.clone()),
1321 );
1322 continue;
1323 }
1324 // v7.24 (round-16 A) — order the collected items per the
1325 // aggregate-internal ORDER BY before finalize consumes
1326 // them.
1327 let st_sorted;
1328 let st_final: &AggState =
1329 if !agg_specs[i].order_by.is_empty() && st.item_keys.len() == st.items.len() {
1330 let mut idx: Vec<usize> = (0..st.items.len()).collect();
1331 let ob = &agg_specs[i].order_by;
1332 idx.sort_by(|&x, &y| cmp_order_keys(ob, &st.item_keys[x], &st.item_keys[y]));
1333 let mut sorted = st.clone();
1334 sorted.items = idx.iter().map(|&j| st.items[j].clone()).collect();
1335 st_sorted = sorted;
1336 &st_sorted
1337 } else {
1338 st
1339 };
1340 // Ordered-set aggregates compute from the sorted items + the
1341 // direct fraction; everything else uses the running state.
1342 let v = if is_within_group_name(&agg_specs[i].name) {
1343 finalize_ordered_set(
1344 &agg_specs[i].name,
1345 st_final,
1346 direct_arg_vals[i].as_ref(),
1347 agg_specs[i].order_by.first(),
1348 )
1349 } else {
1350 finalize(&agg_specs[i].name, st_final)
1351 };
1352 values.push(v);
1353 }
1354 synth_rows.push(Row::new(values));
1355 }
1356 Ok(synth_rows)
1357}
1358
1359/// (3) Rewrite the user's SELECT items + HAVING to reference the
1360/// synthetic columns, filter groups by HAVING, and project each
1361/// surviving group into an output row. The synth rows ride alongside
1362/// (`kept_synth`) so post-LIMIT deferred subqueries can evaluate later.
1363#[allow(clippy::too_many_lines)]
1364fn project_groups(
1365 synth_rows: Vec<Row>,
1366 stmt: &SelectStatement,
1367 group_exprs: &[Expr],
1368 agg_specs: &[AggSpec],
1369 synth_schema: &[ColumnSchema],
1370 correlated_eval: Option<CorrelatedEval<'_>>,
1371) -> Result<Projection, EvalError> {
1372 // Rewrite the user's SELECT items + ORDER BY to reference synthetic
1373 // columns. After rewriting, every remaining `Expr::Column` must
1374 // resolve against the synthetic schema (i.e. must have been a GROUP
1375 // BY expression).
1376 let columns: Vec<ColumnSchema> = stmt
1377 .items
1378 .iter()
1379 .map(|item| match item {
1380 SelectItem::Wildcard => Err(EvalError::TypeMismatch {
1381 detail: "SELECT * with aggregates is not supported".into(),
1382 }),
1383 SelectItem::Expr { expr, alias } => {
1384 let rewritten = rewrite_expr(expr, group_exprs, agg_specs);
1385 let name = alias.clone().unwrap_or_else(|| expr.to_string());
1386 Ok(ColumnSchema::new(
1387 name,
1388 agg_or_group_type(&rewritten, synth_schema),
1389 true,
1390 ))
1391 }
1392 })
1393 .collect::<Result<_, _>>()?;
1394
1395 // Project per synthetic row. HAVING filters out groups *before*
1396 // we keep the projected row — same semantics as PG: HAVING runs
1397 // against the aggregated row (so `HAVING count(*) > 1` works) and
1398 // sees only group-by'd columns plus aggregate values.
1399 let synth_ctx = EvalContext::new(synth_schema, None);
1400 let having_rewritten = stmt
1401 .having
1402 .as_ref()
1403 .map(|h| rewrite_expr(h, group_exprs, agg_specs));
1404 // v7.30 (phase 3e-1) - rewrite SELECT items ONCE. This ran per
1405 // GROUP (23.5k x 9 items of AST cloning = ~48% of the inbox
1406 // query in sampled stacks); the rewrite is group-independent.
1407 // Stable addresses also let the per-expression subquery plans
1408 // (v7.29 3c) hit across groups instead of rebuilding.
1409 let items_rewritten: alloc::vec::Vec<Option<Expr>> = stmt
1410 .items
1411 .iter()
1412 .map(|item| match item {
1413 SelectItem::Expr { expr, .. } => Some(rewrite_expr(expr, group_exprs, agg_specs)),
1414 SelectItem::Wildcard => None,
1415 })
1416 .collect();
1417 // v7.31 (perf — PG lesson #1): subquery-bearing select items
1418 // deferred to post-LIMIT, when no sort/filter key can observe
1419 // them. ORDER BY rewrites are hoisted here so the safety check
1420 // and the sort below share one rewrite pass.
1421 let order_rewritten: Vec<Expr> = stmt
1422 .order_by
1423 .iter()
1424 .map(|o| rewrite_expr(&o.expr, group_exprs, agg_specs))
1425 .collect();
1426 let defer_enabled = correlated_eval.is_some()
1427 && !stmt.distinct
1428 && !having_rewritten
1429 .as_ref()
1430 .is_some_and(crate::expr_has_subquery)
1431 && !order_rewritten.iter().any(crate::expr_has_subquery);
1432 let deferred: Vec<(usize, Expr)> = if defer_enabled {
1433 items_rewritten
1434 .iter()
1435 .enumerate()
1436 .filter_map(|(i, r)| {
1437 r.as_ref()
1438 .filter(|e| crate::expr_has_subquery(e))
1439 .map(|e| (i, e.clone()))
1440 })
1441 .collect()
1442 } else {
1443 Vec::new()
1444 };
1445 // v7.32 (architecture v2, P2) — compile the per-group synth-row
1446 // expressions ONCE. The projection / HAVING here run per GROUP
1447 // (24k for the inbox shape) × per item; the rewritten exprs are
1448 // mostly `Column(__agg_N)` / `Column(__grp_K)` against the synth
1449 // schema — flat step programs, no tree walk per group.
1450 let having_compiled = having_rewritten
1451 .as_ref()
1452 .filter(|h| eval::fully_compilable(h))
1453 .map(|h| eval::compile_expr(h, &synth_ctx));
1454 let items_compiled: Vec<Option<eval::CompiledExpr>> = items_rewritten
1455 .iter()
1456 .enumerate()
1457 .map(|(i, r)| {
1458 r.as_ref()
1459 .filter(|e| !deferred.iter().any(|(c, _)| *c == i) && eval::fully_compilable(e))
1460 .map(|e| eval::compile_expr(e, &synth_ctx))
1461 })
1462 .collect();
1463 let mut kept_synth: Vec<Row> = Vec::new();
1464 let mut out_rows: Vec<Row> = Vec::new();
1465 let mut stack: Vec<Value> = Vec::new();
1466 for srow in synth_rows {
1467 if let Some(hc) = &having_compiled {
1468 let cond = eval::eval_compiled(hc, &srow, &synth_ctx, &mut stack)?;
1469 if !matches!(cond, Value::Bool(true)) {
1470 continue;
1471 }
1472 } else if let Some(h) = &having_rewritten {
1473 let cond = match correlated_eval {
1474 Some(f) if crate::expr_has_subquery(h) => f(h, &srow, &synth_ctx)?,
1475 _ => eval::eval_expr(h, &srow, &synth_ctx)?,
1476 };
1477 if !matches!(cond, Value::Bool(true)) {
1478 continue;
1479 }
1480 }
1481 let mut values: Vec<Value> = Vec::with_capacity(columns.len());
1482 for (i, rewritten) in items_rewritten.iter().enumerate() {
1483 let Some(rewritten) = rewritten else { continue };
1484 if deferred.iter().any(|(c, _)| *c == i) {
1485 values.push(Value::Null);
1486 continue;
1487 }
1488 values.push(if let Some(cc) = &items_compiled[i] {
1489 eval::eval_compiled(cc, &srow, &synth_ctx, &mut stack)?
1490 } else {
1491 match correlated_eval {
1492 Some(f) if crate::expr_has_subquery(rewritten) => {
1493 f(rewritten, &srow, &synth_ctx)?
1494 }
1495 _ => eval::eval_expr(rewritten, &srow, &synth_ctx)?,
1496 }
1497 });
1498 }
1499 kept_synth.push(srow);
1500 out_rows.push(Row::new(values));
1501 }
1502 Ok(Projection {
1503 columns,
1504 out_rows,
1505 kept_synth,
1506 deferred,
1507 order_rewritten,
1508 })
1509}
1510
1511/// (4) Sort the projected output by the rewritten ORDER BY keys. The
1512/// synth rows ride through the sort so deferred subqueries evaluate
1513/// against the surviving groups after the caller's LIMIT truncation.
1514fn sort_synth_by_order_by(
1515 synth_schema: &[ColumnSchema],
1516 order_by: &[spg_sql::ast::OrderBy],
1517 order_rewritten: &[Expr],
1518 mut kept_synth: Vec<Row>,
1519 mut out_rows: Vec<Row>,
1520 correlated_eval: Option<CorrelatedEval<'_>>,
1521) -> Result<(Vec<Row>, Vec<Row>), EvalError> {
1522 let synth_ctx = EvalContext::new(synth_schema, None);
1523 // v6.4.0 — multi-key ORDER BY on aggregate output. Each key
1524 // gets its own rewrite + per-key DESC flag. (Rewrites hoisted
1525 // above as `order_rewritten` — shared with the deferral
1526 // safety check.)
1527 let keys_meta: Vec<(bool, Option<bool>)> =
1528 order_by.iter().map(|o| (o.desc, o.nulls_first)).collect();
1529 // P2: compile order-by keys once (per-group sort keys are
1530 // the same `__agg_N` / `__grp_K` shape as the projection).
1531 let order_compiled: Vec<Option<eval::CompiledExpr>> = order_rewritten
1532 .iter()
1533 .map(|e| {
1534 Some(e)
1535 .filter(|e| eval::fully_compilable(e))
1536 .map(|e| eval::compile_expr(e, &synth_ctx))
1537 })
1538 .collect();
1539 // The synth row rides through the sort so deferred exprs can
1540 // evaluate against the surviving groups after the caller's
1541 // LIMIT truncation.
1542 let mut keystack: Vec<Value> = Vec::new();
1543 let mut tagged: Vec<(Vec<Value>, Row, Row)> = Vec::with_capacity(kept_synth.len());
1544 for (s, o) in kept_synth.into_iter().zip(out_rows) {
1545 let mut keys = Vec::with_capacity(order_rewritten.len());
1546 for (e, oc) in order_rewritten.iter().zip(&order_compiled) {
1547 keys.push(if let Some(oc) = oc {
1548 eval::eval_compiled(oc, &s, &synth_ctx, &mut keystack)?
1549 } else {
1550 match correlated_eval {
1551 Some(f) if crate::expr_has_subquery(e) => f(e, &s, &synth_ctx)?,
1552 _ => eval::eval_expr(e, &s, &synth_ctx)?,
1553 }
1554 });
1555 }
1556 tagged.push((keys, s, o));
1557 }
1558 tagged.sort_by(|a, b| {
1559 use core::cmp::Ordering;
1560 for (i, (ka, kb)) in a.0.iter().zip(b.0.iter()).enumerate() {
1561 let (desc, nf) = keys_meta[i];
1562 let cmp = crate::order_by_value_cmp(desc, nf, ka, kb);
1563 if cmp != Ordering::Equal {
1564 return cmp;
1565 }
1566 }
1567 Ordering::Equal
1568 });
1569 kept_synth = Vec::with_capacity(tagged.len());
1570 out_rows = Vec::with_capacity(tagged.len());
1571 for (_, s, o) in tagged {
1572 kept_synth.push(s);
1573 out_rows.push(o);
1574 }
1575 Ok((kept_synth, out_rows))
1576}
1577
1578/// v7.17.0 — walk the statement again to validate the positional
1579/// arity of every aggregate call site. Done after AST collection
1580/// rather than inside `collect_aggregates` so the collector stays
1581/// infallible; callers in `run()` can do a single early-error
1582/// exit before any per-row work.
1583fn validate_agg_arities(stmt: &SelectStatement, _specs: &[AggSpec]) -> Result<(), EvalError> {
1584 fn walk(e: &Expr) -> Result<(), EvalError> {
1585 if let Expr::FunctionCall { name, args } = e {
1586 let lower = name.to_ascii_lowercase();
1587 let expected: Option<usize> = match lower.as_str() {
1588 "count_star" => Some(0),
1589 "count" | "sum" | "avg" | "min" | "max" | "array_agg"
1590 // v7.17.0 — boolean aggregates also take exactly
1591 // one arg. `every` is an alias normalised inside
1592 // collect_aggregates / rewrite_expr.
1593 | "bool_and" | "bool_or" | "every"
1594 // v7.32 (round-29) — statistical + bitwise aggregates
1595 // + single-arg JSON aggregate.
1596 | "stddev" | "stddev_samp" | "stddev_pop"
1597 | "variance" | "var_samp" | "var_pop"
1598 | "bit_and" | "bit_or" | "bit_xor"
1599 | "json_agg" | "jsonb_agg" => Some(1),
1600 // v7.32 (round-29) — two-argument aggregates: string_agg,
1601 // the regression family f(Y, X), and json_object_agg.
1602 "string_agg"
1603 | "covar_pop" | "covar_samp" | "corr"
1604 | "regr_count" | "regr_avgx" | "regr_avgy" | "regr_slope"
1605 | "regr_intercept" | "regr_r2" | "regr_sxx" | "regr_syy" | "regr_sxy"
1606 | "json_object_agg" | "jsonb_object_agg" => Some(2),
1607 _ => None,
1608 };
1609 if let Some(want) = expected
1610 && args.len() != want
1611 {
1612 return Err(EvalError::TypeMismatch {
1613 detail: alloc::format!("{lower}() takes {want} arg(s), got {}", args.len()),
1614 });
1615 }
1616 for a in args {
1617 walk(a)?;
1618 }
1619 } else if let Expr::Binary { lhs, rhs, .. } = e {
1620 walk(lhs)?;
1621 walk(rhs)?;
1622 } else if let Expr::Unary { expr, .. }
1623 | Expr::Cast { expr, .. }
1624 | Expr::IsNull { expr, .. } = e
1625 {
1626 walk(expr)?;
1627 }
1628 Ok(())
1629 }
1630 for item in &stmt.items {
1631 if let SelectItem::Expr { expr, .. } = item {
1632 walk(expr)?;
1633 }
1634 }
1635 for o in &stmt.order_by {
1636 walk(&o.expr)?;
1637 }
1638 if let Some(h) = &stmt.having {
1639 walk(h)?;
1640 }
1641 Ok(())
1642}
1643
1644/// v7.33 (array_agg argmax) — recognise `(array_agg(x ORDER BY y))[1]`,
1645/// the argmax/argmin idiom: a non-DISTINCT ordered `array_agg`
1646/// subscripted by the constant 1. Returns `(value_arg, order_by,
1647/// filter)` on a match. When matched, the whole per-group array build +
1648/// sort + materialise is replaced by a running first-by-order scalar
1649/// accumulator and the subscript node is consumed (replaced by the
1650/// synthetic column). collect_aggregates and rewrite_expr share this one
1651/// matcher so their `__agg_<i>` assignment stays in lockstep.
1652fn first_ordered_array_agg(e: &Expr) -> Option<(&Expr, &[spg_sql::ast::OrderBy], Option<&Expr>)> {
1653 let Expr::ArraySubscript { target, index } = e else {
1654 return None;
1655 };
1656 if !matches!(
1657 index.as_ref(),
1658 Expr::Literal(spg_sql::ast::Literal::Integer(1))
1659 ) {
1660 return None;
1661 }
1662 let Expr::AggregateOrdered {
1663 call,
1664 order_by,
1665 distinct,
1666 filter,
1667 } = target.as_ref()
1668 else {
1669 return None;
1670 };
1671 if *distinct || order_by.is_empty() {
1672 return None;
1673 }
1674 let Expr::FunctionCall { name, args } = call.as_ref() else {
1675 return None;
1676 };
1677 if !name.eq_ignore_ascii_case("array_agg") || args.len() != 1 {
1678 return None;
1679 }
1680 Some((&args[0], order_by, filter.as_deref()))
1681}
1682
1683fn collect_aggregates(e: &Expr, out: &mut Vec<AggSpec>) {
1684 match e {
1685 // v7.24 (round-16 A) — ordered aggregate: register the inner
1686 // call's spec with the ordering attached.
1687 Expr::AggregateOrdered {
1688 call,
1689 order_by,
1690 distinct,
1691 filter,
1692 } => {
1693 if let Expr::FunctionCall { name, args } = call.as_ref() {
1694 let lower = name.to_ascii_lowercase();
1695 if is_aggregate_name(&lower) {
1696 let canonical = if lower == "every" {
1697 "bool_and".to_string()
1698 } else {
1699 lower
1700 };
1701 // Ordered-set aggregates (`percentile_cont(f)
1702 // WITHIN GROUP (ORDER BY x)`) take the value to
1703 // aggregate from the sort spec and the in-parens
1704 // arg as the direct (fraction) argument.
1705 let ordered_set = is_within_group_name(&canonical);
1706 let (arg, direct_arg) = if ordered_set {
1707 (
1708 order_by.first().map(|o| o.expr.clone()),
1709 args.first().cloned(),
1710 )
1711 } else {
1712 (args.first().cloned(), None)
1713 };
1714 let spec = AggSpec {
1715 name: canonical.clone(),
1716 arg,
1717 arg2: if agg_uses_second_arg(&canonical) {
1718 args.get(1).cloned()
1719 } else {
1720 None
1721 },
1722 distinct: *distinct,
1723 order_by: order_by.clone(),
1724 filter: filter.as_deref().cloned(),
1725 direct_arg,
1726 first_ordered: false,
1727 };
1728 if !out.iter().any(|s| {
1729 s.name == spec.name
1730 && s.arg == spec.arg
1731 && s.arg2 == spec.arg2
1732 && s.distinct == spec.distinct
1733 && s.order_by == spec.order_by
1734 && s.filter == spec.filter
1735 && s.direct_arg == spec.direct_arg
1736 && s.first_ordered == spec.first_ordered
1737 }) {
1738 out.push(spec);
1739 }
1740 return;
1741 }
1742 }
1743 collect_aggregates(call, out);
1744 for o in order_by {
1745 collect_aggregates(&o.expr, out);
1746 }
1747 }
1748 Expr::FunctionCall { name, args } => {
1749 let lower = name.to_ascii_lowercase();
1750 if is_aggregate_name(&lower) {
1751 let arg = if lower == "count_star" {
1752 None
1753 } else {
1754 args.first().cloned()
1755 };
1756 // v7.17.0 — second positional arg for
1757 // `string_agg(value, separator)`; v7.32 — also the
1758 // regression family `f(Y, X)` and `json_object_agg`.
1759 let arg2 = if agg_uses_second_arg(&lower) {
1760 args.get(1).cloned()
1761 } else {
1762 None
1763 };
1764 // v7.17.0 — `every` is the SQL-standard alias for
1765 // `bool_and`; collapse at collection time so
1766 // update_state / finalize need only one arm.
1767 let canonical = if lower == "every" {
1768 "bool_and".to_string()
1769 } else {
1770 lower
1771 };
1772 let spec = AggSpec {
1773 name: canonical,
1774 arg: arg.clone(),
1775 arg2: arg2.clone(),
1776 distinct: false,
1777 order_by: Vec::new(),
1778 filter: None,
1779 direct_arg: None,
1780 first_ordered: false,
1781 };
1782 if !out.iter().any(|s| {
1783 s.name == spec.name
1784 && s.arg == spec.arg
1785 && s.arg2 == spec.arg2
1786 && !s.distinct
1787 && s.order_by == spec.order_by
1788 && s.filter.is_none()
1789 && !s.first_ordered
1790 }) {
1791 out.push(spec);
1792 }
1793 // Don't recurse into the arg — nested aggregates are
1794 // illegal in standard SQL.
1795 } else {
1796 for a in args {
1797 collect_aggregates(a, out);
1798 }
1799 }
1800 }
1801 Expr::Binary { lhs, rhs, .. } => {
1802 collect_aggregates(lhs, out);
1803 collect_aggregates(rhs, out);
1804 }
1805 Expr::Unary { expr, .. } | Expr::Cast { expr, .. } | Expr::IsNull { expr, .. } => {
1806 collect_aggregates(expr, out);
1807 }
1808 Expr::Like { expr, pattern, .. } => {
1809 collect_aggregates(expr, out);
1810 collect_aggregates(pattern, out);
1811 }
1812 Expr::InList { expr, list, .. } => {
1813 collect_aggregates(expr, out);
1814 for item in list {
1815 collect_aggregates(item, out);
1816 }
1817 }
1818 Expr::Extract { source, .. } => collect_aggregates(source, out),
1819 // v4.10 subquery + v4.12 window / Literal / Column —
1820 // non-recursing leaves for the aggregate collector.
1821 Expr::ScalarSubquery(_)
1822 | Expr::Exists { .. }
1823 | Expr::InSubquery { .. }
1824 | Expr::WindowFunction { .. }
1825 | Expr::Literal(_)
1826 | Expr::Placeholder(_)
1827 | Expr::Column(_) => {}
1828 // v7.10.10 — recurse into array constructor children +
1829 // subscript / ANY/ALL operands.
1830 Expr::Array(items) => {
1831 for elem in items {
1832 collect_aggregates(elem, out);
1833 }
1834 }
1835 Expr::ArraySubscript { target, index } => {
1836 // v7.33 (array_agg argmax) — `(array_agg(x ORDER BY y))[1]`
1837 // collects as a first_ordered spec; the subscript is consumed
1838 // here (do NOT recurse into the array_agg, or it would also
1839 // register a plain full-array spec).
1840 if let Some((arg, order_by, filter)) = first_ordered_array_agg(e) {
1841 let spec = AggSpec {
1842 name: "array_agg".to_string(),
1843 arg: Some(arg.clone()),
1844 arg2: None,
1845 distinct: false,
1846 order_by: order_by.to_vec(),
1847 filter: filter.cloned(),
1848 direct_arg: None,
1849 first_ordered: true,
1850 };
1851 if !out.iter().any(|s| {
1852 s.name == spec.name
1853 && s.arg == spec.arg
1854 && s.order_by == spec.order_by
1855 && s.filter == spec.filter
1856 && s.first_ordered
1857 }) {
1858 out.push(spec);
1859 }
1860 return;
1861 }
1862 collect_aggregates(target, out);
1863 collect_aggregates(index, out);
1864 }
1865 Expr::AnyAll { expr, array, .. } => {
1866 collect_aggregates(expr, out);
1867 collect_aggregates(array, out);
1868 }
1869 Expr::Case {
1870 operand,
1871 branches,
1872 else_branch,
1873 } => {
1874 if let Some(o) = operand {
1875 collect_aggregates(o, out);
1876 }
1877 for (w, t) in branches {
1878 collect_aggregates(w, out);
1879 collect_aggregates(t, out);
1880 }
1881 if let Some(e) = else_branch {
1882 collect_aggregates(e, out);
1883 }
1884 }
1885 }
1886}
1887
1888fn update_state(
1889 st: &mut AggState,
1890 name: &str,
1891 v: &Value,
1892 arg2: Option<&Value>,
1893 order_keys: Option<Vec<Value>>,
1894) -> Result<(), EvalError> {
1895 let is_null = matches!(v, Value::Null);
1896 match name {
1897 "count_star" => st.count += 1,
1898 "count" => {
1899 if !is_null {
1900 st.count += 1;
1901 }
1902 }
1903 "sum" | "avg" => {
1904 if is_null {
1905 return Ok(());
1906 }
1907 st.count += 1;
1908 match v {
1909 Value::Int(n) => st.sum_int += i64::from(*n),
1910 Value::BigInt(n) => st.sum_int += *n,
1911 Value::Float(x) => {
1912 st.use_float = true;
1913 st.sum_float += *x;
1914 }
1915 other => {
1916 return Err(EvalError::TypeMismatch {
1917 detail: format!("sum/avg need numeric, got {:?}", other.data_type()),
1918 });
1919 }
1920 }
1921 }
1922 "min" => {
1923 if is_null {
1924 return Ok(());
1925 }
1926 match &st.extreme {
1927 None => st.extreme = Some(v.clone()),
1928 Some(cur) => {
1929 if value_cmp(v, cur) == core::cmp::Ordering::Less {
1930 st.extreme = Some(v.clone());
1931 }
1932 }
1933 }
1934 }
1935 "max" => {
1936 if is_null {
1937 return Ok(());
1938 }
1939 match &st.extreme {
1940 None => st.extreme = Some(v.clone()),
1941 Some(cur) => {
1942 if value_cmp(v, cur) == core::cmp::Ordering::Greater {
1943 st.extreme = Some(v.clone());
1944 }
1945 }
1946 }
1947 }
1948 // v7.17.0 — string_agg(value, separator). NULL value is
1949 // skipped (PG aggregate-skip-null). Separator captured
1950 // from the latest row that flows through; matches PG's
1951 // semantics of evaluating the separator per row but using
1952 // the last value at finalize time (in practice it's
1953 // constant). count is bumped so we can distinguish "empty
1954 // group → NULL" from "all-NULL group → NULL".
1955 "string_agg" => {
1956 if let Some(sep) = arg2
1957 && let Value::Text(s) = sep
1958 {
1959 st.separator = Some(s.clone());
1960 }
1961 if is_null {
1962 return Ok(());
1963 }
1964 if let Value::Text(s) = v {
1965 st.items.push(Value::Text(s.clone()));
1966 if let Some(k) = order_keys {
1967 st.item_keys.push(k);
1968 }
1969 st.count += 1;
1970 } else {
1971 return Err(EvalError::TypeMismatch {
1972 detail: format!("string_agg requires text value, got {:?}", v.data_type()),
1973 });
1974 }
1975 }
1976 // v7.17.0 — array_agg(value). Unlike string_agg, NULL
1977 // elements are KEPT in the array (PG behaviour); the
1978 // result is NULL only when ZERO rows fed in. Element type
1979 // is locked from the first row's value type; subsequent
1980 // rows must match (PG also rejects mixed-type array_agg).
1981 "array_agg" => {
1982 st.items.push(v.clone());
1983 if let Some(k) = order_keys {
1984 st.item_keys.push(k);
1985 }
1986 st.count += 1;
1987 }
1988 // v7.17.0 — bool_and(p): TRUE iff every non-NULL input is
1989 // TRUE. NULL skipped; running accumulator stays at TRUE
1990 // until the first non-NULL FALSE.
1991 "bool_and" => {
1992 if is_null {
1993 return Ok(());
1994 }
1995 let b = match v {
1996 Value::Bool(b) => *b,
1997 other => {
1998 return Err(EvalError::TypeMismatch {
1999 detail: format!("bool_and requires bool, got {:?}", other.data_type()),
2000 });
2001 }
2002 };
2003 st.bool_acc = Some(st.bool_acc.map_or(b, |acc| acc && b));
2004 }
2005 // v7.17.0 — bool_or(p): TRUE iff any non-NULL input is
2006 // TRUE. NULL skipped.
2007 "bool_or" => {
2008 if is_null {
2009 return Ok(());
2010 }
2011 let b = match v {
2012 Value::Bool(b) => *b,
2013 other => {
2014 return Err(EvalError::TypeMismatch {
2015 detail: format!("bool_or requires bool, got {:?}", other.data_type()),
2016 });
2017 }
2018 };
2019 st.bool_acc = Some(st.bool_acc.map_or(b, |acc| acc || b));
2020 }
2021 // v7.32 (round-29) — variance / stddev family. Accumulate the
2022 // running sum (sum_float) and sum of squares (sum_sq) over the
2023 // non-NULL numeric inputs; finalize divides by n or n-1.
2024 "stddev" | "stddev_samp" | "stddev_pop" | "variance" | "var_samp" | "var_pop" => {
2025 if is_null {
2026 return Ok(());
2027 }
2028 let x = match v {
2029 Value::Int(n) => f64::from(*n),
2030 Value::SmallInt(n) => f64::from(*n),
2031 Value::BigInt(n) => *n as f64,
2032 Value::Float(x) => *x,
2033 other => {
2034 return Err(EvalError::TypeMismatch {
2035 detail: format!("{name} needs numeric, got {:?}", other.data_type()),
2036 });
2037 }
2038 };
2039 st.count += 1;
2040 st.sum_float += x;
2041 st.sum_sq += x * x;
2042 }
2043 // v7.32 (round-29) — bitwise aggregates over integer inputs.
2044 "bit_and" | "bit_or" | "bit_xor" => {
2045 if is_null {
2046 return Ok(());
2047 }
2048 let n = match v {
2049 Value::Int(n) => i64::from(*n),
2050 Value::SmallInt(n) => i64::from(*n),
2051 Value::BigInt(n) => *n,
2052 other => {
2053 return Err(EvalError::TypeMismatch {
2054 detail: format!("{name} needs integer, got {:?}", other.data_type()),
2055 });
2056 }
2057 };
2058 st.bit_acc = Some(match (st.bit_acc, name) {
2059 (None, _) => n,
2060 (Some(acc), "bit_and") => acc & n,
2061 (Some(acc), "bit_or") => acc | n,
2062 (Some(acc), _) => acc ^ n, // bit_xor
2063 });
2064 }
2065 // v7.32 (round-29) — WITHIN GROUP aggregates (ordered-set +
2066 // hypothetical-set) collect the sort value (NULLs ignored, per
2067 // PG) into `items`, sorted at finalize by the parallel
2068 // `item_keys`.
2069 n if is_within_group_name(n) => {
2070 if is_null {
2071 return Ok(());
2072 }
2073 st.items.push(v.clone());
2074 if let Some(k) = order_keys {
2075 st.item_keys.push(k);
2076 }
2077 st.count += 1;
2078 }
2079 // v7.32 (round-29) — regression family f(Y, X). Only rows with
2080 // BOTH inputs non-NULL contribute (PG semantics). `v` is Y,
2081 // `arg2` is X.
2082 n if is_regression_name(n) => {
2083 let (Some(y), Some(x)) = (agg_value_to_f64(v), arg2.and_then(agg_value_to_f64)) else {
2084 return Ok(()); // NULL (or non-numeric) in either input
2085 };
2086 st.reg_n += 1;
2087 st.reg_sx += x;
2088 st.reg_sy += y;
2089 st.reg_sxx += x * x;
2090 st.reg_syy += y * y;
2091 st.reg_sxy += x * y;
2092 }
2093 // v7.32 (round-29) — json_agg / jsonb_agg collect every input
2094 // (NULL becomes JSON null, per PG) in row order.
2095 "json_agg" | "jsonb_agg" => {
2096 st.items.push(v.clone());
2097 st.count += 1;
2098 }
2099 // v7.32 (round-29) — json_object_agg(key, value): keys in
2100 // `items`, values in `aux_items`. A NULL key is skipped (PG
2101 // raises; we drop it rather than abort the whole query).
2102 "json_object_agg" | "jsonb_object_agg" => {
2103 if is_null {
2104 return Ok(());
2105 }
2106 st.items.push(v.clone());
2107 st.aux_items.push(arg2.cloned().unwrap_or(Value::Null));
2108 st.count += 1;
2109 }
2110 _ => unreachable!("non-aggregate {name} in update_state"),
2111 }
2112 Ok(())
2113}
2114
2115#[allow(clippy::cast_precision_loss)]
2116fn finalize(name: &str, st: &AggState) -> Value {
2117 match name {
2118 "count" | "count_star" => Value::BigInt(st.count),
2119 "sum" => {
2120 if st.count == 0 {
2121 Value::Null
2122 } else if st.use_float {
2123 Value::Float(st.sum_float + (st.sum_int as f64))
2124 } else {
2125 Value::BigInt(st.sum_int)
2126 }
2127 }
2128 "avg" => {
2129 if st.count == 0 {
2130 Value::Null
2131 } else {
2132 let total = if st.use_float {
2133 st.sum_float + (st.sum_int as f64)
2134 } else {
2135 st.sum_int as f64
2136 };
2137 Value::Float(total / (st.count as f64))
2138 }
2139 }
2140 "min" | "max" => st.extreme.clone().unwrap_or(Value::Null),
2141 // v7.17.0 — string_agg: join all collected text items with
2142 // the captured separator. Empty / all-NULL group → NULL
2143 // (PG semantics).
2144 "string_agg" => {
2145 if st.items.is_empty() {
2146 return Value::Null;
2147 }
2148 let sep = st.separator.clone().unwrap_or_default();
2149 let mut out = String::new();
2150 for (i, item) in st.items.iter().enumerate() {
2151 if i > 0 {
2152 out.push_str(&sep);
2153 }
2154 if let Value::Text(s) = item {
2155 out.push_str(s);
2156 }
2157 }
2158 Value::Text(out)
2159 }
2160 // v7.17.0 — array_agg: collect into a typed array. NULL
2161 // elements are preserved per PG. Result type is decided
2162 // by the first non-NULL element seen (or Text fallback
2163 // when the whole group is NULL — PG would surface the
2164 // declared input type, but SPG hasn't yet wired the
2165 // aggregate's static input-type from `describe`).
2166 "array_agg" => {
2167 if st.items.is_empty() {
2168 return Value::Null;
2169 }
2170 let probe = st.items.iter().find(|v| !v.is_null());
2171 match probe.and_then(spg_storage::Value::data_type) {
2172 Some(DataType::Int) | Some(DataType::SmallInt) => {
2173 let items: Vec<Option<i32>> = st
2174 .items
2175 .iter()
2176 .map(|v| match v {
2177 Value::Int(n) => Some(*n),
2178 Value::SmallInt(n) => Some(i32::from(*n)),
2179 _ => None,
2180 })
2181 .collect();
2182 Value::IntArray(items)
2183 }
2184 Some(DataType::BigInt) => {
2185 let items: Vec<Option<i64>> = st
2186 .items
2187 .iter()
2188 .map(|v| match v {
2189 Value::BigInt(n) => Some(*n),
2190 _ => None,
2191 })
2192 .collect();
2193 Value::BigIntArray(items)
2194 }
2195 _ => {
2196 let items: Vec<Option<String>> = st
2197 .items
2198 .iter()
2199 .map(|v| match v {
2200 Value::Text(s) => Some(s.clone()),
2201 Value::Null => None,
2202 other => Some(format!("{other:?}")),
2203 })
2204 .collect();
2205 Value::TextArray(items)
2206 }
2207 }
2208 }
2209 // v7.17.0 — bool_and / bool_or finalize: lazy-init pattern
2210 // means `None` is exactly "empty group or all-NULL", which
2211 // PG surfaces as SQL NULL.
2212 "bool_and" | "bool_or" => st.bool_acc.map_or(Value::Null, Value::Bool),
2213 // v7.32 (round-29) — variance / stddev. PG: `variance` ==
2214 // `var_samp`, `stddev` == `stddev_samp`. samp needs n >= 2
2215 // (n < 2 → NULL); pop needs n >= 1 (n == 1 → 0).
2216 "variance" | "var_samp" | "var_pop" | "stddev" | "stddev_samp" | "stddev_pop" => {
2217 let n = st.count;
2218 if n == 0 {
2219 return Value::Null;
2220 }
2221 let nf = n as f64;
2222 // Sum of squared deviations from the mean.
2223 let ss = st.sum_sq - (st.sum_float * st.sum_float) / nf;
2224 let pop = name.ends_with("_pop");
2225 let denom = if pop { nf } else { nf - 1.0 };
2226 if denom <= 0.0 {
2227 // var_samp / stddev (samp) with n == 1 → NULL.
2228 return Value::Null;
2229 }
2230 let var = (ss / denom).max(0.0); // clamp fp noise below 0
2231 if name.starts_with("stddev") {
2232 Value::Float(crate::eval::f64_sqrt(var))
2233 } else {
2234 Value::Float(var)
2235 }
2236 }
2237 // v7.32 (round-29) — bitwise aggregates: None (empty / all-NULL)
2238 // → SQL NULL.
2239 "bit_and" | "bit_or" | "bit_xor" => st.bit_acc.map_or(Value::Null, Value::BigInt),
2240 // v7.32 (round-29) — regression family. `regr_count` is the
2241 // paired n; everything else is NULL over an empty set. Terms
2242 // are the mean-centred sums of squares / cross-products.
2243 "regr_count" => Value::BigInt(st.reg_n),
2244 "covar_pop" | "covar_samp" | "corr" | "regr_avgx" | "regr_avgy" | "regr_slope"
2245 | "regr_intercept" | "regr_r2" | "regr_sxx" | "regr_syy" | "regr_sxy" => {
2246 let n = st.reg_n;
2247 if n == 0 {
2248 return Value::Null;
2249 }
2250 let nf = n as f64;
2251 let sxx = st.reg_sxx - st.reg_sx * st.reg_sx / nf;
2252 let syy = st.reg_syy - st.reg_sy * st.reg_sy / nf;
2253 let sxy = st.reg_sxy - st.reg_sx * st.reg_sy / nf;
2254 let avgx = st.reg_sx / nf;
2255 let avgy = st.reg_sy / nf;
2256 let out = match name {
2257 "regr_avgx" => Some(avgx),
2258 "regr_avgy" => Some(avgy),
2259 "regr_sxx" => Some(sxx),
2260 "regr_syy" => Some(syy),
2261 "regr_sxy" => Some(sxy),
2262 "covar_pop" => Some(sxy / nf),
2263 "covar_samp" => (n >= 2).then(|| sxy / (nf - 1.0)),
2264 "regr_slope" => (sxx != 0.0).then(|| sxy / sxx),
2265 "regr_intercept" => (sxx != 0.0).then(|| avgy - (sxy / sxx) * avgx),
2266 "corr" => {
2267 let d = sxx * syy;
2268 (d > 0.0).then(|| sxy / crate::eval::f64_sqrt(d))
2269 }
2270 // PG: NULL when sxx==0; 1 when syy==0 (and sxx>0).
2271 "regr_r2" => {
2272 if sxx == 0.0 {
2273 None
2274 } else if syy == 0.0 {
2275 Some(1.0)
2276 } else {
2277 Some((sxy * sxy) / (sxx * syy))
2278 }
2279 }
2280 _ => None,
2281 };
2282 out.map_or(Value::Null, Value::Float)
2283 }
2284 // v7.32 (round-29) — json_agg / jsonb_agg: a JSON array of every
2285 // collected element in row order; empty set → SQL NULL.
2286 "json_agg" | "jsonb_agg" => {
2287 if st.items.is_empty() {
2288 return Value::Null;
2289 }
2290 let mut out = String::from("[");
2291 for (i, item) in st.items.iter().enumerate() {
2292 if i > 0 {
2293 out.push_str(", ");
2294 }
2295 out.push_str(&crate::json::value_to_json_text(item));
2296 }
2297 out.push(']');
2298 Value::Json(out)
2299 }
2300 // v7.32 (round-29) — json_object_agg: a JSON object built from
2301 // the parallel key (`items`) / value (`aux_items`) streams.
2302 "json_object_agg" | "jsonb_object_agg" => {
2303 if st.items.is_empty() {
2304 return Value::Null;
2305 }
2306 let mut out = String::from("{");
2307 for (i, key) in st.items.iter().enumerate() {
2308 if i > 0 {
2309 out.push_str(", ");
2310 }
2311 // Object keys are always JSON strings (PG coerces).
2312 let key_text = match key {
2313 Value::Text(s) | Value::Json(s) => s.clone(),
2314 other => crate::json::value_to_json_text(other),
2315 };
2316 out.push_str(&crate::json::value_to_json_text(&Value::Text(key_text)));
2317 out.push_str(": ");
2318 let val = st.aux_items.get(i).unwrap_or(&Value::Null);
2319 out.push_str(&crate::json::value_to_json_text(val));
2320 }
2321 out.push('}');
2322 Value::Json(out)
2323 }
2324 // Ordered-set aggregates are finalized in `run` (they need the
2325 // sorted items + the direct fraction argument), never here.
2326 _ => unreachable!(),
2327 }
2328}
2329
2330/// v7.32 (round-29) — numeric coercion for the percentile interpolation.
2331fn agg_value_to_f64(v: &Value) -> Option<f64> {
2332 match v {
2333 Value::Int(n) => Some(f64::from(*n)),
2334 Value::SmallInt(n) => Some(f64::from(*n)),
2335 Value::BigInt(n) => Some(*n as f64),
2336 Value::Float(x) => Some(*x),
2337 _ => None,
2338 }
2339}
2340
2341/// v7.32 (round-29) — finalize a WITHIN GROUP aggregate. `st.items` is
2342/// already sorted by the `WITHIN GROUP (ORDER BY …)` spec. `direct` is
2343/// the evaluated direct argument: the fraction for `percentile_*`, the
2344/// hypothetical value for the hypothetical-set family (`rank` etc.),
2345/// and unused by `mode`. `order` is the (single) sort key, needed by
2346/// the hypothetical-set family to compare in the sort direction.
2347#[allow(
2348 clippy::cast_precision_loss,
2349 clippy::cast_possible_truncation,
2350 clippy::cast_sign_loss
2351)]
2352fn finalize_ordered_set(
2353 name: &str,
2354 st: &AggState,
2355 direct: Option<&Value>,
2356 order: Option<&spg_sql::ast::OrderBy>,
2357) -> Value {
2358 let fraction = direct;
2359 let items = &st.items;
2360 if items.is_empty() {
2361 // A hypothetical row ranks first over an empty group; the
2362 // distribution functions are 0 / divide-by-(n+1).
2363 return match name {
2364 "rank" | "dense_rank" => Value::BigInt(1),
2365 "percent_rank" => Value::Float(0.0),
2366 "cume_dist" => Value::Float(1.0),
2367 _ => Value::Null,
2368 };
2369 }
2370 let n = items.len();
2371 match name {
2372 // v7.32 (round-29) — hypothetical-set: the rank the direct value
2373 // would have if inserted into the group, in the sort direction.
2374 "rank" | "dense_rank" | "percent_rank" | "cume_dist" => {
2375 let Some(h) = fraction else {
2376 return Value::Null;
2377 };
2378 let (desc, nulls_first) = order.map_or((false, None), |o| (o.desc, o.nulls_first));
2379 let mut before = 0usize; // sort strictly before h
2380 let mut before_or_eq = 0usize; // sort before-or-peer with h
2381 let mut distinct_before = 0usize;
2382 let mut last_before: Option<&Value> = None;
2383 for it in items {
2384 match crate::order_by_value_cmp(desc, nulls_first, it, h) {
2385 core::cmp::Ordering::Less => {
2386 before += 1;
2387 before_or_eq += 1;
2388 if last_before
2389 .is_none_or(|p| value_cmp(p, it) != core::cmp::Ordering::Equal)
2390 {
2391 distinct_before += 1;
2392 last_before = Some(it);
2393 }
2394 }
2395 core::cmp::Ordering::Equal => before_or_eq += 1,
2396 core::cmp::Ordering::Greater => {}
2397 }
2398 }
2399 let nn = n as f64;
2400 match name {
2401 "rank" => Value::BigInt((before + 1) as i64),
2402 "dense_rank" => Value::BigInt((distinct_before + 1) as i64),
2403 "percent_rank" => Value::Float(before as f64 / nn),
2404 "cume_dist" => Value::Float((before_or_eq as f64 + 1.0) / (nn + 1.0)),
2405 _ => unreachable!(),
2406 }
2407 }
2408 // Most frequent value; equal values are adjacent in the sorted
2409 // run, and a frequency tie resolves to the earliest run (the
2410 // smallest value under an ascending sort), matching PG.
2411 "mode" => {
2412 let (mut best_i, mut best_cnt) = (0usize, 1usize);
2413 let (mut run_i, mut run_cnt) = (0usize, 1usize);
2414 for i in 1..n {
2415 if value_cmp(&items[i], &items[run_i]) == core::cmp::Ordering::Equal {
2416 run_cnt += 1;
2417 } else {
2418 run_i = i;
2419 run_cnt = 1;
2420 }
2421 if run_cnt > best_cnt {
2422 best_cnt = run_cnt;
2423 best_i = run_i;
2424 }
2425 }
2426 items[best_i].clone()
2427 }
2428 // The first value whose cumulative fraction reaches `f`.
2429 "percentile_disc" => {
2430 let f = fraction
2431 .and_then(agg_value_to_f64)
2432 .unwrap_or(0.0)
2433 .clamp(0.0, 1.0);
2434 let idx = if f <= 0.0 {
2435 0
2436 } else {
2437 (crate::eval::f64_ceil(f * n as f64) as usize)
2438 .saturating_sub(1)
2439 .min(n - 1)
2440 };
2441 items[idx].clone()
2442 }
2443 // Linear interpolation between the two bracketing values.
2444 "percentile_cont" => {
2445 let f = fraction
2446 .and_then(agg_value_to_f64)
2447 .unwrap_or(0.0)
2448 .clamp(0.0, 1.0);
2449 let Some(nums) = items
2450 .iter()
2451 .map(agg_value_to_f64)
2452 .collect::<Option<Vec<f64>>>()
2453 else {
2454 return Value::Null; // non-numeric ordered set
2455 };
2456 if n == 1 {
2457 return Value::Float(nums[0]);
2458 }
2459 let rank = f * (n as f64 - 1.0);
2460 let lo = crate::eval::f64_floor(rank) as usize;
2461 let hi = crate::eval::f64_ceil(rank) as usize;
2462 let frac = rank - lo as f64;
2463 Value::Float(nums[lo] + (nums[hi] - nums[lo]) * frac)
2464 }
2465 _ => unreachable!(),
2466 }
2467}
2468
2469fn infer_agg_type(spec: &AggSpec, schema_cols: &[ColumnSchema]) -> DataType {
2470 // v7.26 (round-20 C) — the argument's statically-derived shape
2471 // types MIN/MAX/SUM/array_agg properly; RowDescription used to
2472 // report TEXT for these, breaking every sqlx typed decode.
2473 let arg_ty = spec
2474 .arg
2475 .as_ref()
2476 .and_then(|a| crate::describe::describe_expr(a, schema_cols))
2477 .map(|shape| shape.ty);
2478 // v7.33 (array_agg argmax) — `(array_agg(x ORDER BY y))[1]` yields the
2479 // ELEMENT type (x), not the array type.
2480 if spec.first_ordered {
2481 return arg_ty.unwrap_or(DataType::Text);
2482 }
2483 match spec.name.as_str() {
2484 "count" | "count_star" => DataType::BigInt,
2485 "sum" => match arg_ty {
2486 Some(DataType::Float) => DataType::Float,
2487 _ => DataType::BigInt,
2488 },
2489 "avg" => DataType::Float,
2490 // v7.17.0 — string_agg always returns TEXT.
2491 "string_agg" => DataType::Text,
2492 "array_agg" => match arg_ty {
2493 Some(DataType::Int | DataType::SmallInt) => DataType::IntArray,
2494 Some(DataType::BigInt) => DataType::BigIntArray,
2495 _ => DataType::TextArray,
2496 },
2497 // v7.17.0 — boolean aggregates always return BOOL (nullable
2498 // — empty / all-NULL group → NULL).
2499 "bool_and" | "bool_or" => DataType::Bool,
2500 // v7.32 (round-29) — variance / stddev are floating point;
2501 // percentile_cont interpolates to float; the regression family
2502 // (except regr_count) is floating point.
2503 "stddev" | "stddev_samp" | "stddev_pop" | "variance" | "var_samp" | "var_pop"
2504 | "percentile_cont" | "covar_pop" | "covar_samp" | "corr" | "regr_avgx" | "regr_avgy"
2505 | "regr_slope" | "regr_intercept" | "regr_r2" | "regr_sxx" | "regr_syy" | "regr_sxy" => {
2506 DataType::Float
2507 }
2508 // v7.32 (round-29) — bitwise aggregates, regr_count, and the
2509 // integer hypothetical-set ranks return an integer.
2510 "bit_and" | "bit_or" | "bit_xor" | "regr_count" | "rank" | "dense_rank" => DataType::BigInt,
2511 // v7.32 (round-29) — hypothetical-set distribution functions.
2512 "percent_rank" | "cume_dist" => DataType::Float,
2513 // v7.32 (round-29) — JSON aggregates return JSON.
2514 "json_agg" | "jsonb_agg" | "json_object_agg" | "jsonb_object_agg" => DataType::Json,
2515 // min/max, percentile_disc, mode, and anything pass-through:
2516 // the argument's shape (for ordered-set aggs `spec.arg` is the
2517 // WITHIN GROUP value expression).
2518 _ => arg_ty.unwrap_or(DataType::Text),
2519 }
2520}
2521
2522fn agg_or_group_type(e: &Expr, synth: &[ColumnSchema]) -> DataType {
2523 if let Expr::Column(c) = e
2524 && let Some(s) = synth.iter().find(|s| s.name == c.name)
2525 {
2526 return s.ty;
2527 }
2528 // v7.26 (round-20 C) — compound expressions over aggregates
2529 // (COALESCE(BOOL_OR(…), false), (array_agg(…))[1], CASE …)
2530 // derive their shape statically against the synth schema; the
2531 // old Text fallback broke sqlx typed decodes of exactly these
2532 // columns.
2533 crate::describe::describe_expr(e, synth)
2534 .map(|shape| shape.ty)
2535 .unwrap_or(DataType::Text)
2536}
2537
2538fn rewrite_expr(e: &Expr, group_exprs: &[Expr], aggs: &[AggSpec]) -> Expr {
2539 // v7.33 (array_agg argmax) — `(array_agg(x ORDER BY y))[1]` rewrites
2540 // to its first_ordered synth column, consuming the subscript. Checked
2541 // before the AggregateOrdered/recursion arms (which would otherwise
2542 // rewrite the inner array_agg and leave the subscript). Same matcher
2543 // as collect_aggregates, so the spec it finds is the one collected.
2544 if let Some((arg, order_by, filter)) = first_ordered_array_agg(e) {
2545 let arg_owned = Some(arg.clone());
2546 let filter_owned = filter.cloned();
2547 for (i, spec) in aggs.iter().enumerate() {
2548 if spec.first_ordered
2549 && spec.name == "array_agg"
2550 && spec.arg == arg_owned
2551 && spec.order_by == *order_by
2552 && spec.filter == filter_owned
2553 {
2554 return Expr::Column(spg_sql::ast::ColumnName {
2555 qualifier: None,
2556 name: format!("__agg_{i}"),
2557 });
2558 }
2559 }
2560 }
2561 // v7.24 (round-16 A) — ordered aggregate: match on the inner
2562 // call PLUS the ordering keys.
2563 if let Expr::AggregateOrdered {
2564 call,
2565 order_by,
2566 distinct,
2567 filter,
2568 } = e
2569 && let Expr::FunctionCall { name, args } = call.as_ref()
2570 {
2571 let lower = name.to_ascii_lowercase();
2572 if is_aggregate_name(&lower) {
2573 let canonical: &str = if lower == "every" { "bool_and" } else { &lower };
2574 // Mirror collect_aggregates: ordered-set aggregates take the
2575 // value from the sort spec and the in-parens arg as direct.
2576 let (arg, direct_arg) = if is_within_group_name(canonical) {
2577 (
2578 order_by.first().map(|o| o.expr.clone()),
2579 args.first().cloned(),
2580 )
2581 } else {
2582 (args.first().cloned(), None)
2583 };
2584 let arg2 = if agg_uses_second_arg(canonical) {
2585 args.get(1).cloned()
2586 } else {
2587 None
2588 };
2589 let filter_owned = filter.as_deref().cloned();
2590 for (i, spec) in aggs.iter().enumerate() {
2591 if spec.name == canonical
2592 && spec.arg == arg
2593 && spec.arg2 == arg2
2594 && spec.distinct == *distinct
2595 && spec.order_by == *order_by
2596 && spec.filter == filter_owned
2597 && spec.direct_arg == direct_arg
2598 {
2599 return Expr::Column(spg_sql::ast::ColumnName {
2600 qualifier: None,
2601 name: format!("__agg_{i}"),
2602 });
2603 }
2604 }
2605 }
2606 }
2607 // Match aggregate FunctionCalls first — they sit outside group_by.
2608 if let Expr::FunctionCall { name, args } = e {
2609 let lower = name.to_ascii_lowercase();
2610 if is_aggregate_name(&lower) {
2611 let arg = if lower == "count_star" {
2612 None
2613 } else {
2614 args.first().cloned()
2615 };
2616 // v7.17.0 — match the spec we registered for
2617 // string_agg(value, separator) on the full pair; v7.32 also
2618 // the regression family and json_object_agg.
2619 let arg2 = if agg_uses_second_arg(&lower) {
2620 args.get(1).cloned()
2621 } else {
2622 None
2623 };
2624 // v7.17.0 — `every` collapses into `bool_and` at
2625 // collection; mirror that here so the rewrite finds
2626 // the matching synth column.
2627 let canonical: &str = if lower == "every" {
2628 "bool_and"
2629 } else {
2630 lower.as_str()
2631 };
2632 for (i, spec) in aggs.iter().enumerate() {
2633 if spec.name == canonical
2634 && spec.arg == arg
2635 && spec.arg2 == arg2
2636 && !spec.distinct
2637 && spec.order_by.is_empty()
2638 {
2639 return Expr::Column(spg_sql::ast::ColumnName {
2640 qualifier: None,
2641 name: format!("__agg_{i}"),
2642 });
2643 }
2644 }
2645 }
2646 }
2647 // Match a group_by expression by AST equality.
2648 for (i, g) in group_exprs.iter().enumerate() {
2649 if g == e {
2650 return Expr::Column(spg_sql::ast::ColumnName {
2651 qualifier: None,
2652 name: format!("__grp_{i}"),
2653 });
2654 }
2655 }
2656 // Recurse into children.
2657 match e {
2658 Expr::AggregateOrdered {
2659 call,
2660 order_by,
2661 distinct,
2662 filter,
2663 } => Expr::AggregateOrdered {
2664 call: Box::new(rewrite_expr(call, group_exprs, aggs)),
2665 distinct: *distinct,
2666 order_by: order_by
2667 .iter()
2668 .map(|o| spg_sql::ast::OrderBy {
2669 expr: rewrite_expr(&o.expr, group_exprs, aggs),
2670 desc: o.desc,
2671 nulls_first: o.nulls_first,
2672 })
2673 .collect(),
2674 // The filter is evaluated against SOURCE rows during
2675 // accumulation, never against synth rows — keep it as-is.
2676 filter: filter.clone(),
2677 },
2678 Expr::Binary { lhs, op, rhs } => Expr::Binary {
2679 lhs: Box::new(rewrite_expr(lhs, group_exprs, aggs)),
2680 op: *op,
2681 rhs: Box::new(rewrite_expr(rhs, group_exprs, aggs)),
2682 },
2683 Expr::Unary { op, expr } => Expr::Unary {
2684 op: *op,
2685 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
2686 },
2687 Expr::Cast { expr, target } => Expr::Cast {
2688 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
2689 target: *target,
2690 },
2691 Expr::IsNull { expr, negated } => Expr::IsNull {
2692 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
2693 negated: *negated,
2694 },
2695 Expr::FunctionCall { name, args } => Expr::FunctionCall {
2696 name: name.clone(),
2697 args: args
2698 .iter()
2699 .map(|a| rewrite_expr(a, group_exprs, aggs))
2700 .collect(),
2701 },
2702 Expr::Like {
2703 expr,
2704 pattern,
2705 negated,
2706 case_insensitive,
2707 } => Expr::Like {
2708 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
2709 pattern: Box::new(rewrite_expr(pattern, group_exprs, aggs)),
2710 negated: *negated,
2711 case_insensitive: *case_insensitive,
2712 },
2713 Expr::Extract { field, source } => Expr::Extract {
2714 field: *field,
2715 source: Box::new(rewrite_expr(source, group_exprs, aggs)),
2716 },
2717 // v7.25.2 (round-19 A) — subquery nodes: rewrite group-key
2718 // references INSIDE the body to `__grp_N` so the correlated
2719 // resolver can substitute them against the synthesised group
2720 // row (aggs are NOT matched inside the body — a COUNT in the
2721 // subquery is the subquery's own aggregate).
2722 Expr::ScalarSubquery(s) => {
2723 Expr::ScalarSubquery(Box::new(rewrite_group_keys_in_select(s, group_exprs)))
2724 }
2725 Expr::Exists { subquery, negated } => Expr::Exists {
2726 subquery: Box::new(rewrite_group_keys_in_select(subquery, group_exprs)),
2727 negated: *negated,
2728 },
2729 Expr::InSubquery {
2730 expr,
2731 subquery,
2732 negated,
2733 } => Expr::InSubquery {
2734 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
2735 subquery: Box::new(rewrite_group_keys_in_select(subquery, group_exprs)),
2736 negated: *negated,
2737 },
2738 // v4.12 window / Literal / Column — clone-pass (these don't
2739 // participate in aggregate rewrite).
2740 Expr::WindowFunction { .. } | Expr::Literal(_) | Expr::Placeholder(_) | Expr::Column(_) => {
2741 e.clone()
2742 }
2743 // v7.10.10 — recurse children for array nodes.
2744 Expr::Array(items) => Expr::Array(
2745 items
2746 .iter()
2747 .map(|elem| rewrite_expr(elem, group_exprs, aggs))
2748 .collect(),
2749 ),
2750 Expr::ArraySubscript { target, index } => Expr::ArraySubscript {
2751 target: Box::new(rewrite_expr(target, group_exprs, aggs)),
2752 index: Box::new(rewrite_expr(index, group_exprs, aggs)),
2753 },
2754 Expr::AnyAll {
2755 expr,
2756 op,
2757 array,
2758 is_any,
2759 } => Expr::AnyAll {
2760 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
2761 op: *op,
2762 array: Box::new(rewrite_expr(array, group_exprs, aggs)),
2763 is_any: *is_any,
2764 },
2765 Expr::InList {
2766 expr,
2767 list,
2768 negated,
2769 } => Expr::InList {
2770 expr: Box::new(rewrite_expr(expr, group_exprs, aggs)),
2771 list: list
2772 .iter()
2773 .map(|item| rewrite_expr(item, group_exprs, aggs))
2774 .collect(),
2775 negated: *negated,
2776 },
2777 Expr::Case {
2778 operand,
2779 branches,
2780 else_branch,
2781 } => Expr::Case {
2782 operand: operand
2783 .as_deref()
2784 .map(|o| Box::new(rewrite_expr(o, group_exprs, aggs))),
2785 branches: branches
2786 .iter()
2787 .map(|(w, t)| {
2788 (
2789 rewrite_expr(w, group_exprs, aggs),
2790 rewrite_expr(t, group_exprs, aggs),
2791 )
2792 })
2793 .collect(),
2794 else_branch: else_branch
2795 .as_deref()
2796 .map(|e| Box::new(rewrite_expr(e, group_exprs, aggs))),
2797 },
2798 }
2799}
2800
2801/// v7.25.2 (round-19 A) — rewrite group-key references inside a
2802/// subquery body to `__grp_N` synthetic columns (aggregates are
2803/// not touched: empty spec list). Runs through the canonical
2804/// Select walker so every expression slot is covered.
2805fn rewrite_group_keys_in_select(
2806 s: &spg_sql::ast::SelectStatement,
2807 group_exprs: &[Expr],
2808) -> spg_sql::ast::SelectStatement {
2809 let mut out = s.clone();
2810 let _ = crate::walk_select_exprs_mut(&mut out, &mut |e| {
2811 *e = rewrite_expr(e, group_exprs, &[]);
2812 Ok(())
2813 });
2814 out
2815}
2816
2817/// Canonical string key for a tuple of group values. Used as map key.
2818/// Per-value group-key encoding (shared by owned and borrowed paths).
2819fn encode_one(out: &mut String, v: &Value) {
2820 use core::fmt::Write;
2821 match v {
2822 Value::Null => out.push_str("N|"),
2823 // v7.36 (perf — mailrs Phase 1) — switch the integer / float
2824 // encoders to `write!`. `n.to_string()` allocates a fresh
2825 // `String` per cell just to push its bytes into the
2826 // (already-cleared) reuse buffer — for the 25 k-row JOIN
2827 // probe in `count_messages` that's 25 k heap allocs per
2828 // query. `write!(&mut String, ...)` formats straight into
2829 // the buffer; no intermediate alloc.
2830 Value::SmallInt(n) => {
2831 let _ = write!(out, "s{n}|");
2832 }
2833 Value::Int(n) => {
2834 let _ = write!(out, "I{n}|");
2835 }
2836 Value::BigInt(n) => {
2837 let _ = write!(out, "B{n}|");
2838 }
2839 Value::Float(x) => {
2840 let _ = write!(out, "F{x}|");
2841 }
2842 Value::Bool(b) => {
2843 out.push(if *b { 'T' } else { 'f' });
2844 out.push('|');
2845 }
2846 Value::Text(s) => {
2847 out.push('S');
2848 out.push_str(s);
2849 out.push('|');
2850 }
2851 Value::Vector(v) => {
2852 out.push('V');
2853 for x in v {
2854 out.push_str(&x.to_string());
2855 out.push(',');
2856 }
2857 out.push('|');
2858 }
2859 // v6.0.1: GROUP BY on a `VECTOR(N) USING SQ8` column.
2860 // Two cells with byte-identical `(min, max, bytes)`
2861 // share the same group; equivalence is byte-equality
2862 // (same as f32 grouping today — neither path tries to
2863 // normalise nan/-0).
2864 Value::Sq8Vector(q) => {
2865 out.push('Q');
2866 out.push_str(&q.min.to_string());
2867 out.push('@');
2868 out.push_str(&q.max.to_string());
2869 out.push(':');
2870 for b in &q.bytes {
2871 out.push_str(&b.to_string());
2872 out.push(',');
2873 }
2874 out.push('|');
2875 }
2876 // v6.0.3: GROUP BY on a `VECTOR(N) USING HALF` column.
2877 // Byte-equality over the raw u16 bits; matches the SQ8
2878 // path's byte-key model.
2879 Value::HalfVector(h) => {
2880 out.push('H');
2881 for b in &h.bytes {
2882 out.push_str(&b.to_string());
2883 out.push(',');
2884 }
2885 out.push('|');
2886 }
2887 Value::Numeric { scaled, scale } => {
2888 out.push('D');
2889 out.push_str(&scaled.to_string());
2890 out.push('@');
2891 out.push_str(&scale.to_string());
2892 out.push('|');
2893 }
2894 Value::Date(d) => {
2895 out.push('d');
2896 out.push_str(&d.to_string());
2897 out.push('|');
2898 }
2899 Value::Timestamp(t) => {
2900 out.push('t');
2901 out.push_str(&t.to_string());
2902 out.push('|');
2903 }
2904 Value::Interval { months, micros } => {
2905 out.push('i');
2906 out.push_str(&months.to_string());
2907 out.push('m');
2908 out.push_str(µs.to_string());
2909 out.push('|');
2910 }
2911 Value::Json(s) => {
2912 out.push('j');
2913 out.push_str(s);
2914 out.push('|');
2915 }
2916 // v7.5.0 — Value is #[non_exhaustive] for downstream
2917 // forward-compat. Any future variant lacking explicit
2918 // handling here will share a debug-derived group key,
2919 // which is observably wrong but won't crash.
2920 _ => {
2921 out.push('?');
2922 out.push_str(&format!("{v:?}"));
2923 out.push('|');
2924 }
2925 }
2926}
2927
2928/// v7.30 (perf campaign) - encode from borrowed cells without
2929/// materialising an owned Vec<Value> first.
2930pub(crate) fn encode_key_refs(vals: &[&Value]) -> String {
2931 let mut out = String::new();
2932 for v in vals {
2933 encode_one(&mut out, v);
2934 }
2935 out
2936}
2937
2938/// v7.31 (perf 3e) — encode into a caller-owned scratch buffer.
2939/// The per-row key paths (group hash, DISTINCT set, join build/
2940/// probe) ran 24k+ String allocations per query through the
2941/// allocator just to LOOK UP a map; the scratch form allocates
2942/// only when a map actually has to take ownership (vacant insert).
2943pub(crate) fn encode_key_refs_into(vals: &[&Value], out: &mut String) {
2944 out.clear();
2945 for v in vals {
2946 encode_one(out, v);
2947 }
2948}
2949
2950pub(crate) fn encode_key(vals: &[Value]) -> String {
2951 let mut out = String::new();
2952 for v in vals {
2953 encode_one(&mut out, v);
2954 }
2955 out
2956}
2957
2958#[allow(clippy::cast_precision_loss)]
2959fn value_cmp(a: &Value, b: &Value) -> core::cmp::Ordering {
2960 use core::cmp::Ordering::Equal;
2961 match (a, b) {
2962 (Value::Null, Value::Null) => Equal,
2963 (Value::Null, _) => core::cmp::Ordering::Greater, // NULLs last
2964 (_, Value::Null) => core::cmp::Ordering::Less,
2965 (Value::Int(x), Value::Int(y)) => x.cmp(y),
2966 (Value::BigInt(x), Value::BigInt(y)) => x.cmp(y),
2967 (Value::Int(x), Value::BigInt(y)) => i64::from(*x).cmp(y),
2968 (Value::BigInt(x), Value::Int(y)) => x.cmp(&i64::from(*y)),
2969 (Value::Float(x), Value::Float(y)) => x.partial_cmp(y).unwrap_or(Equal),
2970 (Value::Int(x), Value::Float(y)) => f64::from(*x).partial_cmp(y).unwrap_or(Equal),
2971 (Value::Float(x), Value::Int(y)) => x.partial_cmp(&f64::from(*y)).unwrap_or(Equal),
2972 (Value::BigInt(x), Value::Float(y)) => (*x as f64).partial_cmp(y).unwrap_or(Equal),
2973 (Value::Float(x), Value::BigInt(y)) => x.partial_cmp(&(*y as f64)).unwrap_or(Equal),
2974 (Value::Text(x), Value::Text(y)) => x.cmp(y),
2975 (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
2976 _ => Equal,
2977 }
2978}