This is an automated email from the ASF dual-hosted git repository.
github-bot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 28290baa3f Impl `gather_filters_for_pushdown` for
`CoalescePartitionsExec` (#18046)
28290baa3f is described below
commit 28290baa3f87682497f4d6e9f5a58b57e3e9767f
Author: xudong.w <[email protected]>
AuthorDate: Tue Oct 14 14:51:27 2025 +0800
Impl `gather_filters_for_pushdown` for `CoalescePartitionsExec` (#18046)
* Impl gather_filters_for_pushdown for CoalescePartitionsExec
* add tests
---
.../physical_optimizer/filter_pushdown/mod.rs | 134 ++++++++++++++++++++-
.../physical_optimizer/filter_pushdown/util.rs | 2 +-
.../physical-plan/src/coalesce_partitions.rs | 12 ++
3 files changed, 146 insertions(+), 2 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown/mod.rs
b/datafusion/core/tests/physical_optimizer/filter_pushdown/mod.rs
index f002355f66..614378db94 100644
--- a/datafusion/core/tests/physical_optimizer/filter_pushdown/mod.rs
+++ b/datafusion/core/tests/physical_optimizer/filter_pushdown/mod.rs
@@ -33,7 +33,11 @@ use datafusion::{
prelude::{ParquetReadOptions, SessionConfig, SessionContext},
scalar::ScalarValue,
};
+use datafusion_catalog::memory::DataSourceExec;
use datafusion_common::config::ConfigOptions;
+use datafusion_datasource::{
+ file_groups::FileGroup, file_scan_config::FileScanConfigBuilder,
PartitionedFile,
+};
use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_expr::ScalarUDF;
use datafusion_functions::math::random::RandomFunc;
@@ -60,6 +64,8 @@ use futures::StreamExt;
use object_store::{memory::InMemory, ObjectStore};
use util::{format_plan_for_test, OptimizationTest, TestNode, TestScanBuilder};
+use crate::physical_optimizer::filter_pushdown::util::TestSource;
+
mod util;
#[test]
@@ -834,6 +840,132 @@ async fn
test_topk_dynamic_filter_pushdown_multi_column_sort() {
assert!(stream.next().await.is_none());
}
+#[tokio::test]
+async fn test_topk_filter_passes_through_coalesce_partitions() {
+ // Create multiple batches for different partitions
+ let batches = vec![
+ record_batch!(
+ ("a", Utf8, ["aa", "ab"]),
+ ("b", Utf8, ["bd", "bc"]),
+ ("c", Float64, [1.0, 2.0])
+ )
+ .unwrap(),
+ record_batch!(
+ ("a", Utf8, ["ac", "ad"]),
+ ("b", Utf8, ["bb", "ba"]),
+ ("c", Float64, [2.0, 1.0])
+ )
+ .unwrap(),
+ ];
+
+ // Create a source that supports all batches
+ let source = Arc::new(TestSource::new(true, batches));
+
+ let base_config = FileScanConfigBuilder::new(
+ ObjectStoreUrl::parse("test://").unwrap(),
+ Arc::clone(&schema()),
+ source,
+ )
+ .with_file_groups(vec![
+ // Partition 0
+ FileGroup::new(vec![PartitionedFile::new("test1.parquet", 123)]),
+ // Partition 1
+ FileGroup::new(vec![PartitionedFile::new("test2.parquet", 123)]),
+ ])
+ .build();
+
+ let scan = DataSourceExec::from_data_source(base_config);
+
+ // Add CoalescePartitionsExec to merge the two partitions
+ let coalesce = Arc::new(CoalescePartitionsExec::new(scan)) as Arc<dyn
ExecutionPlan>;
+
+ // Add SortExec with TopK
+ let plan = Arc::new(
+ SortExec::new(
+ LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("b", &schema()).unwrap(),
+ SortOptions::new(true, false),
+ )])
+ .unwrap(),
+ coalesce,
+ )
+ .with_fetch(Some(1)),
+ ) as Arc<dyn ExecutionPlan>;
+
+ // Test optimization - the filter SHOULD pass through
CoalescePartitionsExec
+ // if it properly implements from_children (not all_unsupported)
+ insta::assert_snapshot!(
+ OptimizationTest::new(Arc::clone(&plan),
FilterPushdown::new_post_optimization(), true),
+ @r"
+ OptimizationTest:
+ input:
+ - SortExec: TopK(fetch=1), expr=[b@1 DESC NULLS LAST],
preserve_partitioning=[false]
+ - CoalescePartitionsExec
+ - DataSourceExec: file_groups={2 groups: [[test1.parquet],
[test2.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
+ output:
+ Ok:
+ - SortExec: TopK(fetch=1), expr=[b@1 DESC NULLS LAST],
preserve_partitioning=[false]
+ - CoalescePartitionsExec
+ - DataSourceExec: file_groups={2 groups: [[test1.parquet],
[test2.parquet]]}, projection=[a, b, c], file_type=test,
pushdown_supported=true, predicate=DynamicFilter [ empty ]
+ "
+ );
+}
+
+#[tokio::test]
+async fn test_topk_filter_passes_through_coalesce_batches() {
+ let batches = vec![
+ record_batch!(
+ ("a", Utf8, ["aa", "ab"]),
+ ("b", Utf8, ["bd", "bc"]),
+ ("c", Float64, [1.0, 2.0])
+ )
+ .unwrap(),
+ record_batch!(
+ ("a", Utf8, ["ac", "ad"]),
+ ("b", Utf8, ["bb", "ba"]),
+ ("c", Float64, [2.0, 1.0])
+ )
+ .unwrap(),
+ ];
+
+ let scan = TestScanBuilder::new(schema())
+ .with_support(true)
+ .with_batches(batches)
+ .build();
+
+ let coalesce_batches =
+ Arc::new(CoalesceBatchesExec::new(scan, 1024)) as Arc<dyn
ExecutionPlan>;
+
+ // Add SortExec with TopK
+ let plan = Arc::new(
+ SortExec::new(
+ LexOrdering::new(vec![PhysicalSortExpr::new(
+ col("b", &schema()).unwrap(),
+ SortOptions::new(true, false),
+ )])
+ .unwrap(),
+ coalesce_batches,
+ )
+ .with_fetch(Some(1)),
+ ) as Arc<dyn ExecutionPlan>;
+
+ insta::assert_snapshot!(
+ OptimizationTest::new(Arc::clone(&plan),
FilterPushdown::new_post_optimization(), true),
+ @r"
+ OptimizationTest:
+ input:
+ - SortExec: TopK(fetch=1), expr=[b@1 DESC NULLS LAST],
preserve_partitioning=[false]
+ - CoalesceBatchesExec: target_batch_size=1024
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[a, b, c], file_type=test, pushdown_supported=true
+ output:
+ Ok:
+ - SortExec: TopK(fetch=1), expr=[b@1 DESC NULLS LAST],
preserve_partitioning=[false]
+ - CoalesceBatchesExec: target_batch_size=1024
+ - DataSourceExec: file_groups={1 group: [[test.parquet]]},
projection=[a, b, c], file_type=test, pushdown_supported=true,
predicate=DynamicFilter [ empty ]
+ "
+ );
+}
+
#[tokio::test]
async fn test_hashjoin_dynamic_filter_pushdown() {
use datafusion_common::JoinType;
@@ -1662,7 +1794,7 @@ async fn test_topk_dynamic_filter_pushdown_integration() {
ctx.sql(
r"
COPY (
- SELECT 1372708800 + value AS t
+ SELECT 1372708800 + value AS t
FROM generate_series(0, 99999)
ORDER BY t
) TO 'memory:///1.parquet'
diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs
b/datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs
index ffc2607f57..f05f3f0028 100644
--- a/datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs
+++ b/datafusion/core/tests/physical_optimizer/filter_pushdown/util.rs
@@ -115,7 +115,7 @@ pub struct TestSource {
}
impl TestSource {
- fn new(support: bool, batches: Vec<RecordBatch>) -> Self {
+ pub fn new(support: bool, batches: Vec<RecordBatch>) -> Self {
Self {
support,
metrics: ExecutionPlanMetricsSet::new(),
diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs
b/datafusion/physical-plan/src/coalesce_partitions.rs
index 323a844ae3..5869c51b26 100644
--- a/datafusion/physical-plan/src/coalesce_partitions.rs
+++ b/datafusion/physical-plan/src/coalesce_partitions.rs
@@ -28,11 +28,14 @@ use super::{
Statistics,
};
use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType};
+use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase};
use crate::projection::{make_with_child, ProjectionExec};
use crate::{DisplayFormatType, ExecutionPlan, Partitioning};
+use datafusion_common::config::ConfigOptions;
use datafusion_common::{internal_err, Result};
use datafusion_execution::TaskContext;
+use datafusion_physical_expr::PhysicalExpr;
/// Merge execution plan executes partitions in parallel and combines them
into a single
/// partition. No guarantees are made about the order of the resulting
partition.
@@ -260,6 +263,15 @@ impl ExecutionPlan for CoalescePartitionsExec {
cache: self.cache.clone(),
}))
}
+
+ fn gather_filters_for_pushdown(
+ &self,
+ _phase: FilterPushdownPhase,
+ parent_filters: Vec<Arc<dyn PhysicalExpr>>,
+ _config: &ConfigOptions,
+ ) -> Result<FilterDescription> {
+ FilterDescription::from_children(parent_filters, &self.children())
+ }
}
#[cfg(test)]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
