Skip to content

Commit 3eb7cf7

Browse files
adriangbclaude
andcommitted
perf: collapse partitioned hash join InList dynamic filters into one IN list
In partitioned mode the hash join pushes a routed filter: CASE hash_repartition % N WHEN i THEN bounds_i AND key IN (list_i) ... END When every non-empty build partition pushes an InList and no partition is canceled, the routing is redundant. Routing is a deterministic function of the join keys, so a build row with key K is in the partition that a probe row with key K routes to. A test of K against the union of all lists therefore accepts the same rows as the routed CASE. The per-partition bounds reject no additional rows either, because every key in a list is inside the bounds of its partition. The filter is now `key IN (union)`. This removes the per-row routing hash from the probe side, and the pruning code can use an InList (up to `max_in_list_size`), which it cannot do with a CASE. The union is capped at 1 MiB. Each partition's list is limited independently by `hash_join_inlist_pushdown_max_size`, so the union grows with the partition count; past the cap the routed CASE, where a probe row checks only one list, stays in use. Partitions that push a hash table, and builds with a canceled partition, also keep the CASE. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1 parent 95bb0a0 commit 3eb7cf7

4 files changed

Lines changed: 209 additions & 8 deletions

File tree

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1010,7 +1010,7 @@ async fn optimize_and_collect_pushdown_plan(
10101010
// Not portable to sqllogictest: this test pins `PartitionMode::Partitioned`
10111011
// by hand-wiring `RepartitionExec(Hash, 12)` on both join sides. A SQL
10121012
// INNER JOIN over small parquet inputs plans as `CollectLeft`, so the
1013-
// per-partition CASE filter this test exercises is not reachable via SQL.
1013+
// partitioned filter this test exercises is not reachable via SQL.
10141014
#[tokio::test]
10151015
async fn test_hashjoin_dynamic_filter_pushdown_partitioned() {
10161016
// Rough sketch of the MRE we're trying to recreate:
@@ -1165,7 +1165,7 @@ async fn test_hashjoin_dynamic_filter_pushdown_partitioned() {
11651165
- RepartitionExec: partitioning=Hash([a@0, b@1], 12), input_partitions=1
11661166
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
11671167
- RepartitionExec: partitioning=Hash([a@0, b@1], 12), input_partitions=1
1168-
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ CASE hash_repartition % 12 WHEN 5 THEN a@0 >= ab AND a@0 <= ab AND b@1 >= bb AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:ab,c1:bb}]) WHEN 8 THEN a@0 >= aa AND a@0 <= aa AND b@1 >= ba AND b@1 <= ba AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}]) ELSE false END ]
1168+
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ struct(a@0, b@1) IN (SET) ([{c0:ab,c1:bb}, {c0:aa,c1:ba}]) ]
11691169
"
11701170
);
11711171

@@ -1369,6 +1369,11 @@ async fn test_hashjoin_dynamic_filter_pushdown_range_partitioned() {
13691369
config.execution.parquet.pushdown_filters = true;
13701370
config.optimizer.enable_dynamic_filter_pushdown = true;
13711371
config.optimizer.preserve_file_partitions = 1;
1372+
// Push hash table lookups instead of `InList`s so the filter keeps the
1373+
// `range_partition` routing: an all-`InList` build collapses into one `InList`.
1374+
config
1375+
.optimizer
1376+
.hash_join_inlist_pushdown_max_distinct_values = 0;
13721377
let (plan, batches) = optimize_and_collect_pushdown_plan(plan, config).await;
13731378

13741379
// Now check what our filter looks like
@@ -1381,7 +1386,7 @@ async fn test_hashjoin_dynamic_filter_pushdown_range_partitioned() {
13811386
- RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1
13821387
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true
13831388
- RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1
1384-
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN a@0 >= aa AND a@0 <= aa AND b@1 >= ba AND b@1 <= ba AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}]) WHEN 1 THEN a@0 >= ab AND a@0 <= ab AND b@1 >= bb AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:ab,c1:bb}]) ELSE false END ]
1389+
- DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN a@0 >= aa AND a@0 <= aa AND b@1 >= ba AND b@1 <= ba AND hash_lookup WHEN 1 THEN a@0 >= ab AND a@0 <= ab AND b@1 >= bb AND b@1 <= bb AND hash_lookup ELSE false END ]
13851390
"
13861391
);
13871392

