Encapsulate ProjectionMapping as a struct

Author: alamb

datafusion/physical-expr/src/equivalence.rs

@@ -42,8 +42,55 @@ use indexmap::IndexMap;
 pub type EquivalenceClass = Vec<Arc<dyn PhysicalExpr>>;
 
 /// Stores the mapping between source expressions and target expressions for a
-/// projection. Indices in the vector corresponds to the indices after projection.
-pub type ProjectionMapping = Vec<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)>;
+/// projection.
+#[derive(Debug, Clone)]
+pub struct ProjectionMapping {
+    /// `(source expression)` --> `(target expression)`
+    /// Indices in the vector corresponds to the indices after projection.
+    inner: Vec<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)>,
+}
+
+impl ProjectionMapping {
+    /// Constructs the mapping between a projection's input and output
+    pub fn try_new(
+        expr: &[(Arc<dyn PhysicalExpr>, String)],
+        input_schema: &SchemaRef,
+    ) -> Result<Self> {
+        // Construct a map from the input expressions to the output expression of the projection:
+        let mut inner = vec![];
+        for (expr_idx, (expression, name)) in expr.iter().enumerate() {
+            let target_expr = Arc::new(Column::new(name, expr_idx)) as _;
+
+            let source_expr = expression.clone().transform_down(&|e| match e
+                .as_any()
+                .downcast_ref::<Column>(
+            ) {
+                Some(col) => {
+                    // Sometimes, expression and its name in the input_schema doesn't match.
+                    // This can cause problems. Hence in here we make sure that expression name
+                    // matches with the name in the inout_schema.
+                    // Conceptually, source_expr and expression should be same.
+                    let idx = col.index();
+                    let matching_input_field = input_schema.field(idx);
+                    let matching_input_column =
+                        Column::new(matching_input_field.name(), idx);
+                    Ok(Transformed::Yes(Arc::new(matching_input_column)))
+                }
+                None => Ok(Transformed::No(e)),
+            })?;
+
+            inner.push((source_expr, target_expr));
+        }
+        Ok(Self { inner })
+    }
+
+    /// Iterate over pairs of (source, target) expressions
+    pub fn iter(
+        &self,
+    ) -> impl Iterator<Item = &(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)> + '_ {
+        self.inner.iter()
+    }
+}
 
 /// An `EquivalenceGroup` is a collection of `EquivalenceClass`es where each
 /// class represents a distinct equivalence class in a relation.
