Add Projection struct w/ helper methods to manipulate projections (#18176)

I'm hoping these will help with
#14993. For now they're just
an internal refactor + exposing some functionality as `pub` and lots of
new tests.

Written with AI assistance.

Author: adriangb

datafusion/core/src/physical_planner.rs

@@ -3063,7 +3063,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: [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 expected = "exprs: [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 }";
 
         assert_contains!(format!("{execution_plan:?}"), expected);
 

datafusion/physical-expr/src/equivalence/class.rs

@@ -20,10 +20,10 @@ use std::ops::Deref;
 use std::sync::Arc;
 use std::vec::IntoIter;
 
-use super::projection::ProjectionTargets;
 use super::ProjectionMapping;
 use crate::expressions::Literal;
 use crate::physical_expr::add_offset_to_expr;
+use crate::projection::ProjectionTargets;
 use crate::{PhysicalExpr, PhysicalExprRef, PhysicalSortExpr, PhysicalSortRequirement};
 
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};

datafusion/physical-expr/src/equivalence/mod.rs

@@ -25,12 +25,13 @@ use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 
 mod class;
 mod ordering;
-mod projection;
 mod properties;
 
 pub use class::{AcrossPartitions, ConstExpr, EquivalenceClass, EquivalenceGroup};
 pub use ordering::OrderingEquivalenceClass;
-pub use projection::{project_ordering, project_orderings, ProjectionMapping};
+// Re-export for backwards compatibility, we recommend importing from
+// datafusion_physical_expr::projection instead
+pub use crate::projection::{project_ordering, project_orderings, ProjectionMapping};
 pub use properties::{
     calculate_union, join_equivalence_properties, EquivalenceProperties,
 };
@@ -61,7 +62,7 @@ mod tests {
 
     use arrow::compute::SortOptions;
     use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-    use datafusion_common::{plan_err, Result};
+    use datafusion_common::Result;
     use datafusion_physical_expr_common::sort_expr::PhysicalSortRequirement;
 
     /// Converts a string to a physical sort expression
@@ -95,31 +96,6 @@ mod tests {
         sort_expr
     }
 
-    pub fn output_schema(
-        mapping: &ProjectionMapping,
-        input_schema: &Arc<Schema>,
-    ) -> Result<SchemaRef> {
-        // Calculate output schema:
-        let mut fields = vec![];
-        for (source, targets) in mapping.iter() {
-            let data_type = source.data_type(input_schema)?;
-            let nullable = source.nullable(input_schema)?;
-            for (target, _) in targets.iter() {
-                let Some(column) = target.as_any().downcast_ref::<Column>() else {
-                    return plan_err!("Expects to have column");
-                };
-                fields.push(Field::new(column.name(), data_type.clone(), nullable));
-            }
-        }
-
-        let output_schema = Arc::new(Schema::new_with_metadata(
-            fields,
-            input_schema.metadata().clone(),
-        ));
-
-        Ok(output_schema)
-    }
-
     // Generate a schema which consists of 8 columns (a, b, c, d, e, f, g, h)
     pub fn create_test_schema() -> Result<SchemaRef> {
         let a = Field::new("a", DataType::Int32, true);

datafusion/physical-expr/src/equivalence/properties/dependency.rs

@@ -387,11 +387,11 @@ mod tests {
 
     use super::*;
     use crate::equivalence::tests::{
-        convert_to_sort_reqs, create_test_params, create_test_schema, output_schema,
-        parse_sort_expr,
+        convert_to_sort_reqs, create_test_params, create_test_schema, parse_sort_expr,
     };
     use crate::equivalence::{convert_to_sort_exprs, ProjectionMapping};
     use crate::expressions::{col, BinaryExpr, CastExpr, Column};
+    use crate::projection::tests::output_schema;
     use crate::{ConstExpr, EquivalenceProperties, ScalarFunctionExpr};
 
     use arrow::compute::SortOptions;

datafusion/physical-expr/src/lib.rs

@@ -37,6 +37,7 @@ pub mod intervals;
 mod partitioning;
 mod physical_expr;
 pub mod planner;
+pub mod projection;
 mod scalar_function;
 pub mod simplifier;
 pub mod statistics;