diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index 4e56cf013c8f7..edf3a82d1109d 100644 --- a/datafusion/physical-plan/src/joins/array_map.rs +++ b/datafusion/physical-plan/src/joins/array_map.rs @@ -24,7 +24,7 @@ use crate::joins::chain::traverse_chain; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; use arrow::buffer::BooleanBuffer; use arrow::datatypes::ArrowNumericType; -use datafusion_common::{Result, ScalarValue, internal_err}; +use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err, internal_err}; /// A macro to downcast only supported integer types (up to 64-bit) and invoke a generic function. /// @@ -172,7 +172,10 @@ impl ArrayMap { /// Note: This function processes only the non-null values in the input `array`, /// ignoring any rows where the key is `NULL`. /// - pub(crate) fn try_new(array: &ArrayRef, min_val: u64, max_val: u64) -> Result { + /// # Note + /// This is public for internal testing purposes only and is not + /// guaranteed to be stable across versions. + pub fn try_new(array: &ArrayRef, min_val: u64, max_val: u64) -> Result { let range = Self::calculate_range(min_val, max_val); if range >= usize::MAX as u64 { return internal_err!("ArrayMap key range is too large to be allocated."); @@ -416,6 +419,47 @@ impl ArrayMap { }); Ok(BooleanArray::new(buffer, None)) } + + #[cfg(feature = "proto")] + pub(crate) fn to_proto_membership_only( + &self, + ) -> datafusion_proto_models::protobuf::ArrayMapMembership { + use datafusion_proto_models::protobuf; + + let bits = BooleanBuffer::collect_bool(self.data.len(), |i| self.data[i] != 0); + let presence = bits.sliced().as_slice().to_vec(); + protobuf::ArrayMapMembership { + offset: self.offset, + num_slots: self.data.len() as u64, + presence, + } + } + + #[cfg(feature = "proto")] + pub(crate) fn try_from_proto_membership_only( + node: &datafusion_proto_models::protobuf::ArrayMapMembership, + ) -> Result { + use arrow::util::bit_iterator::BitIndexIterator; + + // Assert that i < num_slots is a valid index into node.presence + assert_eq_or_internal_err!( + node.presence.len() as u64, + node.num_slots.div_ceil(8) + ); + + let mut data = vec![0u32; node.num_slots as usize]; + let mut num_of_distinct_key = 0; + for slot in BitIndexIterator::new(&node.presence, 0, node.num_slots as usize) { + data[slot] = 1; + num_of_distinct_key += 1; + } + Ok(Self { + data, + offset: node.offset, + next: Vec::new(), + num_of_distinct_key, + }) + } } #[cfg(test)] 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 60a25fc2efcff..24082abb0d0d3 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 @@ -19,6 +19,7 @@ use std::{fmt::Display, hash::Hash, sync::Arc}; +use arrow::array::BooleanArray; use arrow::{ array::{ArrayRef, UInt64Array}, datatypes::{DataType, Schema}, @@ -33,8 +34,10 @@ use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::physical_expr::{ DynHash, PhysicalExpr, PhysicalExprRef, }; +use hashbrown::HashTable; use crate::joins::Map; +use crate::joins::array_map::ArrayMap; /// RandomState wrapper that preserves the seed used to create it. /// @@ -254,15 +257,17 @@ impl HashExpr { /// Physical expression that checks join keys in a [`Map`] (hash table or array map). /// -/// Returns a [`BooleanArray`](arrow::array::BooleanArray) indicating if join keys (from `on_columns`) exist in the map. +/// Returns a [`BooleanArray`] indicating if join keys (from `on_columns`) exist in the map. // TODO: rename to MapLookupExpr pub struct HashTableLookupExpr { /// Columns in the ON clause used to compute the join key for lookups on_columns: Vec, /// Random state for hashing (with seeds preserved for serialization) random_state: SeededRandomState, - /// Map to check against (hash table or array map) - map: Arc, + /// Map to check against. Deserialized expressions hold a membership-only + /// variant of [`HashTableLookupExprMap`], which supports nothing beyond + /// membership checks. + map: HashTableLookupExprMap, /// Description for display description: String, } @@ -286,7 +291,7 @@ impl HashTableLookupExpr { Self { on_columns, random_state, - map, + map: HashTableLookupExprMap::Normal(map), description, } } @@ -309,30 +314,16 @@ impl Hash for HashTableLookupExpr { self.on_columns.dyn_hash(state); self.description.hash(state); self.random_state.seed().hash(state); - // Note that we compare hash_map by pointer equality. - // Actually comparing the contents of the hash maps would be expensive. - // The way these hash maps are used in actuality is that HashJoinExec creates - // one per partition per query execution, thus it is never possible for two different - // hash maps to have the same content in practice. - // Theoretically this is a public API and users could create identical hash maps, - // but that seems unlikely and not worth paying the cost of deep comparison all the time. - Arc::as_ptr(&self.map).hash(state); + self.map.hash(state); } } impl PartialEq for HashTableLookupExpr { fn eq(&self, other: &Self) -> bool { - // Note that we compare hash_map by pointer equality. - // Actually comparing the contents of the hash maps would be expensive. - // The way these hash maps are used in actuality is that HashJoinExec creates - // one per partition per query execution, thus it is never possible for two different - // hash maps to have the same content in practice. - // Theoretically this is a public API and users could create identical hash maps, - // but that seems unlikely and not worth paying the cost of deep comparison all the time. self.on_columns == other.on_columns && self.description == other.description && self.random_state.seed() == other.random_state.seed() - && Arc::ptr_eq(&self.map, &other.map) + && self.map == other.map } } @@ -353,12 +344,12 @@ impl PhysicalExpr for HashTableLookupExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(HashTableLookupExpr::new( - children, - self.random_state.clone(), - Arc::clone(&self.map), - self.description.clone(), - ))) + Ok(Arc::new(HashTableLookupExpr { + on_columns: children, + random_state: self.random_state.clone(), + map: self.map.clone(), + description: self.description.clone(), + })) } fn data_type(&self, _input_schema: &Schema) -> Result { @@ -373,49 +364,59 @@ impl PhysicalExpr for HashTableLookupExpr { // Evaluate columns let join_keys = evaluate_columns(&self.on_columns, batch)?; - match self.map.as_ref() { - Map::HashMap(map) => { + match &self.map { + HashTableLookupExprMap::Normal(normal) => match &**normal { + Map::HashMap(map) => { + with_hashes(&join_keys, self.random_state.random_state(), |hashes| { + let array = map.contain_hashes(hashes); + Ok(ColumnarValue::Array(Arc::new(array))) + }) + } + Map::ArrayMap(map) => { + let array = map.contain_keys(&join_keys)?; + Ok(ColumnarValue::Array(Arc::new(array))) + } + }, + HashTableLookupExprMap::MembershipOnlyHashMap(map) => { with_hashes(&join_keys, self.random_state.random_state(), |hashes| { let array = map.contain_hashes(hashes); Ok(ColumnarValue::Array(Arc::new(array))) }) } - Map::ArrayMap(map) => { + HashTableLookupExprMap::MembershipOnlyArrayMap(map) => { let array = map.contain_keys(&join_keys)?; Ok(ColumnarValue::Array(Arc::new(array))) } } } + #[cfg(feature = "proto")] fn try_to_proto( &self, - _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, ) -> Result> { use datafusion_proto_models::protobuf; use datafusion_proto_models::protobuf::physical_expr_node::ExprType; // HashTableLookupExpr holds a runtime Arc (the build-side hash - // table) that cannot be serialized, so it is replaced with lit(true). - // - // Dynamic filtering is a performance optimisation only — replacing the - // lookup with lit(true) preserves correctness by allowing all rows - // through. - // - // If a plan is serialized before execution, HashTableLookupExpr is not - // yet present in the dynamic filter expression. - // - // If a plan is serialized after execution, any runtime-created - // HashTableLookupExpr is replaced during serialization. Re-executing - // the plan requires reset_state(), after which HashJoinExec rebuilds - // fresh dynamic filters at runtime. - let value = datafusion_proto_common::ScalarValue { - value: Some(datafusion_proto_common::scalar_value::Value::BoolValue( - true, - )), - }; + // table). This can be serialized, but only in a way that maintains + // the set membership of the Map. A round-tripped map is not useable for + // anything beyond expression evaluation + + let on_columns = ctx.encode_children_expressions(&self.on_columns)?; + let map = try_map_to_proto_membership_only(&self.map)?; + + let expr = + ExprType::HashTableLookupExpr(protobuf::PhysicalHashTableLookupExprNode { + on_columns, + seed0: self.random_state.seed, + description: self.description.clone(), + map: Some(map), + }); + Ok(Some(protobuf::PhysicalExprNode { expr_id: None, - expr_type: Some(ExprType::Literal(value)), + expr_type: Some(expr), })) } fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -423,6 +424,96 @@ impl PhysicalExpr for HashTableLookupExpr { } } +/// Encode the map in its membership-only proto form. +#[cfg(feature = "proto")] +fn try_map_to_proto_membership_only( + map: &HashTableLookupExprMap, +) -> Result +{ + use datafusion_proto_models::protobuf; + + match map { + HashTableLookupExprMap::Normal(normal) => match &**normal { + Map::ArrayMap(array_map) => Ok( + protobuf::physical_hash_table_lookup_expr_node::Map::ArrayMapMembership( + array_map.to_proto_membership_only(), + ), + ), + Map::HashMap(hash_map) => Ok( + protobuf::physical_hash_table_lookup_expr_node::Map::HashMapMembership( + protobuf::HashMapMembership { + build_hashes: hash_map.hashes(), + }, + ), + ), + }, + HashTableLookupExprMap::MembershipOnlyHashMap(hash_map) => Ok( + protobuf::physical_hash_table_lookup_expr_node::Map::HashMapMembership( + protobuf::HashMapMembership { + build_hashes: hash_map.hashes(), + }, + ), + ), + HashTableLookupExprMap::MembershipOnlyArrayMap(array_map) => Ok( + protobuf::physical_hash_table_lookup_expr_node::Map::ArrayMapMembership( + array_map.0.to_proto_membership_only(), + ), + ), + } +} + +#[cfg(feature = "proto")] +impl HashTableLookupExpr { + /// Reconstruct a [`HashTableLookupExpr`] from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`], the exact inverse of what + /// [`PhysicalExpr::try_to_proto`] produces, so every expression's + /// `try_from_proto` shares one signature. Child sub-expressions are + /// decoded recursively via [`PhysicalExprDecodeCtx::decode`]. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + /// [`PhysicalExpr::try_to_proto`]: datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto + /// [`PhysicalExprDecodeCtx::decode`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::decode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf::{ + physical_expr_node::ExprType, physical_hash_table_lookup_expr_node::Map, + }; + + let hash_table_lookup_expr = match &node.expr_type { + Some(ExprType::HashTableLookupExpr(h)) => h, + _ => return internal_err!("PhysicalExprNode is not a HashTableLookupExpr"), + }; + let on_columns = + ctx.decode_children_expressions(&hash_table_lookup_expr.on_columns)?; + + let map = match &hash_table_lookup_expr.map { + Some(Map::HashMapMembership(membership)) => { + HashTableLookupExprMap::MembershipOnlyHashMap(Arc::new( + JoinHashMembershipMap::new(&membership.build_hashes), + )) + } + Some(Map::ArrayMapMembership(membership)) => { + HashTableLookupExprMap::MembershipOnlyArrayMap(Arc::new( + JoinMembershipArrayMap(ArrayMap::try_from_proto_membership_only( + membership, + )?), + )) + } + None => return internal_err!("HashTableLookupExpr has no map"), + }; + + Ok(Arc::new(HashTableLookupExpr { + on_columns, + random_state: SeededRandomState::with_seed(hash_table_lookup_expr.seed0), + map, + description: hash_table_lookup_expr.description.clone(), + })) + } +} + fn evaluate_columns( columns: &[PhysicalExprRef], batch: &RecordBatch, @@ -434,6 +525,92 @@ fn evaluate_columns( .collect() } +#[derive(Clone)] +enum HashTableLookupExprMap { + /// A regular build-side hash map + Normal(Arc), + /// A membership-checks only version of a hashmap, only constructed via expression deserialization + MembershipOnlyHashMap(Arc), + /// A membership-checks only version of an array map, only constructed via expression deserialization + MembershipOnlyArrayMap(Arc), +} +impl Hash for HashTableLookupExprMap { + // Note that we compare hash_map by pointer equality. + // Actually comparing the contents of the hash maps would be expensive. + // The way these hash maps are used in actuality is that HashJoinExec creates + // one per partition per query execution, thus it is never possible for two different + // hash maps to have the same content in practice. + // Theoretically this is a public API and users could create identical hash maps, + // but that seems unlikely and not worth paying the cost of deep comparison all the time. + + fn hash(&self, state: &mut H) { + match self { + Self::Normal(a) => Arc::as_ptr(a).hash(state), + Self::MembershipOnlyHashMap(a) => Arc::as_ptr(a).hash(state), + Self::MembershipOnlyArrayMap(a) => Arc::as_ptr(a).hash(state), + } + } +} +impl PartialEq for HashTableLookupExprMap { + // Note that we compare hash_map by pointer equality. + // Actually comparing the contents of the hash maps would be expensive. + // The way these hash maps are used in actuality is that HashJoinExec creates + // one per partition per query execution, thus it is never possible for two different + // hash maps to have the same content in practice. + // Theoretically this is a public API and users could create identical hash maps, + // but that seems unlikely and not worth paying the cost of deep comparison all the time. + + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Normal(a), Self::Normal(b)) => Arc::ptr_eq(a, b), + (Self::MembershipOnlyHashMap(a), Self::MembershipOnlyHashMap(b)) => { + Arc::ptr_eq(a, b) + } + (Self::MembershipOnlyArrayMap(a), Self::MembershipOnlyArrayMap(b)) => { + Arc::ptr_eq(a, b) + } + _ => false, + } + } +} + +/// Membership-only join map reconstructed from serialized distinct hashes. +/// Supports `contain_hashes` lookups only; it has no build-side rows. +struct JoinHashMembershipMap { + map: HashTable<(u64, ())>, +} +impl JoinHashMembershipMap { + fn new(hashes: &[u64]) -> Self { + let mut map = HashTable::with_capacity(hashes.len()); + + // Wire input is trusted to contain distinct hashes. A duplicate would + // insert a duplicate entry, which wastes memory and skews `len()` but + // cannot affect membership results. + for h in hashes { + map.insert_unique(*h, (*h, ()), |(h, _)| *h); + } + + Self { map } + } + + fn contain_hashes(&self, hashes: &[u64]) -> BooleanArray { + crate::joins::join_hash_map::contain_hashes(&self.map, hashes) + } + + fn hashes(&self) -> Vec { + self.map.iter().map(|(h, _)| *h).collect() + } +} + +/// Wrapper type for ArrayMap to restrict it to membership checks only +struct JoinMembershipArrayMap(ArrayMap); +impl JoinMembershipArrayMap { + #[inline] + fn contain_keys(&self, keys: &[ArrayRef]) -> Result { + self.0.contain_keys(keys) + } +} + #[cfg(test)] mod tests { use super::*; @@ -552,6 +729,8 @@ mod tests { #[cfg(feature = "proto")] mod proto_tests { use super::*; + use crate::joins::join_hash_map::JoinHashMapU64; + use arrow::array::Int64Array; use arrow::datatypes::{DataType, Field}; use datafusion_common::internal_datafusion_err; use datafusion_physical_expr_common::physical_expr::proto_decode::{ @@ -713,6 +892,465 @@ mod tests { "{err}" ); } + + #[test] + fn hash_table_lookup_expr_try_to_proto_hash_map_membership() { + let build_hashes = [7u64, 42, 9999]; + let map = Arc::new(Map::HashMap(Box::new(join_hash_map_with_hashes( + &build_hashes, + )))); + let expr = HashTableLookupExpr::new( + vec![Arc::new(Column::new("a", 0))], + SeededRandomState::with_seed(42), + map, + "hash_lookup".to_string(), + ); + let encoder = TestEncoder; + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let proto = expr.try_to_proto(&ctx).unwrap().unwrap(); + + assert_eq!(proto.expr_id, None); + let node = match proto.expr_type.unwrap() { + protobuf::physical_expr_node::ExprType::HashTableLookupExpr(node) => node, + other => panic!("expected HashTableLookupExpr, got {other:?}"), + }; + assert_eq!(node.seed0, 42); + assert_eq!(node.description, "hash_lookup"); + assert_eq!(node.on_columns.len(), 1); + let membership = match node.map.unwrap() { + protobuf::physical_hash_table_lookup_expr_node::Map::HashMapMembership( + m, + ) => m, + other => panic!("expected HashMapMembership, got {other:?}"), + }; + // Hash table iteration order is arbitrary; compare as a sorted set. + let mut got = membership.build_hashes; + got.sort_unstable(); + assert_eq!(got, build_hashes); + } + + #[test] + fn hash_table_lookup_expr_try_to_proto_array_map_membership() { + let build: ArrayRef = Arc::new(Int64Array::from(vec![10i64, 12, 10, 15])); + let array_map = ArrayMap::try_new(&build, 10, 15).unwrap(); + let expr = HashTableLookupExpr::new( + vec![Arc::new(Column::new("a", 0))], + SeededRandomState::with_seed(42), + Arc::new(Map::ArrayMap(array_map)), + "hash_lookup".to_string(), + ); + let encoder = TestEncoder; + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let proto = expr.try_to_proto(&ctx).unwrap().unwrap(); + + let node = match proto.expr_type.unwrap() { + protobuf::physical_expr_node::ExprType::HashTableLookupExpr(node) => node, + other => panic!("expected HashTableLookupExpr, got {other:?}"), + }; + let membership = match node.map.unwrap() { + protobuf::physical_hash_table_lookup_expr_node::Map::ArrayMapMembership( + m, + ) => m, + other => panic!("expected ArrayMapMembership, got {other:?}"), + }; + assert_eq!(membership.offset, 10); + assert_eq!(membership.num_slots, 6); + // Keys 10, 12, 15 occupy slots 0, 2, 5 (LSB-first bit order). + assert_eq!(membership.presence, vec![0b0010_0101u8]); + } + + #[test] + fn hash_table_lookup_expr_try_from_proto_hash_map_membership() { + let schema = lookup_schema(); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + let proto = lookup_expr_proto( + protobuf::physical_hash_table_lookup_expr_node::Map::HashMapMembership( + protobuf::HashMapMembership { + build_hashes: hashes_for(&[1, 3], 42), + }, + ), + ); + + let expr = HashTableLookupExpr::try_from_proto(&proto, &ctx).unwrap(); + let expr = expr.downcast_ref::().unwrap(); + + assert_eq!(expr.random_state.seed(), 42); + assert_eq!(expr.description, "hash_lookup"); + assert_eq!(expr.on_columns.len(), 1); + assert!(matches!( + expr.map, + HashTableLookupExprMap::MembershipOnlyHashMap(_) + )); + // Probe keys are hashed with the deserialized seed, so 1 and 3 + // (whose hashes were serialized) match and 2 and 4 do not. + // Under force_hash_collisions every value hashes identically, + // so the filter degenerates to all-true and this doesn't hold. + #[cfg(not(feature = "force_hash_collisions"))] + assert_eq!( + eval_lookup(expr, &probe_batch(&[1, 2, 3, 4])), + [true, false, true, false] + ); + } + + #[test] + fn hash_table_lookup_expr_try_from_proto_array_map_membership() { + let schema = lookup_schema(); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + let proto = lookup_expr_proto( + protobuf::physical_hash_table_lookup_expr_node::Map::ArrayMapMembership( + protobuf::ArrayMapMembership { + offset: 10, + num_slots: 6, + presence: vec![0b0010_0101u8], + }, + ), + ); + + let expr = HashTableLookupExpr::try_from_proto(&proto, &ctx).unwrap(); + let expr = expr.downcast_ref::().unwrap(); + + assert!(matches!( + expr.map, + HashTableLookupExprMap::MembershipOnlyArrayMap(_) + )); + // In-range hits (10, 12, 15), in-range misses (11, 14), and + // out-of-range probes on both sides (9, 16). + assert_eq!( + eval_lookup(expr, &probe_batch(&[9, 10, 11, 12, 14, 15, 16])), + [false, true, false, true, false, true, false] + ); + } + + #[test] + fn hash_table_lookup_expr_try_from_proto_rejects_wrong_node_type() { + let schema = Schema::empty(); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + let proto = column_node("a", 0); + + let err = HashTableLookupExpr::try_from_proto(&proto, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalExprNode is not a HashTableLookupExpr"), + "{err}" + ); + } + + #[test] + fn hash_table_lookup_expr_try_from_proto_rejects_missing_map() { + let schema = lookup_schema(); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + let proto = protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some( + protobuf::physical_expr_node::ExprType::HashTableLookupExpr( + protobuf::PhysicalHashTableLookupExprNode { + on_columns: vec![column_node("a", 0)], + seed0: 42, + description: "hash_lookup".to_string(), + map: None, + }, + ), + ), + }; + + let err = HashTableLookupExpr::try_from_proto(&proto, &ctx).unwrap_err(); + assert!( + err.to_string().contains("HashTableLookupExpr has no map"), + "{err}" + ); + } + + #[test] + fn hash_table_lookup_expr_try_from_proto_rejects_bad_presence_length() { + let schema = lookup_schema(); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + // 6 slots need exactly 1 presence byte; send 2. + let proto = lookup_expr_proto( + protobuf::physical_hash_table_lookup_expr_node::Map::ArrayMapMembership( + protobuf::ArrayMapMembership { + offset: 10, + num_slots: 6, + presence: vec![0b0010_0101u8, 0], + }, + ), + ); + + assert!(HashTableLookupExpr::try_from_proto(&proto, &ctx).is_err()); + } + + #[test] + fn hash_table_lookup_expr_roundtrip_hash_map() { + let build_hashes = hashes_for(&[1, 3, 5], 42); + let map = Arc::new(Map::HashMap(Box::new(join_hash_map_with_hashes( + &build_hashes, + )))); + let expr = HashTableLookupExpr::new( + vec![Arc::new(Column::new("a", 0))], + SeededRandomState::with_seed(42), + map, + "hash_lookup".to_string(), + ); + + let encoder = TestEncoder; + let proto = expr + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .unwrap(); + let schema = lookup_schema(); + let decoder = TestDecoder; + let decoded = HashTableLookupExpr::try_from_proto( + &proto, + &test_decode_ctx(&schema, &decoder), + ) + .unwrap(); + + let batch = probe_batch(&[0, 1, 2, 3, 4, 5, 6]); + #[cfg(not(feature = "force_hash_collisions"))] + assert_eq!( + eval_lookup(&expr, &batch), + [false, true, false, true, false, true, false] + ); + assert_eq!( + eval_lookup(&expr, &batch), + eval_lookup(decoded.as_ref(), &batch) + ); + } + + #[test] + fn hash_table_lookup_expr_roundtrip_array_map_sign_crossing_range() { + // A build-side range that crosses zero wraps mid-range in the + // u64 key domain (-5 maps to a slot below 0's slot); roundtrip + // must preserve membership across the wrap. + let build: ArrayRef = Arc::new(Int64Array::from(vec![-5i64, 0, 5])); + let array_map = ArrayMap::try_new(&build, (-5i64) as u64, 5).unwrap(); + let expr = HashTableLookupExpr::new( + vec![Arc::new(Column::new("a", 0))], + SeededRandomState::with_seed(42), + Arc::new(Map::ArrayMap(array_map)), + "hash_lookup".to_string(), + ); + + let encoder = TestEncoder; + let proto = expr + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .unwrap(); + let schema = lookup_schema(); + let decoder = TestDecoder; + let decoded = HashTableLookupExpr::try_from_proto( + &proto, + &test_decode_ctx(&schema, &decoder), + ) + .unwrap(); + + let batch = probe_batch(&[-6, -5, -1, 0, 1, 5, 6]); + assert_eq!( + eval_lookup(&expr, &batch), + [false, true, false, true, false, true, false] + ); + assert_eq!( + eval_lookup(&expr, &batch), + eval_lookup(decoded.as_ref(), &batch) + ); + + // Re-encoding the membership-only map is byte-identical. + let reencoded = decoded + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .unwrap(); + assert_eq!(proto, reencoded); + } + + #[test] + fn hash_table_lookup_expr_roundtrip_hash_map_u64() { + let build_hashes = hashes_for(&[2, 4], 42); + let map = Arc::new(Map::HashMap(Box::new(join_hash_map_u64_with_hashes( + &build_hashes, + )))); + let expr = HashTableLookupExpr::new( + vec![Arc::new(Column::new("a", 0))], + SeededRandomState::with_seed(42), + map, + "hash_lookup".to_string(), + ); + + let encoder = TestEncoder; + let proto = expr + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .unwrap(); + let schema = lookup_schema(); + let decoder = TestDecoder; + let decoded = HashTableLookupExpr::try_from_proto( + &proto, + &test_decode_ctx(&schema, &decoder), + ) + .unwrap(); + + let batch = probe_batch(&[1, 2, 3, 4, 5]); + #[cfg(not(feature = "force_hash_collisions"))] + assert_eq!( + eval_lookup(&expr, &batch), + [false, true, false, true, false] + ); + assert_eq!( + eval_lookup(&expr, &batch), + eval_lookup(decoded.as_ref(), &batch) + ); + } + + #[test] + fn hash_table_lookup_expr_roundtrip_multi_column() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ]); + // Build side has key pairs (1, 10) and (7, 70) + let build_a: ArrayRef = Arc::new(Int64Array::from(vec![1i64, 7])); + let build_b: ArrayRef = Arc::new(Int64Array::from(vec![10i64, 70])); + let mut build_hashes = vec![0u64; 2]; + create_hashes( + &[build_a, build_b], + SeededRandomState::with_seed(42).random_state(), + &mut build_hashes, + ) + .unwrap(); + let map = Arc::new(Map::HashMap(Box::new(join_hash_map_with_hashes( + &build_hashes, + )))); + let expr = HashTableLookupExpr::new( + vec![Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1))], + SeededRandomState::with_seed(42), + map, + "hash_lookup".to_string(), + ); + + let encoder = TestEncoder; + let proto = expr + .try_to_proto(&PhysicalExprEncodeCtx::new(&encoder)) + .unwrap() + .unwrap(); + match proto.expr_type.as_ref().unwrap() { + protobuf::physical_expr_node::ExprType::HashTableLookupExpr(node) => { + assert_eq!(node.on_columns.len(), 2) + } + other => panic!("expected HashTableLookupExpr, got {other:?}"), + } + let decoder = TestDecoder; + let decoded = HashTableLookupExpr::try_from_proto( + &proto, + &test_decode_ctx(&schema, &decoder), + ) + .unwrap(); + + // Present pairs (1, 10) and (7, 70) match; the cross-pairings + // (1, 70) and (7, 10) must not. + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Int64Array::from(vec![1i64, 1, 7, 7])), + Arc::new(Int64Array::from(vec![10i64, 70, 10, 70])), + ], + ) + .unwrap(); + #[cfg(not(feature = "force_hash_collisions"))] + assert_eq!(eval_lookup(&expr, &batch), [true, false, false, true]); + assert_eq!( + eval_lookup(&expr, &batch), + eval_lookup(decoded.as_ref(), &batch) + ); + } + + fn lookup_schema() -> Schema { + Schema::new(vec![Field::new("a", DataType::Int64, false)]) + } + + fn probe_batch(values: &[i64]) -> RecordBatch { + RecordBatch::try_new( + Arc::new(lookup_schema()), + vec![Arc::new(Int64Array::from(values.to_vec()))], + ) + .unwrap() + } + + fn eval_lookup(expr: &dyn PhysicalExpr, batch: &RecordBatch) -> Vec { + let array = expr + .evaluate(batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + let bools = array.as_any().downcast_ref::().unwrap(); + bools.iter().map(|v| v.unwrap()).collect() + } + + /// Build a `JoinHashMapU32` containing exactly the given distinct hashes. + fn join_hash_map_with_hashes(hashes: &[u64]) -> JoinHashMapU32 { + let mut table = HashTable::with_capacity(hashes.len()); + for (i, h) in hashes.iter().enumerate() { + table.insert_unique(*h, (*h, i as u32 + 1), |(h, _)| *h); + } + JoinHashMapU32::new(table, vec![0; hashes.len()]) + } + + /// Build a `JoinHashMapU64` containing exactly the given distinct hashes. + fn join_hash_map_u64_with_hashes(hashes: &[u64]) -> JoinHashMapU64 { + let mut table = HashTable::with_capacity(hashes.len()); + for (i, h) in hashes.iter().enumerate() { + table.insert_unique(*h, (*h, i as u64 + 1), |(h, _)| *h); + } + JoinHashMapU64::new(table, vec![0; hashes.len()]) + } + + /// Hash `values` the same way `HashTableLookupExpr::evaluate` hashes + /// probe keys, so tests can construct build hashes that match. + fn hashes_for(values: &[i64], seed: u64) -> Vec { + let array: ArrayRef = Arc::new(Int64Array::from(values.to_vec())); + let mut buf = vec![0u64; values.len()]; + create_hashes( + &[array], + SeededRandomState::with_seed(seed).random_state(), + &mut buf, + ) + .unwrap(); + buf + } + + fn column_node(name: &str, index: u32) -> protobuf::PhysicalExprNode { + protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: name.to_string(), + index, + }, + )), + } + } + + fn lookup_expr_proto( + map: protobuf::physical_hash_table_lookup_expr_node::Map, + ) -> protobuf::PhysicalExprNode { + protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some( + protobuf::physical_expr_node::ExprType::HashTableLookupExpr( + protobuf::PhysicalHashTableLookupExprNode { + on_columns: vec![column_node("a", 0)], + seed0: 42, + description: "hash_lookup".to_string(), + map: Some(map), + }, + ), + ), + } + } } #[test] diff --git a/datafusion/physical-plan/src/joins/join_hash_map.rs b/datafusion/physical-plan/src/joins/join_hash_map.rs index 454cc916aeb12..6788cc445a038 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -139,6 +139,16 @@ pub trait JoinHashMapType: Send + Sync { /// Returns the number of entries in the join hash map. fn len(&self) -> usize; + + /// Returns the distinct join-key hashes stored in this map, one per + /// entry, in unspecified order. + /// + /// Used to serialize [`HashTableLookupExpr`] dynamic filters: the + /// returned set must contain exactly the hashes that + /// [`contain_hashes`](Self::contain_hashes) would report as present. + /// + /// [`HashTableLookupExpr`]: crate::joins::HashTableLookupExpr + fn hashes(&self) -> Vec; } pub struct JoinHashMapU32 { @@ -219,6 +229,10 @@ impl JoinHashMapType for JoinHashMapU32 { fn len(&self) -> usize { self.map.len() } + + fn hashes(&self) -> Vec { + self.map.iter().map(|(hash, _)| *hash).collect() + } } pub struct JoinHashMapU64 { @@ -299,6 +313,10 @@ impl JoinHashMapType for JoinHashMapU64 { fn len(&self) -> usize { self.map.len() } + + fn hashes(&self) -> Vec { + self.map.iter().map(|(hash, _)| *hash).collect() + } } use crate::joins::MapOffset; diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index e4f7e2e123e0e..cc57da51df337 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -17,6 +17,7 @@ //! DataFusion Join implementations +pub use array_map::ArrayMap; use arrow::array::BooleanBufferBuilder; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; @@ -49,7 +50,6 @@ mod join_filter; /// and is not guaranteed to be stable across versions. pub mod join_hash_map; -use array_map::ArrayMap; use utils::JoinHashMapType; /// The build-side map of a hash join, indexing build rows by join key. diff --git a/datafusion/physical-plan/src/joins/stream_join_utils.rs b/datafusion/physical-plan/src/joins/stream_join_utils.rs index 05a56d241102e..059e70fa8d490 100644 --- a/datafusion/physical-plan/src/joins/stream_join_utils.rs +++ b/datafusion/physical-plan/src/joins/stream_join_utils.rs @@ -112,6 +112,10 @@ impl JoinHashMapType for PruningJoinHashMap { fn len(&self) -> usize { self.map.len() } + + fn hashes(&self) -> Vec { + self.map.iter().map(|(hash, _)| *hash).collect() + } } /// The `PruningJoinHashMap` is similar to a regular `JoinHashMap`, but with diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 43a90264c2b1f..32b6ad349352e 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1051,6 +1051,7 @@ message PhysicalExprNode { PhysicalLambdaExprNode lambda = 25; PhysicalLambdaVariableExprNode lambda_variable = 26; PhysicalRangeExprNode range_expr = 27; + PhysicalHashTableLookupExprNode hash_table_lookup_expr = 28; } } @@ -1203,6 +1204,49 @@ message PhysicalHashExprNode { string description = 6; } +// Serialized form of `HashTableLookupExpr`: a dynamic-filter expression that +// tests probe-side join keys for membership in a hash join's build side. +// +// The build-side map is encoded membership-only: the deserialized expression +// supports membership checks but cannot serve as a join's build map. +message PhysicalHashTableLookupExprNode { + // Probe-side key columns evaluated to produce lookup keys. + repeated PhysicalExprNode on_columns = 1; + // Seed for the hash function applied to `on_columns` when probing a + // `HashMapMembership` map. Hashes are only comparable between identical + // DataFusion builds: the hash function (ahash) is not stable across + // versions or platforms, and a mismatch silently drops join rows. + uint64 seed0 = 2; + // Display string for EXPLAIN output; preserved verbatim across the roundtrip. + string description = 3; + oneof map { + HashMapMembership hash_map_membership = 4; + ArrayMapMembership array_map_membership = 5; + } +} + +// Membership-only encoding of a hash join's build-side hash table: the set +// of distinct join-key hashes present on the build side. +message HashMapMembership { + // Distinct 64-bit key hashes, computed with `seed0`. Order is unspecified. + repeated fixed64 build_hashes = 1; +} + +// Membership-only encoding of an `ArrayMap` (single-column integer join keys +// within a bounded range), as a presence bitmap over the key range. +message ArrayMapMembership { + // Minimum build-side key value as a raw wrapped u64 (two's complement bit + // pattern for signed key types). Bitmap slot `i` corresponds to key + // `offset + i`, computed with wrapping arithmetic. + uint64 offset = 1; + // Width of the key range, i.e. the number of bitmap slots. `presence` + // must be exactly ceil(num_slots / 8) bytes or decoding fails. + uint64 num_slots = 2; + // Presence bitmap, LSB-first within each byte (Arrow validity-buffer bit + // order): bit `i` set means key `offset + i` exists on the build side. + bytes presence = 3; +} + message PhysicalRangeExprNode { repeated PhysicalSortExprNode sort_expr = 1; repeated PhysicalRangeSplitPoint split_point = 2; diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 908f9752b7f18..c99436bfa3f33 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1445,6 +1445,144 @@ impl<'de> serde::Deserialize<'de> for AnalyzedLogicalPlanType { deserializer.deserialize_struct("datafusion.AnalyzedLogicalPlanType", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for ArrayMapMembership { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.offset != 0 { + len += 1; + } + if self.num_slots != 0 { + len += 1; + } + if !self.presence.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ArrayMapMembership", len)?; + if self.offset != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("offset", ToString::to_string(&self.offset).as_str())?; + } + if self.num_slots != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("numSlots", ToString::to_string(&self.num_slots).as_str())?; + } + if !self.presence.is_empty() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("presence", pbjson::private::base64::encode(&self.presence).as_str())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ArrayMapMembership { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "offset", + "num_slots", + "numSlots", + "presence", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Offset, + NumSlots, + Presence, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "offset" => Ok(GeneratedField::Offset), + "numSlots" | "num_slots" => Ok(GeneratedField::NumSlots), + "presence" => Ok(GeneratedField::Presence), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ArrayMapMembership; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ArrayMapMembership") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut offset__ = None; + let mut num_slots__ = None; + let mut presence__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Offset => { + if offset__.is_some() { + return Err(serde::de::Error::duplicate_field("offset")); + } + offset__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::NumSlots => { + if num_slots__.is_some() { + return Err(serde::de::Error::duplicate_field("numSlots")); + } + num_slots__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::Presence => { + if presence__.is_some() { + return Err(serde::de::Error::duplicate_field("presence")); + } + presence__ = + Some(map_.next_value::<::pbjson::private::BytesDeserialize<_>>()?.0) + ; + } + } + } + Ok(ArrayMapMembership { + offset: offset__.unwrap_or_default(), + num_slots: num_slots__.unwrap_or_default(), + presence: presence__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.ArrayMapMembership", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for ArrowScanExecNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -9302,6 +9440,101 @@ impl<'de> serde::Deserialize<'de> for HashJoinExecNode { deserializer.deserialize_struct("datafusion.HashJoinExecNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for HashMapMembership { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.build_hashes.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.HashMapMembership", len)?; + if !self.build_hashes.is_empty() { + struct_ser.serialize_field("buildHashes", &self.build_hashes.iter().map(ToString::to_string).collect::>())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for HashMapMembership { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "build_hashes", + "buildHashes", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + BuildHashes, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "buildHashes" | "build_hashes" => Ok(GeneratedField::BuildHashes), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = HashMapMembership; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.HashMapMembership") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut build_hashes__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::BuildHashes => { + if build_hashes__.is_some() { + return Err(serde::de::Error::duplicate_field("buildHashes")); + } + build_hashes__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } + } + } + Ok(HashMapMembership { + build_hashes: build_hashes__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.HashMapMembership", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for HashRepartition { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -18492,6 +18725,9 @@ impl serde::Serialize for PhysicalExprNode { physical_expr_node::ExprType::RangeExpr(v) => { struct_ser.serialize_field("rangeExpr", v)?; } + physical_expr_node::ExprType::HashTableLookupExpr(v) => { + struct_ser.serialize_field("hashTableLookupExpr", v)?; + } } } struct_ser.end() @@ -18549,6 +18785,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "lambdaVariable", "range_expr", "rangeExpr", + "hash_table_lookup_expr", + "hashTableLookupExpr", ]; #[allow(clippy::enum_variant_names)] @@ -18579,6 +18817,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { Lambda, LambdaVariable, RangeExpr, + HashTableLookupExpr, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -18626,6 +18865,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "lambda" => Ok(GeneratedField::Lambda), "lambdaVariable" | "lambda_variable" => Ok(GeneratedField::LambdaVariable), "rangeExpr" | "range_expr" => Ok(GeneratedField::RangeExpr), + "hashTableLookupExpr" | "hash_table_lookup_expr" => Ok(GeneratedField::HashTableLookupExpr), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -18830,6 +19070,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { return Err(serde::de::Error::duplicate_field("rangeExpr")); } expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::RangeExpr) +; + } + GeneratedField::HashTableLookupExpr => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("hashTableLookupExpr")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::HashTableLookupExpr) ; } } @@ -19311,6 +19558,173 @@ impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { deserializer.deserialize_struct("datafusion.PhysicalHashRepartition", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for PhysicalHashTableLookupExprNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.on_columns.is_empty() { + len += 1; + } + if self.seed0 != 0 { + len += 1; + } + if !self.description.is_empty() { + len += 1; + } + if self.map.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalHashTableLookupExprNode", len)?; + if !self.on_columns.is_empty() { + struct_ser.serialize_field("onColumns", &self.on_columns)?; + } + if self.seed0 != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("seed0", ToString::to_string(&self.seed0).as_str())?; + } + if !self.description.is_empty() { + struct_ser.serialize_field("description", &self.description)?; + } + if let Some(v) = self.map.as_ref() { + match v { + physical_hash_table_lookup_expr_node::Map::HashMapMembership(v) => { + struct_ser.serialize_field("hashMapMembership", v)?; + } + physical_hash_table_lookup_expr_node::Map::ArrayMapMembership(v) => { + struct_ser.serialize_field("arrayMapMembership", v)?; + } + } + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalHashTableLookupExprNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "on_columns", + "onColumns", + "seed0", + "description", + "hash_map_membership", + "hashMapMembership", + "array_map_membership", + "arrayMapMembership", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + OnColumns, + Seed0, + Description, + HashMapMembership, + ArrayMapMembership, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "onColumns" | "on_columns" => Ok(GeneratedField::OnColumns), + "seed0" => Ok(GeneratedField::Seed0), + "description" => Ok(GeneratedField::Description), + "hashMapMembership" | "hash_map_membership" => Ok(GeneratedField::HashMapMembership), + "arrayMapMembership" | "array_map_membership" => Ok(GeneratedField::ArrayMapMembership), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalHashTableLookupExprNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalHashTableLookupExprNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut on_columns__ = None; + let mut seed0__ = None; + let mut description__ = None; + let mut map__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::OnColumns => { + if on_columns__.is_some() { + return Err(serde::de::Error::duplicate_field("onColumns")); + } + on_columns__ = Some(map_.next_value()?); + } + GeneratedField::Seed0 => { + if seed0__.is_some() { + return Err(serde::de::Error::duplicate_field("seed0")); + } + seed0__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::Description => { + if description__.is_some() { + return Err(serde::de::Error::duplicate_field("description")); + } + description__ = Some(map_.next_value()?); + } + GeneratedField::HashMapMembership => { + if map__.is_some() { + return Err(serde::de::Error::duplicate_field("hashMapMembership")); + } + map__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_hash_table_lookup_expr_node::Map::HashMapMembership) +; + } + GeneratedField::ArrayMapMembership => { + if map__.is_some() { + return Err(serde::de::Error::duplicate_field("arrayMapMembership")); + } + map__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_hash_table_lookup_expr_node::Map::ArrayMapMembership) +; + } + } + } + Ok(PhysicalHashTableLookupExprNode { + on_columns: on_columns__.unwrap_or_default(), + seed0: seed0__.unwrap_or_default(), + description: description__.unwrap_or_default(), + map: map__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalHashTableLookupExprNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PhysicalHigherOrderUdfNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index ba00577ab9a1b..0bf4f6e40b387 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1557,7 +1557,7 @@ pub struct PhysicalExprNode { pub expr_id: ::core::option::Option, #[prost( oneof = "physical_expr_node::ExprType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28" )] pub expr_type: ::core::option::Option, } @@ -1622,6 +1622,8 @@ pub mod physical_expr_node { LambdaVariable(super::PhysicalLambdaVariableExprNode), #[prost(message, tag = "27")] RangeExpr(super::PhysicalRangeExprNode), + #[prost(message, tag = "28")] + HashTableLookupExpr(super::PhysicalHashTableLookupExprNode), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1862,6 +1864,64 @@ pub struct PhysicalHashExprNode { #[prost(string, tag = "6")] pub description: ::prost::alloc::string::String, } +/// Serialized form of `HashTableLookupExpr`: a dynamic-filter expression that +/// tests probe-side join keys for membership in a hash join's build side. +/// +/// The build-side map is encoded membership-only: the deserialized expression +/// supports membership checks but cannot serve as a join's build map. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalHashTableLookupExprNode { + /// Probe-side key columns evaluated to produce lookup keys. + #[prost(message, repeated, tag = "1")] + pub on_columns: ::prost::alloc::vec::Vec, + /// Seed for the hash function applied to `on_columns` when probing a + /// `HashMapMembership` map. Hashes are only comparable between identical + /// DataFusion builds: the hash function (ahash) is not stable across + /// versions or platforms, and a mismatch silently drops join rows. + #[prost(uint64, tag = "2")] + pub seed0: u64, + /// Display string for EXPLAIN output; preserved verbatim across the roundtrip. + #[prost(string, tag = "3")] + pub description: ::prost::alloc::string::String, + #[prost(oneof = "physical_hash_table_lookup_expr_node::Map", tags = "4, 5")] + pub map: ::core::option::Option, +} +/// Nested message and enum types in `PhysicalHashTableLookupExprNode`. +pub mod physical_hash_table_lookup_expr_node { + #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum Map { + #[prost(message, tag = "4")] + HashMapMembership(super::HashMapMembership), + #[prost(message, tag = "5")] + ArrayMapMembership(super::ArrayMapMembership), + } +} +/// Membership-only encoding of a hash join's build-side hash table: the set +/// of distinct join-key hashes present on the build side. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct HashMapMembership { + /// Distinct 64-bit key hashes, computed with `seed0`. Order is unspecified. + #[prost(fixed64, repeated, tag = "1")] + pub build_hashes: ::prost::alloc::vec::Vec, +} +/// Membership-only encoding of an `ArrayMap` (single-column integer join keys +/// within a bounded range), as a presence bitmap over the key range. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ArrayMapMembership { + /// Minimum build-side key value as a raw wrapped u64 (two's complement bit + /// pattern for signed key types). Bitmap slot `i` corresponds to key + /// `offset + i`, computed with wrapping arithmetic. + #[prost(uint64, tag = "1")] + pub offset: u64, + /// Width of the key range, i.e. the number of bitmap slots. `presence` + /// must be exactly ceil(num_slots / 8) bytes or decoding fails. + #[prost(uint64, tag = "2")] + pub num_slots: u64, + /// Presence bitmap, LSB-first within each byte (Arrow validity-buffer bit + /// order): bit `i` set means key `offset + i` exists on the build side. + #[prost(bytes = "vec", tag = "3")] + pub presence: ::prost::alloc::vec::Vec, +} #[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalRangeExprNode { #[prost(message, repeated, tag = "1")] diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 06105be806cfc..520679ac82424 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -38,7 +38,7 @@ use datafusion_physical_plan::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn, }; -use datafusion_physical_plan::joins::HashExpr; +use datafusion_physical_plan::joins::{HashExpr, HashTableLookupExpr}; use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; use datafusion_physical_plan::repartition::RangeExpr; use datafusion_physical_plan::windows::{create_window_expr, schema_add_window_field}; @@ -348,6 +348,9 @@ pub fn parse_physical_expr_with_converter( } ExprType::LikeExpr(_) => LikeExpr::try_from_proto(proto, &decode_ctx)?, ExprType::HashExpr(_) => HashExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::HashTableLookupExpr(_) => { + HashTableLookupExpr::try_from_proto(proto, &decode_ctx)? + } ExprType::RangeExpr(_) => RangeExpr::try_from_proto(proto, &decode_ctx)?, ExprType::ScalarSubquery(_) => { let results = ctx.scalar_subquery_results().ok_or_else(|| { diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs index 518b4a62ce072..20e45b38578c7 100644 --- a/datafusion/proto/tests/cases/plans/exprs.rs +++ b/datafusion/proto/tests/cases/plans/exprs.rs @@ -23,7 +23,6 @@ use arrow::datatypes::Fields; use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema}; use datafusion::logical_expr::Operator; -use datafusion::physical_expr::expressions::Literal; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::expressions::{ BinaryExpr, Column, PhysicalSortExpr, binary, col, like, lit, @@ -31,18 +30,15 @@ use datafusion::physical_plan::expressions::{ use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion::physical_plan::repartition::RangeExpr; -use datafusion::physical_plan::{ - ExecutionPlan, PhysicalExpr, RangePartitioning, SplitPoint, -}; +use datafusion::physical_plan::{PhysicalExpr, RangePartitioning, SplitPoint}; use datafusion::prelude::SessionContext; use datafusion::scalar::ScalarValue; use datafusion_common::Result; use datafusion_proto::physical_plan::{ - AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalProtoConverterExtension, }; use datafusion_proto::protobuf; -use datafusion_proto::protobuf::PhysicalPlanNode; use std::sync::Arc; use std::vec; @@ -95,23 +91,28 @@ fn roundtrip_like() -> Result<()> { roundtrip_test(plan) } -/// Test that HashTableLookupExpr serializes to lit(true) +/// Test that HashTableLookupExpr roundtrips through a full plan. /// -/// HashTableLookupExpr contains a runtime hash table that cannot be serialized. -/// The serialization code replaces it with lit(true) which is safe because -/// it's a performance optimization filter, not a correctness requirement. +/// The build-side map is serialized as a membership-only encoding (distinct +/// hashes for hash maps, a presence bitmap for array maps) and reconstructed +/// on deserialization as a map that supports only membership checks. #[test] -fn roundtrip_hash_table_lookup_expr_to_lit() -> Result<()> { - use datafusion::physical_plan::joins::join_hash_map::JoinHashMapU32; +fn roundtrip_hash_table_lookup_expr() -> Result<()> { + use datafusion::physical_plan::joins::join_hash_map::{ + JoinHashMapType, JoinHashMapU32, + }; use datafusion::physical_plan::joins::{HashTableLookupExpr, Map}; // Create a simple schema and input plan let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); let input = Arc::new(EmptyExec::new(schema.clone())); - // Create a HashTableLookupExpr - it will be replaced with lit(true) during serialization - let hash_map = Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(0)))); let on_columns = vec![col("col", &schema)?]; + // Populate the map so the roundtrip carries a real membership payload + let build_hashes: Vec = vec![100, 200, 300]; + let mut join_map = JoinHashMapU32::with_capacity(build_hashes.len()); + join_map.update_from_iter(Box::new(build_hashes.iter().enumerate()), 0); + let hash_map = Arc::new(Map::HashMap(Box::new(join_map))); let lookup_expr: Arc = Arc::new(HashTableLookupExpr::new( on_columns, datafusion::physical_plan::joins::SeededRandomState::with_seed(0), @@ -121,28 +122,33 @@ fn roundtrip_hash_table_lookup_expr_to_lit() -> Result<()> { // Create a filter with the lookup expression let filter = Arc::new(FilterExec::try_new(lookup_expr, input)?); + roundtrip_test(filter) +} - // Serialize - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; +/// Roundtrip a plan whose HashTableLookupExpr carries the ArrayMap +/// membership encoding (dense integer keys within a bounded range). +#[test] +fn roundtrip_hash_table_lookup_expr_array_map() -> Result<()> { + use datafusion::arrow::array::{ArrayRef, Int64Array}; + use datafusion::physical_plan::joins::{ArrayMap, HashTableLookupExpr, Map}; - let proto: PhysicalPlanNode = - PhysicalPlanNode::try_from_physical_plan(filter.clone(), &codec) - .expect("serialization should succeed"); + let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone())); - // Deserialize - let result: Arc = proto - .try_into_physical_plan(&ctx.task_ctx(), &codec) - .expect("deserialization should succeed"); + // Keys {10, 12, 15} over the range [10, 15], with a duplicate key + let build_keys: ArrayRef = Arc::new(Int64Array::from(vec![10i64, 12, 10, 15])); + let array_map = ArrayMap::try_new(&build_keys, 10, 15)?; - // The deserialized plan should have lit(true) instead of HashTableLookupExpr - // Verify the filter predicate is a Literal(true) - let result_filter = result.downcast_ref::().unwrap(); - let predicate = result_filter.predicate(); - let literal = predicate.downcast_ref::().unwrap(); - assert_eq!(*literal.value(), ScalarValue::Boolean(Some(true))); + let on_columns = vec![col("col", &schema)?]; + let lookup_expr: Arc = Arc::new(HashTableLookupExpr::new( + on_columns, + datafusion::physical_plan::joins::SeededRandomState::with_seed(0), + Arc::new(Map::ArrayMap(array_map)), + "test_lookup".to_string(), + )); - Ok(()) + let filter = Arc::new(FilterExec::try_new(lookup_expr, input)?); + roundtrip_test(filter) } #[test] diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 7d0c19c846ca0..23bcc8ad5e29b 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -1219,3 +1219,28 @@ let table_opts = TableParquetOptions::try_from(&proto_table_opts)?; ``` See [issue #24019](https://github.com/apache/datafusion/issues/24019) for details. + +### `JoinHashMapType` has a new required method `hashes` + +To support serializing hash join dynamic filters (`HashTableLookupExpr`), the +`datafusion_physical_plan::joins::join_hash_map::JoinHashMapType` trait (also +re-exported as `joins::utils::JoinHashMapType`) has a new required method +`hashes`, which returns the distinct join-key hashes stored in the map. + +**Who is affected:** + +- Users with custom `JoinHashMapType` implementations. + +**Migration guide:** + +Implement `hashes` to return the distinct join-key hashes stored in the map, +in any order — exactly the set that `contain_hashes` reports as present. For +`hashbrown`-backed implementations this is typically: + +```rust,ignore +fn hashes(&self) -> Vec { + self.map.iter().map(|(hash, _)| *hash).collect() +} +``` + +See [issue #24277](https://github.com/apache/datafusion/issues/24277) for details.