fix: preserve nested projection expressions - #24686

Codecov Report

❌ Patch coverage is 69.38776% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.58%. Comparing base (17547e2) to head (e6f887b).
⚠️ Report is 3 commits behind head on main.

@@            Coverage Diff             @@
##             main   #24686      +/-   ##
==========================================
- Coverage   81.58%   81.58%   -0.01%     
==========================================
  Files        1123     1123              
  Lines      406743   406778      +35     
  Branches   406743   406778      +35     
==========================================
+ Hits       331840   331862      +22     
+ Misses      55464    55462       -2     
- Partials    19439    19454      +15     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

@discord9,

Thanks for working on this. I took another look at the latest version and re-verified the behavior against the current head. I think there are still a few issues we need to address before merging.

The main concern is that removing the iterative merge avoids the reported failure for the shallow case, but it does not address the underlying structural-equality fast path. That path can still remove a projection that needs to be evaluated. Removing the loop also introduces a projection-collapse regression for deeper chains.

I left the specific changes I think are needed inline. In short, I suggest keeping the iterative merge, guarding the equality fast path so it only elides projections that are safe to elide, and extending the regression coverage to deeper chains and the default optimizer pass count.

Thanks again for digging into this.

Thanks for the review, I think fa440be should address most of the review comment? basically just remove fast path and add some tests now

The iterative whole-chain merge is restored:

datafusion/datafusion/optimizer/src/optimize_projections/mod.rs

Lines 524 to 548 in fa440be

	fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Projection>> {
	// Collapse the whole chain in one pass; otherwise an N-deep chain needs
	// N outer optimizer passes to fully fold.
	let mut current = proj;
	let mut transformed_any = false;
	loop {
	let Transformed {
	data, transformed, ..
	} = merge_consecutive_projections_one_level(current)?;
	current = data;
	if !transformed {
	break;
	}
	transformed_any = true;
	}
	Ok(if transformed_any {
	Transformed::yes(current)
	} else {
	Transformed::no(current)
	})
	}
	fn merge_consecutive_projections_one_level(
	proj: Projection,
	) -> Result<Transformed<Projection>> {

The structural-equality fast path is removed entirely rather than guarded. Consecutive projections now proceed through the normal referral and rewrite logic:

datafusion/datafusion/optimizer/src/optimize_projections/mod.rs

Lines 555 to 568 in fa440be

	let LogicalPlan::Projection(prev_projection) = input.as_ref() else {
	return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no);
	};
	// Count usages (referrals) of each projection expression in its input fields:
	let mut column_referral_map = HashMap::<&Column, usize>::new();
	for expr in &expr {
	expr.add_column_ref_counts(&mut column_referral_map);
	}
	// If an expression is non-trivial (KeepInPlace) and appears more than once, do not merge
	// them as consecutive projections will benefit from a compute-once approach.
	// For details, see: https://github.com/apache/datafusion/issues/8296
	if column_referral_map.into_iter().any(|(col, usage)| {

A focused test verifies that two structurally equal, non-idempotent i + 1 AS i projections retain both evaluations:

datafusion/datafusion/optimizer/src/optimize_projections/mod.rs

Lines 1202 to 1218 in fa440be

	#[test]
	fn merge_structurally_equal_non_idempotent_projections() -> Result<()> {
	let schema = Schema::new(vec![Field::new("i", DataType::Int32, false)]);
	let projection = || col("i").add(lit(1)).alias("i");
	let plan = table_scan(TableReference::none(), &schema, None)?
	.project(vec![projection()])?
	.project(vec![projection()])?
	.build()?;
	assert_optimized_plan_equal!(
	plan,
	@r"
	Projection: ?table?.i + Int32(1) + Int32(1) AS i
	TableScan: ?table? projection=[i]
	"
	)
	}

A 12-level test runs with max_passes = 1 and asserts both that all 12 additions remain and that the chain collapses to one Projection:

datafusion/datafusion/optimizer/src/optimize_projections/mod.rs

Lines 1220 to 1242 in fa440be

	#[test]
	fn merge_deep_projection_chain_in_one_pass() -> Result<()> {
	let schema = Schema::new(vec![Field::new("i", DataType::Int32, false)]);
	let mut plan = table_scan(TableReference::none(), &schema, None)?
	.project(vec![col("i").add(lit(1)).alias("i")])?
	.build()?;
	for _ in 1..12 {
	plan = LogicalPlanBuilder::from(plan)
	.project(vec![col("i").add(lit(1)).alias("i")])?
	.build()?;
	}
	let optimizer = Optimizer::with_rules(vec![Arc::new(OptimizeProjections::new())]);
	let optimized = optimizer.optimize(
	plan,
	&OptimizerContext::new().with_max_passes(1),
	observe,
	)?;
	let plan_string = format!("{optimized}");
	assert_eq!(12, plan_string.matches("Int32(1)").count());
	assert_eq!(1, plan_string.matches("Projection:").count());
	Ok(())
	}

Metadata-bearing Alias(Column) is covered and its field metadata is asserted to survive the normal merge path:

datafusion/datafusion/optimizer/src/optimize_projections/mod.rs

Lines 1244 to 1263 in fa440be

	#[test]
	fn merge_columns_and_metadata_alias() -> Result<()> {
	let metadata =
	datafusion_common::metadata::FieldMetadata::from(HashMap::from([(
	"key".to_string(),
	"value".to_string(),
	)]));
	let plan = LogicalPlanBuilder::from(test_table_scan()?)
	.project(vec![col("a")])?
	.project(vec![col("a").alias_with_metadata("a", Some(metadata))])?
	.build()?;
	let optimized = optimize(plan)?;
	assert_eq!(
	"value",
	optimized.schema().field(0).metadata().get("key").unwrap()
	);
	assert_eq!(1, format!("{optimized}").matches("Projection:").count());
	Ok(())
	}

The SQL regression now has six anonymous i + 1 AS i layers and checks both max_passes = 1 and the default optimizer configuration:

datafusion/datafusion/sqllogictest/test_files/projection.slt

Lines 220 to 286 in fa440be

	# Regression: one optimizer pass must preserve anonymous nested projections.
	statement ok
	CREATE TABLE nested_projection(i INT);
	statement ok
	INSERT INTO nested_projection VALUES (3), (4), (5);
	statement ok
	SET datafusion.optimizer.max_passes = 1;
	query I rowsort
	SELECT i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM nested_projection
	)
	)
	)
	)
	)
	)
	----
	10
	11
	9
	statement ok
	RESET datafusion.optimizer.max_passes;
	query I rowsort
	SELECT i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM (
	SELECT i + 1 AS i
	FROM nested_projection
	)
	)
	)
	)
	)
	)
	----
	10
	11
	9
	statement ok
	DROP TABLE nested_projection;

The extra pass-through Projection: t.a, t.ts was reverted; the integration expectation is back to the tighter plan from main.

🚀
@discord9
Thank you for your contribution.