diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 763db758cb7..86efdf2523e 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -157,6 +157,7 @@ pub struct WriteStrategyBuilder { compressor: CompressorConfig, row_block_size: usize, field_writers: HashMap>, + field_zoned_options: HashMap, allow_encodings: Option>, flat_strategy: Option>, probe_compressor: Option>, @@ -174,6 +175,7 @@ impl Default for WriteStrategyBuilder { compressor: CompressorConfig::BtrBlocks(BtrBlocksCompressorBuilder::default()), row_block_size: 8192, field_writers: HashMap::new(), + field_zoned_options: HashMap::new(), allow_encodings: Some(ALLOWED_ENCODINGS.clone()), flat_strategy: None, probe_compressor: None, @@ -216,6 +218,20 @@ 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. /// /// The configured flat leaf strategy is wrapped in a [`LayoutStrategyEncodingValidator`] @@ -260,7 +276,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 { @@ -332,36 +348,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() - }, - ); - - // 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, - }, - ); + 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() + }); + + 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 3a4d80902f3..a877c36e5f6 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 048905dade6..4ace2a57eeb 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -13,21 +13,6 @@ use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; 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::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::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_error::VortexError; use vortex_error::VortexResult; use vortex_io::session::RuntimeSessionExt; @@ -40,7 +25,7 @@ use crate::LayoutStrategy; 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; @@ -52,6 +37,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, @@ -195,91 +181,3 @@ impl LayoutStrategy for ZonedStrategy { self.child.buffered_bytes() + self.stats.buffered_bytes() } } - -fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[AggregateFnRef]> { - let (max, min) = 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()), - ), - }; - - let mut aggregate_fns = vec![max, min]; - 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::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::aggregate_fn::fns::sum::Sum; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - use vortex_array::extension::datetime::TimeUnit; - use vortex_array::extension::datetime::Timestamp; - - use super::*; - - #[test] - fn default_aggregates_bound_variable_length_min_max() { - let aggregate_fns = default_zoned_aggregate_fns( - &DType::Utf8(Nullability::NonNullable), - &vortex_array::array_session(), - ); - - 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 default_aggregates_keep_fixed_width_min_max_exact() { - let aggregate_fns = - default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session()); - - assert!(aggregate_fns[0].is::()); - assert!(aggregate_fns[1].is::()); - 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::()) - ); - } -}