From 6062e448a48fbf3c4b69a0abfc860e19a430ad47 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 23 Aug 2026 20:54:52 -0400 Subject: [PATCH] feat(sql): give lagInFrame/leadInFrame their own WindowFuncKind variants (#267) Mapping ClickHouse's frame-respecting lagInFrame/leadInFrame onto the existing Lag/Lead variants would silently ignore the window frame clause, since ANSI LAG/LEAD (and DataFusion's native lag/lead) are defined to ignore it entirely. Add distinct WindowFuncKind::LagInFrame/LeadInFrame variants instead, so the frame clause is never conflated away. DataFusion's planner rejects lagInFrame/leadInFrame outright as unknown functions -- unlike lag/lead, they need a stub WindowUDF registered (mirroring the existing CLICKHOUSE_BUILTINS/CLICKHOUSE_SCALAR_BUILTINS mechanism), so this adds a CLICKHOUSE_WINDOW_BUILTINS catalog table alongside them. WindowFuncKind still has no frame representation, so these lower and behave like plain Lag/Lead today -- the tag is correct, the frame-respecting behavior isn't modeled yet (#231). Updates the two pinned SQL-corpus tests (bgp_analytics, bgp_jan2024_workload) to reflect the new coverage. Co-Authored-By: Claude Sonnet 5 --- crates/frontend-sql/src/sql/mod.rs | 83 ++++++++++++++- .../tests/bgp_analytics/bgp_analytics.rs | 59 ++++++----- .../bgp_jan2024_workload.rs | 15 ++- crates/frontend-sql/tests/sql_lowering.rs | 36 +++++++ crates/sql-function-catalog/src/lib.rs | 100 +++++++++++++++++- crates/types/src/pre_asap/query_expr.rs | 12 +++ 6 files changed, 272 insertions(+), 33 deletions(-) diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index bb2c708a..c4611ef3 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -26,7 +26,7 @@ use std::rc::Rc; use std::sync::Arc; -use datafusion::arrow::datatypes::DataType as ArrowDataType; +use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field}; use datafusion::catalog_common::MemorySchemaProvider; use datafusion::common::config::ConfigOptions; use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; @@ -36,9 +36,11 @@ 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::function::{PartitionEvaluatorArgs, WindowUDFFieldArgs}; use datafusion::logical_expr::{ - self, lit, AggregateUDF, Case, Distinct, Expr, JoinType, LogicalPlan, ScalarUDF, ScalarUDFImpl, - Signature, SimpleAggregateUDF, TypeSignature, Volatility, WindowFunctionDefinition, + self, lit, AggregateUDF, Case, Distinct, Expr, JoinType, LogicalPlan, PartitionEvaluator, + ScalarUDF, ScalarUDFImpl, Signature, SimpleAggregateUDF, TypeSignature, Volatility, + WindowFunctionDefinition, WindowUDF, WindowUDFImpl, }; use datafusion::optimizer::analyzer::function_rewrite::ApplyFunctionRewrites; use datafusion::optimizer::{AnalyzerRule, OptimizerConfig}; @@ -221,6 +223,16 @@ impl<'a> SqlLowerer<'a> { builtin.arity, )); } + // Register a stub `WindowUDF` for every catalog-listed ClickHouse-only + // *window* builtin — same reason as the two loops above, but with no + // rewrite step to follow: `lower_window_func_kind` already maps each + // name directly to its own `WindowFuncKind` variant (issue #267). + for builtin in asap_sql_function_catalog::CLICKHOUSE_WINDOW_BUILTINS { + ctx.register_udwf(clickhouse_window_builtin_stub_udwf( + builtin.name, + builtin.arity, + )); + } Ok(ctx) } @@ -1036,6 +1048,68 @@ impl ScalarUDFImpl for ClickHouseScalarBuiltinStub { } } +// ── ClickHouse window-builtin compatibility ───────────────────────────────── +// +// The window counterpart of the scalar mechanism above: a stub `WindowUDF` +// registered purely so DataFusion's planner accepts the call name during +// `SqlToRel` conversion. No rewrite step follows — `lower_window_func_kind` +// already maps each `asap_sql_function_catalog::CLICKHOUSE_WINDOW_BUILTINS` +// name directly to its own `WindowFuncKind` variant (issue #267). + +/// A stub `WindowUDF` for one `CLICKHOUSE_WINDOW_BUILTINS` entry, registered +/// purely so DataFusion's planner can resolve the function name inside an +/// `OVER (...)` clause. This front end only ever uses DataFusion for +/// planning/type-checking, never physical execution, so +/// `partition_evaluator` (which physical execution alone would call) is +/// unreachable in practice. +fn clickhouse_window_builtin_stub_udwf(name: &'static str, arity: Arity) -> WindowUDF { + WindowUDF::from(ClickHouseWindowBuiltinStub { + name, + signature: arity_to_signature(arity), + }) +} + +/// A stub `WindowUDFImpl` carrying only what DataFusion's planner needs: +/// name, arity-only [`Signature`], and a field type derived from the first +/// argument (matching `lag`/`lead`'s own "output type = input type" +/// behavior). `partition_evaluator` is left `unimplemented!()` — see +/// [`clickhouse_window_builtin_stub_udwf`]'s doc for why that is unreachable. +#[derive(Debug)] +struct ClickHouseWindowBuiltinStub { + name: &'static str, + signature: Signature, +} + +impl WindowUDFImpl for ClickHouseWindowBuiltinStub { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn name(&self) -> &str { + self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn field(&self, field_args: WindowUDFFieldArgs) -> datafusion::common::Result { + let dtype = field_args.get_input_type(0).unwrap_or(ArrowDataType::Null); + Ok(Field::new(field_args.name(), dtype, true)) + } + + fn partition_evaluator( + &self, + _partition_evaluator_args: PartitionEvaluatorArgs, + ) -> datafusion::common::Result> { + let name = self.name; + unimplemented!( + "{name} has no partition evaluator: this front end never runs DataFusion's \ + physical planner, only SqlToRel + the unoptimized LogicalPlan" + ) + } +} + /// 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 @@ -1641,6 +1715,9 @@ fn lower_window_func_kind(fun: &WindowFunctionDefinition) -> Result Ok(WindowFuncKind::DenseRank), "lag" => Ok(WindowFuncKind::Lag), "lead" => Ok(WindowFuncKind::Lead), + // ClickHouse: frame-respecting variants, not plain Lag/Lead (#267). + "laginframe" => Ok(WindowFuncKind::LagInFrame), + "leadinframe" => Ok(WindowFuncKind::LeadInFrame), "first_value" => Ok(WindowFuncKind::FirstValue), "last_value" => Ok(WindowFuncKind::LastValue), "nth_value" => Ok(WindowFuncKind::NthValue(None)), diff --git a/crates/frontend-sql/tests/bgp_analytics/bgp_analytics.rs b/crates/frontend-sql/tests/bgp_analytics/bgp_analytics.rs index 0a817480..ae27e31e 100644 --- a/crates/frontend-sql/tests/bgp_analytics/bgp_analytics.rs +++ b/crates/frontend-sql/tests/bgp_analytics/bgp_analytics.rs @@ -7,23 +7,25 @@ //! attribution, MOAS detection, prefix deaggregation, RIB visibility/churn. //! //! Unlike the DQC and netflow SQL corpora, this one does **not** pin full -//! coverage: only queries 1, 2, 6, 12 (plain `COUNT`/`GROUP BY`/`ORDER -//! BY`/`LIMIT`, plus `uniqExact` -- rewritten to `COUNT(DISTINCT ...)` by +//! coverage: only queries 1, 2, 6, 11, 12 (plain `COUNT`/`GROUP BY`/`ORDER +//! BY`/`LIMIT`, `uniqExact` -- rewritten to `COUNT(DISTINCT ...)` by //! DataFusion's own `Analyzer` before `lower_plan` runs, see -//! `UniqExactRewrite` in `sql/mod.rs`) lower end to end. The other 11 fail -//! for two distinct reasons -- distinct `DataFusionError` variants, not just -//! "some `SqlError::DataFusion`" -- and `EXPECTED` below pins each query to -//! its specific variant *and* a snippet of the error message, so a query -//! silently drifting from one failure mode to the other (e.g. a grammar gap -//! getting "fixed" into an unknown-function error, or vice versa) fails the -//! test even though the aggregate 4/9/2 split wouldn't otherwise move: -//! - 9 queries hit `DataFusionError::Plan` ("unknown function"): they use -//! ClickHouse-only builtins (`toIntervalMinute`, `lagInFrame`, -//! `isIPAddressInRange`, `arrayJoin`, `arrayFilter`, ...) that parse fine -//! under the ClickHouse dialect but have no DataFusion planner equivalent -//! registered. (`toStartOfInterval` itself is registered -- issue #230 -- -//! but query 7 nests an unregistered `toIntervalMinute(...)` call inside -//! it, so it still lands here, just one function name deeper.) +//! `UniqExactRewrite` in `sql/mod.rs` -- and `lagInFrame`, via a stub +//! `WindowUDF` and its own `WindowFuncKind` variant, issue #267) lower end +//! to end. The other 10 fail for two distinct reasons -- distinct +//! `DataFusionError` variants, not just "some `SqlError::DataFusion`" -- +//! and `EXPECTED` below pins each query to its specific variant *and* a +//! snippet of the error message, so a query silently drifting from one +//! failure mode to the other (e.g. a grammar gap getting "fixed" into an +//! unknown-function error, or vice versa) fails the test even though the +//! aggregate 5/8/2 split wouldn't otherwise move: +//! - 8 queries hit `DataFusionError::Plan` ("unknown function"): they use +//! ClickHouse-only builtins (`toIntervalMinute`, `isIPAddressInRange`, +//! `arrayJoin`, `arrayFilter`, ...) that parse fine under the ClickHouse +//! dialect but have no DataFusion planner equivalent registered. +//! (`toStartOfInterval` itself is registered -- issue #230 -- but query +//! 7 nests an unregistered `toIntervalMinute(...)` call inside it, so it +//! still lands here, just one function name deeper.) //! - 2 queries (14, 15) hit `DataFusionError::SQL` (a `ParserError`): they //! use ClickHouse grammar the vendored sqlparser doesn't implement at //! all -- a scalar/tuple `WITH AS ` binding, and the @@ -132,9 +134,11 @@ const EXPECTED: &[Expected] = &[ Expected::UnknownFunction("arrayfilter"), // 8 Expected::UnknownFunction("isipaddressinrange"), // 9 Expected::UnknownFunction("arrayfilter"), // 10 - Expected::UnknownFunction("laginframe"), // 11 - Expected::Lowered, // 12 - Expected::UnknownFunction("arrayjoin"), // 13 + // `lagInFrame` now has its own stub `WindowUDF` + `WindowFuncKind` + // variant (issue #267), so this lowers end to end. + Expected::Lowered, // 11 + Expected::Lowered, // 12 + Expected::UnknownFunction("arrayjoin"), // 13 Expected::UnsupportedGrammar("Expected: identifier, found: ("), // 14 Expected::UnsupportedGrammar("Expected: a list of columns in parentheses"), // 15 ]; @@ -192,19 +196,20 @@ async fn corpus_lowering_matches_the_pinned_per_query_outcome() { } eprintln!("bgp-analytics SQL corpus: {t:?}"); - // Coverage ratchet -- queries 1, 2, 6, 12 lower (6 and 12's `uniqExact` - // calls are rewritten to `COUNT(DISTINCT ...)` before `lower_plan` ever - // sees them, see `UniqExactRewrite` in `sql/mod.rs`); further - // ClickHouse-builtin UDF support or vendored-grammar fixes should move the - // affected query's entry in `EXPECTED` to `Lowered` (raising `t.lowered` - // here) rather than just this summary. + // Coverage ratchet -- queries 1, 2, 6, 11, 12 lower (6 and 12's + // `uniqExact` calls are rewritten to `COUNT(DISTINCT ...)` before + // `lower_plan` ever sees them, see `UniqExactRewrite` in `sql/mod.rs`; + // 11's `lagInFrame` resolves via a stub `WindowUDF`, issue #267); + // further ClickHouse-builtin UDF support or vendored-grammar fixes + // should move the affected query's entry in `EXPECTED` to `Lowered` + // (raising `t.lowered` here) rather than just this summary. assert_eq!( - t.lowered, 4, + t.lowered, 5, "BGP analytics SQL coverage changed -- update EXPECTED if support for \ a ClickHouse builtin or grammar gap was added: {t:?}" ); assert_eq!( - t.unknown_function, 9, + t.unknown_function, 8, "BGP analytics SQL coverage changed: {t:?}" ); assert_eq!( 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 0ab185d8..a7e4afa3 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 @@ -187,8 +187,19 @@ async fn corpus_lowering_matches_the_pinned_aggregate_tally() { // failure for all 3 corpus occurrences: every one has both arguments as // bare columns of the scanned table, so all 3 lower end to end with no // companion gap (unlike `splitByChar`'s array-indexing gap, #230). - expect(Category::Lowered, 148); - expect(Category::Plan, 37); + // + // `lagInFrame` support (issue #267 -- a stub `WindowUDF` plus its own + // `WindowFuncKind::LagInFrame` variant, not conflated with `Lag`) clears + // the "unknown function: laginframe" `Plan` failure for 3 of the 6 + // corpus occurrences (q024, q025, q171). The other 3 (q056, q114, q170) + // still land in `Plan`, but now for an unrelated, pre-existing reason + // each: q056/q114 also call `dateDiff`, q170's outer query also calls + // `toString` -- neither is a registered `CLICKHOUSE_SCALAR_BUILTINS` + // entry, so the query still fails at the first unknown-function name it + // hits, just no longer `laginframe`. Out of scope for #267, same as + // `splitByChar`'s array-indexing companion gap above. + expect(Category::Lowered, 151); + expect(Category::Plan, 34); expect(Category::Schema, 0); expect(Category::Parse, 0); // One query that used to fail at `uniqExact` (`Plan`) now clears that diff --git a/crates/frontend-sql/tests/sql_lowering.rs b/crates/frontend-sql/tests/sql_lowering.rs index 34cc62cb..dd9828cd 100644 --- a/crates/frontend-sql/tests/sql_lowering.rs +++ b/crates/frontend-sql/tests/sql_lowering.rs @@ -1554,3 +1554,39 @@ async fn arg_max_rejects_a_non_column_argument() { "got {err}" ); } + +// ── Issue #267: lagInFrame/leadInFrame get distinct WindowFuncKind variants, +// not conflated with ANSI Lag/Lead ────────────────────────────────────────── + +#[tokio::test] +async fn lag_in_frame_lowers_to_its_own_kind_not_lag() { + let qe = lower_clickhouse( + "SELECT service, lagInFrame(bytes) OVER (PARTITION BY service ORDER BY ts) \ + FROM metrics", + ) + .await; + let win = find_windowfunc(&qe).expect("expected a SQLWindowFunc node"); + let QueryExpr::SQLWindowFunc { func, args, .. } = win else { + unreachable!(); + }; + assert_eq!(*func, WindowFuncKind::LagInFrame); + assert_eq!( + args, + &vec![QueryExpr::Column(3)], + "lagInFrame(bytes) → arg col 3" + ); +} + +#[tokio::test] +async fn lead_in_frame_lowers_to_its_own_kind_not_lead() { + let qe = lower_clickhouse( + "SELECT service, leadInFrame(bytes) OVER (PARTITION BY service ORDER BY ts) \ + FROM metrics", + ) + .await; + let win = find_windowfunc(&qe).expect("expected a SQLWindowFunc node"); + let QueryExpr::SQLWindowFunc { func, .. } = win else { + unreachable!(); + }; + assert_eq!(*func, WindowFuncKind::LeadInFrame); +} diff --git a/crates/sql-function-catalog/src/lib.rs b/crates/sql-function-catalog/src/lib.rs index 82be1fea..39645717 100644 --- a/crates/sql-function-catalog/src/lib.rs +++ b/crates/sql-function-catalog/src/lib.rs @@ -16,7 +16,7 @@ //! outright (`reducer_col` in `asap-frontend-sql`) rather than trying to //! typecheck it. //! -//! Three tables, matching the three problems this replaces: +//! Four tables, matching the four problems this replaces: //! //! - [`NATIVE_FUNCTIONS`] -- aggregate names DataFusion's own planner already //! resolves (`sum`, `avg`, `approx_percentile_cont`, ...). [`lookup_native`] @@ -45,6 +45,12 @@ //! lowers generically (`asap-frontend-sql::sql::expr`'s //! `Expr::ScalarFunction` arm), so a stub `ScalarUDF` registered for the //! name is the whole fix (issue #230). +//! - [`CLICKHOUSE_WINDOW_BUILTINS`] -- ClickHouse-only *window* names +//! DataFusion doesn't know at all (`lagInFrame`, `leadInFrame`). No +//! [`RewriteKind`] here either: `asap-frontend-sql::sql:: +//! lower_window_func_kind` already maps each name directly to its own +//! `WindowFuncKind` variant, so a stub `WindowUDF` registered for the name +//! is the whole fix (issue #267). //! //! Generating these tables from a live introspectable source -- ClickHouse's //! `system.functions`, DataFusion's own in-process UDF/UDAF registry -- the @@ -467,6 +473,53 @@ pub fn lookup_clickhouse_scalar_builtin(name: &str) -> Option<&'static ClickHous CLICKHOUSE_SCALAR_BUILTINS.iter().find(|b| b.name == name) } +/// One [`CLICKHOUSE_WINDOW_BUILTINS`] entry -- just `{ name, arity }`, no +/// [`RewriteKind`]: unlike an aggregate call, `lower_window_func_kind` +/// (`asap-frontend-sql::sql`) already maps each of these names directly to +/// its own `WindowFuncKind` variant, so once DataFusion's planner accepts +/// the name at all -- via a stub `WindowUDF`, see +/// `asap-frontend-sql::sql::clickhouse_window_builtin_stub_udwf` -- no +/// rewrite step is needed either. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ClickHouseWindowBuiltin { + /// Lowercase function name, matching the name a stub `WindowUDF` is + /// registered under (DataFusion resolves a SQL call to it + /// case-insensitively, but reports it back lowercase). + pub name: &'static str, + pub arity: Arity, +} + +/// ClickHouse-only *window* builtin names DataFusion's planner has no native +/// equivalent for at all -- each needs a stub `WindowUDF` registered so the +/// planner accepts the call. `lagInFrame`/`leadInFrame` (issue #267) differ +/// from ANSI `LAG`/`LEAD` (which DataFusion's native `lag`/`lead` already +/// model) by respecting the window frame bounds instead of ignoring them -- +/// conflating the two under one `WindowFuncKind` would make the frame clause +/// silently meaningless depending on which name got you there, so each gets +/// its own catalog entry and its own `WindowFuncKind` variant +/// (`LagInFrame`/`LeadInFrame`). `WindowFuncKind` itself still has no frame +/// representation, so the frame-respecting behavior isn't modeled yet either +/// -- see issue #231. +pub const CLICKHOUSE_WINDOW_BUILTINS: &[ClickHouseWindowBuiltin] = &[ + // lagInFrame(x[, offset[, default]]) -- like LAG, but NULL/default past + // the frame boundary instead of reaching arbitrarily far back. + ClickHouseWindowBuiltin { + name: "laginframe", + arity: Arity::Range { min: 1, max: 3 }, + }, + // leadInFrame(x[, offset[, default]]) -- the LEAD counterpart. + ClickHouseWindowBuiltin { + name: "leadinframe", + arity: Arity::Range { min: 1, max: 3 }, + }, +]; + +/// Look up a ClickHouse-only window builtin by name (case-sensitive, see +/// [`lookup_native`]). +pub fn lookup_clickhouse_window_builtin(name: &str) -> Option<&'static ClickHouseWindowBuiltin> { + CLICKHOUSE_WINDOW_BUILTINS.iter().find(|b| b.name == name) +} + #[cfg(test)] mod tests { use super::*; @@ -589,4 +642,49 @@ mod tests { ); } } + + #[test] + fn clickhouse_window_builtin_lookup_finds_every_listed_name() { + for b in CLICKHOUSE_WINDOW_BUILTINS { + let found = + lookup_clickhouse_window_builtin(b.name).expect("listed name must be found"); + assert_eq!(found.name, b.name); + assert_eq!(found.arity, b.arity); + } + assert_eq!( + lookup_clickhouse_window_builtin("not_a_real_function"), + None + ); + } + + #[test] + fn clickhouse_window_builtin_names_are_lowercase() { + for b in CLICKHOUSE_WINDOW_BUILTINS { + assert_eq!(b.name, b.name.to_lowercase(), "not lowercase: {}", b.name); + } + } + + /// Window builtins live in their own namespace from the aggregate and + /// scalar tables (see `clickhouse_scalar_builtins_do_not_shadow_native_or_ + /// aggregate_names`'s doc for why this documents rather than enforces). + #[test] + fn clickhouse_window_builtins_do_not_shadow_other_names() { + for b in CLICKHOUSE_WINDOW_BUILTINS { + assert!( + lookup_native(b.name).is_none(), + "{} listed as both a native aggregate and a ClickHouse window builtin", + b.name + ); + assert!( + lookup_clickhouse_builtin(b.name).is_none(), + "{} listed as both a ClickHouse aggregate and window builtin", + b.name + ); + assert!( + lookup_clickhouse_scalar_builtin(b.name).is_none(), + "{} listed as both a ClickHouse scalar and window builtin", + b.name + ); + } + } } diff --git a/crates/types/src/pre_asap/query_expr.rs b/crates/types/src/pre_asap/query_expr.rs index 6d01ede2..b4f19141 100644 --- a/crates/types/src/pre_asap/query_expr.rs +++ b/crates/types/src/pre_asap/query_expr.rs @@ -321,6 +321,16 @@ pub enum WindowFuncKind { DenseRank, Lag, Lead, + /// ClickHouse `lagInFrame`/`leadInFrame`: unlike [`Lag`](Self::Lag)/[`Lead`](Self::Lead), + /// these respect the window frame bounds (NULL/default past the frame edge) + /// rather than reaching arbitrarily far back/forward. Kept as distinct + /// variants so the frame clause is never silently discarded by conflating + /// them with `Lag`/`Lead` (#267). `WindowFuncKind` still has no frame + /// representation, so today these lower and behave exactly like + /// `Lag`/`Lead` — the tag is correct, the frame-respecting behavior isn't + /// implemented yet. See #231 for modeling window frames properly. + LagInFrame, + LeadInFrame, FirstValue, LastValue, /// `NTH_VALUE(expr, n)` — `n` is resolved from the (literal) 2nd argument. @@ -1129,6 +1139,8 @@ impl QueryExpr { // Navigation funcs: arg type, nullable (boundary rows are NULL). WindowFuncKind::Lag | WindowFuncKind::Lead + | WindowFuncKind::LagInFrame + | WindowFuncKind::LeadInFrame | WindowFuncKind::FirstValue | WindowFuncKind::LastValue | WindowFuncKind::NthValue(_) => (arg_dtype(), true),