diff --git a/Cargo.lock b/Cargo.lock index 2448cb9aff8..70d7eb4f883 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11009,6 +11009,7 @@ dependencies = [ "rstest", "vortex-alp", "vortex-array", + "vortex-bench-support", "vortex-buffer", "vortex-error", "vortex-fastlanes", diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 9085390b67b..9fb317ec8c8 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -39,6 +39,7 @@ rand = { workspace = true } rstest = { workspace = true } vortex-alp = { path = "../alp" } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-bench-support = { workspace = true } vortex-fastlanes = { path = ".", features = ["_test-harness"] } [features] @@ -48,6 +49,10 @@ _test-harness = ["dep:rand"] name = "bitpacking_take" harness = false +[[bench]] +name = "bitpacking_list_contains" +harness = false + [[bench]] name = "canonicalize_bench" harness = false diff --git a/encodings/fastlanes/benches/bitpacking_list_contains.rs b/encodings/fastlanes/benches/bitpacking_list_contains.rs new file mode 100644 index 00000000000..72b4c1beada --- /dev/null +++ b/encodings/fastlanes/benches/bitpacking_list_contains.rs @@ -0,0 +1,245 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Measures compressed constant-list membership. +//! +//! FastLanes evaluates constant lists with at most four distinct non-null members during unpacking. +//! Mid-size lists use repeated packed comparisons. Larger lists decode once at a threshold that +//! depends on the physical integer width and array length. Every path runs on each real CPU feature +//! leg in CodSpeed. +//! To recalculate the thresholds, temporarily replace `min_decode_source_members` with a constant. +//! Return `usize::MAX` to force repeated comparisons. Return `5` to force decode-once. +//! +//! Run with `cargo bench -p vortex-fastlanes --bench bitpacking_list_contains`. + +#![expect(clippy::unwrap_used)] + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hint::black_box; +use std::sync::Arc; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::Alignment; +use vortex_buffer::BufferMut; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArray; +use vortex_fastlanes::BitPackedData; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +trait BenchInt: IntegerPType + Copy + Into { + fn from_counter(value: u64) -> Self; +} + +impl BenchInt for u8 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u16 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u32 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u64 { + fn from_counter(value: u64) -> Self { + value + } +} + +#[derive(Clone, Copy)] +struct PackedCase { + name: &'static str, + ptype: PType, + bit_width: u8, + len: usize, + member_count: usize, + member_stride: u64, +} + +impl Display for PackedCase { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "{}_{}_w{}_n{}", + self.name, self.ptype, self.bit_width, self.len + ) + } +} + +const fn strided_case( + name: &'static str, + ptype: PType, + bit_width: u8, + len: usize, + count: usize, + stride: u64, +) -> PackedCase { + PackedCase { + name, + ptype, + bit_width, + len, + member_count: count, + member_stride: stride, + } +} + +const PACKED_CASES: &[PackedCase] = &[ + strided_case("direct_u8_m4", PType::U8, 6, 65_536, 4, 2), + strided_case("direct_u16_m4", PType::U16, 12, 65_536, 4, 2), + strided_case("direct_u32_m4", PType::U32, 20, 65_536, 4, 2), + strided_case("direct_u64_m4", PType::U64, 40, 65_536, 4, 2), + strided_case("generic_u8_m29", PType::U8, 6, 65_536, 29, 2), + strided_case("decode_u8_m30", PType::U8, 6, 65_536, 30, 2), + strided_case("generic_u16_m24", PType::U16, 8, 65_536, 24, 2), + strided_case("decode_u16_m25", PType::U16, 8, 65_536, 25, 2), + strided_case("generic_u32_m12", PType::U32, 8, 65_536, 12, 2), + strided_case("decode_u32_m13", PType::U32, 8, 65_536, 13, 2), + strided_case("decode_u64_m5", PType::U64, 40, 65_536, 5, 2), + strided_case("short_direct_u32_m4", PType::U32, 10, 1_024, 4, 2), + strided_case("short_generic_u8_m9", PType::U8, 6, 8_192, 9, 2), + strided_case("short_decode_u8_m10", PType::U8, 6, 8_192, 10, 2), + strided_case("short_generic_u16_m9", PType::U16, 8, 8_192, 9, 2), + strided_case("short_decode_u16_m10", PType::U16, 8, 8_192, 10, 2), + strided_case("short_generic_u32_m10", PType::U32, 8, 16_384, 10, 2), + strided_case("short_decode_u32_m11", PType::U32, 8, 16_384, 11, 2), + strided_case("longer_generic_u8_m10", PType::U8, 6, 16_384, 10, 2), + strided_case("longer_generic_u16_m10", PType::U16, 8, 16_384, 10, 2), + strided_case("longer_generic_u32_m11", PType::U32, 8, 32_768, 11, 2), + strided_case("short_direct_u64_m4", PType::U64, 8, 8_192, 4, 2), + strided_case("short_decode_u64_m5", PType::U64, 8, 8_192, 5, 2), +]; + +fn page_aligned(array: BitPackedArray) -> BitPackedArray { + let ptype = array.dtype().as_ptype(); + let parts = BitPacked::into_parts(array); + BitPacked::try_new( + parts.packed.ensure_aligned(Alignment::new(4_096)).unwrap(), + ptype, + parts.validity, + parts.patches, + parts.bit_width, + parts.len, + parts.offset, + ) + .unwrap() +} + +fn generated_values(case: PackedCase, members: &[u64]) -> Vec { + let domain_size = 1u64 << case.bit_width; + let mut state = 0x9E37_79B9_7F4A_7C15u64; + (0..case.len) + .map(|_| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + let is_hit = (state >> 32).is_multiple_of(2); + if is_hit { + let member_index = + usize::try_from(state % u64::try_from(members.len()).unwrap()).unwrap(); + members[member_index] + } else { + let mut candidate = state.rotate_left(17) % domain_size; + while members.contains(&candidate) { + candidate = (candidate + 1) % domain_size; + } + candidate + } + }) + .collect() +} + +fn list_scalar(members: &[u64]) -> Scalar { + Scalar::list( + Arc::new(DType::Primitive(T::PTYPE, Nullability::NonNullable)), + members + .iter() + .map(|value| T::from_counter(*value).into()) + .collect(), + Nullability::NonNullable, + ) +} + +fn packed_input( + case: PackedCase, +) -> (BitPackedArray, Scalar, BoolArray, VortexSession) { + let session = array_session(); + vortex_fastlanes::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let members = (0..case.member_count) + .map(|index| u64::try_from(index).unwrap() * case.member_stride) + .collect::>(); + let generated = generated_values(case, &members); + let expected = BoolArray::from_iter(generated.iter().map(|value| members.contains(value))); + let values: BufferMut = generated.into_iter().map(T::from_counter).collect(); + let packed = page_aligned( + BitPackedData::encode( + &PrimitiveArray::new(values.freeze(), Validity::NonNullable).into_array(), + case.bit_width, + &mut ctx, + ) + .unwrap(), + ); + (packed, list_scalar::(&members), expected, session) +} + +fn bench_packed_current(bencher: Bencher, case: PackedCase) { + let (packed, list, expected, session) = packed_input::(case); + let contains = packed + .into_array() + .apply(&list_contains(lit(list), root())) + .unwrap(); + let mut ctx = session.create_execution_ctx(); + let actual = contains.clone().execute::(&mut ctx).unwrap(); + assert_arrays_eq!(actual, expected, &mut ctx); + bencher + .counter(ItemsCount::new(case.len)) + .bench_local(|| black_box(contains.clone().execute::(&mut ctx).unwrap())); +} + +macro_rules! dispatch_packed { + ($bencher:expr, $case:expr, $function:ident) => { + match $case.ptype { + PType::U8 => $function::($bencher, $case), + PType::U16 => $function::($bencher, $case), + PType::U32 => $function::($bencher, $case), + PType::U64 => $function::($bencher, $case), + _ => unreachable!("benchmark case uses an unsigned integer type"), + } + }; +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = PACKED_CASES)] +fn packed_current(bencher: Bencher, case: PackedCase) { + dispatch_packed!(bencher, case, bench_packed_current); +} diff --git a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs index f27384b898c..82a43647fc0 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Fused compare kernel for [`BitPackedArray`] against a constant. +//! Fused predicate kernel for [`BitPackedArray`]. //! //! Where [`super::stream_predicate`] unpacks a full 1024-element FastLanes block into a scratch //! buffer and *then* folds a predicate over it, this path hands the comparison down into the -//! FastLanes [`BitPackingCompare::unchecked_unpack_cmp`] kernel, which compares each value against -//! the constant *as it is unpacked*, accumulating the boolean results straight into a 1024-bit +//! FastLanes [`BitPackingCompare::unchecked_unpack_cmp`] kernel, which evaluates each value +//! *as it is unpacked*, accumulating the boolean results straight into a 1024-bit //! mask (`[u64; 16]`) in transposed FastLanes lane order - one register-resident word per lane, no //! `[bool; 1024]` or `[T; 1024]` scratch. A single SIMD [`transpose_bits`] per block then rotates //! that mask into logical row order. @@ -21,7 +21,7 @@ //! slot with no per-block temporary and only one shared scratch `[u64; 16]`. The leading `offset` //! garbage rows are represented as the final [`BitBuffer`] bit offset, which naturally handles //! sub-byte slices without copy-aligning. Inline patches are spliced in afterwards by overwriting -//! the bits at the patched indices with `cmp(patch_value, rhs)`. +//! the bits at the patched indices with the predicate result. //! //! [`BitPackedArray`]: crate::BitPackedArray //! [`BitBuffer`]: vortex_buffer::BitBuffer @@ -70,6 +70,46 @@ pub(super) fn stream_compare_fused( cmp: F, ctx: &mut ExecutionCtx, ) -> VortexResult +where + T: NativePType + + BitPackedIter + + FastLanesComparable::Physical>, + ::Physical: BitPacking + NativePType + BitPackingCompare, + F: Fn(T, T) -> bool + Copy, +{ + stream_compare_fused_inner(array, rhs, nullability, cmp, ctx) +} + +/// Evaluates `predicate` while FastLanes unpacks each value. +pub(super) fn stream_predicate_fused( + array: ArrayView<'_, BitPacked>, + nullability: Nullability, + predicate: F, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativePType + + BitPackedIter + + FastLanesComparable::Physical>, + ::Physical: BitPacking + NativePType + BitPackingCompare, + F: Fn(T) -> bool + Copy, +{ + stream_compare_fused_inner( + array, + T::default(), + nullability, + move |value, _| predicate(value), + ctx, + ) +} + +fn stream_compare_fused_inner( + array: ArrayView<'_, BitPacked>, + rhs: T, + nullability: Nullability, + cmp: F, + ctx: &mut ExecutionCtx, +) -> VortexResult where T: NativePType + BitPackedIter @@ -84,7 +124,7 @@ where // A degenerate width has no packed payload for the fused kernel to consume; defer to the scalar // streaming predicate, which handles every layout (including the empty array). if len == 0 || bit_width == 0 { - return stream_predicate::(array, nullability, move |v| cmp(v, rhs), ctx); + return stream_predicate::(array, nullability, move |value| cmp(value, rhs), ctx); } // Over-allocate to whole 1024-bit blocks in padded coordinates so every block - including the @@ -119,12 +159,12 @@ where let mut bits = BitBufferMut::from_buffer(words.into_byte_buffer(), offset, len); - // Patched indices hold placeholder packed values, so their fused result is meaningless; - // overwrite each with the comparison against the real patch value. + // Patched indices hold placeholder packed values, so their fused result is meaningless. + // Overwrite each result with the predicate for the real patch value. // TODO(joe): apply patches per `packed_chunked`. if let Some(p) = array.patches() { let p_idx = p.indices().clone().execute::(ctx)?; - // TODO(joe): push down cmp?? + // TODO(joe): push down the predicate. let p_val = p.values().clone().execute::(ctx)?; let p_off = p.offset(); match_each_unsigned_integer_ptype!(p_idx.ptype(), |I| { diff --git a/encodings/fastlanes/src/bitpacking/compute/list_contains/mod.rs b/encodings/fastlanes/src/bitpacking/compute/list_contains/mod.rs new file mode 100644 index 00000000000..d708c60a89e --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/compute/list_contains/mod.rs @@ -0,0 +1,157 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use fastlanes::BitPacking; +use fastlanes::BitPackingCompare; +use fastlanes::FastLanesComparable; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::evaluate_prepared_integer_membership; +use vortex_array::arrays::primitive::integer_membership_binary_search_min; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::dtype::PhysicalPType; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::list_contains::IntegerMembership; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; + +use super::compare_fused::stream_predicate_fused; +use crate::BitPacked; +use crate::unpack_iter::BitPacked as BitPackedIter; + +const MAX_FUSED_DISTINCT_MEMBERS: usize = 4; +const SHORT_ARRAY_MAX_ROWS_8_16: usize = 8_192; +const SHORT_ARRAY_MAX_ROWS_32: usize = 16_384; +fn min_decode_source_members(ptype: PType, len: usize) -> usize { + // The generic fallback scans the packed child once per source member. Decode before repeated + // packed scans become more expensive than one decode plus Primitive membership evaluation. + let short_array_max_rows = if ptype.bit_width() == 32 { + SHORT_ARRAY_MAX_ROWS_32 + } else { + SHORT_ARRAY_MAX_ROWS_8_16 + }; + if len <= short_array_max_rows && ptype.bit_width() < 64 { + return integer_membership_binary_search_min(ptype); + } + match ptype.bit_width() { + 8 => 30, + 16 => 25, + 32 => 13, + 64 => 5, + _ => 5, + } +} + +impl ListContainsElementKernel for BitPacked { + fn list_contains( + list: &ArrayRef, + element: ArrayView<'_, Self>, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + list_contains_compressed(list, element, ctx) + } +} + +fn list_contains_compressed( + list: &ArrayRef, + element: ArrayView<'_, BitPacked>, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let nullability = list.dtype().nullability() | element.dtype().nullability(); + + match_each_integer_ptype!(element.dtype().as_ptype(), |T| { + list_contains_typed::(list, element, nullability, ctx) + }) +} + +fn list_contains_typed( + list: &ArrayRef, + element: ArrayView<'_, BitPacked>, + nullability: vortex_array::dtype::Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult> +where + T: IntegerPType + + BitPackedIter + + FastLanesComparable::Physical>, + ::Physical: BitPacking + NativePType + BitPackingCompare, +{ + let Some(membership) = IntegerMembership::::try_from_constant_list(list, element.dtype())? + else { + return Ok(None); + }; + if membership.members().len() > MAX_FUSED_DISTINCT_MEMBERS { + if membership.non_null_source_len() + < min_decode_source_members(element.dtype().as_ptype(), element.len()) + { + return Ok(None); + } + // The generic list implementation expands membership into one comparison per source + // member. Each comparison scans the packed child. Decode once before applying the + // Primitive membership policy when repeated packed scans become more expensive. + let primitive = element.array().clone().execute::(ctx)?; + return evaluate_prepared_integer_membership(membership, primitive.as_view(), nullability) + .map(Some); + } + + let result = match membership.members() { + [] => BoolArray::new( + BitBuffer::new_unset(element.len()), + element.validity()?.union_nullability(nullability), + ) + .into_array(), + [member] => { + let member = *member; + stream_predicate_fused::( + element, + nullability, + move |value| value.is_eq(member), + ctx, + )? + } + [first, second] => { + let (first, second) = (*first, *second); + stream_predicate_fused::( + element, + nullability, + move |value| value.is_eq(first) | value.is_eq(second), + ctx, + )? + } + [first, second, third] => { + let (first, second, third) = (*first, *second, *third); + stream_predicate_fused::( + element, + nullability, + move |value| value.is_eq(first) | value.is_eq(second) | value.is_eq(third), + ctx, + )? + } + [first, second, third, fourth] => { + let (first, second, third, fourth) = (*first, *second, *third, *fourth); + stream_predicate_fused::( + element, + nullability, + move |value| { + value.is_eq(first) + | value.is_eq(second) + | value.is_eq(third) + | value.is_eq(fourth) + }, + ctx, + )? + } + _ => return Ok(None), + }; + Ok(Some(result)) +} + +#[cfg(test)] +mod tests; diff --git a/encodings/fastlanes/src/bitpacking/compute/list_contains/tests.rs b/encodings/fastlanes/src/bitpacking/compute/list_contains/tests.rs new file mode 100644 index 00000000000..435de238896 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/compute/list_contains/tests.rs @@ -0,0 +1,351 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::slice::SliceKernel; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::PValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementKernel; +#[cfg(not(codspeed))] +use vortex_array::test_harness::trace::trace_op; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +use crate::BitPacked; +use crate::BitPackedArray; +use crate::BitPackedArrayExt; +use crate::BitPackedData; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + +fn member_list( + values: impl IntoIterator>, + member_nullability: Nullability, +) -> Scalar +where + T: NativePType + Into, +{ + let member_dtype = DType::Primitive(T::PTYPE, member_nullability); + let members = values + .into_iter() + .map(|value| { + value + .map(|value| Scalar::primitive(value, member_nullability)) + .unwrap_or_else(|| Scalar::null(member_dtype.clone())) + }) + .collect(); + Scalar::list(Arc::new(member_dtype), members, Nullability::NonNullable) +} + +fn list_array(list: Scalar, len: usize) -> ArrayRef { + ConstantArray::new(list, len).into_array() +} + +fn execute_direct( + list: &ArrayRef, + element: &BitPackedArray, + ctx: &mut vortex_array::ExecutionCtx, +) -> VortexResult { + ::list_contains(list, element.as_view(), ctx)? + .ok_or_else(|| vortex_err!("BitPacked list_contains kernel declined a supported input"))? + .execute::(ctx) +} + +macro_rules! integer_type_test { + ($name:ident, $T:ty, $bit_width:expr) => { + #[test] + fn $name() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..2_048) + .map(|value| (value % 64) as $T) + .collect::>(); + let members = [1 as $T, 3 as $T, 63 as $T]; + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), $bit_width, &mut ctx)?; + let list = list_array( + member_list(members.into_iter().map(Some), Nullability::NonNullable), + packed.len(), + ); + + let actual = execute_direct(&list, &packed, &mut ctx)?; + let expected = + BoolArray::from_iter(values.into_iter().map(|value| members.contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + }; +} + +integer_type_test!(test_integer_type_u8, u8, 6); +integer_type_test!(test_integer_type_u16, u16, 6); +integer_type_test!(test_integer_type_u32, u32, 6); +integer_type_test!(test_integer_type_u64, u64, 6); +integer_type_test!(test_integer_type_i8, i8, 6); +integer_type_test!(test_integer_type_i16, i16, 6); +integer_type_test!(test_integer_type_i32, i32, 6); +integer_type_test!(test_integer_type_i64, i64, 6); + +#[rstest] +#[case::one(vec![3])] +#[case::two(vec![3, 7])] +#[case::three(vec![3, 7, 11])] +#[case::four(vec![3, 7, 11, 15])] +#[case::duplicate_source(vec![3, 3, 7, 7, 11, 11, 15, 15, 15])] +fn test_member_cardinalities(#[case] members: Vec) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..2_048).map(|value| value % 128).collect::>(); + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + let list = list_array( + member_list(members.iter().copied().map(Some), Nullability::NonNullable), + packed.len(), + ); + + let actual = execute_direct(&list, &packed, &mut ctx)?; + let expected = BoolArray::from_iter(values.into_iter().map(|value| members.contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::generic_five((0..5).map(|value| value * 2).collect())] +#[case::decoded_many((0..32).map(|value| value * 2).collect())] +fn test_many_member_public_expression_paths(#[case] members: Vec) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..4_096).map(|value| value % 128).collect::>(); + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + let expression = list_contains( + lit(member_list( + members.iter().copied().map(Some), + Nullability::NonNullable, + )), + root(), + ); + + let actual = packed + .into_array() + .apply(&expression)? + .execute::(&mut ctx)?; + let expected = BoolArray::from_iter(values.into_iter().map(|value| members.contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_many_member_kernel_policy() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = [0i32, 7, 99]; + let primitive = PrimitiveArray::from_iter(values); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + let decode_threshold = + super::min_decode_source_members(vortex_array::dtype::PType::I32, packed.len()); + + for (member_count, expected_supported) in + [(decode_threshold - 1, false), (decode_threshold, true)] + { + let member_count = i32::try_from(member_count).vortex_expect("member count fits in an i32"); + let list = list_array( + member_list((0..member_count).map(Some), Nullability::NonNullable), + packed.len(), + ); + let actual = ::list_contains( + &list, + packed.as_view(), + &mut ctx, + )?; + + assert_eq!(actual.is_some(), expected_supported); + if let Some(actual) = actual { + let expected = + BoolArray::from_iter(values.map(|value| (0..member_count).contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + } + } + Ok(()) +} + +#[rstest] +#[case::present([true; 128], vec![0])] +#[case::absent([false; 128], vec![1])] +fn test_zero_bit_width( + #[case] expected: [bool; 128], + #[case] members: Vec, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_iter([0i32; 128]); + let packed = BitPackedData::encode(&primitive.into_array(), 0, &mut ctx)?; + let list = list_array( + member_list(members.into_iter().map(Some), Nullability::NonNullable), + packed.len(), + ); + + let actual = execute_direct(&list, &packed, &mut ctx)?; + let expected = BoolArray::from_iter(expected); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_sliced_patched_array() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..5_000) + .map(|index| { + if index % 97 == 0 { + 100_000 + index + } else { + index % 100 + } + }) + .collect::>(); + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + assert!(packed.patches().is_some(), "test setup requires patches"); + let range = 333..4_333; + let sliced = ::slice(packed.as_view(), range.clone(), &mut ctx)? + .ok_or_else(|| vortex_err!("BitPacked slice kernel declined a supported input"))?; + let members = [3, 100_388]; + let list = list_array( + member_list(members.into_iter().map(Some), Nullability::NonNullable), + sliced.len(), + ); + + let actual = ::list_contains( + &list, + sliced.as_::(), + &mut ctx, + )? + .ok_or_else(|| vortex_err!("BitPacked list_contains kernel declined a sliced input"))? + .execute::(&mut ctx)?; + let expected = BoolArray::from_iter(values[range].iter().map(|value| members.contains(value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::nullable_needles( + vec![Some(1), Some(3)], + Nullability::NonNullable, + vec![Some(1), None, Some(2)], + vec![Some(true), None, Some(false)], +)] +#[case::nullable_members( + vec![Some(1), None, Some(3)], + Nullability::Nullable, + vec![Some(1), Some(2), Some(3)], + vec![Some(true), Some(false), Some(true)], +)] +#[case::all_null_members( + vec![None, None], + Nullability::Nullable, + vec![Some(1), None, Some(2)], + vec![Some(false), None, Some(false)], +)] +fn test_null_semantics( + #[case] members: Vec>, + #[case] member_nullability: Nullability, + #[case] values: Vec>, + #[case] expected: Vec>, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_option_iter(values); + let packed = BitPackedData::encode(&primitive.into_array(), 3, &mut ctx)?; + let list = list_array(member_list(members, member_nullability), packed.len()); + + let actual = execute_direct(&list, &packed, &mut ctx)?; + let expected = BoolArray::from_iter(expected); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_wrong_integer_type_declines_without_panic() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_iter([1i32, 2, 3]); + let packed = BitPackedData::encode(&primitive.into_array(), 2, &mut ctx)?; + let list = list_array( + member_list([Some(1i64), Some(3)], Nullability::NonNullable), + packed.len(), + ); + + let result = + ::list_contains(&list, packed.as_view(), &mut ctx)?; + assert!(result.is_none()); + Ok(()) +} + +#[test] +fn test_nonconstant_list_declines() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let primitive = PrimitiveArray::from_iter([1i32, 2, 3]); + let packed = BitPackedData::encode(&primitive.into_array(), 2, &mut ctx)?; + let list = ListArray::from_iter_slow::( + vec![vec![1i32], vec![2], vec![3]], + Arc::new(DType::Primitive(i32::PTYPE, Nullability::NonNullable)), + )? + .into_array(); + + let result = + ::list_contains(&list, packed.as_view(), &mut ctx)?; + assert!(result.is_none()); + Ok(()) +} + +#[test] +#[cfg(not(codspeed))] +fn test_registered_kernel_executes_through_expression() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0..2_048).map(|value| value % 128).collect::>(); + let primitive = PrimitiveArray::from_iter(values.iter().copied()); + let packed = BitPackedData::encode(&primitive.into_array(), 7, &mut ctx)?; + let members = [0, 99]; + let expression = list_contains( + lit(member_list( + members.into_iter().map(Some), + Nullability::NonNullable, + )), + root(), + ); + let contains = packed.into_array().apply(&expression)?; + + let traced = trace_op(|| contains.execute::(&mut ctx))?; + let trace = traced.trace.to_string(); + let applied = trace + .lines() + .filter(|line| { + line.contains("child_execute_parent session[") + && line.contains("slot=1") + && line.contains("parent=vortex.list.contains") + && line.contains("child=fastlanes.bitpacked") + }) + .collect::>(); + // A silent fallback preserves values but loses compressed-domain execution. + assert_eq!(applied.len(), 1, "{trace}"); + + let expected = BoolArray::from_iter(values.into_iter().map(|value| members.contains(&value))); + assert_arrays_eq!(traced.output, expected, &mut ctx); + Ok(()) +} diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index 38f86f781bb..f5986711d73 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -7,6 +7,7 @@ mod compare; mod compare_fused; mod filter; pub(crate) mod is_constant; +pub(crate) mod list_contains; mod slice; mod stream_predicate; mod take; diff --git a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs index eb0dd9b7a23..9a0add2130b 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs @@ -16,6 +16,8 @@ use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_array::scalar_fn::fns::list_contains::ListContains; +use vortex_array::scalar_fn::fns::list_contains::ListContainsElementExecuteAdaptor; use vortex_session::VortexSession; use crate::BitPacked; @@ -36,4 +38,9 @@ pub(crate) fn initialize(session: &VortexSession) { kernels.register_execute_parent_kernel(Filter.id(), BitPacked, FilterExecuteAdaptor(BitPacked)); kernels.register_execute_parent_kernel(Slice.id(), BitPacked, SliceExecuteAdaptor(BitPacked)); kernels.register_execute_parent_kernel(Dict.id(), BitPacked, TakeExecuteAdaptor(BitPacked)); + kernels.register_execute_parent_kernel( + ListContains.id(), + BitPacked, + ListContainsElementExecuteAdaptor(BitPacked), + ); } diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index 80ffcad24cd..d2350a1c35a 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -6,9 +6,9 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::list_contains::ListContainsElementReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::array::Sequence; @@ -23,11 +23,19 @@ impl ListContainsElementReduce for Sequence { let Some(list_scalar) = list.as_constant() else { return Ok(None); }; + let DType::List(member_dtype, _) = list.dtype() else { + return Ok(None); + }; + if !member_dtype.eq_ignore_nullability(element.dtype()) { + return Ok(None); + } - let list_elements = list_scalar - .as_list() - .elements() - .vortex_expect("non-null element (checked in entry)"); + let Some(list_elements) = list_scalar.as_list().elements() else { + return Ok(None); + }; + if list_elements.is_empty() { + return Ok(None); + } let nullability = list.dtype().nullability() | element.dtype().nullability(); @@ -65,10 +73,13 @@ mod tests { use std::sync::Arc; use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Constant; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType::I32; use vortex_array::expr::list_contains; @@ -139,4 +150,32 @@ mod tests { let expected = BoolArray::from_iter([Some(true), Some(true), Some(true)]); assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } + + #[rstest] + #[case::null_list( + Scalar::null(DType::List(Arc::new(I32.into()), Nullability::Nullable)), + [None, None, None] + )] + #[case::empty_list( + Scalar::list(Arc::new(I32.into()), vec![], Nullability::Nullable), + [Some(false), Some(false), Some(false)] + )] + fn test_constant_list_semantics( + #[case] list_scalar: Scalar, + #[case] expected: [Option; 3], + ) { + let array = Sequence::try_new_typed(1i32, 1, Nullability::NonNullable, 3) + .unwrap() + .into_array(); + let expr = list_contains(lit(list_scalar), root()); + + let result = array.apply(&expr).unwrap(); + + assert!(result.is::()); + assert_arrays_eq!( + result, + BoolArray::from_iter(expected), + &mut SESSION.create_execution_ctx() + ); + } } diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index fb9ea651ad0..cd3dd8c11fe 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -290,6 +290,10 @@ harness = false name = "list_length" harness = false +[[bench]] +name = "list_contains" +harness = false + [[bench]] name = "list_sum" harness = false diff --git a/vortex-array/benches/list_contains.rs b/vortex-array/benches/list_contains.rs new file mode 100644 index 00000000000..9e67e13678b --- /dev/null +++ b/vortex-array/benches/list_contains.rs @@ -0,0 +1,195 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares the Primitive constant-list membership dispatch paths. +//! +//! Primitive arrays use direct comparisons for up to four distinct members. They use binary +//! search from 10 members for 8- and 16-bit integers. The 32- and 64-bit thresholds are 11 and 13 +//! members. Every path runs on each real CPU feature leg in CodSpeed. +//! To recalculate the thresholds, run this benchmark twice with temporary policy constants. Use a +//! high cutoff to force generic evaluation. Use `5` to force binary search above four members. +//! +//! Run with `cargo bench -p vortex-array --bench list_contains`. + +#![expect(clippy::unwrap_used)] + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hint::black_box; +use std::sync::Arc; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +trait BenchInt: IntegerPType + Copy + Into { + fn from_counter(value: u64) -> Self; +} + +impl BenchInt for u8 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u16 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u32 { + fn from_counter(value: u64) -> Self { + Self::try_from(value).unwrap() + } +} + +impl BenchInt for u64 { + fn from_counter(value: u64) -> Self { + value + } +} + +#[derive(Clone, Copy)] +struct PrimitiveCase { + name: &'static str, + len: usize, + member_count: usize, +} + +impl Display for PrimitiveCase { + fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "{}_m{}_n{}", + self.name, self.member_count, self.len + ) + } +} + +const fn primitive_case(name: &'static str, len: usize, member_count: usize) -> PrimitiveCase { + PrimitiveCase { + name, + len, + member_count, + } +} + +const LONG_M1: PrimitiveCase = primitive_case("long", 65_536, 1); +const LONG_M4: PrimitiveCase = primitive_case("long", 65_536, 4); +const LONG_M9: PrimitiveCase = primitive_case("long", 65_536, 9); +const LONG_M10: PrimitiveCase = primitive_case("long", 65_536, 10); +const LONG_M11: PrimitiveCase = primitive_case("long", 65_536, 11); +const LONG_M12: PrimitiveCase = primitive_case("long", 65_536, 12); +const LONG_M13: PrimitiveCase = primitive_case("long", 65_536, 13); +const LONG_M32: PrimitiveCase = primitive_case("long", 65_536, 32); +const SHORT_M11: PrimitiveCase = primitive_case("short", 1_024, 11); +const SHORT_M13: PrimitiveCase = primitive_case("short", 1_024, 13); + +const CURRENT_10: &[PrimitiveCase] = &[LONG_M1, LONG_M4, LONG_M9, LONG_M10, LONG_M32, SHORT_M11]; +const CURRENT_11: &[PrimitiveCase] = &[LONG_M1, LONG_M4, LONG_M10, LONG_M11, LONG_M32, SHORT_M11]; +const CURRENT_13: &[PrimitiveCase] = &[LONG_M1, LONG_M4, LONG_M12, LONG_M13, LONG_M32, SHORT_M13]; + +fn primitive_input( + case: PrimitiveCase, +) -> (PrimitiveArray, Scalar, BoolArray, VortexSession) { + let members = (0..case.member_count) + .map(|index| T::from_counter(u64::try_from(index).unwrap() * 2)) + .collect::>(); + let domain_bits = T::PTYPE.bit_width().min(12); + let domain_size = 1u64 << domain_bits; + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let generated = (0..case.len) + .map(|_| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + if (state >> 32).is_multiple_of(2) { + let member_index = + usize::try_from(state % u64::try_from(members.len()).unwrap()).unwrap(); + members[member_index] + } else { + let mut candidate = state.rotate_left(17) % domain_size; + while members.contains(&T::from_counter(candidate)) { + candidate = (candidate + 1) % domain_size; + } + T::from_counter(candidate) + } + }) + .collect::>(); + let expected = BoolArray::from_iter(generated.iter().map(|value| members.contains(value))); + let list = Scalar::list( + Arc::new(DType::Primitive(T::PTYPE, Nullability::NonNullable)), + members.iter().copied().map(Into::into).collect(), + Nullability::NonNullable, + ); + ( + PrimitiveArray::new::(generated, Validity::NonNullable), + list, + expected, + array_session(), + ) +} + +fn bench_current(bencher: Bencher, case: PrimitiveCase) { + let (array, list, expected, session) = primitive_input::(case); + let expression = list_contains(lit(list), root()); + let mut ctx = session.create_execution_ctx(); + let actual = array + .clone() + .into_array() + .apply(&expression) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + assert_arrays_eq!(actual, expected, &mut ctx); + + bencher.counter(ItemsCount::new(case.len)).bench_local(|| { + black_box( + array + .clone() + .into_array() + .apply(&expression) + .unwrap() + .execute::(&mut ctx) + .unwrap(), + ) + }); +} + +macro_rules! primitive_benchmarks { + ($type_name:ident, $ty:ty, $current:ident) => { + mod $type_name { + use super::*; + + #[vortex_bench_support::cpu_features] + #[divan::bench(args = $current)] + fn current(bencher: Bencher, case: PrimitiveCase) { + bench_current::<$ty>(bencher, case); + } + } + }; +} + +primitive_benchmarks!(u8_cases, u8, CURRENT_10); +primitive_benchmarks!(u16_cases, u16, CURRENT_10); +primitive_benchmarks!(u32_cases, u32, CURRENT_11); +primitive_benchmarks!(u64_cases, u64, CURRENT_13); diff --git a/vortex-array/src/arrays/primitive/compute/list_contains.rs b/vortex-array/src/arrays/primitive/compute/list_contains.rs new file mode 100644 index 00000000000..eb3965b5167 --- /dev/null +++ b/vortex-array/src/arrays/primitive/compute/list_contains.rs @@ -0,0 +1,329 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ArrayView; +use crate::ExecutionCtx; +use crate::arrays::Primitive; +use crate::dtype::IntegerPType; +use crate::dtype::PType; +use crate::match_each_integer_ptype; +use crate::scalar_fn::fns::list_contains::IntegerMembership; +use crate::scalar_fn::fns::list_contains::ListContainsElementKernel; +use crate::scalar_fn::fns::list_contains::constant_list_scalar_contains; + +/// Returns the source-member count where Primitive integer membership uses binary search. +#[doc(hidden)] +pub fn integer_membership_binary_search_min(ptype: PType) -> usize { + // The generic implementation evaluates one equality expression per source member. Use the + // prepared set once binary search becomes faster than the expression tree. + match ptype.bit_width() { + 8 | 16 => 10, + 32 => 11, + 64 => 13, + _ => 13, + } +} + +impl ListContainsElementKernel for Primitive { + fn list_contains( + list: &ArrayRef, + element: ArrayView<'_, Self>, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + evaluate_constant_list_membership(list, element) + } +} + +fn evaluate_constant_list_membership( + list: &ArrayRef, + element: ArrayView<'_, Primitive>, +) -> VortexResult> { + if !element.ptype().is_int() { + return Ok(None); + } + + let nullability = list.dtype().nullability() | element.dtype().nullability(); + + match_each_integer_ptype!(element.ptype(), |T| { + evaluate_integer_membership::(list, element, nullability) + }) +} + +fn evaluate_integer_membership( + list: &ArrayRef, + element: ArrayView<'_, Primitive>, + nullability: crate::dtype::Nullability, +) -> VortexResult> { + let Some(membership) = IntegerMembership::::try_from_constant_list(list, element.dtype())? + else { + return Ok(None); + }; + evaluate_prepared_integer_membership(membership, element, nullability).map(Some) +} + +/// Evaluates a prepared integer set against a Primitive array. +#[doc(hidden)] +pub fn evaluate_prepared_integer_membership( + membership: IntegerMembership, + element: ArrayView<'_, Primitive>, + nullability: crate::dtype::Nullability, +) -> VortexResult { + if membership.members().len() > 4 + && membership.non_null_source_len() < integer_membership_binary_search_min(element.ptype()) + { + return constant_list_scalar_contains( + &membership.source_list().as_list(), + element.array(), + nullability, + ); + } + membership.evaluate_primitive(element, nullability) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_error::VortexExpect; + + use super::*; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::BoolArray; + use crate::arrays::Constant; + use crate::arrays::ConstantArray; + use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType::F32; + use crate::dtype::PType::I32; + use crate::dtype::PType::I64; + use crate::expr::list_contains; + use crate::expr::lit; + use crate::expr::root; + use crate::scalar::Scalar; + #[cfg(not(codspeed))] + use crate::test_harness::trace::trace_op; + + fn list(values: impl IntoIterator, len: usize) -> ArrayRef { + ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + values + .into_iter() + .map(|value| Scalar::primitive(value, Nullability::NonNullable)) + .collect(), + Nullability::NonNullable, + ), + len, + ) + .into_array() + } + + #[rstest] + #[case::one(vec![3])] + #[case::two(vec![3, 7])] + #[case::three(vec![3, 7, 11])] + #[case::four(vec![3, 7, 11, 15])] + #[case::five((0..5).map(|value| value * 3).collect())] + #[case::eleven((0..11).map(|value| value * 3).collect())] + #[case::many((0..32).map(|value| value * 3).collect())] + #[case::duplicate_heavy((0..32).map(|value| value % 5).collect())] + fn test_membership_plans(#[case] members: Vec) -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let values = [0, 3, 7, 15, 31, 90_000, 310_000]; + let element = PrimitiveArray::from_iter(values); + let expected = BoolArray::from_iter(values.map(|value| members.contains(&value))); + + let actual = ::list_contains( + &list(members, element.len()), + element.as_view(), + &mut ctx, + )? + .vortex_expect("integer constant-list membership is supported"); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::small(5)] + #[case::many(13)] + fn test_i64_membership(#[case] member_count: usize) -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let members = (0..member_count) + .map(|value| i64::try_from(value).vortex_expect("member count fits i64")) + .collect::>(); + let values = [0i64, 11, 99]; + let element = PrimitiveArray::from_iter(values); + let list = ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(I64, Nullability::NonNullable)), + members.iter().copied().map(Scalar::from).collect(), + Nullability::NonNullable, + ), + element.len(), + ) + .into_array(); + + let actual = ::list_contains( + &list, + element.as_view(), + &mut ctx, + )? + .vortex_expect("integer constant-list membership is supported"); + let expected = BoolArray::from_iter(values.map(|value| members.contains(&value))); + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + #[cfg(not(codspeed))] + fn test_registered_kernel_executes_through_expression() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let values = [0i32, 1, 2, 3]; + let element = PrimitiveArray::from_iter(values); + let members = [1, 3]; + let contains = element.into_array().apply(&list_contains( + lit(list(members, values.len()) + .as_constant() + .vortex_expect("constant list")), + root(), + ))?; + + let traced = trace_op(|| contains.execute::(&mut ctx))?; + let trace = traced.trace.to_string(); + let applied = trace + .lines() + .filter(|line| { + line.contains("child_execute_parent session[") + && line.contains("slot=1") + && line.contains("parent=vortex.list.contains") + && line.contains("child=vortex.primitive") + }) + .collect::>(); + // A silent fallback preserves values but loses the membership optimization. + assert_eq!(applied.len(), 1, "{trace}"); + + let expected = BoolArray::from_iter(values.map(|value| members.contains(&value))); + assert_arrays_eq!(traced.output, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_float_falls_back_through_expression() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let values = [1.5f32, 2.5, 3.5]; + let element = PrimitiveArray::from_iter(values); + let members = [1.5f32, 3.5]; + let list = ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(F32, Nullability::NonNullable)), + members.into_iter().map(Scalar::from).collect(), + Nullability::NonNullable, + ), + element.len(), + ) + .into_array(); + let list_scalar = list.as_constant().vortex_expect("list is constant"); + + let actual = element + .into_array() + .apply(&list_contains(lit(list_scalar), root()))? + .execute::(&mut ctx)?; + let expected = BoolArray::from_iter(values.map(|value| members.contains(&value))); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn test_null_needles() -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let element = PrimitiveArray::from_option_iter([Some(1), None, Some(2)]); + let expected = BoolArray::from_iter([Some(true), None, Some(false)]); + + let actual = ::list_contains( + &list([1, 3], element.len()), + element.as_view(), + &mut ctx, + )? + .vortex_expect("integer constant-list membership is supported"); + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::null_list(true)] + #[case::empty_list(false)] + fn test_constant_list_adaptor(#[case] null_list: bool) -> VortexResult<()> { + let member_dtype = DType::Primitive(I32, Nullability::NonNullable); + let list = if null_list { + Scalar::null(DType::List(Arc::new(member_dtype), Nullability::Nullable)) + } else { + Scalar::list(Arc::new(member_dtype), vec![], Nullability::NonNullable) + }; + let needles = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); + + let mut ctx = crate::array_session().create_execution_ctx(); + let contains = needles + .apply(&list_contains(lit(list), root()))? + .execute::(&mut ctx)?; + let expected = if null_list { + BoolArray::from_iter([None, None, None]) + } else { + BoolArray::from_iter([Some(false), Some(false), Some(false)]) + }; + + assert!(contains.is::()); + assert_arrays_eq!(contains, expected, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::mixed( + vec![Some(1), None, Some(3)], + [Some(true), None, Some(true)] + )] + #[case::all_null(vec![None, None], [Some(false), None, Some(false)])] + fn test_nullable_members( + #[case] members: Vec>, + #[case] expected: [Option; 3], + ) -> VortexResult<()> { + let mut ctx = crate::array_session().create_execution_ctx(); + let member_dtype = DType::Primitive(I32, Nullability::Nullable); + let list = ConstantArray::new( + Scalar::list( + Arc::new(member_dtype.clone()), + members + .into_iter() + .map(|member| { + member + .map(|value| Scalar::primitive(value, Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(member_dtype.clone())) + }) + .collect(), + Nullability::NonNullable, + ), + 3, + ) + .into_array(); + let element = PrimitiveArray::from_option_iter([Some(1), None, Some(3)]); + + let actual = ::list_contains( + &list, + element.as_view(), + &mut ctx, + )? + .vortex_expect("integer constant-list membership is supported"); + + assert_arrays_eq!(actual, BoolArray::from_iter(expected), &mut ctx); + Ok(()) + } +} diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index 382b42ee6e2..7769def7df5 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -5,6 +5,9 @@ mod between; mod cast; mod fill_null; mod fixed_width; +mod list_contains; +pub use list_contains::evaluate_prepared_integer_membership; +pub use list_contains::integer_membership_binary_search_min; mod mask; pub(crate) mod rules; mod slice; diff --git a/vortex-array/src/arrays/primitive/mod.rs b/vortex-array/src/arrays/primitive/mod.rs index 748beda8339..b99ec51c0f8 100644 --- a/vortex-array/src/arrays/primitive/mod.rs +++ b/vortex-array/src/arrays/primitive/mod.rs @@ -14,6 +14,10 @@ pub use vtable::PrimitiveArray; pub(crate) mod compute; mod vtable; +#[doc(hidden)] +pub use compute::evaluate_prepared_integer_membership; +#[doc(hidden)] +pub use compute::integer_membership_binary_search_min; pub use compute::rules::PrimitiveMaskedValidityRule; pub use vtable::Primitive; diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 6382ea73794..3f13282c334 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -15,6 +15,8 @@ use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; +use crate::scalar_fn::fns::list_contains::ListContains; +use crate::scalar_fn::fns::list_contains::ListContainsElementExecuteAdaptor; use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; @@ -31,6 +33,11 @@ pub(crate) fn initialize(session: &VortexSession) { Primitive, FillNullExecuteAdaptor(Primitive), ); + kernels.register_execute_parent_kernel( + ListContains.id(), + Primitive, + ListContainsElementExecuteAdaptor(Primitive), + ); kernels.register_execute_parent_kernel(Dict.id(), Primitive, TakeExecuteAdaptor(Primitive)); kernels.register_execute_parent_kernel(Zip.id(), Primitive, ZipExecuteAdaptor(Primitive)); } diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index fb8bfe227aa..16dcd3256a7 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -1100,6 +1100,8 @@ pub fn bound_dynamic( /// Creates an expression that checks if a value is contained in a list. /// /// Returns a boolean array indicating whether the value appears in each list. +/// A null list produces null. An empty list produces false, including for a null value. +/// A null value produces null for a nonempty list. Null list members do not match any value. /// /// ```rust /// # use vortex_array::expr::{list_contains, lit, root}; diff --git a/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs b/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs new file mode 100644 index 00000000000..1e206870a64 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/list_contains/integer_membership.rs @@ -0,0 +1,143 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ArrayView; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Primitive; +use crate::dtype::DType; +use crate::dtype::IntegerPType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; + +/// A prepared integer set for constant-list membership kernels. +/// +/// The set sorts and deduplicates its members. +#[doc(hidden)] +pub struct IntegerMembership { + members: Box<[T]>, + non_null_source_len: usize, + source_list: Scalar, +} + +impl IntegerMembership { + fn new(mut members: Vec, source_list: Scalar) -> Self { + let non_null_source_len = members.len(); + members.sort_unstable(); + members.dedup(); + Self { + members: members.into_boxed_slice(), + non_null_source_len, + source_list, + } + } + + /// Extracts an integer set from a compatible constant list. + pub fn try_from_constant_list( + list: &ArrayRef, + element_dtype: &DType, + ) -> VortexResult> { + let Some(list_scalar) = list.as_constant() else { + return Ok(None); + }; + let DType::List(member_dtype, _) = list.dtype() else { + return Ok(None); + }; + if !member_dtype.eq_ignore_nullability(element_dtype) { + return Ok(None); + } + let Some(elements) = list_scalar.as_list().elements() else { + return Ok(None); + }; + + let members = elements + .iter() + .map(|value| { + value + .as_primitive_opt() + .vortex_expect("list member type was checked") + .try_typed_value::() + }) + .collect::>>>()? + .into_iter() + .flatten() + .collect(); + Ok(Some(Self::new(members, list_scalar))) + } + + /// Returns the prepared members. + pub fn members(&self) -> &[T] { + &self.members + } + + /// Returns the number of non-null source members before deduplication. + #[doc(hidden)] + pub fn non_null_source_len(&self) -> usize { + self.non_null_source_len + } + + pub(crate) fn source_list(&self) -> &Scalar { + &self.source_list + } + + /// Tests whether the prepared set contains `value`. + pub(crate) fn contains(&self, value: T) -> bool { + self.members.binary_search(&value).is_ok() + } + + /// Evaluates this set against a primitive array of the same integer type. + pub(crate) fn evaluate_primitive( + self, + element: ArrayView<'_, Primitive>, + nullability: Nullability, + ) -> VortexResult { + vortex_ensure!( + element.ptype() == T::PTYPE, + "Membership type {} does not match array type {}", + T::PTYPE, + element.ptype(), + ); + let values = element.as_slice::(); + let bits = match self.members() { + [] => BitBuffer::new_unset(values.len()), + [member] => collect_direct(values, move |value| value.is_eq(*member)), + [first, second] => collect_direct(values, move |value| { + value.is_eq(*first) | value.is_eq(*second) + }), + [first, second, third] => collect_direct(values, move |value| { + value.is_eq(*first) | value.is_eq(*second) | value.is_eq(*third) + }), + [first, second, third, fourth] => collect_direct(values, move |value| { + value.is_eq(*first) + | value.is_eq(*second) + | value.is_eq(*third) + | value.is_eq(*fourth) + }), + _ => collect_many(values, &self), + }; + + Ok(BoolArray::new(bits, element.validity()?.union_nullability(nullability)).into_array()) + } +} + +fn collect_direct(values: &[T], mut predicate: impl FnMut(T) -> bool) -> BitBuffer { + BitBuffer::collect_bool_multiversioned(values.len(), |index| { + // SAFETY: collect_bool_multiversioned visits each valid index once. + predicate(unsafe { *values.get_unchecked(index) }) + }) +} + +fn collect_many(values: &[T], membership: &IntegerMembership) -> BitBuffer { + BitBuffer::collect_bool(values.len(), |index| { + // SAFETY: collect_bool visits each valid index once. + let value = unsafe { *values.get_unchecked(index) }; + membership.contains(value) + }) +} diff --git a/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs b/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs index 563600bfeee..fc50c8e70ac 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/kernel.rs @@ -6,16 +6,42 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::IntoArray; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::ConstantArray; use crate::arrays::ScalarFn; use crate::arrays::scalar_fn::ExactScalarFn; use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::arrays::scalar_fn::ScalarFnArrayView; +use crate::dtype::DType; use crate::kernel::ExecuteParentKernel; use crate::optimizer::rules::ArrayParentReduceRule; +use crate::scalar::Scalar; use crate::scalar_fn::fns::list_contains::ListContains as ListContainsExpr; +fn constant_list_result( + list: &ArrayRef, + element_len: usize, + element_nullability: crate::dtype::Nullability, +) -> Option { + let list_scalar = list.as_constant()?; + let DType::List(_, list_nullability) = list.dtype() else { + return None; + }; + let nullability = *list_nullability | element_nullability; + + match list_scalar.as_list().elements() { + None => Some( + ConstantArray::new(Scalar::null(DType::Bool(nullability)), element_len).into_array(), + ), + Some(elements) if elements.is_empty() => { + Some(ConstantArray::new(Scalar::bool(false, nullability), element_len).into_array()) + } + Some(_) => None, + } +} + /// Check list-contains without reading buffers (metadata-only). /// /// This trait dispatches on the **element** (needle) child at index 1 of the `ListContains` @@ -25,6 +51,8 @@ use crate::scalar_fn::fns::list_contains::ListContains as ListContainsExpr; /// A future `ListContainsListReduce` could dispatch on the list side (child 0) for encodings /// with specialized list representations. /// +/// The parent adaptor resolves null and empty constant lists before delegation. +/// /// Return `None` if the operation cannot be resolved from metadata alone. pub trait ListContainsElementReduce: VTable { fn list_contains( @@ -38,6 +66,8 @@ pub trait ListContainsElementReduce: VTable { /// Like [`ListContainsElementReduce`], this dispatches on the **element** (needle) child at /// index 1. Unlike the reduce variant, implementations may read and execute on buffers via /// the provided [`ExecutionCtx`]. +/// +/// The parent adaptor resolves null and empty constant lists before delegation. pub trait ListContainsElementKernel: VTable { fn list_contains( list: &ArrayRef, @@ -70,6 +100,9 @@ where .as_opt::() .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray"); let list = scalar_fn_array.get_child(0); + if let Some(result) = constant_list_result(list, array.len(), array.dtype().nullability()) { + return Ok(Some(result)); + } ::list_contains(list, array) } } @@ -99,6 +132,9 @@ where .as_opt::() .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray"); let list = scalar_fn_array.get_child(0); + if let Some(result) = constant_list_result(list, array.len(), array.dtype().nullability()) { + return Ok(Some(result)); + } ::list_contains(list, array, ctx) } } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index d2508014089..df1e37ef044 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -1,11 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod integer_membership; mod kernel; use std::ops::BitOr; use arrow_buffer::bit_iterator::BitIndexIterator; +pub use integer_membership::IntegerMembership; pub use kernel::*; use num_traits::Zero; use vortex_buffer::BitBuffer; @@ -13,6 +15,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; use vortex_utils::iter::ReduceBalancedIterExt; @@ -56,7 +59,8 @@ impl ListContains { /// /// # Errors /// - /// Returns an error if the children have different lengths or `list` is not a list array. + /// Returns an error if the children have different lengths, `list` is not a list array, or + /// the list member type differs from the needle type. pub fn try_new(list: ArrayRef, needle: ArrayRef) -> VortexResult { ScalarFnArray::try_new(ListContains.bind(EmptyOptions), vec![list, needle]) } @@ -100,16 +104,17 @@ impl ScalarFnVTable for ListContains { let list_dtype = &arg_dtypes[0]; let needle_dtype = &arg_dtypes[1]; - let nullability = match list_dtype { - DType::List(_, list_nullability) => list_nullability, - _ => { - vortex_bail!( - "First argument to ListContains must be a List, got {:?}", - list_dtype - ); - } + let DType::List(member_dtype, list_nullability) = list_dtype else { + vortex_bail!("First argument to ListContains must be a List, got {list_dtype}"); + }; + if !member_dtype.eq_ignore_nullability(needle_dtype) { + vortex_bail!( + "Element type {} of list does not match search value {}", + member_dtype, + needle_dtype + ); } - .bitor(needle_dtype.nullability()); + let nullability = list_nullability.bitor(needle_dtype.nullability()); Ok(DType::Bool(nullability)) } @@ -146,8 +151,7 @@ impl ScalarFnVTable for ListContains { fn compute_contains_scalar(list: &Scalar, needle: &Scalar) -> VortexResult { let nullability = list.dtype().nullability() | needle.dtype().nullability(); - // Handle null list or null needle - if list.is_null() || needle.is_null() { + if list.is_null() { return Ok(Scalar::null(DType::Bool(nullability))); } @@ -155,6 +159,12 @@ fn compute_contains_scalar(list: &Scalar, needle: &Scalar) -> VortexResult(ctx)?; + return list_false_if_empty_else_null(&list_array, nullability, ctx); + } + if let Some(value_scalar) = value.as_constant() { list_contains_scalar(array, &value_scalar, nullability, ctx) } else if let Some(list_scalar) = array.as_constant() { @@ -196,7 +211,7 @@ fn compute_list_contains( } /// There is a constant list scalar (haystack) being compared to an array of needles. -fn constant_list_scalar_contains( +pub(crate) fn constant_list_scalar_contains( list_scalar: &ListScalar, values: &ArrayRef, nullability: Nullability, @@ -206,8 +221,13 @@ fn constant_list_scalar_contains( let len = values.len(); let false_scalar = Scalar::bool(false, nullability); + if elements.is_empty() { + return Ok(ConstantArray::new(false_scalar, len).into_array()); + } + let result = elements .iter() + .filter(|element| !element.is_null()) .map(|element| { Binary::try_new( ConstantArray::new(element.clone(), len).into_array(), @@ -221,7 +241,12 @@ fn constant_list_scalar_contains( .into_iter() .try_reduce_balanced(|acc, res| acc.binary(res, Operator::Or))?; - Ok(result.unwrap_or_else(|| ConstantArray::new(false_scalar, len).into_array())) + let result = result.unwrap_or_else(|| ConstantArray::new(false_scalar, len).into_array()); + if values.dtype().is_nullable() { + result.mask(values.validity()?.to_array(len)) + } else { + Ok(result) + } } /// Returns a [`BoolArray`] where each bit represents if a list contains the scalar. @@ -244,6 +269,9 @@ fn list_contains_scalar( // Must return false when a list is empty (but valid), or null when the list itself is null. return list_false_or_null(&list_array, nullability); } + if value.is_null() { + return list_false_if_empty_else_null(&list_array, nullability, ctx); + } let rhs = ConstantArray::new(value.clone(), elems.len()); let matching_elements = @@ -266,13 +294,7 @@ fn list_contains_scalar( list_false_or_null(&list_array, nullability) } // No elements match, and all comparisons are valid (result in `false`). - Some(false) => { - // False, but match the nullability to the input list array. - Ok( - ConstantArray::new(Scalar::bool(false, nullability), list_array.len()) - .into_array(), - ) - } + Some(false) => list_false_or_null(&list_array, nullability), // All elements match, and all comparisons are valid (result in `true`). Some(true) => { // True, unless the list itself is empty or NULL. @@ -294,9 +316,9 @@ fn list_contains_scalar( // Process based on the offset and size types. let list_matches = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { match_each_unsigned_integer_ptype!(sizes.ptype(), |S| { - process_matches::(matches, list_array.len(), offsets, sizes) + process_matches::(&matches, list_array.len(), offsets, sizes, ctx) }) - }); + })?; Ok(BoolArray::new( list_matches, @@ -305,33 +327,58 @@ fn list_contains_scalar( .into_array()) } +/// Returns false for valid empty lists and null for all other lists. +fn list_false_if_empty_else_null( + list_array: &ListViewArray, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let sizes = list_array.sizes().clone().execute::(ctx)?; + let empty = match_each_integer_ptype!(sizes.ptype(), |S| { + Mask::from_iter(sizes.as_slice::().iter().map(|size| size.is_zero())) + }); + let valid = list_array.validity()?.execute_mask(list_array.len(), ctx)? & ∅ + + Ok(BoolArray::new( + BitBuffer::new_unset(list_array.len()), + Validity::from_mask(valid, nullability), + ) + .into_array()) +} + /// Returns a [`BitBuffer`] where each bit represents if a list contains the scalar, derived from a /// [`BoolArray`] of matches on the child elements array. fn process_matches( - matches: BoolArray, + matches: &BoolArray, list_array_len: usize, offsets: PrimitiveArray, sizes: PrimitiveArray, -) -> BitBuffer + ctx: &mut ExecutionCtx, +) -> VortexResult where O: IntegerPType, S: IntegerPType, { let offsets_slice = offsets.as_slice::(); let sizes_slice = sizes.as_slice::(); - let bits = matches.bit_buffer_view(); + let value_bits = matches.to_bit_buffer(); + let valid_matches = match matches.validity()? { + Validity::NonNullable | Validity::AllValid => value_bits, + Validity::AllInvalid => BitBuffer::new_unset(matches.len()), + validity => value_bits & validity.execute_mask(matches.len(), ctx)?.into_bit_buffer(), + }; - (0..list_array_len) + Ok((0..list_array_len) .map(|i| { let offset = offsets_slice[i].as_(); let size = sizes_slice[i].as_(); // BitIndexIterator yields indices of true bits only. If `.next()` returns // `Some(_)`, at least one element in this list's range matches. - let mut set_bits = BitIndexIterator::new(bits.inner(), offset, size); + let mut set_bits = BitIndexIterator::new(valid_matches.inner(), offset, size); set_bits.next().is_some() }) - .collect::() + .collect::()) } /// Returns a `Bool` array with `false` for lists that are valid, @@ -414,6 +461,8 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::Dict; + use crate::arrays::DictArray; use crate::arrays::ListArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; @@ -435,6 +484,7 @@ mod tests { use crate::scalar::Scalar; use crate::scalar_fn::fns::list_contains::BoolArray; use crate::scalar_fn::fns::list_contains::ConstantArray; + use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_contains::ListViewArray; use crate::scalar_fn::fns::list_contains::PrimitiveArray; use crate::stats::StatsSession; @@ -540,8 +590,10 @@ mod tests { ); } - #[test] - pub fn test_nullable() { + #[rstest] + #[case::match_present(2, Some(true))] + #[case::match_absent(4, Some(false))] + pub fn test_nullable(#[case] needle: i32, #[case] expected_first: Option) { let arr = ListArray::try_new( PrimitiveArray::from_iter(vec![1, 1, 2, 2, 2]).into_array(), PrimitiveArray::from_iter(vec![0, 5, 5]).into_array(), @@ -550,18 +602,13 @@ mod tests { .unwrap() .into_array(); - let expr = list_contains(root(), lit(2)); + let expr = list_contains(root(), lit(needle)); let item = arr.apply(&expr).unwrap(); - assert_eq!( - item.execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), - Scalar::bool(true, Nullability::Nullable) - ); - assert!( - !item - .is_valid(1, &mut array_session().create_execution_ctx()) - .unwrap() + assert_arrays_eq!( + item, + BoolArray::from_iter([expected_first, None]), + &mut array_session().create_execution_ctx() ); } @@ -587,6 +634,54 @@ mod tests { ); } + #[test] + fn test_return_type_rejects_mismatched_member_type() { + let list = ConstantArray::new( + Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + vec![], + Nullability::NonNullable, + ), + 1, + ) + .into_array(); + let needle = + ConstantArray::new(Scalar::utf8("needle", Nullability::NonNullable), 1).into_array(); + + let error = ListContains::try_new(list, needle).unwrap_err(); + + assert!( + error + .to_string() + .contains("Element type i32 of list does not match search value utf8") + ); + } + + #[test] + fn test_dictionary_needles_preserve_dictionary_pushdown() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); + let codes = PrimitiveArray::from_iter([0u8, 1, 2, 0]).into_array(); + let needles = DictArray::try_new(codes, values)?.into_array(); + let list = Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + vec![1.into(), 3.into()], + Nullability::NonNullable, + ); + let contains = needles.apply(&list_contains(lit(list), root()))?; + + // Dictionary preservation avoids materializing repeated needle values. + assert!(contains.is::()); + let actual = contains.execute::(&mut ctx)?; + + assert_arrays_eq!( + actual, + BoolArray::from_iter([true, false, true, true]), + &mut ctx + ); + Ok(()) + } + #[test] pub fn list_falsification() -> VortexResult<()> { let expr = list_contains( @@ -639,36 +734,47 @@ mod tests { assert_eq!(expr2.to_string(), "vortex.list.contains($, 42i32)"); } - #[test] - pub fn test_constant_scalars() { - let arr = test_array(); - - // Both list and needle are constants - should use scalar optimization - let list_scalar = Scalar::list( - Arc::new(DType::Primitive(I32, Nullability::NonNullable)), - vec![1.into(), 2.into(), 3.into()], - Nullability::NonNullable, - ); - - // Test contains true - let expr = list_contains(lit(list_scalar.clone()), lit(2i32)); - let result = arr.clone().apply(&expr).unwrap(); - assert_eq!( - result - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), - Scalar::bool(true, Nullability::NonNullable) - ); + #[rstest] + #[case::present(false, vec![1, 2, 3], Some(2), Some(true))] + #[case::absent(false, vec![1, 2, 3], Some(42), Some(false))] + #[case::null_list(true, vec![], Some(1), None)] + #[case::empty_list_null_needle(false, vec![], None, Some(false))] + #[case::nonempty_list_null_needle(false, vec![1], None, None)] + fn test_constant_scalar_null_semantics( + #[case] null_list: bool, + #[case] members: Vec, + #[case] needle: Option, + #[case] expected: Option, + ) -> VortexResult<()> { + let member_dtype = DType::Primitive(I32, Nullability::NonNullable); + let list_dtype = DType::List(Arc::new(member_dtype.clone()), Nullability::Nullable); + let list = if null_list { + Scalar::null(list_dtype) + } else { + Scalar::list( + Arc::new(member_dtype), + members.into_iter().map(Scalar::from).collect(), + Nullability::Nullable, + ) + }; + let needle = needle + .map(|value| Scalar::primitive(value, Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(DType::Primitive(I32, Nullability::Nullable))); + let expected = expected + .map(|value| Scalar::bool(value, Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(DType::Bool(Nullability::Nullable))); + + let contains = ListContains::try_new( + ConstantArray::new(list, 1).into_array(), + ConstantArray::new(needle, 1).into_array(), + )? + .into_array(); - // Test contains false - let expr = list_contains(lit(list_scalar), lit(42i32)); - let result = arr.apply(&expr).unwrap(); assert_eq!( - result - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), - Scalar::bool(false, Nullability::NonNullable) + contains.execute_scalar(0, &mut array_session().create_execution_ctx())?, + expected ); + Ok(()) } // -- Tests migrated from compute/list_contains.rs -- @@ -749,7 +855,7 @@ mod tests { #[case( null_strings(vec![vec![], vec![None, None], vec![None, None, None]]), None, - bool_array(vec![false, true, true], Validity::AllInvalid) + BoolArray::from_iter([Some(false), None, None]) )] #[case( null_strings(vec![vec![], vec![None, None], vec![None, None, None]]), @@ -777,23 +883,47 @@ mod tests { assert_arrays_eq!(result, expected, &mut ctx); } - #[test] - fn test_constant_list() { + #[rstest] + #[case::empty( + Vec::>::new(), + [Some(false), Some(false), Some(false)] + )] + #[case::nonempty( + vec![Some("a"), Some("c")], + [Some(true), None, Some(false)] + )] + #[case::all_null( + vec![None, None], + [Some(false), None, Some(false)] + )] + fn test_constant_string_list_nullable_needles( + #[case] members: Vec>, + #[case] expected: [Option; 3], + ) { let mut ctx = array_session().create_execution_ctx(); - let list_array = ConstantArray::new( - Scalar::list( - Arc::new(DType::Primitive(I32, Nullability::NonNullable)), - vec![1i32.into(), 2i32.into(), 3i32.into()], - Nullability::NonNullable, - ), - 2, + let member_dtype = DType::Utf8(Nullability::Nullable); + let list = Scalar::list( + Arc::new(member_dtype.clone()), + members + .into_iter() + .map(|member| { + member + .map(|value| Scalar::utf8(value, Nullability::Nullable)) + .unwrap_or_else(|| Scalar::null(member_dtype.clone())) + }) + .collect(), + Nullability::NonNullable, + ); + let needles = VarBinArray::from_iter( + [Some("a"), None, Some("b")], + DType::Utf8(Nullability::Nullable), ) .into_array(); - let expr = list_contains(root(), lit(2i32)); - let contains = list_array.apply(&expr).unwrap(); - let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(contains, expected, &mut ctx); + let result = needles.apply(&list_contains(lit(list), root())).unwrap(); + let expected = BoolArray::from_iter(expected); + + assert_arrays_eq!(result, expected, &mut ctx); } #[test] @@ -818,6 +948,27 @@ mod tests { assert_arrays_eq!(contains, expected, &mut ctx); } + #[test] + fn test_nonconstant_all_null_needles() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let lists = ListArray::try_new( + PrimitiveArray::from_iter([1i32]).into_array(), + PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(), + Validity::Array(BoolArray::from(BitBuffer::from(vec![true, true, false])).into_array()), + )? + .into_array(); + let needles = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); + + let contains = ListContains::try_new(lists, needles)?.into_array(); + + assert_arrays_eq!( + contains, + BoolArray::from_iter([Some(false), None, None]), + &mut ctx + ); + Ok(()) + } + #[test] fn test_list_array_element() { let mut ctx = array_session().create_execution_ctx(); @@ -887,8 +1038,9 @@ mod tests { ); assert_arrays_eq!(result, expected, &mut ctx); - // Searching for non-null - let expr2 = list_contains(root(), lit(42i32)); + // Null primitive payloads default to zero. Searching for zero verifies that invalid + // comparison values do not become matches. + let expr2 = list_contains(root(), lit(0i32)); let result2 = list_array.into_array().apply(&expr2).unwrap(); let expected2 = BoolArray::from_iter([false, false, false]);