[Enhancement] Don't repartition ProjectionExec when it does not compu…

…te anything

ProjectionExec can have the following two types of computations:
1. reorder/rename
2. other computations like col1 + col2
For reorder/rename, ProjectionExec will not benefit from repartition,
we should disable the repartition if all exprs are reorder and rename.

In this pr, we introduce `would_benefit` to ProjectionExec, if it is true,
then ProjectionExec would benefit from partitions, benefits_from_input_partitioning in
ProjectionExec will return true. Otherwise, benefits_from_input_partitioning will return
false. would_benefit will be false if only if all exprs are column_expr.

Signed-off-by: xyz <a997647204@gmail.com>

datafusion/core/src/physical_plan/projection.rs

@@ -61,6 +61,9 @@ pub struct ProjectionExec {
     alias_map: HashMap<Column, Vec<Column>>,
     /// Execution metrics
     metrics: ExecutionPlanMetricsSet,
+    /// If 'would_benefit` is true, ProjectionExec would benefit
+    /// from additional partitions and thus repatitioning is considered.
+    would_benefit: bool,
 }
 
 impl ProjectionExec {
@@ -69,6 +72,13 @@ impl ProjectionExec {
         expr: Vec<(Arc<dyn PhysicalExpr>, String)>,
         input: Arc<dyn ExecutionPlan>,
     ) -> Result<Self> {
+        let mut all_column_expr = true;
+        // If expressions are all column_expr, then projection would not benefit
+        // from the repartition, this is because all computations in this projection
+        // are reorder or rename.
+        for (e, _) in expr.iter() {
+            all_column_expr &= e.is_column_expr()
+        }
         let input_schema = input.schema();
 
         let fields: Result<Vec<Field>> = expr
@@ -134,6 +144,7 @@ impl ProjectionExec {
             output_ordering,
             alias_map,
             metrics: ExecutionPlanMetricsSet::new(),
+            would_benefit: !all_column_expr,
         })
     }
 
@@ -221,6 +232,10 @@ impl ExecutionPlan for ProjectionExec {
         )?))
     }
 
+    fn benefits_from_input_partitioning(&self) -> bool {
+        self.would_benefit
+    }
+
     fn execute(
         &self,
         partition: usize,
@@ -384,8 +399,21 @@ mod tests {
     use crate::scalar::ScalarValue;
     use crate::test::{self};
     use crate::test_util;
+    use datafusion_expr::Operator;
+    use datafusion_physical_expr::expressions::binary;
     use futures::future;
 
+    // Create a binary expression without coercion. Used here when we do not want to coerce the expressions
+    // to valid types. Usage can result in an execution (after plan) error.
+    fn binary_simple(
+        l: Arc<dyn PhysicalExpr>,
+        op: Operator,
+        r: Arc<dyn PhysicalExpr>,
+        input_schema: &Schema,
+    ) -> Arc<dyn PhysicalExpr> {
+        binary(l, op, r, input_schema).unwrap()
+    }
+
     #[tokio::test]
     async fn project_first_column() -> Result<()> {
         let session_ctx = SessionContext::new();
@@ -425,6 +453,38 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn project_input_not_partitioning() -> Result<()> {
+        let schema = test_util::aggr_test_schema();
+
+        let partitions = 4;
+        let csv = test::scan_partitioned_csv(partitions)?;
+
+        // pick column c1 and name it column c1 in the output schema
+        let projection =
+            ProjectionExec::try_new(vec![(col("c1", &schema)?, "c1".to_string())], csv)?;
+        assert!(!projection.would_benefit);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn project_input_partitioning() -> Result<()> {
+        let schema = test_util::aggr_test_schema();
+
+        let partitions = 4;
+        let csv = test::scan_partitioned_csv(partitions)?;
+
+        let c1 = col("c2", &schema).unwrap();
+        let c2 = col("c9", &schema).unwrap();
+        let c1_plus_c2 = binary_simple(c1, Operator::Plus, c2, &schema);
+
+        let projection =
+            ProjectionExec::try_new(vec![(c1_plus_c2, "c2 + c9".to_string())], csv)?;
+
+        assert!(projection.would_benefit);
+        Ok(())
+    }
+
     #[tokio::test]
     async fn project_no_column() -> Result<()> {
         let session_ctx = SessionContext::new();

datafusion/physical-expr/src/expressions/column.rs

@@ -109,6 +109,11 @@ impl PhysicalExpr for Column {
         let col_bounds = context.column_boundaries[self.index].clone();
         context.with_boundaries(col_bounds)
     }
+
+    /// indict whether this expr is a column expr
+    fn is_column_expr(&self) -> bool {
+        true
+    }
 }
 
 impl PartialEq<dyn Any> for Column {

datafusion/physical-expr/src/physical_expr.rs

@@ -81,6 +81,11 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + PartialEq<dyn Any> {
     fn analyze(&self, context: AnalysisContext) -> AnalysisContext {
         context
     }
+
+    /// indict whether this expr is a column expr
+    fn is_column_expr(&self) -> bool {
+        false
+    }
 }
 
 /// The shared context used during the analysis of an expression. Includes