Skip to content

Commit f830b03

Browse files
zzylolclaude
andauthored
feat(sql): lower ClickHouse argMax/argMin to AggIntent::Extension (#242)
argMax(arg, val)/argMin(arg, val) ("arg's value from the row where val is maximal/minimal") are two-column, row-selecting aggregates -- unlike every existing AggIntent reducer (Sum/Min/Max/Avg/...), which folds one column to a value derived from itself, these return a *different* column's value, selected by which row extremizes a second column. No existing AggIntent shape fits. Design decision (issue #232 asked for research before picking): a repo-wide grep for any existing arg-max/arg-min concept (PromQL front end, other SQL dialects, docs) found nothing -- only the 3 corpus occurrences this closes. sql-function-catalog's own KNOWN_UNMAPPED_NATIVE_FUNCTIONS already documents the same "two-column reducer" and "value from a particular row" gaps for corr/covar/regr_* and first_value/last_value/nth_value, for the identical reason: AggIntent's input_col()/resolve_agg_intent/aggregate_output_schema machinery is single-column by construction. Per AggIntent::Extension's own bar ("core only grows for intents >= 2 deployment models actually use"), with exactly one consumer found, this lowers to AggIntent::Extension { ext_kind: "arg_max" | "arg_min", payload } rather than earning first-class ArgMax/ArgMin core variants. This required zero changes to crates/types: Extension already has complete plumbing (resolve, output_column, requires/is_per_series, and downstream PassThrough handling in asap-aware-mapping), which is itself evidence the Extension escape hatch is doing its job. Implementation: - sql-function-catalog: new RewriteKind::PassThrough for a CLICKHOUSE_BUILTINS entry with no native DataFusion shape to rewrite to at all (argmax/argmin, arity 2) -- a new pattern alongside the existing CountDistinct/CountIfToSum rewrite targets. - frontend-sql: ClickHouseBuiltinRewrite treats PassThrough as a no-op, so the call reaches lower_agg_intent under its own ClickHouse name. New lower_arg_selector() validates both arguments as bare columns (reducer_col's existing "no expression arguments" rule, issue #115, generalized to 2 args) and builds the Extension node directly; core never resolves Extension's payload, so both columns are kept as ColumnRef names rather than positionally-bound ColumnIds. - Documented, not fixed (matching #230's splitByChar-array-indexing precedent): DerivedCols::rewrite_agg only special-cases an aggregate's *first* argument, so val would be dropped from a Project inserted for some unrelated reason in the same query. Does not affect the corpus today -- all 3 argMax uses are bare columns. Corpus tally (bgp_jan2024_workload): Lowered 145 -> 148, Plan 40 -> 37, every other category unchanged -- all 3 argMax occurrences move cleanly to Lowered with no companion gap. Closes #232 Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 86670ef commit f830b03

6 files changed

Lines changed: 202 additions & 10 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/frontend-sql/Cargo.toml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,9 @@ edition = "2021"
1111
asap-types = { path = "../types" }
1212
asap-sql-function-catalog = { path = "../sql-function-catalog" }
1313
datafusion = "43"
14+
# `AggIntent::Extension.payload` for ClickHouse's argMax/argMin (issue #232)
15+
# -- opaque JSON, matching how `asap-types` itself carries it.
16+
serde_json = "1"
1417

1518
[dev-dependencies]
1619
tokio = { version = "1", features = ["rt", "macros", "rt-multi-thread"] }

‎crates/frontend-sql/src/sql/mod.rs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1063,6 +1063,10 @@ impl FunctionRewrite for ClickHouseBuiltinRewrite {
10631063
return Ok(Transformed::no(Expr::AggregateFunction(f)));
10641064
};
10651065
let rewritten = match builtin.rewrite {
1066+
// No native DataFusion shape to become — leave the call exactly
1067+
// as DataFusion's planner parsed it. `lower_agg_intent` handles
1068+
// the ClickHouse name (`argMax`/`argMin`) directly (issue #232).
1069+
RewriteKind::PassThrough => return Ok(Transformed::no(Expr::AggregateFunction(f))),
10661070
// `f(args...)` -> `count(args...) DISTINCT` — `lower_agg_intent`
10671071
// already maps `count` + `DISTINCT` to `AggIntent::Cardinality`.
10681072
RewriteKind::CountDistinct => AggregateFunction::new_udf(
@@ -1120,6 +1124,14 @@ fn lower_agg_intent(expr: &Expr) -> Result<AggIntent<ColumnRef>, LoweringError>
11201124
Expr::Alias(a) => lower_agg_intent(&a.expr),
11211125
Expr::AggregateFunction(agg_fn) => {
11221126
let name = agg_fn.func.name().to_lowercase();
1127+
// ClickHouse's row-selecting `argMax`/`argMin` — `RewriteKind::
1128+
// PassThrough` in the catalog, so the call reaches here under its
1129+
// own name rather than a native DataFusion aggregate. Handled
1130+
// before the `NATIVE_FUNCTIONS` lookup below since neither name
1131+
// is in that table (issue #232).
1132+
if let Some(intent) = lower_arg_selector(&name, &agg_fn.args)? {
1133+
return Ok(intent);
1134+
}
11231135
let semantic = asap_sql_function_catalog::lookup_native(&name)
11241136
.ok_or_else(|| LoweringError::UnsupportedAggregate(name.clone()))?;
11251137
// The canonical intent algebra has no DISTINCT modifier for the
@@ -1194,6 +1206,62 @@ fn lower_agg_intent(expr: &Expr) -> Result<AggIntent<ColumnRef>, LoweringError>
11941206
}
11951207
}
11961208

1209+
/// ClickHouse's row-selecting `argMax(arg, val)` / `argMin(arg, val)` —
1210+
/// "return `arg`'s value from the row where `val` is maximal/minimal".
1211+
/// `Some(name)` for `"argmax"`/`"argmin"`, `None` for every other name (the
1212+
/// caller falls through to the ordinary `NATIVE_FUNCTIONS` path).
1213+
///
1214+
/// Unlike every existing `AggIntent` reducer (`Sum`/`Min`/`Max`/`Avg`/…),
1215+
/// which folds *one* column to a value derived from itself, this is a
1216+
/// two-column, row-selecting aggregate: it returns a *different* column's
1217+
/// value, selected by which row maximizes/minimizes a second column. No
1218+
/// existing `AggIntent` shape fits, and — per its own doc comment's
1219+
/// "core only grows for intents ≥2 deployment models actually use" bar —
1220+
/// a repo-wide search (PromQL front end, the other SQL dialects, docs) found
1221+
/// no second deployment model wanting this shape, so this lowers to
1222+
/// `AggIntent::Extension` rather than earning a first-class `ArgMax`/`ArgMin`
1223+
/// core variant (issue #232). Core treats `Extension` opaquely: both columns
1224+
/// are kept only as validated bare-column names in `payload` (`reducer_col`'s
1225+
/// same "no expression arguments" rule, issue #115) — they are **not** run
1226+
/// through `resolve_agg_intent`'s positional `ColumnRef` -> `ColumnId`
1227+
/// binding the way a real reducer's `col` is, since `Extension` carries no
1228+
/// typed column field for core to resolve. A deployment model that wants to
1229+
/// actually bind/execute `argMax`/`argMin` is expected to re-derive that
1230+
/// itself from `payload`, per `AggIntent::Extension`'s own doc.
1231+
///
1232+
/// Known scope gap, not chased down here (matching issue #230's precedent
1233+
/// for other structural-only gaps like `splitByChar` array-indexing):
1234+
/// `DerivedCols::rewrite_agg` (this module) only materializes/passes through
1235+
/// an aggregate's *first* argument when deriving columns beneath the
1236+
/// `Aggregate` node, so `val` here would be silently dropped from a `Project`
1237+
/// DataFusion's planner inserts for some *other* reason in the same query
1238+
/// (e.g. a sibling `SUM(a * b)`). Every corpus `argMax` use has both
1239+
/// arguments as bare columns of the scanned table and no such `Project` is
1240+
/// inserted, so this does not affect `bgp_jan2024_workload` today.
1241+
fn lower_arg_selector(
1242+
name: &str,
1243+
args: &[Expr],
1244+
) -> Result<Option<AggIntent<ColumnRef>>, LoweringError> {
1245+
let ext_kind = match name {
1246+
"argmax" => "arg_max",
1247+
"argmin" => "arg_min",
1248+
_ => return Ok(None),
1249+
};
1250+
let [arg, val] = args else {
1251+
unreachable!(
1252+
"{name}'s stub signature (asap_sql_function_catalog::CLICKHOUSE_BUILTINS) fixes \
1253+
its arity at 2 -- the planner already rejected any other argument count before \
1254+
lower_agg_intent runs"
1255+
);
1256+
};
1257+
let arg_col = reducer_col(name, std::slice::from_ref(arg))?;
1258+
let val_col = reducer_col(name, std::slice::from_ref(val))?;
1259+
Ok(Some(AggIntent::Extension {
1260+
ext_kind: ext_kind.to_string(),
1261+
payload: serde_json::json!({ "arg_col": arg_col, "val_col": val_col }),
1262+
}))
1263+
}
1264+
11971265
/// The name an `IN (subquery)`'s key column is projected under, so the join
11981266
/// predicate cannot bind it to a same-named column of the outer relation.
11991267
const IN_SUBQUERY_KEY: &str = "__asap_in_key";

‎crates/frontend-sql/tests/bgp_jan2024_workload/bgp_jan2024_workload.rs‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -182,8 +182,13 @@ async fn corpus_lowering_matches_the_pinned_aggregate_tally() {
182182
// caveat" section) -- getting these functions' *names* to lower to a
183183
// structurally correct `FunctionCall` node is what's in scope here, not
184184
// array/map indexing or interval-literal conversion.
185-
expect(Category::Lowered, 145);
186-
expect(Category::Plan, 40);
185+
// `argMax` support (issue #232 -- `AggIntent::Extension`, catalog-driven
186+
// `RewriteKind::PassThrough`) clears the "unknown function: argmax" `Plan`
187+
// failure for all 3 corpus occurrences: every one has both arguments as
188+
// bare columns of the scanned table, so all 3 lower end to end with no
189+
// companion gap (unlike `splitByChar`'s array-indexing gap, #230).
190+
expect(Category::Lowered, 148);
191+
expect(Category::Plan, 37);
187192
expect(Category::Schema, 0);
188193
expect(Category::Parse, 0);
189194
// One query that used to fail at `uniqExact` (`Plan`) now clears that

‎crates/frontend-sql/tests/sql_lowering.rs‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1466,3 +1466,91 @@ async fn count_if_composes_with_group_by() {
14661466
assert_eq!(*by, GroupKeys::by(vec![0]));
14671467
assert!(matches!(measures.as_slice(), [AggIntent::Sum { .. }]));
14681468
}
1469+
1470+
// ── Issue #232: argMax/argMin -- AggIntent::Extension, not a first-class
1471+
// core variant. A repo-wide search (PromQL front end, other SQL dialects,
1472+
// docs) turned up no second deployment model wanting this two-column,
1473+
// row-selecting shape, so per `AggIntent::Extension`'s own "core only grows
1474+
// for intents ≥2 deployment models actually use" bar, it stays an opaque
1475+
// `Extension` rather than a new `ArgMax`/`ArgMin` core variant. Unlike
1476+
// `countIf`/`uniqExact`, there is no native DataFusion aggregate shape to
1477+
// rewrite to (`RewriteKind::PassThrough`) -- `lower_agg_intent` builds the
1478+
// `AggIntent` directly from the ClickHouse name. ─────────────────────────
1479+
1480+
#[tokio::test]
1481+
async fn arg_max_lowers_to_an_extension_intent() {
1482+
// No existing `AggIntent` reducer fits: every one folds one column to a
1483+
// value derived from itself, while `argMax(arg, val)` returns a
1484+
// *different* column's value, selected by which row maximizes a second.
1485+
let qe = lower_clickhouse(
1486+
"SELECT service, argMax(service, latency) AS busiest FROM metrics GROUP BY service",
1487+
)
1488+
.await;
1489+
let (by, measures) = find_aggregate(&qe).expect("expected an Aggregate");
1490+
assert_eq!(
1491+
*by,
1492+
GroupKeys::by(vec![1]),
1493+
"grouped by `service` (schema index 1)"
1494+
);
1495+
assert!(
1496+
matches!(
1497+
measures.as_slice(),
1498+
[AggIntent::Extension { ext_kind, .. }] if ext_kind == "arg_max"
1499+
),
1500+
"expected Extension {{ ext_kind: \"arg_max\", .. }}, got {measures:?}"
1501+
);
1502+
}
1503+
1504+
#[tokio::test]
1505+
async fn arg_min_lowers_to_its_own_extension_kind() {
1506+
let qe = lower_clickhouse("SELECT argMin(service, latency) FROM metrics").await;
1507+
let (by, measures) = find_aggregate(&qe).expect("expected an Aggregate");
1508+
assert!(by.is_empty());
1509+
assert!(
1510+
matches!(
1511+
measures.as_slice(),
1512+
[AggIntent::Extension { ext_kind, .. }] if ext_kind == "arg_min"
1513+
),
1514+
"expected Extension {{ ext_kind: \"arg_min\", .. }}, got {measures:?}"
1515+
);
1516+
}
1517+
1518+
#[tokio::test]
1519+
async fn arg_max_payload_preserves_both_column_names() {
1520+
// Core never resolves an `Extension`'s payload, so both columns are kept
1521+
// as validated bare-column `ColumnRef`s in `payload`, not run through
1522+
// positional `ColumnId` binding -- see `lower_arg_selector`'s doc.
1523+
let qe = lower_clickhouse("SELECT argMax(service, latency) AS m FROM metrics").await;
1524+
let (_, measures) = find_aggregate(&qe).expect("expected an Aggregate");
1525+
let AggIntent::Extension { payload, .. } = &measures[0] else {
1526+
panic!("expected an Extension intent, got {:?}", measures[0]);
1527+
};
1528+
let named = |key: &str| {
1529+
payload
1530+
.get(key)
1531+
.and_then(|c| c.get("Named"))
1532+
.and_then(|n| n.as_str())
1533+
.map(str::to_string)
1534+
};
1535+
assert_eq!(named("arg_col"), Some("service".to_string()));
1536+
assert_eq!(named("val_col"), Some("latency".to_string()));
1537+
}
1538+
1539+
#[tokio::test]
1540+
async fn arg_max_rejects_a_non_column_argument() {
1541+
// Same "bare column only" rule as every other reducer (`reducer_col`,
1542+
// issue #115) -- an expression argument is rejected, not silently
1543+
// dropped or materialized into the wrong column.
1544+
let err = lower_sql_dialect(
1545+
"SELECT argMax(service, latency * 2) FROM metrics",
1546+
&catalog(),
1547+
SqlDialect::ClickhouseSQL,
1548+
AccuracyTarget::Exact,
1549+
)
1550+
.await
1551+
.expect_err("argMax over a non-column expression must be rejected");
1552+
assert!(
1553+
format!("{err}").contains("non-column expression"),
1554+
"got {err}"
1555+
);
1556+
}

‎crates/sql-function-catalog/src/lib.rs‎

Lines changed: 35 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,18 @@
2626
//! must be a bare column" rule are call-site logic, not per-function data,
2727
//! and stay in `asap-frontend-sql`.
2828
//! - [`CLICKHOUSE_BUILTINS`] -- ClickHouse-only *aggregate* names DataFusion
29-
//! doesn't know at all (`uniqExact`, `countIf`). Each entry additionally
30-
//! carries a [`RewriteKind`]: the native DataFusion aggregate shape the
31-
//! call rewrites to before `lower_agg_intent` (or DataFusion's own physical
32-
//! planner) ever has to understand the ClickHouse name itself. This is
33-
//! what generalizes `uniqExact`'s old bespoke `UniqExactRewrite` +
34-
//! `uniq_exact_udaf` pair (issue #221): a new builtin that rewrites to an
35-
//! already-handled shape is a new entry in this table, not a new
36-
//! `FunctionRewrite` impl and a new stub-`AggregateUDF` constructor.
29+
//! doesn't know at all (`uniqExact`, `countIf`, `argMax`, `argMin`). Each
30+
//! entry additionally carries a [`RewriteKind`]: either the native
31+
//! DataFusion aggregate shape the call rewrites to before
32+
//! `lower_agg_intent` (or DataFusion's own physical planner) ever has to
33+
//! understand the ClickHouse name itself -- this is what generalizes
34+
//! `uniqExact`'s old bespoke `UniqExactRewrite` + `uniq_exact_udaf` pair
35+
//! (issue #221): a new builtin that rewrites to an already-handled shape is
36+
//! a new entry in this table, not a new `FunctionRewrite` impl and a new
37+
//! stub-`AggregateUDF` constructor -- or [`RewriteKind::PassThrough`] for a
38+
//! builtin with no native shape to rewrite to at all (`argMax`/`argMin`,
39+
//! issue #232): the call survives unchanged and `lower_agg_intent` handles
40+
//! the ClickHouse name itself directly.
3741
//! - [`CLICKHOUSE_SCALAR_BUILTINS`] -- ClickHouse-only *scalar* names
3842
//! DataFusion doesn't know at all (`splitByChar`, `toDate`, `match`,
3943
//! the `toStartOf*` family, `startsWith`, `positionCaseInsensitive`). No
@@ -229,6 +233,16 @@ pub enum RewriteKind {
229233
/// for any reducer over an expression (issue #110) -- so no new
230234
/// `AggIntent` variant or lowering path is needed either.
231235
CountIfToSum,
236+
/// No native DataFusion aggregate shape to rewrite to at all -- the call
237+
/// survives unchanged (`ClickHouseBuiltinRewrite` is a no-op for it) and
238+
/// reaches `lower_agg_intent`, which builds an `AggIntent` directly from
239+
/// the ClickHouse name itself. `argMax`/`argMin` (issue #232): a
240+
/// two-column, row-selecting aggregate ("return `arg`'s value from the
241+
/// row where `val` is maximal") that fits no existing single-column
242+
/// `AggIntent` reducer shape and has no native DataFusion equivalent to
243+
/// borrow one from -- unlike `CountDistinct`/`CountIfToSum`, there is no
244+
/// "already-handled shape" for this rewrite to target.
245+
PassThrough,
232246
}
233247

234248
/// One [`CLICKHOUSE_BUILTINS`] entry.
@@ -259,6 +273,19 @@ pub const CLICKHOUSE_BUILTINS: &[ClickHouseBuiltin] = &[
259273
arity: Arity::Exact(1),
260274
rewrite: RewriteKind::CountIfToSum,
261275
},
276+
// argMax(arg, val) -- "arg's value from the row where val is maximal".
277+
// No native DataFusion aggregate shape to rewrite to (issue #232), so
278+
// `lower_agg_intent` builds `AggIntent::Extension` from these directly.
279+
ClickHouseBuiltin {
280+
name: "argmax",
281+
arity: Arity::Exact(2),
282+
rewrite: RewriteKind::PassThrough,
283+
},
284+
ClickHouseBuiltin {
285+
name: "argmin",
286+
arity: Arity::Exact(2),
287+
rewrite: RewriteKind::PassThrough,
288+
},
262289
];
263290

264291
/// Aggregate function names DataFusion's own planner resolves out of the box

0 commit comments

Comments
 (0)