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/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 3b571eed279dd..fe579bc45f96c 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -26,7 +26,7 @@ use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, }; use datafusion_common::{ - Column, DFSchemaRef, Diagnostic, HashMap, Result, ScalarValue, + 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; @@ -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,20 +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)) } } -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) - { - projection.push(expr); +/// 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 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 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); + } + } + + /// 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 + } + + pub(crate) fn into_exprs(self) -> Vec { + self.exprs } } /// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection @@ -685,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 { @@ -719,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 { @@ -743,14 +776,67 @@ 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, 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::{ + DedupedProjection, resolve_positions_to_exprs, rewrite_recursive_unnest_bottom_up, + }; use indexmap::IndexMap; + /// [`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_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 + 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>>, @@ -800,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"))) @@ -835,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") @@ -867,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") @@ -907,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")); @@ -931,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)"),] ); @@ -964,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)") @@ -1021,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")); @@ -1049,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)")] ); @@ -1081,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