You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/14 04:50:57 UTC
(datafusion) 01/01: Reduce repetitive string formatting in push_projection_dedupl (#25236)
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-25236-38d58ed09db21afe0bc224607457d86d4ed62e03
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit b30cfd7d1fb4b5a9c097a26c682eba7bb1f79aa2
Author: Wyatt Herkamp <wh...@gmail.com>
AuthorDate: Mon Sep 14 04:50:34 2026 +0000
Reduce repetitive string formatting in push_projection_dedupl (#25236)
## Which issue does this PR close?
- Closes #25235.
## Rationale for this change
Calling schema_name to string is expensive
## What changes are included in this PR?
Uses a hashset of names to prevent reformatting the name
```
logical_unnest_plus_200_columns
time: [1.0161 ms 1.0183 ms 1.0205 ms]
change: [−79.252% −79.058% −78.900%] (p = 0.00 < 0.05)
Performance has improved.
Found 5 outliers among 100 measurements (5.00%)
4 (4.00%) high mild
1 (1.00%) high severe
logical_unnest_plus_1000_columns
time: [14.751 ms 14.799 ms 14.849 ms]
change: [−86.220% −86.145% −86.079%] (p = 0.00 < 0.05)
Performance has improved.
Found 9 outliers among 100 measurements (9.00%)
1 (1.00%) low mild
6 (6.00%) high mild
2 (2.00%) high severe
logical_unnest_in_200_exprs
time: [3.1697 ms 3.1815 ms 3.1953 ms]
change: [−71.260% −71.088% −70.910%] (p = 0.00 < 0.05)
Performance has improved.
Found 6 outliers among 100 measurements (6.00%)
3 (3.00%) high mild
3 (3.00%) high severe
```
## What is the testing strategy for this PR?
- Added benchmarks
- Added tests directed to `DedupedProjection`
- Added a sql logic test for a case @xudong963 identified in my initial
fix
## Are there any user-facing changes?
None
---
datafusion/core/benches/sql_planner_extended.rs | 59 +++++++-
datafusion/sql/src/select.rs | 16 +--
datafusion/sql/src/utils.rs | 154 ++++++++++++++++-----
datafusion/sqllogictest/test_files/distinct_on.slt | 22 +++
4 files changed, 207 insertions(+), 44 deletions(-)
diff --git a/datafusion/core/benches/sql_planner_extended.rs b/datafusion/core/benches/sql_planner_extended.rs
index 5bea9860c4..96a424e184 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<Int32>` column, used to benchmark `UNNEST` planning.
+fn register_list_table(ctx: &SessionContext, name: &str, num_columns: usize) {
+ let mut fields: Vec<Field> = (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::<Vec<_>>()
+ .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::<Vec<_>>()
+ .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 86916c5693..79ab852d3d 100644
--- a/datafusion/sql/src/select.rs
+++ b/datafusion/sql/src/select.rs
@@ -24,10 +24,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;
@@ -662,7 +662,7 @@ impl<S: ContextProvider> 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());
@@ -727,7 +727,7 @@ impl<S: ContextProvider> 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;
@@ -830,7 +830,7 @@ impl<S: ContextProvider> 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,
@@ -865,7 +865,7 @@ impl<S: ContextProvider> 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 0f0b5fe44c..882d6fa681 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<String> {
pub(crate) fn rewrite_recursive_unnests_bottom_up(
input: &LogicalPlan,
unnest_placeholder_columns: &mut IndexMap<Column, Option<Vec<ColumnUnnestList>>>,
- inner_projection_exprs: &mut Vec<Expr>,
+ inner_projection_exprs: &mut DedupedProjection,
original_exprs: &[Expr],
) -> Result<Vec<Expr>> {
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<Unnest>,
consecutive_unnest: Vec<Option<Unnest>>,
- inner_projection_exprs: &'a mut Vec<Expr>,
+ inner_projection_exprs: &'a mut DedupedProjection,
columns_unnestings: &'a mut IndexMap<Column, Option<Vec<ColumnUnnestList>>>,
transformed_root_exprs: Option<Vec<Expr>>,
}
@@ -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
@@ -655,20 +651,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<top most unnest>
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: 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<Expr>,
+ names: HashSet<String>,
+}
+
+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<Expr> {
+ self.exprs
}
}
/// The context is we want to rewrite unnest() into InnerProjection->Unnest->OuterProjection
@@ -683,7 +716,7 @@ fn push_projection_dedupl(projection: &mut Vec<Expr>, expr: Expr) {
pub(crate) fn rewrite_recursive_unnest_bottom_up(
input: &LogicalPlan,
unnest_placeholder_columns: &mut IndexMap<Column, Option<Vec<ColumnUnnestList>>>,
- inner_projection_exprs: &mut Vec<Expr>,
+ inner_projection_exprs: &mut DedupedProjection,
original_expr: &Expr,
) -> Result<Vec<Expr>> {
let mut rewriter = RecursiveUnnestRewriter {
@@ -717,13 +750,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 {
@@ -741,14 +774,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<String> = 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<Column, Option<Vec<ColumnUnnestList>>>,
@@ -798,7 +884,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")))
@@ -833,7 +919,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")
@@ -865,7 +951,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")
@@ -905,7 +991,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"));
@@ -929,7 +1015,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)"),]
);
@@ -962,7 +1048,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)")
@@ -1019,7 +1105,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"));
@@ -1047,7 +1133,7 @@ mod tests {
);
assert_eq!(
- inner_projection_exprs,
+ inner_projection_exprs.exprs().to_vec(),
vec![col("struct_list").alias("__unnest_placeholder(struct_list)")]
);
@@ -1079,7 +1165,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 d54a310918..d477d30b78 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
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org