Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
28 changes: 0 additions & 28 deletions datafusion/core/tests/sql/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,34 +21,6 @@ use super::*;
use datafusion_common::{ParamValues, ScalarValue, metadata::ScalarAndMetadata};
use insta::assert_snapshot;

#[tokio::test]
async fn snowflake_nested_window_functions_execute_in_stages() -> Result<()> {
let config =
SessionConfig::new().set_str("datafusion.sql_parser.dialect", "Snowflake");
let ctx = SessionContext::new_with_config(config);

let results = ctx
.sql(
"SELECT column1 AS id, \
SUM(SUM(column1) OVER ()) OVER () AS nested_sum \
FROM VALUES (1), (2), (3) ORDER BY id",
)
.await?
.collect()
.await?;

assert_snapshot!(batches_to_sort_string(&results), @r"
+----+------------+
| id | nested_sum |
+----+------------+
| 1 | 18 |
| 2 | 18 |
| 3 | 18 |
+----+------------+
");
Ok(())
}

#[tokio::test]
async fn test_list_query_parameters() -> Result<()> {
let tmp_dir = TempDir::new()?;
Expand Down
166 changes: 28 additions & 138 deletions datafusion/sql/src/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ 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 @@ -40,8 +39,7 @@ 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::{
NamePreserver, normalize_col, normalize_col_with_schemas_and_ambiguity_check,
normalize_sorts,
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 @@ -89,39 +87,6 @@ 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 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 @@ -386,7 +351,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
plan,
select_exprs: mut select_exprs_post_aggr,
having_expr: having_expr_post_aggr,
qualify_expr: mut qualify_expr_post_aggr,
qualify_expr: 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 @@ -436,71 +401,11 @@ 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 plan_nested_windows = self.context_provider.options().sql_parser.dialect
== SqlDialect::Snowflake
&& window_func_exprs.iter().try_fold(false, |found, expr| {
Ok::<_, datafusion_common::DataFusionError>(
found || contains_nested_window_function(expr)?,
)
})?;

// 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 a window call inside another window call. A logical
// Window node cannot evaluate that shape directly, so materialize each
// innermost level and rebase its parents onto the generated columns.
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 @@ -532,52 +437,37 @@ 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 {
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)),
// 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 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
14 changes: 6 additions & 8 deletions datafusion/sql/tests/sql_integration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2035,17 +2035,15 @@ fn select_nested_window_function_snowflake() {
let mut config_options = datafusion_common::config::ConfigOptions::new();
config_options.sql_parser.dialect = datafusion_common::config::Dialect::Snowflake;

let plan = logical_plan_with_config(
let err = logical_plan_with_config(
"SELECT sum(sum(age) OVER ()) OVER () FROM person",
config_options,
)
.unwrap();
assert_snapshot!(plan, @r"
Projection: sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
WindowAggr: windowExpr=[[sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]
WindowAggr: windowExpr=[[sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]
TableScan: person
");
.expect_err("Snowflake rejects nested window function calls");
assert_snapshot!(
err.strip_backtrace(),
@"Error during planning: Window function calls cannot be nested: 'sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'"
);
}

#[test]
Expand Down
Loading