‎datafusion/physical-plan/src/joins/hash_join/exec.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8641,6 +8641,12 @@ mod tests {
86418641
.options_mut()
86428642
.optimizer
86438643
.enable_dynamic_filter_pushdown = true;
8644+
// Push hash table lookups instead of `InList`s: when every partition pushes
8645+
// an `InList`, the routed `CASE` collapses into a single `InList`.
8646+
session_config
8647+
.options_mut()
8648+
.optimizer
8649+
.hash_join_inlist_pushdown_max_distinct_values = 0;
86448650
let task_ctx =
86458651
Arc::new(TaskContext::default().with_session_config(session_config));
86468652

‎datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs‎

Lines changed: 191 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ use crate::joins::hash_join::partitioned_hash_eval::{
3333
};
3434
use crate::repartition::RangeExpr;
3535
use arrow::array::ArrayRef;
36+
use arrow::compute::concat;
3637
use arrow::datatypes::{DataType, Field, Schema};
3738
use datafusion_common::config::ConfigOptions;
3839
use datafusion_common::{
@@ -272,6 +273,15 @@ pub(crate) struct SharedBuildAccumulator {
272273
null_aware: bool,
273274
}
274275

276+
/// Ceiling on the combined size of the per-partition `InList` arrays that
277+
/// [`SharedBuildAccumulator::union_inlist_filter`] will concatenate.
278+
///
279+
/// Each partition's list is independently capped by
280+
/// `hash_join_inlist_pushdown_max_size`, so without a combined cap the union
281+
/// grows with the partition count. Past this size, keeping the routed `CASE`
282+
/// (where each probe row only probes one list) is the cheaper shape.
283+
const MAX_UNIONED_INLIST_BYTES: usize = 1024 * 1024;
284+
275285
/// Strategy for filter pushdown (decided at collection time)
276286
#[derive(Clone)]
277287
pub(crate) enum PushdownStrategy {
@@ -657,7 +667,19 @@ impl SharedBuildAccumulator {
657667
/// Builds one routed probe-side filter from finalized partitioned build data.
658668
/// Empty partitions reject their routed rows, while canceled partitions stay
659669
/// permissive because their build contents are unknown.
670+
///
671+
/// When every non-empty partition pushes an `InList`, the routed `CASE` is
672+
/// replaced by one `InList` over the union of the lists. See
673+
/// [`Self::union_inlist_filter`].
660674
fn build_partitioned_filter(&self, partitions: Vec<PartitionStatus>) -> Result<()> {
675+
if let Some((filter_expr, keys_have_null)) =
676+
self.union_inlist_filter(&partitions)?
677+
{
678+
return self
679+
.dynamic_filter
680+
.update(self.preserve_probe_nulls(filter_expr, keys_have_null)?);
681+
}
682+
661683
let mut partition_filters = Vec::with_capacity(partitions.len());
662684
let mut real_partition_ids = Vec::new();
663685
let mut empty_partition_ids = Vec::new();
@@ -798,6 +820,73 @@ impl SharedBuildAccumulator {
798820
.update(self.preserve_probe_nulls(filter_expr, keys_have_null)?)
799821
}
800822

823+
/// Collapses an all-`InList` partitioned build into one `InList` over the
824+
/// union of the per-partition lists, instead of a `CASE` that routes each
825+
/// probe row to the list of its partition.
826+
///
827+
/// This is exact, not a relaxation: routing is a deterministic function of
828+
/// the key columns, so every build row with key `K` is in the partition that
829+
/// a probe row with key `K` routes to. A test of `K` against the union thus
830+
/// accepts the same rows as the routed `CASE`. The result does not compute
831+
/// the routing hash for each probe row, and unlike a `CASE`, pruning can use
832+
/// an `InList`.
833+
///
834+
/// The per-partition bounds are dropped: every key in a partition's list is
835+
/// inside that partition's bounds, so the bounds reject no additional rows.
836+
///
837+
/// Returns the filter and whether any build key is NULL, or `None` when the
838+
/// collapse does not apply: fewer than two partitions have rows, a partition
839+
/// is canceled or pushes a hash table, the lists have different types, or
840+
/// the union is larger than [`MAX_UNIONED_INLIST_BYTES`].
841+
fn union_inlist_filter(
842+
&self,
843+
partitions: &[PartitionStatus],
844+
) -> Result<Option<(Arc<dyn PhysicalExpr>, bool)>> {
845+
let mut arrays: Vec<&ArrayRef> = Vec::with_capacity(partitions.len());
846+
let mut total_bytes = 0;
847+
let mut keys_have_null = false;
848+
for partition in partitions {
849+
let PartitionStatus::Reported(PartitionData {
850+
pushdown,
851+
keys_have_null: partition_keys_have_null,
852+
..
853+
}) = partition
854+
else {
855+
return Ok(None);
856+
};
857+
let values = match pushdown {
858+
PushdownStrategy::InList(values) => values,
859+
PushdownStrategy::Empty => continue,
860+
PushdownStrategy::Map(_) => return Ok(None),
861+
};
862+
total_bytes += values.get_array_memory_size();
863+
if total_bytes > MAX_UNIONED_INLIST_BYTES
864+
|| arrays
865+
.first()
866+
.is_some_and(|first| first.data_type() != values.data_type())
867+
{
868+
return Ok(None);
869+
}
870+
keys_have_null |= partition_keys_have_null;
871+
arrays.push(values);
872+
}
873+
874+
// With zero or one non-empty partition, `build_partitioned_filter`
875+
// already skips the `CASE`.
876+
if arrays.len() < 2 {
877+
return Ok(None);
878+
}
879+
880+
let union = concat(&arrays.iter().map(|a| a.as_ref()).collect::<Vec<_>>())?;
881+
let filter_expr = create_membership_predicate(
882+
&self.on_right,
883+
PushdownStrategy::InList(union),
884+
&HASH_JOIN_SEED,
885+
self.probe_schema.as_ref(),
886+
)?;
887+
Ok(filter_expr.map(|expr| (expr, keys_have_null)))
888+
}
889+
801890
/// Keeps probe rows with a NULL key when the join semantics need them.
802891
///
803892
/// The build-side predicate drops probe rows whose key is NULL. A null-aware join
@@ -910,6 +999,7 @@ pub(super) fn completed_partitions_for_test(acc: &SharedBuildAccumulator) -> usi
910999
mod tests {
9111000
use super::*;
9121001

1002+
use crate::joins::join_hash_map::JoinHashMapU32;
9131003
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array};
9141004
use arrow::compute::SortOptions;
9151005
use arrow::record_batch::RecordBatch;
@@ -985,6 +1075,12 @@ mod tests {
9851075
PushdownStrategy::InList(Arc::new(Int32Array::from(values.to_vec())) as ArrayRef)
9861076
}
9871077

1078+
fn map_pushdown() -> PushdownStrategy {
1079+
PushdownStrategy::Map(Arc::new(Map::HashMap(Box::new(
1080+
JoinHashMapU32::with_capacity(1),
1081+
))))
1082+
}
1083+
9881084
fn bounds(min: i32, max: i32) -> PartitionBounds {
9891085
PartitionBounds::new(vec![ColumnBounds::new(
9901086
ScalarValue::Int32(Some(min)),
@@ -1142,6 +1238,77 @@ mod tests {
11421238
assert!(expr.downcast_ref::<CaseExpr>().is_none());
11431239
}
11441240

1241+
#[test]
1242+
fn partitioned_all_inlist_collapses_to_a_single_union_inlist() {
1243+
let acc = make_partitioned_expr_accumulator_for_test(3);
1244+
1245+
acc.build_filter(FinalizeInput::Partitioned(vec![
1246+
reported(in_list(&[1, 4]), bounds(1, 4)),
1247+
reported(in_list(&[2, 5]), bounds(2, 5)),
1248+
reported(PushdownStrategy::Empty, no_bounds()),
1249+
]))
1250+
.unwrap();
1251+
1252+
// Routing is a function of the key, so a probe key can only match the
1253+
// list of the partition it routes to: the union is exact, and neither the
1254+
// `CASE` nor the per-partition bounds are necessary.
1255+
let expr = current_expr(&acc);
1256+
assert_in_list_column_values(&expr, "probe_key", 0, &[1, 4, 2, 5]);
1257+
}
1258+
1259+
#[test]
1260+
fn partitioned_mixed_strategies_keep_the_routing_case() {
1261+
let acc = make_partitioned_expr_accumulator_for_test(2);
1262+
1263+
acc.build_filter(FinalizeInput::Partitioned(vec![
1264+
reported(in_list(&[1, 2]), bounds(1, 2)),
1265+
reported(map_pushdown(), bounds(3, 4)),
1266+
]))
1267+
.unwrap();
1268+
1269+
// One partition needs a hash table lookup, so routing is necessary.
1270+
let expr = current_expr(&acc);
1271+
assert_eq!(case_expr(&expr).when_then_expr().len(), 2);
1272+
}
1273+
1274+
#[test]
1275+
fn partitioned_canceled_partition_keeps_the_routing_case() {
1276+
let acc = make_partitioned_expr_accumulator_for_test(3);
1277+
1278+
acc.build_filter(FinalizeInput::Partitioned(vec![
1279+
reported(in_list(&[1]), no_bounds()),
1280+
reported(in_list(&[2]), no_bounds()),
1281+
PartitionStatus::CanceledUnknown,
1282+
]))
1283+
.unwrap();
1284+
1285+
// The canceled partition's keys are unknown, so the union is incomplete
1286+
// and the rows routed to that partition must stay permissive.
1287+
let expr = current_expr(&acc);
1288+
let case = case_expr(&expr);
1289+
assert_eq!(case.when_then_expr().len(), 2);
1290+
assert_literal_bool(
1291+
case.else_expr().expect("expected permissive fallback"),
1292+
true,
1293+
);
1294+
}
1295+
1296+
#[test]
1297+
fn partitioned_oversized_inlist_union_keeps_the_routing_case() {
1298+
let acc = make_partitioned_expr_accumulator_for_test(2);
1299+
let half = MAX_UNIONED_INLIST_BYTES / size_of::<i32>() / 2 + 1;
1300+
let values = (0..half as i32).collect::<Vec<_>>();
1301+
1302+
acc.build_filter(FinalizeInput::Partitioned(vec![
1303+
reported(in_list(&values), no_bounds()),
1304+
reported(in_list(&values), no_bounds()),
1305+
]))
1306+
.unwrap();
1307+
1308+
let expr = current_expr(&acc);
1309+
assert_eq!(case_expr(&expr).when_then_expr().len(), 2);
1310+
}
1311+
11451312
#[test]
11461313
fn partitioned_canceled_unknown_partitions_keep_unknown_routes_permissive() {
11471314
let acc = make_partitioned_expr_accumulator_for_test(2);
@@ -1497,8 +1664,9 @@ mod tests {
14971664
fn partitioned_null_keys_in_one_partition_widen_whole_routed_filter() {
14981665
let acc = null_equal_partitioned_accumulator(2);
14991666

1667+
// One partition pushes a hash table, so the filter keeps the routed `CASE`.
15001668
acc.build_filter(FinalizeInput::Partitioned(vec![
1501-
reported(in_list(&[1]), no_bounds()),
1669+
reported(map_pushdown(), no_bounds()),
15021670
reported_with_null_keys(in_list(&[2]), no_bounds()),
15031671
]))
15041672
.unwrap();
@@ -1518,6 +1686,28 @@ mod tests {
15181686
);
15191687
}
15201688

1689+
/// The collapsed `InList` must be widened by the NULL flag of every partition,
1690+
/// not only of the first one.
1691+
#[test]
1692+
fn partitioned_null_keys_in_one_partition_widen_union_inlist() {
1693+
let acc = null_equal_partitioned_accumulator(2);
1694+
1695+
acc.build_filter(FinalizeInput::Partitioned(vec![
1696+
reported(in_list(&[1]), no_bounds()),
1697+
reported_with_null_keys(in_list(&[2]), no_bounds()),
1698+
]))
1699+
.unwrap();
1700+
1701+
let expr = current_expr(&acc);
1702+
assert_top_binary_op(&expr, Operator::Or);
1703+
let widened = binary_expr(&expr);
1704+
assert!(
1705+
widened.left().downcast_ref::<IsNullExpr>().is_some(),
1706+
"expected the IS NULL disjunct first, got: {expr}"
1707+
);
1708+
assert_in_list_column_values(widened.right(), "probe_key", 0, &[1, 2]);
1709+
}
1710+
15211711
/// A canceled partition's build content is unknown, so it may hold a NULL key:
15221712
/// the aggregated flag must be permissive even when no partition reported one.
15231713
#[test]

‎datafusion/sqllogictest/test_files/range_partitioning.slt‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1834,10 +1834,10 @@ SELECT range_key, SUM(value) FROM (
18341834

18351835
##########
18361836
# TEST 48: Hash Join Dynamic Filter Pushdown on Compatible Range Inputs
1837-
# The Parquet-backed probe accepts the partition-routed dynamic filter. Matching
1837+
# The Parquet-backed probe accepts the partitioned dynamic filter. Matching
18381838
# Range split points keep build filter i aligned with probe partition i. The
1839-
# build has rows only in partitions 0 and 2, so the sparse runtime CASE includes
1840-
# only those branches and lets empty partitions fall through to false.
1839+
# build has rows only in partitions 0 and 2, and both push an IN list, so the
1840+
# routed CASE collapses into one IN list over the union of their keys.
18411841
##########
18421842

18431843
statement ok
@@ -1878,7 +1878,7 @@ JOIN range_partitioned p ON b.range_key = p.range_key;
18781878
Plan with Metrics
18791879
01)HashJoinExec: mode=Partitioned<slt:ignore>metrics=[output_rows=2,<slt:ignore>]
18801880
02)--DataSourceExec: <slt:ignore>file_type=parquet, predicate=range_key@0 = 5 OR range_key@0 = 20<slt:ignore>metrics=[output_rows=2,<slt:ignore>]
1881-
03)--DataSourceExec: <slt:ignore>file_type=parquet, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN range_key@0 >= 5 AND range_key@0 <= 5 AND range_key@0 IN (SET) ([5]) WHEN 2 THEN range_key@0 >= 20 AND range_key@0 <= 20 AND range_key@0 IN (SET) ([20]) ELSE false END ]<slt:ignore>metrics=[output_rows=2,<slt:ignore>]
1881+
03)--DataSourceExec: <slt:ignore>file_type=parquet, predicate=DynamicFilter [ range_key@0 IN (SET) ([5, 20]) ]<slt:ignore>metrics=[output_rows=2,<slt:ignore>]
18821882

18831883
query III
18841884
SELECT b.range_key, b.value, p.value

0 commit comments

Comments
 (0)