tschwarzinger commented on code in PR #26249:
URL: https://github.com/apache/datafusion/pull/26249#discussion_r4256749966
##########
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>,
+}
+
+impl PartitionedLookupExpr {
+ pub(crate) fn new(
+ on_columns: Vec<PhysicalExprRef>,
+ routing_state: SeededRandomState,
+ lookup_state: SeededRandomState,
+ partitions: Vec<PartitionLookup>,
+ ) -> Self {
+ Self {
+ on_columns,
+ routing_state,
+ lookup_state,
+ partitions,
+ }
+ }
+}
+
+impl std::fmt::Debug for PartitionedLookupExpr {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ let cols = self
+ .on_columns
+ .iter()
+ .map(|e| e.to_string())
+ .collect::<Vec<_>>()
+ .join(", ");
+ write!(
+ f,
+ "partitioned_hash_lookup({cols}, {} partitions)",
+ self.partitions.len()
+ )
+ }
+}
+
+impl Display for PartitionedLookupExpr {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(f, "partitioned_hash_lookup({})", self.partitions.len())
+ }
+}
+
+impl Hash for PartitionedLookupExpr {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.on_columns.dyn_hash(state);
+ self.routing_state.seed().hash(state);
+ self.lookup_state.seed().hash(state);
+ for partition in &self.partitions {
+ match partition {
+ PartitionLookup::Map(map) => Arc::as_ptr(map).hash(state),
+ PartitionLookup::Const(value) => value.hash(state),
+ }
+ }
+ }
+}
+
+impl PartialEq for PartitionedLookupExpr {
+ fn eq(&self, other: &Self) -> bool {
+ self.on_columns == other.on_columns
+ && self.routing_state.seed() == other.routing_state.seed()
+ && self.lookup_state.seed() == other.lookup_state.seed()
+ && self.partitions.len() == other.partitions.len()
+ && self
+ .partitions
+ .iter()
+ .zip(&other.partitions)
+ .all(|(a, b)| match (a, b) {
Review Comment:
Maybe implementing `PartialEq` for `PartitionLookup` makes sense.
##########
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(_)))
Review Comment:
Maybe we can add
`let lookup_expression_beneficial = lookups.iter().any(|l| matches!(l,
PartitionLookup::Map(_)))`
or something similar to make it more obvious why this check is in place. At
first I thought this is a correctness check but it seems like the construction
of the `lookup` vector guarantees that the expression can be created correctly.
##########
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>,
+}
+
+impl PartitionedLookupExpr {
+ pub(crate) fn new(
+ on_columns: Vec<PhysicalExprRef>,
+ routing_state: SeededRandomState,
+ lookup_state: SeededRandomState,
+ partitions: Vec<PartitionLookup>,
+ ) -> Self {
+ Self {
+ on_columns,
+ routing_state,
+ lookup_state,
+ partitions,
+ }
+ }
+}
+
+impl std::fmt::Debug for PartitionedLookupExpr {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ let cols = self
+ .on_columns
+ .iter()
+ .map(|e| e.to_string())
+ .collect::<Vec<_>>()
+ .join(", ");
+ write!(
+ f,
+ "partitioned_hash_lookup({cols}, {} partitions)",
+ self.partitions.len()
+ )
+ }
+}
+
+impl Display for PartitionedLookupExpr {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(f, "partitioned_hash_lookup({})", self.partitions.len())
+ }
+}
+
+impl Hash for PartitionedLookupExpr {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.on_columns.dyn_hash(state);
+ self.routing_state.seed().hash(state);
+ self.lookup_state.seed().hash(state);
+ for partition in &self.partitions {
+ match partition {
+ PartitionLookup::Map(map) => Arc::as_ptr(map).hash(state),
+ PartitionLookup::Const(value) => value.hash(state),
+ }
+ }
+ }
+}
+
+impl PartialEq for PartitionedLookupExpr {
+ fn eq(&self, other: &Self) -> bool {
+ self.on_columns == other.on_columns
+ && self.routing_state.seed() == other.routing_state.seed()
+ && self.lookup_state.seed() == other.lookup_state.seed()
+ && self.partitions.len() == other.partitions.len()
+ && self
+ .partitions
+ .iter()
+ .zip(&other.partitions)
+ .all(|(a, b)| match (a, b) {
+ (PartitionLookup::Map(a), PartitionLookup::Map(b)) => {
+ Arc::ptr_eq(a, b)
+ }
+ (PartitionLookup::Const(a), PartitionLookup::Const(b)) =>
a == b,
+ _ => false,
+ })
+ }
+}
+
+impl Eq for PartitionedLookupExpr {}
+
+impl PhysicalExpr for PartitionedLookupExpr {
+ fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+ Ok(DataType::Boolean)
+ }
+
+ fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+ Ok(false)
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ let keys = evaluate_columns(&self.on_columns, batch)?;
+ let num_rows = batch.num_rows();
+ let num_partitions = self.partitions.len() as u64;
+
+ // Same routing as RepartitionExec: hash(keys) % partitions.
+ let mut routing_hashes = vec![0u64; num_rows];
Review Comment:
Would be interesting if we can re-use this allocation for subsequent batches
as this is a performance-critical expression. But this could be an additional
PR so we can easily compare the effect.
##########
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,
+ };
+ merged_bounds = match (merged_bounds,
real_partition_ids.len()) {
+ (None, 1) => Some(bounds.clone()),
+ (Some(merged), _) => merged.union(&bounds),
Review Comment:
Is it expected that the union fails here or would that be a bug?
In case of the latter, maybe returning an error would be better than
silently dropping it.
--
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]
