From 3bcc8530e75315989ba86d7a070975a190aea737 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 23 Jul 2026 14:57:46 -0400 Subject: [PATCH 1/7] Add opt-in Bloom skipping index Signed-off-by: Connor Tsui --- vortex-file/src/strategy.rs | 72 ++-- vortex-file/tests/bloom_skip_index.rs | 247 ++++++++++++++ .../layouts/zoned/aggregates/bloom_filter.rs | 308 ++++++++++++++++++ .../src/layouts/zoned/aggregates/min_max.rs | 70 ++++ .../src/layouts/zoned/aggregates/mod.rs | 96 ++++++ vortex-layout/src/layouts/zoned/mod.rs | 2 + .../src/layouts/zoned/skip_index/bloom.rs | 269 +++++++++++++++ .../src/layouts/zoned/skip_index/mod.rs | 59 ++++ vortex-layout/src/layouts/zoned/writer.rs | 3 +- 9 files changed, 1099 insertions(+), 27 deletions(-) create mode 100644 vortex-file/tests/bloom_skip_index.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/min_max.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/mod.rs create mode 100644 vortex-layout/src/layouts/zoned/skip_index/bloom.rs create mode 100644 vortex-layout/src/layouts/zoned/skip_index/mod.rs diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 9d4dbb90610..942ad91bc3a 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -59,6 +59,7 @@ pub struct WriteStrategyBuilder { row_block_size: usize, data_block_target_bytes: Option, field_writers: HashMap>, + field_zoned_options: HashMap, allow_encodings: Option>, flat_strategy: Option>, probe_compressor: Option>, @@ -78,6 +79,7 @@ impl Default for WriteStrategyBuilder { data_block_target_bytes: Some(ONE_MEG), field_writers: HashMap::new(), allow_encodings: None, + field_zoned_options: HashMap::new(), flat_strategy: None, probe_compressor: None, use_list_layout: use_experimental_list_layout(), @@ -128,7 +130,21 @@ impl WriteStrategyBuilder { self } - /// Override the allowed array encodings for file writing. + /// Override only the zoned-statistics options for a field while retaining the default + /// repartitioning, dictionary, compression, buffering, and flat-layout pipeline. + /// + /// This can attach custom per-zone aggregates without changing the physical data strategy for + /// the field. + pub fn with_field_zoned_options( + mut self, + field: impl Into, + options: ZonedLayoutOptions, + ) -> Self { + self.field_zoned_options.insert(field.into(), options); + self + } + + /// Override the allowed array encodings for normalization. /// /// The configured flat leaf strategy is wrapped in a [`LayoutStrategyEncodingValidator`] /// that recursively checks every chunk before passing it to the leaf writer. [`build`](Self::build) @@ -175,7 +191,7 @@ impl WriteStrategyBuilder { /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides /// applied. - pub fn build(self) -> Arc { + pub fn build(mut self) -> Arc { let flat: Arc = if let Some(flat) = self.flat_strategy { flat } else { @@ -261,36 +277,40 @@ impl WriteStrategyBuilder { let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"); - // 2. calculate stats for each row group - let stats = ZonedStrategy::new( - dict, - compress_then_flat.clone(), - ZonedLayoutOptions { - block_size: row_block_size, - ..Default::default() - }, - ); + let column_writer = |options: ZonedLayoutOptions| -> Arc { + // 2. calculate stats for each row group + let stats = + ZonedStrategy::new(dict.clone(), compress_then_flat.clone(), options.clone()); - // 1. repartition each column to fixed row counts - let repartition = RepartitionStrategy::new( - stats, - RepartitionWriterOptions { - // No minimum block size in bytes - block_size_minimum: 0, - // Always repartition into 8K row blocks - block_len_multiple: self.row_block_size, - block_size_target: None, - canonicalize: false, - }, - ); + // 1. repartition each column to fixed row counts + Arc::new(RepartitionStrategy::new( + stats, + RepartitionWriterOptions { + // No minimum block size in bytes + block_size_minimum: 0, + block_len_multiple: options.block_size.get(), + block_size_target: None, + canonicalize: false, + }, + )) + }; + let repartition = column_writer(ZonedLayoutOptions { + block_size: row_block_size, + ..Default::default() + }); + + for (field, options) in self.field_zoned_options { + self.field_writers + .entry(field) + .or_insert_with(|| column_writer(options)); + } // 0. start with splitting columns let validity_strategy = CollectStrategy::new(compress_then_flat.clone()); // Take any field overrides from the builder and apply them to the final strategy. - let mut table_strategy = - TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) - .with_field_writers(self.field_writers); + let mut table_strategy = TableStrategy::new(Arc::new(validity_strategy), repartition) + .with_field_writers(self.field_writers); if self.use_list_layout { // We need a closure here to enable recursive application of list layout. diff --git a/vortex-file/tests/bloom_skip_index.rs b/vortex-file/tests/bloom_skip_index.rs new file mode 100644 index 00000000000..6e1b6c51986 --- /dev/null +++ b/vortex-file/tests/bloom_skip_index.rs @@ -0,0 +1,247 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! End-to-end coverage for the zoned Bloom skipping index. +//! +//! Bloom indexes are optional extensions rather than part of the default file layout. This test +//! exercises the complete opt-in lifecycle: +//! +//! 1. register the index with a write session and request it for one field; +//! 2. persist one Bloom filter per zone and reopen the file with a fresh registered session; +//! 3. prove that equality predicates prune zones while returning the same rows as a full scan; and +//! 4. reopen the indexed file with an unregistered, allow-unknown session to verify that the index +//! is ignorable. +//! +//! The input is intentionally hostile to ordinary min/max pruning. Zone `z` contains values whose +//! remainder modulo [`NZONES`] is `z`, so both [`HIT`] and [`MISS`] lie inside every zone's +//! min/max range. `MISS` is then removed from its zone without changing that range. Consequently, +//! pruning either value requires the Bloom filter rather than the built-in range statistics. + +#![expect(clippy::expect_used)] + +use std::num::NonZeroU8; +use std::num::NonZeroUsize; +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::eq; +use vortex_array::expr::get_item; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::field_path; +use vortex_array::stream::ArrayStreamExt; +use vortex_error::VortexResult; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_file::WriteStrategyBuilder; +use vortex_io::session::RuntimeSession; +use vortex_layout::LayoutStrategy; +use vortex_layout::layouts::zoned::skip_index::SkipIndex; +use vortex_layout::layouts::zoned::skip_index::bloom::BloomOptions; +use vortex_layout::layouts::zoned::skip_index::bloom::BloomSkipIndex; +use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; +use vortex_layout::session::LayoutSession; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +const ZONE_LEN: usize = 256; +const NZONES: usize = 4; +const HIT: i64 = 502; +const MISS: i64 = 503; + +fn bloom() -> BloomSkipIndex { + // A deliberately roomy filter keeps this correctness test's false-positive outcome + // deterministic. False positives are valid Bloom behavior, but false negatives are not. + BloomSkipIndex::new(BloomOptions::new( + NonZeroUsize::new(1024).expect("1024 is non-zero"), + NonZeroU8::new(5).expect("5 is non-zero"), + )) +} + +fn session(index: Option<&dyn SkipIndex>) -> VortexSession { + let session = vortex_array::array_session() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + + // Registration installs the persisted aggregate, membership probe, and equality rewrite. + // Callers must do this independently for the sessions that write and read an indexed file. + if let Some(index) = index { + index.register(&session); + } + session +} + +fn data() -> ArrayRef { + data_with_shape(ZONE_LEN, NZONES, Some(MISS)) +} + +fn data_with_shape(zone_len: usize, nzones: usize, missing: Option) -> ArrayRef { + let chunks = (0..nzones) + .map(|zone| { + let mut values = (0..zone_len) + .map(|row| i64::try_from(row * nzones + zone).expect("test value fits i64")) + .collect::>(); + if let Some(missing) = missing + && usize::try_from(missing).expect("missing value is non-negative") % nzones == zone + { + // Leave a hole inside every zone's min/max range so a MISS cannot be pruned by the + // ordinary range stats. The bloom must provide the proof. + values[usize::try_from(missing).expect("missing value is non-negative") / nzones] = + i64::try_from(zone_len * nzones + zone).expect("replacement fits i64"); + } + StructArray::from_fields(&[("id", PrimitiveArray::from_iter(values).into_array())]) + .expect("valid test struct") + .into_array() + }) + .collect::>(); + ChunkedArray::try_new( + chunks, + DType::struct_( + [("id", DType::Primitive(PType::I64, Nullability::NonNullable))], + Nullability::NonNullable, + ), + ) + .expect("valid chunked test data") + .into_array() +} + +fn filter(value: i64) -> Expression { + eq(get_item("id", root()), lit(value)) +} + +fn strategy( + session: &VortexSession, + index: Option<&dyn SkipIndex>, + zone_len: usize, +) -> VortexResult> { + let mut options = ZonedLayoutOptions { + block_size: NonZeroUsize::new(zone_len).expect("zone length is non-zero"), + ..Default::default() + }; + if let Some(index) = index { + // Adding the aggregate to these field-specific zoned options is the explicit write-side + // opt-in. Registering the index in the session alone does not change the file layout. + options = options.with_skip_index(index, &PType::I64.into(), session)?; + } + Ok(WriteStrategyBuilder::default() + .with_field_zoned_options(field_path!(id), options) + .build()) +} + +async fn scan(file: &vortex_file::VortexFile, value: i64) -> VortexResult { + file.scan()? + .with_filter(filter(value)) + .into_array_stream()? + .read_all() + .await +} + +async fn write_file( + session: &VortexSession, + input: &ArrayRef, + index: Option<&dyn SkipIndex>, + zone_len: usize, +) -> VortexResult> { + let mut bytes = Vec::new(); + session + .write_options() + .with_strategy(strategy(session, index, zone_len)?) + .write(&mut bytes, input.to_array_stream()) + .await?; + Ok(bytes) +} + +#[expect(clippy::tests_outside_test_module)] +#[tokio::test] +async fn bloom_roundtrip_prunes_and_unknown_reader_matches_full_scan() -> VortexResult<()> { + let index = bloom(); + let write_session = session(Some(&index)); + let input = data(); + let bytes = write_file(&write_session, &input, Some(&index), ZONE_LEN).await?; + + // Reconstruct every read-side extension from a fresh session rather than accidentally relying + // on state retained by the writer. + let read_session = session(Some(&index)); + let file = read_session.open_options().open_buffer(bytes.clone())?; + let reader = file.layout_reader()?; + let row_count = file.row_count(); + + // HIT is present only in zone 2. Since it falls within every zone's min/max range, the exact + // one-zone mask proves that the Bloom falsifier participated in pruning. + let hit_mask = reader + .pruning_evaluation( + &(0..row_count), + &filter(HIT), + Mask::new_true(usize::try_from(row_count)?), + )? + .await?; + assert_eq!(hit_mask.true_count(), ZONE_LEN); + assert!(hit_mask.iter().take(2 * ZONE_LEN).all(|keep| !keep)); + assert!( + hit_mask + .iter() + .skip(2 * ZONE_LEN) + .take(ZONE_LEN) + .all(|keep| keep) + ); + assert!(hit_mask.iter().skip(3 * ZONE_LEN).all(|keep| !keep)); + + // MISS was removed while remaining inside every zone's min/max range. Only the Bloom filters + // can prove that all four zones are absent. + let miss_mask = reader + .pruning_evaluation( + &(0..row_count), + &filter(MISS), + Mask::new_true(usize::try_from(row_count)?), + )? + .await?; + assert!( + miss_mask.all_false(), + "an absent value should prune every zone" + ); + + // An allow-unknown reader without Bloom registration bypasses the unavailable zone map and + // scans the data child. This both supplies the reference result and verifies that an optional + // index does not become a hard read-time dependency. + let full_scan_session = session(None); + full_scan_session.allow_unknown(); + let full_scan_file = full_scan_session.open_options().open_buffer(bytes)?; + + let indexed_hit = scan(&file, HIT).await?; + let full_scan_hit = scan(&full_scan_file, HIT).await?; + // A Bloom filter may retain extra zones, but it must never change query results. + assert_arrays_eq!( + indexed_hit, + full_scan_hit, + &mut read_session.create_execution_ctx() + ); + let expected_hit = + StructArray::from_fields(&[("id", PrimitiveArray::from_iter([HIT]).into_array())])? + .into_array(); + assert_arrays_eq!( + full_scan_hit, + expected_hit, + &mut read_session.create_execution_ctx() + ); + + let indexed_miss = scan(&file, MISS).await?; + let full_scan_miss = scan(&full_scan_file, MISS).await?; + assert_arrays_eq!( + indexed_miss, + full_scan_miss, + &mut read_session.create_execution_ctx() + ); + assert_eq!(full_scan_miss.len(), 0); + Ok(()) +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter.rs new file mode 100644 index 00000000000..d1d3707beab --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter.rs @@ -0,0 +1,308 @@ +//! Bloom-filter aggregate for zoned layouts. + +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::num::NonZeroU8; +use std::num::NonZeroUsize; + +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +/// Bloom-filter tuning persisted as aggregate metadata. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct BloomOptions { + bytes: NonZeroUsize, + hashes: NonZeroU8, +} + +impl BloomOptions { + /// Create bloom options with a fixed number of bytes and hash probes per zone. + pub fn new(bytes: NonZeroUsize, hashes: NonZeroU8) -> Self { + Self { bytes, hashes } + } + + /// Bytes stored for each zone. + pub fn bytes(&self) -> NonZeroUsize { + self.bytes + } + + /// Hash probes performed for each inserted or tested value. + pub fn hashes(&self) -> NonZeroU8 { + self.hashes + } +} + +impl Default for BloomOptions { + fn default() -> Self { + Self { + // Eight bits per row at the default 8192-row zone size. + bytes: NonZeroUsize::new(8192).unwrap_or(NonZeroUsize::MIN), + hashes: NonZeroU8::new(5).unwrap_or(NonZeroU8::MIN), + } + } +} + +impl Display for BloomOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "bytes={},hashes={}", self.bytes, self.hashes) + } +} + +/// Aggregate that stores one fixed-size Bloom bitset as a `Binary` scalar for every zone. +#[derive(Clone, Debug)] +pub(in crate::layouts::zoned) struct BloomFilter; + +/// In-memory Bloom accumulator. Only the bitset is persisted. +pub(in crate::layouts::zoned) struct BloomPartial { + bits: Vec, + hashes: u8, +} + +impl AggregateFnVTable for BloomFilter { + type Options = BloomOptions; + type Partial = BloomPartial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.bloom_filter.i64.v1"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + let bytes = u32::try_from(options.bytes.get())?; + let mut metadata = bytes.to_le_bytes().to_vec(); + metadata.push(options.hashes.get()); + Ok(Some(metadata)) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure!(metadata.len() == 5, "invalid bloom metadata length"); + let bytes = u32::from_le_bytes([metadata[0], metadata[1], metadata[2], metadata[3]]); + Ok(BloomOptions::new( + NonZeroUsize::new(bytes as usize) + .ok_or_else(|| vortex_err!("bloom byte length must be non-zero"))?, + NonZeroU8::new(metadata[4]) + .ok_or_else(|| vortex_err!("bloom hash count must be non-zero"))?, + )) + } + + fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + matches!(input_dtype, DType::Primitive(PType::I64, _)) + .then_some(DType::Binary(Nullability::NonNullable)) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + } + + fn empty_partial( + &self, + options: &Self::Options, + _input_dtype: &DType, + ) -> VortexResult { + Ok(BloomPartial { + bits: vec![0; options.bytes.get()], + hashes: options.hashes.get(), + }) + } + + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + if other.is_null() { + return Ok(()); + } + let other = other + .as_binary() + .value() + .ok_or_else(|| vortex_err!("non-null bloom partial has no bytes"))?; + vortex_ensure!( + partial.bits.len() == other.len(), + "bloom partial length mismatch" + ); + for (dst, src) in partial.bits.iter_mut().zip(other.as_slice()) { + *dst |= *src; + } + Ok(()) + } + + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + Ok(Scalar::binary( + partial.bits.clone(), + Nullability::NonNullable, + )) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.bits.fill(0); + } + + fn is_saturated(&self, partial: &Self::Partial) -> bool { + partial.bits.iter().all(|byte| *byte == u8::MAX) + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match batch { + Columnar::Constant(constant) => { + if let Some(value) = i64_value(constant.scalar())? { + bloom_insert(&mut partial.bits, value, partial.hashes); + } + } + Columnar::Canonical(canonical) => { + let primitive = canonical.as_primitive(); + let values = primitive.as_slice::(); + let validity = primitive.validity()?.execute_mask(values.len(), ctx)?; + for (&value, valid) in values.iter().zip(validity.iter()) { + if valid { + bloom_insert(&mut partial.bits, value, partial.hashes); + } + } + } + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef) -> VortexResult { + Ok(partials) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + self.to_scalar(partial) + } +} + +pub(in crate::layouts::zoned) fn i64_value(scalar: &Scalar) -> VortexResult> { + if scalar.is_null() { + return Ok(None); + } + scalar + .as_primitive_opt() + .and_then(|primitive| primitive.typed_value::()) + .map(Some) + .ok_or_else(|| vortex_err!("bloom value must be i64")) +} + +fn bloom_insert(bits: &mut [u8], value: i64, hashes: u8) { + bloom_insert_hash( + bits, + splitmix64(value as u64 ^ 0x243f_6a88_85a3_08d3), + hashes, + ); +} + +fn bloom_insert_hash(bits: &mut [u8], hash: u64, hashes: u8) { + for (byte, bit) in bloom_positions(hash, bits.len(), hashes) { + bits[byte] |= 1 << bit; + } +} + +pub(in crate::layouts::zoned) fn bloom_contains(bits: &[u8], value: i64, hashes: u8) -> bool { + bloom_contains_hash( + bits, + splitmix64(value as u64 ^ 0x243f_6a88_85a3_08d3), + hashes, + ) +} + +fn bloom_contains_hash(bits: &[u8], hash: u64, hashes: u8) -> bool { + bloom_positions(hash, bits.len(), hashes).all(|(byte, bit)| bits[byte] & (1 << bit) != 0) +} + +fn bloom_positions(hash: u64, bytes: usize, hashes: u8) -> impl Iterator { + let h1 = hash; + let h2 = splitmix64(h1 ^ 0x1319_8a2e_0370_7344) | 1; + let bit_len = u64::try_from(bytes).unwrap_or(u64::MAX).saturating_mul(8); + (0..u64::from(hashes)).map(move |probe| { + let position = h1 + .wrapping_add(probe.wrapping_mul(h2)) + .wrapping_rem(bit_len); + // `position / 8` is less than `bytes`, which is already a `usize`. + let byte = usize::try_from(position / 8).unwrap_or_default(); + let bit = u32::try_from(position % 8).unwrap_or_default(); + (byte, bit) + }) +} + +fn splitmix64(mut value: u64) -> u64 { + value = value.wrapping_add(0x9e37_79b9_7f4a_7c15); + value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + value ^ (value >> 31) +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU8; + use std::num::NonZeroUsize; + + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_error::VortexResult; + + use super::BloomFilter; + use super::BloomOptions; + use super::bloom_contains; + + fn small_options() -> BloomOptions { + BloomOptions::new( + NonZeroUsize::new(64).expect("64 is non-zero"), + NonZeroU8::new(3).expect("3 is non-zero"), + ) + } + + #[test] + fn roundtrips_options_and_membership() -> VortexResult<()> { + let session = vortex_array::array_session(); + let options = small_options(); + let metadata = BloomFilter + .serialize(&options)? + .expect("bloom is serializable"); + assert_eq!(BloomFilter.deserialize(&metadata, &session)?, options); + + let mut ctx = session.create_execution_ctx(); + let mut accumulator = Accumulator::try_new( + BloomFilter, + options.clone(), + DType::Primitive(PType::I64, Nullability::NonNullable), + )?; + accumulator.accumulate( + &PrimitiveArray::from_iter([10i64, 20, 30]).into_array(), + &mut ctx, + )?; + let state = accumulator.finish()?; + let bytes = state.as_binary().value().expect("bloom state is non-null"); + assert!(bloom_contains(bytes.as_slice(), 10, options.hashes.get())); + assert!(bloom_contains(bytes.as_slice(), 20, options.hashes.get())); + assert!(!bloom_contains(bytes.as_slice(), 999, options.hashes.get())); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/min_max.rs b/vortex-layout/src/layouts/zoned/aggregates/min_max.rs new file mode 100644 index 00000000000..0403d5fb404 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/min_max.rs @@ -0,0 +1,70 @@ +//! Min/max aggregate selection for zoned layouts. + +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; +use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions; +use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; +use vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions; +use vortex_array::aggregate_fn::fns::max::Max; +use vortex_array::aggregate_fn::fns::min::Min; +use vortex_array::dtype::DType; + +use super::super::schema::default_bounded_stat_max_bytes; + +pub(super) fn min_max_aggregate_fns(dtype: &DType) -> [AggregateFnRef; 2] { + match dtype { + DType::Utf8(_) | DType::Binary(_) => [ + BoundedMax.bind(BoundedMaxOptions { + max_bytes: default_bounded_stat_max_bytes(), + }), + BoundedMin.bind(BoundedMinOptions { + max_bytes: default_bounded_stat_max_bytes(), + }), + ], + _ => [ + Max.bind(NumericalAggregateOpts::skip_nans()), + Min.bind(NumericalAggregateOpts::skip_nans()), + ], + } +} + +#[cfg(test)] +mod tests { + use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; + use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; + use vortex_array::aggregate_fn::fns::max::Max; + use vortex_array::aggregate_fn::fns::min::Min; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + + use super::default_bounded_stat_max_bytes; + use super::min_max_aggregate_fns; + + #[test] + fn variable_length_min_max_are_bounded() { + let aggregate_fns = min_max_aggregate_fns(&DType::Utf8(Nullability::NonNullable)); + + assert_eq!( + aggregate_fns[0].as_::().max_bytes, + default_bounded_stat_max_bytes() + ); + assert_eq!( + aggregate_fns[1].as_::().max_bytes, + default_bounded_stat_max_bytes() + ); + } + + #[test] + fn fixed_width_min_max_are_exact() { + let aggregate_fns = min_max_aggregate_fns(&PType::I32.into()); + + assert!(aggregate_fns[0].is::()); + assert!(aggregate_fns[1].is::()); + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/mod.rs new file mode 100644 index 00000000000..b90d9cd6773 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/mod.rs @@ -0,0 +1,96 @@ +//! Aggregate functions selected by the zoned layout. + +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::nan_count::NanCount; +use vortex_array::aggregate_fn::fns::null_count::NullCount; +use vortex_array::aggregate_fn::fns::sum::Sum; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::dtype::DType; +use vortex_session::VortexSession; + +pub(in crate::layouts::zoned) mod bloom_filter; +mod min_max; + +pub(in crate::layouts::zoned) use bloom_filter::BloomFilter; +pub(in crate::layouts::zoned) use bloom_filter::bloom_contains; +pub(in crate::layouts::zoned) use bloom_filter::i64_value; +use min_max::min_max_aggregate_fns; + +pub(super) fn default_zoned_aggregate_fns( + dtype: &DType, + session: &VortexSession, +) -> Arc<[AggregateFnRef]> { + let mut aggregate_fns = Vec::from(min_max_aggregate_fns(dtype)); + if Sum + .return_dtype(&NumericalAggregateOpts::skip_nans(), dtype) + .is_some() + { + aggregate_fns.push(Sum.bind(NumericalAggregateOpts::skip_nans())); + } + aggregate_fns.push(NanCount.bind(EmptyOptions)); + aggregate_fns.push(NullCount.bind(EmptyOptions)); + + // Stats from geo extension types are discovered from the registry at runtime instead. + aggregate_fns.extend(session.aggregate_fns().zone_stat_defaults(dtype)); + + aggregate_fns.into() +} + +#[cfg(test)] +mod tests { + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::fns::sum::Sum; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::extension::datetime::Timestamp; + + use super::BloomFilter; + use super::bloom_filter::BloomOptions; + use super::default_zoned_aggregate_fns; + + #[test] + fn default_aggregates_exclude_bloom_filter() { + let aggregate_fns = + default_zoned_aggregate_fns(&PType::I64.into(), &vortex_array::array_session()); + let bloom = BloomFilter.bind(BloomOptions::default()); + + assert!( + aggregate_fns + .iter() + .all(|aggregate_fn| aggregate_fn != &bloom) + ); + } + + #[test] + fn default_aggregates_include_sum_for_numeric_dtype() { + let aggregate_fns = + default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session()); + + assert!(aggregate_fns[2].is::()); + } + + #[test] + fn default_aggregates_skip_sum_for_non_summable_dtype() { + let dtype = DType::Extension( + Timestamp::new(TimeUnit::Microseconds, Nullability::Nullable).erased(), + ); + let aggregate_fns = default_zoned_aggregate_fns(&dtype, &vortex_array::array_session()); + + assert!( + aggregate_fns + .iter() + .all(|aggregate_fn| !aggregate_fn.is::()) + ); + } +} diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index bc3d7d0626e..4948570ea36 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -11,10 +11,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod aggregates; mod builder; mod pruning; mod reader; mod schema; +pub mod skip_index; pub mod writer; pub mod zone_map; diff --git a/vortex-layout/src/layouts/zoned/skip_index/bloom.rs b/vortex-layout/src/layouts/zoned/skip_index/bloom.rs new file mode 100644 index 00000000000..7c6985126f3 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/skip_index/bloom.rs @@ -0,0 +1,269 @@ +//! Bloom skipping index for equality predicates. + +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::varbinview::VarBinViewArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::is_root; +use vortex_array::expr::not; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::session::ScalarFnSessionExt; +use vortex_array::stats::rewrite::StatsRewriteCtx; +use vortex_array::stats::rewrite::StatsRewriteRule; +use vortex_array::stats::session::StatsSessionExt; +use vortex_array::stats::stat; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::super::aggregates::BloomFilter; +use super::super::aggregates::bloom_contains; +pub use super::super::aggregates::bloom_filter::BloomOptions; +use super::super::aggregates::i64_value; +use super::SkipIndex; + +/// Bloom skipping index for `i64` equality predicates. +#[derive(Clone, Debug, Default)] +pub struct BloomSkipIndex { + options: BloomOptions, +} + +impl BloomSkipIndex { + /// Create an index with explicit Bloom tuning. + pub fn new(options: BloomOptions) -> Self { + Self { options } + } + + /// The persisted Bloom options. + pub fn options(&self) -> &BloomOptions { + &self.options + } +} + +impl SkipIndex for BloomSkipIndex { + fn aggregate_fn(&self, input_dtype: &DType) -> Option { + BloomFilter + .return_dtype(&self.options, input_dtype) + .map(|_| BloomFilter.bind(self.options.clone())) + } + + fn register(&self, session: &VortexSession) { + session.aggregate_fns().register(BloomFilter); + session.scalar_fns().register(BloomContains); + session.stats().register_rewrite(BloomEqRewrite { + options: self.options.clone(), + }); + } +} + +/// Probe scalar function: test one `i64` literal against each binary Bloom state. +#[derive(Clone, Debug)] +struct BloomContains; + +impl ScalarFnVTable for BloomContains { + type Options = BloomOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.bloom_contains.i64.v1"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + BloomFilter.serialize(options) + } + + fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { + BloomFilter.deserialize(metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(2) + } + + fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("filter"), + 1 => ChildName::from("needle"), + _ => unreachable!("bloom_contains has exactly two children"), + } + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + vortex_ensure!( + matches!(args[0], DType::Binary(_)), + "bloom filter must be Binary" + ); + vortex_ensure!( + matches!(args[1], DType::Primitive(PType::I64, _)), + "bloom needle must be i64" + ); + Ok(DType::Bool(args[0].nullability() | args[1].nullability())) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filters = args.get(0)?.execute::(ctx)?; + let needle_array = args.get(1)?; + let needle = needle_array + .as_constant() + .ok_or_else(|| vortex_err!("bloom needle must be constant"))?; + let Some(needle) = i64_value(&needle)? else { + return Ok(ConstantArray::new( + Scalar::null(DType::Bool(Nullability::Nullable)), + args.row_count(), + ) + .into_array()); + }; + + let validity = filters.varbinview_validity(); + let valid = validity.execute_mask(filters.len(), ctx)?; + let mut possible = Vec::with_capacity(filters.len()); + for (idx, is_valid) in valid.iter().enumerate() { + if is_valid { + let filter = filters.bytes_at(idx); + vortex_ensure!( + filter.len() == options.bytes().get(), + "stored bloom byte length does not match options" + ); + possible.push(bloom_contains( + filter.as_slice(), + needle, + options.hashes().get(), + )); + } else { + possible.push(false); + } + } + Ok(BoolArray::new(BitBuffer::from_iter(possible), validity).into_array()) + } + + fn is_null_sensitive(&self, _options: &Self::Options) -> bool { + false + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +/// Equality rewrite that turns a Bloom miss into a zone falsifier. +#[derive(Clone, Debug)] +struct BloomEqRewrite { + options: BloomOptions, +} + +impl StatsRewriteRule for BloomEqRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + Binary.id() + } + + fn falsify( + &self, + expr: &Expression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + if *expr.as_::() != Operator::Eq { + return Ok(None); + } + + let (column, literal) = if is_root(expr.child(0)) && expr.child(1).is::() { + (expr.child(0), expr.child(1)) + } else if is_root(expr.child(1)) && expr.child(0).is::() { + (expr.child(1), expr.child(0)) + } else { + return Ok(None); + }; + if !matches!(ctx.return_dtype(column)?, DType::Primitive(PType::I64, _)) + || literal.as_::().is_null() + { + return Ok(None); + } + + let filter = stat(column.clone(), BloomFilter.bind(self.options.clone())); + let contains = BloomContains.new_expr(self.options.clone(), [filter, literal.clone()]); + Ok(Some(not(contains))) + } +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU8; + use std::num::NonZeroUsize; + use std::sync::Arc; + + use vortex_array::arrays::StructArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::eq; + use vortex_array::expr::lit; + use vortex_array::expr::root; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + + use super::BloomOptions; + use super::BloomSkipIndex; + use super::SkipIndex; + use crate::layouts::zoned::zone_map::ZoneMap; + + fn small_options() -> BloomOptions { + BloomOptions::new( + NonZeroUsize::new(64).expect("64 is non-zero"), + NonZeroU8::new(3).expect("3 is non-zero"), + ) + } + + #[test] + fn missing_stat_stays_inconclusive() -> VortexResult<()> { + let session = vortex_array::array_session(); + let index = BloomSkipIndex::new(small_options()); + index.register(&session); + let predicate = eq(root(), lit(42i64)); + let proof = predicate + .falsify( + &DType::Primitive(PType::I64, Nullability::NonNullable), + &session, + )? + .expect("equality has a bloom proof"); + + let zone_map = ZoneMap::try_new( + DType::Primitive(PType::I64, Nullability::NonNullable), + StructArray::try_new(Vec::<&str>::new().into(), vec![], 2, Validity::NonNullable)?, + Arc::new([]), + 8, + 16, + )?; + assert!(zone_map.prune(&proof, &session)?.all_false()); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/skip_index/mod.rs b/vortex-layout/src/layouts/zoned/skip_index/mod.rs new file mode 100644 index 00000000000..2def62acab0 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/skip_index/mod.rs @@ -0,0 +1,59 @@ +//! Skipping-index interface and implementations. + +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Debug; +use std::sync::Arc; + +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::dtype::DType; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use super::aggregates::default_zoned_aggregate_fns; +use super::writer::ZonedLayoutOptions; + +pub mod bloom; + +/// One definition that supplies a persisted aggregate and registers every read-side component +/// needed to consult it. +/// +/// The writer helper [`ZonedLayoutOptions::with_skip_index`] is the explicit per-column declaration +/// seam. Readers call [`SkipIndex::register`] on their session before opening the file. +pub trait SkipIndex: Debug + Send + Sync + 'static { + /// The aggregate state to persist for `input_dtype`, or `None` when unsupported. + fn aggregate_fn(&self, input_dtype: &DType) -> Option; + + /// Register the aggregate, optional probe function, and predicate rewrite as one operation. + fn register(&self, session: &VortexSession); +} + +impl ZonedLayoutOptions { + /// Add `index` to this zoned writer while retaining the default min/max-style aggregates. + /// + /// `WriteStrategyBuilder::with_field_zoned_options` can install the configured options for one + /// field while retaining the default data layout pipeline. + pub fn with_skip_index( + mut self, + index: &I, + input_dtype: &DType, + session: &VortexSession, + ) -> VortexResult { + let aggregate_fn = index + .aggregate_fn(input_dtype) + .ok_or_else(|| vortex_err!("skip index does not support input dtype {input_dtype}"))?; + + let mut aggregate_fns = self + .aggregate_fns + .take() + .unwrap_or_else(|| default_zoned_aggregate_fns(input_dtype, session)) + .to_vec(); + if !aggregate_fns.iter().any(|stored| stored == &aggregate_fn) { + aggregate_fns.push(aggregate_fn); + } + self.aggregate_fns = Some(Arc::from(aggregate_fns)); + Ok(self) + } +} diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 4151679e6c3..d4dcad7e8a6 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -38,7 +38,7 @@ use crate::LayoutWriterContext; use crate::layouts::zoned::AggregateStatsAccumulator; use crate::layouts::zoned::ZonedLayout; use crate::layouts::zoned::aggregate_partials; -use crate::layouts::zoned::schema::default_bounded_stat_max_bytes; +use crate::layouts::zoned::aggregates::default_zoned_aggregate_fns; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; @@ -50,6 +50,7 @@ use crate::sequence::SequentialStreamExt; /// /// The input stream is assumed to already be partitioned into one chunk per zone, except /// possibly the final partial zone. +#[derive(Clone)] pub struct ZonedLayoutOptions { /// The size of a statistics block pub block_size: NonZeroUsize, From a9e0550c93223daba977962b10cafce576c12370 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Wed, 12 Aug 2026 22:25:19 +0200 Subject: [PATCH 2/7] Update bloom filter to use Split Block Bloom Filters (SBBFs) and implement constant/canonical accumulators. --- Cargo.lock | 1 + Cargo.toml | 1 + vortex-file/tests/bloom_skip_index.rs | 6 +- vortex-layout/Cargo.toml | 1 + .../layouts/zoned/aggregates/bloom_filter.rs | 308 -------------- .../aggregates/bloom_filter/canonical/bool.rs | 114 ++++++ .../bloom_filter/canonical/decimal.rs | 91 ++++ .../bloom_filter/canonical/extension.rs | 91 ++++ .../aggregates/bloom_filter/canonical/mod.rs | 51 +++ .../bloom_filter/canonical/primitive.rs | 286 +++++++++++++ .../bloom_filter/canonical/varbin.rs | 136 ++++++ .../zoned/aggregates/bloom_filter/constant.rs | 193 +++++++++ .../zoned/aggregates/bloom_filter/mod.rs | 387 ++++++++++++++++++ .../zoned/aggregates/bloom_filter/partial.rs | 140 +++++++ .../src/layouts/zoned/aggregates/mod.rs | 5 +- vortex-layout/src/layouts/zoned/mod.rs | 8 + .../src/layouts/zoned/skip_index/bloom.rs | 211 +++++++--- 17 files changed, 1665 insertions(+), 365 deletions(-) delete mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/decimal.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/varbin.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs create mode 100644 vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs diff --git a/Cargo.lock b/Cargo.lock index a92a0f5be59..8bdbeca0cb3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10344,6 +10344,7 @@ dependencies = [ "termtree", "tokio", "tracing", + "twox-hash", "vortex-array", "vortex-arrow", "vortex-btrblocks", diff --git a/Cargo.toml b/Cargo.toml index 65452f18300..dfda9f8a959 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -281,6 +281,7 @@ tpchgen-arrow = { version = "2.0.2", git = "https://github.com/clflushopt/tpchge tracing = { version = "0.1.41", default-features = false } tracing-perfetto = "0.1.5" tracing-subscriber = "0.3" +twox-hash = "2.1.2" url = "2.5.7" uuid = { version = "1.23", features = ["js"] } wasm-bindgen-futures = "0.4.58" diff --git a/vortex-file/tests/bloom_skip_index.rs b/vortex-file/tests/bloom_skip_index.rs index 6e1b6c51986..6b8371a5e4b 100644 --- a/vortex-file/tests/bloom_skip_index.rs +++ b/vortex-file/tests/bloom_skip_index.rs @@ -19,7 +19,6 @@ #![expect(clippy::expect_used)] -use std::num::NonZeroU8; use std::num::NonZeroUsize; use std::sync::Arc; @@ -62,10 +61,7 @@ const MISS: i64 = 503; fn bloom() -> BloomSkipIndex { // A deliberately roomy filter keeps this correctness test's false-positive outcome // deterministic. False positives are valid Bloom behavior, but false negatives are not. - BloomSkipIndex::new(BloomOptions::new( - NonZeroUsize::new(1024).expect("1024 is non-zero"), - NonZeroU8::new(5).expect("5 is non-zero"), - )) + BloomSkipIndex::new(BloomOptions::default()) } fn session(index: Option<&dyn SkipIndex>) -> VortexSession { diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index f772b9ab639..bae7cbd0efc 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -39,6 +39,7 @@ sketches-ddsketch = { workspace = true } termtree = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } +twox-hash = { workspace = true } vortex-array = { workspace = true } vortex-arrow = { workspace = true } vortex-btrblocks = { workspace = true } diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter.rs deleted file mode 100644 index d1d3707beab..00000000000 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter.rs +++ /dev/null @@ -1,308 +0,0 @@ -//! Bloom-filter aggregate for zoned layouts. - -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fmt; -use std::fmt::Display; -use std::fmt::Formatter; -use std::num::NonZeroU8; -use std::num::NonZeroUsize; - -use vortex_array::ArrayRef; -use vortex_array::Columnar; -use vortex_array::ExecutionCtx; -use vortex_array::aggregate_fn::AggregateFnId; -use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::scalar::Scalar; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; -use vortex_session::VortexSession; -use vortex_session::registry::CachedId; - -/// Bloom-filter tuning persisted as aggregate metadata. -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub struct BloomOptions { - bytes: NonZeroUsize, - hashes: NonZeroU8, -} - -impl BloomOptions { - /// Create bloom options with a fixed number of bytes and hash probes per zone. - pub fn new(bytes: NonZeroUsize, hashes: NonZeroU8) -> Self { - Self { bytes, hashes } - } - - /// Bytes stored for each zone. - pub fn bytes(&self) -> NonZeroUsize { - self.bytes - } - - /// Hash probes performed for each inserted or tested value. - pub fn hashes(&self) -> NonZeroU8 { - self.hashes - } -} - -impl Default for BloomOptions { - fn default() -> Self { - Self { - // Eight bits per row at the default 8192-row zone size. - bytes: NonZeroUsize::new(8192).unwrap_or(NonZeroUsize::MIN), - hashes: NonZeroU8::new(5).unwrap_or(NonZeroU8::MIN), - } - } -} - -impl Display for BloomOptions { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - write!(f, "bytes={},hashes={}", self.bytes, self.hashes) - } -} - -/// Aggregate that stores one fixed-size Bloom bitset as a `Binary` scalar for every zone. -#[derive(Clone, Debug)] -pub(in crate::layouts::zoned) struct BloomFilter; - -/// In-memory Bloom accumulator. Only the bitset is persisted. -pub(in crate::layouts::zoned) struct BloomPartial { - bits: Vec, - hashes: u8, -} - -impl AggregateFnVTable for BloomFilter { - type Options = BloomOptions; - type Partial = BloomPartial; - - fn id(&self) -> AggregateFnId { - static ID: CachedId = CachedId::new("vortex.bloom_filter.i64.v1"); - *ID - } - - fn serialize(&self, options: &Self::Options) -> VortexResult>> { - let bytes = u32::try_from(options.bytes.get())?; - let mut metadata = bytes.to_le_bytes().to_vec(); - metadata.push(options.hashes.get()); - Ok(Some(metadata)) - } - - fn deserialize( - &self, - metadata: &[u8], - _session: &VortexSession, - ) -> VortexResult { - vortex_ensure!(metadata.len() == 5, "invalid bloom metadata length"); - let bytes = u32::from_le_bytes([metadata[0], metadata[1], metadata[2], metadata[3]]); - Ok(BloomOptions::new( - NonZeroUsize::new(bytes as usize) - .ok_or_else(|| vortex_err!("bloom byte length must be non-zero"))?, - NonZeroU8::new(metadata[4]) - .ok_or_else(|| vortex_err!("bloom hash count must be non-zero"))?, - )) - } - - fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { - matches!(input_dtype, DType::Primitive(PType::I64, _)) - .then_some(DType::Binary(Nullability::NonNullable)) - } - - fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { - self.return_dtype(options, input_dtype) - } - - fn empty_partial( - &self, - options: &Self::Options, - _input_dtype: &DType, - ) -> VortexResult { - Ok(BloomPartial { - bits: vec![0; options.bytes.get()], - hashes: options.hashes.get(), - }) - } - - fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { - if other.is_null() { - return Ok(()); - } - let other = other - .as_binary() - .value() - .ok_or_else(|| vortex_err!("non-null bloom partial has no bytes"))?; - vortex_ensure!( - partial.bits.len() == other.len(), - "bloom partial length mismatch" - ); - for (dst, src) in partial.bits.iter_mut().zip(other.as_slice()) { - *dst |= *src; - } - Ok(()) - } - - fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - Ok(Scalar::binary( - partial.bits.clone(), - Nullability::NonNullable, - )) - } - - fn reset(&self, partial: &mut Self::Partial) { - partial.bits.fill(0); - } - - fn is_saturated(&self, partial: &Self::Partial) -> bool { - partial.bits.iter().all(|byte| *byte == u8::MAX) - } - - fn accumulate( - &self, - partial: &mut Self::Partial, - batch: &Columnar, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - match batch { - Columnar::Constant(constant) => { - if let Some(value) = i64_value(constant.scalar())? { - bloom_insert(&mut partial.bits, value, partial.hashes); - } - } - Columnar::Canonical(canonical) => { - let primitive = canonical.as_primitive(); - let values = primitive.as_slice::(); - let validity = primitive.validity()?.execute_mask(values.len(), ctx)?; - for (&value, valid) in values.iter().zip(validity.iter()) { - if valid { - bloom_insert(&mut partial.bits, value, partial.hashes); - } - } - } - } - Ok(()) - } - - fn finalize(&self, partials: ArrayRef) -> VortexResult { - Ok(partials) - } - - fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { - self.to_scalar(partial) - } -} - -pub(in crate::layouts::zoned) fn i64_value(scalar: &Scalar) -> VortexResult> { - if scalar.is_null() { - return Ok(None); - } - scalar - .as_primitive_opt() - .and_then(|primitive| primitive.typed_value::()) - .map(Some) - .ok_or_else(|| vortex_err!("bloom value must be i64")) -} - -fn bloom_insert(bits: &mut [u8], value: i64, hashes: u8) { - bloom_insert_hash( - bits, - splitmix64(value as u64 ^ 0x243f_6a88_85a3_08d3), - hashes, - ); -} - -fn bloom_insert_hash(bits: &mut [u8], hash: u64, hashes: u8) { - for (byte, bit) in bloom_positions(hash, bits.len(), hashes) { - bits[byte] |= 1 << bit; - } -} - -pub(in crate::layouts::zoned) fn bloom_contains(bits: &[u8], value: i64, hashes: u8) -> bool { - bloom_contains_hash( - bits, - splitmix64(value as u64 ^ 0x243f_6a88_85a3_08d3), - hashes, - ) -} - -fn bloom_contains_hash(bits: &[u8], hash: u64, hashes: u8) -> bool { - bloom_positions(hash, bits.len(), hashes).all(|(byte, bit)| bits[byte] & (1 << bit) != 0) -} - -fn bloom_positions(hash: u64, bytes: usize, hashes: u8) -> impl Iterator { - let h1 = hash; - let h2 = splitmix64(h1 ^ 0x1319_8a2e_0370_7344) | 1; - let bit_len = u64::try_from(bytes).unwrap_or(u64::MAX).saturating_mul(8); - (0..u64::from(hashes)).map(move |probe| { - let position = h1 - .wrapping_add(probe.wrapping_mul(h2)) - .wrapping_rem(bit_len); - // `position / 8` is less than `bytes`, which is already a `usize`. - let byte = usize::try_from(position / 8).unwrap_or_default(); - let bit = u32::try_from(position % 8).unwrap_or_default(); - (byte, bit) - }) -} - -fn splitmix64(mut value: u64) -> u64 { - value = value.wrapping_add(0x9e37_79b9_7f4a_7c15); - value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); - value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); - value ^ (value >> 31) -} - -#[cfg(test)] -mod tests { - use std::num::NonZeroU8; - use std::num::NonZeroUsize; - - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::Accumulator; - use vortex_array::aggregate_fn::AggregateFnVTable; - use vortex_array::aggregate_fn::DynAccumulator; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_error::VortexResult; - - use super::BloomFilter; - use super::BloomOptions; - use super::bloom_contains; - - fn small_options() -> BloomOptions { - BloomOptions::new( - NonZeroUsize::new(64).expect("64 is non-zero"), - NonZeroU8::new(3).expect("3 is non-zero"), - ) - } - - #[test] - fn roundtrips_options_and_membership() -> VortexResult<()> { - let session = vortex_array::array_session(); - let options = small_options(); - let metadata = BloomFilter - .serialize(&options)? - .expect("bloom is serializable"); - assert_eq!(BloomFilter.deserialize(&metadata, &session)?, options); - - let mut ctx = session.create_execution_ctx(); - let mut accumulator = Accumulator::try_new( - BloomFilter, - options.clone(), - DType::Primitive(PType::I64, Nullability::NonNullable), - )?; - accumulator.accumulate( - &PrimitiveArray::from_iter([10i64, 20, 30]).into_array(), - &mut ctx, - )?; - let state = accumulator.finish()?; - let bytes = state.as_binary().value().expect("bloom state is non-null"); - assert!(bloom_contains(bytes.as_slice(), 10, options.hashes.get())); - assert!(bloom_contains(bytes.as_slice(), 20, options.hashes.get())); - assert!(!bloom_contains(bytes.as_slice(), 999, options.hashes.get())); - Ok(()) - } -} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs new file mode 100644 index 00000000000..747fd353457 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ExecutionCtx; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::bool::BoolArrayExt; +use vortex_error::VortexResult; +use vortex_mask::AllOr; + +use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; + +/// Similar to [vortex_array::aggregate_fn::fns::min_max::accumulate_bool] +pub(super) fn accumulate_bool( + array: &BoolArray, + partial: &mut BloomPartial, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let mask = array.validity()?.execute_mask(array.len(), ctx)?; + let bits = array.bit_buffer_view(); + + let (true_count, valid_count) = match mask.bit_buffer() { + AllOr::None => return Ok(()), + AllOr::All => (bits.true_count() as u64, array.len() as u64), + AllOr::Some(validity) => { + let masked = bits.to_bit_buffer() & validity; + (masked.true_count() as u64, validity.true_count() as u64) + } + }; + + if true_count > 0 { + partial.insert(true); + } + if true_count < valid_count { + partial.insert(false); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::arrays::BoolArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::build_filter; + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::setup; + + #[rstest] + #[case::inserts_each_valid_boolean_value( + &[Some(true), Some(false)], + Nullability::NonNullable, + true, + true + )] + #[case::inserts_only_false( + &[Some(false), Some(false)], + Nullability::NonNullable, + false, + true + )] + #[case::inserts_only_true( + &[Some(true), Some(true)], + Nullability::NonNullable, + true, + false + )] + #[case::ignores_null_boolean_values( + &[Some(true), None, Some(true)], + Nullability::Nullable, + true, + false + )] + #[case::all_null_booleans_leave_the_filter_empty( + &[None, None], + Nullability::Nullable, + false, + false + )] + fn membership( + #[case] values: &[Option], + #[case] nullability: Nullability, + #[case] expect_true: bool, + #[case] expect_false: bool, + ) -> VortexResult<()> { + let ctx = setup()?; + let array = match nullability { + Nullability::NonNullable => BoolArray::from_iter( + values + .iter() + .copied() + .collect::>>() + .ok_or_else(|| vortex_err!("non-null test case contains a null"))?, + ), + Nullability::Nullable => BoolArray::from_iter(values.iter().copied()), + }; + let bloom_filter = build_filter(array.into_array(), DType::Bool(nullability), ctx)?; + + assert_eq!( + bloom_filter.contains_valid_scalar(&Scalar::bool(true, nullability))?, + expect_true + ); + assert_eq!( + bloom_filter.contains_valid_scalar(&Scalar::bool(false, nullability))?, + expect_false + ); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/decimal.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/decimal.rs new file mode 100644 index 00000000000..bfd81716fc5 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/decimal.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::Array; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::Decimal; +use vortex_array::match_each_decimal_value_type; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; + +pub(super) fn accumulate_decimal( + array: &Array, + partial: &mut BloomPartial, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + match_each_decimal_value_type!(array.values_type(), |D| { + match array.validity()?.execute_mask(array.len(), ctx)? { + Mask::AllTrue(_) => { + array + .buffer::() + .iter() + .for_each(|value| partial.insert(value)); + } + Mask::AllFalse(_) => {} + Mask::Values(v) => { + array + .buffer::() + .iter() + .zip(v.bit_buffer().iter()) + .for_each(|(value, valid)| { + if valid { + partial.insert(value) + } + }); + } + } + }); + + Ok(()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::NativeDecimalType; + use vortex_array::dtype::Nullability; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::build_filter; + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::setup; + + #[rstest] + #[case(3u8, 0i8, &[1i8, 2, 3], 99i8)] + #[case(5u8, 1i8, &[10i16, 20, 30], 99i16)] + #[case(9u8, 2i8, &[1000i32, 2000, 3000], 99999i32)] + #[case(18u8, 2i8, &[1000i64, 2000, 3000], 99999i64)] + #[case(10u8, 2i8, &[1000i128, 2000, 3000], 99999i128)] + fn membership( + #[case] precision: u8, + #[case] scale: i8, + #[case] present: &[T], + #[case] absent: T, + ) -> VortexResult<()> + where + T: Copy + Into + NativeDecimalType, + { + let ctx = setup()?; + let decimal_dtype = DecimalDType::new(precision, scale); + let dtype = DType::Decimal(decimal_dtype, Nullability::NonNullable); + let values: DecimalArray = DecimalArray::from_iter(present.iter().copied(), decimal_dtype); + let bloom_filter = build_filter(values.into_array(), dtype, ctx)?; + + for &v in present { + let scalar = Scalar::decimal(v.into(), decimal_dtype, Nullability::NonNullable); + assert!(bloom_filter.contains_valid_scalar(&scalar)?); + } + + let absent_scalar = Scalar::decimal(absent.into(), decimal_dtype, Nullability::NonNullable); + assert!(!bloom_filter.contains_valid_scalar(&absent_scalar)?); + + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs new file mode 100644 index 00000000000..08762093e86 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_error::VortexResult; + +use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; + +pub(super) fn accumulate_extension( + array: &ExtensionArray, + partial: &mut BloomPartial, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let storage = array.storage_array().clone(); + let canonical = storage.execute::(ctx)?; + super::accumulate_canonical(&canonical, partial, ctx) +} + +#[cfg(test)] +mod tests { + use vortex_array::IntoArray; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::extension::datetime::Timestamp; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::build_filter; + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::setup; + + #[test] + fn hashes_extension_values_through_storage() -> VortexResult<()> { + let ctx = setup()?; + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); + let bloom_filter = build_filter( + ExtensionArray::new( + ext_dtype.clone(), + PrimitiveArray::from_iter([1_000i64, 2_000, 3_000]).into_array(), + ) + .into_array(), + DType::Extension(ext_dtype.clone()), + ctx, + )?; + + for value in [1_000i64, 2_000, 3_000] { + let scalar = Scalar::extension_ref( + ext_dtype.clone(), + Scalar::primitive(value, Nullability::NonNullable), + ); + assert!(bloom_filter.contains_valid_scalar(&scalar)?); + } + + let absent = Scalar::extension_ref( + ext_dtype, + Scalar::primitive(4_000i64, Nullability::NonNullable), + ); + assert!(!bloom_filter.contains_valid_scalar(&absent)?); + Ok(()) + } + + #[test] + fn ignores_null_extension_values() -> VortexResult<()> { + let ctx = setup()?; + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased(); + let bloom_filter = build_filter( + ExtensionArray::new( + ext_dtype.clone(), + PrimitiveArray::from_option_iter([Some(1_000i64), None, Some(3_000)]).into_array(), + ) + .into_array(), + DType::Extension(ext_dtype.clone()), + ctx, + )?; + + let present = Scalar::extension_ref( + ext_dtype.clone(), + Scalar::primitive(1_000i64, Nullability::Nullable), + ); + let null_slot_value = + Scalar::extension_ref(ext_dtype, Scalar::primitive(0i64, Nullability::Nullable)); + assert!(bloom_filter.contains_valid_scalar(&present)?); + assert!(!bloom_filter.contains_valid_scalar(&null_slot_value)?); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs new file mode 100644 index 00000000000..bee291cd548 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +mod bool; +mod decimal; +mod extension; +mod primitive; +mod varbin; + +use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; +use crate::layouts::zoned::aggregates::bloom_filter::canonical::bool::accumulate_bool; +use crate::layouts::zoned::aggregates::bloom_filter::canonical::decimal::accumulate_decimal; +use crate::layouts::zoned::aggregates::bloom_filter::canonical::extension::accumulate_extension; +use crate::layouts::zoned::aggregates::bloom_filter::canonical::primitive::accumulate_primitive; +use crate::layouts::zoned::aggregates::bloom_filter::canonical::varbin::accumulate_varbin; + +pub(super) fn accumulate_canonical( + canonical: &Canonical, + partial: &mut BloomPartial, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + match canonical { + Canonical::Bool(array) => accumulate_bool(array, partial, ctx)?, + Canonical::Primitive(array) => accumulate_primitive(array, partial, ctx)?, + Canonical::Decimal(array) => accumulate_decimal(array, partial, ctx)?, + Canonical::VarBinView(array) => accumulate_varbin(array, partial, ctx)?, + Canonical::Extension(array) => accumulate_extension(array, partial, ctx)?, + + // Nulls are skipped and are not included in any Bloom filter. + Canonical::Null(_) => {} + + // TODO (joacoc): pending canonical + Canonical::Struct(_) + | Canonical::List(_) + | Canonical::FixedSizeList(_) + | Canonical::Variant(_) + | Canonical::Union(_) => { + vortex_bail!( + "Unsupported canonical type for bloom filter: {}", + canonical.dtype() + ) + } + } + + Ok(()) +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs new file mode 100644 index 00000000000..24b2418ffa0 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs @@ -0,0 +1,286 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::Array; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::Primitive; +use vortex_array::dtype::PType; +use vortex_array::match_each_float_ptype; +use vortex_array::match_each_integer_ptype; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; + +pub(super) fn accumulate_primitive( + array: &Array, + partial: &mut BloomPartial, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + // TODO (joacoc): What about a single new macro that separates both: + // match... { + // Floats => .. + // Integers => .. + // } + match array.ptype() { + PType::F16 | PType::F32 | PType::F64 => { + match_each_float_ptype!(array.ptype(), |T| { + let slice = array.as_slice::(); + + match array.validity()?.execute_mask(slice.len(), ctx)? { + Mask::AllTrue(_) => { + for &value in slice { + partial.insert(value.to_bits()); + } + } + Mask::AllFalse(_) => {} + Mask::Values(mask_values) => { + for &(start, end) in mask_values.slices() { + for &value in &slice[start..end] { + partial.insert(value.to_bits()); + } + } + } + }; + + // TODO (joacoc): What about using density threshold? + // const INSERT_SLICES_DENSITY_THRESHOLD: f64 = 0.8; + // let slice = array.as_slice::(); + // match array + // .validity()? + // .execute_mask(slice.len(), ctx)? + // .threshold_iter(INSERT_SLICES_DENSITY_THRESHOLD) + // { + // AllOr::None => {} + // AllOr::All => { + // for &value in slice { + // partial.insert(value.to_bits()); + // } + // } + // AllOr::Some(MaskIter::Slices(slices)) => { + // for &(start, end) in slices { + // for &value in &slice[start..end] { + // partial.insert(value.to_bits()); + // } + // } + // } + // AllOr::Some(MaskIter::Indices(indices)) => { + // for &idx in indices { + // partial.insert(slice[idx].to_bits()); + // } + // } + // } + }); + } + + _ => { + match_each_integer_ptype!(array.ptype(), |T| { + let slice = array.as_slice::(); + + match array.validity()?.execute_mask(slice.len(), ctx)? { + Mask::AllTrue(_) => { + for &value in slice { + partial.insert(value); + } + } + Mask::AllFalse(_) => {} + Mask::Values(mask_values) => { + for &(start, end) in mask_values.slices() { + for &value in &slice[start..end] { + partial.insert(value); + } + } + } + }; + + // TODO (joacoc): What about using density threshold? + // let slice = array.as_slice::(); + // match array + // .validity()? + // .execute_mask(slice.len(), ctx)? + // .threshold_iter(INSERT_SLICES_DENSITY_THRESHOLD) + // { + // AllOr::None => {} + // AllOr::All => { + // for &value in slice { + // partial.insert(value); + // } + // } + // AllOr::Some(MaskIter::Slices(slices)) => { + // for &(start, end) in slices { + // for &value in &slice[start..end] { + // partial.insert(value); + // } + // } + // } + // AllOr::Some(MaskIter::Indices(indices)) => { + // for &idx in indices { + // partial.insert(slice[idx]); + // } + // } + // } + }); + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::NativePType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + #[cfg(test)] + use vortex_array::scalar::PValue; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::build_filter; + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::setup; + + #[rstest] + #[case(&[10i8, 20, 30], 99i8)] + #[case(&[10i16, 20, 30], 99i16)] + #[case(&[10i32, 20, 30], 999i32)] + #[case(&[10i64, 20, 30], 999i64)] + #[case(&[10.0f32, 20.0, 30.0], 999.0f32)] + #[case(&[10.0f64, 20.0, 30.0], 999.0f64)] + fn membership(#[case] present: &[T], #[case] absent: T) -> VortexResult<()> + where + T: Copy + NativePType + Into, + { + let ctx = setup()?; + let values = PrimitiveArray::from_iter(present.iter().copied()); + let bloom_filter = build_filter( + values.into_array(), + DType::Primitive(T::PTYPE, Nullability::NonNullable), + ctx, + )?; + for &v in present { + let scalar = Scalar::primitive(v, Nullability::NonNullable); + assert!(bloom_filter.contains_valid_scalar(&scalar)?); + } + let scalar = Scalar::primitive(absent, Nullability::NonNullable); + assert!(!bloom_filter.contains_valid_scalar(&scalar)?); + Ok(()) + } + + #[rstest] + #[case(&[10i8, 20, 30, 40, 50])] + fn validity(#[case] present: &[T]) -> VortexResult<()> + where + T: Copy + NativePType + Into, + { + { + let ctx = setup()?; + let all_valid = PrimitiveArray::from_option_iter(present.iter().map(|&v| Some(v))); + let bloom_filter = build_filter( + all_valid.into_array(), + DType::Primitive(T::PTYPE, Nullability::Nullable), + ctx, + )?; + for &v in present { + let scalar = Scalar::primitive(v, Nullability::Nullable); + assert!(bloom_filter.contains_valid_scalar(&scalar)?); + } + } + + // Mask::AllFalse: every position null. Nothing should ever be found — including the + // `T::default()` filler that `from_option_iter` writes into null slots internally. + { + let ctx = setup()?; + let all_invalid = PrimitiveArray::from_option_iter(present.iter().map(|_| None::)); + let bloom_filter = build_filter( + all_invalid.into_array(), + DType::Primitive(T::PTYPE, Nullability::Nullable), + ctx, + )?; + for &v in present { + let scalar = Scalar::primitive(v, Nullability::Nullable); + assert!(!bloom_filter.contains_valid_scalar(&scalar)?); + } + } + + // Mask::Values: alternating valid/null. Only the valid positions should be members, + // and the default filler at null slots must not leak through as a false member. + { + let ctx = setup()?; + let mixed: Vec> = present + .iter() + .enumerate() + .map(|(i, &v)| if i % 2 == 0 { Some(v) } else { None }) + .collect(); + let bloom_filter = build_filter( + PrimitiveArray::from_option_iter(mixed).into_array(), + DType::Primitive(T::PTYPE, Nullability::Nullable), + ctx, + )?; + for (i, &v) in present.iter().enumerate() { + if i % 2 == 0 { + let scalar = Scalar::primitive(v, Nullability::Nullable); + assert!(bloom_filter.contains_valid_scalar(&scalar)?); + } + } + } + + Ok(()) + } + + // The idea is here is to test two different NaN float values. + // + // Given that the bloom filter uses bits, it means that one + // NaN value could be present but others not. + // + // So the following test does the following, insert a NaN + // value, update the NaN value to another NaN value with a different + // set of bits, and then check that it doesn't exists. + #[test] + fn nan_bit_patterns_are_distinct_members() -> VortexResult<()> { + let ctx = setup()?; + let canonical_nan = f64::NAN; + + let present = [1.0_f64, canonical_nan, 3.0]; + let values = PrimitiveArray::from_iter(present); + let bloom_filter = build_filter( + values.into_array(), + DType::Primitive(PType::F64, Nullability::NonNullable), + ctx, + )?; + + assert!( + bloom_filter + .contains_valid_scalar(&Scalar::primitive(1.0_f64, Nullability::NonNullable))? + ); + assert!( + bloom_filter + .contains_valid_scalar(&Scalar::primitive(3.0_f64, Nullability::NonNullable))? + ); + assert!( + bloom_filter.contains_valid_scalar(&Scalar::primitive( + canonical_nan, + Nullability::NonNullable + ))? + ); + + // Check that another NaN doesn't exists. + // Update the canonical + let other_nan = f64::from_bits(canonical_nan.to_bits() ^ 0x1); + assert!(other_nan.is_nan()); + + assert!( + !bloom_filter + .contains_valid_scalar(&Scalar::primitive(other_nan, Nullability::NonNullable))? + ); + assert!( + !bloom_filter + .contains_valid_scalar(&Scalar::primitive(999.0_f64, Nullability::NonNullable))? + ); + + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/varbin.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/varbin.rs new file mode 100644 index 00000000000..a12ae2d3453 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/varbin.rs @@ -0,0 +1,136 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::Array; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::varbinview::BinaryView; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; + +pub(super) fn accumulate_varbin( + array: &Array, + partial: &mut BloomPartial, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + // Utility function to process views in both validity cases. + // Other agg. functions handle both cases together and always + // check the validity, even when all values are valid. + // I don't think there is much difference. + let mut process_view = |view: &BinaryView, buffers: &Vec<&Buffer>| { + if view.is_inlined() { + partial.insert(view.as_inlined().value()); + } else { + let view_ref = view.as_view(); + let value = &buffers[view_ref.buffer_index as usize][view_ref.as_range()]; + partial.insert(value); + } + }; + + match array.validity()?.execute_mask(array.len(), ctx)? { + Mask::AllTrue(_) => { + let buffers = array + .data_buffers() + .iter() + .map(|b| b.as_host()) + .collect::>(); + array + .views() + .iter() + .for_each(|view| process_view(view, &buffers)); + } + Mask::AllFalse(_) => todo!(), + Mask::Values(mask_values) => { + let views_iter = array.views().iter(); + let buffers = array + .data_buffers() + .iter() + .map(|b| b.as_host()) + .collect::>(); + + views_iter + .zip(mask_values.bit_buffer()) + .for_each(|(view, valid)| { + if valid { + process_view(view, &buffers) + } + }); + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::build_filter; + use crate::layouts::zoned::aggregates::bloom_filter::test_utils::setup; + + #[rstest] + #[case::inlined(&["a", "Lorem", "ipsum"], "neverever")] // inlined vs non-inline + #[case::buffered( + &[ + "Lorem ipsum dolor sit amet, consectetur adipiscing elit", + "sed do eiusmod tempor incididunt ut labore et dolore magna aliqua", + ], + "Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip", + )] + fn roundtrips_options_and_membership_varbin( + #[case] present: &[&str], + #[case] absent: &str, + ) -> VortexResult<()> { + let ctx = setup()?; + let dtype = DType::Utf8(Nullability::NonNullable); + let batch = VarBinViewArray::from_iter_str(present.iter().copied()); + let bloom_filter = build_filter(batch.into_array(), dtype, ctx)?; + + for &v in present { + let scalar = Scalar::binary(v.as_bytes().to_vec(), Nullability::NonNullable); + assert!(bloom_filter.contains_valid_scalar(&scalar)?); + } + + let absent_scalar = Scalar::binary(absent.as_bytes().to_vec(), Nullability::NonNullable); + assert!(!bloom_filter.contains_valid_scalar(&absent_scalar)?); + Ok(()) + } + + #[test] + fn roundtrips_options_and_membership_varbin_mixed_with_nulls() -> VortexResult<()> { + let ctx = setup()?; + let dtype = DType::Utf8(Nullability::Nullable); + let values = VarBinViewArray::from_iter( + vec![ + Some("short"), + None, + Some("Lorem ipsum dolor sit amet, consectetur adipiscing elit"), + None, + ], + dtype.clone(), + ); + let bloom_filter = build_filter(values.into_array(), dtype, ctx)?; + + let present = Scalar::binary(b"short".to_vec(), Nullability::NonNullable); + assert!(bloom_filter.contains_valid_scalar(&present)?); + + let present_long = Scalar::binary( + b"Lorem ipsum dolor sit amet, consectetur adipiscing elit".to_vec(), + Nullability::NonNullable, + ); + assert!(bloom_filter.contains_valid_scalar(&present_long)?); + + let absent = Scalar::binary(b"never ever".to_vec(), Nullability::NonNullable); + assert!(!bloom_filter.contains_valid_scalar(&absent)?); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs new file mode 100644 index 00000000000..34fc9a4469b --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs @@ -0,0 +1,193 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::match_each_float_ptype; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; + +pub(super) fn accumulate_constant( + constant: &ConstantArray, + partial: &mut BloomPartial, +) -> VortexResult<()> { + let scalar = constant.scalar(); + + // Omit NULL values on purpose. + if scalar.is_null() { + return Ok(()); + } + + partial.insert_hash(partial.hash_valid_scalar(scalar)?); + Ok(()) +} + +impl BloomPartial { + /// Scalar values must be valid otherwise the function will raise an err. + /// + /// This function is used by both, for accumulating scalars, + /// but also to get a scalar value membership. + pub(in crate::layouts::zoned) fn hash_valid_scalar( + &self, + scalar: &Scalar, + ) -> VortexResult { + if scalar.is_null() { + return Err(vortex_err!("cannot hash invalid scalars in bloom filter")); + } + + Ok(match scalar.dtype() { + DType::Extension(_) => { + self.hash_valid_scalar(&scalar.as_extension().to_storage_scalar())? + } + DType::Bool(_) => self.hash( + scalar + .as_bool() + .value() + .vortex_expect("non-null boolean value"), + ), + DType::Primitive(ptype, _) => match ptype { + PType::F16 | PType::F32 | PType::F64 => { + match_each_float_ptype!(ptype, |T| { + let value = scalar + .as_primitive() + .typed_value::() + .vortex_expect("non-null primitive value"); + self.hash(value.to_bits()) + }) + } + _ => match_each_integer_ptype!(ptype, |T| { + let value = scalar + .as_primitive() + .typed_value::() + .vortex_expect("non-null primitive value"); + self.hash(value) + }), + }, + DType::Decimal(..) => { + let decimal = scalar + .as_decimal() + .decimal_value() + .vortex_expect("non-null decimal value"); + match decimal { + DecimalValue::I8(v) => self.hash(v), + DecimalValue::I16(v) => self.hash(v), + DecimalValue::I32(v) => self.hash(v), + DecimalValue::I64(v) => self.hash(v), + DecimalValue::I128(v) => self.hash(v), + DecimalValue::I256(v) => self.hash(v), + } + } + DType::Utf8(_) => { + let buffer = scalar + .as_utf8() + .value() + .vortex_expect("non-null utf8 value"); + self.hash(buffer.as_bytes()) + } + DType::Binary(_) => { + let buffer = scalar + .as_binary() + .value() + .vortex_expect("non-null binary value"); + self.hash(buffer.as_slice()) + } + other => { + return Err(vortex_err!("bloom filter does not support dtype {other}")); + } + }) + } + + pub(in crate::layouts::zoned) fn contains_valid_scalar( + &self, + scalar: &Scalar, + ) -> VortexResult { + let hash = self.hash_valid_scalar(scalar)?; + Ok(self.find_hash(hash)) + } +} + +#[cfg(test)] +mod tests { + + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::arrays::ConstantArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::extension::datetime::TimeUnit; + use vortex_array::extension::datetime::Timestamp; + use vortex_array::scalar::Scalar; + use vortex_error::VortexResult; + + use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; + use crate::layouts::zoned::aggregates::bloom_filter::BloomOptions; + use crate::layouts::zoned::aggregates::bloom_filter::constant::accumulate_constant; + + #[test] + fn nulls_are_omitted() { + let bloom = BloomFilter; + let mut zone_partial = bloom + .empty_partial( + &BloomOptions::default(), + &DType::Primitive(PType::I32, Nullability::Nullable), + ) + .unwrap(); + + assert!( + accumulate_constant( + &ConstantArray::new( + Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), + 1, + ), + &mut zone_partial, + ) + .is_ok(), + "expected to return Ok() for null scalars" + ) + } + + #[test] + fn null_raises_error_on_hash() { + let bloom = BloomFilter; + let zone_partial = bloom + .empty_partial( + &BloomOptions::default(), + &DType::Primitive(PType::I32, Nullability::Nullable), + ) + .unwrap(); + + assert!( + zone_partial + .contains_valid_scalar(&Scalar::null(DType::Primitive( + PType::I32, + Nullability::Nullable + ))) + .is_err(), + "expected to return Err() for null scalars" + ) + } + + #[test] + fn valid_extension_is_a_member() -> VortexResult<()> { + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); + let scalar = Scalar::extension_ref( + ext_dtype.clone(), + Scalar::primitive(1_000i64, Nullability::NonNullable), + ); + let bloom = BloomFilter; + let mut zone_partial = + bloom.empty_partial(&BloomOptions::default(), &DType::Extension(ext_dtype))?; + + accumulate_constant(&ConstantArray::new(scalar.clone(), 1), &mut zone_partial)?; + + assert!(zone_partial.contains_valid_scalar(&scalar)?); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs new file mode 100644 index 00000000000..d7fd32b7e3a --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -0,0 +1,387 @@ +//! Bloom-filter aggregate for zoned layouts. + +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt; +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hash; +use std::num::NonZeroUsize; + +use vortex_array::ArrayRef; +use vortex_array::Columnar; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +mod canonical; +mod partial; + +pub(in crate::layouts::zoned) mod constant; +pub use partial::BloomPartial; + +use crate::layouts::zoned::skip_index::bloom::is_bloom_valid_dtype; + +// (joacoc) Opted for blocks_count as the main way to tune +// the Bloom filter, but there are also other ways, like +// setting the false-positive probability (FPP) and +// number of distinct values (NDV). Given the statistics +// collected from a random subset of values, they could be +// used to determine the number of blocks automatically, +// which I think is more aligned with how Vortex works with +// encoders. But for now, this is a simple approach that +// pushes the block-count selection back to the user +// and otherwise defaults to [DEFAULT_BLOCKS_COUNT]. +// +// In case that the blocks count stays as is, +// a guide on how to select the number of blocks +// would be great. + +/// Bloom-filter tuning persisted as aggregate metadata. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct BloomOptions { + /// Number of blocks in the split block Bloom filter (SBBF). + /// + /// The filter is partitioned into 256-bit blocks. More blocks reduce the + /// number of distinct values assigned to each block, reducing the + /// false-positive rate at the cost of increased filter size. + /// + /// Defaults to: [DEFAULT_BLOCKS_COUNT]. + /// + /// ### Block size and memory usage + /// + /// As a reference, you can use the following table + /// to understand the relationship between block size and memory usage + /// for a **single zone**. + /// + /// | `blocks_count` | Memory | + /// | --------------: | ----------: | + /// | 8 | **256 B** | + /// | 256 | **8 KiB** | + /// | 8192 | **256 KiB** | + /// | 65,536 | **2 MiB** | + /// | 1,048,576 | **32 MiB** | + blocks_count: NonZeroUsize, +} + +impl BloomOptions { + pub fn new(blocks_count: NonZeroUsize) -> Self { + Self { blocks_count } + } + + pub fn blocks(&self) -> NonZeroUsize { + self.blocks_count + } +} + +/// The default value is derived from the default [WriteStrategyBuilder::row_block_size] +const DEFAULT_BLOCKS_COUNT: usize = 256; + +impl Default for BloomOptions { + fn default() -> Self { + Self { + blocks_count: NonZeroUsize::new(DEFAULT_BLOCKS_COUNT).expect("valid blocks size"), + } + } +} + +impl Display for BloomOptions { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "blocks={}", self.blocks_count) + } +} + +/// Aggregate that stores one fixed-size Bloom block as a `Binary` scalar for every zone. +/// The bloom filter only stores valid scalar values, and does not consider +/// `NULL` values as part of the filter. If you need to know if a zone has null values +/// you should use the NULL count. +/// +/// Empty unit struct, in accordance to [AggregateFnVTable] definition. +#[derive(Clone, Debug)] +pub struct BloomFilter; + +impl AggregateFnVTable for BloomFilter { + type Options = BloomOptions; + type Partial = BloomPartial; + + fn id(&self) -> AggregateFnId { + static ID: CachedId = CachedId::new("vortex.bloom_filter.v1"); + *ID + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + let blocks = u32::try_from(options.blocks_count.get())?; + let metadata = blocks.to_le_bytes().to_vec(); + Ok(Some(metadata)) + } + + fn deserialize( + &self, + metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + vortex_ensure_eq!(metadata.len(), 4, "invalid bloom metadata length"); + let blocks = u32::from_le_bytes([metadata[0], metadata[1], metadata[2], metadata[3]]); + Ok(BloomOptions::new( + NonZeroUsize::new(blocks as usize) + .ok_or_else(|| vortex_err!("bloom blocks length must be non-zero"))?, + )) + } + + /// Returns [Binary(Nullability::NonNullable)] when input [DType] is valid. + /// + /// The [BloomFilter] is serialized/deserialized as a sequence of bytes + /// that represents the filter state. + /// + /// An empty filter rather than being NULL is represented + /// by a zero-initialized byte sequence (`0x0..0`) + fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option { + is_bloom_valid_dtype(input_dtype).then_some(DType::Binary(Nullability::NonNullable)) + } + + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { + self.return_dtype(options, input_dtype) + } + + /// Returns an empty Bloom filter with all blocks zero-initialized. + fn empty_partial(&self, options: &Self::Options, _: &DType) -> VortexResult { + Ok(BloomPartial { + blocks: vec![[0u32; 8]; options.blocks_count.get()], + }) + } + + // Combination happens by doing an OR between both filters bits + fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()> { + if other.is_null() { + return Ok(()); + } + + let bytes = other + .as_binary() + .value() + .ok_or_else(|| vortex_err!("non-null bloom partial has no bytes"))?; + + let other = BloomPartial::try_from(bytes.as_slice())?; + + vortex_ensure_eq!( + partial.blocks.len(), + other.blocks.len(), + "bloom partial block count mismatch — partials built with different blocks_count" + ); + + for (dst, src) in partial.blocks.iter_mut().zip(other.blocks.iter()) { + for i in 0..8 { + dst[i] |= src[i]; + } + } + + Ok(()) + } + + /// Returns the non-nullable binary representation of a bloom filter + /// + /// Basically turns each block into a single byte sequence. + fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { + let mut bytes = Vec::with_capacity(partial.blocks.len() * 8 * size_of::()); + bytes.extend( + partial + .blocks + .iter() + .flatten() + .flat_map(|block_seq| block_seq.to_le_bytes()), + ); + + Ok(Scalar::binary(bytes, Nullability::NonNullable)) + } + + fn reset(&self, partial: &mut Self::Partial) { + partial.blocks.fill([0; 8]); + } + + fn is_saturated(&self, partial: &Self::Partial) -> bool { + partial.blocks.iter().all(|byte| *byte == [u32::MAX; 8]) + } + + fn accumulate( + &self, + partial: &mut Self::Partial, + batch: &Columnar, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match batch { + Columnar::Constant(constant) => constant::accumulate_constant(constant, partial)?, + Columnar::Canonical(canonical) => { + canonical::accumulate_canonical(canonical, partial, ctx)? + } + } + Ok(()) + } + + fn finalize(&self, partials: ArrayRef) -> VortexResult { + Ok(partials) + } + + fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult { + self.to_scalar(partial) + } +} + +// The following functions are utils/useful for tests in canonical and constants. + +#[cfg(test)] +pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + use vortex_error::vortex_ensure; + + use super::*; + + pub fn setup() -> VortexResult { + let session = vortex_array::array_session(); + let options = BloomOptions::default(); + let metadata = BloomFilter + .serialize(&options)? + .expect("bloom is serializable"); + assert_eq!(BloomFilter.deserialize(&metadata, &session)?, options); + + let ctx = session.create_execution_ctx(); + + Ok(ctx) + } + + pub fn extract_bloom_blocks(state: &Scalar) -> VortexResult> { + let bytes = state.as_binary().value().expect("bloom state is non-null"); + vortex_ensure!( + bytes.len() % (8 * size_of::()) == 0, + "invalid bloom state length" + ); + let mut blocks = Vec::with_capacity(bytes.len() / 32); + for block_bytes in bytes.chunks_exact(32) { + let mut block = [0u32; 8]; + for (word, word_bytes) in block.iter_mut().zip(block_bytes.chunks_exact(4)) { + *word = u32::from_le_bytes( + word_bytes + .try_into() + .expect("chunks_exact(4) always produces 4 bytes"), + ); + } + blocks.push(block); + } + Ok(blocks) + } + + pub fn build_filter( + batch: ArrayRef, + dtype: DType, + mut ctx: ExecutionCtx, + ) -> VortexResult { + let mut accumulator = Accumulator::try_new(BloomFilter, BloomOptions::default(), dtype)?; + accumulator.accumulate(&batch.into_array(), &mut ctx)?; + let state = accumulator.finish()?; + let blocks = extract_bloom_blocks(&state)?; + let bloom_filter = BloomPartial::from(blocks); + + Ok(bloom_filter) + } + + #[test] + fn saturation_false_when_empty() -> VortexResult<()> { + let options = BloomOptions::default(); + let partial = + BloomFilter.empty_partial(&options, &DType::Binary(Nullability::NonNullable))?; + assert!(!BloomFilter.is_saturated(&partial)); + Ok(()) + } + + #[test] + fn saturation_true_when_every_block_is_full() { + let blocks = vec![[u32::MAX; 8]; 4]; + let partial = BloomPartial::from(blocks); + + assert!(BloomFilter.is_saturated(&partial)); + } + + #[test] + fn combine_partials_rejects_mismatched_block_counts() -> VortexResult<()> { + let mut smaller = BloomFilter.empty_partial( + &BloomOptions::new(NonZeroUsize::new(4).unwrap()), + &DType::Binary(Nullability::NonNullable), + )?; + let bigger = BloomFilter.empty_partial( + &BloomOptions::default(), + &DType::Binary(Nullability::NonNullable), + )?; + + let bigger_scalar = BloomFilter.to_scalar(&bigger)?; + let result = BloomFilter.combine_partials(&mut smaller, bigger_scalar); + + assert!( + result.is_err(), + "combining partials built with different blocks_count must fail loudly, not corrupt state" + ); + Ok(()) + } + + #[test] + fn combine_partials_unions_two_disjoint_partials() -> VortexResult<()> { + let mut partial = BloomFilter.empty_partial( + &BloomOptions::default(), + &DType::Binary(Nullability::NonNullable), + )?; + for i in 0..50i64 { + partial.insert(i); + } + + let mut secondary_partial = BloomFilter.empty_partial( + &BloomOptions::default(), + &DType::Binary(Nullability::NonNullable), + )?; + for i in 50..100i64 { + secondary_partial.insert(i); + } + + // The following expected works because seed is equal for all. + // If the seed is different for both partials, then this will fail. + let mut expected = BloomFilter.empty_partial( + &BloomOptions::default(), + &DType::Binary(Nullability::NonNullable), + )?; + for i in 0..100i64 { + expected.insert(i); + } + + let secondary_partial_as_scalar = BloomFilter.to_scalar(&secondary_partial)?; + BloomFilter.combine_partials(&mut partial, secondary_partial_as_scalar)?; + + assert_eq!( + partial.blocks, expected.blocks, + "merging via combine_partials should equal a single filter built from the union of inputs" + ); + + for i in 0..100i64 { + let hash = partial.hash(i); + assert!(partial.find_hash(hash), "value {i} missing after merge"); + } + + for i in 101..200i64 { + let hash = partial.hash(i); + assert!( + !partial.find_hash(hash), + "value {i} shouldn't be present after" + ); + } + + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs new file mode 100644 index 00000000000..ead7f2224e7 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs @@ -0,0 +1,140 @@ +//! Split block Bloom filters implementation for vortex. +//! +//! [Split block Bloom filters]: https://arxiv.org/pdf/2101.01719 + +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::hash::Hash; +use std::hash::Hasher; + +use twox_hash::XxHash3_64; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +const BLOCK_SIZE: usize = 8 * size_of::(); + +pub struct BloomPartial { + pub(super) blocks: Vec<[u32; 8]>, +} + +impl BloomPartial { + #[inline] + pub fn len(&self) -> usize { + self.blocks.len() + } + + #[inline] + pub(super) fn insert(&mut self, value: T) + where + T: Hash, + { + let hash = self.hash(value); + self.insert_hash(hash); + } + + #[inline] + pub(super) fn insert_hash(&mut self, hash: u64) { + self.add_hash(hash); + } + + #[inline] + pub(super) fn hash(&self, value: T) -> u64 + where + T: Hash, + { + // > Since the seed is optional, it can be 0. + // Ref: https://github.com/Cyan4973/xxHash/blob/v0.8.3/doc/xxhash_spec.md#step-1-initialize-internal-accumulators + let mut hasher = XxHash3_64::with_seed(0); + value.hash(&mut hasher); + hasher.finish() + } + + /// Hash should be u64 or u32? + fn add_hash(&mut self, hash: u64) { + let idx = self.block_index(hash, self.blocks.len()) as usize; + // Block idx already consumed the hash, + let mask = self.make_mask(hash as u32); + for i in 0..8 { + self.blocks[idx][i] |= mask[i]; + } + } + + /// Check whether a hash is (probably) present in the filter. + pub(super) fn find_hash(&self, hash: u64) -> bool { + let idx = self.block_index(hash, self.blocks.len()) as usize; + let mask = self.make_mask(hash as u32); + + for i in 0..8 { + if self.blocks[idx][i] & mask[i] != mask[i] { + return false; + } + } + + true + } + + /// Takes a hash value and creates a mask with one bit set in each 32-bit lane. + /// These are the bits to set or check when accessing the block. + /// + /// The following code is SIMD friendly and will get vectorized + /// by the compiler automatically (for releases). + fn make_mask(&self, hash: u32) -> [u32; 8] { + let mut out = [0u32; 8]; + + // Set eight odd constants for multiply-shift hashing + let rehash: [u32; 8] = [ + 0x47b6137b, 0x44974d91, 0x8824ad5b, 0xa2b7289d, 0x705495c7, 0x2df1424b, 0x9efc4947, + 0x5c6bfb31, + ]; + + for i in 0..8 { + // Shift all data right, reducing the hash values from 32 bits to five bits. + // Those five bits represent an index in [0, 31) + let y = hash.wrapping_mul(rehash[i]) >> 27; + + // Set a bit in each lane based on using the [0, 32) data as shift values. + out[i] = 1u32 << y; + } + + out + } + + #[inline] + fn block_index(&self, hash: u64, blocks_count: usize) -> u64 { + ((hash >> 32) * (blocks_count as u64)) >> 32 + } +} + +#[cfg(test)] +impl From> for BloomPartial { + fn from(value: Vec<[u32; 8]>) -> Self { + BloomPartial { blocks: value } + } +} + +impl TryFrom<&[u8]> for BloomPartial { + type Error = vortex_error::VortexError; + + /// Reconstruct a partial from its serialized byte representation + /// (the same layout produced by `to_scalar`). + fn try_from(bytes: &[u8]) -> VortexResult { + vortex_ensure!( + !bytes.is_empty() && bytes.len() % BLOCK_SIZE == 0, + "invalid bloom filter byte length: {}", + bytes.len() + ); + + let blocks = bytes + .chunks_exact(BLOCK_SIZE) + .map(|chunk| { + let mut block = [0u32; 8]; + for (word, wb) in block.iter_mut().zip(chunk.chunks_exact(4)) { + *word = u32::from_le_bytes(wb.try_into().unwrap()); + } + block + }) + .collect(); + Ok(BloomPartial { blocks }) + } +} diff --git a/vortex-layout/src/layouts/zoned/aggregates/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/mod.rs index b90d9cd6773..0ea90a5ba5b 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/mod.rs @@ -20,9 +20,6 @@ use vortex_session::VortexSession; pub(in crate::layouts::zoned) mod bloom_filter; mod min_max; -pub(in crate::layouts::zoned) use bloom_filter::BloomFilter; -pub(in crate::layouts::zoned) use bloom_filter::bloom_contains; -pub(in crate::layouts::zoned) use bloom_filter::i64_value; use min_max::min_max_aggregate_fns; pub(super) fn default_zoned_aggregate_fns( @@ -55,7 +52,7 @@ mod tests { use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; - use super::BloomFilter; + use super::bloom_filter::BloomFilter; use super::bloom_filter::BloomOptions; use super::default_zoned_aggregate_fns; diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 4948570ea36..96893dfeac7 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -512,6 +512,8 @@ mod tests { use crate::LayoutBuildContext; use crate::children::OwnedLayoutChildren; use crate::layouts::flat::FlatLayout; + use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; + use crate::layouts::zoned::aggregates::bloom_filter::BloomOptions; use crate::segments::SegmentId; fn aggregate_spec(aggregate_fn: AggregateFnRef) -> AggregateSpecProto { @@ -530,6 +532,12 @@ mod tests { aggregate_spec(Min.bind(NumericalAggregateOpts::skip_nans())), ]), })] + #[case::bloom(ZonedMetadata { + zone_len: 1, + aggregate_specs: Arc::new([ + aggregate_spec(BloomFilter.bind(BloomOptions::default())), + ]), + })] fn test_metadata_serialization(#[case] metadata: ZonedMetadata) { let serialized = metadata.clone().serialize(); assert_eq!(serialized[0], ZONED_METADATA_PROTO_VERSION); diff --git a/vortex-layout/src/layouts/zoned/skip_index/bloom.rs b/vortex-layout/src/layouts/zoned/skip_index/bloom.rs index 7c6985126f3..71c8fb29cc9 100644 --- a/vortex-layout/src/layouts/zoned/skip_index/bloom.rs +++ b/vortex-layout/src/layouts/zoned/skip_index/bloom.rs @@ -11,16 +11,12 @@ use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrays::BoolArray; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::VarBinViewArrayExt; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; use vortex_array::expr::Expression; use vortex_array::expr::is_root; use vortex_array::expr::not; -use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::ExecutionArgs; @@ -42,25 +38,25 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use super::super::aggregates::BloomFilter; -use super::super::aggregates::bloom_contains; -pub use super::super::aggregates::bloom_filter::BloomOptions; -use super::super::aggregates::i64_value; use super::SkipIndex; +pub use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; +pub use crate::layouts::zoned::aggregates::bloom_filter::BloomOptions; +pub use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; -/// Bloom skipping index for `i64` equality predicates. +/// Bloom skip index for constant-equality predicates. +/// +/// TODO(joacoc): Add documentation about the Bloom skip index +/// and how it works here. #[derive(Clone, Debug, Default)] pub struct BloomSkipIndex { options: BloomOptions, } impl BloomSkipIndex { - /// Create an index with explicit Bloom tuning. pub fn new(options: BloomOptions) -> Self { Self { options } } - /// The persisted Bloom options. pub fn options(&self) -> &BloomOptions { &self.options } @@ -82,7 +78,7 @@ impl SkipIndex for BloomSkipIndex { } } -/// Probe scalar function: test one `i64` literal against each binary Bloom state. +/// Probe scalar function: test one literal against each binary Bloom state. #[derive(Clone, Debug)] struct BloomContains; @@ -90,7 +86,7 @@ impl ScalarFnVTable for BloomContains { type Options = BloomOptions; fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("vortex.bloom_contains.i64.v1"); + static ID: CachedId = CachedId::new("vortex.bloom_contains.v1"); *ID } @@ -106,6 +102,9 @@ impl ScalarFnVTable for BloomContains { Arity::Exact(2) } + /// Only two children are expected. + /// The first child represents the filter as a byte sequence, + /// while the second child represents the literal value to search for (the needle). fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { match child_idx { 0 => ChildName::from("filter"), @@ -120,9 +119,10 @@ impl ScalarFnVTable for BloomContains { "bloom filter must be Binary" ); vortex_ensure!( - matches!(args[1], DType::Primitive(PType::I64, _)), - "bloom needle must be i64" + is_bloom_valid_dtype(&args[1]), + "bloom needle must be bool, primitive, decimal, utf8, binary or extension" ); + Ok(DType::Bool(args[0].nullability() | args[1].nullability())) } @@ -133,37 +133,43 @@ impl ScalarFnVTable for BloomContains { ctx: &mut ExecutionCtx, ) -> VortexResult { let filters = args.get(0)?.execute::(ctx)?; + + // The Bloom search is performed using valid scalars [BloomPartial::contains_valid_scalar]. + // + // If the needle accepts an array of values, e.g., Array[1, 2, 3], + // the following code should be updated. let needle_array = args.get(1)?; let needle = needle_array .as_constant() .ok_or_else(|| vortex_err!("bloom needle must be constant"))?; - let Some(needle) = i64_value(&needle)? else { - return Ok(ConstantArray::new( - Scalar::null(DType::Bool(Nullability::Nullable)), - args.row_count(), - ) - .into_array()); - }; let validity = filters.varbinview_validity(); let valid = validity.execute_mask(filters.len(), ctx)?; + + // Quick return if the needle is invalid. + if !needle.is_valid() { + let possible = vec![false; filters.len()]; + return Ok(BoolArray::new(BitBuffer::from_iter(possible), validity).into_array()); + } + let mut possible = Vec::with_capacity(filters.len()); for (idx, is_valid) in valid.iter().enumerate() { - if is_valid { - let filter = filters.bytes_at(idx); - vortex_ensure!( - filter.len() == options.bytes().get(), - "stored bloom byte length does not match options" - ); - possible.push(bloom_contains( - filter.as_slice(), - needle, - options.hashes().get(), - )); - } else { + if !is_valid { possible.push(false); + continue; } + + let bytes = filters.bytes_at(idx); + let partial = BloomPartial::try_from(bytes.as_slice())?; + + vortex_ensure!( + partial.len() == options.blocks().get(), + "stored bloom length does not match options" + ); + + possible.push(partial.contains_valid_scalar(&needle)?); } + Ok(BoolArray::new(BitBuffer::from_iter(possible), validity).into_array()) } @@ -187,6 +193,9 @@ impl StatsRewriteRule for BloomEqRewrite { Binary.id() } + /// Only works for root literal comparisons and valid [DTypes]. + /// + /// E.g. `eq(root(), lit(5i32))` or `eq(lit(5i32), root())` fn falsify( &self, expr: &Expression, @@ -203,9 +212,8 @@ impl StatsRewriteRule for BloomEqRewrite { } else { return Ok(None); }; - if !matches!(ctx.return_dtype(column)?, DType::Primitive(PType::I64, _)) - || literal.as_::().is_null() - { + + if !is_bloom_valid_dtype(&ctx.return_dtype(column)?) || literal.as_::().is_null() { return Ok(None); } @@ -215,38 +223,64 @@ impl StatsRewriteRule for BloomEqRewrite { } } +/// Returns true if the type is valid for the bloom index to acc/contain. +/// +/// This is defined by the available implementations in +/// [crate::layouts::zoned::aggregates::bloom::constant] and +/// [crate::layouts::zoned::aggregates::bloom::canonical] +pub(in crate::layouts::zoned) fn is_bloom_valid_dtype(dtype: &DType) -> bool { + match dtype { + DType::Extension(ext) => is_bloom_valid_dtype(ext.storage_dtype()), + DType::Bool(_) + | DType::Primitive(..) + | DType::Decimal(..) + | DType::Utf8(_) + | DType::Binary(_) => true, + _ => false, + } +} + #[cfg(test)] mod tests { - use std::num::NonZeroU8; - use std::num::NonZeroUsize; use std::sync::Arc; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::Columnar; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateFnVTable; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::expr::Expression; use vortex_array::expr::eq; + use vortex_array::expr::gt_eq; use vortex_array::expr::lit; use vortex_array::expr::root; + use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; + use vortex_buffer::buffer; use vortex_error::VortexResult; - use super::BloomOptions; use super::BloomSkipIndex; use super::SkipIndex; + use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; + use crate::layouts::zoned::aggregates::bloom_filter::BloomOptions; use crate::layouts::zoned::zone_map::ZoneMap; - - fn small_options() -> BloomOptions { - BloomOptions::new( - NonZeroUsize::new(64).expect("64 is non-zero"), - NonZeroU8::new(3).expect("3 is non-zero"), - ) - } + use crate::test::SESSION; #[test] fn missing_stat_stays_inconclusive() -> VortexResult<()> { let session = vortex_array::array_session(); - let index = BloomSkipIndex::new(small_options()); + let index = BloomSkipIndex::new(BloomOptions::default()); index.register(&session); let predicate = eq(root(), lit(42i64)); let proof = predicate @@ -266,4 +300,85 @@ mod tests { assert!(zone_map.prune(&proof, &session)?.all_false()); Ok(()) } + + /// Similar zone map tests as the ones in [crate::layouts::zoned::tests] + /// but using BloomFilter rather than max/min zones. + fn build_bloom_zone_map(dtype: DType, batch: ArrayRef) -> ZoneMap { + let bloom = BloomFilter; + let options = BloomOptions::default(); + + // If index is not registered there will be no warning, + // but the index will return false for everything. + // + // (joacoc) should be considered a warning for missing aggregatefns? + BloomSkipIndex::new(options.clone()).register(&SESSION); + let mut ctx = SESSION.create_execution_ctx(); + + let mut zone_filter = bloom.empty_partial(&options, &dtype).unwrap(); + bloom + .accumulate( + &mut zone_filter, + &Columnar::Canonical(batch.execute::(&mut ctx).unwrap()), + &mut ctx, + ) + .unwrap(); + + let zone_filter_as_scalar = bloom.to_scalar(&zone_filter).unwrap(); + let zone_filter_as_bytes = zone_filter_as_scalar.as_binary().value().unwrap().to_vec(); + let zone_filter_as_varbin = + VarBinArray::from_nullable_bytes(vec![Some(zone_filter_as_bytes.as_slice())]); + + let bloom = BloomFilter.bind(options); + let zone_filter_struct = StructArray::from_fields(&[( + bloom.clone().to_string(), + zone_filter_as_varbin.into_array(), + )]) + .unwrap(); + + ZoneMap::try_new(dtype, zone_filter_struct, Arc::new([bloom]), 1, 10).unwrap() + } + + fn assert_prune(zone_map: &ZoneMap, dtype: &DType, expr: Expression, expected: [bool; 1]) { + let mut ctx = SESSION.create_execution_ctx(); + let pruning_expr = expr.falsify(dtype, &SESSION).unwrap().unwrap(); + let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); + assert_arrays_eq!(mask.into_array(), BoolArray::from_iter(expected), &mut ctx); + } + + #[rstest] + #[case::equals_value_not_in_batch(eq(root(), lit(99i32)), [true])] + #[case::equals_value_in_batch(eq(root(), lit(5i32)), [false])] + #[case::gt_eq_not_supported_by_bloom(gt_eq(root(), lit(4i32)), [false])] + #[case::null_never_prunes( + gt_eq(root(), lit(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)))), + [false] + )] + fn test_zone_map_prunes_with_bloom_filter_i32( + #[case] expr: Expression, + #[case] expected: [bool; 1], + ) -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::Nullable); + let batch = PrimitiveArray::new(buffer![5i32, 6i32, 7i32], Validity::AllValid).into_array(); + let zone_map = build_bloom_zone_map(dtype.clone(), batch); + assert_prune(&zone_map, &dtype, expr, expected); + Ok(()) + } + + #[rstest] + #[case::equals_value_not_in_batch(eq(root(), lit("zz")), [true])] + #[case::equals_value_in_batch(eq(root(), lit("london")), [false])] + fn test_zone_map_prunes_with_bloom_filter_varbin( + #[case] expr: Expression, + #[case] expected: [bool; 1], + ) -> VortexResult<()> { + let dtype = DType::Utf8(Nullability::NonNullable); + let batch = VarBinArray::from_iter( + [Some("london"), Some("hamburg"), Some("newyork")], + dtype.clone(), + ) + .into_array(); + let zone_map = build_bloom_zone_map(dtype.clone(), batch); + assert_prune(&zone_map, &dtype, expr, expected); + Ok(()) + } } From 8d81b781e1250bafbae006d9cb13d1e8b0722610 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Wed, 12 Aug 2026 22:51:00 +0200 Subject: [PATCH 3/7] (doc) update documentation comments --- .../zoned/aggregates/bloom_filter/mod.rs | 21 +++++++------------ .../zoned/aggregates/bloom_filter/partial.rs | 2 +- .../src/layouts/zoned/skip_index/bloom.rs | 2 -- 3 files changed, 8 insertions(+), 17 deletions(-) diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index d7fd32b7e3a..408be9d0c5a 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -31,18 +31,11 @@ pub use partial::BloomPartial; use crate::layouts::zoned::skip_index::bloom::is_bloom_valid_dtype; -// (joacoc) Opted for blocks_count as the main way to tune -// the Bloom filter, but there are also other ways, like -// setting the false-positive probability (FPP) and -// number of distinct values (NDV). Given the statistics -// collected from a random subset of values, they could be -// used to determine the number of blocks automatically, -// which I think is more aligned with how Vortex works with -// encoders. But for now, this is a simple approach that -// pushes the block-count selection back to the user -// and otherwise defaults to [DEFAULT_BLOCKS_COUNT]. -// -// In case that the blocks count stays as is, +// 1. (joacoc) Opted for blocks_count as a simpler way to tune +// the Bloom filter, even though there are other ways. +// 2. (joacoc) I think the optimal could be using statistics, similar +// to how vortex encoder selection works. +// 3. (joacoc) In case that the tune/options stays as is, // a guide on how to select the number of blocks // would be great. @@ -51,12 +44,12 @@ use crate::layouts::zoned::skip_index::bloom::is_bloom_valid_dtype; pub struct BloomOptions { /// Number of blocks in the split block Bloom filter (SBBF). /// + /// Defaults to: [DEFAULT_BLOCKS_COUNT]. + /// /// The filter is partitioned into 256-bit blocks. More blocks reduce the /// number of distinct values assigned to each block, reducing the /// false-positive rate at the cost of increased filter size. /// - /// Defaults to: [DEFAULT_BLOCKS_COUNT]. - /// /// ### Block size and memory usage /// /// As a reference, you can use the following table diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs index ead7f2224e7..232078b2ace 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs @@ -60,7 +60,7 @@ impl BloomPartial { } } - /// Check whether a hash is (probably) present in the filter. + /// Checks whether a hash is (probably) present in the filter. pub(super) fn find_hash(&self, hash: u64) -> bool { let idx = self.block_index(hash, self.blocks.len()) as usize; let mask = self.make_mask(hash as u32); diff --git a/vortex-layout/src/layouts/zoned/skip_index/bloom.rs b/vortex-layout/src/layouts/zoned/skip_index/bloom.rs index 71c8fb29cc9..55f98126959 100644 --- a/vortex-layout/src/layouts/zoned/skip_index/bloom.rs +++ b/vortex-layout/src/layouts/zoned/skip_index/bloom.rs @@ -134,8 +134,6 @@ impl ScalarFnVTable for BloomContains { ) -> VortexResult { let filters = args.get(0)?.execute::(ctx)?; - // The Bloom search is performed using valid scalars [BloomPartial::contains_valid_scalar]. - // // If the needle accepts an array of values, e.g., Array[1, 2, 3], // the following code should be updated. let needle_array = args.get(1)?; From 7d0495faf53808ed9bc5f14b708bf47fa9209521 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Thu, 13 Aug 2026 13:03:40 +0200 Subject: [PATCH 4/7] Move hash/insert/contains scalar into partial.rs, add more documentation comments and fix primitive test for validation --- .../aggregates/bloom_filter/canonical/bool.rs | 1 + .../bloom_filter/canonical/extension.rs | 1 + .../bloom_filter/canonical/primitive.rs | 106 +++++++------ .../zoned/aggregates/bloom_filter/constant.rs | 95 +----------- .../zoned/aggregates/bloom_filter/mod.rs | 9 +- .../zoned/aggregates/bloom_filter/partial.rs | 140 ++++++++++++++++-- 6 files changed, 197 insertions(+), 155 deletions(-) diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs index 747fd353457..7d85d0c687d 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs @@ -109,6 +109,7 @@ mod tests { bloom_filter.contains_valid_scalar(&Scalar::bool(false, nullability))?, expect_false ); + Ok(()) } } diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs index 08762093e86..7ad21137ff8 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs @@ -16,6 +16,7 @@ pub(super) fn accumulate_extension( ) -> VortexResult<()> { let storage = array.storage_array().clone(); let canonical = storage.execute::(ctx)?; + super::accumulate_canonical(&canonical, partial, ctx) } diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs index 24b2418ffa0..dd5308bffa6 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs @@ -170,61 +170,75 @@ mod tests { Ok(()) } + /// The following three test will check if the validity cases + /// are correctly implemented for the three branches. #[rstest] #[case(&[10i8, 20, 30, 40, 50])] - fn validity(#[case] present: &[T]) -> VortexResult<()> + fn validity_all_true(#[case] present: &[T]) -> VortexResult<()> where T: Copy + NativePType + Into, { - { - let ctx = setup()?; - let all_valid = PrimitiveArray::from_option_iter(present.iter().map(|&v| Some(v))); - let bloom_filter = build_filter( - all_valid.into_array(), - DType::Primitive(T::PTYPE, Nullability::Nullable), - ctx, - )?; - for &v in present { - let scalar = Scalar::primitive(v, Nullability::Nullable); - assert!(bloom_filter.contains_valid_scalar(&scalar)?); - } + let ctx = setup()?; + let all_valid = PrimitiveArray::from_option_iter(present.iter().map(|&v| Some(v))); + let bloom_filter = build_filter( + all_valid.into_array(), + DType::Primitive(T::PTYPE, Nullability::Nullable), + ctx, + )?; + + for &v in present { + let scalar = Scalar::primitive(v, Nullability::Nullable); + assert!(bloom_filter.contains_valid_scalar(&scalar)?); } - // Mask::AllFalse: every position null. Nothing should ever be found — including the - // `T::default()` filler that `from_option_iter` writes into null slots internally. - { - let ctx = setup()?; - let all_invalid = PrimitiveArray::from_option_iter(present.iter().map(|_| None::)); - let bloom_filter = build_filter( - all_invalid.into_array(), - DType::Primitive(T::PTYPE, Nullability::Nullable), - ctx, - )?; - for &v in present { - let scalar = Scalar::primitive(v, Nullability::Nullable); - assert!(!bloom_filter.contains_valid_scalar(&scalar)?); - } + Ok(()) + } + + #[rstest] + #[case(&[10i8, 20, 30, 40, 50])] + fn validity_all_false(#[case] present: &[T]) -> VortexResult<()> + where + T: Copy + NativePType + Into, + { + let ctx = setup()?; + let all_invalid = PrimitiveArray::from_option_iter(present.iter().map(|_| None::)); + let bloom_filter = build_filter( + all_invalid.into_array(), + DType::Primitive(T::PTYPE, Nullability::Nullable), + ctx, + )?; + + for &v in present { + let scalar = Scalar::primitive(v, Nullability::Nullable); + assert!(!bloom_filter.contains_valid_scalar(&scalar)?); } - // Mask::Values: alternating valid/null. Only the valid positions should be members, - // and the default filler at null slots must not leak through as a false member. - { - let ctx = setup()?; - let mixed: Vec> = present - .iter() - .enumerate() - .map(|(i, &v)| if i % 2 == 0 { Some(v) } else { None }) - .collect(); - let bloom_filter = build_filter( - PrimitiveArray::from_option_iter(mixed).into_array(), - DType::Primitive(T::PTYPE, Nullability::Nullable), - ctx, - )?; - for (i, &v) in present.iter().enumerate() { - if i % 2 == 0 { - let scalar = Scalar::primitive(v, Nullability::Nullable); - assert!(bloom_filter.contains_valid_scalar(&scalar)?); - } + Ok(()) + } + + #[rstest] + #[case(&[10i8, 20, 30, 40, 50])] + fn validity_mixed(#[case] present: &[T]) -> VortexResult<()> + where + T: Copy + NativePType + Into, + { + let ctx = setup()?; + let mixed: Vec> = present + .iter() + .enumerate() + .map(|(i, &v)| if i % 2 == 0 { Some(v) } else { None }) + .collect(); + + let bloom_filter = build_filter( + PrimitiveArray::from_option_iter(mixed).into_array(), + DType::Primitive(T::PTYPE, Nullability::Nullable), + ctx, + )?; + + for (i, &v) in present.iter().enumerate() { + if i % 2 == 0 { + let scalar = Scalar::primitive(v, Nullability::Nullable); + assert!(bloom_filter.contains_valid_scalar(&scalar)?); } } diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs index 34fc9a4469b..0100e36071f 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs @@ -2,15 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::arrays::ConstantArray; -use vortex_array::dtype::DType; -use vortex_array::dtype::PType; -use vortex_array::match_each_float_ptype; -use vortex_array::match_each_integer_ptype; -use vortex_array::scalar::DecimalValue; -use vortex_array::scalar::Scalar; -use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_err; use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; @@ -25,92 +17,9 @@ pub(super) fn accumulate_constant( return Ok(()); } - partial.insert_hash(partial.hash_valid_scalar(scalar)?); - Ok(()) -} - -impl BloomPartial { - /// Scalar values must be valid otherwise the function will raise an err. - /// - /// This function is used by both, for accumulating scalars, - /// but also to get a scalar value membership. - pub(in crate::layouts::zoned) fn hash_valid_scalar( - &self, - scalar: &Scalar, - ) -> VortexResult { - if scalar.is_null() { - return Err(vortex_err!("cannot hash invalid scalars in bloom filter")); - } - - Ok(match scalar.dtype() { - DType::Extension(_) => { - self.hash_valid_scalar(&scalar.as_extension().to_storage_scalar())? - } - DType::Bool(_) => self.hash( - scalar - .as_bool() - .value() - .vortex_expect("non-null boolean value"), - ), - DType::Primitive(ptype, _) => match ptype { - PType::F16 | PType::F32 | PType::F64 => { - match_each_float_ptype!(ptype, |T| { - let value = scalar - .as_primitive() - .typed_value::() - .vortex_expect("non-null primitive value"); - self.hash(value.to_bits()) - }) - } - _ => match_each_integer_ptype!(ptype, |T| { - let value = scalar - .as_primitive() - .typed_value::() - .vortex_expect("non-null primitive value"); - self.hash(value) - }), - }, - DType::Decimal(..) => { - let decimal = scalar - .as_decimal() - .decimal_value() - .vortex_expect("non-null decimal value"); - match decimal { - DecimalValue::I8(v) => self.hash(v), - DecimalValue::I16(v) => self.hash(v), - DecimalValue::I32(v) => self.hash(v), - DecimalValue::I64(v) => self.hash(v), - DecimalValue::I128(v) => self.hash(v), - DecimalValue::I256(v) => self.hash(v), - } - } - DType::Utf8(_) => { - let buffer = scalar - .as_utf8() - .value() - .vortex_expect("non-null utf8 value"); - self.hash(buffer.as_bytes()) - } - DType::Binary(_) => { - let buffer = scalar - .as_binary() - .value() - .vortex_expect("non-null binary value"); - self.hash(buffer.as_slice()) - } - other => { - return Err(vortex_err!("bloom filter does not support dtype {other}")); - } - }) - } + partial.insert_valid_scalar(scalar)?; - pub(in crate::layouts::zoned) fn contains_valid_scalar( - &self, - scalar: &Scalar, - ) -> VortexResult { - let hash = self.hash_valid_scalar(scalar)?; - Ok(self.find_hash(hash)) - } + Ok(()) } #[cfg(test)] diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index 408be9d0c5a..533420375f6 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -363,16 +363,11 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { ); for i in 0..100i64 { - let hash = partial.hash(i); - assert!(partial.find_hash(hash), "value {i} missing after merge"); + assert!(partial.contains(i), "value {i} missing after merge"); } for i in 101..200i64 { - let hash = partial.hash(i); - assert!( - !partial.find_hash(hash), - "value {i} shouldn't be present after" - ); + assert!(!partial.contains(i), "value {i} shouldn't be present after"); } Ok(()) diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs index 232078b2ace..363216042b1 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs @@ -1,4 +1,4 @@ -//! Split block Bloom filters implementation for vortex. +//! Split block Bloom filters (SBBF) implementation for vortex. //! //! [Split block Bloom filters]: https://arxiv.org/pdf/2101.01719 @@ -9,16 +9,33 @@ use std::hash::Hash; use std::hash::Hasher; use twox_hash::XxHash3_64; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::match_each_float_ptype; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +/// Block size in bits (8 * 4 = 32 bits) const BLOCK_SIZE: usize = 8 * size_of::(); +/// Represents a Split block Bloom Filter filter for a single layout zone. pub struct BloomPartial { pub(super) blocks: Vec<[u32; 8]>, } +/// The following Split block Bloom filter (SBBF) implementation +/// is a translation of the original paper's names and values, +/// with slight changes to let the Rust compiler generate +/// optimized, vectorized code for `make_mask`, `add_hash`, and `find_hash`. impl BloomPartial { + /// Returns the blocks len. + /// + /// Matches [BloomOptions::blocks_count] #[inline] pub fn len(&self) -> usize { self.blocks.len() @@ -34,12 +51,12 @@ impl BloomPartial { } #[inline] - pub(super) fn insert_hash(&mut self, hash: u64) { + fn insert_hash(&mut self, hash: u64) { self.add_hash(hash); } #[inline] - pub(super) fn hash(&self, value: T) -> u64 + fn hash(&self, value: T) -> u64 where T: Hash, { @@ -50,10 +67,9 @@ impl BloomPartial { hasher.finish() } - /// Hash should be u64 or u32? fn add_hash(&mut self, hash: u64) { let idx = self.block_index(hash, self.blocks.len()) as usize; - // Block idx already consumed the hash, + let mask = self.make_mask(hash as u32); for i in 0..8 { self.blocks[idx][i] |= mask[i]; @@ -61,7 +77,7 @@ impl BloomPartial { } /// Checks whether a hash is (probably) present in the filter. - pub(super) fn find_hash(&self, hash: u64) -> bool { + fn find_hash(&self, hash: u64) -> bool { let idx = self.block_index(hash, self.blocks.len()) as usize; let mask = self.make_mask(hash as u32); @@ -76,9 +92,6 @@ impl BloomPartial { /// Takes a hash value and creates a mask with one bit set in each 32-bit lane. /// These are the bits to set or check when accessing the block. - /// - /// The following code is SIMD friendly and will get vectorized - /// by the compiler automatically (for releases). fn make_mask(&self, hash: u32) -> [u32; 8] { let mut out = [0u32; 8]; @@ -106,6 +119,115 @@ impl BloomPartial { } } +/// The following implementation provides a simpler access for scalars. +impl BloomPartial { + /// Returns the hash of the scalar's underlying value. + /// Returns an error if the [Scalar] is invalid or its [DType] is unsupported. + /// + /// For example, `Scalar(Primitive(I32(54)))` is hashed as `hash(54)`. + pub(in crate::layouts::zoned) fn hash_valid_scalar( + &self, + scalar: &Scalar, + ) -> VortexResult { + if scalar.is_null() { + return Err(vortex_err!("cannot hash invalid scalars in bloom filter")); + } + + Ok(match scalar.dtype() { + DType::Extension(_) => { + self.hash_valid_scalar(&scalar.as_extension().to_storage_scalar())? + } + DType::Bool(_) => self.hash( + scalar + .as_bool() + .value() + .vortex_expect("non-null boolean value"), + ), + DType::Primitive(ptype, _) => match ptype { + PType::F16 | PType::F32 | PType::F64 => { + match_each_float_ptype!(ptype, |T| { + let value = scalar + .as_primitive() + .typed_value::() + .vortex_expect("non-null primitive value"); + self.hash(value.to_bits()) + }) + } + _ => match_each_integer_ptype!(ptype, |T| { + let value = scalar + .as_primitive() + .typed_value::() + .vortex_expect("non-null primitive value"); + self.hash(value) + }), + }, + DType::Decimal(..) => { + let decimal = scalar + .as_decimal() + .decimal_value() + .vortex_expect("non-null decimal value"); + match decimal { + DecimalValue::I8(v) => self.hash(v), + DecimalValue::I16(v) => self.hash(v), + DecimalValue::I32(v) => self.hash(v), + DecimalValue::I64(v) => self.hash(v), + DecimalValue::I128(v) => self.hash(v), + DecimalValue::I256(v) => self.hash(v), + } + } + DType::Utf8(_) => { + let buffer = scalar + .as_utf8() + .value() + .vortex_expect("non-null utf8 value"); + self.hash(buffer.as_bytes()) + } + DType::Binary(_) => { + let buffer = scalar + .as_binary() + .value() + .vortex_expect("non-null binary value"); + self.hash(buffer.as_slice()) + } + other => { + return Err(vortex_err!("bloom filter does not support dtype {other}")); + } + }) + } + + /// Returns true if the underlying value of a [Scalar] may be present. + /// Returns an error if the [Scalar] is invalid or its [DType] is unsupported. + pub(in crate::layouts::zoned) fn contains_valid_scalar( + &self, + scalar: &Scalar, + ) -> VortexResult { + let hash = self.hash_valid_scalar(scalar)?; + Ok(self.find_hash(hash)) + } + + /// Inserts the underlying value of a [Scalar] if it is valid. + /// Returns an error if the [Scalar] is invalid or its [DType] is unsupported. + pub(in crate::layouts::zoned) fn insert_valid_scalar( + &mut self, + scalar: &Scalar, + ) -> VortexResult<()> { + let hash = self.hash_valid_scalar(scalar)?; + Ok(self.insert_hash(hash)) + } +} + +#[cfg(test)] +impl BloomPartial { + #[inline] + pub(super) fn contains(&self, value: T) -> bool + where + T: Hash, + { + let hash = self.hash(value); + self.find_hash(hash) + } +} + #[cfg(test)] impl From> for BloomPartial { fn from(value: Vec<[u32; 8]>) -> Self { From 43b064dd446de86d56ccbed7afde4af82fd51b7d Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Thu, 13 Aug 2026 14:02:11 +0200 Subject: [PATCH 5/7] (improvement) remove unwrap from try_from when deserializing partial from bytes, and use constant for block_size --- .../layouts/zoned/aggregates/bloom_filter/mod.rs | 13 +++++++------ .../zoned/aggregates/bloom_filter/partial.rs | 13 ++++++++----- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index 533420375f6..f7bd4b8cbb5 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -17,6 +17,7 @@ use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; @@ -29,6 +30,7 @@ mod partial; pub(in crate::layouts::zoned) mod constant; pub use partial::BloomPartial; +use crate::layouts::zoned::aggregates::bloom_filter::partial::BLOCK_SIZE; use crate::layouts::zoned::skip_index::bloom::is_bloom_valid_dtype; // 1. (joacoc) Opted for blocks_count as a simpler way to tune @@ -82,7 +84,8 @@ const DEFAULT_BLOCKS_COUNT: usize = 256; impl Default for BloomOptions { fn default() -> Self { Self { - blocks_count: NonZeroUsize::new(DEFAULT_BLOCKS_COUNT).expect("valid blocks size"), + blocks_count: NonZeroUsize::new(DEFAULT_BLOCKS_COUNT) + .vortex_expect("valid blocks size"), } } } @@ -184,7 +187,7 @@ impl AggregateFnVTable for BloomFilter { /// /// Basically turns each block into a single byte sequence. fn to_scalar(&self, partial: &Self::Partial) -> VortexResult { - let mut bytes = Vec::with_capacity(partial.blocks.len() * 8 * size_of::()); + let mut bytes = Vec::with_capacity(partial.blocks.len() * BLOCK_SIZE); bytes.extend( partial .blocks @@ -239,6 +242,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { use vortex_error::vortex_ensure; use super::*; + use crate::layouts::zoned::aggregates::bloom_filter::partial::BLOCK_SIZE; pub fn setup() -> VortexResult { let session = vortex_array::array_session(); @@ -255,10 +259,7 @@ pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { pub fn extract_bloom_blocks(state: &Scalar) -> VortexResult> { let bytes = state.as_binary().value().expect("bloom state is non-null"); - vortex_ensure!( - bytes.len() % (8 * size_of::()) == 0, - "invalid bloom state length" - ); + vortex_ensure!(bytes.len() % BLOCK_SIZE == 0, "invalid bloom state length"); let mut blocks = Vec::with_capacity(bytes.len() / 32); for block_bytes in bytes.chunks_exact(32) { let mut block = [0u32; 8]; diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs index 363216042b1..6684ddab8f0 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs @@ -21,7 +21,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; /// Block size in bits (8 * 4 = 32 bits) -const BLOCK_SIZE: usize = 8 * size_of::(); +pub(super) const BLOCK_SIZE: usize = 8 * size_of::(); /// Represents a Split block Bloom Filter filter for a single layout zone. pub struct BloomPartial { @@ -251,12 +251,15 @@ impl TryFrom<&[u8]> for BloomPartial { .chunks_exact(BLOCK_SIZE) .map(|chunk| { let mut block = [0u32; 8]; - for (word, wb) in block.iter_mut().zip(chunk.chunks_exact(4)) { - *word = u32::from_le_bytes(wb.try_into().unwrap()); + for (lane, lane_bytes) in block.iter_mut().zip(chunk.chunks_exact(4)) { + *lane = u32::from_le_bytes(lane_bytes.try_into().map_err(|_| { + vortex_err!("invalid bloom filter word length: {}", lane_bytes.len()) + })?); } - block + Ok(block) }) - .collect(); + .collect::>>()?; + Ok(BloomPartial { blocks }) } } From f6ee474edbcd5aa9e90015aef3e7733e8b7f9950 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Fri, 14 Aug 2026 16:18:12 +0200 Subject: [PATCH 6/7] Restore the writer and file strategy. Another pull request will handle those cases. Also remove aggregates/min_max.rs and aggregates/mod.rs, and re-add them once the integration is done (they're currently in the writer) so the skip index integration PR should re-add them. This commit also adds a new composable type to the bloom filter (map) but yet not supported --- vortex-file/src/strategy.rs | 72 ++-- vortex-file/tests/bloom_skip_index.rs | 243 ----------- .../aggregates/bloom_filter/canonical/mod.rs | 3 +- .../zoned/aggregates/bloom_filter/mod.rs | 27 +- .../src/layouts/zoned/aggregates/min_max.rs | 70 ---- .../src/layouts/zoned/aggregates/mod.rs | 87 ---- vortex-layout/src/layouts/zoned/mod.rs | 1 - .../src/layouts/zoned/skip_index/bloom.rs | 382 ------------------ .../src/layouts/zoned/skip_index/mod.rs | 59 --- vortex-layout/src/layouts/zoned/writer.rs | 3 +- 10 files changed, 46 insertions(+), 901 deletions(-) delete mode 100644 vortex-file/tests/bloom_skip_index.rs delete mode 100644 vortex-layout/src/layouts/zoned/aggregates/min_max.rs delete mode 100644 vortex-layout/src/layouts/zoned/skip_index/bloom.rs delete mode 100644 vortex-layout/src/layouts/zoned/skip_index/mod.rs diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 942ad91bc3a..9d4dbb90610 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -59,7 +59,6 @@ pub struct WriteStrategyBuilder { row_block_size: usize, data_block_target_bytes: Option, field_writers: HashMap>, - field_zoned_options: HashMap, allow_encodings: Option>, flat_strategy: Option>, probe_compressor: Option>, @@ -79,7 +78,6 @@ impl Default for WriteStrategyBuilder { data_block_target_bytes: Some(ONE_MEG), field_writers: HashMap::new(), allow_encodings: None, - field_zoned_options: HashMap::new(), flat_strategy: None, probe_compressor: None, use_list_layout: use_experimental_list_layout(), @@ -130,21 +128,7 @@ impl WriteStrategyBuilder { self } - /// Override only the zoned-statistics options for a field while retaining the default - /// repartitioning, dictionary, compression, buffering, and flat-layout pipeline. - /// - /// This can attach custom per-zone aggregates without changing the physical data strategy for - /// the field. - pub fn with_field_zoned_options( - mut self, - field: impl Into, - options: ZonedLayoutOptions, - ) -> Self { - self.field_zoned_options.insert(field.into(), options); - self - } - - /// Override the allowed array encodings for normalization. + /// Override the allowed array encodings for file writing. /// /// The configured flat leaf strategy is wrapped in a [`LayoutStrategyEncodingValidator`] /// that recursively checks every chunk before passing it to the leaf writer. [`build`](Self::build) @@ -191,7 +175,7 @@ impl WriteStrategyBuilder { /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides /// applied. - pub fn build(mut self) -> Arc { + pub fn build(self) -> Arc { let flat: Arc = if let Some(flat) = self.flat_strategy { flat } else { @@ -277,40 +261,36 @@ impl WriteStrategyBuilder { let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"); - let column_writer = |options: ZonedLayoutOptions| -> Arc { - // 2. calculate stats for each row group - let stats = - ZonedStrategy::new(dict.clone(), compress_then_flat.clone(), options.clone()); - - // 1. repartition each column to fixed row counts - Arc::new(RepartitionStrategy::new( - stats, - RepartitionWriterOptions { - // No minimum block size in bytes - block_size_minimum: 0, - block_len_multiple: options.block_size.get(), - block_size_target: None, - canonicalize: false, - }, - )) - }; - let repartition = column_writer(ZonedLayoutOptions { - block_size: row_block_size, - ..Default::default() - }); + // 2. calculate stats for each row group + let stats = ZonedStrategy::new( + dict, + compress_then_flat.clone(), + ZonedLayoutOptions { + block_size: row_block_size, + ..Default::default() + }, + ); - for (field, options) in self.field_zoned_options { - self.field_writers - .entry(field) - .or_insert_with(|| column_writer(options)); - } + // 1. repartition each column to fixed row counts + let repartition = RepartitionStrategy::new( + stats, + RepartitionWriterOptions { + // No minimum block size in bytes + block_size_minimum: 0, + // Always repartition into 8K row blocks + block_len_multiple: self.row_block_size, + block_size_target: None, + canonicalize: false, + }, + ); // 0. start with splitting columns let validity_strategy = CollectStrategy::new(compress_then_flat.clone()); // Take any field overrides from the builder and apply them to the final strategy. - let mut table_strategy = TableStrategy::new(Arc::new(validity_strategy), repartition) - .with_field_writers(self.field_writers); + let mut table_strategy = + TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) + .with_field_writers(self.field_writers); if self.use_list_layout { // We need a closure here to enable recursive application of list layout. diff --git a/vortex-file/tests/bloom_skip_index.rs b/vortex-file/tests/bloom_skip_index.rs deleted file mode 100644 index 6b8371a5e4b..00000000000 --- a/vortex-file/tests/bloom_skip_index.rs +++ /dev/null @@ -1,243 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! End-to-end coverage for the zoned Bloom skipping index. -//! -//! Bloom indexes are optional extensions rather than part of the default file layout. This test -//! exercises the complete opt-in lifecycle: -//! -//! 1. register the index with a write session and request it for one field; -//! 2. persist one Bloom filter per zone and reopen the file with a fresh registered session; -//! 3. prove that equality predicates prune zones while returning the same rows as a full scan; and -//! 4. reopen the indexed file with an unregistered, allow-unknown session to verify that the index -//! is ignorable. -//! -//! The input is intentionally hostile to ordinary min/max pruning. Zone `z` contains values whose -//! remainder modulo [`NZONES`] is `z`, so both [`HIT`] and [`MISS`] lie inside every zone's -//! min/max range. `MISS` is then removed from its zone without changing that range. Consequently, -//! pruning either value requires the Bloom filter rather than the built-in range statistics. - -#![expect(clippy::expect_used)] - -use std::num::NonZeroUsize; -use std::sync::Arc; - -use vortex_array::ArrayRef; -use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; -use vortex_array::arrays::ChunkedArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::StructArray; -use vortex_array::assert_arrays_eq; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::expr::Expression; -use vortex_array::expr::eq; -use vortex_array::expr::get_item; -use vortex_array::expr::lit; -use vortex_array::expr::root; -use vortex_array::field_path; -use vortex_array::stream::ArrayStreamExt; -use vortex_error::VortexResult; -use vortex_file::OpenOptionsSessionExt; -use vortex_file::WriteOptionsSessionExt; -use vortex_file::WriteStrategyBuilder; -use vortex_io::session::RuntimeSession; -use vortex_layout::LayoutStrategy; -use vortex_layout::layouts::zoned::skip_index::SkipIndex; -use vortex_layout::layouts::zoned::skip_index::bloom::BloomOptions; -use vortex_layout::layouts::zoned::skip_index::bloom::BloomSkipIndex; -use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; -use vortex_layout::session::LayoutSession; -use vortex_mask::Mask; -use vortex_session::VortexSession; - -const ZONE_LEN: usize = 256; -const NZONES: usize = 4; -const HIT: i64 = 502; -const MISS: i64 = 503; - -fn bloom() -> BloomSkipIndex { - // A deliberately roomy filter keeps this correctness test's false-positive outcome - // deterministic. False positives are valid Bloom behavior, but false negatives are not. - BloomSkipIndex::new(BloomOptions::default()) -} - -fn session(index: Option<&dyn SkipIndex>) -> VortexSession { - let session = vortex_array::array_session() - .with::() - .with::(); - vortex_file::register_default_encodings(&session); - - // Registration installs the persisted aggregate, membership probe, and equality rewrite. - // Callers must do this independently for the sessions that write and read an indexed file. - if let Some(index) = index { - index.register(&session); - } - session -} - -fn data() -> ArrayRef { - data_with_shape(ZONE_LEN, NZONES, Some(MISS)) -} - -fn data_with_shape(zone_len: usize, nzones: usize, missing: Option) -> ArrayRef { - let chunks = (0..nzones) - .map(|zone| { - let mut values = (0..zone_len) - .map(|row| i64::try_from(row * nzones + zone).expect("test value fits i64")) - .collect::>(); - if let Some(missing) = missing - && usize::try_from(missing).expect("missing value is non-negative") % nzones == zone - { - // Leave a hole inside every zone's min/max range so a MISS cannot be pruned by the - // ordinary range stats. The bloom must provide the proof. - values[usize::try_from(missing).expect("missing value is non-negative") / nzones] = - i64::try_from(zone_len * nzones + zone).expect("replacement fits i64"); - } - StructArray::from_fields(&[("id", PrimitiveArray::from_iter(values).into_array())]) - .expect("valid test struct") - .into_array() - }) - .collect::>(); - ChunkedArray::try_new( - chunks, - DType::struct_( - [("id", DType::Primitive(PType::I64, Nullability::NonNullable))], - Nullability::NonNullable, - ), - ) - .expect("valid chunked test data") - .into_array() -} - -fn filter(value: i64) -> Expression { - eq(get_item("id", root()), lit(value)) -} - -fn strategy( - session: &VortexSession, - index: Option<&dyn SkipIndex>, - zone_len: usize, -) -> VortexResult> { - let mut options = ZonedLayoutOptions { - block_size: NonZeroUsize::new(zone_len).expect("zone length is non-zero"), - ..Default::default() - }; - if let Some(index) = index { - // Adding the aggregate to these field-specific zoned options is the explicit write-side - // opt-in. Registering the index in the session alone does not change the file layout. - options = options.with_skip_index(index, &PType::I64.into(), session)?; - } - Ok(WriteStrategyBuilder::default() - .with_field_zoned_options(field_path!(id), options) - .build()) -} - -async fn scan(file: &vortex_file::VortexFile, value: i64) -> VortexResult { - file.scan()? - .with_filter(filter(value)) - .into_array_stream()? - .read_all() - .await -} - -async fn write_file( - session: &VortexSession, - input: &ArrayRef, - index: Option<&dyn SkipIndex>, - zone_len: usize, -) -> VortexResult> { - let mut bytes = Vec::new(); - session - .write_options() - .with_strategy(strategy(session, index, zone_len)?) - .write(&mut bytes, input.to_array_stream()) - .await?; - Ok(bytes) -} - -#[expect(clippy::tests_outside_test_module)] -#[tokio::test] -async fn bloom_roundtrip_prunes_and_unknown_reader_matches_full_scan() -> VortexResult<()> { - let index = bloom(); - let write_session = session(Some(&index)); - let input = data(); - let bytes = write_file(&write_session, &input, Some(&index), ZONE_LEN).await?; - - // Reconstruct every read-side extension from a fresh session rather than accidentally relying - // on state retained by the writer. - let read_session = session(Some(&index)); - let file = read_session.open_options().open_buffer(bytes.clone())?; - let reader = file.layout_reader()?; - let row_count = file.row_count(); - - // HIT is present only in zone 2. Since it falls within every zone's min/max range, the exact - // one-zone mask proves that the Bloom falsifier participated in pruning. - let hit_mask = reader - .pruning_evaluation( - &(0..row_count), - &filter(HIT), - Mask::new_true(usize::try_from(row_count)?), - )? - .await?; - assert_eq!(hit_mask.true_count(), ZONE_LEN); - assert!(hit_mask.iter().take(2 * ZONE_LEN).all(|keep| !keep)); - assert!( - hit_mask - .iter() - .skip(2 * ZONE_LEN) - .take(ZONE_LEN) - .all(|keep| keep) - ); - assert!(hit_mask.iter().skip(3 * ZONE_LEN).all(|keep| !keep)); - - // MISS was removed while remaining inside every zone's min/max range. Only the Bloom filters - // can prove that all four zones are absent. - let miss_mask = reader - .pruning_evaluation( - &(0..row_count), - &filter(MISS), - Mask::new_true(usize::try_from(row_count)?), - )? - .await?; - assert!( - miss_mask.all_false(), - "an absent value should prune every zone" - ); - - // An allow-unknown reader without Bloom registration bypasses the unavailable zone map and - // scans the data child. This both supplies the reference result and verifies that an optional - // index does not become a hard read-time dependency. - let full_scan_session = session(None); - full_scan_session.allow_unknown(); - let full_scan_file = full_scan_session.open_options().open_buffer(bytes)?; - - let indexed_hit = scan(&file, HIT).await?; - let full_scan_hit = scan(&full_scan_file, HIT).await?; - // A Bloom filter may retain extra zones, but it must never change query results. - assert_arrays_eq!( - indexed_hit, - full_scan_hit, - &mut read_session.create_execution_ctx() - ); - let expected_hit = - StructArray::from_fields(&[("id", PrimitiveArray::from_iter([HIT]).into_array())])? - .into_array(); - assert_arrays_eq!( - full_scan_hit, - expected_hit, - &mut read_session.create_execution_ctx() - ); - - let indexed_miss = scan(&file, MISS).await?; - let full_scan_miss = scan(&full_scan_file, MISS).await?; - assert_arrays_eq!( - indexed_miss, - full_scan_miss, - &mut read_session.create_execution_ctx() - ); - assert_eq!(full_scan_miss.len(), 0); - Ok(()) -} diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs index bee291cd548..4be8b75025e 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs @@ -39,7 +39,8 @@ pub(super) fn accumulate_canonical( | Canonical::List(_) | Canonical::FixedSizeList(_) | Canonical::Variant(_) - | Canonical::Union(_) => { + | Canonical::Union(_) + | Canonical::Map(_) => { vortex_bail!( "Unsupported canonical type for bloom filter: {}", canonical.dtype() diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs index f7bd4b8cbb5..3ca7f104775 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -31,15 +31,6 @@ pub(in crate::layouts::zoned) mod constant; pub use partial::BloomPartial; use crate::layouts::zoned::aggregates::bloom_filter::partial::BLOCK_SIZE; -use crate::layouts::zoned::skip_index::bloom::is_bloom_valid_dtype; - -// 1. (joacoc) Opted for blocks_count as a simpler way to tune -// the Bloom filter, even though there are other ways. -// 2. (joacoc) I think the optimal could be using statistics, similar -// to how vortex encoder selection works. -// 3. (joacoc) In case that the tune/options stays as is, -// a guide on how to select the number of blocks -// would be great. /// Bloom-filter tuning persisted as aggregate metadata. #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -231,8 +222,24 @@ impl AggregateFnVTable for BloomFilter { } } -// The following functions are utils/useful for tests in canonical and constants. +/// Returns true if the type is valid for the bloom index to acc/contain. +/// +/// This is defined by the available implementations in +/// [crate::layouts::zoned::aggregates::bloom::constant] and +/// [crate::layouts::zoned::aggregates::bloom::canonical] +fn is_bloom_valid_dtype(dtype: &DType) -> bool { + match dtype { + DType::Extension(ext) => is_bloom_valid_dtype(ext.storage_dtype()), + DType::Bool(_) + | DType::Primitive(..) + | DType::Decimal(..) + | DType::Utf8(_) + | DType::Binary(_) => true, + _ => false, + } +} +// The following functions are utils/useful for tests in canonical and constants. #[cfg(test)] pub(in crate::layouts::zoned::aggregates::bloom_filter) mod test_utils { use vortex_array::IntoArray; diff --git a/vortex-layout/src/layouts/zoned/aggregates/min_max.rs b/vortex-layout/src/layouts/zoned/aggregates/min_max.rs deleted file mode 100644 index 0403d5fb404..00000000000 --- a/vortex-layout/src/layouts/zoned/aggregates/min_max.rs +++ /dev/null @@ -1,70 +0,0 @@ -//! Min/max aggregate selection for zoned layouts. - -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::AggregateFnVTableExt; -use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; -use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions; -use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; -use vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions; -use vortex_array::aggregate_fn::fns::max::Max; -use vortex_array::aggregate_fn::fns::min::Min; -use vortex_array::dtype::DType; - -use super::super::schema::default_bounded_stat_max_bytes; - -pub(super) fn min_max_aggregate_fns(dtype: &DType) -> [AggregateFnRef; 2] { - match dtype { - DType::Utf8(_) | DType::Binary(_) => [ - BoundedMax.bind(BoundedMaxOptions { - max_bytes: default_bounded_stat_max_bytes(), - }), - BoundedMin.bind(BoundedMinOptions { - max_bytes: default_bounded_stat_max_bytes(), - }), - ], - _ => [ - Max.bind(NumericalAggregateOpts::skip_nans()), - Min.bind(NumericalAggregateOpts::skip_nans()), - ], - } -} - -#[cfg(test)] -mod tests { - use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; - use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; - use vortex_array::aggregate_fn::fns::max::Max; - use vortex_array::aggregate_fn::fns::min::Min; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - - use super::default_bounded_stat_max_bytes; - use super::min_max_aggregate_fns; - - #[test] - fn variable_length_min_max_are_bounded() { - let aggregate_fns = min_max_aggregate_fns(&DType::Utf8(Nullability::NonNullable)); - - assert_eq!( - aggregate_fns[0].as_::().max_bytes, - default_bounded_stat_max_bytes() - ); - assert_eq!( - aggregate_fns[1].as_::().max_bytes, - default_bounded_stat_max_bytes() - ); - } - - #[test] - fn fixed_width_min_max_are_exact() { - let aggregate_fns = min_max_aggregate_fns(&PType::I32.into()); - - assert!(aggregate_fns[0].is::()); - assert!(aggregate_fns[1].is::()); - } -} diff --git a/vortex-layout/src/layouts/zoned/aggregates/mod.rs b/vortex-layout/src/layouts/zoned/aggregates/mod.rs index 0ea90a5ba5b..0c0f15610c6 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/mod.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/mod.rs @@ -3,91 +3,4 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - -use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::AggregateFnVTableExt; -use vortex_array::aggregate_fn::EmptyOptions; -use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::nan_count::NanCount; -use vortex_array::aggregate_fn::fns::null_count::NullCount; -use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; -use vortex_array::dtype::DType; -use vortex_session::VortexSession; - pub(in crate::layouts::zoned) mod bloom_filter; -mod min_max; - -use min_max::min_max_aggregate_fns; - -pub(super) fn default_zoned_aggregate_fns( - dtype: &DType, - session: &VortexSession, -) -> Arc<[AggregateFnRef]> { - let mut aggregate_fns = Vec::from(min_max_aggregate_fns(dtype)); - if Sum - .return_dtype(&NumericalAggregateOpts::skip_nans(), dtype) - .is_some() - { - aggregate_fns.push(Sum.bind(NumericalAggregateOpts::skip_nans())); - } - aggregate_fns.push(NanCount.bind(EmptyOptions)); - aggregate_fns.push(NullCount.bind(EmptyOptions)); - - // Stats from geo extension types are discovered from the registry at runtime instead. - aggregate_fns.extend(session.aggregate_fns().zone_stat_defaults(dtype)); - - aggregate_fns.into() -} - -#[cfg(test)] -mod tests { - use vortex_array::aggregate_fn::AggregateFnVTableExt; - use vortex_array::aggregate_fn::fns::sum::Sum; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::extension::datetime::TimeUnit; - use vortex_array::extension::datetime::Timestamp; - - use super::bloom_filter::BloomFilter; - use super::bloom_filter::BloomOptions; - use super::default_zoned_aggregate_fns; - - #[test] - fn default_aggregates_exclude_bloom_filter() { - let aggregate_fns = - default_zoned_aggregate_fns(&PType::I64.into(), &vortex_array::array_session()); - let bloom = BloomFilter.bind(BloomOptions::default()); - - assert!( - aggregate_fns - .iter() - .all(|aggregate_fn| aggregate_fn != &bloom) - ); - } - - #[test] - fn default_aggregates_include_sum_for_numeric_dtype() { - let aggregate_fns = - default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session()); - - assert!(aggregate_fns[2].is::()); - } - - #[test] - fn default_aggregates_skip_sum_for_non_summable_dtype() { - let dtype = DType::Extension( - Timestamp::new(TimeUnit::Microseconds, Nullability::Nullable).erased(), - ); - let aggregate_fns = default_zoned_aggregate_fns(&dtype, &vortex_array::array_session()); - - assert!( - aggregate_fns - .iter() - .all(|aggregate_fn| !aggregate_fn.is::()) - ); - } -} diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 96893dfeac7..b699c58cea0 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -16,7 +16,6 @@ mod builder; mod pruning; mod reader; mod schema; -pub mod skip_index; pub mod writer; pub mod zone_map; diff --git a/vortex-layout/src/layouts/zoned/skip_index/bloom.rs b/vortex-layout/src/layouts/zoned/skip_index/bloom.rs deleted file mode 100644 index 55f98126959..00000000000 --- a/vortex-layout/src/layouts/zoned/skip_index/bloom.rs +++ /dev/null @@ -1,382 +0,0 @@ -//! Bloom skipping index for equality predicates. - -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::AggregateFnVTableExt; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; -use vortex_array::arrays::BoolArray; -use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::varbinview::VarBinViewArrayExt; -use vortex_array::dtype::DType; -use vortex_array::expr::Expression; -use vortex_array::expr::is_root; -use vortex_array::expr::not; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; -use vortex_array::scalar_fn::ExecutionArgs; -use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::binary::Binary; -use vortex_array::scalar_fn::fns::literal::Literal; -use vortex_array::scalar_fn::fns::operators::Operator; -use vortex_array::scalar_fn::session::ScalarFnSessionExt; -use vortex_array::stats::rewrite::StatsRewriteCtx; -use vortex_array::stats::rewrite::StatsRewriteRule; -use vortex_array::stats::session::StatsSessionExt; -use vortex_array::stats::stat; -use vortex_buffer::BitBuffer; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; -use vortex_session::VortexSession; -use vortex_session::registry::CachedId; - -use super::SkipIndex; -pub use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; -pub use crate::layouts::zoned::aggregates::bloom_filter::BloomOptions; -pub use crate::layouts::zoned::aggregates::bloom_filter::BloomPartial; - -/// Bloom skip index for constant-equality predicates. -/// -/// TODO(joacoc): Add documentation about the Bloom skip index -/// and how it works here. -#[derive(Clone, Debug, Default)] -pub struct BloomSkipIndex { - options: BloomOptions, -} - -impl BloomSkipIndex { - pub fn new(options: BloomOptions) -> Self { - Self { options } - } - - pub fn options(&self) -> &BloomOptions { - &self.options - } -} - -impl SkipIndex for BloomSkipIndex { - fn aggregate_fn(&self, input_dtype: &DType) -> Option { - BloomFilter - .return_dtype(&self.options, input_dtype) - .map(|_| BloomFilter.bind(self.options.clone())) - } - - fn register(&self, session: &VortexSession) { - session.aggregate_fns().register(BloomFilter); - session.scalar_fns().register(BloomContains); - session.stats().register_rewrite(BloomEqRewrite { - options: self.options.clone(), - }); - } -} - -/// Probe scalar function: test one literal against each binary Bloom state. -#[derive(Clone, Debug)] -struct BloomContains; - -impl ScalarFnVTable for BloomContains { - type Options = BloomOptions; - - fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("vortex.bloom_contains.v1"); - *ID - } - - fn serialize(&self, options: &Self::Options) -> VortexResult>> { - BloomFilter.serialize(options) - } - - fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { - BloomFilter.deserialize(metadata, session) - } - - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(2) - } - - /// Only two children are expected. - /// The first child represents the filter as a byte sequence, - /// while the second child represents the literal value to search for (the needle). - fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName { - match child_idx { - 0 => ChildName::from("filter"), - 1 => ChildName::from("needle"), - _ => unreachable!("bloom_contains has exactly two children"), - } - } - - fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { - vortex_ensure!( - matches!(args[0], DType::Binary(_)), - "bloom filter must be Binary" - ); - vortex_ensure!( - is_bloom_valid_dtype(&args[1]), - "bloom needle must be bool, primitive, decimal, utf8, binary or extension" - ); - - Ok(DType::Bool(args[0].nullability() | args[1].nullability())) - } - - fn execute( - &self, - options: &Self::Options, - args: &dyn ExecutionArgs, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let filters = args.get(0)?.execute::(ctx)?; - - // If the needle accepts an array of values, e.g., Array[1, 2, 3], - // the following code should be updated. - let needle_array = args.get(1)?; - let needle = needle_array - .as_constant() - .ok_or_else(|| vortex_err!("bloom needle must be constant"))?; - - let validity = filters.varbinview_validity(); - let valid = validity.execute_mask(filters.len(), ctx)?; - - // Quick return if the needle is invalid. - if !needle.is_valid() { - let possible = vec![false; filters.len()]; - return Ok(BoolArray::new(BitBuffer::from_iter(possible), validity).into_array()); - } - - let mut possible = Vec::with_capacity(filters.len()); - for (idx, is_valid) in valid.iter().enumerate() { - if !is_valid { - possible.push(false); - continue; - } - - let bytes = filters.bytes_at(idx); - let partial = BloomPartial::try_from(bytes.as_slice())?; - - vortex_ensure!( - partial.len() == options.blocks().get(), - "stored bloom length does not match options" - ); - - possible.push(partial.contains_valid_scalar(&needle)?); - } - - Ok(BoolArray::new(BitBuffer::from_iter(possible), validity).into_array()) - } - - fn is_null_sensitive(&self, _options: &Self::Options) -> bool { - false - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false - } -} - -/// Equality rewrite that turns a Bloom miss into a zone falsifier. -#[derive(Clone, Debug)] -struct BloomEqRewrite { - options: BloomOptions, -} - -impl StatsRewriteRule for BloomEqRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - Binary.id() - } - - /// Only works for root literal comparisons and valid [DTypes]. - /// - /// E.g. `eq(root(), lit(5i32))` or `eq(lit(5i32), root())` - fn falsify( - &self, - expr: &Expression, - ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { - if *expr.as_::() != Operator::Eq { - return Ok(None); - } - - let (column, literal) = if is_root(expr.child(0)) && expr.child(1).is::() { - (expr.child(0), expr.child(1)) - } else if is_root(expr.child(1)) && expr.child(0).is::() { - (expr.child(1), expr.child(0)) - } else { - return Ok(None); - }; - - if !is_bloom_valid_dtype(&ctx.return_dtype(column)?) || literal.as_::().is_null() { - return Ok(None); - } - - let filter = stat(column.clone(), BloomFilter.bind(self.options.clone())); - let contains = BloomContains.new_expr(self.options.clone(), [filter, literal.clone()]); - Ok(Some(not(contains))) - } -} - -/// Returns true if the type is valid for the bloom index to acc/contain. -/// -/// This is defined by the available implementations in -/// [crate::layouts::zoned::aggregates::bloom::constant] and -/// [crate::layouts::zoned::aggregates::bloom::canonical] -pub(in crate::layouts::zoned) fn is_bloom_valid_dtype(dtype: &DType) -> bool { - match dtype { - DType::Extension(ext) => is_bloom_valid_dtype(ext.storage_dtype()), - DType::Bool(_) - | DType::Primitive(..) - | DType::Decimal(..) - | DType::Utf8(_) - | DType::Binary(_) => true, - _ => false, - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use rstest::rstest; - use vortex_array::ArrayRef; - use vortex_array::Canonical; - use vortex_array::Columnar; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::aggregate_fn::AggregateFnVTable; - use vortex_array::aggregate_fn::AggregateFnVTableExt; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::StructArray; - use vortex_array::arrays::VarBinArray; - use vortex_array::assert_arrays_eq; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::expr::Expression; - use vortex_array::expr::eq; - use vortex_array::expr::gt_eq; - use vortex_array::expr::lit; - use vortex_array::expr::root; - use vortex_array::scalar::Scalar; - use vortex_array::validity::Validity; - use vortex_buffer::buffer; - use vortex_error::VortexResult; - - use super::BloomSkipIndex; - use super::SkipIndex; - use crate::layouts::zoned::aggregates::bloom_filter::BloomFilter; - use crate::layouts::zoned::aggregates::bloom_filter::BloomOptions; - use crate::layouts::zoned::zone_map::ZoneMap; - use crate::test::SESSION; - - #[test] - fn missing_stat_stays_inconclusive() -> VortexResult<()> { - let session = vortex_array::array_session(); - let index = BloomSkipIndex::new(BloomOptions::default()); - index.register(&session); - let predicate = eq(root(), lit(42i64)); - let proof = predicate - .falsify( - &DType::Primitive(PType::I64, Nullability::NonNullable), - &session, - )? - .expect("equality has a bloom proof"); - - let zone_map = ZoneMap::try_new( - DType::Primitive(PType::I64, Nullability::NonNullable), - StructArray::try_new(Vec::<&str>::new().into(), vec![], 2, Validity::NonNullable)?, - Arc::new([]), - 8, - 16, - )?; - assert!(zone_map.prune(&proof, &session)?.all_false()); - Ok(()) - } - - /// Similar zone map tests as the ones in [crate::layouts::zoned::tests] - /// but using BloomFilter rather than max/min zones. - fn build_bloom_zone_map(dtype: DType, batch: ArrayRef) -> ZoneMap { - let bloom = BloomFilter; - let options = BloomOptions::default(); - - // If index is not registered there will be no warning, - // but the index will return false for everything. - // - // (joacoc) should be considered a warning for missing aggregatefns? - BloomSkipIndex::new(options.clone()).register(&SESSION); - let mut ctx = SESSION.create_execution_ctx(); - - let mut zone_filter = bloom.empty_partial(&options, &dtype).unwrap(); - bloom - .accumulate( - &mut zone_filter, - &Columnar::Canonical(batch.execute::(&mut ctx).unwrap()), - &mut ctx, - ) - .unwrap(); - - let zone_filter_as_scalar = bloom.to_scalar(&zone_filter).unwrap(); - let zone_filter_as_bytes = zone_filter_as_scalar.as_binary().value().unwrap().to_vec(); - let zone_filter_as_varbin = - VarBinArray::from_nullable_bytes(vec![Some(zone_filter_as_bytes.as_slice())]); - - let bloom = BloomFilter.bind(options); - let zone_filter_struct = StructArray::from_fields(&[( - bloom.clone().to_string(), - zone_filter_as_varbin.into_array(), - )]) - .unwrap(); - - ZoneMap::try_new(dtype, zone_filter_struct, Arc::new([bloom]), 1, 10).unwrap() - } - - fn assert_prune(zone_map: &ZoneMap, dtype: &DType, expr: Expression, expected: [bool; 1]) { - let mut ctx = SESSION.create_execution_ctx(); - let pruning_expr = expr.falsify(dtype, &SESSION).unwrap().unwrap(); - let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter(expected), &mut ctx); - } - - #[rstest] - #[case::equals_value_not_in_batch(eq(root(), lit(99i32)), [true])] - #[case::equals_value_in_batch(eq(root(), lit(5i32)), [false])] - #[case::gt_eq_not_supported_by_bloom(gt_eq(root(), lit(4i32)), [false])] - #[case::null_never_prunes( - gt_eq(root(), lit(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)))), - [false] - )] - fn test_zone_map_prunes_with_bloom_filter_i32( - #[case] expr: Expression, - #[case] expected: [bool; 1], - ) -> VortexResult<()> { - let dtype = DType::Primitive(PType::I32, Nullability::Nullable); - let batch = PrimitiveArray::new(buffer![5i32, 6i32, 7i32], Validity::AllValid).into_array(); - let zone_map = build_bloom_zone_map(dtype.clone(), batch); - assert_prune(&zone_map, &dtype, expr, expected); - Ok(()) - } - - #[rstest] - #[case::equals_value_not_in_batch(eq(root(), lit("zz")), [true])] - #[case::equals_value_in_batch(eq(root(), lit("london")), [false])] - fn test_zone_map_prunes_with_bloom_filter_varbin( - #[case] expr: Expression, - #[case] expected: [bool; 1], - ) -> VortexResult<()> { - let dtype = DType::Utf8(Nullability::NonNullable); - let batch = VarBinArray::from_iter( - [Some("london"), Some("hamburg"), Some("newyork")], - dtype.clone(), - ) - .into_array(); - let zone_map = build_bloom_zone_map(dtype.clone(), batch); - assert_prune(&zone_map, &dtype, expr, expected); - Ok(()) - } -} diff --git a/vortex-layout/src/layouts/zoned/skip_index/mod.rs b/vortex-layout/src/layouts/zoned/skip_index/mod.rs deleted file mode 100644 index 2def62acab0..00000000000 --- a/vortex-layout/src/layouts/zoned/skip_index/mod.rs +++ /dev/null @@ -1,59 +0,0 @@ -//! Skipping-index interface and implementations. - -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fmt::Debug; -use std::sync::Arc; - -use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::dtype::DType; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_session::VortexSession; - -use super::aggregates::default_zoned_aggregate_fns; -use super::writer::ZonedLayoutOptions; - -pub mod bloom; - -/// One definition that supplies a persisted aggregate and registers every read-side component -/// needed to consult it. -/// -/// The writer helper [`ZonedLayoutOptions::with_skip_index`] is the explicit per-column declaration -/// seam. Readers call [`SkipIndex::register`] on their session before opening the file. -pub trait SkipIndex: Debug + Send + Sync + 'static { - /// The aggregate state to persist for `input_dtype`, or `None` when unsupported. - fn aggregate_fn(&self, input_dtype: &DType) -> Option; - - /// Register the aggregate, optional probe function, and predicate rewrite as one operation. - fn register(&self, session: &VortexSession); -} - -impl ZonedLayoutOptions { - /// Add `index` to this zoned writer while retaining the default min/max-style aggregates. - /// - /// `WriteStrategyBuilder::with_field_zoned_options` can install the configured options for one - /// field while retaining the default data layout pipeline. - pub fn with_skip_index( - mut self, - index: &I, - input_dtype: &DType, - session: &VortexSession, - ) -> VortexResult { - let aggregate_fn = index - .aggregate_fn(input_dtype) - .ok_or_else(|| vortex_err!("skip index does not support input dtype {input_dtype}"))?; - - let mut aggregate_fns = self - .aggregate_fns - .take() - .unwrap_or_else(|| default_zoned_aggregate_fns(input_dtype, session)) - .to_vec(); - if !aggregate_fns.iter().any(|stored| stored == &aggregate_fn) { - aggregate_fns.push(aggregate_fn); - } - self.aggregate_fns = Some(Arc::from(aggregate_fns)); - Ok(self) - } -} diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index d4dcad7e8a6..4151679e6c3 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -38,7 +38,7 @@ use crate::LayoutWriterContext; use crate::layouts::zoned::AggregateStatsAccumulator; use crate::layouts::zoned::ZonedLayout; use crate::layouts::zoned::aggregate_partials; -use crate::layouts::zoned::aggregates::default_zoned_aggregate_fns; +use crate::layouts::zoned::schema::default_bounded_stat_max_bytes; use crate::segments::SegmentSinkRef; use crate::sequence::SendableSequentialStream; use crate::sequence::SequencePointer; @@ -50,7 +50,6 @@ use crate::sequence::SequentialStreamExt; /// /// The input stream is assumed to already be partitioned into one chunk per zone, except /// possibly the final partial zone. -#[derive(Clone)] pub struct ZonedLayoutOptions { /// The size of a statistics block pub block_size: NonZeroUsize, From 72f3d527747e189160c269e49f3353e08143a568 Mon Sep 17 00:00:00 2001 From: Joaquin Colacci Date: Fri, 14 Aug 2026 18:12:35 +0200 Subject: [PATCH 7/7] DCO Remediation Commit for Joaquin Colacci I, Joaquin Colacci , hereby add my Signed-off-by to this commit: a9e0550c93223daba977962b10cafce576c12370 I, Joaquin Colacci , hereby add my Signed-off-by to this commit: 8d81b781e1250bafbae006d9cb13d1e8b0722610 I, Joaquin Colacci , hereby add my Signed-off-by to this commit: 7d0495faf53808ed9bc5f14b708bf47fa9209521 I, Joaquin Colacci , hereby add my Signed-off-by to this commit: 43b064dd446de86d56ccbed7afde4af82fd51b7d I, Joaquin Colacci , hereby add my Signed-off-by to this commit: f6ee474edbcd5aa9e90015aef3e7733e8b7f9950 Signed-off-by: Joaquin Colacci --- .../src/layouts/zoned/aggregates/bloom_filter/partial.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs index 6684ddab8f0..2f3d8fdd8f1 100644 --- a/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs @@ -28,7 +28,7 @@ pub struct BloomPartial { pub(super) blocks: Vec<[u32; 8]>, } -/// The following Split block Bloom filter (SBBF) implementation +/// The following split block Bloom filter implementation /// is a translation of the original paper's names and values, /// with slight changes to let the Rust compiler generate /// optimized, vectorized code for `make_mask`, `add_hash`, and `find_hash`.