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-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/canonical/bool.rs b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs new file mode 100644 index 00000000000..7d85d0c687d --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/bool.rs @@ -0,0 +1,115 @@ +// 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..7ad21137ff8 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/extension.rs @@ -0,0 +1,92 @@ +// 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..4be8b75025e --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/mod.rs @@ -0,0 +1,52 @@ +// 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(_) + | Canonical::Map(_) => { + 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..dd5308bffa6 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/canonical/primitive.rs @@ -0,0 +1,300 @@ +// 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(()) + } + + /// 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_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)?); + } + + 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)?); + } + + 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)?); + } + } + + 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..0100e36071f --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/constant.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::arrays::ConstantArray; +use vortex_error::VortexResult; + +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_valid_scalar(scalar)?; + + Ok(()) +} + +#[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..3ca7f104775 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/mod.rs @@ -0,0 +1,383 @@ +//! 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::VortexExpect; +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::aggregates::bloom_filter::partial::BLOCK_SIZE; + +/// 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). + /// + /// 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. + /// + /// ### 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) + .vortex_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() * BLOCK_SIZE); + 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) + } +} + +/// 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; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::Accumulator; + use vortex_array::aggregate_fn::DynAccumulator; + 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(); + 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() % 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]; + 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 { + assert!(partial.contains(i), "value {i} missing after merge"); + } + + for i in 101..200i64 { + 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 new file mode 100644 index 00000000000..2f3d8fdd8f1 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/bloom_filter/partial.rs @@ -0,0 +1,265 @@ +//! Split block Bloom filters (SBBF) 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_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) +pub(super) 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 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() + } + + #[inline] + pub(super) fn insert(&mut self, value: T) + where + T: Hash, + { + let hash = self.hash(value); + self.insert_hash(hash); + } + + #[inline] + fn insert_hash(&mut self, hash: u64) { + self.add_hash(hash); + } + + #[inline] + 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() + } + + fn add_hash(&mut self, hash: u64) { + let idx = self.block_index(hash, self.blocks.len()) as usize; + + let mask = self.make_mask(hash as u32); + for i in 0..8 { + self.blocks[idx][i] |= mask[i]; + } + } + + /// Checks whether a hash is (probably) present in the filter. + 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. + 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 + } +} + +/// 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 { + 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 (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()) + })?); + } + Ok(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 new file mode 100644 index 00000000000..0c0f15610c6 --- /dev/null +++ b/vortex-layout/src/layouts/zoned/aggregates/mod.rs @@ -0,0 +1,6 @@ +//! Aggregate functions selected by the zoned layout. + +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +pub(in crate::layouts::zoned) mod bloom_filter; diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index bc3d7d0626e..b699c58cea0 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -11,6 +11,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod aggregates; mod builder; mod pruning; mod reader; @@ -510,6 +511,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 { @@ -528,6 +531,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);