adriangb commented on code in PR #26249:
URL: https://github.com/apache/datafusion/pull/26249#discussion_r4300168101
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -685,6 +712,20 @@ impl SharedBuildAccumulator {
}) => {
real_partition_ids.push(partition_id);
keys_have_null |= partition_keys_have_null;
+ lookups = match (&pushdown, lookups) {
+ (PushdownStrategy::Map(map, _), Some(mut lookups))
+ if matches!(map.as_ref(), Map::HashMap(_)) =>
+ {
+
lookups.push(PartitionLookup::Map(Arc::clone(map)));
+ Some(lookups)
+ }
+ _ => None,
Review Comment:
This PR falls back to CASE at high `target_partitions`. That is where the
CASE path costs the most: its cost grows about linearly with N (microbenchmark:
about 6 / 50 / 150 ns/row at N=1 / 16 / 128).
Since this arm returns `None` as soon as one partition is not a hash map,
and at high N each partition gets few keys, so some partitions pick InList (the
per-partition 150-distinct threshold). Then the whole filter falls back to
`CASE hash_repartition % N`.
MRE (uses `probe.parquet` from the other comment):
```sql
SET datafusion.execution.target_partitions = 128;
SET datafusion.optimizer.hash_join_single_partition_threshold = 0;
SET datafusion.optimizer.hash_join_single_partition_threshold_rows = 0;
EXPLAIN ANALYZE SELECT count(*)
FROM (SELECT value AS k FROM generate_series(1, 19200)) b
JOIN probe p ON b.k * 97 = p.k;
```
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -729,6 +773,22 @@ impl SharedBuildAccumulator {
// Only one build partition contains rows, so its filter covers
every
// possible probe match without routing.
Arc::clone(&partition_filters[real_partition_ids[0]])
+ } else if let Some(lookups) = lookups.filter(|lookups| {
+ self.probe_range_partitioning.is_none()
+ && lookups.iter().any(|l| matches!(l, PartitionLookup::Map(_)))
+ }) {
+ // Every partition is a hash table, so route rows to the right
table
+ // directly instead of through a CASE over per-partition filters.
+ let membership = Arc::new(PartitionedLookupExpr::new(
+ self.on_right.clone(),
+ self.repartition_random_state.clone(),
+ HASH_JOIN_SEED,
+ lookups,
+ )) as Arc<dyn PhysicalExpr>;
+ let bounds = merged_bounds
Review Comment:
The merged bounds ignore canceled partitions; is that intended? It looks
safe because a canceled partition’s rows have no consumer.
A `CanceledUnknown` partition pushes `PartitionLookup::Const(true)` but adds
nothing to the bounds. The bounds predicate is then ANDed for every probe row.
So a row routed to the canceled partition is removed if its key is outside the
other partitions' range. That row can have a match in the canceled partition.
The CASE path keeps these rows, because each partition has its own bounds
and a canceled partition gives `true`.
Possible fix (needs checking): do not use the merged bounds when a partition
is canceled.
```suggestion
let bounds = merged_bounds.filter(|_| !has_canceled_unknown)
```
##########
datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs:
##########
@@ -458,6 +456,209 @@ impl PhysicalExpr for HashTableLookupExpr {
}
}
+/// Build-side lookup for one partition of a `Partitioned` hash join.
+pub(crate) enum PartitionLookup {
+ /// The partition's hash table; probe rows routed here are looked up in it.
+ Map(Arc<Map>),
+ /// The partition is known empty (`false`) or its contents are unknown
(`true`).
+ Const(bool),
+}
+
+/// Routes each probe row to the hash table of the build partition that would
+/// contain it and looks the row up there.
+///
+/// This replaces `CASE hash(keys) % n WHEN i THEN lookup_i ... END` for
+/// partitioned joins. `CaseExpr` evaluates each branch by filtering the batch
+/// down to that branch's rows and merging the partial results, which copies
+/// every column of the batch once per branch. This expression computes the
+/// routing and lookup hashes once, buckets row indices by partition and only
+/// gathers the 8-byte hashes, so the batch columns are never copied.
+pub(crate) struct PartitionedLookupExpr {
+ on_columns: Vec<PhysicalExprRef>,
+ routing_state: SeededRandomState,
+ lookup_state: SeededRandomState,
+ partitions: Vec<PartitionLookup>,
Review Comment:
Related to the pruning regression (see the comment on `shared_bounds.rs`):
this expression has no equivalent of `HashTableLookupExpr::pruning_bitmap()`,
so `PruningPredicate` cannot use it.
Suggestion:
- Store one combined `Option<Arc<KeyRangeBitmap>>` here. It is the union of
the per-partition bitmaps. If any partition has no bitmap (or is
`Const(true)`), use `None`.
- Expose it with `pruning_bitmap()`, and let `pruning_predicate.rs` handle
`PartitionedLookupExpr` the same way as `HashTableLookupExpr`.
```rust
pub(crate) struct PartitionedLookupExpr {
on_columns: Vec<PhysicalExprRef>,
routing_state: SeededRandomState,
lookup_state: SeededRandomState,
partitions: Vec<PartitionLookup>,
/// Union of the per-partition key-range bitmaps, for row group and page
pruning.
pruning_bitmap: Option<Arc<KeyRangeBitmap>>,
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
