From bb7a6e3e9f45d438b463584f1822b71cd53f7431 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Fri, 6 Nov 2026 21:26:21 +0000 Subject: [PATCH] fix: preserve constant equivalence properties through projection (#27123) ## Which issue does this PR close? Closes #26624. ## Rationale for this change With `prefer_hash_join = false`, a query can fail physical plan validation when a filter makes a computed sort or join key constant. Embedding a projection into the sort-merge join loses the constant property after sorting has already relied on it. ## What changes are included in this PR? Preserve each projected equivalence class's constant metadata. A projected member has the same value as its source member on every row, and projection keeps rows in the same partitions. Its known value and uniform or partition-local status therefore still hold. Outer joins clear constants on the NULL-padded side before projecting their output. Partition merges clear `Heterogeneous` constants. Also preserve scalar-subquery leaves during indirect projection. They read a shared subquery result without depending on the input schema. The arithmetic `ORDER BY` reproduction needs the metadata copy; the scalar-subquery join reproduction needs both changes. ## What is the testing strategy for this PR? SQL cases in `joins.slt` cover arithmetic, `upper()`, left joins, and the scalar-subquery join. A NULL-padded left join checks that matched values sort before the NULL rows. The metadata unit test covers known, unknown, and partition-local constants, dropped arguments, and partition merging. The original SQL cases and metadata test fail before the fix and pass afterward. Validation includes the extended workspace tests, full lint suite, and strict all-features Clippy. ## Are there any user-facing changes? Affected queries plan and execute successfully. No public API changes. Needs a backport to `branch-56` for the 56.0.0 regression tracker (#24461). --------- Co-authored-by: Andrew Lamb Co-authored-by: Claude Fable 5.1 (cherry picked from commit ad8578a9df4593265b416cace5d1d2a78be1cdd3) --- .../physical-expr/src/equivalence/class.rs | 54 +++++++++- .../sqllogictest/test_files/aggregate.slt | 4 +- datafusion/sqllogictest/test_files/joins.slt | 100 ++++++++++++++++++ 3 files changed, 154 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index 352d690a54b82..a94a38f47f75d 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -12,6 +12,7 @@ use super::ProjectionMapping; use crate::expressions::{Column, Literal}; use crate::physical_expr::add_offset_to_expr; use crate::projection::ProjectionTargets; +use crate::scalar_subquery::ScalarSubqueryExpr; use crate::{PhysicalExpr, PhysicalExprRef, PhysicalSortExpr, PhysicalSortRequirement}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; @@ -631,8 +632,8 @@ impl EquivalenceGroup { aug_mapping: &AugmentedMapping, expr: &Arc, ) -> Option> { - // Literals don't need to be projected - if expr.downcast_ref::().is_some() { + // These leaves do not read columns from the projected input schema. + if expr.is::() || expr.is::() { return Some(Arc::clone(expr)); } @@ -742,7 +743,11 @@ impl EquivalenceGroup { pub fn project(&self, mapping: &ProjectionMapping) -> Self { let projected_classes = self.iter().map(|cls| { let new_exprs = self.project_expressions(mapping, cls.iter()); - EquivalenceClass::new(new_exprs.flatten()) + let mut projected = EquivalenceClass::new(new_exprs.flatten()); + // Projection changes column references, not constant values or how + // they are distributed across partitions. + projected.constant = projected.constant.or_else(|| cls.constant.clone()); + projected }); // The key is the source expression, and the value is the equivalence @@ -1334,6 +1339,49 @@ mod tests { Ok(()) } + #[test] + fn test_project_preserves_constant_metadata() -> Result<()> { + let schema = abc_schema(); + let expr = binary(col("b", &schema)?, Operator::Divide, lit(2i32), &schema)?; + let mapping = ProjectionMapping::from_indices(&[1], &schema)?; + let dropped_mapping = ProjectionMapping::from_indices(&[0], &schema)?; + let projected_b: Arc = Arc::new(Column::new("b", 0)); + let projected_expr: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&projected_b), + Operator::Divide, + lit(2i32), + )); + + for across in [ + AcrossPartitions::Uniform(Some(ScalarValue::Int32(Some(2)))), + AcrossPartitions::Uniform(None), + AcrossPartitions::Heterogeneous, + ] { + let mut group = EquivalenceGroup::default(); + // The class has no literal member, so `EquivalenceClass::new` cannot + // restore the constant. + group.add_constant(ConstExpr::new(Arc::clone(&expr), across.clone())); + + let mut projected = group.project(&mapping); + assert_eq!( + projected.is_expr_constant(&projected_expr), + Some(across.clone()) + ); + // A constant expression does not imply that its argument is constant. + assert_eq!(projected.is_expr_constant(&projected_b), None); + assert!(group.project(&dropped_mapping).is_empty()); + + let heterogeneous = across == AcrossPartitions::Heterogeneous; + assert_eq!(projected.clear_per_partition_constants(), heterogeneous); + assert_eq!( + projected.is_expr_constant(&projected_expr), + (!heterogeneous).then_some(across) + ); + } + + Ok(()) + } + #[test] fn test_project_constant_of_equivalent_source() -> Result<()> { // - columns: [a, b, c]. diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 605e54d1543cb..1e12ee4194a49 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -3093,7 +3093,9 @@ SELECT g, percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, 0.25 AS 2 26 3 27 -query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values 0.5 and 0.9 +# The partial states are merged in whichever order the partitions finish, +# so the two differing values can be reported in either order +query error DataFusion error: Execution error: Percentile value for 'PERCENTILE_CONT' must be constant across the aggregation, found differing values (0\.5 and 0\.9|0\.9 and 0\.5) SELECT g, percentile_cont(x, m) FROM (SELECT value % 4 AS g, value AS x, CASE WHEN value <= 50 THEN 0.5 ELSE 0.9 END AS m FROM generate_series(1, 100)) GROUP BY g query IR diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index b795b050f919f..d2be4f690022b 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -6343,3 +6343,103 @@ DROP TABLE projection_networks; statement ok set datafusion.execution.target_partitions = 4; + +# A computed constant must survive the projection embedded into a sort-merge +# join, even when only the expression's argument is projected. +# https://github.com/apache/datafusion/issues/26624 +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +statement ok +CREATE TABLE constant_projection (a INT, b INT, s VARCHAR) AS VALUES +(1, 1, 'x'), (2, 0, 'X'), (3, 1, 'y'), (4, NULL, NULL), (5, 1, 'X'); + +query II +SELECT t1.b + 1 AS x, t2.a +FROM constant_projection t1 JOIN constant_projection t2 ON t1.a = t2.a +WHERE t1.b + 1 = 2 +ORDER BY x, t2.a; +---- +2 1 +2 3 +2 5 + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT t1.b + 1 AS x, t2.a +FROM constant_projection t1 JOIN constant_projection t2 ON t1.a = t2.a +WHERE t1.b + 1 = 2 +ORDER BY x, t2.a; +---- +physical_plan +01)SortPreservingMergeExec: [x@0 ASC NULLS LAST, a@1 ASC NULLS LAST] +02)--SortExec: expr=[a@1 ASC NULLS LAST], preserve_partitioning=[true] +03)----ProjectionExec: expr=[CAST(b@0 AS Int64) + 1 as x, a@1 as a] +04)------SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)], projection=[b@1, a@2] +05)--------SortExec: expr=[a@0 ASC], preserve_partitioning=[true] +06)----------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1 +07)------------FilterExec: CAST(b@1 AS Int64) + 1 = 2 +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] +09)--------SortExec: expr=[a@0 ASC], preserve_partitioning=[true] +10)----------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1 +11)------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +reset datafusion.explain.physical_plan_only; + +query II +SELECT t1.b + 1 AS x, t2.a +FROM constant_projection t1 LEFT JOIN constant_projection t2 ON t1.a = t2.a +WHERE t1.b + 1 = 2 +ORDER BY x, t2.a; +---- +2 1 +2 3 +2 5 + +# NULL padding invalidates the right side's constant, so x still needs sorting. +query II +SELECT t2.b + 1 AS x, t1.a +FROM constant_projection t1 +LEFT JOIN (SELECT * FROM constant_projection WHERE b + 1 = 2) t2 ON t1.a = t2.a +ORDER BY x, t1.a; +---- +2 1 +2 3 +2 5 +NULL 2 +NULL 4 + +query TI +SELECT upper(t1.s) AS x, t2.a +FROM constant_projection t1 JOIN constant_projection t2 ON t1.a = t2.a +WHERE upper(t1.s) = 'X' +ORDER BY x, t2.a; +---- +X 1 +X 2 +X 5 + +statement ok +DROP TABLE constant_projection; + +# The expression must also survive when it contains a scalar-subquery leaf. +statement ok +CREATE TABLE subquery_projection (a INT, b INT) AS VALUES (1, 1), (2, 0); + +query III +SELECT * FROM ( + SELECT t1.b + (SELECT max(a) FROM subquery_projection) AS x + FROM subquery_projection t1 JOIN subquery_projection t2 ON t1.a = t2.a +) s JOIN subquery_projection t3 ON s.x = t3.a +WHERE s.x = 2; +---- +2 2 0 + +statement ok +DROP TABLE subquery_projection; + +statement ok +reset datafusion.optimizer.prefer_hash_join;