Consolidate projection check

Author: alamb

datafusion/core/src/physical_optimizer/projection_pushdown.rs

@@ -97,6 +97,13 @@ pub fn remove_unnecessary_projections(
         if is_projection_removable(projection) {
             return Ok(Transformed::Yes(projection.input().clone()));
         }
+
+        // If the projection does not narrow the the schema, we should not try
+        // to push it down
+        if projection.expr().len() >= projection.input().schema().fields().len() {
+            return Ok(Transformed::No(plan));
+        }
+
         // If it does, check if we can push it under its child(ren):
         let input = projection.input().as_any();
         if let Some(csv) = input.downcast_ref::<CsvExec>() {
@@ -227,11 +234,6 @@ fn try_swapping_with_output_req(
     projection: &ProjectionExec,
     output_req: &OutputRequirementExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-    // If the projection does not narrow the the schema, we should not try to push it down:
-    if projection.expr().len() >= projection.input().schema().fields().len() {
-        return Ok(None);
-    }
-
     let mut updated_sort_reqs = vec![];
     // None or empty_vec can be treated in the same way.
     if let Some(reqs) = &output_req.required_input_ordering()[0] {
@@ -277,10 +279,6 @@ fn try_swapping_with_output_req(
 fn try_swapping_with_coalesce_partitions(
     projection: &ProjectionExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-    // If the projection does not narrow the the schema, we should not try to push it down:
-    if projection.expr().len() >= projection.input().schema().fields().len() {
-        return Ok(None);
-    }
     // CoalescePartitionsExec always has a single child, so zero indexing is safe.
     make_with_child(projection, &projection.input().children()[0])
         .map(|e| Some(Arc::new(CoalescePartitionsExec::new(e)) as _))
@@ -292,10 +290,6 @@ fn try_swapping_with_filter(
     projection: &ProjectionExec,
     filter: &FilterExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-    // If the projection does not narrow the the schema, we should not try to push it down:
-    if projection.expr().len() >= projection.input().schema().fields().len() {
-        return Ok(None);
-    }
     // Each column in the predicate expression must exist after the projection.
     let Some(new_predicate) = update_expr(filter.predicate(), projection.expr(), false)?
     else {
@@ -313,11 +307,6 @@ fn try_swapping_with_repartition(
     projection: &ProjectionExec,
     repartition: &RepartitionExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-    // If the projection does not narrow the the schema, we should not try to push it down.
-    if projection.expr().len() >= projection.input().schema().fields().len() {
-        return Ok(None);
-    }
-
     // If pushdown is not beneficial or applicable, break it.
     if projection.benefits_from_input_partitioning()[0] || !all_columns(projection.expr())
     {
@@ -355,11 +344,6 @@ fn try_swapping_with_sort(
     projection: &ProjectionExec,
     sort: &SortExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-    // If the projection does not narrow the the schema, we should not try to push it down.
-    if projection.expr().len() >= projection.input().schema().fields().len() {
-        return Ok(None);
-    }
-
     let mut updated_exprs = vec![];
     for sort in sort.expr() {
         let Some(new_expr) = update_expr(&sort.expr, projection.expr(), false)? else {
@@ -384,11 +368,6 @@ fn try_swapping_with_sort_preserving_merge(
     projection: &ProjectionExec,
     spm: &SortPreservingMergeExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-    // If the projection does not narrow the the schema, we should not try to push it down.
-    if projection.expr().len() >= projection.input().schema().fields().len() {
-        return Ok(None);
-    }
-
     let mut updated_exprs = vec![];
     for sort in spm.expr() {
         let Some(updated_expr) = update_expr(&sort.expr, projection.expr(), false)?