This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25181-1e09a2a3d76c03f469e80bed712a6bada45e77c8 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 6574a8c3b2c93ef9689a677c1cc1c0a4a844617f Author: Shayan Gh <[email protected]> AuthorDate: Mon Sep 21 09:00:24 2026 +0000 fix(substrait): consume chained window functions whose default names collide (#25181) ## Which issue does this PR close? - Closes #23007. ## Rationale for this change A logical plan that contains two window aggregates with the same default schema name, distinguished only by aliases, executes fine in DataFusion but cannot be consumed back after a Substrait round trip. The consumer fails with: ``` SchemaError(DuplicateUnqualifiedField { name: "avg(data.b) PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW" }) ``` DataFusion's Substrait producer uses positional projections and does not preserve intermediate window aliases: a `ProjectRel` carries expressions and an output mapping. The aliases that kept the two columns apart in DataFusion (`AS avg1`, `AS avg2`) are dropped on the wire. When the consumer rebuilds the outer projection, the inherited window column and the new window expression both come back under the same default name, and `Window::try_new` rejects the plan. The SQL from the issue reproduces it, and so does a minimal plan built with `LogicalPlanBuilder` that stacks two `row_number()` windows. ## What changes are included in this PR? Two small changes, one per crate: - **`datafusion-expr`**: `group_window_expr_by_sort_keys` (used by `LogicalPlanBuilder::window_plan`) now looks through a single `Expr::Alias` to derive the sort key and keeps the aliased expression in its group, so `window_plan` can build a `Window` node whose output field carries an alias. `LogicalPlanBuilder::window`, filter pushdown, and the physical planner already handle aliased window expressions, so this makes `window_plan` consistent with them. Nested aliases are still rejected, matching what filter pushdown tolerates. - **`datafusion-substrait` consumer**: `from_project_rel` reserves the input schema's field names in a `NameTracker` (new `reserve_schema` helper), aliases any new window expression whose default name would collide, and rewrites the projection to reference the alias while preserving the projection's own output names via `NamePreserver`. Plans without a collision are built exactly as before. ## What is the testing strategy for this PR? - `datafusion-expr`: a unit test (`test_group_window_expr_by_sort_keys_aliased_window_expr`) asserting that an aliased window expression is grouped by the inner function's sort key and kept aliased. It failed on `main` with `Impossibly got non-window expr`. - `datafusion-substrait`: two round-trip tests in `roundtrip_logical_plan.rs`, `stacked_windows_with_same_default_name_via_builder` (minimal `LogicalPlanBuilder` form) and `chained_windows_with_same_default_name` (the SQL from the issue). Both failed before this change with the `DuplicateUnqualifiedField` error above. They assert Arrow schema equality (names, types, nullability) and execute the consumed plan. They do not compare plan text, because the consumer must synthesize an alias for the colliding window, so the text legitimately differs. They compare Arrow schemas rather than full `DFSchema`s because qualifiers and inferred functional dependencies can change during reconstruction: DataFusion's Substrait producer does not preserve subquery aliases, and the rebuilt `Window` nodes infer dependencies the original plan did not record. - The full `datafusion-substrait` test run passes unchanged (273 passed, 6 ignored across the lib, integration, and doc tests), confirming the non-colliding path is unaffected. This change was developed with AI assistance. I have reviewed and understand every line and stand behind it. ## Are there any user-facing changes? No API changes. Plans consumed from Substrait that previously failed with `DuplicateUnqualifiedField` now consume successfully. Synthesized aliases of the form `<default name>__temp__N` may appear in the decoded logical plan; final output column names are preserved. Plans without such a collision are unchanged. --------- Co-authored-by: AI assistants <[email protected]> --- datafusion/expr/src/utils.rs | 74 ++++++++++++++---- .../src/logical_plan/consumer/rel/project_rel.rs | 60 +++++++++++++++ .../substrait/src/logical_plan/consumer/utils.rs | 8 +- .../tests/cases/roundtrip_logical_plan.rs | 88 ++++++++++++++++++++++ 4 files changed, 213 insertions(+), 17 deletions(-) diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 9462fa1f08..07af8b1f56 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -616,28 +616,45 @@ pub fn compare_sort_expr( Ordering::Equal } -/// Group a slice of window expression expr by their order by expressions +/// Group window expressions by their sort keys, preserving any outer alias. pub fn group_window_expr_by_sort_keys( window_expr: impl IntoIterator<Item = Expr>, ) -> Result<Vec<(WindowSortKey, Vec<Expr>)>> { let mut result = vec![]; - window_expr.into_iter().try_for_each(|expr| match &expr { - Expr::WindowFunction(window_fun) => { - let WindowFunctionParams{ partition_by, order_by, ..} = &window_fun.as_ref().params; - let sort_key = generate_sort_key(partition_by, order_by)?; - if let Some((_, values)) = result.iter_mut().find( - |group: &&mut (WindowSortKey, Vec<Expr>)| matches!(group, (key, _) if *key == sort_key), - ) { - values.push(expr); - } else { - result.push((sort_key, vec![expr])) - } - Ok(()) + + window_expr.into_iter().try_for_each(|expr| { + // Read the window's settings through one alias. + // Keep `expr` intact so its output name is preserved. + let inner = match &expr { + Expr::Alias(alias) => alias.expr.as_ref(), + _ => &expr, + }; + + let Expr::WindowFunction(window_fun) = inner else { + return internal_err!("Impossibly got non-window expr {expr:?}"); + }; + + let WindowFunctionParams { + partition_by, + order_by, + .. + } = &window_fun.as_ref().params; + + let sort_key = generate_sort_key(partition_by, order_by)?; + + if let Some((_, values)) = result.iter_mut().find( + |group: &&mut (WindowSortKey, Vec<Expr>)| { + matches!(group, (key, _) if *key == sort_key) + }, + ) { + values.push(expr); + } else { + result.push((sort_key, vec![expr])); } - other => internal_err!( - "Impossibly got non-window expr {other:?}" - ), + + Ok(()) })?; + Ok(result) } @@ -1567,6 +1584,31 @@ mod tests { Ok(()) } + #[test] + fn test_group_window_expr_by_sort_keys_aliased_window_expr() -> Result<()> { + let age_asc = Sort::new(col("age"), true, true); + let max1 = Expr::from(WindowFunction::new( + WindowFunctionDefinition::AggregateUDF(max_udaf()), + vec![col("name")], + )) + .order_by(vec![age_asc.clone()]) + .build() + .unwrap(); + // The same window function under an alias, as the Substrait consumer + // produces when a window column's default name collides with an + // input column. It must be grouped by the inner function's sort key + // and kept aliased. + let max1_aliased = max1.clone().alias("max_name"); + + let result = + group_window_expr_by_sort_keys(vec![max1.clone(), max1_aliased.clone()])?; + + let key = vec![(age_asc, false)]; + let expected: Vec<(WindowSortKey, Vec<Expr>)> = + vec![(key, vec![max1, max1_aliased])]; + assert_eq!(expected, result); + Ok(()) + } #[test] fn test_group_window_expr_by_sort_keys() -> Result<()> { let age_asc = Sort::new(col("age"), true, true); diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs index 5aea6c809b..fa93add0ff 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs @@ -18,8 +18,10 @@ use crate::logical_plan::consumer::SubstraitConsumer; use crate::logical_plan::consumer::utils::NameTracker; use async_recursion::async_recursion; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::common::{Column, not_impl_err}; use datafusion::logical_expr::builder::project; +use datafusion::logical_expr::expr_rewriter::NamePreserver; use datafusion::logical_expr::utils::find_window_exprs; use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; use std::collections::HashSet; @@ -65,6 +67,37 @@ pub async fn from_project_rel( } let input = if !window_exprs.is_empty() { + // Window outputs must have unique names across the input schema + // and the new window expressions. + let mut window_names = NameTracker::new(); + window_names.reserve_schema(&original_schema); + + let mut aliased_columns: Vec<(Expr, Expr)> = vec![]; + let window_exprs = window_exprs + .into_iter() + .map(|window_expr| { + let named = + window_names.get_uniquely_named_expr(window_expr.clone())?; + + if let Expr::Alias(alias) = &named { + aliased_columns.push(( + window_expr, + Expr::Column(Column::from_name(&alias.name)), + )); + } + + Ok(named) + }) + .collect::<datafusion::common::Result<Vec<_>>>()?; + + // References to renamed windows must point to their new output columns. + if !aliased_columns.is_empty() { + explicit_exprs = explicit_exprs + .into_iter() + .map(|expr| reference_aliased_windows(expr, &aliased_columns)) + .collect::<datafusion::common::Result<Vec<_>>>()?; + } + LogicalPlanBuilder::window_plan(input, window_exprs)? } else { input @@ -81,3 +114,30 @@ pub async fn from_project_rel( not_impl_err!("Projection without an input is not supported") } } + +/// Reference renamed window outputs while preserving the projection's +/// original output name. +fn reference_aliased_windows( + expr: Expr, + aliased_columns: &[(Expr, Expr)], +) -> datafusion::common::Result<Expr> { + let saved_name = NamePreserver::new_for_projection().save(&expr); + + let rewritten = expr + .transform_down(|node| { + match aliased_columns + .iter() + .find(|(window_expr, _)| *window_expr == node) + { + Some((_, column)) => Ok(Transformed::new( + column.clone(), + true, + TreeNodeRecursion::Jump, + )), + None => Ok(Transformed::no(node)), + } + })? + .data; + + Ok(saved_name.restore(rewritten)) +} diff --git a/datafusion/substrait/src/logical_plan/consumer/utils.rs b/datafusion/substrait/src/logical_plan/consumer/utils.rs index 824c79452d..04fafdc691 100644 --- a/datafusion/substrait/src/logical_plan/consumer/utils.rs +++ b/datafusion/substrait/src/logical_plan/consumer/utils.rs @@ -18,7 +18,7 @@ use crate::logical_plan::consumer::SubstraitConsumer; use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UnionFields}; use datafusion::common::{ - DFSchema, DFSchemaRef, TableReference, exec_err, not_impl_err, + Column, DFSchema, DFSchemaRef, TableReference, exec_err, not_impl_err, substrait_datafusion_err, substrait_err, }; use datafusion::logical_expr::expr::Sort; @@ -424,6 +424,12 @@ impl NameTracker { } } + pub(super) fn reserve_schema(&mut self, schema: &DFSchema) { + for (qualifier, field) in schema.iter() { + self.insert(&Expr::Column(Column::from((qualifier, field)))); + } + } + /// Check if the expression would cause a conflict either in: /// 1. validate_unique_names (duplicate schema_name) /// 2. DFSchema::check_names (ambiguous reference) diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 813d0ed6c3..43bb94746f 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -1684,6 +1684,94 @@ async fn simple_window_function() -> Result<()> { roundtrip("SELECT RANK() OVER (PARTITION BY a ORDER BY b), d, sum(b) OVER (PARTITION BY a) FROM data;").await } +#[tokio::test] +async fn stacked_windows_with_same_default_name_via_builder() -> Result<()> { + // Substrait projections are positional and drop DataFusion's aliases, so the + // inherited `rn1` column and the new window both come back under the same + // default name. The consumer must keep them apart. + // See https://github.com/apache/datafusion/issues/23007 + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::functions_window::expr_fn::row_number; + use datafusion::logical_expr::ExprFunctionExt; + + let ctx = create_context().await?; + let scan = ctx.table("data").await?.into_optimized_plan()?; + let row_number_by_a = || { + row_number() + .order_by(vec![col("a").sort(true, false)]) + .build() + }; + // The filter between the two windows makes `rn1` and `rn2` differ for the + // surviving row. A consumer that resolved `rn2` to the inherited window + // column would then return wrong values instead of passing unnoticed. + let plan = LogicalPlanBuilder::from(scan) + .window(vec![row_number_by_a()?.alias("rn1")])? + .filter(col("a").gt(lit(1i64)))? + .window(vec![row_number_by_a()?.alias("rn2")])? + .build()?; + + let plan2 = substrait_roundtrip(&plan, &ctx).await?; + // Compare output fields; qualifiers and functional dependencies + // can differ after a Substrait round trip. + assert_eq!(plan.schema().as_arrow(), plan2.schema().as_arrow()); + + let expected = DataFrame::new(ctx.state(), plan) + .select_columns(&["a", "rn1", "rn2"])? + .collect() + .await?; + let actual = DataFrame::new(ctx.state(), plan2) + .select_columns(&["a", "rn1", "rn2"])? + .collect() + .await?; + let actual = pretty_format_batches(&actual)?.to_string(); + assert_eq!(pretty_format_batches(&expected)?.to_string(), actual); + assert_snapshot!( + actual, + @r" + +---+-----+-----+ + | a | rn1 | rn2 | + +---+-----+-----+ + | 3 | 2 | 1 | + +---+-----+-----+ + " + ); + Ok(()) +} + +#[tokio::test] +async fn chained_windows_with_same_default_name() -> Result<()> { + // SQL form from https://github.com/apache/datafusion/issues/23007. The two + // `avg` windows differ only by aliases that Substrait does not carry. As in + // roundtrip_self_join, the consumer must synthesize an alias, so the plan + // text differs; verify schema and executability instead. + let ctx = create_context().await?; + let plan = ctx + .sql( + "SELECT a, b, avg1, avg2 FROM ( + SELECT a, b, avg1, + row_number() OVER () AS seq2, + avg(b) OVER (PARTITION BY a ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS avg2 + FROM ( + SELECT a, b, avg1 FROM ( + SELECT a, b, + row_number() OVER () AS seq1, + avg(b) OVER (PARTITION BY a ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS avg1 + FROM data + ) t1 ORDER BY seq1 + ) t2 + ) t3 ORDER BY seq2", + ) + .await? + .into_optimized_plan()?; + + let plan2 = substrait_roundtrip(&plan, &ctx).await?; + // Compare output fields; qualifiers and functional dependencies + // can differ after a Substrait round trip. + assert_eq!(plan.schema().as_arrow(), plan2.schema().as_arrow()); + DataFrame::new(ctx.state(), plan2).show().await?; + Ok(()) +} + #[tokio::test] async fn window_with_rows() -> Result<()> { roundtrip("SELECT sum(b) OVER (PARTITION BY a ROWS BETWEEN 2 PRECEDING AND CURRENT ROW) FROM data;").await?; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
