diff --git a/Cargo.lock b/Cargo.lock index 9776deea..7a77f876 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -336,6 +336,7 @@ dependencies = [ name = "asap-frontend-sql" version = "0.1.0" dependencies = [ + "asap-sql-function-catalog", "asap-types", "datafusion", "serde", @@ -352,6 +353,10 @@ dependencies = [ "asap-types", ] +[[package]] +name = "asap-sql-function-catalog" +version = "0.1.0" + [[package]] name = "asap-types" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 50fbf85d..907f3ac2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "crates/types", + "crates/sql-function-catalog", "crates/asap-aware-mapping", "crates/frontend-promql", "crates/frontend-sql", diff --git a/crates/frontend-sql/Cargo.toml b/crates/frontend-sql/Cargo.toml index ed059716..91bc4137 100644 --- a/crates/frontend-sql/Cargo.toml +++ b/crates/frontend-sql/Cargo.toml @@ -4,9 +4,12 @@ version = "0.1.0" edition = "2021" # SQL front end: L1 (parse + plan via DataFusion) → L2 relational, then the -# shared L2→L3 converter — both in asap-types. Pulls DataFusion only — never promql-parser. +# shared L2→L3 converter — both in asap-types. Pulls DataFusion, plus the +# independent SQL function-name catalog (asap-sql-function-catalog, issue +# #225) it consults when lowering an aggregate call — never promql-parser. [dependencies] asap-types = { path = "../types" } +asap-sql-function-catalog = { path = "../sql-function-catalog" } datafusion = "43" [dev-dependencies] diff --git a/crates/frontend-sql/src/lib.rs b/crates/frontend-sql/src/lib.rs index 7f961535..755e1e05 100644 --- a/crates/frontend-sql/src/lib.rs +++ b/crates/frontend-sql/src/lib.rs @@ -36,8 +36,10 @@ pub async fn lower_sql( /// Lower a single SQL query string under an explicit [`SqlDialect`]. /// /// `ClickhouseSQL` parses via sqlparser's vendored `ClickHouseDialect` -/// (array-lambda syntax, `arr[-1]` indexing) — it does not teach DataFusion's -/// planner any ClickHouse-only builtin functions, which still fail to plan. +/// (array-lambda syntax, `arr[-1]` indexing). It also teaches DataFusion's +/// planner the ClickHouse-only builtin functions listed in +/// `asap_sql_function_catalog::CLICKHOUSE_BUILTINS` (`uniqExact`, `countIf`) +/// — every other ClickHouse-only builtin still fails to plan. /// `ElasticSQL` has no vendored parser and always returns `UnsupportedDialect`. pub async fn lower_sql_dialect( query: &str, diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index 9fe602d5..cc1cebab 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -33,16 +33,18 @@ use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::common::{Column as DfColumn, DFSchema, ScalarValue as DfScalarValue}; use datafusion::datasource::MemTable; use datafusion::functions_aggregate::count::count_udaf; +use datafusion::functions_aggregate::sum::sum_udaf; use datafusion::logical_expr::expr::AggregateFunction; use datafusion::logical_expr::expr_rewriter::FunctionRewrite; use datafusion::logical_expr::{ - self, AggregateUDF, Distinct, Expr, JoinType, LogicalPlan, Signature, SimpleAggregateUDF, - Volatility, WindowFunctionDefinition, + self, lit, AggregateUDF, Case, Distinct, Expr, JoinType, LogicalPlan, Signature, + SimpleAggregateUDF, TypeSignature, Volatility, WindowFunctionDefinition, }; use datafusion::optimizer::analyzer::function_rewrite::ApplyFunctionRewrites; use datafusion::optimizer::{AnalyzerRule, OptimizerConfig}; use datafusion::prelude::{SessionConfig, SessionContext}; +use asap_sql_function_catalog::{AggSemantic, Arity, RewriteKind}; use asap_types::pre_asap::agg_intent::AggIntent; use asap_types::pre_asap::query_expr::{ GroupKeys, Predicate, ProjectItem, Reduction, SortKey, Source, @@ -116,9 +118,11 @@ impl<'a> SqlLowerer<'a> { /// Parse under a specific SQL dialect (e.g. `ClickhouseSQL`, which maps to /// sqlparser's vendored `ClickHouseDialect` — array-lambda syntax and /// `arr[-1]` indexing parse under it that don't parse generically). This - /// only changes *parsing*: ClickHouse-only builtin functions (`uniqExact`, - /// `countIf`, …) are still unknown to DataFusion's planner and still fail - /// there, and `ElasticSQL` has no vendored parser at all. + /// only changes *parsing*: a ClickHouse-only builtin function not listed + /// in `asap_sql_function_catalog::CLICKHOUSE_BUILTINS` (`uniqExact` and + /// `countIf` are; most of ClickHouse's builtin surface isn't yet) is + /// still unknown to DataFusion's planner and still fails there, and + /// `ElasticSQL` has no vendored parser at all. pub fn with_dialect(catalog: &'a SqlCatalog, dialect: SqlDialect) -> Self { Self { catalog, dialect } } @@ -135,20 +139,23 @@ impl<'a> SqlLowerer<'a> { /// Runs `ApplyFunctionRewrites` — the single `AnalyzerRule` DataFusion's /// own `Analyzer` uses internally to apply `FunctionRewrite`s, called /// directly rather than through `Analyzer::execute_and_check` — over the - /// raw parsed plan before lowering, carrying only `UniqExactRewrite`. - /// `ctx.sql(...).into_unoptimized_plan()` alone returns `SqlToRel`'s - /// output untouched, and a `FunctionRewrite` only ever runs as part of - /// this rule, so calling it directly is unavoidable to make the rewrite - /// fire. Its `analyze()` already does a full `transform_up_with_subqueries` - /// over the whole plan, so it needs no wrapping `Analyzer` at all — - /// deliberately not `Analyzer::execute_and_check` (whether with the - /// default 5-rule analyzer or an empty one carrying just this rewrite): - /// that method runs an unconditional post-check (`check_plan`, hardcoded, - /// not itself a rule) that isn't wanted here — e.g. it independently - /// rejects a multi-column `IN (subquery)` before `lower_in_subquery`'s - /// own arity check would. Going straight to `ApplyFunctionRewrites` - /// avoids that entirely: zero behavior change for every query that - /// doesn't call `uniqExact`. + /// raw parsed plan before lowering, carrying only + /// `ClickHouseBuiltinRewrite` (catalog-driven, see its own doc — it + /// covers every `asap_sql_function_catalog::CLICKHOUSE_BUILTINS` entry, + /// not just one). `ctx.sql(...).into_unoptimized_plan()` alone returns + /// `SqlToRel`'s output untouched, and a `FunctionRewrite` only ever runs + /// as part of this rule, so calling it directly is unavoidable to make + /// the rewrite fire. Its `analyze()` already does a full + /// `transform_up_with_subqueries` over the whole plan, so it needs no + /// wrapping `Analyzer` at all — deliberately not + /// `Analyzer::execute_and_check` (whether with the default 5-rule + /// analyzer or an empty one carrying just this rewrite): that method + /// runs an unconditional post-check (`check_plan`, hardcoded, not itself + /// a rule) that isn't wanted here — e.g. it independently rejects a + /// multi-column `IN (subquery)` before `lower_in_subquery`'s own arity + /// check would. Going straight to `ApplyFunctionRewrites` avoids that + /// entirely: zero behavior change for every query that doesn't call a + /// catalog-listed ClickHouse builtin. pub async fn lower( &self, sql: &str, @@ -157,7 +164,7 @@ impl<'a> SqlLowerer<'a> { let ctx = self.build_context()?; let df = ctx.sql(sql).await?; let plan = df.into_unoptimized_plan(); - let rewriter = ApplyFunctionRewrites::new(vec![Arc::new(UniqExactRewrite)]); + let rewriter = ApplyFunctionRewrites::new(vec![Arc::new(ClickHouseBuiltinRewrite)]); let plan = rewriter.analyze(plan, ctx.state().options())?; let _guard = AccuracyGuard::install(accuracy.clone()); self.lower_plan(&plan) @@ -193,11 +200,14 @@ impl<'a> SqlLowerer<'a> { let mem_table = MemTable::try_new(arrow_schema, vec![])?; ctx.register_table(name.as_str(), Arc::new(mem_table))?; } - // Register a stub `uniqexact` UDAF purely so DataFusion's planner can - // resolve the function name during parsing — `lower()` rewrites every - // call site to `COUNT(DISTINCT ...)` via `UniqExactRewrite` before - // `lower_plan` sees it. - ctx.register_udaf(uniq_exact_udaf()); + // Register a stub `AggregateUDF` for every catalog-listed + // ClickHouse-only builtin, purely so DataFusion's planner can + // resolve its name during parsing — `lower()` rewrites every call + // site to a native DataFusion aggregate via `ClickHouseBuiltinRewrite` + // before `lower_plan` sees it. + for builtin in asap_sql_function_catalog::CLICKHOUSE_BUILTINS { + ctx.register_udaf(clickhouse_builtin_stub_udaf(builtin.name, builtin.arity)); + } Ok(ctx) } @@ -867,41 +877,63 @@ impl<'a> SqlLowerer<'a> { } // ── ClickHouse-builtin compatibility, taught to DataFusion itself ────────────── - -/// A stub `uniqexact` UDAF, registered purely so DataFusion's planner can -/// resolve the function name during `SqlToRel` conversion (it errors on an -/// unknown function before a rewrite ever gets a chance to run). Every call -/// site is replaced by `UniqExactRewrite` — via the `Analyzer` `lower()` runs -/// after parsing — before physical planning could ever ask this UDAF for an -/// `Accumulator`, so `accumulator` is unreachable. -fn uniq_exact_udaf() -> AggregateUDF { +// +// Generalized over `asap_sql_function_catalog::CLICKHOUSE_BUILTINS` (issue +// #225): adding support for one more ClickHouse-only builtin DataFusion +// doesn't know at all is a catalog data entry (name, arity, `RewriteKind`) +// plus, only if its rewrite target is a genuinely new shape, one match arm +// in `ClickHouseBuiltinRewrite::rewrite` below — never a new stub-UDAF +// constructor or a new `FunctionRewrite`-implementing type. `uniqExact` +// (issue #221) and `countIf` both go through this one mechanism. + +/// A stub `AggregateUDF` for one `CLICKHOUSE_BUILTINS` entry, registered +/// purely so DataFusion's planner can resolve the function name during +/// `SqlToRel` conversion (it errors on an unknown function before a rewrite +/// ever gets a chance to run). Every call site is replaced by +/// `ClickHouseBuiltinRewrite` — via the `Analyzer` `lower()` runs after +/// parsing — before physical planning could ever ask this UDAF for an +/// `Accumulator`, so `accumulator` is unreachable for every catalog entry. +fn clickhouse_builtin_stub_udaf(name: &'static str, arity: Arity) -> AggregateUDF { AggregateUDF::from(SimpleAggregateUDF::new_with_signature( - "uniqexact", - Signature::any(1, Volatility::Immutable), + name, + arity_to_signature(arity), ArrowDataType::Int64, - Arc::new(|_| { + Arc::new(move |_| { // ponytail: dead code by construction (see doc comment above) — - // a real accumulator would just reimplement COUNT(DISTINCT). + // a real accumulator would just reimplement whatever native + // shape `ClickHouseBuiltinRewrite` rewrites this call to. unimplemented!( - "uniqexact has no accumulator: every call site is rewritten to \ - COUNT(DISTINCT ...) before physical planning" + "{name} has no accumulator: every call site is rewritten to a native \ + DataFusion aggregate before physical planning" ) }), vec![], )) } -/// Rewrites `uniqexact(x)` to DataFusion's own `count(x) DISTINCT` — so -/// ClickHouse's exact-distinct-count builtin becomes an ordinary DataFusion -/// aggregate before the plan ever reaches `lower_agg_intent`, which already -/// maps `count` + `distinct` to `AggIntent::Cardinality` and needs no -/// ClickHouse-specific name of its own. +/// A catalog [`Arity`] as the DataFusion `Signature` a stub UDAF is +/// registered with. +fn arity_to_signature(arity: Arity) -> Signature { + match arity { + Arity::Exact(n) => Signature::any(n, Volatility::Immutable), + Arity::Range { min, max } => Signature::one_of( + (min..=max).map(TypeSignature::Any).collect(), + Volatility::Immutable, + ), + } +} + +/// Rewrites every `asap_sql_function_catalog::CLICKHOUSE_BUILTINS` call to +/// the native DataFusion aggregate shape its entry's `RewriteKind` names — +/// so a ClickHouse-only builtin DataFusion doesn't know at all becomes an +/// ordinary DataFusion aggregate before the plan ever reaches +/// `lower_agg_intent`, which needs no ClickHouse-specific name of its own. #[derive(Debug)] -struct UniqExactRewrite; +struct ClickHouseBuiltinRewrite; -impl FunctionRewrite for UniqExactRewrite { +impl FunctionRewrite for ClickHouseBuiltinRewrite { fn name(&self) -> &str { - "uniqexact -> count distinct" + "clickhouse builtin -> native DataFusion aggregate" } fn rewrite( @@ -910,21 +942,49 @@ impl FunctionRewrite for UniqExactRewrite { _schema: &DFSchema, _config: &ConfigOptions, ) -> datafusion::common::Result> { - match expr { - Expr::AggregateFunction(f) if f.func.name() == "uniqexact" => { - Ok(Transformed::yes(Expr::AggregateFunction( - AggregateFunction::new_udf( - count_udaf(), - f.args, - true, // DISTINCT - f.filter, - f.order_by, - f.null_treatment, - ), - ))) + let Expr::AggregateFunction(f) = expr else { + return Ok(Transformed::no(expr)); + }; + let Some(builtin) = asap_sql_function_catalog::lookup_clickhouse_builtin(f.func.name()) + else { + return Ok(Transformed::no(Expr::AggregateFunction(f))); + }; + let rewritten = match builtin.rewrite { + // `f(args...)` -> `count(args...) DISTINCT` — `lower_agg_intent` + // already maps `count` + `DISTINCT` to `AggIntent::Cardinality`. + RewriteKind::CountDistinct => AggregateFunction::new_udf( + count_udaf(), + f.args, + true, + f.filter, + f.order_by, + f.null_treatment, + ), + // `f(cond)` -> `sum(CASE WHEN cond THEN 1 ELSE 0 END)` — see + // `RewriteKind::CountIfToSum`'s doc for why a plain `count(...) + // FILTER (WHERE cond)` doesn't work here (`AggIntent::Count` + // never consults its argument). + RewriteKind::CountIfToSum => { + let cond = f.args.into_iter().next().expect( + "countif's stub signature fixes its arity at 1 -- the planner \ + already rejected any other argument count before this rewrite runs", + ); + let indicator = Expr::Case(Case::new( + None, + vec![(Box::new(cond), Box::new(lit(1i64)))], + Some(Box::new(lit(0i64))), + )); + AggregateFunction::new_udf( + sum_udaf(), + vec![indicator], + false, + f.filter, + f.order_by, + f.null_treatment, + ) } - other => Ok(Transformed::no(other)), - } + }; + Ok(Transformed::yes(Expr::AggregateFunction(rewritten))) } } @@ -932,20 +992,28 @@ impl FunctionRewrite for UniqExactRewrite { /// Map a DataFusion aggregate expression directly to the canonical /// [`AggIntent`] — issue #179's "dedicated function → canonical -/// intent directly" front-end construction, no `AggFunc` intermediate. -/// `resolve_root` resolves `col` to a positional `ColumnId`; the output name -/// (DataFusion's own, e.g. `"sum(metrics.bytes)"`) is threaded separately as -/// `Aggregate.output_names`, not carried here. +/// intent directly" front-end construction, no `AggFunc` intermediate. The +/// name → semantic mapping itself lives in `asap_sql_function_catalog` +/// (issue #225) as flat data (`NATIVE_FUNCTIONS`); what stays here is +/// call-site logic that isn't a function of the name alone — the DISTINCT +/// modifier rule, the "reducer argument must be a bare column" rule +/// (`reducer_col`), φ extraction from a literal argument, and the ambient +/// `AccuracyTarget`. `resolve_root` resolves `col` to a positional +/// `ColumnId`; the output name (DataFusion's own, e.g. +/// `"sum(metrics.bytes)"`) is threaded separately as `Aggregate.output_names`, +/// not carried here. fn lower_agg_intent(expr: &Expr) -> Result, LoweringError> { match expr { Expr::Alias(a) => lower_agg_intent(&a.expr), Expr::AggregateFunction(agg_fn) => { let name = agg_fn.func.name().to_lowercase(); + let semantic = asap_sql_function_catalog::lookup_native(&name) + .ok_or_else(|| LoweringError::UnsupportedAggregate(name.clone()))?; // The canonical intent algebra has no DISTINCT modifier for the // value reducers; only // COUNT(DISTINCT) maps (to Cardinality). Reject DISTINCT elsewhere // rather than silently lowering `SUM(DISTINCT x)` as `SUM(x)`. - if agg_fn.distinct && name != "count" { + if agg_fn.distinct && !matches!(semantic, AggSemantic::Count) { return Err(LoweringError::UnsupportedAggregate(format!( "DISTINCT {name}" ))); @@ -959,61 +1027,52 @@ fn lower_agg_intent(expr: &Expr) -> Result, LoweringError> let col = |args: &[Expr]| -> Result, LoweringError> { reducer_col(&name, args).map(Some) }; - Ok(match name.as_str() { - "count" if agg_fn.distinct => AggIntent::Cardinality { + Ok(match semantic { + AggSemantic::Count if agg_fn.distinct => AggIntent::Cardinality { col: col(&agg_fn.args)?, accuracy: current_accuracy(), }, - "count" => AggIntent::Count { + AggSemantic::Count => AggIntent::Count { accuracy: current_accuracy(), }, - "sum" => AggIntent::Sum { + AggSemantic::Sum => AggIntent::Sum { col: col(&agg_fn.args)?, }, - "min" => AggIntent::Min { + AggSemantic::Min => AggIntent::Min { col: col(&agg_fn.args)?, }, - "max" => AggIntent::Max { + AggSemantic::Max => AggIntent::Max { col: col(&agg_fn.args)?, }, - "avg" | "mean" => AggIntent::Avg { + AggSemantic::Avg => AggIntent::Avg { col: col(&agg_fn.args)?, }, - "stddev" | "stddev_samp" => AggIntent::StdDev { + AggSemantic::StdDev { population } => AggIntent::StdDev { col: col(&agg_fn.args)?, - population: false, + population, }, - "stddev_pop" => AggIntent::StdDev { + AggSemantic::Variance { population } => AggIntent::Variance { col: col(&agg_fn.args)?, - population: true, - }, - "var" | "variance" | "var_samp" => AggIntent::Variance { - col: col(&agg_fn.args)?, - population: false, - }, - "var_pop" => AggIntent::Variance { - col: col(&agg_fn.args)?, - population: true, - }, - "approx_percentile_cont" | "percentile_cont" => AggIntent::Quantile { - col: col(&agg_fn.args)?, - q: extract_percentile_q(&agg_fn.args)?, - accuracy: current_accuracy(), + population, }, - // `median(c)` is the φ=0.5 quantile. As with `approx_distinct` and - // `approx_percentile_cont`, the `approx_` prefix does not force an - // approximation: the sketch-vs-exact choice is the AccuracyTarget's - // (see `plan::boundary`), so both spellings share one intent (#111). - "median" | "approx_median" => AggIntent::Quantile { + // `fixed_q = Some(0.5)` is `median`/`approx_median`. As with + // `approx_distinct` and `approx_percentile_cont`, the + // `approx_` prefix does not force an approximation: the + // sketch-vs-exact choice is the AccuracyTarget's (see + // `plan::boundary`), so both spellings share one intent + // (#111). + AggSemantic::Quantile { fixed_q } => AggIntent::Quantile { col: col(&agg_fn.args)?, - q: 0.5, + q: match fixed_q { + Some(q) => q, + None => extract_percentile_q(&agg_fn.args)?, + }, accuracy: current_accuracy(), }, - "approx_distinct" => AggIntent::Cardinality { + AggSemantic::Cardinality => AggIntent::Cardinality { col: col(&agg_fn.args)?, accuracy: current_accuracy(), }, - _ => return Err(LoweringError::UnsupportedAggregate(name)), }) } _ => Err(LoweringError::UnsupportedAggregate(format!("{expr:?}"))), diff --git a/crates/frontend-sql/tests/bgp_jan2024_workload/bgp_jan2024_workload.rs b/crates/frontend-sql/tests/bgp_jan2024_workload/bgp_jan2024_workload.rs index 505f14c9..c3e73fe9 100644 --- a/crates/frontend-sql/tests/bgp_jan2024_workload/bgp_jan2024_workload.rs +++ b/crates/frontend-sql/tests/bgp_jan2024_workload/bgp_jan2024_workload.rs @@ -160,15 +160,20 @@ async fn corpus_lowering_matches_the_pinned_aggregate_tally() { gap was added/removed: {tally:?}" ); }; - expect(Category::Lowered, 85); - expect(Category::Plan, 105); + // `countIf` support (issue #225 -- `ClickHouseBuiltinRewrite`, catalog- + // driven the same way `uniqExact` (#221) already was) clears the + // "unknown function: countif" `Plan` failure for 13 queries: 12 now + // lower end to end, and a 13th plans far enough to hit a second, + // pre-existing gap (map/array index access, `NotImplemented`). + expect(Category::Lowered, 97); + expect(Category::Plan, 92); expect(Category::Schema, 0); expect(Category::Parse, 0); // One query that used to fail at `uniqExact` (`Plan`) now clears that - // hurdle -- `UniqExactRewrite` in `sql/mod.rs` rewrites it to + // hurdle -- `ClickHouseBuiltinRewrite` in `sql/mod.rs` rewrites it to // `COUNT(DISTINCT ...)` before `lower_plan` runs -- and plans far enough // to hit a second, pre-existing gap: map/array index access. - expect(Category::NotImplemented, 4); + expect(Category::NotImplemented, 5); expect(Category::UnsupportedFeature, 6); expect(Category::Other, 0); } diff --git a/crates/frontend-sql/tests/sql_lowering.rs b/crates/frontend-sql/tests/sql_lowering.rs index 4f0dd1dd..a9ae6521 100644 --- a/crates/frontend-sql/tests/sql_lowering.rs +++ b/crates/frontend-sql/tests/sql_lowering.rs @@ -5,13 +5,14 @@ //! the shared `resolve_root` produces the positional, resolved canonical //! tree (the same resolver the PromQL path uses). -use asap_frontend_sql::{lower_sql, SqlCatalog}; +use asap_frontend_sql::{lower_sql, lower_sql_dialect, SqlCatalog}; use asap_types::pre_asap::schema::{Column, DataType, Schema}; use asap_types::pre_asap::{ AggIntent, CompareOp, GroupKeys, JoinKind, QueryExpr, Reduction, ScalarValue, Source, WindowFuncKind, }; use asap_types::types::AccuracyTarget; +use asap_types::workload::SqlDialect; fn col(name: &str, dtype: DataType) -> Column { Column::new(name, dtype, false) @@ -1392,3 +1393,76 @@ async fn array_agg_is_deliberately_rejected() { "expected a clean UnsupportedAggregate, got {err}" ); } + +// ── Issue #225: catalog-driven ClickHouse builtins (countIf, generalizing +// uniqExact from #221) ─────────────────────────────────────────────────── + +async fn lower_clickhouse(sql: &str) -> QueryExpr { + lower_sql_dialect( + sql, + &catalog(), + SqlDialect::ClickhouseSQL, + AccuracyTarget::Exact, + ) + .await + .unwrap_or_else(|e| panic!("lower failed for {sql:?}: {e}")) +} + +#[tokio::test] +async fn count_if_lowers_to_a_sum_over_a_derived_indicator_column() { + // ClickHouse's `countIf(cond)` has no DataFusion equivalent at all, so it + // goes through the same stub-UDAF + catalog-driven `FunctionRewrite` + // mechanism `uniqExact` (#221) does — rewritten, before `lower_agg_intent` + // ever runs, to `sum(CASE WHEN cond THEN 1 ELSE 0 END)`. Not a plain + // `count(...) FILTER (WHERE cond)`: `AggIntent::Count` never consults its + // argument (always a row count), so the filter would be silently dropped; + // summing a 0/1 indicator keeps `cond` observable through the ordinary + // `Sum` path instead. + let qe = lower_clickhouse("SELECT countIf(bytes > 100) AS big FROM metrics").await; + let (by, measures) = find_aggregate(&qe).expect("expected an Aggregate"); + assert!(by.is_empty()); + assert!( + matches!(measures.as_slice(), [AggIntent::Sum { col: Some(_) }]), + "expected a Sum bound to the derived indicator column, got {measures:?}" + ); + let (_, materialized) = reducer_input_names(&qe); + assert!( + materialized, + "the indicator expression must be materialized in a Project beneath the Aggregate" + ); +} + +#[tokio::test] +async fn two_count_ifs_with_different_conditions_stay_distinct_reducers() { + // The corpus pattern (`countIf(operation = 'A'), countIf(operation = 'W')` + // in one GROUP BY) needs each call's own condition to survive as its own + // derived column, not collapse onto a shared one. + let qe = lower_clickhouse( + "SELECT service, countIf(bytes > 100) AS big, countIf(bytes <= 100) AS small \ + FROM metrics GROUP BY service", + ) + .await; + let (by, measures) = find_aggregate(&qe).expect("expected an Aggregate"); + assert_eq!(*by, GroupKeys::by(vec![0])); + assert!( + matches!( + measures.as_slice(), + [ + AggIntent::Sum { col: Some(a) }, + AggIntent::Sum { col: Some(b) } + ] if a != b + ), + "expected two distinct Sum reducers, got {measures:?}" + ); +} + +#[tokio::test] +async fn count_if_composes_with_group_by() { + let qe = lower_clickhouse( + "SELECT service, countIf(bytes > 100) AS big FROM metrics GROUP BY service", + ) + .await; + let (by, measures) = find_aggregate(&qe).expect("expected an Aggregate"); + assert_eq!(*by, GroupKeys::by(vec![0])); + assert!(matches!(measures.as_slice(), [AggIntent::Sum { .. }])); +} diff --git a/crates/sql-function-catalog/Cargo.toml b/crates/sql-function-catalog/Cargo.toml new file mode 100644 index 00000000..6abf84b9 --- /dev/null +++ b/crates/sql-function-catalog/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "asap-sql-function-catalog" +version = "0.1.0" +edition = "2021" + +# The (SQL function name, arity) -> canonical AggIntent-shaped semantic +# catalog `asap-frontend-sql` consults when lowering an aggregate call, plus +# the ClickHouse-only-builtin -> native-DataFusion-shape rewrite table +# (issue #225). No internal deps -- like `asap-types`, this is meant to be +# consultable on its own (a future extraction tool populating it, a test +# asserting a name is covered, ...) without pulling in DataFusion or the +# front end that uses it. `asap-frontend-sql` depends on this crate, never +# the other way around. +[dependencies] diff --git a/crates/sql-function-catalog/src/lib.rs b/crates/sql-function-catalog/src/lib.rs new file mode 100644 index 00000000..c60cc52d --- /dev/null +++ b/crates/sql-function-catalog/src/lib.rs @@ -0,0 +1,319 @@ +//! SQL aggregate-function name catalog: independent, hand-maintained data +//! mapping a resolved function name to (arity, canonical semantic) -- +//! deliberately pulled out of `asap-frontend-sql::sql::lower_agg_intent` +//! (issue #225) so the mapping is its own small, independently-testable +//! artifact rather than buried in a hand-written `match`. +//! +//! Design lifted from [`tobilg/polyglot`](https://github.com/tobilg/polyglot)'s +//! `polyglot-sql-function-catalogs` crate: per-dialect function data +//! (existence, arity/overloads), deliberately *not* depending on whatever +//! consumes it (there, the SQL-parsing crate; here, `asap-frontend-sql` and +//! DataFusion) -- exposed to the consumer instead. Scope is equally +//! deliberately shallow, matching that project's own choice: no argument or +//! return *type* modeling, just existence + arity + which canonical +//! semantic a call maps to. `AggIntent` construction itself doesn't need +//! more than that either -- it already rejects a non-column argument +//! outright (`reducer_col` in `asap-frontend-sql`) rather than trying to +//! typecheck it. +//! +//! Two tables, matching the two problems this replaces: +//! +//! - [`NATIVE_FUNCTIONS`] -- names DataFusion's own planner already resolves +//! (`sum`, `avg`, `approx_percentile_cont`, ...). [`lookup_native`] maps +//! one to the [`AggSemantic`] `lower_agg_intent` builds an `AggIntent` +//! from. The DISTINCT-modifier rule ("`COUNT DISTINCT` alone maps, to +//! `Cardinality`; reject DISTINCT elsewhere") and the "reducer argument +//! must be a bare column" rule are call-site logic, not per-function data, +//! and stay in `asap-frontend-sql`. +//! - [`CLICKHOUSE_BUILTINS`] -- ClickHouse-only names DataFusion doesn't +//! know at all (`uniqExact`, `countIf`). Each entry additionally carries a +//! [`RewriteKind`]: the native DataFusion aggregate shape the call +//! rewrites to before `lower_agg_intent` (or DataFusion's own physical +//! planner) ever has to understand the ClickHouse name itself. This is +//! what generalizes `uniqExact`'s old bespoke `UniqExactRewrite` + +//! `uniq_exact_udaf` pair (issue #221): a new builtin that rewrites to an +//! already-handled shape is a new entry in this table, not a new +//! `FunctionRewrite` impl and a new stub-`AggregateUDF` constructor. +//! +//! Generating these tables from a live introspectable source -- ClickHouse's +//! `system.functions`, DataFusion's own in-process UDF/UDAF registry -- the +//! way `polyglot`'s `tools/*/extract_functions.py` do, is a deliberately +//! deferred follow-up (issue #225, item 3: it needs a decision on where such +//! extraction tooling would actually run, e.g. whether a live ClickHouse +//! instance in CI/dev is a given). Until then these are ordinary hand-edited +//! Rust consts. + +/// How many arguments a catalog entry's function accepts. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Arity { + /// Exactly `n` arguments. + Exact(usize), + /// Between `min` and `max` arguments, inclusive. + Range { min: usize, max: usize }, +} + +/// The canonical semantic a native function name maps to -- the +/// classification [`lower_agg_intent`](../asap_frontend_sql/index.html) +/// switches on to build the real `AggIntent`. +/// +/// Deliberately *not* `asap_types::pre_asap::agg_intent::AggIntent` itself: +/// most `AggIntent` variants carry call-site-only state that isn't a +/// function of the name alone -- the ambient `AccuracyTarget` (thread-local, +/// not catalog data), φ pulled from a call's literal 2nd argument, whether +/// `DISTINCT` was written. This enum is only the per-name discriminant those +/// call sites key off; building the actual `AggIntent` is +/// `asap-frontend-sql`'s job. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum AggSemantic { + /// `COUNT(*)` / `COUNT(x)` -- ignores its argument (always a row count). + /// `COUNT(DISTINCT x)` is the one exception the call site special-cases + /// into `Cardinality` instead; that combination is not its own catalog + /// entry (it is the same name, `count`, with a modifier). + Count, + Sum, + Min, + Max, + Avg, + /// Sample stddev unless `population`. + StdDev { + population: bool, + }, + /// Sample variance unless `population`. + Variance { + population: bool, + }, + /// φ-quantile. `fixed_q = Some(0.5)` for a name that always means the + /// median (`median`, `approx_median`); `None` for a name whose φ is a + /// literal argument the call site must read (`approx_percentile_cont`, + /// `percentile_cont`). + Quantile { + fixed_q: Option, + }, + /// Approximate/exact distinct-value count (`approx_distinct`, and + /// `COUNT DISTINCT` via the call-site modifier above). + Cardinality, +} + +/// One [`NATIVE_FUNCTIONS`] entry. +#[derive(Debug, Clone, Copy)] +pub struct NativeFunction { + /// Lowercase function name, as DataFusion's `Expr::AggregateFunction` + /// reports it (`agg_fn.func.name()`). + pub name: &'static str, + pub arity: Arity, + pub semantic: AggSemantic, +} + +/// Every aggregate function name `SqlDialect::DataFusionSQL` resolves out of +/// the box, mapped to its canonical semantic. A name appears more than once +/// when DataFusion (or this front end) accepts more than one spelling for +/// the same semantic (`avg`/`mean`, `stddev`/`stddev_samp`, ...). +pub const NATIVE_FUNCTIONS: &[NativeFunction] = &[ + NativeFunction { + name: "count", + arity: Arity::Range { min: 0, max: 1 }, + semantic: AggSemantic::Count, + }, + NativeFunction { + name: "sum", + arity: Arity::Exact(1), + semantic: AggSemantic::Sum, + }, + NativeFunction { + name: "min", + arity: Arity::Exact(1), + semantic: AggSemantic::Min, + }, + NativeFunction { + name: "max", + arity: Arity::Exact(1), + semantic: AggSemantic::Max, + }, + NativeFunction { + name: "avg", + arity: Arity::Exact(1), + semantic: AggSemantic::Avg, + }, + NativeFunction { + name: "mean", + arity: Arity::Exact(1), + semantic: AggSemantic::Avg, + }, + NativeFunction { + name: "stddev", + arity: Arity::Exact(1), + semantic: AggSemantic::StdDev { population: false }, + }, + NativeFunction { + name: "stddev_samp", + arity: Arity::Exact(1), + semantic: AggSemantic::StdDev { population: false }, + }, + NativeFunction { + name: "stddev_pop", + arity: Arity::Exact(1), + semantic: AggSemantic::StdDev { population: true }, + }, + NativeFunction { + name: "var", + arity: Arity::Exact(1), + semantic: AggSemantic::Variance { population: false }, + }, + NativeFunction { + name: "variance", + arity: Arity::Exact(1), + semantic: AggSemantic::Variance { population: false }, + }, + NativeFunction { + name: "var_samp", + arity: Arity::Exact(1), + semantic: AggSemantic::Variance { population: false }, + }, + NativeFunction { + name: "var_pop", + arity: Arity::Exact(1), + semantic: AggSemantic::Variance { population: true }, + }, + NativeFunction { + name: "approx_percentile_cont", + arity: Arity::Exact(2), + semantic: AggSemantic::Quantile { fixed_q: None }, + }, + NativeFunction { + name: "percentile_cont", + arity: Arity::Exact(2), + semantic: AggSemantic::Quantile { fixed_q: None }, + }, + NativeFunction { + name: "median", + arity: Arity::Exact(1), + semantic: AggSemantic::Quantile { fixed_q: Some(0.5) }, + }, + NativeFunction { + name: "approx_median", + arity: Arity::Exact(1), + semantic: AggSemantic::Quantile { fixed_q: Some(0.5) }, + }, + NativeFunction { + name: "approx_distinct", + arity: Arity::Exact(1), + semantic: AggSemantic::Cardinality, + }, +]; + +/// The native DataFusion aggregate shape a [`ClickHouseBuiltin`] call +/// rewrites to, before `lower_agg_intent` (or DataFusion's own physical +/// planner) ever has to know the ClickHouse name existed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RewriteKind { + /// `f(args...)` -> `count(args...) DISTINCT` -- ClickHouse's exact + /// distinct-count family. `lower_agg_intent` already maps + /// `count` + `DISTINCT` to `AggIntent::Cardinality`, so nothing further + /// is needed once the call wears DataFusion's own name. + CountDistinct, + /// `f(cond)` -> `sum(CASE WHEN cond THEN 1 ELSE 0 END)` -- ClickHouse's + /// conditional-count family. Not a plain `count(...) FILTER (WHERE + /// cond)`: `AggIntent::Count` never consults its argument (it always + /// means "row count"), so a per-call *filtered* count needs a shape + /// whose value actually depends on `cond` to survive `lower_agg_intent` + /// unchanged. Summing a 0/1 indicator does, and lands on the existing + /// `Sum` path -- including the general non-column-argument + /// materialization `asap-frontend-sql`'s `lower_aggregate` already does + /// for any reducer over an expression (issue #110) -- so no new + /// `AggIntent` variant or lowering path is needed either. + CountIfToSum, +} + +/// One [`CLICKHOUSE_BUILTINS`] entry. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ClickHouseBuiltin { + /// Lowercase function name, matching the name a stub `AggregateUDF` is + /// registered under (DataFusion resolves a SQL call to it + /// case-insensitively, but reports it back lowercase). + pub name: &'static str, + pub arity: Arity, + pub rewrite: RewriteKind, +} + +/// ClickHouse-only builtin aggregate names DataFusion's planner has no +/// native equivalent for at all -- each needs a stub `AggregateUDF` +/// registered (so the planner accepts the call) and a +/// [`RewriteKind`]-driven rewrite (so it becomes something DataFusion +/// natively understands before physical planning). See the module doc and +/// `asap-frontend-sql::sql::ClickHouseBuiltinRewrite`. +pub const CLICKHOUSE_BUILTINS: &[ClickHouseBuiltin] = &[ + ClickHouseBuiltin { + name: "uniqexact", + arity: Arity::Exact(1), + rewrite: RewriteKind::CountDistinct, + }, + ClickHouseBuiltin { + name: "countif", + arity: Arity::Exact(1), + rewrite: RewriteKind::CountIfToSum, + }, +]; + +/// Look up a native function name's canonical semantic (case-sensitive -- +/// callers normalize case first, as `asap-frontend-sql` already does via +/// `.to_lowercase()`). +pub fn lookup_native(name: &str) -> Option { + NATIVE_FUNCTIONS + .iter() + .find(|f| f.name == name) + .map(|f| f.semantic) +} + +/// Look up a ClickHouse-only builtin by name (case-sensitive, see +/// [`lookup_native`]). +pub fn lookup_clickhouse_builtin(name: &str) -> Option<&'static ClickHouseBuiltin> { + CLICKHOUSE_BUILTINS.iter().find(|b| b.name == name) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn native_lookup_finds_every_listed_name() { + for f in NATIVE_FUNCTIONS { + assert_eq!(lookup_native(f.name), Some(f.semantic), "name: {}", f.name); + } + assert_eq!(lookup_native("not_a_real_function"), None); + } + + #[test] + fn clickhouse_builtin_lookup_finds_every_listed_name() { + for b in CLICKHOUSE_BUILTINS { + let found = lookup_clickhouse_builtin(b.name).expect("listed name must be found"); + assert_eq!(found.name, b.name); + assert_eq!(found.rewrite, b.rewrite); + } + assert_eq!(lookup_clickhouse_builtin("not_a_real_function"), None); + } + + #[test] + fn native_and_clickhouse_names_are_lowercase() { + for f in NATIVE_FUNCTIONS { + assert_eq!(f.name, f.name.to_lowercase(), "not lowercase: {}", f.name); + } + for b in CLICKHOUSE_BUILTINS { + assert_eq!(b.name, b.name.to_lowercase(), "not lowercase: {}", b.name); + } + } + + /// The two tables are disjoint -- a ClickHouse-only name has no native + /// DataFusion equivalent by construction (that's the whole reason it + /// needs a stub + rewrite), so it should never also appear as a name + /// DataFusion already resolves. + #[test] + fn clickhouse_builtins_do_not_shadow_native_names() { + for b in CLICKHOUSE_BUILTINS { + assert!( + lookup_native(b.name).is_none(), + "{} listed as both native and ClickHouse-only", + b.name + ); + } + } +}