perf: embed projections in ASOF joins

datafusion/physical-plan/src/joins/asof_join.rs

@@ -110,6 +110,7 @@ use crate::metrics::{
     BaselineMetrics, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricsSet,
     RecordOutput, Time,
 };
+use crate::projection::{EmbeddedProjection, ProjectionExec, try_embed_projection};
 use crate::statistics::{ChildStats, StatisticsArgs};
 use crate::stream::RecordBatchStreamAdapter;
 use crate::{
@@ -246,6 +247,18 @@ impl AsOfJoinExec {
         })
     }
 
+    /// Returns this join emitting only the columns in `projection`, in that order.
+    /// The indices address the join's own schema, before any projection.
+    pub fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> {
+        Self::try_new(
+            Arc::clone(&self.left),
+            Arc::clone(&self.right),
+            self.on.clone(),
+            self.match_condition.clone(),
+            projection,
+        )
+    }
+
     fn compute_properties(
         left: &Arc<dyn ExecutionPlan>,
         join_schema: &SchemaRef,
@@ -289,6 +302,12 @@ impl AsOfJoinExec {
     }
 }
 
+impl EmbeddedProjection for AsOfJoinExec {
+    fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> {
+        self.with_projection(projection)
+    }
+}
+
 impl DisplayAs for AsOfJoinExec {
     fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> std::fmt::Result {
         let on = self
@@ -382,6 +401,16 @@ impl ExecutionPlan for AsOfJoinExec {
         vec![&self.left, &self.right]
     }
 
+    fn try_swapping_with_projection(
+        &self,
+        projection: &ProjectionExec,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+        if self.projection.is_some() {
+            return Ok(None);
+        }
+        try_embed_projection(projection, self)
+    }
+
     fn apply_expressions(
         &self,
         f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>,
@@ -1315,6 +1344,7 @@ mod tests {
     use datafusion_execution::runtime_env::RuntimeEnvBuilder;
     use datafusion_expr::ColumnarValue;
     use datafusion_physical_expr::expressions::{BinaryExpr, CastExpr};
+    use datafusion_physical_expr::projection::ProjectionExpr;
     use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
     use insta::assert_snapshot;
 
@@ -1492,6 +1522,95 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn embeds_output_projection() -> Result<()> {
+        let exec = Arc::new(test_exec()?.with_projection(None)?);
+        let input = Arc::clone(&exec) as Arc<dyn ExecutionPlan>;
+        let projection = ProjectionExec::try_new(
+            [
+                ProjectionExpr {
+                    expr: Arc::new(PhysicalColumn::new("id", 2)),
+                    alias: "id".to_string(),
+                },
+                ProjectionExpr {
+                    expr: Arc::new(PhysicalColumn::new("price", 5)),
+                    alias: "price".to_string(),
+                },
+            ],
+            input,
+        )?;
+
+        let embedded = exec
+            .try_swapping_with_projection(&projection)?
+            .expect("projection should be embedded");
+        let embedded_exec = embedded
+            .downcast_ref::<AsOfJoinExec>()
+            .expect("identity projection should be removed");
+        assert_eq!(embedded_exec.projection.as_deref(), Some(&[2, 5][..]));
+        assert_eq!(
+            embedded_exec
+                .schema()
+                .fields()
+                .iter()
+                .map(|field| field.name().as_str())
+                .collect::<Vec<_>>(),
+            vec!["id", "price"]
+        );
+
+        let batches = collect(embedded, Arc::new(TaskContext::default())).await?;
+        assert_snapshot!(batches_to_sort_string(&batches), @r"
+        +----+-------+
+        | id | price |
+        +----+-------+
+        | 0  |       |
+        | 1  |       |
+        | 2  |       |
+        | 3  | 40    |
+        | 4  | 60    |
+        | 5  | 101   |
+        | 6  |       |
+        +----+-------+
+        ");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn empty_projection_preserves_row_count() -> Result<()> {
+        let exec = Arc::new(test_exec()?.with_projection(None)?);
+        let input = Arc::clone(&exec) as Arc<dyn ExecutionPlan>;
+        let projection = ProjectionExec::try_new(Vec::<ProjectionExpr>::new(), input)?;
+
+        let embedded = exec
+            .try_swapping_with_projection(&projection)?
+            .expect("empty projection should be embedded");
+        let embedded_exec = embedded
+            .downcast_ref::<AsOfJoinExec>()
+            .expect("empty projection should remove ProjectionExec");
+        assert_eq!(embedded_exec.projection.as_deref(), Some(&[][..]));
+        assert!(embedded_exec.schema().fields().is_empty());
+
+        let batches = collect(embedded, Arc::new(TaskContext::default())).await?;
+        assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 7);
+        assert!(batches.iter().all(|batch| batch.num_columns() == 0));
+        Ok(())
+    }
+
+    #[test]
+    fn declines_projection_when_already_embedded() -> Result<()> {
+        let exec = test_exec()?;
+        let input = Arc::clone(&exec) as Arc<dyn ExecutionPlan>;
+        let projection = ProjectionExec::try_new(
+            [ProjectionExpr {
+                expr: Arc::new(PhysicalColumn::new("id", 2)),
+                alias: "id".to_string(),
+            }],
+            input,
+        )?;
+
+        assert!(exec.try_swapping_with_projection(&projection)?.is_none());
+        Ok(())
+    }
+
     fn exec_without_equality_keys(
         left_times: Vec<i64>,
         right_times: Vec<i64>,