feat: expression rewriting in execution plans

Introduces `physical_expressions` and `with_physical_expressions`
methods to the `ExecutionPlan` trait. These methods provide an interface
to inspect and replace physical expressions within execution plan nodes.

The commit also implements these methods for various ExecutionPlans and
introduces a new physical expression invariant check to ensure
consistency between expression retrieval and replacement.

Author: LLDay

datafusion/core/src/physical_planner.rs

@@ -2780,6 +2780,15 @@ impl DefaultPhysicalPlanner {
         // to verify that the plan is executable.
         InvariantChecker(InvariantLevel::Executable).check(&new_plan)?;
 
+        #[cfg(debug_assertions)]
+        {
+            use datafusion_physical_plan::execution_plan::check_physical_expressions;
+            new_plan.apply(|p| {
+                check_physical_expressions(Arc::clone(p))?;
+                Ok(TreeNodeRecursion::Continue)
+            })?;
+        }
+
         debug!(
             "Optimized physical plan:\n{}\n",
             displayable(new_plan.as_ref()).indent(false)

datafusion/core/tests/physical_optimizer/pushdown_utils.rs

@@ -156,6 +156,16 @@ impl FileSource for TestSource {
         })
     }
 
+    fn with_filter_and_projection(
+        &self,
+        filter: Option<Arc<dyn PhysicalExpr>>,
+        _projection: ProjectionExprs,
+    ) -> Result<Option<Arc<dyn FileSource>>> {
+        let mut conf = self.clone();
+        conf.predicate = filter;
+        Ok(Some(Arc::new(conf)))
+    }
+
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource-arrow/src/source.rs

@@ -38,6 +38,7 @@ use datafusion_datasource::{TableSchema, as_file_source};
 
 use arrow::buffer::Buffer;
 use arrow::ipc::reader::{FileDecoder, FileReader, StreamReader};
+use datafusion_common::assert_or_internal_err;
 use datafusion_common::error::Result;
 use datafusion_common::exec_datafusion_err;
 use datafusion_common::tree_node::TreeNodeRecursion;
@@ -51,6 +52,7 @@ use datafusion_physical_plan::projection::ProjectionExprs;
 
 use datafusion_datasource::file_stream::FileOpenFuture;
 use datafusion_datasource::file_stream::FileOpener;
+use datafusion_physical_plan::PhysicalExpr;
 use futures::StreamExt;
 use itertools::Itertools;
 use object_store::{GetOptions, GetRange, GetResultPayload, ObjectStore, ObjectStoreExt};
@@ -400,9 +402,7 @@ impl FileSource for ArrowSource {
 
     fn apply_expressions(
         &self,
-        f: &mut dyn FnMut(
-            &dyn datafusion_physical_plan::PhysicalExpr,
-        ) -> Result<TreeNodeRecursion>,
+        f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
         // Visit projection expressions
         let mut tnr = TreeNodeRecursion::Continue;
@@ -411,6 +411,20 @@ impl FileSource for ArrowSource {
         }
         Ok(tnr)
     }
+
+    fn with_filter_and_projection(
+        &self,
+        filter: Option<Arc<dyn PhysicalExpr>>,
+        projection: ProjectionExprs,
+    ) -> Result<Option<Arc<dyn FileSource>>> {
+        assert_or_internal_err!(filter.is_none(), "filter should not be defined");
+
+        let mut conf = self.clone();
+        conf.projection =
+            SplitProjection::new(self.table_schema.file_schema(), &projection);
+
+        Ok(Some(Arc::new(conf)))
+    }
 }
 
 /// `FileOpener` wrapper for both Arrow IPC file and stream formats

datafusion/datasource-avro/src/source.rs

@@ -22,6 +22,7 @@ use std::sync::Arc;
 
 use crate::avro_to_arrow::Reader as AvroReader;
 
+use datafusion_common::assert_or_internal_err;
 use datafusion_common::error::Result;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_datasource::TableSchema;
@@ -33,6 +34,7 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering;
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion_physical_plan::projection::ProjectionExprs;
 
+use datafusion_physical_plan::PhysicalExpr;
 use object_store::ObjectStore;
 
 /// AvroSource holds the extra configuration that is necessary for opening avro files
@@ -123,6 +125,20 @@ impl FileSource for AvroSource {
         Some(&self.projection.source)
     }
 
