From 3483b453ac6c8821140813eeae6f8a4ae9f0b71a Mon Sep 17 00:00:00 2001 From: rjhallsted Date: Thu, 13 Aug 2026 10:48:21 -0600 Subject: [PATCH 1/8] add proto and serialization --- .../physical-plan/src/joins/array_map.rs | 17 + .../joins/hash_join/partitioned_hash_eval.rs | 64 ++- .../physical-plan/src/joins/join_hash_map.rs | 10 + .../src/joins/stream_join_utils.rs | 4 + .../proto-models/proto/datafusion.proto | 19 + .../proto-models/src/generated/pbjson.rs | 414 ++++++++++++++++++ .../proto-models/src/generated/prost.rs | 39 +- 7 files changed, 545 insertions(+), 22 deletions(-) diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index 4e56cf013c8f7..c82ec752b3f1d 100644 --- a/datafusion/physical-plan/src/joins/array_map.rs +++ b/datafusion/physical-plan/src/joins/array_map.rs @@ -416,6 +416,23 @@ impl ArrayMap { }); Ok(BooleanArray::new(buffer, None)) } + + #[cfg(feature = "proto")] + pub fn to_proto_membership_only( + &self, + ) -> datafusion_proto_models::protobuf::ArrayMapMembershipNode { + use datafusion_proto_models::protobuf; + + let bits = BooleanBuffer::collect_bool(self.data.len(), |i| self.data[i] != 0); + // SAFETY: This always succeeds. None of the 3 error conditions of + // into_vec() are present. + let presence = bits.sliced().into_vec().unwrap(); + protobuf::ArrayMapMembershipNode { + offset: self.offset, + num_slots: self.data.len() as u64, + presence, + } + } } #[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..b3fa28884b999 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 @@ -262,6 +262,8 @@ pub struct HashTableLookupExpr { /// Random state for hashing (with seeds preserved for serialization) random_state: SeededRandomState, /// Map to check against (hash table or array map) + /// Safety: Post-deserialization, this is only safe for membership checks and + /// should not be modified further. map: Arc, /// Description for display description: String, @@ -386,36 +388,33 @@ impl PhysicalExpr for HashTableLookupExpr { } } } + #[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::HashLookupExpr(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_id: self.expression_id(), + expr_type: Some(expr), })) } fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -423,6 +422,29 @@ impl PhysicalExpr for HashTableLookupExpr { } } +#[cfg(feature = "proto")] +fn try_map_to_proto_membership_only( + map: &Map, +) -> Result +{ + use datafusion_proto_models::protobuf; + + match map { + 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::HashSetMapMembershipNode { + build_hashes: hash_map.hashes(), + }, + ), + ), + } +} + fn evaluate_columns( columns: &[PhysicalExprRef], batch: &RecordBatch, diff --git a/datafusion/physical-plan/src/joins/join_hash_map.rs b/datafusion/physical-plan/src/joins/join_hash_map.rs index 454cc916aeb12..29849c5d0920f 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -139,6 +139,8 @@ pub trait JoinHashMapType: Send + Sync { /// Returns the number of entries in the join hash map. fn len(&self) -> usize; + + fn hashes(&self) -> Vec; } pub struct JoinHashMapU32 { @@ -219,6 +221,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 +305,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/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..ad70fc7d44eca 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_lookup_expr = 28; } } @@ -1203,6 +1204,24 @@ message PhysicalHashExprNode { string description = 6; } +message PhysicalHashTableLookupExprNode { + repeated PhysicalExprNode on_columns = 1; + uint64 seed0 = 2; + string description = 3; + oneof map { + HashSetMapMembership hash_map_membership = 4; + ArrayMap array_map = 5; + } +} + +message HashSetMapMembershipNode { + repeated fixed64 build_hashes = 1; +} + +message ArrayMap { + repeated fixed64 keys = 1; +} + 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..93be6ce1648e8 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 ArrayMapMembershipNode { + #[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.ArrayMapMembershipNode", 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 ArrayMapMembershipNode { + #[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 = ArrayMapMembershipNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ArrayMapMembershipNode") + } + + 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(ArrayMapMembershipNode { + offset: offset__.unwrap_or_default(), + num_slots: num_slots__.unwrap_or_default(), + presence: presence__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.ArrayMapMembershipNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for ArrowScanExecNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -9416,6 +9554,101 @@ impl<'de> serde::Deserialize<'de> for HashRepartition { deserializer.deserialize_struct("datafusion.HashRepartition", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for HashSetMapMembershipNode { + #[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.HashSetMapMembershipNode", 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 HashSetMapMembershipNode { + #[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 = HashSetMapMembershipNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.HashSetMapMembershipNode") + } + + 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(HashSetMapMembershipNode { + build_hashes: build_hashes__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.HashSetMapMembershipNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for HigherOrderUdfExprNode { #[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::HashLookupExpr(v) => { + struct_ser.serialize_field("hashLookupExpr", v)?; + } } } struct_ser.end() @@ -18549,6 +18785,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "lambdaVariable", "range_expr", "rangeExpr", + "hash_lookup_expr", + "hashLookupExpr", ]; #[allow(clippy::enum_variant_names)] @@ -18579,6 +18817,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { Lambda, LambdaVariable, RangeExpr, + HashLookupExpr, } 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), + "hashLookupExpr" | "hash_lookup_expr" => Ok(GeneratedField::HashLookupExpr), _ => 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::HashLookupExpr => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("hashLookupExpr")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::HashLookupExpr) ; } } @@ -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..caf2d52ea7ebd 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")] + HashLookupExpr(super::PhysicalHashTableLookupExprNode), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1863,6 +1865,41 @@ pub struct PhysicalHashExprNode { pub description: ::prost::alloc::string::String, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalHashTableLookupExprNode { + #[prost(message, repeated, tag = "1")] + pub on_columns: ::prost::alloc::vec::Vec, + #[prost(uint64, tag = "2")] + pub seed0: u64, + #[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::HashSetMapMembershipNode), + #[prost(message, tag = "5")] + ArrayMapMembership(super::ArrayMapMembershipNode), + } +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct HashSetMapMembershipNode { + #[prost(fixed64, repeated, tag = "1")] + pub build_hashes: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ArrayMapMembershipNode { + #[prost(uint64, tag = "1")] + pub offset: u64, + #[prost(uint64, tag = "2")] + pub num_slots: u64, + #[prost(bytes = "vec", tag = "3")] + pub presence: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalRangeExprNode { #[prost(message, repeated, tag = "1")] pub sort_expr: ::prost::alloc::vec::Vec, From ba316c41b395cc00bd9798126803ea7f1ffbbb0b Mon Sep 17 00:00:00 2001 From: rjhallsted Date: Thu, 13 Aug 2026 13:51:41 -0600 Subject: [PATCH 2/8] add deserializatoin --- .../physical-plan/src/joins/array_map.rs | 24 ++- .../joins/hash_join/partitioned_hash_eval.rs | 169 +++++++++++++++--- .../physical-plan/src/joins/join_hash_map.rs | 70 ++++++++ .../proto-models/proto/datafusion.proto | 14 +- .../proto-models/src/generated/pbjson.rs | 164 ++++++++--------- .../proto-models/src/generated/prost.rs | 10 +- .../proto/src/physical_plan/from_proto.rs | 5 +- 7 files changed, 337 insertions(+), 119 deletions(-) diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index c82ec752b3f1d..3417d3577ff76 100644 --- a/datafusion/physical-plan/src/joins/array_map.rs +++ b/datafusion/physical-plan/src/joins/array_map.rs @@ -420,19 +420,39 @@ impl ArrayMap { #[cfg(feature = "proto")] pub fn to_proto_membership_only( &self, - ) -> datafusion_proto_models::protobuf::ArrayMapMembershipNode { + ) -> datafusion_proto_models::protobuf::ArrayMapMembership { use datafusion_proto_models::protobuf; let bits = BooleanBuffer::collect_bool(self.data.len(), |i| self.data[i] != 0); // SAFETY: This always succeeds. None of the 3 error conditions of // into_vec() are present. let presence = bits.sliced().into_vec().unwrap(); - protobuf::ArrayMapMembershipNode { + protobuf::ArrayMapMembership { offset: self.offset, num_slots: self.data.len() as u64, presence, } } + + #[cfg(feature = "proto")] + pub fn try_from_proto_membership_only( + node: &datafusion_proto_models::protobuf::ArrayMapMembership, + ) -> Result { + use arrow_data::bit_iterator::BitIndexIterator; + + 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 b3fa28884b999..b9eb43366b3f5 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. /// @@ -264,7 +267,7 @@ pub struct HashTableLookupExpr { /// Map to check against (hash table or array map) /// Safety: Post-deserialization, this is only safe for membership checks and /// should not be modified further. - map: Arc, + map: Arc, /// Description for display description: String, } @@ -288,7 +291,7 @@ impl HashTableLookupExpr { Self { on_columns, random_state, - map, + map: Arc::new(HashTableLookupExprMap::Normal(map)), description, } } @@ -355,12 +358,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: Arc::clone(&self.map), + description: self.description.clone(), + })) } fn data_type(&self, _input_schema: &Schema) -> Result { @@ -376,13 +379,25 @@ impl PhysicalExpr for HashTableLookupExpr { let join_keys = evaluate_columns(&self.on_columns, batch)?; match self.map.as_ref() { - Map::HashMap(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))) } @@ -405,12 +420,13 @@ impl PhysicalExpr for HashTableLookupExpr { let on_columns = ctx.encode_children_expressions(&self.on_columns)?; let map = try_map_to_proto_membership_only(&self.map)?; - let expr = ExprType::HashLookupExpr(protobuf::PhysicalHashTableLookupExprNode { - on_columns, - seed0: self.random_state.seed, - description: self.description.clone(), - map: Some(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: self.expression_id(), @@ -424,24 +440,88 @@ impl PhysicalExpr for HashTableLookupExpr { #[cfg(feature = "proto")] fn try_map_to_proto_membership_only( - map: &Map, + map: &HashTableLookupExprMap, ) -> Result { use datafusion_proto_models::protobuf; match map { - Map::ArrayMap(array_map) => Ok( - protobuf::physical_hash_table_lookup_expr_node::Map::ArrayMapMembership( - array_map.to_proto_membership_only(), + 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( + 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::HashSetMapMembershipNode { + 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)) => { + Arc::new(HashTableLookupExprMap::MembershipOnlyHashMap( + JoinHashMembershipMap::new(&membership.build_hashes), + )) + } + Some(Map::ArrayMapMembership(membership)) => Arc::new( + HashTableLookupExprMap::MembershipOnlyArrayMap(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(), + })) } } @@ -456,6 +536,49 @@ fn evaluate_columns( .collect() } +enum HashTableLookupExprMap { + /// A regular build-side hash map + Normal(Arc), + /// A membership-checks only version of a hashmap, only constructed via expression deserialization + MembershipOnlyHashMap(JoinHashMembershipMap), + /// A membership-checks only version of an array map, only constructed via expression deserialization + MembershipOnlyArrayMap(JoinMembershipArrayMap), +} + +/// 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()); + + 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::*; diff --git a/datafusion/physical-plan/src/joins/join_hash_map.rs b/datafusion/physical-plan/src/joins/join_hash_map.rs index 29849c5d0920f..26f79fe3b59a2 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -311,6 +311,76 @@ impl JoinHashMapType for JoinHashMapU64 { } } +/// Membership-only join map reconstructed from serialized distinct hashes. +/// Supports `contain_hashes` lookups only; it has no build-side rows. +pub struct JoinHashMembershipMap { + map: HashTable<(u64, ())>, +} +impl JoinHashMembershipMap { + pub fn new(hashes: &[u64]) -> Self { + let mut map = HashTable::with_capacity(hashes.len()); + + for h in hashes { + map.insert_unique(*h, (*h, ()), |(h, _)| *h); + } + + Self { map } + } +} + +impl JoinHashMapType for JoinHashMembershipMap { + fn extend_zero(&mut self, _len: usize) { + unimplemented!(); + } + + fn update_from_iter<'a>( + &mut self, + _iter: Box + Send + 'a>, + _deleted_offset: usize, + ) { + unimplemented!(); + } + + fn get_matched_indices<'a>( + &self, + _iter: Box + 'a>, + _deleted_offset: Option, + ) -> (Vec, Vec) { + unimplemented!(); + } + + fn get_matched_indices_with_limit_offset( + &self, + _hash_values: &[u64], + _valid_keys: Option<&NullBuffer>, + _limit: usize, + _offset: MapOffset, + _input_indices: &mut Vec, + _match_indices: &mut Vec, + ) -> Option { + unimplemented!() + } + + /// Returns a BooleanArray indicating which of the provided hashes exist in the map. + fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray { + contain_hashes(&self.map, hash_values) + } + + /// Returns `true` if the join hash map contains no entries. + fn is_empty(&self) -> bool { + self.map.is_empty() + } + + /// Returns the number of entries in the join hash map. + fn len(&self) -> usize { + self.map.len() + } + + fn hashes(&self) -> Vec { + self.map.iter().map(|(h, _)| *h).collect() + } +} + use crate::joins::MapOffset; use crate::joins::chain::traverse_chain; diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index ad70fc7d44eca..14f8bd2243038 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1051,7 +1051,7 @@ message PhysicalExprNode { PhysicalLambdaExprNode lambda = 25; PhysicalLambdaVariableExprNode lambda_variable = 26; PhysicalRangeExprNode range_expr = 27; - PhysicalHashTableLookupExprNode hash_lookup_expr = 28; + PhysicalHashTableLookupExprNode hash_table_lookup_expr = 28; } } @@ -1209,17 +1209,19 @@ message PhysicalHashTableLookupExprNode { uint64 seed0 = 2; string description = 3; oneof map { - HashSetMapMembership hash_map_membership = 4; - ArrayMap array_map = 5; + HashMapMembership hash_map_membership = 4; + ArrayMapMembership array_map_membership = 5; } } -message HashSetMapMembershipNode { +message HashMapMembership { repeated fixed64 build_hashes = 1; } -message ArrayMap { - repeated fixed64 keys = 1; +message ArrayMapMembership { + uint64 offset = 1; + uint64 num_slots = 2; + bytes presence = 3; } message PhysicalRangeExprNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 93be6ce1648e8..c99436bfa3f33 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -1445,7 +1445,7 @@ impl<'de> serde::Deserialize<'de> for AnalyzedLogicalPlanType { deserializer.deserialize_struct("datafusion.AnalyzedLogicalPlanType", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for ArrayMapMembershipNode { +impl serde::Serialize for ArrayMapMembership { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -1462,7 +1462,7 @@ impl serde::Serialize for ArrayMapMembershipNode { if !self.presence.is_empty() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.ArrayMapMembershipNode", len)?; + 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)] @@ -1481,7 +1481,7 @@ impl serde::Serialize for ArrayMapMembershipNode { struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for ArrayMapMembershipNode { +impl<'de> serde::Deserialize<'de> for ArrayMapMembership { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where @@ -1532,13 +1532,13 @@ impl<'de> serde::Deserialize<'de> for ArrayMapMembershipNode { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = ArrayMapMembershipNode; + type Value = ArrayMapMembership; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.ArrayMapMembershipNode") + formatter.write_str("struct datafusion.ArrayMapMembership") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { @@ -1573,14 +1573,14 @@ impl<'de> serde::Deserialize<'de> for ArrayMapMembershipNode { } } } - Ok(ArrayMapMembershipNode { + Ok(ArrayMapMembership { offset: offset__.unwrap_or_default(), num_slots: num_slots__.unwrap_or_default(), presence: presence__.unwrap_or_default(), }) } } - deserializer.deserialize_struct("datafusion.ArrayMapMembershipNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.ArrayMapMembership", FIELDS, GeneratedVisitor) } } impl serde::Serialize for ArrowScanExecNode { @@ -9440,7 +9440,7 @@ impl<'de> serde::Deserialize<'de> for HashJoinExecNode { deserializer.deserialize_struct("datafusion.HashJoinExecNode", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for HashRepartition { +impl serde::Serialize for HashMapMembership { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -9448,41 +9448,30 @@ impl serde::Serialize for HashRepartition { { use serde::ser::SerializeStruct; let mut len = 0; - if !self.hash_expr.is_empty() { - len += 1; - } - if self.partition_count != 0 { + if !self.build_hashes.is_empty() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.HashRepartition", len)?; - if !self.hash_expr.is_empty() { - struct_ser.serialize_field("hashExpr", &self.hash_expr)?; - } - if self.partition_count != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("partitionCount", ToString::to_string(&self.partition_count).as_str())?; + 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 HashRepartition { +impl<'de> serde::Deserialize<'de> for HashMapMembership { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "hash_expr", - "hashExpr", - "partition_count", - "partitionCount", + "build_hashes", + "buildHashes", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - HashExpr, - PartitionCount, + BuildHashes, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -9504,8 +9493,7 @@ impl<'de> serde::Deserialize<'de> for HashRepartition { E: serde::de::Error, { match value { - "hashExpr" | "hash_expr" => Ok(GeneratedField::HashExpr), - "partitionCount" | "partition_count" => Ok(GeneratedField::PartitionCount), + "buildHashes" | "build_hashes" => Ok(GeneratedField::BuildHashes), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -9515,46 +9503,39 @@ impl<'de> serde::Deserialize<'de> for HashRepartition { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = HashRepartition; + type Value = HashMapMembership; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.HashRepartition") + formatter.write_str("struct datafusion.HashMapMembership") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut hash_expr__ = None; - let mut partition_count__ = None; + let mut build_hashes__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::HashExpr => { - if hash_expr__.is_some() { - return Err(serde::de::Error::duplicate_field("hashExpr")); - } - hash_expr__ = Some(map_.next_value()?); - } - GeneratedField::PartitionCount => { - if partition_count__.is_some() { - return Err(serde::de::Error::duplicate_field("partitionCount")); + GeneratedField::BuildHashes => { + if build_hashes__.is_some() { + return Err(serde::de::Error::duplicate_field("buildHashes")); } - partition_count__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + build_hashes__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) ; } } } - Ok(HashRepartition { - hash_expr: hash_expr__.unwrap_or_default(), - partition_count: partition_count__.unwrap_or_default(), + Ok(HashMapMembership { + build_hashes: build_hashes__.unwrap_or_default(), }) } } - deserializer.deserialize_struct("datafusion.HashRepartition", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.HashMapMembership", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for HashSetMapMembershipNode { +impl serde::Serialize for HashRepartition { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -9562,30 +9543,41 @@ impl serde::Serialize for HashSetMapMembershipNode { { use serde::ser::SerializeStruct; let mut len = 0; - if !self.build_hashes.is_empty() { + if !self.hash_expr.is_empty() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.HashSetMapMembershipNode", len)?; - if !self.build_hashes.is_empty() { - struct_ser.serialize_field("buildHashes", &self.build_hashes.iter().map(ToString::to_string).collect::>())?; + if self.partition_count != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.HashRepartition", len)?; + if !self.hash_expr.is_empty() { + struct_ser.serialize_field("hashExpr", &self.hash_expr)?; + } + if self.partition_count != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("partitionCount", ToString::to_string(&self.partition_count).as_str())?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for HashSetMapMembershipNode { +impl<'de> serde::Deserialize<'de> for HashRepartition { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "build_hashes", - "buildHashes", + "hash_expr", + "hashExpr", + "partition_count", + "partitionCount", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - BuildHashes, + HashExpr, + PartitionCount, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -9607,7 +9599,8 @@ impl<'de> serde::Deserialize<'de> for HashSetMapMembershipNode { E: serde::de::Error, { match value { - "buildHashes" | "build_hashes" => Ok(GeneratedField::BuildHashes), + "hashExpr" | "hash_expr" => Ok(GeneratedField::HashExpr), + "partitionCount" | "partition_count" => Ok(GeneratedField::PartitionCount), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -9617,36 +9610,43 @@ impl<'de> serde::Deserialize<'de> for HashSetMapMembershipNode { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = HashSetMapMembershipNode; + type Value = HashRepartition; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.HashSetMapMembershipNode") + formatter.write_str("struct datafusion.HashRepartition") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut build_hashes__ = None; + let mut hash_expr__ = None; + let mut partition_count__ = 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")); + GeneratedField::HashExpr => { + if hash_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("hashExpr")); } - build_hashes__ = - Some(map_.next_value::>>()? - .into_iter().map(|x| x.0).collect()) + hash_expr__ = Some(map_.next_value()?); + } + GeneratedField::PartitionCount => { + if partition_count__.is_some() { + return Err(serde::de::Error::duplicate_field("partitionCount")); + } + partition_count__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } } } - Ok(HashSetMapMembershipNode { - build_hashes: build_hashes__.unwrap_or_default(), + Ok(HashRepartition { + hash_expr: hash_expr__.unwrap_or_default(), + partition_count: partition_count__.unwrap_or_default(), }) } } - deserializer.deserialize_struct("datafusion.HashSetMapMembershipNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.HashRepartition", FIELDS, GeneratedVisitor) } } impl serde::Serialize for HigherOrderUdfExprNode { @@ -18725,8 +18725,8 @@ impl serde::Serialize for PhysicalExprNode { physical_expr_node::ExprType::RangeExpr(v) => { struct_ser.serialize_field("rangeExpr", v)?; } - physical_expr_node::ExprType::HashLookupExpr(v) => { - struct_ser.serialize_field("hashLookupExpr", v)?; + physical_expr_node::ExprType::HashTableLookupExpr(v) => { + struct_ser.serialize_field("hashTableLookupExpr", v)?; } } } @@ -18785,8 +18785,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "lambdaVariable", "range_expr", "rangeExpr", - "hash_lookup_expr", - "hashLookupExpr", + "hash_table_lookup_expr", + "hashTableLookupExpr", ]; #[allow(clippy::enum_variant_names)] @@ -18817,7 +18817,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { Lambda, LambdaVariable, RangeExpr, - HashLookupExpr, + HashTableLookupExpr, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -18865,7 +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), - "hashLookupExpr" | "hash_lookup_expr" => Ok(GeneratedField::HashLookupExpr), + "hashTableLookupExpr" | "hash_table_lookup_expr" => Ok(GeneratedField::HashTableLookupExpr), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -19072,11 +19072,11 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::RangeExpr) ; } - GeneratedField::HashLookupExpr => { + GeneratedField::HashTableLookupExpr => { if expr_type__.is_some() { - return Err(serde::de::Error::duplicate_field("hashLookupExpr")); + return Err(serde::de::Error::duplicate_field("hashTableLookupExpr")); } - expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::HashLookupExpr) + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::HashTableLookupExpr) ; } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index caf2d52ea7ebd..fc5487316d22b 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1623,7 +1623,7 @@ pub mod physical_expr_node { #[prost(message, tag = "27")] RangeExpr(super::PhysicalRangeExprNode), #[prost(message, tag = "28")] - HashLookupExpr(super::PhysicalHashTableLookupExprNode), + HashTableLookupExpr(super::PhysicalHashTableLookupExprNode), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1880,18 +1880,18 @@ pub mod physical_hash_table_lookup_expr_node { #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] pub enum Map { #[prost(message, tag = "4")] - HashMapMembership(super::HashSetMapMembershipNode), + HashMapMembership(super::HashMapMembership), #[prost(message, tag = "5")] - ArrayMapMembership(super::ArrayMapMembershipNode), + ArrayMapMembership(super::ArrayMapMembership), } } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] -pub struct HashSetMapMembershipNode { +pub struct HashMapMembership { #[prost(fixed64, repeated, tag = "1")] pub build_hashes: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] -pub struct ArrayMapMembershipNode { +pub struct ArrayMapMembership { #[prost(uint64, tag = "1")] pub offset: u64, #[prost(uint64, tag = "2")] 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(|| { From de235eafa215b5380242627647aeea39dcb37531 Mon Sep 17 00:00:00 2001 From: rjhallsted Date: Thu, 13 Aug 2026 14:30:52 -0600 Subject: [PATCH 3/8] Fix hashing --- .../physical-plan/src/joins/array_map.rs | 12 ++- .../joins/hash_join/partitioned_hash_eval.rs | 84 ++++++++++++------- 2 files changed, 67 insertions(+), 29 deletions(-) diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index 3417d3577ff76..f01b6a6f549f1 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. /// @@ -440,6 +440,16 @@ impl ArrayMap { ) -> Result { use arrow_data::bit_iterator::BitIndexIterator; + // Assert that i < num_slots is a valid index into node.presence + if node.num_slots.is_multiple_of(8) { + assert_eq_or_internal_err!(node.presence.len() as u64, node.num_slots / 8); + } else { + assert_eq_or_internal_err!( + node.presence.len() as u64, + (node.num_slots / 8) + 1 + ); + } + 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) { 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 b9eb43366b3f5..acf08341e4135 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 @@ -267,7 +267,7 @@ pub struct HashTableLookupExpr { /// Map to check against (hash table or array map) /// Safety: Post-deserialization, this is only safe for membership checks and /// should not be modified further. - map: Arc, + map: HashTableLookupExprMap, /// Description for display description: String, } @@ -291,7 +291,7 @@ impl HashTableLookupExpr { Self { on_columns, random_state, - map: Arc::new(HashTableLookupExprMap::Normal(map)), + map: HashTableLookupExprMap::Normal(map), description, } } @@ -314,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 } } @@ -361,7 +347,7 @@ impl PhysicalExpr for HashTableLookupExpr { Ok(Arc::new(HashTableLookupExpr { on_columns: children, random_state: self.random_state.clone(), - map: Arc::clone(&self.map), + map: self.map.clone(), description: self.description.clone(), })) } @@ -378,7 +364,7 @@ impl PhysicalExpr for HashTableLookupExpr { // Evaluate columns let join_keys = evaluate_columns(&self.on_columns, batch)?; - match self.map.as_ref() { + match &self.map { HashTableLookupExprMap::Normal(normal) => match &**normal { Map::HashMap(map) => { with_hashes(&join_keys, self.random_state.random_state(), |hashes| { @@ -504,15 +490,17 @@ impl HashTableLookupExpr { let map = match &hash_table_lookup_expr.map { Some(Map::HashMapMembership(membership)) => { - Arc::new(HashTableLookupExprMap::MembershipOnlyHashMap( + HashTableLookupExprMap::MembershipOnlyHashMap(Arc::new( JoinHashMembershipMap::new(&membership.build_hashes), )) } - Some(Map::ArrayMapMembership(membership)) => Arc::new( - HashTableLookupExprMap::MembershipOnlyArrayMap(JoinMembershipArrayMap( - ArrayMap::try_from_proto_membership_only(membership)?, - )), - ), + Some(Map::ArrayMapMembership(membership)) => { + HashTableLookupExprMap::MembershipOnlyArrayMap(Arc::new( + JoinMembershipArrayMap(ArrayMap::try_from_proto_membership_only( + membership, + )?), + )) + } None => return internal_err!("HashTableLookupExpr has no map"), }; @@ -536,13 +524,53 @@ 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(JoinHashMembershipMap), + MembershipOnlyHashMap(Arc), /// A membership-checks only version of an array map, only constructed via expression deserialization - MembershipOnlyArrayMap(JoinMembershipArrayMap), + 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. From 657e7691be510be870bef35df8401255f4d1ee19 Mon Sep 17 00:00:00 2001 From: rjhallsted Date: Fri, 14 Aug 2026 12:31:13 -0600 Subject: [PATCH 4/8] tests --- .../physical-plan/src/joins/array_map.rs | 16 +- .../joins/hash_join/partitioned_hash_eval.rs | 349 +++++++++++++++++- .../physical-plan/src/joins/join_hash_map.rs | 70 ---- datafusion/proto/tests/cases/plans/exprs.rs | 45 +-- 4 files changed, 363 insertions(+), 117 deletions(-) diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index f01b6a6f549f1..ad964c3741eba 100644 --- a/datafusion/physical-plan/src/joins/array_map.rs +++ b/datafusion/physical-plan/src/joins/array_map.rs @@ -424,9 +424,7 @@ impl ArrayMap { use datafusion_proto_models::protobuf; let bits = BooleanBuffer::collect_bool(self.data.len(), |i| self.data[i] != 0); - // SAFETY: This always succeeds. None of the 3 error conditions of - // into_vec() are present. - let presence = bits.sliced().into_vec().unwrap(); + let presence = bits.sliced().as_slice().to_vec(); protobuf::ArrayMapMembership { offset: self.offset, num_slots: self.data.len() as u64, @@ -441,14 +439,10 @@ impl ArrayMap { use arrow_data::bit_iterator::BitIndexIterator; // Assert that i < num_slots is a valid index into node.presence - if node.num_slots.is_multiple_of(8) { - assert_eq_or_internal_err!(node.presence.len() as u64, node.num_slots / 8); - } else { - assert_eq_or_internal_err!( - node.presence.len() as u64, - (node.num_slots / 8) + 1 - ); - } + 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; 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 acf08341e4135..a53192ecc4587 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 @@ -257,7 +257,7 @@ 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 @@ -725,6 +725,7 @@ mod tests { #[cfg(feature = "proto")] mod proto_tests { use super::*; + use arrow::array::Int64Array; use arrow::datatypes::{DataType, Field}; use datafusion_common::internal_datafusion_err; use datafusion_physical_expr_common::physical_expr::proto_decode::{ @@ -886,6 +887,352 @@ 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. + 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]); + 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); + } + + 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()]) + } + + /// 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 26f79fe3b59a2..29849c5d0920f 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -311,76 +311,6 @@ impl JoinHashMapType for JoinHashMapU64 { } } -/// Membership-only join map reconstructed from serialized distinct hashes. -/// Supports `contain_hashes` lookups only; it has no build-side rows. -pub struct JoinHashMembershipMap { - map: HashTable<(u64, ())>, -} -impl JoinHashMembershipMap { - pub fn new(hashes: &[u64]) -> Self { - let mut map = HashTable::with_capacity(hashes.len()); - - for h in hashes { - map.insert_unique(*h, (*h, ()), |(h, _)| *h); - } - - Self { map } - } -} - -impl JoinHashMapType for JoinHashMembershipMap { - fn extend_zero(&mut self, _len: usize) { - unimplemented!(); - } - - fn update_from_iter<'a>( - &mut self, - _iter: Box + Send + 'a>, - _deleted_offset: usize, - ) { - unimplemented!(); - } - - fn get_matched_indices<'a>( - &self, - _iter: Box + 'a>, - _deleted_offset: Option, - ) -> (Vec, Vec) { - unimplemented!(); - } - - fn get_matched_indices_with_limit_offset( - &self, - _hash_values: &[u64], - _valid_keys: Option<&NullBuffer>, - _limit: usize, - _offset: MapOffset, - _input_indices: &mut Vec, - _match_indices: &mut Vec, - ) -> Option { - unimplemented!() - } - - /// Returns a BooleanArray indicating which of the provided hashes exist in the map. - fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray { - contain_hashes(&self.map, hash_values) - } - - /// Returns `true` if the join hash map contains no entries. - fn is_empty(&self) -> bool { - self.map.is_empty() - } - - /// Returns the number of entries in the join hash map. - fn len(&self) -> usize { - self.map.len() - } - - fn hashes(&self) -> Vec { - self.map.iter().map(|(h, _)| *h).collect() - } -} - use crate::joins::MapOffset; use crate::joins::chain::traverse_chain; diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs index 518b4a62ce072..c22cfb49e8068 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,13 +91,13 @@ 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<()> { +fn roundtrip_hash_table_lookup_expr() -> Result<()> { use datafusion::physical_plan::joins::join_hash_map::JoinHashMapU32; use datafusion::physical_plan::joins::{HashTableLookupExpr, Map}; @@ -109,9 +105,9 @@ fn roundtrip_hash_table_lookup_expr_to_lit() -> Result<()> { 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)?]; + // Create a HashTableLookupExpr with a HashMap + let hash_map = Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(0)))); let lookup_expr: Arc = Arc::new(HashTableLookupExpr::new( on_columns, datafusion::physical_plan::joins::SeededRandomState::with_seed(0), @@ -121,28 +117,7 @@ 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)?); - - // Serialize - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - - let proto: PhysicalPlanNode = - PhysicalPlanNode::try_from_physical_plan(filter.clone(), &codec) - .expect("serialization should succeed"); - - // Deserialize - let result: Arc = proto - .try_into_physical_plan(&ctx.task_ctx(), &codec) - .expect("deserialization should succeed"); - - // 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))); - - Ok(()) + roundtrip_test(filter) } #[test] From 3a1e5a6f3202109dbaf3bb0e4c6838786781bf9c Mon Sep 17 00:00:00 2001 From: rjhallsted Date: Fri, 14 Aug 2026 13:48:22 -0600 Subject: [PATCH 5/8] cleanup --- .../physical-plan/src/joins/array_map.rs | 14 +- .../joins/hash_join/partitioned_hash_eval.rs | 128 +++++++++++++++++- .../physical-plan/src/joins/join_hash_map.rs | 8 ++ datafusion/physical-plan/src/joins/mod.rs | 2 +- .../proto-models/proto/datafusion.proto | 23 ++++ .../proto-models/src/generated/prost.rs | 23 ++++ datafusion/proto/tests/cases/plans/exprs.rs | 37 ++++- 7 files changed, 223 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index ad964c3741eba..667d7a14b32b7 100644 --- a/datafusion/physical-plan/src/joins/array_map.rs +++ b/datafusion/physical-plan/src/joins/array_map.rs @@ -169,10 +169,16 @@ impl ArrayMap { /// Creates a new [`ArrayMap`] from the given array of join keys. /// + /// `min_val` and `max_val` must be the actual minimum and maximum key + /// values in `array`, in the wrapped-u64 key domain (see the type docs). + /// /// 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."); @@ -418,7 +424,7 @@ impl ArrayMap { } #[cfg(feature = "proto")] - pub fn to_proto_membership_only( + pub(crate) fn to_proto_membership_only( &self, ) -> datafusion_proto_models::protobuf::ArrayMapMembership { use datafusion_proto_models::protobuf; @@ -433,10 +439,10 @@ impl ArrayMap { } #[cfg(feature = "proto")] - pub fn try_from_proto_membership_only( + pub(crate) fn try_from_proto_membership_only( node: &datafusion_proto_models::protobuf::ArrayMapMembership, ) -> Result { - use arrow_data::bit_iterator::BitIndexIterator; + use arrow::util::bit_iterator::BitIndexIterator; // Assert that i < num_slots is a valid index into node.presence assert_eq_or_internal_err!( 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 a53192ecc4587..fc3db380d8cc2 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 @@ -213,7 +213,7 @@ impl PhysicalExpr for HashExpr { use datafusion_proto_models::protobuf; let on_columns = ctx.encode_children_expressions(&self.on_columns)?; Ok(Some(protobuf::PhysicalExprNode { - expr_id: None, + expr_id: self.expression_id(), expr_type: Some(protobuf::physical_expr_node::ExprType::HashExpr( protobuf::PhysicalHashExprNode { on_columns, @@ -264,9 +264,9 @@ pub struct HashTableLookupExpr { on_columns: Vec, /// Random state for hashing (with seeds preserved for serialization) random_state: SeededRandomState, - /// Map to check against (hash table or array map) - /// Safety: Post-deserialization, this is only safe for membership checks and - /// should not be modified further. + /// 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, @@ -424,6 +424,15 @@ impl PhysicalExpr for HashTableLookupExpr { } } +/// Encode the map in its membership-only proto form. +/// +/// The encoded size is deliberately uncapped: `build_hashes` grows with the +/// number of distinct join-key hashes on the build side (8 bytes each), and +/// an `ArrayMap` presence bitmap with its key range (1 bit per possible key). +/// This matches the `InList` pushdown path, which serializes its full value +/// list without a limit. If a bound is ever needed, encode `lit(true)` in +/// place of an over-limit map — a weaker filter is always safe here because +/// the join re-verifies matches. #[cfg(feature = "proto")] fn try_map_to_proto_membership_only( map: &HashTableLookupExprMap, @@ -582,6 +591,9 @@ 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); } @@ -725,6 +737,7 @@ 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; @@ -1159,6 +1172,104 @@ mod tests { 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]); + 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(); + 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)]) } @@ -1190,6 +1301,15 @@ mod tests { 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 { diff --git a/datafusion/physical-plan/src/joins/join_hash_map.rs b/datafusion/physical-plan/src/joins/join_hash_map.rs index 29849c5d0920f..6788cc445a038 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -140,6 +140,14 @@ 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; } 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/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 14f8bd2243038..32b6ad349352e 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1204,9 +1204,20 @@ 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; @@ -1214,13 +1225,25 @@ message PhysicalHashTableLookupExprNode { } } +// 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; } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index fc5487316d22b..0bf4f6e40b387 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1864,12 +1864,23 @@ 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")] @@ -1885,17 +1896,29 @@ pub mod physical_hash_table_lookup_expr_node { 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, } diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs index c22cfb49e8068..20e45b38578c7 100644 --- a/datafusion/proto/tests/cases/plans/exprs.rs +++ b/datafusion/proto/tests/cases/plans/exprs.rs @@ -98,7 +98,9 @@ fn roundtrip_like() -> Result<()> { /// on deserialization as a map that supports only membership checks. #[test] fn roundtrip_hash_table_lookup_expr() -> Result<()> { - use datafusion::physical_plan::joins::join_hash_map::JoinHashMapU32; + use datafusion::physical_plan::joins::join_hash_map::{ + JoinHashMapType, JoinHashMapU32, + }; use datafusion::physical_plan::joins::{HashTableLookupExpr, Map}; // Create a simple schema and input plan @@ -106,8 +108,11 @@ fn roundtrip_hash_table_lookup_expr() -> Result<()> { let input = Arc::new(EmptyExec::new(schema.clone())); let on_columns = vec![col("col", &schema)?]; - // Create a HashTableLookupExpr with a HashMap - let hash_map = Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(0)))); + // 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), @@ -120,6 +125,32 @@ fn roundtrip_hash_table_lookup_expr() -> Result<()> { roundtrip_test(filter) } +/// 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 schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone())); + + // 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)?; + + 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(), + )); + + let filter = Arc::new(FilterExec::try_new(lookup_expr, input)?); + roundtrip_test(filter) +} + #[test] fn roundtrip_hash_expr() -> Result<()> { use datafusion::physical_plan::joins::{HashExpr, SeededRandomState}; From ee79f62bb7a3276a15601fc92fc47368de654d4a Mon Sep 17 00:00:00 2001 From: rjhallsted Date: Fri, 14 Aug 2026 14:11:33 -0600 Subject: [PATCH 6/8] more cleanup --- datafusion/physical-plan/src/joins/array_map.rs | 3 --- .../src/joins/hash_join/partitioned_hash_eval.rs | 12 ++---------- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index 667d7a14b32b7..edf3a82d1109d 100644 --- a/datafusion/physical-plan/src/joins/array_map.rs +++ b/datafusion/physical-plan/src/joins/array_map.rs @@ -169,9 +169,6 @@ impl ArrayMap { /// Creates a new [`ArrayMap`] from the given array of join keys. /// - /// `min_val` and `max_val` must be the actual minimum and maximum key - /// values in `array`, in the wrapped-u64 key domain (see the type docs). - /// /// Note: This function processes only the non-null values in the input `array`, /// ignoring any rows where the key is `NULL`. /// 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 fc3db380d8cc2..d7818031eca19 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 @@ -213,7 +213,7 @@ impl PhysicalExpr for HashExpr { use datafusion_proto_models::protobuf; let on_columns = ctx.encode_children_expressions(&self.on_columns)?; Ok(Some(protobuf::PhysicalExprNode { - expr_id: self.expression_id(), + expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::HashExpr( protobuf::PhysicalHashExprNode { on_columns, @@ -415,7 +415,7 @@ impl PhysicalExpr for HashTableLookupExpr { }); Ok(Some(protobuf::PhysicalExprNode { - expr_id: self.expression_id(), + expr_id: None, expr_type: Some(expr), })) } @@ -425,14 +425,6 @@ impl PhysicalExpr for HashTableLookupExpr { } /// Encode the map in its membership-only proto form. -/// -/// The encoded size is deliberately uncapped: `build_hashes` grows with the -/// number of distinct join-key hashes on the build side (8 bytes each), and -/// an `ArrayMap` presence bitmap with its key range (1 bit per possible key). -/// This matches the `InList` pushdown path, which serializes its full value -/// list without a limit. If a bound is ever needed, encode `lit(true)` in -/// place of an over-limit map — a weaker filter is always safe here because -/// the join re-verifies matches. #[cfg(feature = "proto")] fn try_map_to_proto_membership_only( map: &HashTableLookupExprMap, From e4ce61bb0600987f1b6847c9fe45c0cac6dff787 Mon Sep 17 00:00:00 2001 From: rjhallsted Date: Fri, 14 Aug 2026 14:36:58 -0600 Subject: [PATCH 7/8] add upgrade documentation --- .../library-user-guide/upgrading/55.0.0.md | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) 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. From f9d7d4e442af72c19306571a6a9251686a360a8e Mon Sep 17 00:00:00 2001 From: rjhallsted Date: Fri, 14 Aug 2026 15:04:54 -0600 Subject: [PATCH 8/8] gate tests requiring normal hash behavior --- .../src/joins/hash_join/partitioned_hash_eval.rs | 6 ++++++ 1 file changed, 6 insertions(+) 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 d7818031eca19..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 @@ -986,6 +986,9 @@ mod tests { )); // 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] @@ -1109,6 +1112,7 @@ mod tests { .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] @@ -1191,6 +1195,7 @@ mod tests { .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] @@ -1255,6 +1260,7 @@ mod tests { ], ) .unwrap(); + #[cfg(not(feature = "force_hash_collisions"))] assert_eq!(eval_lookup(&expr, &batch), [true, false, false, true]); assert_eq!( eval_lookup(&expr, &batch),