Skip to content

Plan nested LAG and LEAD for Snowflake conditional events #70

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
osipovartem merged 1 commit into from
Sep 6, 2026
Merged

Plan nested LAG and LEAD for Snowflake conditional events #70

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
207 changes: 179 additions & 28 deletions datafusion/sql/src/select.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ use crate::utils::{
};

use arrow::datatypes::DataType;
use datafusion_common::config::Dialect as SqlDialect;
use datafusion_common::error::DataFusionErrorBuilder;
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, not_impl_err, plan_err};
Expand All @@ -39,7 +40,8 @@ use datafusion_expr::builder::get_struct_unnested_columns;
use datafusion_expr::expr::Unnest as UnnestExpr;
use datafusion_expr::expr::{PlannedReplaceSelectItem, WildcardOptions};
use datafusion_expr::expr_rewriter::{
normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_sorts,
NamePreserver, normalize_col, normalize_col_with_schemas_and_ambiguity_check,
normalize_sorts,
};
use datafusion_expr::select_expr::SelectExpr;
use datafusion_expr::utils::{
Expand Down Expand Up @@ -87,6 +89,71 @@ struct RewrittenUnnestExprGroups {
expr_groups: Vec<Vec<Expr>>,
}

fn contains_nested_window_function(expr: &Expr) -> Result<bool> {
let mut found = false;
expr.apply_children(|child| {
child.apply(|nested| {
if matches!(nested, Expr::WindowFunction(_)) {
found = true;
Ok(TreeNodeRecursion::Stop)
} else {
Ok(TreeNodeRecursion::Continue)
}
})
})?;
Ok(found)
}

fn is_snowflake_conditional_event_with_nested_rank(expr: &Expr) -> Result<bool> {
let Expr::WindowFunction(outer) = expr else {
return Ok(false);
};
if !outer
.fun
.name()
.eq_ignore_ascii_case("conditional_true_event")
{
return Ok(false);
}

let mut found_nested_rank = false;
let mut supported = true;
for argument in &outer.params.args {
argument.apply(|nested| {
if let Expr::WindowFunction(inner) = nested {
found_nested_rank = true;
supported &= matches!(
inner.fun.name().to_ascii_lowercase().as_str(),
"lag" | "lead"
) && inner.params.partition_by == outer.params.partition_by
&& inner.params.order_by == outer.params.order_by;
Ok(TreeNodeRecursion::Jump)
} else {
Ok(TreeNodeRecursion::Continue)
}
})?;
}
Ok(found_nested_rank && supported)
}

fn find_innermost_window_exprs<'a>(
exprs: impl IntoIterator<Item = &'a Expr>,
) -> Result<Vec<Expr>> {
let mut window_exprs = Vec::new();
for expr in exprs {
expr.apply(|nested| {
if matches!(nested, Expr::WindowFunction(_))
&& !contains_nested_window_function(nested)?
&& !window_exprs.contains(nested)
{
window_exprs.push(nested.clone());
}
Ok(TreeNodeRecursion::Continue)
})?;
}
Ok(window_exprs)
}

fn flatten_expr_groups(expr_groups: Vec<Vec<Expr>>) -> Vec<Expr> {
expr_groups.into_iter().flatten().collect()
}
Expand Down Expand Up @@ -351,7 +418,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
plan,
select_exprs: mut select_exprs_post_aggr,
having_expr: having_expr_post_aggr,
qualify_expr: qualify_expr_post_aggr,
qualify_expr: mut qualify_expr_post_aggr,
order_by_exprs: mut order_by_rex,
on_exprs: mut on_exprs_post_aggr,
} = if !group_by_exprs.is_empty() || !aggr_exprs.is_empty() {
Expand Down Expand Up @@ -401,11 +468,80 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
.chain(order_by_rex.iter().map(|s| &s.expr))
.chain(on_exprs_post_aggr.iter()),
);
let qualify_had_window_functions = !find_window_exprs(
select_exprs_post_aggr
.iter()
.chain(qualify_expr_post_aggr.iter()),
)
.is_empty();
let (has_nested_windows, all_nested_windows_supported) = window_func_exprs
.iter()
.try_fold((false, true), |state, expr| {
if contains_nested_window_function(expr)? {
Ok::<_, datafusion_common::DataFusionError>((
true,
state.1 && is_snowflake_conditional_event_with_nested_rank(expr)?,
))
} else {
Ok(state)
}
})?;
let plan_nested_windows = self.context_provider.options().sql_parser.dialect
== SqlDialect::Snowflake
&& has_nested_windows
&& all_nested_windows_supported;