+    fn with_filter_and_projection(
+        &self,
+        filter: Option<Arc<dyn PhysicalExpr>>,
+        projection: ProjectionExprs,
+    ) -> Result<Option<Arc<dyn FileSource>>> {
+        assert_or_internal_err!(filter.is_none(), "filter should not be defined");
+
+        let mut conf = self.clone();
+        conf.projection =
+            SplitProjection::new(self.table_schema.file_schema(), &projection);
+
+        Ok(Some(Arc::new(conf)))
+    }
+
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }
@@ -143,9 +159,7 @@ impl FileSource for AvroSource {
 
     fn apply_expressions(
         &self,
-        f: &mut dyn FnMut(
-            &dyn datafusion_physical_plan::PhysicalExpr,
-        ) -> Result<TreeNodeRecursion>,
+        f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
         // Visit projection expressions
         let mut tnr = TreeNodeRecursion::Continue;

datafusion/datasource-csv/src/source.rs

@@ -36,14 +36,14 @@ use datafusion_datasource::{
 use arrow::csv;
 use datafusion_common::config::CsvOptions;
 use datafusion_common::tree_node::TreeNodeRecursion;
-use datafusion_common::{DataFusionError, Result};
+use datafusion_common::{DataFusionError, Result, assert_or_internal_err};
 use datafusion_common_runtime::JoinSet;
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
 use datafusion_execution::TaskContext;
 use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
 use datafusion_physical_plan::{
-    DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+    DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PhysicalExpr,
 };
 
 use crate::file_format::CsvDecoder;
@@ -293,6 +293,20 @@ impl FileSource for CsvSource {
         Some(&self.projection.source)
     }
 
+    fn with_filter_and_projection(
+        &self,
+        filter: Option<Arc<dyn PhysicalExpr>>,
+        projection: ProjectionExprs,
+    ) -> Result<Option<Arc<dyn FileSource>>> {
+        assert_or_internal_err!(filter.is_none(), "filter should not be defined");
+
+        let mut conf = self.clone();
+        conf.projection =
+            SplitProjection::new(self.table_schema.file_schema(), &projection);
+
+        Ok(Some(Arc::new(conf)))
+    }
+
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }
@@ -318,9 +332,7 @@ impl FileSource for CsvSource {
 
     fn apply_expressions(
         &self,
-        f: &mut dyn FnMut(
-            &dyn datafusion_physical_plan::PhysicalExpr,
-        ) -> Result<TreeNodeRecursion>,
+        f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
         // Visit projection expressions
         let mut tnr = TreeNodeRecursion::Continue;

datafusion/datasource-json/src/source.rs

@@ -26,6 +26,7 @@ use std::task::{Context, Poll};
 use crate::file_format::JsonDecoder;
 use crate::utils::{ChannelReader, JsonArrayToNdjsonReader};
 
+use datafusion_common::assert_or_internal_err;
 use datafusion_common::error::{DataFusionError, Result};
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common_runtime::{JoinSet, SpawnedTask};
@@ -37,7 +38,7 @@ use datafusion_datasource::{
     ListingTableUrl, PartitionedFile, RangeCalculation, as_file_source, calculate_range,
 };
 use datafusion_physical_plan::projection::ProjectionExprs;
-use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
+use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties, PhysicalExpr};
 
 use arrow::array::RecordBatch;
 use arrow::json::ReaderBuilder;
@@ -230,6 +231,20 @@ impl FileSource for JsonSource {
         Some(&self.projection.source)
     }
 
+    fn with_filter_and_projection(
+        &self,
+        filter: Option<Arc<dyn PhysicalExpr>>,
+        projection: ProjectionExprs,
+    ) -> Result<Option<Arc<dyn FileSource>>> {
+        assert_or_internal_err!(filter.is_none(), "filter should not be defined");
+
+        let mut conf = self.clone();
+        conf.projection =
+            SplitProjection::new(self.table_schema.file_schema(), &projection);
+
+        Ok(Some(Arc::new(conf)))
+    }
+
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }
@@ -240,9 +255,7 @@ impl FileSource for JsonSource {
 
     fn apply_expressions(
         &self,
-        f: &mut dyn FnMut(
-            &dyn datafusion_physical_plan::PhysicalExpr,
-        ) -> Result<TreeNodeRecursion>,
+        f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
     ) -> Result<TreeNodeRecursion> {
         // Visit projection expressions
         let mut tnr = TreeNodeRecursion::Continue;

datafusion/datasource-parquet/src/source.rs

@@ -604,6 +604,17 @@ impl FileSource for ParquetSource {
         Some(&self.projection)
     }
 
+    fn with_filter_and_projection(
+        &self,
+        filter: Option<Arc<dyn PhysicalExpr>>,
+        projection: ProjectionExprs,
+    ) -> datafusion_common::Result<Option<Arc<dyn FileSource>>> {
+        let mut conf = self.clone();
+        conf.predicate = filter;
+        conf.projection = projection;
+        Ok(Some(Arc::new(conf)))
+    }
+
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource/src/file.rs

@@ -104,6 +104,21 @@ pub trait FileSource: Send + Sync {
         None
     }
 
+    /// Returns new file source with given filter and projection.
+    ///
+    /// This method is primarily used during physical plan rewriting (in
+    /// `ExecutionPlan::with_physical_expressions`) to update the expressions within a file source.
+    ///
+    /// It should NOT be used to pass in projections or filters. That pushdown is handled by
+    /// optimizer rules.
+    fn with_filter_and_projection(
+        &self,
+        _filter: Option<Arc<dyn PhysicalExpr>>,
+        _projection: ProjectionExprs,
+    ) -> Result<Option<Arc<dyn FileSource>>> {
+        Ok(None)
+    }
+
     /// Return execution plan metrics
     fn metrics(&self) -> &ExecutionPlanMetricsSet;
 

datafusion/datasource/src/file_scan_config.rs

@@ -29,7 +29,8 @@ use arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common::{
-    Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err,
+    Constraints, Result, ScalarValue, Statistics, assert_eq_or_internal_err,
+    internal_datafusion_err, internal_err,
 };
 use datafusion_execution::{
     SendableRecordBatchStream, TaskContext, object_store::ObjectStoreUrl,
@@ -38,15 +39,15 @@ use datafusion_expr::Operator;
 
 use datafusion_physical_expr::equivalence::project_orderings;
 use datafusion_physical_expr::expressions::{BinaryExpr, Column};
-use datafusion_physical_expr::projection::ProjectionExprs;
+use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs};
 use datafusion_physical_expr::utils::reassign_expr_columns;
 use datafusion_physical_expr::{EquivalenceProperties, Partitioning, split_conjunction};
 use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 use datafusion_physical_plan::SortOrderPushdownResult;
 use datafusion_physical_plan::coop::cooperative;
-use datafusion_physical_plan::execution_plan::SchedulingType;
+use datafusion_physical_plan::execution_plan::{ReplacePhysicalExpr, SchedulingType};
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType,
     display::{ProjectSchemaDisplay, display_orderings},
@@ -941,6 +942,75 @@ impl DataSource for FileScanConfig {
         // Delegate to the file source
         self.file_source.apply_expressions(f)
     }
+
+    fn physical_expressions<'a>(
+        &'a self,
+    ) -> Option<Box<dyn Iterator<Item = Arc<dyn PhysicalExpr>> + 'a>> {
+        let filter = self.file_source.filter().into_iter();
+        let projection = self
+            .file_source
+            .projection()
+            .into_iter()
+            .flat_map(|p| p.expr_iter());
+
+        Some(Box::new(filter.chain(projection)))
+    }
+
+    fn with_physical_expressions(
+        &self,
+        params: ReplacePhysicalExpr,
+    ) -> Result<Option<Arc<dyn DataSource>>> {
+        let filter_count = self.file_source.filter().iter().len();
+        let projection_count = self
+            .file_source
+            .projection()
+            .map(|p| p.expr_iter().count())
+            .unwrap_or(0);
+
+        let expected_count = filter_count + projection_count;
+        let exprs_count = params.exprs.len();
+
+        assert_eq_or_internal_err!(
+            expected_count,
+            exprs_count,
+            "Inconsistent number of physical expressions for FileScanConfig",
+        );
+
+        let mut filter = None;
+        let mut projection = vec![];
+        let mut expr_iter = params.exprs.into_iter();
+
+        if filter_count > 0 {
+            filter = expr_iter.next();
+        }
+
+        if projection_count > 0 {
+            projection = self
+                .file_source
+                .projection()
+                .expect("should have expressions")
+                .iter()
+                .zip(expr_iter)
+                .map(|(p, expr)| ProjectionExpr::new(expr, p.alias.clone()))
+                .collect();
+        }
+
+        let file_source = self
+            .file_source
+            .with_filter_and_projection(filter, projection.into())?;
+
+        match file_source {
+            Some(file_source) => {
+                let conf_builder: FileScanConfigBuilder = self.clone().into();
+                Ok(Some(Arc::new(
+                    conf_builder.with_source(file_source).build(),
+                )))
+            }
+            None => {
+                internal_err!("file source is not rebuilt")
+            }
+        }
+    }
 }
 
 impl FileScanConfig {