From 97db09bfc3fb4e30d16585a02feb3112ea9bebc0 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 4 Nov 2026 07:40:17 -0300 Subject: [PATCH 1/3] perf: route probe rows to ArrayMap build partitions in the merged filter A `Partitioned` hash join merges its partitions' build maps into one routed `hash_lookup` filter. The routed lookup accepted only hash tables, so a join whose partitions built an `ArrayMap` (dense integer keys) fell back to a `CASE` over one filter per partition. The routed lookup now probes an `ArrayMap` partition with the key values of the rows routed to it, because an `ArrayMap` is keyed by value, not hash. Every non-empty partition can now merge. The `CASE` fallback remains only for canceled partitions. The routed lookup fuzz test now builds an `ArrayMap` for about half of the partitions. Co-Authored-By: Claude Opus 5.5 --- .../joins/hash_join/partitioned_hash_eval.rs | 67 ++++++++++++------- .../src/joins/hash_join/shared_bounds.rs | 9 +-- 2 files changed, 42 insertions(+), 34 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs index 66c7c75f616e6..aac100940873c 100644 --- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs @@ -7,7 +7,8 @@ use std::{fmt::Display, hash::Hash, sync::Arc}; use arrow::{ - array::{ArrayRef, BooleanArray, BooleanBufferBuilder, UInt64Array}, + array::{ArrayRef, BooleanArray, BooleanBufferBuilder, UInt32Array, UInt64Array}, + compute::take, datatypes::{DataType, Schema}, record_batch::RecordBatch, }; @@ -352,8 +353,7 @@ impl HashTableLookupExpr { /// for pruning only, or `None` /// /// A single non-empty map is probed directly without routing, so it may be - /// any [`Map`] and `router` is ignored. Otherwise `router` is required and - /// every map must be a [`Map::HashMap`]. + /// any [`Map`] and `router` is ignored. Otherwise `router` is required. /// /// # Public Only for Internal Use: /// `datafusion-proto` tests require this constructor, but it is not part of @@ -375,21 +375,10 @@ impl HashTableLookupExpr { "a lookup over anything but one non-empty map needs a partition router" ); } - (Some(router), _) => { - if maps - .iter() - .flatten() - .any(|map| !matches!(map.as_ref(), Map::HashMap(_))) - { - return internal_err!( - "routed lookups only support hash table partitions" - ); - } - Lookup::Routed { - router, - maps: maps.into(), - } - } + (Some(router), _) => Lookup::Routed { + router, + maps: maps.into(), + }, }; Ok(Self { on_columns, @@ -441,16 +430,23 @@ impl HashTableLookupExpr { for (partition, map) in maps.iter().enumerate() { // An empty partition holds no key, so its rows stay unmatched. let Some(map) = map else { continue }; - let Map::HashMap(map) = map.as_ref() else { - return internal_err!( - "routed lookups only support hash table partitions" - ); - }; let run = starts[partition]..starts[partition + 1]; if run.is_empty() { continue; } - let hits = map.contain_hashes(&sorted_hashes[run.clone()]); + let hits = match map.as_ref() { + Map::HashMap(map) => map.contain_hashes(&sorted_hashes[run.clone()]), + // An array map is keyed by value, not hash, so it probes + // the run's own key values. + Map::ArrayMap(map) => { + let indices = UInt32Array::from(rows[run.clone()].to_vec()); + let run_keys = keys + .iter() + .map(|key| take(key.as_ref(), &indices, None)) + .collect::, _>>()?; + map.contain_keys(&run_keys)? + } + }; for hit in hits.values().set_indices() { matched.set_bit(rows[run.start + hit] as usize, true); } @@ -639,7 +635,9 @@ fn evaluate_columns( #[cfg(test)] mod tests { use super::*; + use crate::joins::array_map::ArrayMap; use crate::joins::join_hash_map::JoinHashMapU32; + use arrow::array::AsArray; use datafusion_common::ScalarValue; use datafusion_physical_expr::expressions::Column; use std::collections::hash_map::DefaultHasher; @@ -1111,7 +1109,7 @@ mod tests { )])); let batch = RecordBatch::try_new(schema, vec![Arc::clone(probe)])?; let matched = expr.evaluate(&batch)?.into_array(batch.num_rows())?; - Ok(arrow::array::AsArray::as_boolean(&matched).clone()) + Ok(matched.as_boolean().clone()) } /// A hash table of `keys`, hashed with `random_state` as the build side @@ -1125,6 +1123,18 @@ mod tests { Ok(Arc::new(Map::HashMap(Box::new(map)))) } + /// An array map of `keys`, as a dense build side gets. + fn array_map(keys: &ArrayRef) -> Result> { + let keys = keys.as_primitive::(); + let key = |value: Option| { + ArrayMap::key_to_u64(&ScalarValue::Int32(value)).expect("Int32 key") + }; + let min = key(arrow::compute::min(keys)); + let max = key(arrow::compute::max(keys)); + let keys: ArrayRef = Arc::new(keys.clone()); + Ok(Arc::new(Map::ArrayMap(ArrayMap::try_new(&keys, min, max)?))) + } + /// One randomly generated scenario, see [`routed_lookup_fuzz`]. fn run_routed_lookup_scenario(seed: u64) -> Result<()> { use crate::repartition::{BatchPartitioner, REPARTITION_RANDOM_STATE}; @@ -1166,7 +1176,12 @@ mod tests { .partition( RecordBatch::try_new(schema, vec![Arc::clone(&build)])?, |partition, batch| { - maps[partition] = Some(hash_table(batch.column(0), &random_state)?); + // A dense partition builds an array map, others a hash table. + maps[partition] = Some(if rng.random_bool(0.5) { + array_map(batch.column(0))? + } else { + hash_table(batch.column(0), &random_state)? + }); Ok(()) }, )?; diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index 321ca4b10db59..a4e95c320824b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -721,14 +721,7 @@ impl SharedBuildAccumulator { }) => { real_partition_ids.push(partition_id); keys_have_null |= partition_keys_have_null; - let mergeable = match &pushdown { - PushdownStrategy::Map(map, _) => { - matches!(map.as_ref(), Map::HashMap(_)) - } - PushdownStrategy::InList(_) => true, - PushdownStrategy::Empty => false, - }; - members = members.filter(|_| mergeable).map(|mut members| { + members = members.map(|mut members| { members.push((partition_id, pushdown.clone(), bounds.clone())); members }); From 3e1673d18326cb0e5cee5f6c8e2bf20f6b3541d8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 9 Nov 2026 07:54:39 -0600 Subject: [PATCH 2/3] fix: keep gap pruning for partitions without a pruning bitmap Array map partitions now join the merged lookup, and a dense partition usually has no pruning bitmap, so it contributes its whole bounds. The merged pruning bitmap returned `None` when every partition contributed only bounds, so a Range-partitioned join stopped pruning the probe containers that fall in the gaps between partitions. `union` already returns `None` when no bucket can be excluded, so the early return goes. This also adds tests that array map partitions reach the merged filter under hash and range routing. The fuzz test keeps hash tables under `force_hash_collisions`, where a hash table matches every key and an array map would disagree with the reference. The remaining changes are doc wording for array maps, `take_arrays`, and a simpler test helper. Co-Authored-By: Claude Opus 5.5 --- .../joins/hash_join/partitioned_hash_eval.rs | 36 +++--- .../src/joins/hash_join/shared_bounds.rs | 111 ++++++++++++++++-- 2 files changed, 119 insertions(+), 28 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs index 6ac20238f056e..023369e05c515 100644 --- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs @@ -8,7 +8,7 @@ use std::{fmt::Display, hash::Hash, sync::Arc}; use arrow::{ array::{ArrayRef, BooleanArray, BooleanBufferBuilder, UInt32Array, UInt64Array}, - compute::take, + compute::take_arrays, datatypes::{DataType, Schema}, record_batch::RecordBatch, }; @@ -281,10 +281,11 @@ pub struct HashTableLookupExpr { enum Lookup { /// One map holding every build key, probed directly with no routing. Single(Arc), - /// One hash table per build partition of a `Partitioned` join, `None` where - /// the partition is empty. A key is only ever built into the partition its - /// key routes to, so each probe row is looked up in that partition's table - /// alone: no `CASE` over per-partition filters, and no merged table. + /// One map (a hash table or an array map) per build partition of a + /// `Partitioned` join, `None` where the partition is empty. A key is only + /// ever built into the partition its key routes to, so each probe row is + /// looked up in that partition's map alone: no `CASE` over per-partition + /// filters, and no merged map. Routed { router: PartitionRouter, maps: Arc<[Option>]>, @@ -418,7 +419,7 @@ impl HashTableLookupExpr { self.pruning_bitmap.as_ref() } - /// Membership of each row of `keys` in the hash table of its routed partition. + /// Membership of each row of `keys` in the map of its routed partition. fn evaluate_routed( &self, keys: &[ArrayRef], @@ -465,11 +466,7 @@ impl HashTableLookupExpr { // the run's own key values. Map::ArrayMap(map) => { let indices = UInt32Array::from(rows[run.clone()].to_vec()); - let run_keys = keys - .iter() - .map(|key| take(key.as_ref(), &indices, None)) - .collect::, _>>()?; - map.contain_keys(&run_keys)? + map.contain_keys(&take_arrays(keys, &indices, None)?)? } }; for hit in hits.values().set_indices() { @@ -1155,14 +1152,13 @@ mod tests { /// An array map of `keys`, as a dense build side gets. fn array_map(keys: &ArrayRef) -> Result> { - let keys = keys.as_primitive::(); + let values = keys.as_primitive::(); let key = |value: Option| { ArrayMap::key_to_u64(&ScalarValue::Int32(value)).expect("Int32 key") }; - let min = key(arrow::compute::min(keys)); - let max = key(arrow::compute::max(keys)); - let keys: ArrayRef = Arc::new(keys.clone()); - Ok(Arc::new(Map::ArrayMap(ArrayMap::try_new(&keys, min, max)?))) + let min = key(arrow::compute::min(values)); + let max = key(arrow::compute::max(values)); + Ok(Arc::new(Map::ArrayMap(ArrayMap::try_new(keys, min, max)?))) } /// One randomly generated scenario, see [`routed_lookup_fuzz`]. @@ -1206,8 +1202,12 @@ mod tests { .partition( RecordBatch::try_new(schema, vec![Arc::clone(&build)])?, |partition, batch| { - // A dense partition builds an array map, others a hash table. - maps[partition] = Some(if rng.random_bool(0.5) { + // Half the partitions build an array map, as dense keys would. Under + // forced collisions a hash table matches every key, so an array + // map would disagree with the single-table reference below. + let dense = + cfg!(not(feature = "force_hash_collisions")) && rng.random_bool(0.5); + maps[partition] = Some(if dense { array_map(batch.column(0))? } else { hash_table(batch.column(0), &random_state)? diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index f226812dd103f..f083566acfcdb 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -857,8 +857,8 @@ impl SharedBuildAccumulator { .update(self.preserve_probe_nulls(filter_expr, keys_have_null)?) } - /// `members` holds each non-empty partition's id, strategy (an IN list or a - /// hash table) and bounds, out of `num_partitions` partitions. + /// `members` holds each non-empty partition's id, strategy (an IN list, a + /// hash table or an array map) and bounds, out of `num_partitions` partitions. fn merged_membership_filter( &self, members: &[(usize, PushdownStrategy, PartitionBounds)], @@ -965,8 +965,7 @@ impl SharedBuildAccumulator { /// Container pruning reads one bitmap, so merge the partitions' key ranges: /// a hash table's bitmap, an IN list's keys, or else a partition's whole - /// bounds. `None` when no partition narrows its bounds, as the bounds - /// predicate already prunes that much. + /// bounds. `None` when nothing can be excluded. fn merged_pruning_bitmap( &self, members: &[(usize, PushdownStrategy, PartitionBounds)], @@ -991,12 +990,6 @@ impl SharedBuildAccumulator { }), }) .collect::>>()?; - if parts - .iter() - .all(|part| matches!(part, KeyRangePart::Range { .. })) - { - return None; - } KeyRangeBitmap::union(&merged.min, &merged.max, &parts).map(Arc::new) } @@ -1488,6 +1481,104 @@ mod tests { assert_eq!(matched.as_boolean(), &BooleanArray::from(expected.to_vec())); } + /// An array map holding `keys`, as a dense build side gets. + fn array_map_of(keys: &[i32]) -> PushdownStrategy { + use crate::joins::array_map::ArrayMap; + + let bound = |key: Option<&i32>| { + ArrayMap::key_to_u64(&ScalarValue::Int32(key.copied())).unwrap() + }; + let min = bound(keys.iter().min()); + let max = bound(keys.iter().max()); + let keys: ArrayRef = Arc::new(Int32Array::from(keys.to_vec())); + let map = ArrayMap::try_new(&keys, min, max).unwrap(); + let map = Arc::new(Map::ArrayMap(map)); + PushdownStrategy::Map(Arc::new(ReservedMap::unreserved(map)), None) + } + + /// Array map partitions merge into one lookup with hash table partitions. + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn partitioned_array_map_and_hash_map_partitions_merge() { + let acc = make_partitioned_expr_accumulator_for_test(4); + let build: Vec = (0..40).step_by(2).collect(); + + acc.build_filter(FinalizeInput::Partitioned(routed_partitions( + &acc, + 4, + &build, + |partition, keys| { + if partition % 2 == 0 { + array_map_of(keys) + } else { + hash_map_of(keys) + } + }, + ))) + .unwrap(); + + assert!( + current_expr(&acc).downcast_ref::().is_none(), + "expected a merged lookup, got: {}", + current_expr(&acc) + ); + let probe: Vec = (0..40).collect(); + let expected: Vec = probe.iter().map(|key| key % 2 == 0).collect(); + assert_matches(&acc, &probe, &expected); + } + + /// Array map partitions merge under Range routing too. + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn partitioned_range_array_map_partitions_merge() -> Result<()> { + let mut acc = make_partitioned_expr_accumulator_for_test(4); + acc.probe_range_partitioning = Some(RangePartitioning::try_new( + [PhysicalSortExpr::new( + Arc::clone(&acc.on_right[0]), + Default::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + )?); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(array_map_of(&[7, 8, 9]), bounds(7, 9)), + reported(PushdownStrategy::Empty, no_bounds()), + reported(array_map_of(&[27, 28, 29]), bounds(27, 29)), + reported(array_map_of(&[37, 38, 39]), bounds(37, 39)), + ]))?; + + assert!(current_expr(&acc).downcast_ref::().is_none()); + let probe: Vec = (0..50).collect(); + let expected: Vec = probe + .iter() + .map(|key| matches!(key, 7..=9 | 27..=29 | 37..=39)) + .collect(); + assert_matches(&acc, &probe, &expected); + Ok(()) + } + + /// Partitions without a bitmap still leave the gaps between them prunable. + #[test] + fn disjoint_array_map_partitions_keep_gap_pruning() { + let acc = make_partitioned_expr_accumulator_for_test(2); + let members = [ + (0, array_map_of(&[7, 8, 9]), bounds(7, 9)), + (1, array_map_of(&[37, 38, 39]), bounds(37, 39)), + ]; + + let bitmap = acc + .merged_pruning_bitmap(&members, Some(&bounds(7, 39))) + .expect("the gap between the partitions can be pruned"); + + assert!(bitmap.may_contain_range(7, 9)); + assert!(!bitmap.may_contain_range(10, 36)); + } + #[test] fn partitioned_one_real_partition_with_rest_empty_skips_case() { let acc = make_partitioned_expr_accumulator_for_test(3); From c3baea962c970defb55539669a8754b875af639e Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 9 Nov 2026 13:12:44 -0700 Subject: [PATCH 3/3] perf: skip probe hashing in the routed lookup when every map is an array map An array map probes by key value, so when no partition holds a hash map the routed lookup computed per-row hashes that nothing read. The lookup now hashes the probe keys only when some partition is a hash map. The counting sort by partition runs either way, and the sorted hashes are gathered from the sorted rows only when hashes exist. A new test builds every partition as an array map and checks exact matches, including a null probe row. Co-Authored-By: Claude Opus 5.5 --- .../joins/hash_join/partitioned_hash_eval.rs | 83 +++++++++++++++++-- 1 file changed, 77 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs index 023369e05c515..e5352882e9377 100644 --- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs @@ -430,9 +430,15 @@ impl HashTableLookupExpr { let mut partition_ids = Vec::with_capacity(num_rows); router.route(keys, maps.len(), &mut partition_ids)?; - with_hashes(keys, self.random_state.random_state(), |hashes| { + // Only a hash map probes by hash. When every map is an array map, no + // probe reads the hashes, so the lookup skips hashing the keys. + let needs_hashes = maps.iter().flatten().any(|map| { + let map: &Map = map; + matches!(map, Map::HashMap(_)) + }); + let probe = |hashes: Option<&[u64]>| -> Result { // Counting sort of the rows by partition, so each table is probed - // once per batch over a contiguous run of its rows' hashes. + // once per batch over a contiguous run of its rows. let mut starts = vec![0usize; maps.len() + 1]; for &id in &partition_ids { starts[id as usize + 1] += 1; @@ -442,13 +448,15 @@ impl HashTableLookupExpr { } let mut next = starts.clone(); let mut rows = vec![0u32; num_rows]; - let mut sorted_hashes = vec![0u64; num_rows]; - for (row, (&id, &hash)) in partition_ids.iter().zip(hashes).enumerate() { + for (row, &id) in partition_ids.iter().enumerate() { let slot = &mut next[id as usize]; rows[*slot] = row as u32; - sorted_hashes[*slot] = hash; *slot += 1; } + let sorted_hashes: Vec = match hashes { + Some(hashes) => rows.iter().map(|&row| hashes[row as usize]).collect(), + None => Vec::new(), + }; let mut matched = BooleanBufferBuilder::new(num_rows); matched.append_n(num_rows, false); @@ -461,6 +469,8 @@ impl HashTableLookupExpr { } let map: &Map = map; let hits = match map { + // `needs_hashes` is true whenever a hash map exists, so + // `sorted_hashes` holds every row here. Map::HashMap(map) => map.contain_hashes(&sorted_hashes[run.clone()]), // An array map is keyed by value, not hash, so it probes // the run's own key values. @@ -474,7 +484,14 @@ impl HashTableLookupExpr { } } Ok(BooleanArray::new(matched.finish(), None)) - }) + }; + if needs_hashes { + with_hashes(keys, self.random_state.random_state(), |hashes| { + probe(Some(hashes)) + }) + } else { + probe(None) + } } } impl std::fmt::Debug for HashTableLookupExpr { @@ -1161,6 +1178,60 @@ mod tests { Ok(Arc::new(Map::ArrayMap(ArrayMap::try_new(keys, min, max)?))) } + /// With every partition built as an array map, the routed lookup skips + /// hashing and still matches each probe row against its own partition. + #[test] + fn routed_lookup_all_array_maps() -> Result<()> { + use crate::repartition::{BatchPartitioner, REPARTITION_RANDOM_STATE}; + + let column: PhysicalExprRef = Arc::new(Column::new("a", 0)); + let schema = Arc::new(Schema::new(vec![arrow::datatypes::Field::new( + "a", + DataType::Int32, + false, + )])); + let build: ArrayRef = + Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3, 5, 8, 13, 21])); + let num_partitions = 4; + let mut maps = vec![None; num_partitions]; + BatchPartitioner::new_hash_partitioner( + vec![Arc::clone(&column)], + num_partitions, + crate::metrics::Time::new(), + )? + .partition( + RecordBatch::try_new(schema, vec![build])?, + |partition, batch| { + maps[partition] = Some(array_map(batch.column(0))?); + Ok(()) + }, + )?; + let routed = HashTableLookupExpr::new( + vec![column], + SeededRandomState::with_seed(1), + Some(PartitionRouter::Hash(REPARTITION_RANDOM_STATE)), + maps, + "routed".to_string(), + None, + )?; + + let probe: ArrayRef = Arc::new(arrow::array::Int32Array::from(vec![ + Some(0), + Some(1), + None, + Some(4), + Some(5), + Some(13), + Some(20), + Some(21), + ])); + assert_eq!( + lookup(&routed, &probe)?, + BooleanArray::from(vec![false, true, false, false, true, true, false, true]) + ); + Ok(()) + } + /// One randomly generated scenario, see [`routed_lookup_fuzz`]. fn run_routed_lookup_scenario(seed: u64) -> Result<()> { use crate::repartition::{BatchPartitioner, REPARTITION_RANDOM_STATE};