Use a struct for ProjectionExpr (#17398)

Author: adriangb

datafusion/core/src/physical_planner.rs

@@ -45,7 +45,7 @@ use crate::physical_plan::joins::{
     CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec,
 };
 use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
-use crate::physical_plan::projection::ProjectionExec;
+use crate::physical_plan::projection::{ProjectionExec, ProjectionExpr};
 use crate::physical_plan::repartition::RepartitionExec;
 use crate::physical_plan::sorts::sort::SortExec;
 use crate::physical_plan::union::UnionExec;
@@ -2185,17 +2185,25 @@ impl DefaultPhysicalPlanner {
             PlannedExprResult::ExprWithName(physical_exprs),
             input_physical_schema.as_ref(),
         )? {
-            PlanAsyncExpr::Sync(PlannedExprResult::ExprWithName(physical_exprs)) => Ok(
-                Arc::new(ProjectionExec::try_new(physical_exprs, input_exec)?),
-            ),
+            PlanAsyncExpr::Sync(PlannedExprResult::ExprWithName(physical_exprs)) => {
+                let proj_exprs: Vec<ProjectionExpr> = physical_exprs
+                    .into_iter()
+                    .map(|(expr, alias)| ProjectionExpr { expr, alias })
+                    .collect();
+                Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input_exec)?))
+            }
             PlanAsyncExpr::Async(
                 async_map,
                 PlannedExprResult::ExprWithName(physical_exprs),
             ) => {
                 let async_exec =
                     AsyncFuncExec::try_new(async_map.async_exprs, input_exec)?;
+                let proj_exprs: Vec<ProjectionExpr> = physical_exprs
+                    .into_iter()
+                    .map(|(expr, alias)| ProjectionExpr { expr, alias })
+                    .collect();
                 let new_proj_exec =
-                    ProjectionExec::try_new(physical_exprs, Arc::new(async_exec))?;
+                    ProjectionExec::try_new(proj_exprs, Arc::new(async_exec))?;
                 Ok(Arc::new(new_proj_exec))
             }
             _ => internal_err!("Unexpected PlanAsyncExpressions variant"),
@@ -2700,7 +2708,7 @@ mod tests {
         let execution_plan = plan(&logical_plan).await?;
         // verify that the plan correctly adds cast from Int64(1) to Utf8, and the const will be evaluated.
 
-        let expected = "expr: [(BinaryExpr { left: BinaryExpr { left: Column { name: \"c1\", index: 0 }, op: Eq, right: Literal { value: Utf8(\"a\"), field: Field { name: \"lit\", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, fail_on_overflow: false }, op: Or, right: BinaryExpr { left: Column { name: \"c1\", index: 0 }, op: Eq, right: Literal { value: Utf8(\"1\"), field: Field { name: \"lit\", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, fail_on_overflow: false }, fail_on_overflow: false }";
+        let expected = "expr: [ProjectionExpr { expr: BinaryExpr { left: BinaryExpr { left: Column { name: \"c1\", index: 0 }, op: Eq, right: Literal { value: Utf8(\"a\"), field: Field { name: \"lit\", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, fail_on_overflow: false }, op: Or, right: BinaryExpr { left: Column { name: \"c1\", index: 0 }, op: Eq, right: Literal { value: Utf8(\"1\"), field: Field { name: \"lit\", data_type: Utf8, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, fail_on_overflow: false }, fail_on_overflow: false }";
 
         let actual = format!("{execution_plan:?}");
         assert!(actual.contains(expected), "{}", actual);

datafusion/core/tests/execution/coop.rs

@@ -46,7 +46,7 @@ use datafusion_physical_plan::coop::make_cooperative;
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec};
 use datafusion_physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec};
