From 867be3589514b7acb27b52f1c0f25e1a129ecbb8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 14 Jul 2026 10:54:27 -0600 Subject: [PATCH] perf: optimize might_contain_longs in datafusion-comet-spark-expr --- native/spark-expr/Cargo.toml | 4 + .../benches/bloom_filter_might_contain.rs | 114 ++++++++++++++ .../src/bloom_filter/spark_bloom_filter.rs | 141 ++++++++++++++---- native/spark-expr/src/hash_funcs/murmur3.rs | 88 +++++++---- 4 files changed, 288 insertions(+), 59 deletions(-) create mode 100644 native/spark-expr/benches/bloom_filter_might_contain.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..184930596e 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -139,3 +139,7 @@ harness = false [[bench]] name = "to_json" harness = false + +[[bench]] +name = "bloom_filter_might_contain" +harness = false diff --git a/native/spark-expr/benches/bloom_filter_might_contain.rs b/native/spark-expr/benches/bloom_filter_might_contain.rs new file mode 100644 index 0000000000..ee6755c4e9 --- /dev/null +++ b/native/spark-expr/benches/bloom_filter_might_contain.rs @@ -0,0 +1,114 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +use arrow::array::{ArrayRef, Int64Array}; +use arrow::datatypes::{DataType, Field}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::common::ScalarValue; +use datafusion::config::ConfigOptions; +use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion::physical_expr::expressions::Literal; +use datafusion::physical_expr::PhysicalExpr; +use datafusion_comet_spark_expr::murmur3::spark_compatible_murmur3_hash; +use datafusion_comet_spark_expr::BloomFilterMightContain; +use std::hint::black_box; +use std::sync::Arc; + +/// Serialize a Spark V1 bloom filter (`[version][numHashFunctions][numWords][words...]`, all +/// big-endian) containing `items`, mirroring Spark's `BloomFilterImpl.putLong`. +fn serialize_filter(num_hash_functions: u32, num_words: usize, items: &[i64]) -> Vec { + let mut words = vec![0u64; num_words]; + let bit_size = (num_words * 64) as i32; + for item in items { + let h1 = spark_compatible_murmur3_hash(item.to_le_bytes(), 0); + let h2 = spark_compatible_murmur3_hash(item.to_le_bytes(), h1); + for i in 1..=num_hash_functions { + let mut combined = (h1 as i32).wrapping_add((i as i32).wrapping_mul(h2 as i32)); + if combined < 0 { + combined = !combined; + } + let idx = (combined % bit_size) as usize; + words[idx >> 6] |= 1u64 << (idx & 0x3f); + } + } + + let mut bytes = 1i32.to_be_bytes().to_vec(); + bytes.extend_from_slice(&(num_hash_functions as i32).to_be_bytes()); + bytes.extend_from_slice(&(num_words as i32).to_be_bytes()); + for word in words { + bytes.extend_from_slice(&word.to_be_bytes()); + } + bytes +} + +/// A batch of join keys, one in ten of them null, half of them present in the filter. +fn probe_values(num_rows: usize) -> (ArrayRef, Vec) { + let mut values: Vec> = Vec::with_capacity(num_rows); + let mut inserted: Vec = Vec::with_capacity(num_rows / 2); + // Simple LCG so the benchmark data is deterministic without a dependency on an RNG. + let mut state: u64 = 0x2545_f491_4f6c_dd1d; + for row in 0..num_rows { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + if row % 10 == 9 { + values.push(None); + continue; + } + let value = state as i64; + if row % 2 == 0 { + inserted.push(value); + } + values.push(Some(value)); + } + (Arc::new(Int64Array::from(values)), inserted) +} + +fn criterion_benchmark(c: &mut Criterion) { + let num_rows = 8192; + let (values, inserted) = probe_values(num_rows); + let return_field = Arc::new(Field::new("v", DataType::Boolean, true)); + + let mut group = c.benchmark_group("bloom_filter_might_contain"); + // (label, num_hash_functions, num_words): a filter sized as Spark's runtime bloom filter + // join pushdown would size it, and a saturated one where every probe walks all hashes. + for (label, num_hash_functions, num_words) in + [("sparse", 5u32, 16384usize), ("saturated", 8, 64)] + { + let filter_bytes = serialize_filter(num_hash_functions, num_words, &inserted); + let filter: Arc = + Arc::new(Literal::new(ScalarValue::Binary(Some(filter_bytes)))); + let udf = BloomFilterMightContain::try_new(filter).unwrap(); + + group.bench_function(label, |b| { + b.iter(|| { + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::clone(&values))], + number_rows: num_rows, + return_field: Arc::clone(&return_field), + config_options: Arc::new(ConfigOptions::default()), + arg_fields: vec![], + }; + black_box(udf.invoke_with_args(args).unwrap()); + }) + }); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs b/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs index 0c5dbea6c8..50801ffd34 100644 --- a/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs +++ b/native/spark-expr/src/bloom_filter/spark_bloom_filter.rs @@ -15,14 +15,17 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrowNativeTypeOp, BooleanArray, Int64Array}; +use arrow::array::{Array, ArrowNativeTypeOp, BooleanArray, Int64Array}; +use arrow::buffer::BooleanBuffer; use arrow::datatypes::ToByteSlice; use datafusion::common::{DataFusionError, Result as DFResult}; use std::cmp; use crate::bloom_filter::spark_bit_array; use crate::bloom_filter::spark_bit_array::SparkBitArray; -use crate::hash_funcs::murmur3::spark_compatible_murmur3_hash; +use crate::hash_funcs::murmur3::{ + spark_compatible_murmur3_hash, spark_compatible_murmur3_hash_long, +}; const SPARK_BLOOM_FILTER_VERSION_1: i32 = 1; const SPARK_BLOOM_FILTER_VERSION_2: i32 = 2; @@ -71,6 +74,8 @@ pub struct SparkBloomFilter { /// Murmur3 seed. V1 always uses 0; V2 stores a per-filter seed (Spark's /// `BloomFilterImplV2.DEFAULT_SEED` is also 0, so 0 is the common case). seed: i32, + /// Reciprocal of the V1 bit-index modulus. Derived from `bits`, whose size never changes. + modulus: FastMod32, } pub fn optimal_num_hash_functions(expected_items: i32, num_bits: i32) -> i32 { @@ -116,8 +121,10 @@ impl From<&[u8]> for SparkBloomFilter { bits[i as usize] = read_num_be_bytes!(i64, 8, buf[offset..]) as u64; offset += 8; } + let bits = SparkBitArray::new(bits); Self { - bits: SparkBitArray::new(bits), + modulus: FastMod32::for_bit_size(bits.bit_size()), + bits, num_hash_functions: num_hash_functions as u32, version, seed, @@ -136,9 +143,10 @@ impl SparkBloomFilter { seed: i32, ) -> Self { let num_words = spark_bit_array::num_words(num_bits as usize); - let bits = vec![0u64; num_words]; + let bits = SparkBitArray::new(vec![0u64; num_words]); Self { - bits: SparkBitArray::new(bits), + modulus: FastMod32::for_bit_size(bits.bit_size()), + bits, num_hash_functions: num_hash_functions as u32, version, seed: match version { @@ -177,13 +185,8 @@ impl SparkBloomFilter { /// matching `BloomFilterImpl.scatterHashAndSetAllBits` (Spark <= 4.0; still /// available as the V1 codepath in 4.1+). fn scatter_v1(&mut self, h1: u32, h2: u32, set: bool) -> Option { - let bit_size = self.bits.bit_size() as i32; for i in 1..=self.num_hash_functions { - let mut combined_hash = (h1 as i32).add_wrapping((i as i32).mul_wrapping(h2 as i32)); - if combined_hash < 0 { - combined_hash = !combined_hash; - } - let idx = (combined_hash % bit_size) as usize; + let idx = v1_bit_index(h1, h2, i, &self.modulus); if set { self.bits.set(idx); } else if !self.bits.get(idx) { @@ -235,29 +238,40 @@ impl SparkBloomFilter { } } + /// The two murmur3 hashes Spark's `BloomFilter` derives every bit index from: the item + /// hashed with the filter seed, then hashed again with the first hash as seed. + #[inline] + fn hashes_long(&self, item: i64) -> (u32, u32) { + let h1 = spark_compatible_murmur3_hash_long(item, self.seed as u32); + (h1, spark_compatible_murmur3_hash_long(item, h1)) + } + + /// [`Self::hashes_long`] for binary items. + fn hashes_binary(&self, item: &[u8]) -> (u32, u32) { + let h1 = spark_compatible_murmur3_hash(item, self.seed as u32); + (h1, spark_compatible_murmur3_hash(item, h1)) + } + /// Put a long item into the filter. Returns `false`; the original Spark /// `BloomFilter.put` returns whether any bit changed, but no current Comet /// caller uses that, so we don't bother computing it. pub fn put_long(&mut self, item: i64) -> bool { - let h1 = spark_compatible_murmur3_hash(item.to_le_bytes(), self.seed as u32); - let h2 = spark_compatible_murmur3_hash(item.to_le_bytes(), h1); + let (h1, h2) = self.hashes_long(item); self.scatter(h1, h2, true); false } pub fn put_binary(&mut self, item: &[u8]) -> bool { - let h1 = spark_compatible_murmur3_hash(item, self.seed as u32); - let h2 = spark_compatible_murmur3_hash(item, h1); + let (h1, h2) = self.hashes_binary(item); self.scatter(h1, h2, true); false } pub fn might_contain_long(&self, item: i64) -> bool { - let h1 = spark_compatible_murmur3_hash(item.to_le_bytes(), self.seed as u32); - let h2 = spark_compatible_murmur3_hash(item.to_le_bytes(), h1); + let (h1, h2) = self.hashes_long(item); match self.version { SparkBloomFilterVersion::V1 => { - might_contain_long_v1(&self.bits, self.num_hash_functions, h1, h2) + might_contain_long_v1(&self.bits, self.num_hash_functions, h1, h2, &self.modulus) } SparkBloomFilterVersion::V2 => { might_contain_long_v2(&self.bits, self.num_hash_functions, h1, h2) @@ -266,10 +280,37 @@ impl SparkBloomFilter { } pub fn might_contain_longs(&self, items: &Int64Array) -> BooleanArray { - items - .iter() - .map(|v| v.map(|x| self.might_contain_long(x))) - .collect() + // The probe function is the same for every row, so the version is resolved once here + // rather than once per row. + let mut probed = match self.version { + SparkBloomFilterVersion::V1 => self.probe_each(items, |bits, num_hash, h1, h2| { + might_contain_long_v1(bits, num_hash, h1, h2, &self.modulus) + }), + SparkBloomFilterVersion::V2 => self.probe_each(items, might_contain_long_v2), + }; + + // A null row must not report a match, so mask its probe result off. + let nulls = items.nulls(); + if let Some(nulls) = nulls { + probed &= nulls.inner(); + } + BooleanArray::new(probed, nulls.cloned()) + } + + /// Probe `items` with `probe`, one bit per row. Null rows are probed anyway — their slot in + /// the values buffer holds an arbitrary but valid i64 — which keeps the loop branch-free + /// with respect to nullability. The caller masks those results off. + #[inline] + fn probe_each( + &self, + items: &Int64Array, + probe: impl Fn(&SparkBitArray, u32, u32, u32) -> bool, + ) -> BooleanBuffer { + let values = items.values(); + BooleanBuffer::collect_bool(values.len(), |i| { + let (h1, h2) = self.hashes_long(values[i]); + probe(&self.bits, self.num_hash_functions, h1, h2) + }) } /// Number of set bits in the underlying bit array. Mirrors Spark's @@ -341,18 +382,54 @@ impl SparkBloomFilter { } } -fn might_contain_long_v1(bits: &SparkBitArray, num_hash_functions: u32, h1: u32, h2: u32) -> bool { - let bit_size = bits.bit_size() as i32; - for i in 1..=num_hash_functions { - let mut combined_hash = (h1 as i32).add_wrapping((i as i32).mul_wrapping(h2 as i32)); - if combined_hash < 0 { - combined_hash = !combined_hash; - } - if !bits.get((combined_hash % bit_size) as usize) { - return false; +/// Remainder by a fixed 32-bit modulus, computed with multiplications instead of a hardware +/// division (Lemire's method). The modulus is fixed for the life of a filter, so the reciprocal +/// is computed once at construction rather than once per bit index. +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +struct FastMod32 { + divisor: u32, + reciprocal: u64, +} + +impl FastMod32 { + /// The modulus a V1 probe reduces by. Spark truncates the bit size to `int` before taking + /// the remainder, and the dividend is always non-negative, so only the magnitude of the + /// truncated divisor matters. + fn for_bit_size(bit_size: u64) -> Self { + let divisor = (bit_size as i32).unsigned_abs(); + Self { + divisor, + reciprocal: (u64::MAX / divisor as u64).wrapping_add(1), } } - true + + #[inline] + fn rem(&self, n: u32) -> u32 { + let lowbits = self.reciprocal.wrapping_mul(n as u64); + ((lowbits as u128 * self.divisor as u128) >> 64) as u32 + } +} + +/// The `i`th V1 bit index for a hash pair: `combinedHash = h1 + i*h2`, folded to a non-negative +/// value and reduced modulo the bit size. Shared by the read and write paths so they cannot +/// drift apart. +#[inline] +fn v1_bit_index(h1: u32, h2: u32, i: u32, modulus: &FastMod32) -> usize { + let mut combined_hash = (h1 as i32).add_wrapping((i as i32).mul_wrapping(h2 as i32)); + if combined_hash < 0 { + combined_hash = !combined_hash; + } + modulus.rem(combined_hash as u32) as usize +} + +fn might_contain_long_v1( + bits: &SparkBitArray, + num_hash_functions: u32, + h1: u32, + h2: u32, + modulus: &FastMod32, +) -> bool { + (1..=num_hash_functions).all(|i| bits.get(v1_bit_index(h1, h2, i, modulus))) } fn might_contain_long_v2(bits: &SparkBitArray, num_hash_functions: u32, h1: u32, h2: u32) -> bool { diff --git a/native/spark-expr/src/hash_funcs/murmur3.rs b/native/spark-expr/src/hash_funcs/murmur3.rs index dc0f804ab2..8655e6e36a 100644 --- a/native/spark-expr/src/hash_funcs/murmur3.rs +++ b/native/spark-expr/src/hash_funcs/murmur3.rs @@ -68,36 +68,46 @@ pub fn spark_murmur3_hash(args: &[ColumnarValue]) -> Result>(data: T, seed: u32) -> u32 { - #[inline] - fn mix_k1(mut k1: i32) -> i32 { - k1 = k1.mul_wrapping(0xcc9e2d51u32 as i32); - k1 = k1.rotate_left(15); - k1 = k1.mul_wrapping(0x1b873593u32 as i32); - k1 - } +fn mix_k1(mut k1: i32) -> i32 { + k1 = k1.mul_wrapping(0xcc9e2d51u32 as i32); + k1 = k1.rotate_left(15); + k1 = k1.mul_wrapping(0x1b873593u32 as i32); + k1 +} - #[inline] - fn mix_h1(mut h1: i32, k1: i32) -> i32 { - h1 ^= k1; - h1 = h1.rotate_left(13); - h1 = h1.mul_wrapping(5).add_wrapping(0xe6546b64u32 as i32); - h1 - } +#[inline] +fn mix_h1(mut h1: i32, k1: i32) -> i32 { + h1 ^= k1; + h1 = h1.rotate_left(13); + h1 = h1.mul_wrapping(5).add_wrapping(0xe6546b64u32 as i32); + h1 +} - #[inline] - fn fmix(mut h1: i32, len: i32) -> i32 { - h1 ^= len; - h1 ^= (h1 as u32 >> 16) as i32; - h1 = h1.mul_wrapping(0x85ebca6bu32 as i32); - h1 ^= (h1 as u32 >> 13) as i32; - h1 = h1.mul_wrapping(0xc2b2ae35u32 as i32); - h1 ^= (h1 as u32 >> 16) as i32; - h1 - } +#[inline] +fn fmix(mut h1: i32, len: i32) -> i32 { + h1 ^= len; + h1 ^= (h1 as u32 >> 16) as i32; + h1 = h1.mul_wrapping(0x85ebca6bu32 as i32); + h1 ^= (h1 as u32 >> 13) as i32; + h1 = h1.mul_wrapping(0xc2b2ae35u32 as i32); + h1 ^= (h1 as u32 >> 16) as i32; + h1 +} +/// Spark-compatible murmur3 hash of an `i64`, equivalent to +/// `spark_compatible_murmur3_hash(item.to_le_bytes(), seed)` but with the two 4-byte +/// words mixed directly, without going through a byte slice. +#[inline] +pub fn spark_compatible_murmur3_hash_long(item: i64, seed: u32) -> u32 { + let mut h1 = mix_h1(seed as i32, mix_k1(item as i32)); + h1 = mix_h1(h1, mix_k1((item >> 32) as i32)); + fmix(h1, 8) as u32 +} + +/// Spark-compatible murmur3 hash function +#[inline] +pub fn spark_compatible_murmur3_hash>(data: T, seed: u32) -> u32 { #[inline] unsafe fn hash_bytes_by_int(data: &[u8], seed: u32) -> i32 { // safety: data length must be aligned to 4 bytes @@ -194,10 +204,34 @@ mod tests { use arrow::array::{Float32Array, Float64Array}; use std::sync::Arc; - use crate::murmur3::create_murmur3_hashes; + use crate::murmur3::{ + create_murmur3_hashes, spark_compatible_murmur3_hash, spark_compatible_murmur3_hash_long, + }; use crate::test_hashes_with_nulls; use datafusion::arrow::array::{ArrayRef, Int32Array, Int64Array, Int8Array, StringArray}; + #[test] + fn test_hash_long_matches_byte_slice_hash() { + for seed in [0u32, 42, 0xffff_ffff, 12345] { + for item in [ + 0i64, + 1, + -1, + 42, + i64::MIN, + i64::MAX, + 0x0123_4567_89ab_cdef, + -0x0123_4567_89ab_cdef, + ] { + assert_eq!( + spark_compatible_murmur3_hash_long(item, seed), + spark_compatible_murmur3_hash(item.to_le_bytes(), seed), + "item={item}, seed={seed}" + ); + } + } + } + fn test_murmur3_hash>> + 'static>( values: Vec>, expected: Vec,