diff --git a/Cargo.lock b/Cargo.lock index 5b06130ccaa9..99ff2e3f7681 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2209,7 +2209,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -5937,7 +5937,7 @@ dependencies = [ [[package]] name = "greptime-proto" version = "0.1.0" -source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=b6b665d8b6fce4f1f0da49c8c79fb150636df4ff#b6b665d8b6fce4f1f0da49c8c79fb150636df4ff" +source = "git+https://github.com/GreptimeTeam/greptime-proto.git?rev=a26ddccae52214397e03c5646a77e1696ef7eea4#a26ddccae52214397e03c5646a77e1696ef7eea4" dependencies = [ "prost 0.14.1", "prost-types 0.14.1", @@ -6483,7 +6483,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -8962,7 +8962,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -11113,7 +11113,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22505a5c94da8e3b7c2996394d1c933236c4d743e81a410bcca4e6989fc066a4" dependencies = [ "bytes", - "heck 0.4.1", + "heck 0.5.0", "itertools 0.12.1", "log", "multimap", @@ -11133,8 +11133,8 @@ version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac6c3320f9abac597dcbc668774ef006702672474aad53c6d596b62e487b40b1" dependencies = [ - "heck 0.4.1", - "itertools 0.14.0", + "heck 0.5.0", + "itertools 0.10.5", "log", "multimap", "once_cell", @@ -11182,7 +11182,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11195,7 +11195,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9120690fafc389a67ba3803df527d0ec9cbbc9cc45e4cc20b332996dfb672425" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -12629,7 +12629,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -13482,7 +13482,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1961e2ef424c1424204d3a5d6975f934f56b6d50ff5732382d84ebf460e147f7" dependencies = [ - "heck 0.4.1", + "heck 0.5.0", "proc-macro2", "quote", "syn 2.0.117", @@ -13511,7 +13511,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -14503,7 +14503,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix 1.0.7", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -16254,7 +16254,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 1197547bdc63..e615faa3fd7c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -158,7 +158,7 @@ fs2 = "0.4" fst = "0.4.7" futures = "0.3" futures-util = "0.3" -greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "b6b665d8b6fce4f1f0da49c8c79fb150636df4ff" } +greptime-proto = { git = "https://github.com/GreptimeTeam/greptime-proto.git", rev = "a26ddccae52214397e03c5646a77e1696ef7eea4" } hex = "0.4" http = "1" humantime = "2.1" diff --git a/src/client/src/region.rs b/src/client/src/region.rs index 19c173bd05d6..17de962830f2 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -426,6 +426,7 @@ mod test { assert_eq!(msg, "blabla"); } + #[allow(deprecated)] #[test] fn test_build_remote_dyn_filter_request_sets_header_and_body() { let request = build_remote_dyn_filter_update_request( @@ -435,6 +436,7 @@ mod test { payload: vec![1, 2, 3], generation: 7, is_complete: false, + typed_payload: None, }, ); diff --git a/src/common/query/src/request.rs b/src/common/query/src/request.rs index 6dc62b0a89fc..5deadb8816a1 100644 --- a/src/common/query/src/request.rs +++ b/src/common/query/src/request.rs @@ -14,6 +14,7 @@ mod base64_serde; mod initial_remote_dyn_filter_reg; +mod join_hash_bloom; use std::sync::Arc; @@ -39,6 +40,11 @@ pub use self::initial_remote_dyn_filter_reg::{ INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, InitialDynFilterReg, InitialDynFilterRegs, InitialDynFilterSnapshot, }; +pub use self::join_hash_bloom::{ + BLOOM_ENCODER_BITS_PER_HASH, BLOOM_ENCODER_MAX_NUM_BITS_BUDGET, BLOOM_ENCODER_MIN_NUM_BITS, + BLOOM_ENCODER_NUM_PROBES, BLOOM_PROTO_ENVELOPE_OVERHEAD_ESTIMATE, JOIN_HASH_BLOOM_VERSION, + JoinHashBloomPayload, MAX_BLOOM_NUM_PROBES, MAX_BLOOM_RESIDUAL_BYTES, +}; use crate::error::{DynFilterPayloadTooLargeSnafu, Error as CommonQueryError}; pub const DYN_FILTER_PROTOCOL_VERSION: u32 = 1; @@ -61,9 +67,19 @@ pub enum DynFilterPayload { /// A serialized DataFusion [`PhysicalExpr`] encoded as a protobuf /// [`PhysicalExprNode`]. Datafusion(#[serde(with = "base64_serde::bytes")] Vec), + /// A join-hash Bloom filter over DataFusion `u64` join hashes. + JoinHashBloom(JoinHashBloomPayload), } impl DynFilterPayload { + /// Validates payload-level invariants that can be checked without a receiver schema. + pub fn validate(&self) -> DataFusionResult<()> { + match self { + Self::Datafusion(_) => Ok(()), + Self::JoinHashBloom(bloom) => bloom.validate(), + } + } + /// Encodes a DataFusion physical expression into a bounded dynamic filter payload. /// /// Runtime-only hash lookup predicates are degraded to `true` before encoding so @@ -96,7 +112,12 @@ impl DynFilterPayload { input_schema: &datafusion::arrow::datatypes::Schema, max_payload_bytes: usize, ) -> DataFusionResult> { - let Self::Datafusion(bytes) = self; + let Self::Datafusion(bytes) = self else { + return Err(DataFusionError::Plan( + "DynFilterPayload::decode_datafusion_expr called on non-Datafusion payload" + .to_string(), + )); + }; validate_payload_size(bytes.len(), max_payload_bytes).map_err(DataFusionError::from)?; let codec = DefaultPhysicalExtensionCodec {}; let proto = PhysicalExprNode::decode(bytes.as_slice()).map_err(|e| { @@ -108,6 +129,66 @@ impl DynFilterPayload { validate_decoded_payload_expr(&expr, input_schema)?; Ok(expr) } + + /// Returns the proto-encoded byte size of this payload for budget tracking. + /// + /// - `Datafusion`: length of the raw protobuf bytes. + /// - `JoinHashBloom`: `encoded_len()` of the whole typed + /// `RemoteDynFilterPayload` proto that carries the Bloom payload, + /// conservatively computed by converting to proto and measuring. + pub fn encoded_payload_bytes(&self) -> usize { + match self { + Self::Datafusion(bytes) => bytes.len(), + Self::JoinHashBloom(_bloom) => { + let proto = self.to_region_proto_payload(); + proto.encoded_len() + } + } + } + + /// Converts this payload to the corresponding gRPC + /// `RemoteDynFilterPayload` proto. + pub fn to_region_proto_payload(&self) -> api::v1::region::RemoteDynFilterPayload { + match self { + Self::Datafusion(bytes) => api::v1::region::RemoteDynFilterPayload { + kind: Some( + api::v1::region::remote_dyn_filter_payload::Kind::DatafusionPhysicalExpr( + bytes.clone(), + ), + ), + }, + Self::JoinHashBloom(bloom) => api::v1::region::RemoteDynFilterPayload { + kind: Some( + api::v1::region::remote_dyn_filter_payload::Kind::JoinHashBloom( + bloom.to_proto(), + ), + ), + }, + } + } + + /// Constructs a `DynFilterPayload` from a gRPC + /// `RemoteDynFilterPayload` proto. + /// + /// Rejects missing `kind` and validates the payload. + pub fn from_region_proto_payload( + payload: api::v1::region::RemoteDynFilterPayload, + ) -> DataFusionResult { + let kind = payload.kind.ok_or_else(|| { + DataFusionError::Plan("RemoteDynFilterPayload::kind is missing".to_string()) + })?; + + match kind { + api::v1::region::remote_dyn_filter_payload::Kind::DatafusionPhysicalExpr(bytes) => { + Ok(Self::Datafusion(bytes)) + } + api::v1::region::remote_dyn_filter_payload::Kind::JoinHashBloom(bloom) => { + let model = JoinHashBloomPayload::from_proto(&bloom)?; + model.validate()?; + Ok(Self::JoinHashBloom(model)) + } + } + } } fn encode_physical_expr_to_bytes(expr: &Arc) -> DataFusionResult> { diff --git a/src/common/query/src/request/initial_remote_dyn_filter_reg.rs b/src/common/query/src/request/initial_remote_dyn_filter_reg.rs index a4a00d5face1..187abd9d78d5 100644 --- a/src/common/query/src/request/initial_remote_dyn_filter_reg.rs +++ b/src/common/query/src/request/initial_remote_dyn_filter_reg.rs @@ -90,6 +90,15 @@ impl InitialDynFilterRegs { reg.filter_id )); } + + if let Some(snapshot) = ®.initial_snapshot { + snapshot.payload.validate().map_err(|err| { + format!( + "InitialDynFilterRegs contains invalid initial snapshot for filter_id '{}': {}", + reg.filter_id, err + ) + })?; + } } Ok(()) @@ -198,9 +207,7 @@ impl InitialDynFilterSnapshot { } pub fn encoded_payload_bytes(&self) -> usize { - match &self.payload { - DynFilterPayload::Datafusion(bytes) => bytes.len(), - } + self.payload.encoded_payload_bytes() } } @@ -214,6 +221,29 @@ mod tests { use datafusion_common::DataFusionError; use super::*; + use crate::request::join_hash_bloom::JoinHashBloomPayload; + + /// Convenience helper to create a test Bloom payload with minimal fields. + fn make_test_bloom( + num_bits: u64, + bitset: Vec, + residual: Vec, + join_key_child_indices: Vec, + ) -> JoinHashBloomPayload { + JoinHashBloomPayload { + version: 1, + df_seed0: 0, + df_seed1: 0, + df_seed2: 0, + df_seed3: 0, + num_bits, + num_probes: 2, + bitset, + join_key_child_indices, + residual_datafusion_physical_expr: residual, + hash_compat_fingerprint: 0, + } + } #[test] fn initial_dyn_filter_regs_json_round_trip() { @@ -393,4 +423,83 @@ mod tests { assert!(matches!(err, DataFusionError::Plan(_))); } + + // ── Bloom snapshot tests ────────────────────────────────── + + #[test] + fn initial_dyn_filter_snapshot_bloom_encoded_payload_bytes_counts_full_proto() { + let bitset = vec![0u8; 128]; + let bloom = make_test_bloom(1024, bitset.clone(), vec![1, 2, 3], vec![0]); + + let snapshot = + InitialDynFilterSnapshot::new(DynFilterPayload::JoinHashBloom(bloom.clone()), 2, false); + + let encoded_bytes = snapshot.encoded_payload_bytes(); + + // Should be more than bitset + residual alone (includes proto overhead) + assert!( + encoded_bytes > bitset.len() + bloom.residual_datafusion_physical_expr.len(), + "expected > {}, got {encoded_bytes}", + bitset.len() + bloom.residual_datafusion_physical_expr.len() + ); + } + + #[test] + fn initial_dyn_filter_reg_encoded_registration_bytes_include_bloom_snapshot() { + let bloom = make_test_bloom(64, vec![0u8; 8], vec![1], vec![0]); + + let snapshot = + InitialDynFilterSnapshot::new(DynFilterPayload::JoinHashBloom(bloom), 2, false); + + let reg = InitialDynFilterReg::new("filter-bloom", vec![vec![1, 2, 3]]) + .with_initial_snapshot(snapshot); + + let child_bytes = reg.encoded_child_expr_bytes(); + let total_bytes = reg.encoded_registration_bytes(); + + assert_eq!(child_bytes, 3); + // The Bloom payload proto encoded size should be > 0 + assert!( + total_bytes > child_bytes, + "expected total_bytes > {child_bytes}, got {total_bytes}" + ); + } + + #[test] + fn initial_dyn_filter_regs_json_round_trip_with_bloom_snapshot() { + let bloom = make_test_bloom(256, vec![0u8; 32], vec![4, 5, 6], vec![0]); + + let regs = InitialDynFilterRegs::new(vec![ + InitialDynFilterReg::new("filter-bloom", vec![vec![1, 2, 3]]).with_initial_snapshot( + InitialDynFilterSnapshot::new(DynFilterPayload::JoinHashBloom(bloom), 7, true), + ), + ]); + + let encoded = regs.to_extension_value().unwrap(); + let json: serde_json::Value = serde_json::from_str(&encoded).unwrap(); + let decoded = InitialDynFilterRegs::from_extension_value(&encoded).unwrap(); + + assert_eq!( + json["registrations"][0]["initial_snapshot"]["payload"]["kind"], + "join_hash_bloom" + ); + assert_eq!(decoded, regs); + } + + #[test] + fn initial_dyn_filter_regs_from_extension_rejects_invalid_bloom_snapshot() { + let bloom = make_test_bloom(256, vec![0u8; 32], vec![4, 5, 6], vec![]); + + let regs = InitialDynFilterRegs::new(vec![ + InitialDynFilterReg::new("filter-bloom", vec![vec![1, 2, 3]]).with_initial_snapshot( + InitialDynFilterSnapshot::new(DynFilterPayload::JoinHashBloom(bloom), 7, true), + ), + ]); + + let encoded = serde_json::to_string(®s).unwrap(); + let err = InitialDynFilterRegs::from_extension_value(&encoded).unwrap_err(); + + assert!(err.to_string().contains("invalid initial snapshot")); + assert!(err.to_string().contains("join_key_child_indices is empty")); + } } diff --git a/src/common/query/src/request/join_hash_bloom.rs b/src/common/query/src/request/join_hash_bloom.rs new file mode 100644 index 000000000000..2b908b45f162 --- /dev/null +++ b/src/common/query/src/request/join_hash_bloom.rs @@ -0,0 +1,666 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Join-hash Bloom filter payload model for remote dynamic filter updates. +//! +//! The Bloom payload carries a bitset over DataFusion join-hash `u64` values +//! along with the hash seeds, join key child indices, and a residual +//! DataFusion expression where runtime `HashTableLookupExpr` nodes have been +//! replaced by `lit(true)`. Receivers combine the residual expression with a +//! Bloom probe expression via `AND` to obtain the final predicate. + +use std::collections::HashSet; + +use datafusion_common::{DataFusionError, Result as DataFusionResult}; +use serde::{Deserialize, Serialize}; + +use super::base64_serde; + +/// Current wire-format version for join-hash Bloom payloads. +pub const JOIN_HASH_BLOOM_VERSION: u32 = 1; + +/// Upper bound on the number of Bloom hash probes per lookup key. +/// Keeps probe cost bounded and prevents obviously misconfigured payloads. +pub const MAX_BLOOM_NUM_PROBES: u32 = 64; + +/// Upper bound on the residual DataFusion expression bytes carried in a +/// Bloom payload. The residual is an optional safety net; oversized +/// residuals indicate a producer-side budget violation. +pub const MAX_BLOOM_RESIDUAL_BYTES: usize = 8192; + +/// Conservative estimate of proto envelope overhead for a Bloom payload +/// when wrapped inside `RemoteDynFilterPayload`. Used during encoding to +/// avoid overshooting the payload budget. +pub const BLOOM_PROTO_ENVELOPE_OVERHEAD_ESTIMATE: usize = 256; + +/// Default number of Bloom hash probes used by the encoder. +pub const BLOOM_ENCODER_NUM_PROBES: u32 = 5; + +/// Bits allocated per distinct hash in the Bloom filter. +pub const BLOOM_ENCODER_BITS_PER_HASH: u64 = 10; + +/// Minimum Bloom size in bits. +pub const BLOOM_ENCODER_MIN_NUM_BITS: u64 = 64; + +/// Largest `num_bits` the encoder will ever allocate (hard OOM guard). +pub const BLOOM_ENCODER_MAX_NUM_BITS_BUDGET: u64 = 256 * 1024; // 256 KiB = 32 KiB bytes + +/// Serialized Bloom filter payload for remote dynamic filter updates. +/// +/// Carries a Bloom bitset over DataFusion join-hash `u64` values together +/// with the parameters needed to reconstruct the probe-side hash computation +/// and a residual DataFusion expression with `HashTableLookupExpr` replaced +/// by `lit(true)`. +/// +/// Byte fields are base64-encoded in JSON to keep the wire representation +/// compact and safe. +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct JoinHashBloomPayload { + pub version: u32, + pub df_seed0: u64, + pub df_seed1: u64, + pub df_seed2: u64, + pub df_seed3: u64, + pub num_bits: u64, + pub num_probes: u32, + #[serde(with = "base64_serde::bytes")] + pub bitset: Vec, + pub join_key_child_indices: Vec, + #[serde(with = "base64_serde::bytes")] + pub residual_datafusion_physical_expr: Vec, + /// Cross-platform hash-compatibility fingerprint. + /// + /// Computed deterministically by running a `HashExpr` over a canonical + /// `RecordBatch` built from `input_schema`. Zero means "not computed". + pub hash_compat_fingerprint: u64, +} + +impl JoinHashBloomPayload { + /// Fail-closed validation of payload invariants. + /// + /// Rejects: + /// - unsupported version + /// - `num_bits == 0` + /// - `num_bits` not representable as `usize` + /// - `num_probes == 0` or exceeding [`MAX_BLOOM_NUM_PROBES`] + /// - bitset length not equal to `ceil(num_bits / 8)` + /// - empty or duplicate join key child indices + /// - empty residual expression or residual expression exceeding + /// [`MAX_BLOOM_RESIDUAL_BYTES`] + pub fn validate(&self) -> DataFusionResult<()> { + if self.version != JOIN_HASH_BLOOM_VERSION { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: unsupported version {} (expected {})", + self.version, JOIN_HASH_BLOOM_VERSION + ))); + } + + if self.num_bits == 0 { + return Err(DataFusionError::Plan( + "JoinHashBloomPayload: num_bits is zero".to_string(), + )); + } + + if self.num_bits > usize::MAX as u64 { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: num_bits {} exceeds usize::MAX", + self.num_bits + ))); + } + + if self.num_probes == 0 { + return Err(DataFusionError::Plan( + "JoinHashBloomPayload: num_probes is zero".to_string(), + )); + } + + if self.num_probes > MAX_BLOOM_NUM_PROBES { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: num_probes {} exceeds limit {}", + self.num_probes, MAX_BLOOM_NUM_PROBES + ))); + } + + let expected_bitset_len = num_bits_ceil_bytes(self.num_bits as usize); + if self.bitset.len() != expected_bitset_len { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: bitset length {} does not match ceil(num_bits {}/8) = {}", + self.bitset.len(), + self.num_bits, + expected_bitset_len + ))); + } + + if self.join_key_child_indices.is_empty() { + return Err(DataFusionError::Plan( + "JoinHashBloomPayload: join_key_child_indices is empty".to_string(), + )); + } + + let mut seen = HashSet::with_capacity(self.join_key_child_indices.len()); + for &child_index in &self.join_key_child_indices { + if !seen.insert(child_index) { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: duplicate join key child index {}", + child_index + ))); + } + } + + if self.residual_datafusion_physical_expr.is_empty() { + return Err(DataFusionError::Plan( + "JoinHashBloomPayload: residual expression is empty".to_string(), + )); + } + + if self.residual_datafusion_physical_expr.len() > MAX_BLOOM_RESIDUAL_BYTES { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: residual expression {} bytes exceeds budget {}", + self.residual_datafusion_physical_expr.len(), + MAX_BLOOM_RESIDUAL_BYTES + ))); + } + + Ok(()) + } + + /// Converts this model to the corresponding proto `JoinHashBloomPayload`. + pub fn to_proto(&self) -> api::v1::region::JoinHashBloomPayload { + api::v1::region::JoinHashBloomPayload { + version: self.version, + df_seed0: self.df_seed0, + df_seed1: self.df_seed1, + df_seed2: self.df_seed2, + df_seed3: self.df_seed3, + num_bits: self.num_bits, + num_probes: self.num_probes, + bitset: self.bitset.clone(), + join_key_child_indices: self.join_key_child_indices.clone(), + residual_datafusion_physical_expr: self.residual_datafusion_physical_expr.clone(), + hash_compat_fingerprint: self.hash_compat_fingerprint, + } + } + + /// Constructs this model from a proto `JoinHashBloomPayload`. + /// + /// Does *not* validate — call [`Self::validate`] separately. + pub fn from_proto(payload: &api::v1::region::JoinHashBloomPayload) -> DataFusionResult { + Ok(Self { + version: payload.version, + df_seed0: payload.df_seed0, + df_seed1: payload.df_seed1, + df_seed2: payload.df_seed2, + df_seed3: payload.df_seed3, + num_bits: payload.num_bits, + num_probes: payload.num_probes, + bitset: payload.bitset.clone(), + join_key_child_indices: payload.join_key_child_indices.clone(), + residual_datafusion_physical_expr: payload.residual_datafusion_physical_expr.clone(), + hash_compat_fingerprint: payload.hash_compat_fingerprint, + }) + } + + /// Returns the expected bitset byte length for `num_bits`, i.e. + /// `ceil(num_bits / 8)`. + pub fn expected_bitset_bytes(&self) -> usize { + num_bits_ceil_bytes(self.num_bits as usize) + } + + /// Tests whether the given join hash is a member of this Bloom filter. + /// + /// Uses SplitMix64 double-hashing to derive the probe sequence. + /// This is deterministic and false-negative-free for any hash that was + /// previously inserted into the bitset with the same parameters. + #[inline] + pub fn contains_join_hash(&self, hash: u64) -> bool { + bloom_contains(hash, &self.bitset, self.num_bits, self.num_probes) + } + + /// Inserts a join hash into this Bloom filter's bitset in-place. + /// + /// This is used during encoding to build the bitset from distinct join hashes. + #[inline] + pub fn insert_join_hash(&mut self, hash: u64) { + bloom_insert(hash, &mut self.bitset, self.num_bits, self.num_probes) + } + + /// Builds a new Bloom payload from distinct join hashes. + /// + /// Allocates a fresh zeroed bitset of the correct size and inserts every + /// hash. The returned payload passes [`Self::validate`] when `num_bits`, + /// `num_probes`, and the other required fields are set correctly. + /// + /// The `hash_compat_fingerprint` is left at zero; the encoder should set it + /// after construction. + pub fn build_from_hashes( + hashes: impl IntoIterator, + num_bits: u64, + num_probes: u32, + seeds: (u64, u64, u64, u64), + join_key_child_indices: Vec, + residual_datafusion_physical_expr: Vec, + ) -> Self { + let bitset_len = num_bits_ceil_bytes(num_bits as usize); + let mut bitset = vec![0u8; bitset_len]; + bloom_insert_many(hashes, &mut bitset, num_bits, num_probes); + Self { + version: JOIN_HASH_BLOOM_VERSION, + df_seed0: seeds.0, + df_seed1: seeds.1, + df_seed2: seeds.2, + df_seed3: seeds.3, + num_bits, + num_probes, + bitset, + join_key_child_indices, + residual_datafusion_physical_expr, + hash_compat_fingerprint: 0, + } + } + + /// Fallible builder: validates parameters before allocation and insertion. + /// + /// Checks `num_bits`, `num_probes`, join key child indices, and residual before + /// allocating the bitset. The returned payload is already validated via + /// [`Self::validate`]. + pub fn try_build_from_hashes( + hashes: impl IntoIterator, + num_bits: u64, + num_probes: u32, + seeds: (u64, u64, u64, u64), + join_key_child_indices: Vec, + residual_datafusion_physical_expr: Vec, + ) -> DataFusionResult { + // --- validate before allocating --- + + if num_bits == 0 { + return Err(DataFusionError::Plan( + "JoinHashBloomPayload: try_build_from_hashes num_bits is zero".to_string(), + )); + } + if num_bits > BLOOM_ENCODER_MAX_NUM_BITS_BUDGET { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: try_build_from_hashes num_bits {} exceeds budget {}", + num_bits, BLOOM_ENCODER_MAX_NUM_BITS_BUDGET + ))); + } + if num_bits > usize::MAX as u64 { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: try_build_from_hashes num_bits {} exceeds usize::MAX", + num_bits + ))); + } + if num_probes == 0 { + return Err(DataFusionError::Plan( + "JoinHashBloomPayload: try_build_from_hashes num_probes is zero".to_string(), + )); + } + if num_probes > MAX_BLOOM_NUM_PROBES { + return Err(DataFusionError::Plan(format!( + "JoinHashBloomPayload: try_build_from_hashes num_probes {} exceeds limit {}", + num_probes, MAX_BLOOM_NUM_PROBES + ))); + } + if join_key_child_indices.is_empty() { + return Err(DataFusionError::Plan( + "JoinHashBloomPayload: try_build_from_hashes join_key_child_indices is empty" + .to_string(), + )); + } + if residual_datafusion_physical_expr.is_empty() { + return Err(DataFusionError::Plan( + "JoinHashBloomPayload: try_build_from_hashes residual expression is empty" + .to_string(), + )); + } + + // --- allocate and build --- + + let payload = Self::build_from_hashes( + hashes, + num_bits, + num_probes, + seeds, + join_key_child_indices, + residual_datafusion_physical_expr, + ); + + payload.validate()?; + + Ok(payload) + } + + /// Compute a Bloom size recommendation given a number of distinct hashes + /// and a payload byte budget. + /// + /// Returns `None` if the budget is too small to accommodate even the + /// minimum Bloom, the residual, and proto envelope overhead. + pub fn compute_bloom_sizing( + distinct_hash_count: u64, + residual_bytes_len: usize, + max_payload_bytes: usize, + ) -> Option<(u64, u32)> { + let num_bits = std::cmp::max( + BLOOM_ENCODER_MIN_NUM_BITS, + distinct_hash_count.saturating_mul(BLOOM_ENCODER_BITS_PER_HASH), + ); + let num_probes = BLOOM_ENCODER_NUM_PROBES; + + let bitset_bytes = num_bits_ceil_bytes(num_bits as usize); + + // Quick budget check: residual + bitset + envelope overhead must fit + let estimated_total = residual_bytes_len + .saturating_add(bitset_bytes) + .saturating_add(BLOOM_PROTO_ENVELOPE_OVERHEAD_ESTIMATE); + + if estimated_total > max_payload_bytes { + // Shrink num_bits to fit, trading off false-positive rate + let available_for_bitset = max_payload_bytes + .saturating_sub(residual_bytes_len) + .saturating_sub(BLOOM_PROTO_ENVELOPE_OVERHEAD_ESTIMATE); + let max_bits = (available_for_bitset as u64) + .saturating_mul(8) + .min(BLOOM_ENCODER_MAX_NUM_BITS_BUDGET); + if max_bits < BLOOM_ENCODER_MIN_NUM_BITS { + return None; + } + Some((max_bits, num_probes)) + } else if num_bits > BLOOM_ENCODER_MAX_NUM_BITS_BUDGET { + // Cap at hard budget even if there is room + Some((BLOOM_ENCODER_MAX_NUM_BITS_BUDGET, num_probes)) + } else { + Some((num_bits, num_probes)) + } + } +} + +/// `ceil(n / 8)` without floating point. +#[inline] +fn num_bits_ceil_bytes(num_bits: usize) -> usize { + num_bits.div_ceil(8) +} + +// --------------------------------------------------------------------------- +// SplitMix64 double-hash Bloom helpers +// --------------------------------------------------------------------------- + +/// SplitMix64 — a fast, high-quality 64-bit mixing function. +/// +/// Each round avalanches all input bits. This is the same SplitMix64 used by +/// many hash-table and Bloom filter implementations. +#[inline] +fn splitmix64(mut x: u64) -> u64 { + x = x.wrapping_add(0x9e3779b97f4a7c15); + x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); + x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb); + x ^ (x >> 31) +} + +/// Returns the bit position for the k-th probe of the given hash using +/// SplitMix64 double-hashing. +/// +/// - `h1 = splitmix64(hash)` +/// - `delta = splitmix64(h1)` +/// - `bit = (h1 + k * delta) % num_bits` +#[inline] +fn bloom_bit_index(hash: u64, k: u32, num_bits: u64) -> u64 { + let h1 = splitmix64(hash); + let delta = splitmix64(h1); + h1.wrapping_add((k as u64).wrapping_mul(delta)) % num_bits +} + +/// Returns `true` if `hash` is present in the Bloom filter. +/// +/// `bitset` is a byte slice of length `ceil(num_bits / 8)`. Bit 0 of each +/// byte is the least-significant bit. +fn bloom_contains(hash: u64, bitset: &[u8], num_bits: u64, num_probes: u32) -> bool { + for k in 0..num_probes { + let bit = bloom_bit_index(hash, k, num_bits); + let byte_idx = (bit >> 3) as usize; + let bit_mask = 1u8 << (bit & 7); + if bitset[byte_idx] & bit_mask == 0 { + return false; + } + } + true +} + +/// Sets the bits for `hash` in the Bloom filter. +fn bloom_insert(hash: u64, bitset: &mut [u8], num_bits: u64, num_probes: u32) { + for k in 0..num_probes { + let bit = bloom_bit_index(hash, k, num_bits); + let byte_idx = (bit >> 3) as usize; + let bit_mask = 1u8 << (bit & 7); + bitset[byte_idx] |= bit_mask; + } +} + +/// Inserts many hashes and returns the number of *distinct* hashes inserted. +fn bloom_insert_many( + hashes: impl IntoIterator, + bitset: &mut [u8], + num_bits: u64, + num_probes: u32, +) -> u64 { + let mut distinct = 0u64; + let mut seen = std::collections::HashSet::new(); + for hash in hashes { + if seen.insert(hash) { + distinct += 1; + bloom_insert(hash, bitset, num_bits, num_probes); + } + } + distinct +} + +#[cfg(test)] +mod tests { + use super::*; + + fn minimal_valid_payload() -> JoinHashBloomPayload { + JoinHashBloomPayload { + version: JOIN_HASH_BLOOM_VERSION, + df_seed0: 4, + df_seed1: 5, + df_seed2: 6, + df_seed3: 7, + num_bits: 1024, + num_probes: 3, + bitset: vec![0u8; 128], + join_key_child_indices: vec![0, 1], + residual_datafusion_physical_expr: vec![1, 2, 3], + hash_compat_fingerprint: 0, + } + } + + #[test] + fn validate_accepts_valid_payload() { + minimal_valid_payload().validate().unwrap(); + } + + #[test] + fn validate_rejections_table_driven() { + #[track_caller] + fn case(mutate: impl FnOnce(&mut JoinHashBloomPayload), expected_substr: &str) { + let mut p = minimal_valid_payload(); + mutate(&mut p); + let err = p.validate().unwrap_err(); + assert!( + err.to_string().contains(expected_substr), + "expected error containing '{expected_substr}', got: {}", + err + ); + } + + case(|p| p.version = 99, "unsupported version 99"); + case(|p| p.num_bits = 0, "num_bits is zero"); + case(|p| p.num_probes = 0, "num_probes is zero"); + case(|p| p.num_probes = MAX_BLOOM_NUM_PROBES + 1, "exceeds limit"); + case(|p| p.bitset = vec![0u8; 127], "bitset length"); + case( + |p| p.join_key_child_indices = vec![0, 1, 0], + "duplicate join key child index 0", + ); + case( + |p| p.residual_datafusion_physical_expr = vec![0u8; MAX_BLOOM_RESIDUAL_BYTES + 1], + "exceeds budget", + ); + case( + |p| p.join_key_child_indices = vec![], + "join_key_child_indices is empty", + ); + case( + |p| p.residual_datafusion_physical_expr = vec![], + "residual expression is empty", + ); + } + + #[test] + fn validate_num_bits_not_divisible_by_8() { + let mut p = minimal_valid_payload(); + p.num_bits = 1025; + p.bitset = vec![0u8; 129]; // ceil(1025/8) = 129 + p.validate().unwrap(); + + p.bitset = vec![0u8; 128]; + let err = p.validate().unwrap_err(); + assert!(err.to_string().contains("bitset length")); + } + + #[test] + fn compute_bloom_sizing_shrink_path_respects_hard_cap() { + let hard_cap_bytes = num_bits_ceil_bytes(BLOOM_ENCODER_MAX_NUM_BITS_BUDGET as usize); + let max_payload_bytes = hard_cap_bytes + BLOOM_PROTO_ENVELOPE_OVERHEAD_ESTIMATE + 1024; + + let (num_bits, num_probes) = JoinHashBloomPayload::compute_bloom_sizing( + BLOOM_ENCODER_MAX_NUM_BITS_BUDGET, + 0, + max_payload_bytes, + ) + .unwrap(); + + assert_eq!(num_bits, BLOOM_ENCODER_MAX_NUM_BITS_BUDGET); + assert_eq!(num_probes, BLOOM_ENCODER_NUM_PROBES); + + JoinHashBloomPayload::try_build_from_hashes( + [1, 2, 3], + num_bits, + num_probes, + (0, 1, 2, 3), + vec![0], + vec![1], + ) + .unwrap(); + } + + #[test] + fn proto_roundtrip_preserves_fields() { + let original = minimal_valid_payload(); + let proto = original.to_proto(); + let decoded = JoinHashBloomPayload::from_proto(&proto).unwrap(); + assert_eq!(decoded, original); + } + + #[test] + fn json_serde_roundtrip_uses_base64() { + let original = minimal_valid_payload(); + let json = serde_json::to_string(&original).unwrap(); + let decoded: JoinHashBloomPayload = serde_json::from_str(&json).unwrap(); + assert_eq!(decoded, original); + + // Verify bitset and residual are base64-encoded (not raw bytes) + let value: serde_json::Value = serde_json::from_str(&json).unwrap(); + let bitset_str = value["bitset"].as_str().unwrap(); + let residual_str = value["residual_datafusion_physical_expr"].as_str().unwrap(); + // Should be valid base64 + use base64::Engine; + use base64::prelude::BASE64_STANDARD; + BASE64_STANDARD.decode(bitset_str).unwrap(); + BASE64_STANDARD.decode(residual_str).unwrap(); + // And should not be raw byte arrays + assert!(!bitset_str.starts_with('[')); + } + + #[test] + fn json_rejects_invalid_base64_in_bitset() { + let err = serde_json::from_str::( + r#"{"version":1,"df_seed0":0,"df_seed1":0,"df_seed2":0,"df_seed3":0,"num_bits":8,"num_probes":1,"bitset":"!!!not.base64!!!","join_key_child_indices":[],"residual_datafusion_physical_expr":"","hash_compat_fingerprint":0}"#, + ) + .unwrap_err(); + assert!(err.to_string().contains("invalid base64")); + } + + #[test] + fn bloom_insert_contains_no_false_negatives() { + let hashes: Vec = (0..100).map(|i| i * 7 + 13).collect(); + let payload = JoinHashBloomPayload::build_from_hashes( + hashes.clone(), + 1024, + 3, + (1, 2, 3, 4), + vec![0], + vec![1], + ); + + // Every inserted hash must be found + for &h in &hashes { + assert!( + payload.contains_join_hash(h), + "hash {} should be in the Bloom filter", + h + ); + } + + // Hash not inserted should usually be absent (may have false positives) + // but we can at least verify the method doesn't crash + let _ = payload.contains_join_hash(u64::MAX); + } + + #[test] + fn bloom_roundtrip_no_false_negatives_table_driven() { + type Roundtrip = fn(&JoinHashBloomPayload) -> JoinHashBloomPayload; + + fn json_roundtrip(p: &JoinHashBloomPayload) -> JoinHashBloomPayload { + serde_json::from_str(&serde_json::to_string(p).unwrap()).unwrap() + } + + fn proto_roundtrip(p: &JoinHashBloomPayload) -> JoinHashBloomPayload { + JoinHashBloomPayload::from_proto(&p.to_proto()).unwrap() + } + + let hashes = vec![42u64, 12345, 98765]; + let original = JoinHashBloomPayload::build_from_hashes( + hashes.clone(), + 256, + 2, + (0, 0, 0, 0), + vec![0], + vec![1], + ); + + // Test both JSON and proto roundtrip preserve all inserted hashes + let test_cases: &[(&str, Roundtrip)] = + &[("json", json_roundtrip), ("proto", proto_roundtrip)]; + + for (name, decode) in test_cases { + let decoded = decode(&original); + for &h in &hashes { + assert!( + decoded.contains_join_hash(h), + "[{name}] hash {h} should be in the Bloom filter after roundtrip" + ); + } + } + } +} diff --git a/src/datanode/src/region_server/remote_dyn_filter.rs b/src/datanode/src/region_server/remote_dyn_filter.rs index 1317dc050217..2034917fd254 100644 --- a/src/datanode/src/region_server/remote_dyn_filter.rs +++ b/src/datanode/src/region_server/remote_dyn_filter.rs @@ -170,6 +170,7 @@ impl RegionServer { } } + #[allow(deprecated)] async fn handle_remote_dyn_filter_update( &self, query_id: &QueryId, @@ -339,6 +340,7 @@ impl Stream for RemoteDynFilterGuardedStream { } #[cfg(test)] +#[allow(deprecated)] mod tests { use std::assert_matches; use std::collections::{HashMap, HashSet}; @@ -729,6 +731,7 @@ mod tests { payload: vec![1], generation: 1, is_complete: false, + typed_payload: None, }, )), }) @@ -754,6 +757,7 @@ mod tests { payload: Vec::new(), generation: 1, is_complete: false, + typed_payload: None, }, )), }) diff --git a/src/query/src/dist_plan/remote_dyn_filter_registry.rs b/src/query/src/dist_plan/remote_dyn_filter_registry.rs index fdb9d6680459..5abac13a57a8 100644 --- a/src/query/src/dist_plan/remote_dyn_filter_registry.rs +++ b/src/query/src/dist_plan/remote_dyn_filter_registry.rs @@ -430,6 +430,7 @@ async fn fanout_snapshot_for_query( } #[allow(clippy::too_many_arguments)] +#[allow(deprecated)] async fn fanout_update_for_query( query_id: QueryId, region_query_handler: &RegionQueryHandlerRef, @@ -449,6 +450,7 @@ async fn fanout_update_for_query( payload: payload.clone(), generation, is_complete, + typed_payload: None, }; match await_control_rpc_or_lifecycle_close( @@ -903,6 +905,7 @@ mod tests { unreachable!("remote dyn filter registry tests should not execute remote queries") } + #[allow(deprecated)] async fn handle_remote_dyn_filter_update( &self, region_id: RegionId,