fix: projection for CooperativeExec and CoalesceBatchesExec (#19400)

## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> close #19398 ## Rationale for this change see issue #19398 ## What changes are included in this PR? impl `try_swapping_with_projection` for `CooperativeExec` and `CoalesceBatchesExec` ## Are these changes tested? add test case ## Are there any user-facing changes?

Author: haohuaijin

datafusion/core/tests/physical_optimizer/projection_pushdown.rs

@@ -45,7 +45,9 @@ use datafusion_physical_expr_common::sort_expr::{
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::output_requirements::OutputRequirementExec;
 use datafusion_physical_optimizer::projection_pushdown::ProjectionPushdown;
+use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+use datafusion_physical_plan::coop::CooperativeExec;
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter};
 use datafusion_physical_plan::joins::{
@@ -1677,3 +1679,102 @@ fn test_partition_col_projection_pushdown_expr() -> Result<()> {
 
  Ok(())
 }
+
+#[test]
+fn test_coalesce_batches_after_projection() -> Result<()> {
+ let csv = create_simple_csv_exec();
+ let filter = Arc::new(FilterExec::try_new(
+ Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("c", 2)),
+ Operator::Gt,
+ Arc::new(Literal::new(ScalarValue::Int32(Some(0)))),
+ )),
+ csv,
+ )?);
+ let coalesce_batches: Arc<dyn ExecutionPlan> =
+ Arc::new(CoalesceBatchesExec::new(filter, 8192));
+ let projection: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(
+ vec![
+ ProjectionExpr::new(Arc::new(Column::new("a", 0)), "a"),
+ ProjectionExpr::new(Arc::new(Column::new("b", 1)), "b"),
+ ],
+ coalesce_batches,
+ )?);
+
+ let initial = displayable(projection.as_ref()).indent(true).to_string();
+ let actual = initial.trim();
+
+ assert_snapshot!(
+ actual,
+ @r"
+ ProjectionExec: expr=[a@0 as a, b@1 as b]
+ CoalesceBatchesExec: target_batch_size=8192
+ FilterExec: c@2 > 0
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false
+ "
+ );
+
+ let after_optimize =
+ ProjectionPushdown::new().optimize(projection, &ConfigOptions::new())?;
+
+ let after_optimize_string = displayable(after_optimize.as_ref())
+ .indent(true)
+ .to_string();
+ let actual = after_optimize_string.trim();
+
+ // Projection should be pushed down through CoalesceBatchesExec
+ assert_snapshot!(
+ actual,
+ @r"
+ CoalesceBatchesExec: target_batch_size=8192
+ FilterExec: c@2 > 0, projection=[a@0, b@1]
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false
+ "
+ );
+
+ Ok(())
+}
+
+#[test]
+fn test_cooperative_exec_after_projection() -> Result<()> {
+ let csv = create_simple_csv_exec();
+ let cooperative: Arc<dyn ExecutionPlan> = Arc::new(CooperativeExec::new(csv));
+ let projection: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(
+ vec![
+ ProjectionExpr::new(Arc::new(Column::new("a", 0)), "a"),
+ ProjectionExpr::new(Arc::new(Column::new("b", 1)), "b"),
+ ],
+ cooperative,
+ )?);
+
+ let initial = displayable(projection.as_ref()).indent(true).to_string();
+ let actual = initial.trim();
+
+ assert_snapshot!(
+ actual,
+ @r"
+ ProjectionExec: expr=[a@0 as a, b@1 as b]
+ CooperativeExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false
+ "
+ );
+
+ let after_optimize =
+ ProjectionPushdown::new().optimize(projection, &ConfigOptions::new())?;
+
+ let after_optimize_string = displayable(after_optimize.as_ref())
+ .indent(true)
+ .to_string();
+ let actual = after_optimize_string.trim();
+
+ // Projection should be pushed down through CooperativeExec
+ assert_snapshot!(
+ actual,
+ @r"
+ CooperativeExec
+ DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b], file_type=csv, has_header=false
+ "
+ );
+
+ Ok(())
+}

datafusion/physical-plan/src/coalesce_batches.rs

@@ -24,6 +24,7 @@ use std::task::{Context, Poll};
 
 use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use super::{DisplayAs, ExecutionPlanProperties, PlanProperties, Statistics};
+use crate::projection::ProjectionExec;
 use crate::{
  DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream,
 };
@@ -226,6 +227,18 @@ impl ExecutionPlan for CoalesceBatchesExec {
  CardinalityEffect::Equal
  }
 
+ fn try_swapping_with_projection(
+ &self,
+ projection: &ProjectionExec,
+ ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ match self.input.try_swapping_with_projection(projection)? {
+ Some(new_input) => Ok(Some(
+ Arc::new(self.clone()).with_new_children(vec![new_input])?,
+ )),
+ None => Ok(None),
+ }
+ }
+
  fn gather_filters_for_pushdown(
  &self,
  _phase: FilterPushdownPhase,

datafusion/physical-plan/src/coop.rs

@@ -79,6 +79,7 @@ use crate::filter_pushdown::{
  ChildPushdownResult, FilterDescription, FilterPushdownPhase,
  FilterPushdownPropagation,
 };
+use crate::projection::ProjectionExec;
 use crate::{
  DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream,
  SendableRecordBatchStream,
@@ -207,7 +208,7 @@ where
 /// An execution plan decorator that enables cooperative multitasking.
 /// It wraps the streams produced by its input execution plan using the [`make_cooperative`] function,
 /// which makes the stream participate in Tokio cooperative scheduling.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct CooperativeExec {
  input: Arc<dyn ExecutionPlan>,
  properties: PlanProperties,
@@ -298,6 +299,18 @@ impl ExecutionPlan for CooperativeExec {
  Equal
  }
 
+ fn try_swapping_with_projection(
+ &self,
+ projection: &ProjectionExec,
+ ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+ match self.input.try_swapping_with_projection(projection)? {
+ Some(new_input) => Ok(Some(
+ Arc::new(self.clone()).with_new_children(vec![new_input])?,
+ )),
+ None => Ok(None),
+ }
+ }
+
  fn gather_filters_for_pushdown(
  &self,
  _phase: FilterPushdownPhase,