// Process window functions after aggregation as they can reference
// aggregate functions in their body
let plan = if window_func_exprs.is_empty() {
plan
} else if plan_nested_windows {
// Snowflake permits LAG/LEAD inside conditional event functions when
// both calls use the same window specification. Materialize the inner
// rank function before planning the conditional event function.
let mut plan = plan;
loop {
let window_level = find_innermost_window_exprs(
select_exprs_post_aggr
.iter()
.chain(qualify_expr_post_aggr.iter())
.chain(order_by_rex.iter().map(|sort| &sort.expr))
.chain(on_exprs_post_aggr.iter()),
)?;
if window_level.is_empty() {
break;
}

plan = LogicalPlanBuilder::window_plan(plan, window_level.clone())?;
let name_preserver = NamePreserver::new_for_projection();
select_exprs_post_aggr = select_exprs_post_aggr
.iter()
.map(|expr| {
let saved_name = name_preserver.save(expr);
rebase_expr(expr, &window_level, &plan)
.map(|expr| saved_name.restore(expr))
})
.collect::<Result<Vec<_>>>()?;
qualify_expr_post_aggr = qualify_expr_post_aggr
.as_ref()
.map(|expr| rebase_expr(expr, &window_level, &plan))
.transpose()?;
order_by_rex = order_by_rex
.into_iter()
.map(|sort_expr| {
Ok(sort_expr.with_expr(rebase_expr(
&sort_expr.expr,
&window_level,
&plan,
)?))
})
.collect::<Result<Vec<_>>>()?;
on_exprs_post_aggr = on_exprs_post_aggr
.iter()
.map(|expr| rebase_expr(expr, &window_level, &plan))
.collect::<Result<Vec<_>>>()?;
}
plan
} else {
let plan = LogicalPlanBuilder::window_plan(plan, window_func_exprs.clone())?;

Expand Down Expand Up @@ -437,37 +573,52 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
// Process QUALIFY clause after window functions
// QUALIFY filters the results of window functions, similar to how HAVING filters aggregates
let plan = if let Some(qualify_expr) = qualify_expr_post_aggr {
// Validate that QUALIFY is used with window functions in SELECT or QUALIFY
let qualify_window_func_exprs = find_window_exprs(
select_exprs_post_aggr
.iter()
.chain(std::iter::once(&qualify_expr)),
);
if qualify_window_func_exprs.is_empty() {
return plan_err!(
"QUALIFY clause requires window functions in the SELECT list or QUALIFY clause"
if plan_nested_windows {
if !qualify_had_window_functions {
return plan_err!(
"QUALIFY clause requires window functions in the SELECT list or QUALIFY clause"
);
}
self.validate_schema_satisfies_exprs(
plan.schema(),
std::slice::from_ref(&qualify_expr),
)?;
LogicalPlanBuilder::from(plan)
.filter(qualify_expr)?
.build()?
} else {
// Validate that QUALIFY is used with window functions in SELECT or QUALIFY
let qualify_window_func_exprs = find_window_exprs(
select_exprs_post_aggr
.iter()
.chain(std::iter::once(&qualify_expr)),
);
}
if qualify_window_func_exprs.is_empty() {
return plan_err!(
"QUALIFY clause requires window functions in the SELECT list or QUALIFY clause"
);
}

// now attempt to resolve columns and replace with fully-qualified columns
let windows_projection_exprs = window_func_exprs
.iter()
.map(|expr| resolve_columns(expr, &plan))
.collect::<Result<Vec<Expr>>>()?;
// now attempt to resolve columns and replace with fully-qualified columns
let windows_projection_exprs = window_func_exprs
.iter()
.map(|expr| resolve_columns(expr, &plan))
.collect::<Result<Vec<Expr>>>()?;

// Rewrite the qualify expression to reference columns from the window plan
let qualify_expr_post_window =
rebase_expr(&qualify_expr, &windows_projection_exprs, &plan)?;
// Rewrite the qualify expression to reference columns from the window plan
let qualify_expr_post_window =
rebase_expr(&qualify_expr, &windows_projection_exprs, &plan)?;

// Validate that the qualify expression can be resolved from the window plan schema
self.validate_schema_satisfies_exprs(
plan.schema(),
std::slice::from_ref(&qualify_expr_post_window),
)?;
// Validate that the qualify expression can be resolved from the window plan schema
self.validate_schema_satisfies_exprs(
plan.schema(),
std::slice::from_ref(&qualify_expr_post_window),
)?;

LogicalPlanBuilder::from(plan)
.filter(qualify_expr_post_window)?
.build()?
LogicalPlanBuilder::from(plan)
.filter(qualify_expr_post_window)?
.build()?
}
} else {
plan
};
Expand Down
49 changes: 48 additions & 1 deletion datafusion/sql/tests/sql_integration.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ use common::MockContextProvider;
use datafusion_common::{
DFSchema, DataFusionError, Result, ScalarValue, assert_contains,
};
use datafusion_expr::expr_fn::create_udwf;
use datafusion_expr::{
ColumnarValue, CreateIndex, DdlStatement, Expr, HigherOrderFunctionArgs,
HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl,
Expand Down Expand Up @@ -58,7 +59,9 @@ use datafusion_functions_aggregate::{
min_max::{max_udaf, min_udaf},
};
use datafusion_functions_nested::make_array::make_array_udf;
use datafusion_functions_window::{rank::rank_udwf, row_number::row_number_udwf};
use datafusion_functions_window::{
lead_lag::lag_udwf, rank::rank_udwf, row_number::row_number_udwf,
};
use insta::{allow_duplicates, assert_snapshot};
use rstest::rstest;
use sqlparser::dialect::{
Expand Down Expand Up @@ -2046,6 +2049,33 @@ fn select_nested_window_function_snowflake() {
);
}

#[test]
fn select_conditional_event_with_nested_lag_snowflake() {
let plan = snowflake_conditional_event_plan(
"SELECT conditional_true_event(\
lag(age) OVER (PARTITION BY state ORDER BY age) > 1\
) OVER (PARTITION BY state ORDER BY age) FROM person",
)
.unwrap();
let formatted = plan.display_indent().to_string();

assert_eq!(formatted.matches("WindowAggr:").count(), 2);
assert_contains!(&formatted, "conditional_true_event");
assert_contains!(&formatted, "lag(person.age");
}

#[test]
fn select_conditional_event_rejects_mismatched_nested_lag_window() {
let err = snowflake_conditional_event_plan(
"SELECT conditional_true_event(\
lag(age) OVER (PARTITION BY state ORDER BY age) > 1\
) OVER (PARTITION BY age ORDER BY age) FROM person",
)
.expect_err("mismatched nested window specifications should fail");

assert_contains!(err.to_string(), "Window function calls cannot be nested");
}

#[test]
fn select_aggregate_inside_window_function() {
// an aggregate as the argument of a window function is legal: the window
Expand Down Expand Up @@ -4099,6 +4129,23 @@ fn logical_plan_with_config_and_options(
logical_plan_from_state(sql, dialect, options, state)
}

fn snowflake_conditional_event_plan(sql: &str) -> Result<LogicalPlan> {
let conditional_true_event = create_udwf(
"conditional_true_event",
DataType::Boolean,
Arc::new(DataType::Int64),
Volatility::Immutable,
Arc::new(|| unreachable!("planning tests do not execute the window function")),
);
let mut state = mock_session_state()
.with_window_function(lag_udwf())
.with_window_function(Arc::new(conditional_true_event));
state.config_options.sql_parser.dialect =
datafusion_common::config::Dialect::Snowflake;

logical_plan_from_state(sql, &GenericDialect {}, ParserOptions::default(), state)
}

fn mock_session_state() -> MockSessionState {
MockSessionState::default()
.with_scalar_function(Arc::new(unicode::character_length().as_ref().clone()))
Expand Down
Loading