Skip to content

Commit 1e80253

Browse files
adriangbclaude
andcommitted
test: optional dynamic filters end to end
- `parquet/optional_filters.rs`: SQL queries on Parquet files. The hash join, TopK and aggregate dynamic filters arrive at `ParquetSource` as optional conjuncts, and user filters stay required. Results are correct. - `test_hashjoin_optional_dynamic_filter_default_source`: a file source that uses the default `try_pushdown_filter` gets the optional dynamic filter as a normal filter, and the join result is correct. - `test_reordered_parent_filters_are_rejected`: a node that returns its parent filters in a different order gives an error in debug builds. - Proto: the flags survive a round trip, and a reader that does not know the new fields applies all conjuncts as required. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 074fe91 commit 1e80253

6 files changed

Lines changed: 448 additions & 7 deletions

File tree

‎datafusion/core/tests/parquet/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ mod expr_adapter;
5353
mod external_access_plan;
5454
mod file_statistics;
5555
mod filter_pushdown;
56+
mod optional_filters;
5657
mod ordering;
5758
mod page_pruning;
5859
mod row_group_pruning;
Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! End-to-end tests: the dynamic filters of the hash join, TopK and
19+
//! aggregate operators arrive at the Parquet scan as *optional* conjuncts of
20+
//! its [`PhysicalFilter`], and the query results do not change.
21+
//!
22+
//! [`PhysicalFilter`]: datafusion_physical_expr::filter::PhysicalFilter
23+
24+
use std::sync::Arc;
25+
26+
use arrow::array::{Int64Array, RecordBatch};
27+
use arrow::util::pretty::pretty_format_batches;
28+
use arrow_schema::{DataType, Field, Schema};
29+
use datafusion::datasource::physical_plan::ParquetSource;
30+
use datafusion::datasource::source::DataSourceExec;
31+
use datafusion::physical_plan::{ExecutionPlan, collect};
32+
use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext};
33+
use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
34+
use datafusion_datasource::file_scan_config::FileScanConfig;
35+
use insta::assert_snapshot;
36+
use parquet::arrow::ArrowWriter;
37+
use tempfile::TempDir;
38+
39+
/// Write `columns` as a single-file Parquet table `name` and register it.
40+
async fn register_table(
41+
ctx: &SessionContext,
42+
dir: &TempDir,
43+
name: &str,
44+
columns: Vec<(&str, Vec<i64>)>,
45+
) {
46+
let schema = Arc::new(Schema::new(
47+
columns
48+
.iter()
49+
.map(|(name, _)| Field::new(*name, DataType::Int64, false))
50+
.collect::<Vec<_>>(),
51+
));
52+
let arrays = columns
53+
.into_iter()
54+
.map(|(_, values)| Arc::new(Int64Array::from(values)) as _)
55+
.collect();
56+
let batch = RecordBatch::try_new(Arc::clone(&schema), arrays).unwrap();
57+
let path = dir.path().join(format!("{name}.parquet"));
58+
let file = std::fs::File::create(&path).unwrap();
59+
let mut writer = ArrowWriter::try_new(file, schema, None).unwrap();
60+
writer.write(&batch).unwrap();
61+
writer.close().unwrap();
62+
ctx.register_parquet(name, path.to_str().unwrap(), ParquetReadOptions::default())
63+
.await
64+
.unwrap();
65+
}
66+
67+
/// A context with a small `build` table (`k` in 1..=3) and a larger `probe`
68+
/// table (`k` in 0..100, `v = 10 * k`).
69+
async fn context() -> (SessionContext, TempDir) {
70+
let config = SessionConfig::new()
71+
.set_bool("datafusion.execution.parquet.pushdown_filters", true)
72+
.with_target_partitions(2);
73+
let ctx = SessionContext::new_with_config(config);
74+
let dir = TempDir::new().unwrap();
75+
register_table(&ctx, &dir, "build", vec![("k", vec![1, 2, 3])]).await;
76+
let k: Vec<i64> = (0..100).collect();
77+
let v = k.iter().map(|k| k * 10).collect();
78+
register_table(&ctx, &dir, "probe", vec![("k", k), ("v", v)]).await;
79+
(ctx, dir)
80+
}
81+
82+
/// The conjuncts of the filter of each Parquet scan in `plan`, with their
83+
/// optional flag.
84+
fn scan_conjuncts(plan: &Arc<dyn ExecutionPlan>) -> Vec<Vec<(String, bool)>> {
85+
let mut scans = vec![];
86+
plan.apply(|node| {
87+
if let Some(source) = node
88+
.downcast_ref::<DataSourceExec>()
89+
.and_then(|exec| exec.data_source().downcast_ref::<FileScanConfig>())
90+
.and_then(|config| config.file_source().downcast_ref::<ParquetSource>())
91+
{
92+
scans.push(
93+
source
94+
.physical_filter()
95+
.conjuncts()
96+
.iter()
97+
.map(|c| (c.to_string(), c.is_optional()))
98+
.collect(),
99+
);
100+
}
101+
Ok(TreeNodeRecursion::Continue)
102+
})
103+
.unwrap();
104+
scans
105+
}
106+
107+
/// Plan and run `sql`. Return the conjuncts of each Parquet scan and the
108+
/// formatted result.
109+
async fn run(sql: &str) -> (Vec<Vec<(String, bool)>>, String) {
110+
let (ctx, _dir) = context().await;
111+
let plan = ctx
112+
.sql(sql)
113+
.await
114+
.unwrap()
115+
.create_physical_plan()
116+
.await
117+
.unwrap();
118+
let conjuncts = scan_conjuncts(&plan);
119+
let batches = collect(plan, ctx.task_ctx()).await.unwrap();
120+
(
121+
conjuncts,
122+
pretty_format_batches(&batches).unwrap().to_string(),
123+
)
124+
}
125+
126+
fn is_dynamic_filter(conjunct: &(String, bool)) -> bool {
127+
conjunct.0.starts_with("DynamicFilter")
128+
}
129+
130+
#[tokio::test]
131+
async fn hash_join_dynamic_filter_arrives_optional() {
132+
let (scans, result) = run("SELECT p.k, p.v FROM build b JOIN probe p ON b.k = p.k \
133+
WHERE p.v > 10 ORDER BY p.k")
134+
.await;
135+
// The build side scan has no filter. The probe side scan has the
136+
// required user filter and the optional dynamic filter of the join.
137+
let probe = scans
138+
.iter()
139+
.find(|conjuncts| conjuncts.iter().any(is_dynamic_filter))
140+
.expect("a scan with the dynamic filter of the join");
141+
assert!(
142+
probe
143+
.iter()
144+
.all(|conjunct| conjunct.1 == is_dynamic_filter(conjunct)),
145+
"only the dynamic filter is optional: {probe:?}"
146+
);
147+
assert!(
148+
probe
149+
.iter()
150+
.any(|(expr, optional)| expr == "v@1 > 10" && !optional)
151+
);
152+
assert_snapshot!(result, @r"
153+
+---+----+
154+
| k | v |
155+
+---+----+
156+
| 2 | 20 |
157+
| 3 | 30 |
158+
+---+----+
159+
");
160+
}
161+
162+
#[tokio::test]
163+
async fn topk_dynamic_filter_arrives_optional() {
164+
let (scans, result) = run("SELECT v FROM probe ORDER BY v DESC LIMIT 2").await;
165+
assert_eq!(scans.len(), 1);
166+
let [conjunct] = scans[0].as_slice() else {
167+
panic!("expected one conjunct: {scans:?}");
168+
};
169+
assert!(is_dynamic_filter(conjunct) && conjunct.1, "{conjunct:?}");
170+
assert_snapshot!(result, @r"
171+
+-----+
172+
| v |
173+
+-----+
174+
| 990 |
175+
| 980 |
176+
+-----+
177+
");
178+
}
179+
180+
#[tokio::test]
181+
async fn aggregate_dynamic_filter_arrives_optional() {
182+
// The `WHERE` clause prevents an answer from the file statistics.
183+
let (scans, result) = run("SELECT max(v) FROM probe WHERE k > 5").await;
184+
assert_eq!(scans.len(), 1);
185+
let [required, optional] = scans[0].as_slice() else {
186+
panic!("expected two conjuncts: {scans:?}");
187+
};
188+
assert_eq!(required, &("k@0 > 5".to_string(), false));
189+
assert!(is_dynamic_filter(optional) && optional.1, "{optional:?}");
190+
assert_snapshot!(result, @r"
191+
+--------------+
192+
| max(probe.v) |
193+
+--------------+
194+
| 990 |
195+
+--------------+
196+
");
197+
}

‎datafusion/core/tests/physical_optimizer/filter_pushdown.rs‎

Lines changed: 80 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,9 @@ use datafusion_common::{
4040
tree_node::{TreeNode, TreeNodeRecursion},
4141
};
4242
use datafusion_datasource::{
43-
PartitionedFile, file_groups::FileGroup, file_scan_config::FileScanConfigBuilder,
43+
PartitionedFile,
44+
file_groups::FileGroup,
45+
file_scan_config::{FileScanConfig, FileScanConfigBuilder},
4446
};
4547
use datafusion_execution::object_store::ObjectStoreUrl;
4648
use datafusion_expr::ScalarUDF;
@@ -874,6 +876,31 @@ fn test_node_handles_child_pushdown_result() {
874876
);
875877
}
876878

879+
/// The optimizer keeps the properties of each parent filter (for example, the
880+
/// optional flag) by position. A node that returns its parent filters in a
881+
/// different order breaks this rule, and debug builds report an error.
882+
#[test]
883+
fn test_reordered_parent_filters_are_rejected() {
884+
let scan = TestScanBuilder::new(schema()).with_support(true).build();
885+
let predicate = conjunction([
886+
col_lit_predicate("a", "foo", &schema()),
887+
col_lit_predicate("b", "bar", &schema()),
888+
]);
889+
let node = TestNode::new(false, scan, col_lit_predicate("a", "baz", &schema()))
890+
.with_reversed_parent_filters();
891+
let plan = Arc::new(FilterExec::try_new(predicate, Arc::new(node)).unwrap());
892+
let result = FilterPushdown::new().optimize(plan, &ConfigOptions::default());
893+
if cfg!(debug_assertions) {
894+
let err = result.unwrap_err().to_string();
895+
assert!(
896+
err.contains(
897+
"expected TestInsertExec to return the parent filters in input order"
898+
),
899+
"{err}"
900+
);
901+
}
902+
}
903+
877904
// Not portable to sqllogictest: requires manually constructing
878905
// `SortExec(CoalescePartitionsExec(scan))`. A SQL `ORDER BY ... LIMIT` over a
879906
// multi-partition scan plans as `SortPreservingMergeExec(SortExec(scan))`
@@ -1536,6 +1563,58 @@ async fn test_hashjoin_dynamic_filter_pushdown_collect_left() {
15361563
);
15371564
}
15381565

1566+
/// The hash join marks its dynamic filter as optional. `TestSource` does not
1567+
/// override `FileSource::try_pushdown_filter`, so the default method gives it
1568+
/// the dynamic filter through `try_pushdown_filters`, without the flag. The
1569+
/// source applies it as a required filter, and the result is correct.
1570+
#[tokio::test]
1571+
async fn test_hashjoin_optional_dynamic_filter_default_source() {
1572+
let (build_side_schema, build_scan, probe_side_schema, probe_scan) =
1573+
hashjoin_pushdown_scans();
1574+
let on = vec![(
1575+
col("a", &build_side_schema).unwrap(),
1576+
col("a", &probe_side_schema).unwrap(),
1577+
)];
1578+
let plan = Arc::new(
1579+
HashJoinExec::try_new(
1580+
build_scan,
1581+
probe_scan,
1582+
on,
1583+
None,
1584+
&JoinType::Inner,
1585+
None,
1586+
PartitionMode::CollectLeft,
1587+
datafusion_common::NullEquality::NullEqualsNothing,
1588+
false,
1589+
)
1590+
.unwrap(),
1591+
) as Arc<dyn ExecutionPlan>;
1592+
1593+
let mut config = ConfigOptions::default();
1594+
config.execution.parquet.pushdown_filters = true;
1595+
config.optimizer.enable_dynamic_filter_pushdown = true;
1596+
let (plan, batches) = optimize_and_collect_pushdown_plan(plan, config).await;
1597+
1598+
let probe_filter = plan.children()[1]
1599+
.downcast_ref::<DataSourceExec>()
1600+
.unwrap()
1601+
.data_source()
1602+
.downcast_ref::<FileScanConfig>()
1603+
.unwrap()
1604+
.file_source()
1605+
.filter()
1606+
.unwrap();
1607+
insta::assert_snapshot!(probe_filter, @"DynamicFilter [ a@0 >= aa AND a@0 <= ab AND a@0 IN (SET) ([aa, ab]) ]");
1608+
insta::assert_snapshot!(pretty_format_batches(&batches).unwrap(), @r"
1609+
+----+----+-----+----+----+-----+
1610+
| a | b | c | a | b | e |
1611+
+----+----+-----+----+----+-----+
1612+
| aa | ba | 1.0 | aa | ba | 1.0 |
1613+
| ab | bb | 2.0 | ab | bb | 2.0 |
1614+
+----+----+-----+----+----+-----+
1615+
");
1616+
}
1617+
15391618
#[test]
15401619
fn test_hashjoin_parent_filter_pushdown_same_column_names() {
15411620
let build_side_schema = Arc::new(Schema::new(vec![

‎datafusion/core/tests/physical_optimizer/pushdown_utils.rs‎

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -452,6 +452,9 @@ pub(crate) struct TestNode {
452452
inject_filter: bool,
453453
input: Arc<dyn ExecutionPlan>,
454454
predicate: Arc<dyn PhysicalExpr>,
455+
/// Return the parent filters in reverse order (a bug) from
456+
/// `gather_filters_for_pushdown`.
457+
reverse_parent_filters: bool,
455458
}
456459

457460
impl TestNode {
@@ -464,8 +467,16 @@ impl TestNode {
464467
inject_filter,
465468
input,
466469
predicate,
470+
reverse_parent_filters: false,
467471
}
468472
}
473+
474+
/// Make `gather_filters_for_pushdown` return the parent filters in
475+
/// reverse order, which breaks the filter pushdown contract.
476+
pub fn with_reversed_parent_filters(mut self) -> Self {
477+
self.reverse_parent_filters = true;
478+
self
479+
}
469480
}
470481

471482
impl DisplayAs for TestNode {
@@ -497,11 +508,11 @@ impl ExecutionPlan for TestNode {
497508
_: ReplaceChildrenOptions,
498509
) -> Result<Arc<dyn ExecutionPlan>> {
499510
assert_eq!(children.len(), 1);
500-
Ok(Arc::new(TestNode::new(
501-
self.inject_filter,
502-
children[0].clone(),
503-
self.predicate.clone(),
504-
)))
511+
Ok(Arc::new(TestNode {
512+
input: children[0].clone(),
513+
predicate: self.predicate.clone(),
514+
..*self
515+
}))
505516
}
506517

507518
fn with_new_children(
@@ -525,9 +536,12 @@ impl ExecutionPlan for TestNode {
525536
fn gather_filters_for_pushdown(
526537
&self,
527538
_phase: FilterPushdownPhase,
528-
parent_filters: Vec<Arc<dyn PhysicalExpr>>,
539+
mut parent_filters: Vec<Arc<dyn PhysicalExpr>>,
529540
_config: &ConfigOptions,
530541
) -> Result<FilterDescription> {
542+
if self.reverse_parent_filters {
543+
parent_filters.reverse();
544+
}
531545
// Since TestNode marks all parent filters as supported and adds its own filter,
532546
// we use from_child to create a description with all parent filters supported
533547
let child = &self.input;

0 commit comments

Comments
 (0)