Wrap immutable plan parts into Arc

- Closes #19852

Improve performance of query planning and plan state re-set by making node clone cheap.

- Store projection as `Option<Arc<[usize]>>` instead of `Option<Vec<usize>>` in `FilterExec`,
  `HashJoinExec`, `NestedLoopJoinExec`.
- Store exprs as `Arc<[ProjectionExpr]>` instead of Vec in `ProjectionExprs`.
- Store arced aggregation, filter, group by expressions within `AggregateExec`.

Author: askalt

datafusion/common/src/stats.rs

@@ -391,8 +391,13 @@ impl Statistics {
     /// For example, if we had statistics for columns `{"a", "b", "c"}`,
     /// projecting to `vec![2, 1]` would return statistics for columns `{"c",
     /// "b"}`.
-    pub fn project(mut self, projection: Option<&Vec<usize>>) -> Self {
-        let Some(projection) = projection else {
+    pub fn project(self, projection: Option<&impl AsRef<[usize]>>) -> Self {
+        let projection = projection.map(AsRef::as_ref);
+        self.project_impl(projection)
+    }
+
+    fn project_impl(mut self, projection: Option<&[usize]>) -> Self {
+        let Some(projection) = projection.map(AsRef::as_ref) else {
             return self;
         };
 
@@ -410,7 +415,7 @@ impl Statistics {
             .map(Slot::Present)
             .collect();
 
-        for idx in projection {
+        for idx in projection.iter() {
             let next_idx = self.column_statistics.len();
             let slot = std::mem::replace(
                 columns.get_mut(*idx).expect("projection out of bounds"),
@@ -1066,7 +1071,7 @@ mod tests {
 
     #[test]
     fn test_project_none() {
-        let projection = None;
+        let projection: Option<Vec<usize>> = None;
         let stats = make_stats(vec![10, 20, 30]).project(projection.as_ref());
         assert_eq!(stats, make_stats(vec![10, 20, 30]));
     }

datafusion/common/src/utils/mod.rs

@@ -70,10 +70,10 @@ use std::thread::available_parallelism;
 /// ```
 pub fn project_schema(
     schema: &SchemaRef,
-    projection: Option<&Vec<usize>>,
+    projection: Option<&impl AsRef<[usize]>>,
 ) -> Result<SchemaRef> {
     let schema = match projection {
-        Some(columns) => Arc::new(schema.project(columns)?),
+        Some(columns) => Arc::new(schema.project(columns.as_ref())?),
         None => Arc::clone(schema),
     };
     Ok(schema)

datafusion/core/src/physical_planner.rs

@@ -1007,7 +1007,7 @@ impl DefaultPhysicalPlanner {
                         // project the output columns excluding the async functions
                         // The async functions are always appended to the end of the schema.
                         .apply_projection(Some(
-                            (0..input.schema().fields().len()).collect(),
+                            (0..input.schema().fields().len()).collect::<Vec<_>>(),
                         ))?
                         .with_batch_size(session_state.config().batch_size())
                         .build()?

datafusion/core/tests/physical_optimizer/join_selection.rs

@@ -762,7 +762,7 @@ async fn test_hash_join_swap_on_joins_with_projections(
             "ProjectionExec won't be added above if HashJoinExec contains embedded projection",
         );
 
-    assert_eq!(swapped_join.projection, Some(vec![0_usize]));
+    assert_eq!(swapped_join.projection.as_deref().unwrap(), &[0_usize]);
     assert_eq!(swapped.schema().fields.len(), 1);
     assert_eq!(swapped.schema().fields[0].name(), "small_col");
     Ok(())

datafusion/physical-expr/src/projection.rs

@@ -29,7 +29,8 @@ use arrow::datatypes::{Field, Schema, SchemaRef};
 use datafusion_common::stats::{ColumnStatistics, Precision};
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{
-    Result, ScalarValue, assert_or_internal_err, internal_datafusion_err, plan_err,
+    Result, ScalarValue, Statistics, assert_or_internal_err, internal_datafusion_err,
+    plan_err,
 };
 
 use datafusion_physical_expr_common::metrics::ExecutionPlanMetricsSet;
@@ -125,7 +126,8 @@ impl From<ProjectionExpr> for (Arc<dyn PhysicalExpr>, String) {
 /// indices.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ProjectionExprs {
-    exprs: Vec<ProjectionExpr>,
+    /// [`Arc`] used for a cheap clone, which improves physical plan optimization performance.
+    exprs: Arc<[ProjectionExpr]>,
 }
 
 impl std::fmt::Display for ProjectionExprs {
@@ -137,22 +139,24 @@ impl std::fmt::Display for ProjectionExprs {
 
 impl From<Vec<ProjectionExpr>> for ProjectionExprs {
     fn from(value: Vec<ProjectionExpr>) -> Self {
-        Self { exprs: value }
+        Self {
+            exprs: value.into(),
+        }
     }
 }
 
 impl From<&[ProjectionExpr]> for ProjectionExprs {
     fn from(value: &[ProjectionExpr]) -> Self {
         Self {
-            exprs: value.to_vec(),
+            exprs: value.iter().cloned().collect(),
         }
     }
 }
 
 impl FromIterator<ProjectionExpr> for ProjectionExprs {
     fn from_iter<T: IntoIterator<Item = ProjectionExpr>>(exprs: T) -> Self {
         Self {
-            exprs: exprs.into_iter().collect::<Vec<_>>(),
+            exprs: exprs.into_iter().collect(),
         }
     }
 }
@@ -164,12 +168,17 @@ impl AsRef<[ProjectionExpr]> for ProjectionExprs {
 }
 
 impl ProjectionExprs {
-    pub fn new<I>(exprs: I) -> Self
-    where
-        I: IntoIterator<Item = ProjectionExpr>,
-    {
+    /// Make a new [`ProjectionExprs`] from expressions iterator.
+    pub fn new(exprs: impl IntoIterator<Item = ProjectionExpr>) -> Self {
+        Self {
+            exprs: exprs.into_iter().collect(),
+        }
+    }
+
+    /// Make a new [`ProjectionExprs`] from expressions.
+    pub fn from_expressions(exprs: impl Into<Arc<[ProjectionExpr]>>) -> Self {
         Self {
-            exprs: exprs.into_iter().collect::<Vec<_>>(),
+            exprs: exprs.into(),
         }
     }
 
@@ -285,13 +294,14 @@ impl ProjectionExprs {
     {
         let exprs = self
             .exprs
-            .into_iter()
+            .iter()
+            .cloned()
             .map(|mut proj| {
                 proj.expr = f(proj.expr)?;
                 Ok(proj)
             })
-            .collect::<Result<Vec<_>>>()?;
-        Ok(Self::new(exprs))
+            .collect::<Result<Arc<_>>>()?;
+        Ok(Self::from_expressions(exprs))
     }
 
     /// Apply another projection on top of this projection, returning the combined projection.
@@ -361,7 +371,7 @@ impl ProjectionExprs {
     /// applied on top of this projection.
     pub fn try_merge(&self, other: &ProjectionExprs) -> Result<ProjectionExprs> {
         let mut new_exprs = Vec::with_capacity(other.exprs.len());
-        for proj_expr in &other.exprs {
+        for proj_expr in other.exprs.iter() {
             let new_expr = update_expr(&proj_expr.expr, &self.exprs, true)?
                 .ok_or_else(|| {
                     internal_datafusion_err!(
@@ -602,12 +612,12 @@ impl ProjectionExprs {
     /// ```
     pub fn project_statistics(
         &self,
-        mut stats: datafusion_common::Statistics,
+        mut stats: Statistics,
         output_schema: &Schema,
-    ) -> Result<datafusion_common::Statistics> {
+    ) -> Result<Statistics> {
         let mut column_statistics = vec![];
 
-        for proj_expr in &self.exprs {
+        for proj_expr in self.exprs.iter() {
             let expr = &proj_expr.expr;
             let col_stats = if let Some(col) = expr.as_any().downcast_ref::<Column>() {
                 std::mem::take(&mut stats.column_statistics[col.index()])
@@ -754,13 +764,52 @@ impl Projector {
     }
 }
 
-impl IntoIterator for ProjectionExprs {
-    type Item = ProjectionExpr;
-    type IntoIter = std::vec::IntoIter<ProjectionExpr>;
+/// Describes an immutable reference counted projection.
+///
+/// This structure represents projecting a set of columns by index.
+/// [`Arc`] is used to make it cheap to clone.
+pub type ProjectionRef = Arc<[usize]>;
 
-    fn into_iter(self) -> Self::IntoIter {
-        self.exprs.into_iter()
-    }
+/// Combine two projections.
+///
+/// If `p1` is [`None`] then there are no changes.
+/// Otherwise, if passed `p2` is not [`None`] then it is remapped
+/// according to the `p1`. Otherwise, there are no changes.
+///
+/// # Example
+///
+/// If stored projection is [0, 2] and we call `apply_projection([0, 2, 3])`,
+/// then the resulting projection will be [0, 3].
+///
+/// # Error
+///
+/// Returns an internal error if `p1` contains index that is greater than `p2` len.
+///
+pub fn combine_projections(
+    p1: Option<&ProjectionRef>,
+    p2: Option<&ProjectionRef>,
+) -> Result<Option<ProjectionRef>> {
+    let Some(p1) = p1 else {
+        return Ok(None);
+    };
+    let Some(p2) = p2 else {
+        return Ok(Some(Arc::clone(p1)));
+    };
+
+    Ok(Some(
+        p1.iter()
+            .map(|i| {
+                let idx = *i;
+                assert_or_internal_err!(
+                    idx < p2.len(),
+                    "unable to apply projection: index {} is greater than new projection len {}",
+                    idx,
+                    p2.len(),
+                );
+                Ok(p2[*i])
+            })
+            .collect::<Result<Arc<[usize]>>>()?,
+    ))
 }
 
 /// The function operates in two modes:

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -49,7 +49,7 @@ use datafusion_physical_plan::aggregates::{
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::execution_plan::EmissionType;
 use datafusion_physical_plan::joins::{
-    CrossJoinExec, HashJoinExec, PartitionMode, SortMergeJoinExec,
+    CrossJoinExec, HashJoinExec, HashJoinExecBuilder, PartitionMode, SortMergeJoinExec,
 };
 use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr};
 use datafusion_physical_plan::repartition::RepartitionExec;
@@ -305,18 +305,19 @@ pub fn adjust_input_keys_ordering(
                     Vec<(PhysicalExprRef, PhysicalExprRef)>,
                     Vec<SortOptions>,
                 )| {
-                    HashJoinExec::try_new(
+                    HashJoinExecBuilder::new(
                         Arc::clone(left),
                         Arc::clone(right),
                         new_conditions.0,
-                        filter.clone(),
-                        join_type,
-                        // TODO: although projection is not used in the join here, because projection pushdown is after enforce_distribution. Maybe we need to handle it later. Same as filter.
-                        projection.clone(),
-                        PartitionMode::Partitioned,
-                        *null_equality,
-                        *null_aware,
+                        *join_type,
                     )
+                    .with_filter(filter.clone())
+                    // TODO: although projection is not used in the join here, because projection pushdown is after enforce_distribution. Maybe we need to handle it later. Same as filter.
+                    .with_projection_ref(projection.clone())
+                    .with_partition_mode(PartitionMode::Partitioned)
+                    .with_null_equality(*null_equality)
+                    .with_null_aware(*null_aware)
+                    .build()
                     .map(|e| Arc::new(e) as _)
                 };
                 return reorder_partitioned_join_keys(
@@ -638,17 +639,20 @@ pub fn reorder_join_keys_to_inputs(
                     right_keys,
                 } = join_keys;
                 let new_join_on = new_join_conditions(&left_keys, &right_keys);
-                return Ok(Arc::new(HashJoinExec::try_new(
-                    Arc::clone(left),
-                    Arc::clone(right),
-                    new_join_on,
-                    filter.clone(),
-                    join_type,
-                    projection.clone(),
-                    PartitionMode::Partitioned,
-                    *null_equality,
-                    *null_aware,
-                )?));
+                return Ok(Arc::new(
+                    HashJoinExecBuilder::new(
+                        Arc::clone(left),
+                        Arc::clone(right),
+                        new_join_on,
+                        *join_type,
+                    )
+                    .with_filter(filter.clone())
+                    .with_projection_ref(projection.clone())
+                    .with_partition_mode(PartitionMode::Partitioned)
+                    .with_null_equality(*null_equality)
+                    .with_null_aware(*null_aware)
+                    .build()?,
+                ));
             }
         }
     } else if let Some(SortMergeJoinExec {

datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs

@@ -723,7 +723,7 @@ fn handle_hash_join(
         .collect();
 
     let column_indices = build_join_column_index(plan);
-    let projected_indices: Vec<_> = if let Some(projection) = &plan.projection {
+    let projected_indices: Vec<_> = if let Some(projection) = plan.projection.as_ref() {
         projection.iter().map(|&i| &column_indices[i]).collect()
     } else {
         column_indices.iter().collect()

datafusion/physical-optimizer/src/join_selection.rs

@@ -34,7 +34,7 @@ use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_plan::execution_plan::EmissionType;
 use datafusion_physical_plan::joins::utils::ColumnIndex;
 use datafusion_physical_plan::joins::{
-    CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode,
+    CrossJoinExec, HashJoinExec, HashJoinExecBuilder, NestedLoopJoinExec, PartitionMode,
     StreamJoinPartitionMode, SymmetricHashJoinExec,
 };
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
@@ -191,30 +191,18 @@ pub(crate) fn try_collect_left(
             {
                 Ok(Some(hash_join.swap_inputs(PartitionMode::CollectLeft)?))
             } else {
-                Ok(Some(Arc::new(HashJoinExec::try_new(
-                    Arc::clone(left),
-                    Arc::clone(right),
-                    hash_join.on().to_vec(),
-                    hash_join.filter().cloned(),
-                    hash_join.join_type(),
-                    hash_join.projection.clone(),
-                    PartitionMode::CollectLeft,
-                    hash_join.null_equality(),
-                    hash_join.null_aware,
-                )?)))
+                Ok(Some(Arc::new(
+                    HashJoinExecBuilder::from(hash_join)
+                        .with_partition_mode(PartitionMode::CollectLeft)
+                        .build()?,
+                )))
             }
         }
-        (true, false) => Ok(Some(Arc::new(HashJoinExec::try_new(
-            Arc::clone(left),
-            Arc::clone(right),
-            hash_join.on().to_vec(),
-            hash_join.filter().cloned(),
-            hash_join.join_type(),
-            hash_join.projection.clone(),
-            PartitionMode::CollectLeft,
-            hash_join.null_equality(),
-            hash_join.null_aware,
-        )?))),
+        (true, false) => Ok(Some(Arc::new(
+            HashJoinExecBuilder::from(hash_join)
+                .with_partition_mode(PartitionMode::CollectLeft)
+                .build()?,
+        ))),
         (false, true) => {
             // Don't swap null-aware anti joins as they have specific side requirements
             if hash_join.join_type().supports_swap() && !hash_join.null_aware {
@@ -254,17 +242,11 @@ pub(crate) fn partitioned_hash_join(
             PartitionMode::Partitioned
         };
 
-        Ok(Arc::new(HashJoinExec::try_new(
-            Arc::clone(left),
-            Arc::clone(right),
-            hash_join.on().to_vec(),
-            hash_join.filter().cloned(),
-            hash_join.join_type(),
-            hash_join.projection.clone(),
-            partition_mode,
-            hash_join.null_equality(),
-            hash_join.null_aware,
-        )?))
+        Ok(Arc::new(
+            HashJoinExecBuilder::from(hash_join)
+                .with_partition_mode(partition_mode)
+                .build()?,
+        ))
     }
 }
 

datafusion/physical-optimizer/src/projection_pushdown.rs

@@ -135,7 +135,7 @@ fn try_push_down_join_filter(
     );
 
     let new_lhs_length = lhs_rewrite.data.0.schema().fields.len();
-    let projections = match projections {
+    let projections = match projections.as_ref() {
         None => match join.join_type() {
             JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => {
                 // Build projections that ignore the newly projected columns.