From 2930839dabbca42d537ebd202d31fd4cf50ccba3 Mon Sep 17 00:00:00 2001 From: Wyatt Herkamp Date: Sat, 5 Sep 2026 11:38:59 -0400 Subject: [PATCH 1/2] Remove to string comparison in push_projection_dedupl --- .../core/benches/sql_planner_extended.rs | 59 ++++++- datafusion/sql/src/utils.rs | 144 ++++++++++++++++-- 2 files changed, 192 insertions(+), 11 deletions(-) diff --git a/datafusion/core/benches/sql_planner_extended.rs b/datafusion/core/benches/sql_planner_extended.rs index 5bea9860c4be7..96a424e184b46 100644 --- a/datafusion/core/benches/sql_planner_extended.rs +++ b/datafusion/core/benches/sql_planner_extended.rs @@ -18,6 +18,7 @@ use arrow::array::{ArrayRef, RecordBatch}; use arrow_schema::DataType; use arrow_schema::TimeUnit::Nanosecond; +use arrow_schema::{Field, Schema}; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion::prelude::{DataFrame, SessionContext}; use datafusion_catalog::MemTable; @@ -69,6 +70,28 @@ fn register_string_table(ctx: &SessionContext, num_columns: usize, num_rows: usi ctx.register_table("t", Arc::new(table)).unwrap(); } +/// Registers a table `name` with `num_columns` Int32 columns (`a0..aN`) plus an +/// `arr` `List` column, used to benchmark `UNNEST` planning. +fn register_list_table(ctx: &SessionContext, name: &str, num_columns: usize) { + let mut fields: Vec = (0..num_columns) + .map(|i| Field::new(format!("a{i}"), DataType::Int32, true)) + .collect(); + fields.push(Field::new( + "arr", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + )); + let schema = Arc::new(Schema::new(fields)); + let table = MemTable::try_new(schema, vec![vec![]]).unwrap(); + ctx.register_table(name, Arc::new(table)).unwrap(); +} + +/// Create a logical plan from the specified sql (parse + plan only, NO +/// analysis or optimization). Isolates SQL planner cost. +fn logical_plan(ctx: &SessionContext, rt: &Runtime, sql: &str) { + black_box(rt.block_on(ctx.state().create_logical_plan(sql)).unwrap()); +} + /// Build a dataframe for testing logical plan optimization fn build_test_data_frame(ctx: &SessionContext, rt: &Runtime) -> DataFrame { register_string_table(ctx, 100, 1000); @@ -387,8 +410,8 @@ fn criterion_benchmark(c: &mut Criterion) { let case_heavy_left_join_df = build_case_heavy_left_join_df(&case_heavy_ctx, &rt); // really slow :( - let mut group = c.benchmark_group("sample_size_5"); - group.sample_size(5); + let mut group = c.benchmark_group("sample_size_10"); + group.sample_size(10); group.bench_function("logical_plan_optimize", |b| { b.iter(|| { let df_clone = df.clone(); @@ -525,6 +548,38 @@ fn criterion_benchmark(c: &mut Criterion) { ); }) }); + + // UNNEST planning. The unnest rewrite builds an inner projection that is + // deduplicated on every pushed expression, so wide select lists stress it. + // See `push_projection_dedupl` in datafusion/sql/src/utils.rs. + let unnest_ctx = SessionContext::new(); + register_list_table(&unnest_ctx, "t_list200", 200); + register_list_table(&unnest_ctx, "t_list1000", 1000); + + // Bare columns next to an unnest: each column is pushed once. + for num_columns in [200usize, 1000] { + let cols = (0..num_columns) + .map(|i| format!("a{i}")) + .collect::>() + .join(", "); + let query = format!("SELECT unnest(arr), {cols} FROM t_list{num_columns}"); + c.bench_function(&format!("logical_unnest_plus_{num_columns}_columns"), |b| { + b.iter(|| logical_plan(&unnest_ctx, &rt, &query)) + }); + } + + // Columns inside expressions that contain an unnest: exercises both the + // column path and the repeated unnest placeholder alias path. + { + let exprs = (0..200) + .map(|i| format!("unnest(arr) + a{i}")) + .collect::>() + .join(", "); + let query = format!("SELECT {exprs} FROM t_list200"); + c.bench_function("logical_unnest_in_200_exprs", |b| { + b.iter(|| logical_plan(&unnest_ctx, &rt, &query)) + }); + } } criterion_group!(benches, criterion_benchmark); diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 3b571eed279dd..25092a30108f9 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -26,12 +26,12 @@ use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, }; use datafusion_common::{ - Column, DFSchemaRef, Diagnostic, HashMap, Result, ScalarValue, + Column, DFSchemaRef, Diagnostic, HashMap, Result, ScalarValue, TableReference, assert_or_internal_err, exec_datafusion_err, exec_err, internal_err, plan_err, }; use datafusion_expr::builder::get_struct_unnested_columns; use datafusion_expr::expr::{ - Alias, GroupingSet, Unnest, WindowFunction, WindowFunctionParams, + Alias, Cast, GroupingSet, TryCast, Unnest, WindowFunction, WindowFunctionParams, }; use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs}; use datafusion_expr::{ @@ -664,12 +664,75 @@ impl TreeNodeRewriter for RecursiveUnnestRewriter<'_> { } } +/// A key that compares equal exactly when two expressions would produce the +/// same [`Expr::schema_name`], built without formatting or allocating. +/// +/// The rules mirror the `SchemaDisplay` rendering of an [`Expr`]: +/// - [`Expr::Alias`] renders as `[relation.]name`, hiding the aliased expr +/// - [`Expr::Column`] renders as `[relation.]name` +/// - [`Expr::Cast`] / [`Expr::TryCast`] render as their input expression +/// - anything else renders structurally, so structural equality is used +/// +/// The only divergence from comparing formatted names is when two +/// *structurally different* expressions happen to render identically (e.g. +/// `foo(a)` and `foo(CAST(a AS BIGINT))`). Such expressions are kept here, +/// whereas the string comparison silently dropped the second one. Projections +/// containing them are rejected by `validate_unique_names` either way. +#[derive(Debug)] +enum SchemaNameKey<'a> { + /// Alias and Column both render as `[relation.]name` + Named { + relation: Option<&'a TableReference>, + name: &'a str, + }, + /// Any other expr renders structurally, so compare structurally + Other(&'a Expr), +} + +impl<'a> SchemaNameKey<'a> { + fn new(mut expr: &'a Expr) -> Self { + loop { + match expr { + // The schema name of a cast is the schema name of its input + Expr::Cast(Cast { expr: inner, .. }) + | Expr::TryCast(TryCast { expr: inner, .. }) => expr = inner, + Expr::Column(Column { relation, name, .. }) + | Expr::Alias(Alias { relation, name, .. }) => { + return Self::Named { + relation: relation.as_ref(), + name, + }; + } + other => return Self::Other(other), + } + } + } +} + +impl PartialEq for SchemaNameKey<'_> { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + ( + Self::Named { + relation: l_relation, + name: l_name, + }, + Self::Named { + relation: r_relation, + name: r_name, + }, + ) => l_name == r_name && l_relation == r_relation, + (Self::Other(l), Self::Other(r)) => l == r, + _ => false, + } + } +} + +/// Pushes `expr` onto `projection` unless an expression with the same schema +/// name is already present, so the resulting projection has unique field names. fn push_projection_dedupl(projection: &mut Vec, expr: Expr) { - let schema_name = expr.schema_name().to_string(); - if !projection - .iter() - .any(|e| e.schema_name().to_string() == schema_name) - { + let key = SchemaNameKey::new(&expr); + if !projection.iter().any(|e| SchemaNameKey::new(e) == key) { projection.push(expr); } } @@ -743,14 +806,77 @@ mod tests { use arrow::datatypes::{DataType as ArrowDataType, Field, Fields, Schema}; use datafusion_common::{Column, DFSchema, Result}; use datafusion_expr::{ - ColumnUnnestList, EmptyRelation, LogicalPlan, col, lit, unnest, + ColumnUnnestList, EmptyRelation, Expr, LogicalPlan, cast, col, lit, try_cast, + unnest, }; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; - use crate::utils::{resolve_positions_to_exprs, rewrite_recursive_unnest_bottom_up}; + use crate::utils::{ + SchemaNameKey, push_projection_dedupl, resolve_positions_to_exprs, + rewrite_recursive_unnest_bottom_up, + }; use indexmap::IndexMap; + /// `SchemaNameKey` must agree with comparing formatted schema names for + /// every kind of expression the unnest rewrite pushes into the inner + /// projection. + #[test] + fn test_schema_name_key_matches_schema_name() { + let exprs: Vec = vec![ + col("a"), + col("a"), + col("t.a"), + col("b"), + col("x").alias("n"), + col("y").alias("n"), + col("n"), + cast(col("a"), ArrowDataType::Int64), + try_cast(col("a"), ArrowDataType::Int64), + cast(col("b"), ArrowDataType::Int64), + col("a").add(lit(1)), + col("a").add(lit(1)), + col("a").add(lit(2)), + count(col("a")), + count(col("a")), + lit(1), + col("__unnest_placeholder(arr)").alias("__unnest_placeholder(arr)"), + ]; + + for l in &exprs { + for r in &exprs { + let by_key = SchemaNameKey::new(l) == SchemaNameKey::new(r); + let by_name = l.schema_name().to_string() == r.schema_name().to_string(); + assert_eq!( + by_key, by_name, + "key comparison disagrees with schema_name for `{l}` vs `{r}`" + ); + } + } + } + + #[test] + fn test_push_projection_dedupl() { + let mut projection = vec![]; + push_projection_dedupl(&mut projection, col("a")); + push_projection_dedupl(&mut projection, col("a")); + push_projection_dedupl(&mut projection, col("t.a")); + push_projection_dedupl(&mut projection, col("x").alias("n")); + // same alias name as the previous alias, different inner expr + push_projection_dedupl(&mut projection, col("y").alias("n")); + // cast renders as its input, which is already present + push_projection_dedupl(&mut projection, cast(col("a"), ArrowDataType::Int64)); + push_projection_dedupl(&mut projection, col("a").add(lit(1))); + push_projection_dedupl(&mut projection, col("a").add(lit(1))); + push_projection_dedupl(&mut projection, col("a").add(lit(2))); + + let names: Vec = projection + .iter() + .map(|e| e.schema_name().to_string()) + .collect(); + assert_eq!(names, vec!["a", "t.a", "n", "a + Int32(1)", "a + Int32(2)"]); + } + fn column_unnests_eq( l: Vec<&str>, r: &IndexMap>>, From 01215c1ac88d9a94332f39eabfc31c9f117fd66e Mon Sep 17 00:00:00 2001 From: Wyatt Herkamp Date: Mon, 7 Sep 2026 11:40:32 -0400 Subject: [PATCH 2/2] Use HahSet --- datafusion/sql/src/select.rs | 16 +- datafusion/sql/src/utils.rs | 240 ++++++++---------- .../sqllogictest/test_files/distinct_on.slt | 22 ++ 3 files changed, 130 insertions(+), 148 deletions(-) diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index bbd9d203eb124..fe2c77a8662f0 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -23,10 +23,10 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; use crate::query::to_order_by_exprs_with_select; use crate::utils::{ CheckColumnsMustReferenceAggregatePurpose, CheckColumnsSatisfyExprsPurpose, - check_columns_satisfy_exprs, extract_aliases, rebase_expr, resolve_aliases_to_exprs, - resolve_columns, resolve_positions_to_exprs, rewrite_recursive_unnest_bottom_up, - rewrite_recursive_unnests_bottom_up, substitute_top_level_alias, - substitute_top_level_aliases_in_sorts, + DedupedProjection, check_columns_satisfy_exprs, extract_aliases, rebase_expr, + resolve_aliases_to_exprs, resolve_columns, resolve_positions_to_exprs, + rewrite_recursive_unnest_bottom_up, rewrite_recursive_unnests_bottom_up, + substitute_top_level_alias, substitute_top_level_aliases_in_sorts, }; use arrow::datatypes::DataType; @@ -639,7 +639,7 @@ impl SqlToRel<'_, S> { let mut unnest_columns = IndexMap::new(); // from which columns used for projection, before the unnest happen // including non unnest columns and unnest columns - let mut inner_projection_exprs = vec![]; + let mut inner_projection_exprs = DedupedProjection::default(); let mut outer_expr_groups = Vec::with_capacity(intermediate_expr_groups.len()); @@ -704,7 +704,7 @@ impl SqlToRel<'_, S> { } intermediate_plan = LogicalPlanBuilder::from(intermediate_plan) - .project(inner_projection_exprs)? + .project(inner_projection_exprs.into_exprs())? .unnest_columns_with_options(unnest_col_vec, unnest_options)? .build()?; intermediate_expr_groups = outer_expr_groups; @@ -807,7 +807,7 @@ impl SqlToRel<'_, S> { loop { let mut unnest_columns = IndexMap::new(); - let mut inner_projection_exprs = vec![]; + let mut inner_projection_exprs = DedupedProjection::default(); let outer_projection_exprs = rewrite_recursive_unnests_bottom_up( &intermediate_plan, @@ -842,7 +842,7 @@ impl SqlToRel<'_, S> { columns } }; - projection_exprs.extend(inner_projection_exprs); + projection_exprs.extend(inner_projection_exprs.into_exprs()); let mut unnest_col_vec = vec![]; diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 25092a30108f9..fe579bc45f96c 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -26,12 +26,12 @@ use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, }; use datafusion_common::{ - Column, DFSchemaRef, Diagnostic, HashMap, Result, ScalarValue, TableReference, + Column, DFSchemaRef, Diagnostic, HashMap, HashSet, Result, ScalarValue, assert_or_internal_err, exec_datafusion_err, exec_err, internal_err, plan_err, }; use datafusion_expr::builder::get_struct_unnested_columns; use datafusion_expr::expr::{ - Alias, Cast, GroupingSet, TryCast, Unnest, WindowFunction, WindowFunctionParams, + Alias, GroupingSet, Unnest, WindowFunction, WindowFunctionParams, }; use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs}; use datafusion_expr::{ @@ -394,7 +394,7 @@ pub(crate) fn value_to_string(value: &Value) -> Option { pub(crate) fn rewrite_recursive_unnests_bottom_up( input: &LogicalPlan, unnest_placeholder_columns: &mut IndexMap>>, - inner_projection_exprs: &mut Vec, + inner_projection_exprs: &mut DedupedProjection, original_exprs: &[Expr], ) -> Result> { Ok(original_exprs @@ -425,7 +425,7 @@ struct RecursiveUnnestRewriter<'a> { // Useful to detect which child expr is a part of/ not a part of unnest operation top_most_unnest: Option, consecutive_unnest: Vec>, - inner_projection_exprs: &'a mut Vec, + inner_projection_exprs: &'a mut DedupedProjection, columns_unnestings: &'a mut IndexMap>>, transformed_root_exprs: Option>, } @@ -493,10 +493,8 @@ impl RecursiveUnnestRewriter<'_> { struct_allowed, "unnest on struct can only be applied at the root level of select expression" ); - push_projection_dedupl( - self.inner_projection_exprs, - expr_in_unnest.clone().alias(placeholder_name.clone()), - ); + self.inner_projection_exprs + .push(expr_in_unnest.clone().alias(placeholder_name.clone())); self.columns_unnestings .insert(Column::from_name(placeholder_name.clone()), None); Ok(get_struct_unnested_columns(&placeholder_name, inner_fields) @@ -509,10 +507,8 @@ impl RecursiveUnnestRewriter<'_> { | DataType::LargeList(_) | DataType::ListView(_) | DataType::LargeListView(_) => { - push_projection_dedupl( - self.inner_projection_exprs, - expr_in_unnest.clone().alias(placeholder_name.clone()), - ); + self.inner_projection_exprs + .push(expr_in_unnest.clone().alias(placeholder_name.clone())); let post_unnest_expr = col(post_unnest_name.clone()).alias(alias_name); let list_unnesting = self @@ -657,83 +653,57 @@ impl TreeNodeRewriter for RecursiveUnnestRewriter<'_> { // e.g given expr tree unnest(col_a) + col_b, we have to retain projection of col_b // this condition can be checked by maintaining an Option if matches!(&expr, Expr::Column(_)) && self.top_most_unnest.is_none() { - push_projection_dedupl(self.inner_projection_exprs, expr.clone()); + self.inner_projection_exprs.push(expr.clone()); } Ok(Transformed::no(expr)) } } -/// A key that compares equal exactly when two expressions would produce the -/// same [`Expr::schema_name`], built without formatting or allocating. -/// -/// The rules mirror the `SchemaDisplay` rendering of an [`Expr`]: -/// - [`Expr::Alias`] renders as `[relation.]name`, hiding the aliased expr -/// - [`Expr::Column`] renders as `[relation.]name` -/// - [`Expr::Cast`] / [`Expr::TryCast`] render as their input expression -/// - anything else renders structurally, so structural equality is used +/// An inner projection under construction that drops any expression whose +/// [`Expr::schema_name`] is already present, so the finished projection has +/// unique field names. /// -/// The only divergence from comparing formatted names is when two -/// *structurally different* expressions happen to render identically (e.g. -/// `foo(a)` and `foo(CAST(a AS BIGINT))`). Such expressions are kept here, -/// whereas the string comparison silently dropped the second one. Projections -/// containing them are rejected by `validate_unique_names` either way. -#[derive(Debug)] -enum SchemaNameKey<'a> { - /// Alias and Column both render as `[relation.]name` - Named { - relation: Option<&'a TableReference>, - name: &'a str, - }, - /// Any other expr renders structurally, so compare structurally - Other(&'a Expr), +/// The unnest rewrite references inner projection fields *by schema name* (see +/// [`rewrite_recursive_unnest_bottom_up`]), so the deduplication has to be by +/// schema name too. Each pushed expression is rendered once and its name +/// cached, making an `n` expression projection cost `n` renders rather than +/// the `O(n^2)` renders a pairwise comparison would need. +#[derive(Debug, Default)] +pub(crate) struct DedupedProjection { + exprs: Vec, + names: HashSet, } -impl<'a> SchemaNameKey<'a> { - fn new(mut expr: &'a Expr) -> Self { - loop { - match expr { - // The schema name of a cast is the schema name of its input - Expr::Cast(Cast { expr: inner, .. }) - | Expr::TryCast(TryCast { expr: inner, .. }) => expr = inner, - Expr::Column(Column { relation, name, .. }) - | Expr::Alias(Alias { relation, name, .. }) => { - return Self::Named { - relation: relation.as_ref(), - name, - }; - } - other => return Self::Other(other), - } +impl DedupedProjection { + /// Pushes `expr` unless an expression with the same schema name is already + /// present. + fn push(&mut self, expr: Expr) { + if self.names.insert(expr.schema_name().to_string()) { + self.exprs.push(expr); } } -} -impl PartialEq for SchemaNameKey<'_> { - fn eq(&self, other: &Self) -> bool { - match (self, other) { - ( - Self::Named { - relation: l_relation, - name: l_name, - }, - Self::Named { - relation: r_relation, - name: r_name, - }, - ) => l_name == r_name && l_relation == r_relation, - (Self::Other(l), Self::Other(r)) => l == r, - _ => false, + /// Like [`Self::push`], but returns the schema name of `expr`, which the + /// caller needs in order to reference the field from the outer projection. + fn push_returning_name(&mut self, expr: Expr) -> String { + let name = expr.schema_name().to_string(); + if self.names.insert(name.clone()) { + self.exprs.push(expr); } + name + } + + /// The expressions pushed so far, in push order. Only the tests need to + /// inspect a projection while it is still being built; `select.rs` takes + /// the finished `Vec` via [`Self::into_exprs`]. + #[cfg(test)] + fn exprs(&self) -> &[Expr] { + &self.exprs } -} -/// Pushes `expr` onto `projection` unless an expression with the same schema -/// name is already present, so the resulting projection has unique field names. -fn push_projection_dedupl(projection: &mut Vec, expr: Expr) { - let key = SchemaNameKey::new(&expr); - if !projection.iter().any(|e| SchemaNameKey::new(e) == key) { - projection.push(expr); + pub(crate) fn into_exprs(self) -> Vec { + self.exprs } } /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection @@ -748,7 +718,7 @@ fn push_projection_dedupl(projection: &mut Vec, expr: Expr) { pub(crate) fn rewrite_recursive_unnest_bottom_up( input: &LogicalPlan, unnest_placeholder_columns: &mut IndexMap>>, - inner_projection_exprs: &mut Vec, + inner_projection_exprs: &mut DedupedProjection, original_expr: &Expr, ) -> Result> { let mut rewriter = RecursiveUnnestRewriter { @@ -782,13 +752,13 @@ pub(crate) fn rewrite_recursive_unnest_bottom_up( if matches!(&transformed_expr, Expr::Column(_)) || matches!(&transformed_expr, Expr::Wildcard { .. }) { - push_projection_dedupl(inner_projection_exprs, transformed_expr.clone()); + inner_projection_exprs.push(transformed_expr.clone()); Ok(vec![transformed_expr]) } else { // We need to evaluate the expr in the inner projection, // outer projection just select its name - let column_name = transformed_expr.schema_name().to_string(); - push_projection_dedupl(inner_projection_exprs, transformed_expr); + let column_name = + inner_projection_exprs.push_returning_name(transformed_expr); Ok(vec![Expr::Column(Column::from_name(column_name))]) } } else { @@ -806,77 +776,67 @@ mod tests { use arrow::datatypes::{DataType as ArrowDataType, Field, Fields, Schema}; use datafusion_common::{Column, DFSchema, Result}; use datafusion_expr::{ - ColumnUnnestList, EmptyRelation, Expr, LogicalPlan, cast, col, lit, try_cast, - unnest, + ColumnUnnestList, EmptyRelation, LogicalPlan, cast, col, lit, try_cast, unnest, }; use datafusion_functions::core::expr_ext::FieldAccessor; use datafusion_functions_aggregate::expr_fn::count; use crate::utils::{ - SchemaNameKey, push_projection_dedupl, resolve_positions_to_exprs, - rewrite_recursive_unnest_bottom_up, + DedupedProjection, resolve_positions_to_exprs, rewrite_recursive_unnest_bottom_up, }; use indexmap::IndexMap; - /// `SchemaNameKey` must agree with comparing formatted schema names for - /// every kind of expression the unnest rewrite pushes into the inner - /// projection. - #[test] - fn test_schema_name_key_matches_schema_name() { - let exprs: Vec = vec![ - col("a"), - col("a"), - col("t.a"), - col("b"), - col("x").alias("n"), - col("y").alias("n"), - col("n"), - cast(col("a"), ArrowDataType::Int64), - try_cast(col("a"), ArrowDataType::Int64), - cast(col("b"), ArrowDataType::Int64), - col("a").add(lit(1)), - col("a").add(lit(1)), - col("a").add(lit(2)), - count(col("a")), - count(col("a")), - lit(1), - col("__unnest_placeholder(arr)").alias("__unnest_placeholder(arr)"), - ]; - - for l in &exprs { - for r in &exprs { - let by_key = SchemaNameKey::new(l) == SchemaNameKey::new(r); - let by_name = l.schema_name().to_string() == r.schema_name().to_string(); - assert_eq!( - by_key, by_name, - "key comparison disagrees with schema_name for `{l}` vs `{r}`" - ); - } - } - } - + /// [`DedupedProjection`] deduplicates on [`Expr::schema_name`], which + /// hides `CAST`/`TRY_CAST` at *every* depth and hides the expression behind + /// an alias entirely. #[test] - fn test_push_projection_dedupl() { - let mut projection = vec![]; - push_projection_dedupl(&mut projection, col("a")); - push_projection_dedupl(&mut projection, col("a")); - push_projection_dedupl(&mut projection, col("t.a")); - push_projection_dedupl(&mut projection, col("x").alias("n")); + fn test_deduped_projection() { + let mut projection = DedupedProjection::default(); + projection.push(col("a")); + projection.push(col("a")); + projection.push(col("t.a")); + projection.push(col("x").alias("n")); // same alias name as the previous alias, different inner expr - push_projection_dedupl(&mut projection, col("y").alias("n")); - // cast renders as its input, which is already present - push_projection_dedupl(&mut projection, cast(col("a"), ArrowDataType::Int64)); - push_projection_dedupl(&mut projection, col("a").add(lit(1))); - push_projection_dedupl(&mut projection, col("a").add(lit(1))); - push_projection_dedupl(&mut projection, col("a").add(lit(2))); + projection.push(col("y").alias("n")); + // a cast renders as its input, which is already present + projection.push(cast(col("a"), ArrowDataType::Int64)); + projection.push(try_cast(col("a"), ArrowDataType::Int64)); + projection.push(col("a").add(lit(1))); + projection.push(col("a").add(lit(1))); + // casts are hidden at every depth, so these render as `a + Int32(1)` + // too and must not be pushed again + projection.push(cast(col("a"), ArrowDataType::Int64).add(lit(1))); + projection.push(col("a").add(cast(lit(1), ArrowDataType::Int64))); + projection.push(try_cast(col("a"), ArrowDataType::Int64).add(lit(1))); + projection.push(col("a").add(lit(2))); let names: Vec = projection + .exprs() .iter() .map(|e| e.schema_name().to_string()) .collect(); assert_eq!(names, vec!["a", "t.a", "n", "a + Int32(1)", "a + Int32(2)"]); } + /// The name a caller gets back from `push_returning_name` must be the + /// schema name of the pushed expression, whether or not it was retained, + /// since the outer projection references the inner field by that name. + #[test] + fn test_deduped_projection_push_returning_name() { + let mut projection = DedupedProjection::default(); + assert_eq!( + projection.push_returning_name(col("a").add(lit(1))), + "a + Int32(1)" + ); + // dropped as a duplicate, but the name is still the one to reference + assert_eq!( + projection + .push_returning_name(cast(col("a"), ArrowDataType::Int64).add(lit(1))), + "a + Int32(1)" + ); + assert_eq!(projection.exprs().len(), 1); + } + fn column_unnests_eq( l: Vec<&str>, r: &IndexMap>>, @@ -926,7 +886,7 @@ mod tests { }); let mut unnest_placeholder_columns = IndexMap::new(); - let mut inner_projection_exprs = vec![]; + let mut inner_projection_exprs = DedupedProjection::default(); // unnest(unnest(3d_col)) + unnest(unnest(3d_col)) let original_expr = unnest(unnest(col("3d_col"))) @@ -961,7 +921,7 @@ mod tests { // Still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( - inner_projection_exprs, + inner_projection_exprs.exprs().to_vec(), vec![ col("3d_col").alias("__unnest_placeholder(3d_col)"), col("i64_col") @@ -993,7 +953,7 @@ mod tests { // Still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( - inner_projection_exprs, + inner_projection_exprs.exprs().to_vec(), vec![ col("3d_col").alias("__unnest_placeholder(3d_col)"), col("i64_col") @@ -1033,7 +993,7 @@ mod tests { }); let mut unnest_placeholder_columns = IndexMap::new(); - let mut inner_projection_exprs = vec![]; + let mut inner_projection_exprs = DedupedProjection::default(); // unnest(struct_col) let original_expr = unnest(col("struct_col")); @@ -1057,7 +1017,7 @@ mod tests { // Still reference struct_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( - inner_projection_exprs, + inner_projection_exprs.exprs().to_vec(), vec![col("struct_col").alias("__unnest_placeholder(struct_col)"),] ); @@ -1090,7 +1050,7 @@ mod tests { // Still reference array_col in original schema but with alias, // to avoid colliding with the projection on the column itself if any assert_eq!( - inner_projection_exprs, + inner_projection_exprs.exprs().to_vec(), vec![ col("struct_col").alias("__unnest_placeholder(struct_col)"), col("array_col").alias("__unnest_placeholder(array_col)") @@ -1147,7 +1107,7 @@ mod tests { }); let mut unnest_placeholder_columns = IndexMap::new(); - let mut inner_projection_exprs = vec![]; + let mut inner_projection_exprs = DedupedProjection::default(); // An expr with multiple unnest let select_expr1 = unnest(unnest(col("struct_list")).field("subfield1")); @@ -1175,7 +1135,7 @@ mod tests { ); assert_eq!( - inner_projection_exprs, + inner_projection_exprs.exprs().to_vec(), vec![col("struct_list").alias("__unnest_placeholder(struct_list)")] ); @@ -1207,7 +1167,7 @@ mod tests { ); assert_eq!( - inner_projection_exprs, + inner_projection_exprs.exprs().to_vec(), vec![col("struct_list").alias("__unnest_placeholder(struct_list)")] ); diff --git a/datafusion/sqllogictest/test_files/distinct_on.slt b/datafusion/sqllogictest/test_files/distinct_on.slt index 0659b9c208f9c..f0dfb76f7e556 100644 --- a/datafusion/sqllogictest/test_files/distinct_on.slt +++ b/datafusion/sqllogictest/test_files/distinct_on.slt @@ -425,3 +425,25 @@ FROM t GROUP BY a, b ORDER BY a, total DESC; ---- x 1 30 y 2 30 + +# Regression: the inner projection built by the UNNEST rewrite is deduplicated +# by schema name, and `Expr::schema_name` hides casts at *every* depth. The +# DISTINCT ON key `1 + 1` and the ORDER BY key `CAST(1 AS BIGINT) + 1` share +# the same schema name, and both land in that one inner projection, so only one +# of them may be pushed or the projection gets duplicate field names. +query II +SELECT DISTINCT ON (1 + 1) + UNNEST([1, 2]) AS item, + COUNT(*) OVER () AS n +ORDER BY CAST(1 AS BIGINT) + 1; +---- +1 1 + +# Same, via TRY_CAST. +query II +SELECT DISTINCT ON (1 + 1) + UNNEST([1, 2]) AS item, + COUNT(*) OVER () AS n +ORDER BY TRY_CAST(1 AS BIGINT) + 1; +---- +1 1