feat: push scalar subquery filters into scans

datafusion/core/tests/physical_optimizer/filter_pushdown.rs

@@ -18,11 +18,13 @@
 use std::sync::{Arc, LazyLock};
 
 use arrow::{
-    array::{RecordBatch, record_batch},
-    datatypes::{DataType, Field, Schema, SchemaRef},
+    array::{RecordBatch, TimestampNanosecondArray, record_batch},
+    datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit},
     util::pretty::pretty_format_batches,
 };
 use arrow_schema::SortOptions;
+use bytes::Bytes;
+use datafusion::datasource::physical_plan::ParquetSource;
 use datafusion::{
     assert_batches_eq,
     logical_expr::Operator,
@@ -36,6 +38,7 @@ use datafusion::{
 use datafusion_catalog::memory::DataSourceExec;
 use datafusion_common::{
     JoinType,
+    arrow::datatypes::IntervalMonthDayNano,
     config::ConfigOptions,
     tree_node::{TreeNode, TreeNodeRecursion},
 };
@@ -44,6 +47,7 @@ use datafusion_datasource::{
 };
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_expr::ScalarUDF;
+use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex};
 use datafusion_functions::math::random::RandomFunc;
 use datafusion_functions_aggregate::{
     count::count_udaf,
@@ -52,6 +56,7 @@ use datafusion_functions_aggregate::{
 use datafusion_physical_expr::{
     LexOrdering, PhysicalSortExpr,
     expressions::{DynamicFilterPhysicalExpr, col},
+    scalar_subquery::ScalarSubqueryExpr,
     utils::conjunction,
 };
 use datafusion_physical_expr::{
@@ -61,6 +66,7 @@ use datafusion_physical_expr::{
 use datafusion_physical_optimizer::{
     PhysicalOptimizerRule, filter_pushdown::FilterPushdown,
 };
+use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink};
 use datafusion_physical_plan::{
     ExecutionPlan,
     aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy},
@@ -72,6 +78,8 @@ use datafusion_physical_plan::{
     repartition::RepartitionExec,
     sorts::sort::SortExec,
 };
+use object_store::{ObjectStoreExt, path::Path};
+use parquet::{arrow::ArrowWriter, file::properties::WriterProperties};
 
 use super::pushdown_utils::{
     OptimizationTest, TestNode, TestScanBuilder, TestSource, format_plan_for_test,
@@ -3064,6 +3072,174 @@ fn test_hashjoin_dynamic_filter_pushdown_is_used() {
     }
 }
 
+/// A completed scalar-subquery filter is retained for Parquet statistics pruning even
+/// when Parquet row-filter pushdown is disabled.
+#[tokio::test]
+async fn scalar_subquery_dynamic_filter_parquet_pruning_with_pushdown_disabled() {
+    let timestamp = DataType::Timestamp(TimeUnit::Nanosecond, None);
+    let schema = Arc::new(Schema::new(vec![Field::new(
+        "ts",
+        timestamp.clone(),
+        false,
+    )]));
+    let batch = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![Arc::new(TimestampNanosecondArray::from([0, 2_000_000_000]))],
+    )
+    .unwrap();
+
+    let mut parquet_bytes = Vec::new();
+    let properties = WriterProperties::builder()
+        .set_max_row_group_row_count(Some(1))
+        .build();
+    let mut writer =
+        ArrowWriter::try_new(&mut parquet_bytes, Arc::clone(&schema), Some(properties))
+            .unwrap();
+    writer.write(&batch).unwrap();
+    writer.close().unwrap();
+
+    let object_store = Arc::new(InMemory::new());
+    let path = Path::from("scalar-subquery.parquet");
+    let file_size = parquet_bytes.len() as u64;
+    object_store
+        .put(&path, Bytes::from(parquet_bytes).into())
+        .await
+        .unwrap();
+    let scan = DataSourceExec::from_data_source(
+        FileScanConfigBuilder::new(
+            ObjectStoreUrl::parse("test://").unwrap(),
+            Arc::new(ParquetSource::new(Arc::clone(&schema))),
+        )
+        .with_file(PartitionedFile::new("scalar-subquery.parquet", file_size))
+        .build(),
+    ) as Arc<dyn ExecutionPlan>;
+
+    let results = ScalarSubqueryResults::new(1);
+    let producer_batch = RecordBatch::try_new(
+        Arc::clone(&schema),
+        vec![Arc::new(TimestampNanosecondArray::from([1_500_000_000]))],
+    )
+    .unwrap();
+    let producer = datafusion::datasource::memory::MemorySourceConfig::try_new_exec(
+        &[vec![producer_batch]],
+        Arc::clone(&schema),
+        None,
+    )
+    .unwrap();
+    let scalar = Arc::new(ScalarSubqueryExpr::new(
+        timestamp,
+        false,
+        SubqueryIndex::new(0),
+        results.clone(),
+    ));
+    let predicate = Arc::new(BinaryExpr::new(
+        col("ts", &schema).unwrap(),
+        Operator::GtEq,
+        Arc::new(BinaryExpr::new(
+            scalar,
+            Operator::Minus,
+            Arc::new(Literal::new(ScalarValue::IntervalMonthDayNano(Some(
+                IntervalMonthDayNano {
+                    months: 0,
+                    days: 0,
+                    nanoseconds: 500_000_000,
+                },
+            )))),
+        )),
+    )) as Arc<dyn PhysicalExpr>;
+    let plan = Arc::new(ScalarSubqueryExec::new(
+        Arc::new(FilterExec::try_new(predicate, scan).unwrap()),
+        vec![ScalarSubqueryLink {
+            plan: producer,
+            index: SubqueryIndex::new(0),
+        }],
+        results,
+    )) as Arc<dyn ExecutionPlan>;
+
+    let mut config = ConfigOptions::default();
+    config.execution.parquet.pushdown_filters = false;
+    config.execution.parquet.pruning = true;
+    config.optimizer.enable_dynamic_filter_pushdown = true;
+    let optimized = FilterPushdown::new_post_optimization()
+        .optimize(plan, &config)
+        .unwrap();
+    let scalar_exec = optimized
+        .downcast_ref::<ScalarSubqueryExec>()
+        .expect("optimized root should retain ScalarSubqueryExec");
+    assert_eq!(scalar_exec.subqueries().len(), 1);
+    let expression_id = scalar_exec
+        .dynamic_expressions_produced()
+        .first()
+        .expect("scalar subquery producer should retain its dynamic filter")
+        .expression_id()
+        .expect("dynamic filter should have an expression ID");
+    assert!(scalar_exec.input().downcast_ref::<FilterExec>().is_some());
+
+    let mut scan_consumer = None;
+    scalar_exec
+        .input()
+        .apply(|node| {
+            if let Some(scan) = node.downcast_ref::<DataSourceExec>()
+                && let Some((_, parquet)) =
+                    scan.downcast_to_file_source::<ParquetSource>()
+                && let Some(predicate) = parquet.filter()
+            {
+                predicate.apply(|expr| {
+                    if expr.expression_id() == Some(expression_id) {
+                        scan_consumer = Some(Arc::clone(expr));
+                        Ok(TreeNodeRecursion::Stop)
+                    } else {
+                        Ok(TreeNodeRecursion::Continue)
+                    }
+                })?;
+            }
+            Ok(TreeNodeRecursion::Continue)
+        })
+        .unwrap();
+    let scan_consumer =
+        scan_consumer.expect("Parquet predicate should retain the dynamic filter ID");
+
+    let context = SessionContext::new_with_config(SessionConfig::from(config));
+    context.register_object_store(
+        ObjectStoreUrl::parse("test://").unwrap().as_ref(),
+        object_store,
+    );
+    let batches = collect(optimized.clone(), context.task_ctx())
+        .await
+        .unwrap();
+    assert_batches_eq!(
+        &[
+            "+---------------------+",
+            "| ts                  |",
+            "+---------------------+",
+            "| 1970-01-01T00:00:02 |",
+            "+---------------------+",
+        ],
+        &batches
+    );
+
+    let consumer = scan_consumer
+        .downcast_ref::<DynamicFilterPhysicalExpr>()
+        .expect("scan consumer should be a dynamic filter");
+    assert_ne!(consumer.current().unwrap().to_string(), "true");
+    consumer.wait_complete().await;
+
+    let metrics =
+        datafusion::test_util::parquet::TestParquetFile::parquet_metrics(&optimized)
+            .expect("Parquet scan should have metrics");
+    let pruned = match metrics.sum_by_name("row_groups_pruned_statistics").unwrap() {
+        datafusion::physical_plan::metrics::MetricValue::PruningMetrics {
+            pruning_metrics,
+            ..
+        } => pruning_metrics.pruned(),
+        metric => panic!("unexpected statistics pruning metric: {metric:?}"),
+    };
+    assert!(
+        pruned >= 1,
+        "expected statistics pruning, observed {pruned}"
+    );
+}
+
 /// Regression test for https://github.com/apache/datafusion/issues/20109.
 ///
 /// Not portable to sqllogictest: the regression specifically targets the

datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs

@@ -183,6 +183,16 @@ impl Display for DynamicFilterPhysicalExpr {
 }
 
 impl DynamicFilterPhysicalExpr {
+    /// Returns `true` if `other` shares this filter's runtime state.
+    ///
+    /// Derived filters can have distinct outer objects while sharing the same
+    /// state. Conversely, filters reconstructed independently (for example by
+    /// decoding) do not share state even when their expression IDs match.
+    #[doc(hidden)]
+    pub fn shares_runtime_state(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.inner, &other.inner)
+    }
+
     /// Create a new [`DynamicFilterPhysicalExpr`]
     /// from an initial expression and a list of children.
     /// The list of children is provided separately because
@@ -314,15 +324,7 @@ impl DynamicFilterPhysicalExpr {
     /// - When we've computed the probe side's hash table in a HashJoinExec
     /// - After every batch is processed if we update the TopK heap in a SortExec using a TopK approach.
     pub fn update(&self, new_expr: Arc<dyn PhysicalExpr>) -> Result<()> {
-        // Remap the children of the new expression to match the original children
-        // We still do this again in `current()` but doing it preventively here
-        // reduces the work needed in some cases if `current()` is called multiple times
-        // and the same externally facing `PhysicalExpr` is used for both `with_new_children` and `update()`.`
-        let new_expr = Self::remap_children(
-            &self.children,
-            self.remapped_children.as_ref(),
-            new_expr,
-        )?;
+        // Store the canonical expression; `current()` remaps it for each derived consumer.
 
         // Load the current inner, increment generation, and store the new one
         let mut current = self.inner.write();

datafusion/physical-expr/src/scalar_subquery.rs

@@ -29,6 +29,8 @@ use datafusion_expr_common::columnar_value::ColumnarValue;
 use datafusion_expr_common::sort_properties::{ExprProperties, SortProperties};
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 
+use crate::expressions::Literal;
+
 /// A physical expression whose value is provided by a scalar subquery.
 ///
 /// Subquery execution is handled by `ScalarSubqueryExec`, which stores the
@@ -133,6 +135,15 @@ impl PhysicalExpr for ScalarSubqueryExpr {
         Ok(ColumnarValue::Scalar(value))
     }
 
+    fn snapshot(&self) -> Result<Option<Arc<dyn PhysicalExpr>>> {
+        let value = self.results.get(self.index).ok_or_else(|| {
+            internal_datafusion_err!(
+                "ScalarSubqueryExpr snapshotted before the subquery was executed"
+            )
+        })?;
+        Ok(Some(Arc::new(Literal::new(value))))
+    }
+
     fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
         vec![]
     }
@@ -274,6 +285,56 @@ mod tests {
         assert!(result.is_err());
     }
 
+    #[test]
+    fn test_snapshot_pending_and_populated_values() -> Result<()> {
+        let pending_results = ScalarSubqueryResults::new(1);
+        let pending = ScalarSubqueryExpr::new(
+            DataType::Int32,
+            true,
+            SubqueryIndex::new(0),
+            pending_results,
+        );
+        assert!(pending.snapshot().is_err());
+
+        let results = ScalarSubqueryResults::new(2);
+        let non_null = ScalarSubqueryExpr::new(
+            DataType::Int32,
+            false,
+            SubqueryIndex::new(0),
+            results.clone(),
+        );
+        let null = ScalarSubqueryExpr::new(
+            DataType::Utf8,
+            true,
+            SubqueryIndex::new(1),
+            results.clone(),
+        );
+        results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?;
+        results.set(SubqueryIndex::new(1), ScalarValue::Utf8(None))?;
+
+        let non_null_snapshot = non_null
+            .snapshot()?
+            .expect("populated scalar subquery should snapshot");
+        assert_eq!(
+            non_null_snapshot
+                .downcast_ref::<Literal>()
+                .expect("snapshot should be a literal")
+                .value(),
+            &ScalarValue::Int32(Some(42))
+        );
+        let null_snapshot = null
+            .snapshot()?
+            .expect("populated scalar subquery should snapshot");
+        assert_eq!(
+            null_snapshot
+                .downcast_ref::<Literal>()
+                .expect("snapshot should be a literal")
+                .value(),
+            &ScalarValue::Utf8(None)
+        );
+        Ok(())
+    }
+
     #[test]
     fn test_identity_equality() {
         let results = make_results(vec![None, None]);

datafusion/physical-plan/src/proto.rs

@@ -135,6 +135,15 @@ pub trait ExecutionPlanDecode {
         input_schema: &Schema,
     ) -> Result<Arc<dyn PhysicalExpr>>;
 
+    /// Deserialize a physical expression with `results` active for scalar
+    /// subquery expressions in that expression's subtree.
+    fn decode_expr_with_scalar_subquery_results(
+        &self,
+        node: &PhysicalExprNode,
+        input_schema: &Schema,
+        results: ScalarSubqueryResults,
+    ) -> Result<Arc<dyn PhysicalExpr>>;
+
     /// The session task context, used by plans that need the function registry
     /// or session configuration. Never exposes the proto extension codec.
     fn task_ctx(&self) -> &TaskContext;
@@ -294,6 +303,18 @@ impl<'a> ExecutionPlanDecodeCtx<'a> {
         self.decoder.decode_expr(node, input_schema)
     }
 
+    /// Deserialize a physical expression with scalar subquery results scoped
+    /// to this ScalarSubqueryExec.
+    pub fn decode_expr_with_scalar_subquery_results(
+        &self,
+        node: &PhysicalExprNode,
+        input_schema: &Schema,
+        results: ScalarSubqueryResults,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        self.decoder
+            .decode_expr_with_scalar_subquery_results(node, input_schema, results)
+    }
+
     /// Deserialize a required physical expression against `input_schema`.
     pub fn decode_required_expr(
         &self,