-use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
@@ -651,7 +651,7 @@ async fn join_agg_yields(
     // Project only one column (“value” from the left side) because we just want to sum that
     let input_schema = join.schema();
 
-    let proj_expr = vec![(
+    let proj_expr = vec![ProjectionExpr::new(
         Arc::new(Column::new_with_schema("value", &input_schema)?) as _,
         "value".to_string(),
     )];

datafusion/core/tests/physical_optimizer/enforce_distribution.rs

@@ -62,7 +62,7 @@ use datafusion_physical_plan::expressions::col;
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::joins::utils::JoinOn;
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
-use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::union::UnionExec;
 use datafusion_physical_plan::{
@@ -243,7 +243,10 @@ fn projection_exec_with_alias(
 ) -> Arc<dyn ExecutionPlan> {
     let mut exprs = vec![];
     for (column, alias) in alias_pairs.iter() {
-        exprs.push((col(column, &input.schema()).unwrap(), alias.to_string()));
+        exprs.push(ProjectionExpr {
+            expr: col(column, &input.schema()).unwrap(),
+            alias: alias.to_string(),
+        });
     }
     Arc::new(ProjectionExec::try_new(exprs, input).unwrap())
 }
@@ -2207,14 +2210,14 @@ fn repartition_does_not_destroy_sort_more_complex() -> Result<()> {
 #[test]
 fn repartition_transitively_with_projection() -> Result<()> {
     let schema = schema();
-    let proj_exprs = vec![(
-        Arc::new(BinaryExpr::new(
+    let proj_exprs = vec![ProjectionExpr {
+        expr: Arc::new(BinaryExpr::new(
             col("a", &schema)?,
             Operator::Plus,
             col("b", &schema)?,
         )) as _,
-        "sum".to_string(),
-    )];
+        alias: "sum".to_string(),
+    }];
     // non sorted input
     let proj = Arc::new(ProjectionExec::try_new(proj_exprs, parquet_exec())?);
     let sort_key = [PhysicalSortExpr {

datafusion/core/tests/physical_optimizer/join_selection.rs

@@ -235,12 +235,12 @@ async fn test_join_with_swap() {
         .expect("A proj is required to swap columns back to their original order");
 
     assert_eq!(swapping_projection.expr().len(), 2);
-    let (col, name) = &swapping_projection.expr()[0];
-    assert_eq!(name, "big_col");
-    assert_col_expr(col, "big_col", 1);
-    let (col, name) = &swapping_projection.expr()[1];
-    assert_eq!(name, "small_col");
-    assert_col_expr(col, "small_col", 0);
+    let proj_expr = &swapping_projection.expr()[0];
+    assert_eq!(proj_expr.alias, "big_col");
+    assert_col_expr(&proj_expr.expr, "big_col", 1);
+    let proj_expr = &swapping_projection.expr()[1];
+    assert_eq!(proj_expr.alias, "small_col");
+    assert_col_expr(&proj_expr.expr, "small_col", 0);
 
     let swapped_join = swapping_projection
         .input()
@@ -526,12 +526,12 @@ async fn test_nl_join_with_swap(join_type: JoinType) {
         .expect("A proj is required to swap columns back to their original order");
 
     assert_eq!(swapping_projection.expr().len(), 2);
-    let (col, name) = &swapping_projection.expr()[0];
-    assert_eq!(name, "big_col");
-    assert_col_expr(col, "big_col", 1);
-    let (col, name) = &swapping_projection.expr()[1];
-    assert_eq!(name, "small_col");
-    assert_col_expr(col, "small_col", 0);
+    let proj_expr = &swapping_projection.expr()[0];
+    assert_eq!(proj_expr.alias, "big_col");
+    assert_col_expr(&proj_expr.expr, "big_col", 1);
+    let proj_expr = &swapping_projection.expr()[1];
+    assert_eq!(proj_expr.alias, "small_col");
+    assert_col_expr(&proj_expr.expr, "small_col", 0);
 
     let swapped_join = swapping_projection
         .input()

datafusion/core/tests/physical_optimizer/partition_statistics.rs

@@ -47,7 +47,7 @@ mod test {
     use datafusion_physical_plan::joins::CrossJoinExec;
     use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
     use datafusion_physical_plan::placeholder_row::PlaceholderRowExec;
-    use datafusion_physical_plan::projection::ProjectionExec;
+    use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
     use datafusion_physical_plan::repartition::RepartitionExec;
     use datafusion_physical_plan::sorts::sort::SortExec;
     use datafusion_physical_plan::union::UnionExec;
@@ -235,8 +235,10 @@ mod test {
     async fn test_statistics_by_partition_of_projection() -> Result<()> {
         let scan = create_scan_exec_with_statistics(None, Some(2)).await;
         // Add projection execution plan
-        let exprs: Vec<(Arc<dyn PhysicalExpr>, String)> =
-            vec![(Arc::new(Column::new("id", 0)), "id".to_string())];
+        let exprs = vec![ProjectionExpr {
+            expr: Arc::new(Column::new("id", 0)) as Arc<dyn PhysicalExpr>,
+            alias: "id".to_string(),
+        }];
         let projection: Arc<dyn ExecutionPlan> =
             Arc::new(ProjectionExec::try_new(exprs, scan)?);
         let statistics = (0..projection.output_partitioning().partition_count())