diff --git a/codegen/src/main.rs b/codegen/src/main.rs index dd9440dd5..f8922b032 100644 --- a/codegen/src/main.rs +++ b/codegen/src/main.rs @@ -73,6 +73,10 @@ fn generate_iceberg(repo_root: &Path) -> Result<(), Box> ".iceberg.DataFusionPartitioning", "::datafusion_proto::protobuf::Partitioning", ) + .extern_path( + ".iceberg.DataFusionPhysicalExprNode", + "::datafusion_proto::protobuf::PhysicalExprNode", + ) .compile_protos(&[proto_file], &[proto_dir])?; Ok(()) diff --git a/datafusion/distributed-iceberg/README.md b/datafusion/distributed-iceberg/README.md index 5d5c5a4c9..9c5b6507b 100644 --- a/datafusion/distributed-iceberg/README.md +++ b/datafusion/distributed-iceberg/README.md @@ -90,6 +90,14 @@ statistics and task-size estimates. Selected-file column statistics remain a separate change. Task-local eager variants and decoded plans report unknown statistics rather than applying whole-snapshot aggregates to a subset. +Both scans keep the filters DataFusion pushes into them, including the dynamic +filters that joins, top-N sorts and aggregates compute while the query runs. Scans +return `PushedDown::No` for these filters, so DataFusion still checks every row. A join +or aggregate only computes a dynamic filter when a scan below it is a consumer of that +filter. Enable `iceberg.dynamic_filtering_enabled` to make Iceberg scans consumers of the +dynamic filters pushed into them. With it on, `EXPLAIN` shows them on the scan as +`dynamic_filters=`. Scans don't skip data with these filters yet. + ```bash cargo test -p datafusion-distributed-iceberg -p datafusion-distributed-iceberg-test-tools ``` diff --git a/datafusion/distributed-iceberg/src/common/scan_config.rs b/datafusion/distributed-iceberg/src/common/scan_config.rs index 600d6e48c..99ecb7d99 100644 --- a/datafusion/distributed-iceberg/src/common/scan_config.rs +++ b/datafusion/distributed-iceberg/src/common/scan_config.rs @@ -1,10 +1,16 @@ use crate::{IcebergConfig, common::df_err}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::stats::Precision; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion::common::{ColumnStatistics, Statistics}; use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::physical_expr::Partitioning; +use datafusion::physical_expr::expressions::{InListExpr, lit}; +use datafusion::physical_expr::{ + DynamicFilterTracking, Partitioning, PhysicalExpr, conjunction, split_conjunction, +}; +use datafusion::physical_expr_common::physical_expr::snapshot_physical_expr; +use datafusion::physical_plan::apply_expression_roots; use datafusion::physical_plan::limit::LimitStream; use datafusion::physical_plan::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricType, @@ -30,6 +36,8 @@ pub(crate) struct ScanConfig { pub(crate) metrics: ExecutionPlanMetricsSet, pub(crate) iceberg_file_io: FileIO, pub(crate) iceberg_runtime: iceberg::Runtime, + pub(crate) filters: Option>, + pub(crate) dynamic_filtering_enabled: bool, // Clones share the delete loader across this source's local partitions. reader: ArrowReader, } @@ -57,10 +65,77 @@ impl ScanConfig { metrics: Default::default(), iceberg_file_io, iceberg_runtime, + filters: None, + dynamic_filtering_enabled: config.dynamic_filtering_enabled, reader, } } + /// ANDs `filters` onto the scan's filters. + pub(crate) fn push_filters(&mut self, filters: Vec>) { + self.filters = Some(conjunction(self.filters.take().into_iter().chain(filters))); + } + + /// Visits the scan's filters when dynamic filtering is enabled. + pub(crate) fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + if !self.dynamic_filtering_enabled { + return Ok(TreeNodeRecursion::Continue); + } + apply_expression_roots(self.filters.iter(), f) + } + + /// Writes the scan's dynamic filters into `f` when dynamic filtering is enabled. A join filter + /// can list thousands of keys, so each IN list prints as its count. + pub(crate) fn fmt_dynamic_filters(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + if !self.dynamic_filtering_enabled { + return Ok(()); + } + let Some(filters) = &self.filters else { + return Ok(()); + }; + let dynamic_filters: Vec<_> = split_conjunction(filters) + .into_iter() + .filter(|filter| DynamicFilterTracking::classify(filter).contains_dynamic_filter()) + .cloned() + .collect(); + if dynamic_filters.is_empty() { + return Ok(()); + } + write!(f, ", dynamic_filters=")?; + for (i, filter) in dynamic_filters.into_iter().enumerate() { + if i > 0 { + write!(f, " AND ")?; + } + // A dynamic filter is at generation 1 until it gets its first values. + if filter.snapshot_generation() == 1 { + write!(f, "DynamicFilter [ empty ]")?; + continue; + } + let current = snapshot_physical_expr(filter) + .and_then(|current| { + current.transform_up(|expr| match expr.downcast_ref::() { + Some(in_list) => { + let not = if in_list.negated() { "NOT " } else { "" }; + let text = format!( + "{} {not}IN ([{} elements])", + in_list.expr(), + in_list.len() + ); + Ok(Transformed::yes(lit(text))) + } + None => Ok(Transformed::no(expr)), + }) + }) + .map_err(|_| fmt::Error)? + .data; + write!(f, "DynamicFilter [ {current} ]")?; + } + Ok(()) + } + pub(crate) fn open( &self, tasks: FileScanTaskStream, @@ -109,6 +184,8 @@ impl fmt::Debug for ScanConfig { .field("metrics", &self.metrics) .field("iceberg_file_io", &self.iceberg_file_io) .field("iceberg_runtime", &self.iceberg_runtime) + .field("filters", &self.filters) + .field("dynamic_filtering_enabled", &self.dynamic_filtering_enabled) .finish_non_exhaustive() } } diff --git a/datafusion/distributed-iceberg/src/config.rs b/datafusion/distributed-iceberg/src/config.rs index d9c24e850..2f9e38c53 100644 --- a/datafusion/distributed-iceberg/src/config.rs +++ b/datafusion/distributed-iceberg/src/config.rs @@ -18,6 +18,10 @@ extensions_options! { pub column_stats_enabled: bool, default = false /// Discover and retain file tasks during physical planning instead of execution. pub eager_file_planning: bool, default = false + /// (Experimental) Whether Iceberg scans are consumers of the dynamic filters pushed + /// into them, so joins and aggregates compute those filters. Scans don't skip data + /// with them yet. + pub dynamic_filtering_enabled: bool, default = false } } diff --git a/datafusion/distributed-iceberg/src/eager/codec.rs b/datafusion/distributed-iceberg/src/eager/codec.rs index 9aa97478d..b447b4d69 100644 --- a/datafusion/distributed-iceberg/src/eager/codec.rs +++ b/datafusion/distributed-iceberg/src/eager/codec.rs @@ -95,16 +95,17 @@ impl PhysicalExtensionCodec for EagerIcebergCodec { .collect::>>() }) .collect::>>()?; - Ok(DataSourceExec::from_data_source( - EagerIcebergDataSource::from_file_groups( - scan.schema, - groups, - scan.iceberg_file_io, - scan.iceberg_runtime, - scan.fetch, - ctx, - )?, - )) + let mut source = EagerIcebergDataSource::from_file_groups( + scan.schema, + groups, + scan.iceberg_file_io, + scan.iceberg_runtime, + scan.fetch, + ctx, + )?; + source.scan.filters = scan.filters; + source.scan.dynamic_filtering_enabled = scan.dynamic_filtering_enabled; + Ok(DataSourceExec::from_data_source(source)) } fn try_encode( @@ -175,14 +176,22 @@ impl EagerIcebergCodec { proto.custom_file_io.as_deref(), ctx, )?; - Ok(ScanConfig::new( + let filters = proto + .filters + .as_ref() + .map(|filters| proto_converter.proto_to_physical_expr(filters, &schema, &decode_ctx)) + .transpose()?; + let mut scan = ScanConfig::new( schema, partitioning, fetch, iceberg_file_io, self.iceberg_runtime.clone(), ctx, - )) + ); + scan.filters = filters; + scan.dynamic_filtering_enabled = proto.dynamic_filtering_enabled; + Ok(scan) } fn encode_config( @@ -203,6 +212,13 @@ impl EagerIcebergCodec { fetch: source.fetch().map(|value| value as u64), storage_properties: file_io.properties, custom_file_io: file_io.custom, + filters: source + .scan + .filters + .as_ref() + .map(|filters| proto_converter.physical_expr_to_proto(filters, self)) + .transpose()?, + dynamic_filtering_enabled: source.scan.dynamic_filtering_enabled, }) } } diff --git a/datafusion/distributed-iceberg/src/eager/data_source.rs b/datafusion/distributed-iceberg/src/eager/data_source.rs index b4017fb3a..06df942f4 100644 --- a/datafusion/distributed-iceberg/src/eager/data_source.rs +++ b/datafusion/distributed-iceberg/src/eager/data_source.rs @@ -3,12 +3,14 @@ use crate::eager::planned_files::PlannedFiles; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{Statistics, exec_err}; +use datafusion::config::ConfigOptions; use datafusion::datasource::source::DataSource; use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::projection::ProjectionExprs; use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; use datafusion::physical_plan::DisplayFormatType; +use datafusion::physical_plan::filter_pushdown::{FilterPushdownPropagation, PushedDown}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use futures::StreamExt; use iceberg::io::FileIO; @@ -144,9 +146,9 @@ impl EagerIcebergDataSource { impl DataSource for EagerIcebergDataSource { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - Ok(TreeNodeRecursion::Continue) + self.scan.apply_expressions(f) } fn open( &self, @@ -171,7 +173,7 @@ impl DataSource for EagerIcebergDataSource { if let Some(fetch) = self.scan.fetch { write!(f, ", fetch={fetch}")?; } - Ok(()) + self.scan.fmt_dynamic_filters(f) } fn output_partitioning(&self) -> Partitioning { self.scan.partitioning.clone() @@ -200,4 +202,21 @@ impl DataSource for EagerIcebergDataSource { fn metrics(&self) -> ExecutionPlanMetricsSet { self.scan.metrics.clone() } + + fn try_pushdown_filters( + &self, + filters: Vec>, + _config: &ConfigOptions, + ) -> Result>> { + let result = FilterPushdownPropagation::with_parent_pushdown_result(vec![ + PushedDown::No; + filters.len() + ]); + if filters.is_empty() { + return Ok(result); + } + let mut source = self.clone(); + source.scan.push_filters(filters); + Ok(result.with_updated_node(Arc::new(source))) + } } diff --git a/datafusion/distributed-iceberg/src/lazy/codec.rs b/datafusion/distributed-iceberg/src/lazy/codec.rs index 779f6e716..44c3212b2 100644 --- a/datafusion/distributed-iceberg/src/lazy/codec.rs +++ b/datafusion/distributed-iceberg/src/lazy/codec.rs @@ -103,16 +103,24 @@ impl PhysicalExtensionCodec for LazyIcebergCodec { proto.custom_file_io.as_deref(), ctx, )?; + let filters = proto + .filters + .as_ref() + .map(|filters| proto_converter.proto_to_physical_expr(filters, &schema, &decode_ctx)) + .transpose()?; + let mut scan = ScanConfig::new( + schema, + partitioning, + fetch, + iceberg_file_io, + self.iceberg_runtime.clone(), + ctx, + ); + scan.filters = filters; + scan.dynamic_filtering_enabled = proto.dynamic_filtering_enabled; Ok(DataSourceExec::from_data_source(LazyIcebergDataSource { - scan: ScanConfig::new( - schema, - partitioning, - fetch, - iceberg_file_io, - self.iceberg_runtime.clone(), - ctx, - ), + scan, column_stats: None, table_snapshot: None, feed, @@ -150,6 +158,13 @@ impl PhysicalExtensionCodec for LazyIcebergCodec { fetch: source.scan.fetch.map(|value| value as u64), storage_properties: file_io.properties, custom_file_io: file_io.custom, + filters: source + .scan + .filters + .as_ref() + .map(|filters| proto_converter.physical_expr_to_proto(filters, self)) + .transpose()?, + dynamic_filtering_enabled: source.scan.dynamic_filtering_enabled, }; proto.encode(buf).map_err(|error| { proto_error(format!("failed to encode LazyIcebergDataSource: {error}")) @@ -161,12 +176,16 @@ impl PhysicalExtensionCodec for LazyIcebergCodec { mod tests { use datafusion::common::Statistics; use datafusion::datasource::source::DataSource; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{binary, col, lit}; + use datafusion::prelude::SessionConfig; use datafusion_distributed_iceberg_test_tools::create_taxi_table; use datafusion_proto::physical_plan::DefaultPhysicalProtoConverter; use iceberg::io::{FileIOBuilder, LocalFsStorageFactory}; use iceberg::table::Table; use super::*; + use crate::IcebergConfig; use crate::common::df_err; use crate::table_provider::scan::{IcebergDataSourceOptions, plan_scan}; @@ -186,6 +205,10 @@ mod tests { .runtime(iceberg::Runtime::current()) .build() .map_err(df_err)?; + let config = SessionConfig::new().with_option_extension(IcebergConfig { + dynamic_filtering_enabled: true, + ..Default::default() + }); let options = IcebergDataSourceOptions { fetch: Some(10), ..Default::default() @@ -194,9 +217,21 @@ mod tests { table, taxi.schema()?, options, - Arc::new(TaskContext::default()), + Arc::new(TaskContext::default().with_session_config(config.clone())), ) .await?; + let schema = taxi.schema()?; + let filter = binary( + col("fare_amount", &schema)?, + Operator::Gt, + lit(10.0), + &schema, + )?; + let pushed = iceberg_source(&plan)?.try_pushdown_filters(vec![filter], config.options())?; + let Some(pushed_source) = pushed.updated_node else { + return internal_err!("the scan should keep the filter it was offered"); + }; + let plan: Arc = Arc::new(DataSourceExec::new(pushed_source)); let codec = LazyIcebergCodec::new(Arc::new(LocalFsStorageFactory), iceberg::Runtime::current()); let proto_converter = DefaultPhysicalProtoConverter {}; @@ -231,6 +266,11 @@ mod tests { decoded.partition_statistics(None)?.as_ref(), &Statistics::new_unknown(&decoded.scan.schema) ); + let Some(filters) = &decoded.scan.filters else { + return internal_err!("the decoded scan should keep the filter"); + }; + insta::assert_snapshot!(filters, @"fare_amount@8 > 10"); + assert!(decoded.scan.dynamic_filtering_enabled); Ok(()) } diff --git a/datafusion/distributed-iceberg/src/lazy/data_source.rs b/datafusion/distributed-iceberg/src/lazy/data_source.rs index 38686490e..97339bab3 100644 --- a/datafusion/distributed-iceberg/src/lazy/data_source.rs +++ b/datafusion/distributed-iceberg/src/lazy/data_source.rs @@ -126,9 +126,9 @@ impl LazyIcebergDataSource { impl DataSource for LazyIcebergDataSource { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - Ok(TreeNodeRecursion::Continue) + self.scan.apply_expressions(f) } fn open( @@ -149,20 +149,20 @@ impl DataSource for LazyIcebergDataSource { fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "format=iceberg")?; - let Some(feed) = self.feed.inner() else { - return Ok(()); - }; - write!(f, ", table={}", feed.iceberg_table.identifier())?; - if let Some(projection) = &feed.projection { - write!(f, ", projection=[{}]", projection.join(", "))?; + // Workers' copies have no feed, but still show their dynamic filters. + if let Some(feed) = self.feed.inner() { + write!(f, ", table={}", feed.iceberg_table.identifier())?; + if let Some(projection) = &feed.projection { + write!(f, ", projection=[{}]", projection.join(", "))?; + } + if let Some(predicate) = &feed.predicates { + write!(f, ", predicate={predicate}")?; + } + if let Some(fetch) = self.scan.fetch { + write!(f, ", fetch={fetch}")?; + } } - if let Some(predicate) = &feed.predicates { - write!(f, ", predicate={predicate}")?; - } - if let Some(fetch) = self.scan.fetch { - write!(f, ", fetch={fetch}")?; - } - Ok(()) + self.scan.fmt_dynamic_filters(f) } fn output_partitioning(&self) -> Partitioning { @@ -212,12 +212,17 @@ impl DataSource for LazyIcebergDataSource { filters: Vec>, _config: &ConfigOptions, ) -> Result>> { - // TODO: Allow this DataSource to be pushed down filters. Some filters might be more - // straight forward to accept, like simple predicates, but some others might require - // a bit more work, like dynamic filters. - Ok(FilterPushdownPropagation::with_parent_pushdown_result( - vec![PushedDown::No; filters.len()], - )) + // The scan keeps the filters but doesn't apply them, so DataFusion still checks every row. + let result = FilterPushdownPropagation::with_parent_pushdown_result(vec![ + PushedDown::No; + filters.len() + ]); + if filters.is_empty() { + return Ok(result); + } + let mut source = self.clone(); + source.scan.push_filters(filters); + Ok(result.with_updated_node(Arc::new(source))) } fn try_pushdown_sort( diff --git a/datafusion/distributed-iceberg/src/proto/generated/iceberg.rs b/datafusion/distributed-iceberg/src/proto/generated/iceberg.rs index 1a19d3ea9..8e512a298 100644 --- a/datafusion/distributed-iceberg/src/proto/generated/iceberg.rs +++ b/datafusion/distributed-iceberg/src/proto/generated/iceberg.rs @@ -15,6 +15,13 @@ pub struct IcebergDataSource { ::std::collections::HashMap<::prost::alloc::string::String, ::prost::alloc::string::String>, #[prost(bytes = "vec", optional, tag = "6")] pub custom_file_io: ::core::option::Option<::prost::alloc::vec::Vec>, + /// Filters pushed into the scan. A row must pass all of them. + #[prost(message, optional, tag = "7")] + pub filters: ::core::option::Option<::datafusion_proto::protobuf::PhysicalExprNode>, + /// Whether the scan takes part in dynamic filtering. It reports its filters to + /// the operators that fill them in. + #[prost(bool, tag = "8")] + pub dynamic_filtering_enabled: bool, } #[allow(clippy::all)] #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] @@ -36,6 +43,13 @@ pub struct ScanConfig { ::std::collections::HashMap<::prost::alloc::string::String, ::prost::alloc::string::String>, #[prost(bytes = "vec", optional, tag = "5")] pub custom_file_io: ::core::option::Option<::prost::alloc::vec::Vec>, + /// Filters pushed into the scan. A row must pass all of them. + #[prost(message, optional, tag = "6")] + pub filters: ::core::option::Option<::datafusion_proto::protobuf::PhysicalExprNode>, + /// Whether the scan takes part in dynamic filtering. It reports its filters to + /// the operators that fill them in. + #[prost(bool, tag = "7")] + pub dynamic_filtering_enabled: bool, } #[allow(clippy::all)] #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/datafusion/distributed-iceberg/src/proto/iceberg.proto b/datafusion/distributed-iceberg/src/proto/iceberg.proto index b02d978ac..613c4f33c 100644 --- a/datafusion/distributed-iceberg/src/proto/iceberg.proto +++ b/datafusion/distributed-iceberg/src/proto/iceberg.proto @@ -8,6 +8,11 @@ message IcebergDataSource { optional uint64 fetch = 4; map storage_properties = 5; optional bytes custom_file_io = 6; + // Filters pushed into the scan. A row must pass all of them. + DataFusionPhysicalExprNode filters = 7; + // Whether the scan takes part in dynamic filtering. It reports its filters to + // the operators that fill them in. + bool dynamic_filtering_enabled = 8; } message WorkUnitFeed { @@ -17,6 +22,7 @@ message WorkUnitFeed { // The generated Rust fields use DataFusion's canonical protobuf types. message DataFusionSchema {} message DataFusionPartitioning {} +message DataFusionPhysicalExprNode {} message ScanConfig { DataFusionSchema schema = 1; @@ -24,6 +30,11 @@ message ScanConfig { optional uint64 fetch = 3; map storage_properties = 4; optional bytes custom_file_io = 5; + // Filters pushed into the scan. A row must pass all of them. + DataFusionPhysicalExprNode filters = 6; + // Whether the scan takes part in dynamic filtering. It reports its filters to + // the operators that fill them in. + bool dynamic_filtering_enabled = 7; } message AssignedScan { ScanConfig config = 1; diff --git a/datafusion/distributed-iceberg/tests/dynamic_filtering.rs b/datafusion/distributed-iceberg/tests/dynamic_filtering.rs new file mode 100644 index 000000000..09a8d5438 --- /dev/null +++ b/datafusion/distributed-iceberg/tests/dynamic_filtering.rs @@ -0,0 +1,97 @@ +mod tests { + use datafusion::common::Result; + use datafusion::physical_plan::collect; + use datafusion_distributed::{ + DistributedExt, display_plan_ascii, rewrite_distributed_plan_with_dynamic_filters, + }; + use datafusion_distributed_iceberg::IcebergExt; + use datafusion_distributed_iceberg_test_tools::IcebergTestHarness; + use datafusion_distributed_iceberg_test_tools::create_taxi_table; + use std::sync::Arc; + + #[tokio::test] + async fn lazy_planning_scan_gets_the_join_filter() -> Result<()> { + let harness = IcebergTestHarness::builder() + .with_session_extension(|state| state.with_iceberg_integration(Default::default())) + .with_table(create_taxi_table().await?) + .with_config_option("iceberg.eager_file_planning", false) + .with_config_option("iceberg.dynamic_filtering_enabled", true) + .build() + .await?; + + let plan = harness + .physical_plan( + "SELECT COUNT(*) AS trips FROM taxi \ + JOIN (VALUES (DATE '2024-01-10')) AS days(day) ON taxi.pickup_date = days.day \ + WHERE taxi.trip_distance >= 2.0", + ) + .await?; + collect(Arc::clone(&plan), harness.state().task_ctx()).await?; + + insta::assert_snapshot!(display_plan_ascii(plan.as_ref(), false), @" + ProjectionExec: expr=[count(Int64(1))@0 as trips] + AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(day@0, pickup_date@0)], projection=[] + ProjectionExec: expr=[column1@0 as day] + DataSourceExec: partitions=1, partition_sizes=[1] + FilterExec: trip_distance@0 >= 2, projection=[pickup_date@1] + DataSourceExec: format=iceberg, table=default.taxi, projection=[trip_distance, pickup_date], predicate=trip_distance >= 2, dynamic_filters=DynamicFilter [ pickup_date@1 >= 2024-01-10 AND pickup_date@1 <= 2024-01-10 AND pickup_date@1 IN ([1 elements]) ] + "); + Ok(()) + } + + #[tokio::test] + async fn eager_planning_workers_get_the_join_filter() -> Result<()> { + let harness = IcebergTestHarness::builder() + .with_session_extension(|state| state.with_iceberg_integration(Default::default())) + .with_table(create_taxi_table().await?) + .with_config_option("iceberg.eager_file_planning", true) + .with_config_option("iceberg.dynamic_filtering_enabled", true) + .with_workers(2) + .configure_session(|state| { + state + .with_distributed_file_scan_config_bytes_per_partition(1)? + .with_distributed_dynamic_filter_collection(true) + })? + .build() + .await?; + + let plan = harness + .physical_plan( + "SELECT COUNT(*) AS trips FROM taxi \ + JOIN (VALUES (DATE '2024-01-10')) AS days(day) ON taxi.pickup_date = days.day \ + WHERE taxi.trip_distance >= 2.0", + ) + .await?; + let task_ctx = harness.state().task_ctx(); + collect(Arc::clone(&plan), Arc::clone(&task_ctx)).await?; + let plan = rewrite_distributed_plan_with_dynamic_filters(plan, &task_ctx).await?; + + insta::assert_snapshot!(display_plan_ascii(plan.as_ref(), false), @" + ┌───── DistributedExec + │ ProjectionExec: expr=[count(Int64(1))@0 as trips] + │ AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + │ CoalescePartitionsExec + │ [Stage 2] => NetworkCoalesceExec: output_partitions=8, input_tasks=2 + └────────────────────────────────────────────────── + ┌───── Stage 2 ── tasks=2, partitions=8 + │ AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + │ HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(day@0, pickup_date@0)], projection=[] + │ CoalescePartitionsExec + │ [Stage 1] => NetworkBroadcastExec: partitions_per_consumer=1, stage_partitions=2, input_tasks=1 + │ FilterExec: trip_distance@0 >= 2, projection=[pickup_date@1] + │ DistributedLeafExec: + │ t0: DataSourceExec: format=iceberg, planned_files=3, dynamic_filters=DynamicFilter [ pickup_date@1 >= 2024-01-10 AND pickup_date@1 <= 2024-01-10 AND pickup_date@1 IN ([1 elements]) ] + │ t1: DataSourceExec: format=iceberg, planned_files=4, dynamic_filters=DynamicFilter [ pickup_date@1 >= 2024-01-10 AND pickup_date@1 <= 2024-01-10 AND pickup_date@1 IN ([1 elements]) ] + └────────────────────────────────────────────────── + ┌───── Stage 1 ── tasks=1, partitions=2 + │ BroadcastExec: input_partitions=1, consumer_tasks=2, output_partitions=2 + │ ProjectionExec: expr=[column1@0 as day] + │ DataSourceExec: partitions=1, partition_sizes=[1] + └────────────────────────────────────────────────── + "); + Ok(()) + } +}