@@ -328,7 +375,7 @@ impl EquivalenceGroup {
         // TODO: Convert the algorithm below to a version that uses `HashMap`.
         //       once `Arc<dyn PhysicalExpr>` can be stored in `HashMap`.
         let mut new_classes = vec![];
-        for (source, target) in mapping {
+        for (source, target) in mapping.iter() {
             if new_classes.is_empty() {
                 new_classes.push((source, vec![target.clone()]));
             }
@@ -979,7 +1026,7 @@ impl EquivalenceProperties {
             .iter()
             .filter_map(|order| self.eq_group.project_ordering(projection_mapping, order))
             .collect::<Vec<_>>();
-        for (source, target) in projection_mapping {
+        for (source, target) in projection_mapping.iter() {
             let expr_ordering = ExprOrdering::new(source.clone())
                 .transform_up(&|expr| update_ordering(expr, self))
                 .unwrap();

datafusion/physical-plan/src/aggregates/mod.rs

@@ -25,7 +25,7 @@ use crate::aggregates::{
     no_grouping::AggregateStream, row_hash::GroupedHashAggregateStream,
     topk_stream::GroupedTopKAggregateStream,
 };
-use crate::common::calculate_projection_mapping;
+
 use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use crate::windows::{
     get_ordered_partition_by_indices, get_window_mode, PartitionSearchMode,
@@ -60,6 +60,7 @@ mod topk;
 mod topk_stream;
 
 pub use datafusion_expr::AggregateFunction;
+use datafusion_physical_expr::equivalence::ProjectionMapping;
 pub use datafusion_physical_expr::expressions::create_aggregate_expr;
 
 /// Hash aggregate modes
@@ -294,9 +295,8 @@ pub struct AggregateExec {
     /// expressions from protobuf for final aggregate.
     pub input_schema: SchemaRef,
     /// The mapping used to normalize expressions like Partitioning and
-    /// PhysicalSortExpr. The key is the expression from the input schema
-    /// and the value is the expression from the output schema.
-    projection_mapping: Vec<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)>,
+    /// PhysicalSortExpr that maps input to output
+    projection_mapping: ProjectionMapping,
     /// Execution metrics
     metrics: ExecutionPlanMetricsSet,
     required_input_ordering: Option<LexRequirement>,
@@ -535,7 +535,7 @@ impl AggregateExec {
 
         // construct a map from the input expression to the output expression of the Aggregation group by
         let projection_mapping =
-            calculate_projection_mapping(&group_by.expr, &input.schema())?;
+            ProjectionMapping::try_new(&group_by.expr, &input.schema())?;
 
         let required_input_ordering =
             (!new_requirement.is_empty()).then_some(new_requirement);

datafusion/physical-plan/src/common.rs

@@ -31,10 +31,9 @@ use arrow::datatypes::Schema;
 use arrow::ipc::writer::{FileWriter, IpcWriteOptions};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::stats::Precision;
-use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::tree_node::TreeNode;
 use datafusion_common::{plan_err, DataFusionError, Result};
 use datafusion_execution::memory_pool::MemoryReservation;
-use datafusion_physical_expr::equivalence::ProjectionMapping;
 use datafusion_physical_expr::expressions::{BinaryExpr, Column};
 use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
 
@@ -375,38 +374,6 @@ pub fn batch_byte_size(batch: &RecordBatch) -> usize {
     batch.get_array_memory_size()
 }
 
-/// Constructs the mapping between a projection's input and output
-pub fn calculate_projection_mapping(
-    expr: &[(Arc<dyn PhysicalExpr>, String)],
-    input_schema: &Arc<Schema>,
-) -> Result<ProjectionMapping> {
-    // Construct a map from the input expressions to the output expression of the projection:
-    let mut projection_mapping = vec![];
-    for (expr_idx, (expression, name)) in expr.iter().enumerate() {
-        let target_expr = Arc::new(Column::new(name, expr_idx)) as _;
-
-        let source_expr = expression.clone().transform_down(&|e| match e
-            .as_any()
-            .downcast_ref::<Column>()
-        {
-            Some(col) => {
-                // Sometimes, expression and its name in the input_schema doesn't match.
-                // This can cause problems. Hence in here we make sure that expression name
-                // matches with the name in the inout_schema.
-                // Conceptually, source_expr and expression should be same.
-                let idx = col.index();
-                let matching_input_field = input_schema.field(idx);
-                let matching_input_column = Column::new(matching_input_field.name(), idx);
-                Ok(Transformed::Yes(Arc::new(matching_input_column)))
-            }
-            None => Ok(Transformed::No(e)),
-        })?;
-
-        projection_mapping.push((source_expr, target_expr));
-    }
-    Ok(projection_mapping)
-}
-
 #[cfg(test)]
 mod tests {
     use std::ops::Not;

datafusion/physical-plan/src/projection.rs

@@ -29,7 +29,6 @@ use std::task::{Context, Poll};
 use super::expressions::{Column, PhysicalSortExpr};
 use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use super::{DisplayAs, RecordBatchStream, SendableRecordBatchStream, Statistics};
-use crate::common::calculate_projection_mapping;
 use crate::{
     ColumnStatistics, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr,
 };
@@ -42,6 +41,7 @@ use datafusion_execution::TaskContext;
 use datafusion_physical_expr::expressions::{Literal, UnKnownColumn};
 use datafusion_physical_expr::EquivalenceProperties;
 
+use datafusion_physical_expr::equivalence::ProjectionMapping;
 use futures::stream::{Stream, StreamExt};
 use log::trace;
 
@@ -57,9 +57,8 @@ pub struct ProjectionExec {
     /// The output ordering
     output_ordering: Option<Vec<PhysicalSortExpr>>,
     /// The mapping used to normalize expressions like Partitioning and
-    /// PhysicalSortExpr. The key is the expression from the input schema
-    /// and the value is the expression from the output schema.
-    projection_mapping: Vec<(Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>)>,
+    /// PhysicalSortExpr that maps input to output
+    projection_mapping: ProjectionMapping,
     /// Execution metrics
     metrics: ExecutionPlanMetricsSet,
 }
@@ -94,7 +93,7 @@ impl ProjectionExec {
         ));
 
         // construct a map from the input expressions to the output expression of the Projection
-        let projection_mapping = calculate_projection_mapping(&expr, &input_schema)?;
+        let projection_mapping = ProjectionMapping::try_new(&expr, &input_schema)?;
 
         let input_eqs = input.equivalence_properties();
         let project_eqs = input_eqs.project(&projection_mapping, schema.clone());