Copilot commented on code in PR #25647:
URL: https://github.com/apache/datafusion/pull/25647#discussion_r4082001522
##########
datafusion/optimizer/src/optimize_projections/mod.rs:
##########
@@ -353,6 +353,7 @@ fn optimize_projections(
.zip(necessary_children_indices)
.map(|(child, necessary_indices)| {
RequiredIndices::new_from_indices(necessary_indices)
+ .with_projection_beneficial()
.with_plan_exprs(&plan, child.schema())
Review Comment:
`with_projection_beneficial()` is applied unconditionally for all extension
children. This may cause extra `Projection` nodes to be inserted even when (a)
the parent context did *not* consider projections beneficial, or (b)
`necessary_indices` already covers all columns for that child, creating
avoidable plan nodes and work. Consider conditioning this flag (e.g., only set
it when the parent `RequiredIndices` indicates projections are beneficial,
and/or when `necessary_indices` is a true subset of the child’s columns). This
keeps the fix while reducing the risk of widespread plan-shape changes and
unnecessary projections.
##########
datafusion/optimizer/src/optimize_projections/mod.rs:
##########
@@ -1724,6 +1725,119 @@ mod tests {
)
}
+ #[test]
+ fn test_user_defined_logical_plan_above_join() -> Result<()> {
+ let table_scan = test_table_scan()?;
+ let schema = Schema::new(vec![Field::new("c1", DataType::UInt32,
false)]);
+ let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
+
+ // Join test(a, b, c) and test2(c1) on test.a = test2.c1
+ let join_plan = LogicalPlanBuilder::from(table_scan)
+ .join(table2_scan, JoinType::Left, (vec!["a"], vec!["c1"]), None)?
+ .build()?;
+
+ let custom_plan = LogicalPlan::Extension(Extension {
+ node: Arc::new(NoOpUserDefined::new(
+ Arc::clone(join_plan.schema()),
+ Arc::new(join_plan),
+ )),
+ });
+
Review Comment:
These three tests repeat substantial setup (creating scans, building the
same left join, wrapping it in an extension, then projecting). Consider
extracting small helper(s) within the test module (e.g.,
`build_left_join_test_plan()` / `wrap_noop_extension(plan)`) to reduce
duplication and make future additions easier and less error-prone.
##########
datafusion/optimizer/src/optimize_projections/mod.rs:
##########
@@ -1724,6 +1725,119 @@ mod tests {
)
}
+ #[test]
+ fn test_user_defined_logical_plan_above_join() -> Result<()> {
Review Comment:
The added tests cover the positive case where a `Projection` should be
inserted beneath an `Extension` (above a `Join`). To guard against regressions
from unconditionally enabling `projection_beneficial`, add a negative test
asserting no extra `Projection` is inserted when the parent requires *all*
columns from an extension child (or when `necessary_indices` equals the full
set). This will ensure the change doesn’t unintentionally increase projection
churn in common cases.
##########
datafusion/optimizer/src/optimize_projections/mod.rs:
##########
@@ -1724,6 +1725,119 @@ mod tests {
)
}
+ #[test]
+ fn test_user_defined_logical_plan_above_join() -> Result<()> {
+ let table_scan = test_table_scan()?;
+ let schema = Schema::new(vec![Field::new("c1", DataType::UInt32,
false)]);
+ let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
+
+ // Join test(a, b, c) and test2(c1) on test.a = test2.c1
+ let join_plan = LogicalPlanBuilder::from(table_scan)
+ .join(table2_scan, JoinType::Left, (vec!["a"], vec!["c1"]), None)?
+ .build()?;
+
+ let custom_plan = LogicalPlan::Extension(Extension {
+ node: Arc::new(NoOpUserDefined::new(
+ Arc::clone(join_plan.schema()),
+ Arc::new(join_plan),
+ )),
+ });
+
+ // Parent only requires test.a and test.b; join key test2.c1 is not
needed downstream.
+ let plan = LogicalPlanBuilder::from(custom_plan)
+ .project(vec![col("test.a"), col("test.b")])?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Projection: test.a, test.b
+ NoOpUserDefined
+ Projection: test.a, test.b
+ Left Join: test.a = test2.c1
+ TableScan: test projection=[a, b]
+ TableScan: test2 projection=[c1]
+ "
+ )
+ }
+
+ #[test]
+ fn test_user_defined_logical_plan_above_join_with_sort() -> Result<()> {
Review Comment:
These three tests repeat substantial setup (creating scans, building the
same left join, wrapping it in an extension, then projecting). Consider
extracting small helper(s) within the test module (e.g.,
`build_left_join_test_plan()` / `wrap_noop_extension(plan)`) to reduce
duplication and make future additions easier and less error-prone.
##########
datafusion/optimizer/src/optimize_projections/mod.rs:
##########
@@ -1724,6 +1725,119 @@ mod tests {
)
}
+ #[test]
+ fn test_user_defined_logical_plan_above_join() -> Result<()> {
+ let table_scan = test_table_scan()?;
+ let schema = Schema::new(vec![Field::new("c1", DataType::UInt32,
false)]);
+ let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
+
+ // Join test(a, b, c) and test2(c1) on test.a = test2.c1
+ let join_plan = LogicalPlanBuilder::from(table_scan)
+ .join(table2_scan, JoinType::Left, (vec!["a"], vec!["c1"]), None)?
+ .build()?;
+
+ let custom_plan = LogicalPlan::Extension(Extension {
+ node: Arc::new(NoOpUserDefined::new(
+ Arc::clone(join_plan.schema()),
+ Arc::new(join_plan),
+ )),
+ });
+
+ // Parent only requires test.a and test.b; join key test2.c1 is not
needed downstream.
+ let plan = LogicalPlanBuilder::from(custom_plan)
+ .project(vec![col("test.a"), col("test.b")])?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Projection: test.a, test.b
+ NoOpUserDefined
+ Projection: test.a, test.b
+ Left Join: test.a = test2.c1
+ TableScan: test projection=[a, b]
+ TableScan: test2 projection=[c1]
+ "
+ )
+ }
+
+ #[test]
+ fn test_user_defined_logical_plan_above_join_with_sort() -> Result<()> {
+ let table_scan = test_table_scan()?;
+ let schema = Schema::new(vec![Field::new("c1", DataType::UInt32,
false)]);
+ let table2_scan = scan_empty(Some("test2"), &schema, None)?.build()?;
+
+ // Join test(a, b, c) and test2(c1) on test.a = test2.c1
+ let join_plan = LogicalPlanBuilder::from(table_scan)
+ .join(table2_scan, JoinType::Left, (vec!["a"], vec!["c1"]), None)?
+ .build()?;
+
+ let custom_plan = LogicalPlan::Extension(Extension {
+ node: Arc::new(NoOpUserDefined::new(
+ Arc::clone(join_plan.schema()),
+ Arc::new(join_plan),
+ )),
+ });
+
+ // Sort -> Projection -> Extension -> Join
+ let plan = LogicalPlanBuilder::from(custom_plan)
+ .project(vec![col("test.a"), col("test.b")])?
+ .sort(vec![col("test.a").sort(true, false)])?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: test.a ASC NULLS LAST
+ Projection: test.a, test.b
+ NoOpUserDefined
+ Projection: test.a, test.b
+ Left Join: test.a = test2.c1
+ TableScan: test projection=[a, b]
+ TableScan: test2 projection=[c1]
+ "
+ )
+ }
+
+ #[test]
+ fn test_user_defined_logical_plan_multi_input_above_join() -> Result<()> {
Review Comment:
These three tests repeat substantial setup (creating scans, building the
same left join, wrapping it in an extension, then projecting). Consider
extracting small helper(s) within the test module (e.g.,
`build_left_join_test_plan()` / `wrap_noop_extension(plan)`) to reduce
duplication and make future additions easier and less error-prone.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
