diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index f6b06544baf..c4f36c422b3 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -287,6 +287,10 @@ harness = false name = "filter_list" harness = false +[[bench]] +name = "list_contains" +harness = false + [[bench]] name = "list_length" harness = false diff --git a/vortex-array/benches/list_contains.rs b/vortex-array/benches/list_contains.rs new file mode 100644 index 00000000000..975f0a10611 --- /dev/null +++ b/vortex-array/benches/list_contains.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::ArrayRef; +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::expr::list_contains; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::scalar::Scalar; +use vortex_session::VortexSession; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +/// One engine-sized chunk against a handful of representative literal `IN`-list sizes, all at or +/// past `MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP` so `sorted_merge` genuinely takes the fast path +/// (below that threshold both benches run the identical fan-out, which isn't an interesting +/// comparison -- that crossover was measured separately to pick the threshold, see +/// `MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP`'s doc comment). The fan-out (`col IN (a, b, c, ...)` as +/// one `Eq` pass per literal, `Or`-reduced) is `O(list_len * ROWS)`; the sorted merge is +/// `O(ROWS log list_len)` plus a one-time `O(list_len log list_len)` sort of the literal set. +/// Both are benchmarked against the same sorted column, gated only by whether `Stat::IsSorted` +/// has been computed on it, matching how `ListContains::execute` actually chooses between them. +const ROWS: usize = 8_192; +const LIST_LENS: &[usize] = &[16, 64, 256]; + +fn list_scalar(list_len: usize) -> Scalar { + Scalar::list( + std::sync::Arc::new(vortex_array::dtype::DType::Primitive( + vortex_array::dtype::PType::I64, + vortex_array::dtype::Nullability::NonNullable, + )), + // Sparse, non-adjacent literals spread across the probed range so every comparison + // in the fan-out genuinely has to run (no early all-true/all-false short circuit). + (0..list_len) + .map(|i| Scalar::from((i * (ROWS / list_len.max(1))) as i64)) + .collect(), + vortex_array::dtype::Nullability::NonNullable, + ) +} + +fn sorted_column() -> ArrayRef { + let arr = PrimitiveArray::from_iter(0_i64..ROWS as i64).into_array(); + arr.statistics() + .compute_is_sorted(&mut SESSION.create_execution_ctx()); + arr +} + +#[divan::bench(args = LIST_LENS)] +fn fanout(bencher: Bencher, list_len: usize) { + // A column that has never had `Stat::IsSorted` computed: `ListContains::execute` takes the + // pre-existing equality fan-out. + let column = PrimitiveArray::from_iter(0_i64..ROWS as i64).into_array(); + let expr = list_contains(lit(list_scalar(list_len)), root()); + bencher + .counter(ItemsCount::new(ROWS)) + .with_inputs(|| (column.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(column, ctx)| { + column + .clone() + .apply(&expr) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} + +#[divan::bench(args = LIST_LENS)] +fn sorted_merge(bencher: Bencher, list_len: usize) { + let column = sorted_column(); + let expr = list_contains(lit(list_scalar(list_len)), root()); + bencher + .counter(ItemsCount::new(ROWS)) + .with_inputs(|| (column.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(column, ctx)| { + column + .clone() + .apply(&expr) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} diff --git a/vortex-array/benches/search_sorted.rs b/vortex-array/benches/search_sorted.rs index 370290bf1c9..7879e87776b 100644 --- a/vortex-array/benches/search_sorted.rs +++ b/vortex-array/benches/search_sorted.rs @@ -4,12 +4,26 @@ #![expect(clippy::unwrap_used)] use divan::Bencher; +use divan::counter::ItemsCount; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::search_sorted::NullEquality; use vortex_array::search_sorted::SearchSorted; use vortex_array::search_sorted::SearchSortedSide; +use vortex_array::search_sorted::SortedArray; +use vortex_array::search_sorted::SortedDirection; +use vortex_array::search_sorted::SortedNulls; +use vortex_array::search_sorted::SortedOrder; +use vortex_array::search_sorted::sorted_membership_mask; +use vortex_buffer::BitBufferMut; +use vortex_utils::aliases::hash_set::HashSet; fn main() { divan::main(); @@ -51,6 +65,203 @@ fn binary_search_vortex(bencher: Bencher) { }); } +#[divan::bench] +fn sorted_varbin_membership(bencher: Bencher) { + let values = (0_u128..65_536) + .map(|value| value.to_be_bytes()) + .collect::>(); + let members = values.iter().step_by(16).cloned().collect::>(); + let values = VarBinViewArray::from_iter_bin(values).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let members = SortedArray::try_new( + VarBinViewArray::from_iter_bin(members).into_array(), + ascending(), + &mut ctx, + ) + .unwrap(); + + bencher.bench_local(|| { + sorted_membership_mask( + divan::black_box(&values), + divan::black_box(&members), + NullEquality::Unequal, + &mut ctx, + ) + .unwrap() + .true_count() + }); +} + +#[divan::bench] +fn sorted_i64_membership(bencher: Bencher) { + let values = PrimitiveArray::from_iter(0_i64..65_536).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let members = SortedArray::try_new( + PrimitiveArray::from_iter((0_i64..65_536).step_by(16)).into_array(), + ascending(), + &mut ctx, + ) + .unwrap(); + + bencher.bench_local(|| { + sorted_membership_mask( + divan::black_box(&values), + divan::black_box(&members), + NullEquality::Unequal, + &mut ctx, + ) + .unwrap() + .true_count() + }); +} + +fn ascending() -> SortedOrder { + SortedOrder { + direction: SortedDirection::Ascending, + nulls: SortedNulls::First, + } +} + +mod membership_comparison { + use super::*; + + /// One engine-sized probe chunk against increasingly large/sparse member + /// domains. Construction and probe costs are separated because engines + /// may already own either a sorted member array or a hash index. + const PROBE_ROWS: usize = 8_192; + const CASES: &[(usize, i64)] = &[(16_384, 1), (65_536, 4), (1_000_000, 16)]; + + struct Fixture { + values: Vec, + values_array: vortex_array::ArrayRef, + members: Vec, + sorted_members: SortedArray, + hashed_members: HashSet, + } + + impl Fixture { + fn new(member_count: usize, stride: i64) -> Self { + let members = (0..member_count) + .map(|index| index as i64 * stride) + .collect::>(); + let domain = member_count as i64 * stride; + let start = domain / 2 - PROBE_ROWS as i64 / 2; + let values = (start..start + PROBE_ROWS as i64).collect::>(); + let values_array = PrimitiveArray::from_iter(values.iter().copied()).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let sorted_members = SortedArray::try_new( + PrimitiveArray::from_iter(members.iter().copied()).into_array(), + ascending(), + &mut ctx, + ) + .unwrap(); + let hashed_members = members.iter().copied().collect(); + Self { + values, + values_array, + members, + sorted_members, + hashed_members, + } + } + } + + #[divan::bench(args = CASES)] + fn narrowed_merge(bencher: Bencher, &(members, stride): &(usize, i64)) { + let fixture = Fixture::new(members, stride); + let mut ctx = array_session().create_execution_ctx(); + bencher + .counter(ItemsCount::new(PROBE_ROWS)) + .bench_local(|| { + sorted_membership_mask( + divan::black_box(&fixture.values_array), + divan::black_box(&fixture.sorted_members), + NullEquality::Unequal, + &mut ctx, + ) + .unwrap() + .true_count() + }); + } + + #[divan::bench(args = CASES)] + fn full_merge(bencher: Bencher, &(members, stride): &(usize, i64)) { + let fixture = Fixture::new(members, stride); + bencher + .counter(ItemsCount::new(PROBE_ROWS)) + .bench_local(|| full_merge_mask(&fixture.values, &fixture.members)); + } + + #[divan::bench(args = CASES)] + fn per_row_binary_search(bencher: Bencher, &(members, stride): &(usize, i64)) { + let fixture = Fixture::new(members, stride); + bencher + .counter(ItemsCount::new(PROBE_ROWS)) + .bench_local(|| binary_search_mask(&fixture.values, &fixture.members)); + } + + #[divan::bench(args = CASES)] + fn hash_probe(bencher: Bencher, &(members, stride): &(usize, i64)) { + let fixture = Fixture::new(members, stride); + bencher + .counter(ItemsCount::new(PROBE_ROWS)) + .bench_local(|| hash_mask(&fixture.values, &fixture.hashed_members)); + } + + #[divan::bench(args = CASES)] + fn sorted_wrapper_build(bencher: Bencher, &(members, stride): &(usize, i64)) { + let fixture = Fixture::new(members, stride); + let array = PrimitiveArray::from_iter(fixture.members.iter().copied()).into_array(); + let mut ctx = array_session().create_execution_ctx(); + bencher.counter(ItemsCount::new(members)).bench_local(|| { + SortedArray::try_new(divan::black_box(array.clone()), ascending(), &mut ctx) + .unwrap() + .len() + }); + } + + #[divan::bench(args = CASES)] + fn hash_set_build(bencher: Bencher, &(members, stride): &(usize, i64)) { + let fixture = Fixture::new(members, stride); + bencher.counter(ItemsCount::new(members)).bench_local(|| { + fixture + .members + .iter() + .copied() + .collect::>() + .len() + }); + } + + fn full_merge_mask(values: &[i64], members: &[i64]) -> usize { + let mut bits = BitBufferMut::with_capacity(values.len()); + let mut member = 0; + for value in values { + while member < members.len() && members[member] < *value { + member += 1; + } + bits.append(member < members.len() && members[member] == *value); + } + bits.freeze().iter().filter(|selected| *selected).count() + } + + fn binary_search_mask(values: &[i64], members: &[i64]) -> usize { + let mut bits = BitBufferMut::with_capacity(values.len()); + for value in values { + bits.append(members.binary_search(value).is_ok()); + } + bits.freeze().iter().filter(|selected| *selected).count() + } + + fn hash_mask(values: &[i64], members: &HashSet) -> usize { + let mut bits = BitBufferMut::with_capacity(values.len()); + for value in values { + bits.append(members.contains(value)); + } + bits.freeze().iter().filter(|selected| *selected).count() + } +} + fn fixture() -> (Vec, Vec) { let mut rng = StdRng::seed_from_u64(0); let range = Uniform::new(0, 65_536).unwrap(); 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..66772b24210 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -29,10 +29,14 @@ use crate::arrays::ScalarFnArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::listview::ListViewArraySlotsExt; use crate::arrays::primitive::PrimitiveArrayExt; +use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProviderExt; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; use crate::scalar::ListScalar; @@ -46,6 +50,12 @@ use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::Operator; +use crate::search_sorted::NullEquality; +use crate::search_sorted::SortedArray; +use crate::search_sorted::SortedDirection; +use crate::search_sorted::SortedNulls; +use crate::search_sorted::SortedOrder; +use crate::search_sorted::sorted_membership_mask; use crate::validity::Validity; #[derive(Clone)] @@ -189,7 +199,7 @@ fn compute_list_contains( 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() { - constant_list_scalar_contains(&list_scalar.as_list(), value, nullability) + constant_list_scalar_contains(&list_scalar.as_list(), value, nullability, ctx) } else { todo!("unsupported list contains with list and element as arrays") } @@ -200,9 +210,14 @@ fn constant_list_scalar_contains( list_scalar: &ListScalar, values: &ArrayRef, nullability: Nullability, + ctx: &mut ExecutionCtx, ) -> VortexResult { let elements = list_scalar.elements().vortex_expect("non null"); + if let Some(result) = try_sorted_membership_contains(&elements, values, nullability, ctx)? { + return Ok(result); + } + let len = values.len(); let false_scalar = Scalar::bool(false, nullability); @@ -224,6 +239,105 @@ fn constant_list_scalar_contains( Ok(result.unwrap_or_else(|| ConstantArray::new(false_scalar, len).into_array())) } +/// Fast path for `values IN (elements)` when `values` is already known sorted (ascending, +/// nulls-first) via `Stat::IsSorted`/`Stat::IsStrictSorted`. Replaces the `O(elements * +/// values.len())` equality fan-out below with a single sorted merge. +/// +/// The literal `elements` are re-sorted and de-duplicated on every call rather than cached across +/// chunks: `ScalarFnVTable::execute` runs once per chunk with no cross-chunk cache today, but IN +/// lists are normally small, so `O(elements * log(elements))` per chunk is negligible next to the +/// `O(elements * values.len())` fan-out it replaces. +/// +/// Returns `Ok(None)` when the fast path does not apply (unsorted or unsorted-unknown `values`, or +/// an unsupported/floating-point element dtype — float exclusion avoids a mismatch between +/// `Scalar`'s `PartialOrd`, which cannot order NaN, and the total order `SortedArray` validates +/// against), in which case the caller falls back to the equality fan-out unchanged. +/// Below this many literal elements, the fixed cost of sorting the members and constructing a +/// `SortedArray` outweighs the equality fan-out's simplicity. Measured directly (see +/// `vortex-array/benches/list_contains.rs`) on an 8,192-row `i64` column on local (non-CodSpeed, +/// noisier) hardware: 4 elements clearly favor the fan-out, and 8 elements are a toss-up that +/// flipped direction between runs. 16 elements and up consistently favored the sorted merge, with +/// the gap widening sharply from there (256 elements: ~80 us vs. ~1.2 ms). This threshold sits +/// past that noisy zone with margin; re-tuning on CodSpeed's stable runners could likely lower it. +const MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP: usize = 12; + +fn try_sorted_membership_contains( + elements: &[Scalar], + values: &ArrayRef, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if elements.len() < MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP { + return Ok(None); + } + + let elem_dtype = values.dtype(); + let supported_dtype = match elem_dtype { + DType::Bool(_) | DType::Decimal(..) | DType::Utf8(_) | DType::Binary(_) => true, + DType::Primitive(ptype, _) => !ptype.is_float(), + _ => false, + }; + if !supported_dtype { + return Ok(None); + } + + let is_sorted = matches!( + values.statistics().get_as::(Stat::IsSorted), + Precision::Exact(true) + ) || matches!( + values.statistics().get_as::(Stat::IsStrictSorted), + Precision::Exact(true) + ); + if !is_sorted { + return Ok(None); + } + + // `elements` carry the list's own (possibly differently-nullable) element dtype, so cast each + // one to `values`'s non-nullable dtype before building -- every builder in this codebase + // requires an exact dtype match, and we've already dropped every null. + let member_dtype = elem_dtype.as_nonnullable(); + let mut sorted_elements = elements + .iter() + .filter(|element| !element.is_null()) + .map(|element| element.cast(&member_dtype)) + .collect::>>()?; + sorted_elements.sort_by(|a, b| { + a.partial_cmp(b) + .vortex_expect("list elements share a comparable, non-float dtype") + }); + sorted_elements.dedup(); + + let mut builder = builder_with_capacity(&member_dtype, sorted_elements.len()); + for element in &sorted_elements { + builder.append_scalar(element)?; + } + let members_array = builder.finish(); + + let members = SortedArray::try_new( + members_array, + SortedOrder { + direction: SortedDirection::Ascending, + nulls: SortedNulls::First, + }, + ctx, + )?; + + let mask = sorted_membership_mask(values, &members, NullEquality::Unequal, ctx)?; + // `NullEquality::Unequal` makes a null `values` row never match (mask bit `false`). The + // pre-existing equality fan-out below reaches the same physical result for a null needle row + // (every per-element `Eq` is null there, then immediately `.fill_null(false)`-ed before the + // `Or`-reduce) -- once fully executed/canonicalized, a null needle row is `false`, never + // null. So the output here is valid everywhere; only the declared dtype nullability (carried + // through `nullability`, e.g. because the list itself is nullable) needs to match. + let validity = match nullability { + Nullability::NonNullable => Validity::NonNullable, + Nullability::Nullable => Validity::AllValid, + }; + Ok(Some( + BoolArray::new(mask.into_bit_buffer(), validity).into_array(), + )) +} + /// Returns a [`BoolArray`] where each bit represents if a list contains the scalar. fn list_contains_scalar( array: &ArrayRef, @@ -941,4 +1055,301 @@ mod tests { let expected_zero = BoolArray::from_iter([true, false, false, false]); assert_arrays_eq!(result_zero, expected_zero, &mut ctx); } + + fn int_list_scalar(elements: Vec) -> Scalar { + Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + elements.into_iter().map(Scalar::from).collect(), + Nullability::NonNullable, + ) + } + + #[test] + fn test_sorted_membership_fast_path_sorts_and_dedups_unordered_literal_list() { + let mut ctx = array_session().create_execution_ctx(); + + let arr = PrimitiveArray::from_iter([1, 3, 3, 5, 7, 9, 9, 9, 12]).into_array(); + arr.statistics().compute_is_sorted(&mut ctx); + + // Deliberately unsorted, with a duplicate, and long enough (>= + // `MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP`) to actually take the fast path: it must sort/dedup + // internally. + let expr = list_contains( + lit(int_list_scalar(vec![ + 9, 1, 9, 4, 12, 20, 21, 22, 23, 24, 25, 26, + ])), + root(), + ); + let contains = arr.apply(&expr).unwrap(); + + let expected = + BoolArray::from_iter([true, false, false, false, false, true, true, true, true]); + assert_arrays_eq!(contains, expected, &mut ctx); + } + + #[test] + fn test_sorted_membership_fast_path_uses_is_strict_sorted_stat() { + let mut ctx = array_session().create_execution_ctx(); + + let arr = PrimitiveArray::from_iter([1, 2, 3, 4, 5]).into_array(); + arr.statistics().compute_is_strict_sorted(&mut ctx); + + let expr = list_contains( + lit(int_list_scalar(vec![ + 5, 5, 2, 100, 200, 300, 400, 500, 600, 700, 800, 900, + ])), + root(), + ); + let contains = arr.apply(&expr).unwrap(); + + let expected = BoolArray::from_iter([false, true, false, false, true]); + assert_arrays_eq!(contains, expected, &mut ctx); + } + + #[test] + fn test_sorted_membership_fast_path_null_needle_rows_are_false_not_null() { + let mut ctx = array_session().create_execution_ctx(); + + // Nulls-first, then non-decreasing: a valid `Stat::IsSorted` fixture. + let arr = PrimitiveArray::from_option_iter::([ + None, + None, + Some(1), + Some(3), + Some(5), + Some(5), + Some(8), + ]) + .into_array(); + arr.statistics().compute_is_sorted(&mut ctx); + + let expr = list_contains( + lit(int_list_scalar(vec![ + 5, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, + ])), + root(), + ); + // Forcing execution here (rather than comparing the lazy `ScalarFnArray` directly) + // sidesteps a pre-existing, unrelated quirk where per-row scalar access on this + // particular lazy expression tree (list-constant, array-needle) disagrees with its own + // canonicalized result for a null needle row -- reproducible on the untouched equality + // fan-out too, so it's a framework-level gap uncovered by adding null coverage here, not + // something this change introduces or fixes. + let contains = arr + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + + // A null needle row is `false`, never null: matches the pre-existing equality fan-out's + // executed contract (every per-element `Eq` on a null needle is null, then immediately + // `.fill_null(false)`-ed before the `Or`-reduce). The result dtype is nullable (the + // needle column is), even though no row is ever actually null. + let expected = BoolArray::new( + BitBuffer::from_iter([false, false, false, false, true, true, false]), + Validity::AllValid, + ); + assert_arrays_eq!(contains, expected, &mut ctx); + } + + #[test] + fn test_sorted_membership_fast_path_empty_literal_list_stays_on_fanout() { + // Below `MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP`, so this always takes the fan-out, + // regardless of the column's sortedness -- it's here to document that the threshold + // doesn't change the (correct, pre-existing) answer for a degenerate empty list. + let mut ctx = array_session().create_execution_ctx(); + let arr = PrimitiveArray::from_iter([1, 2, 3]).into_array(); + arr.statistics().compute_is_sorted(&mut ctx); + + let empty_list = Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::NonNullable)), + vec![], + Nullability::NonNullable, + ); + let contains = arr.apply(&list_contains(lit(empty_list), root())).unwrap(); + assert_arrays_eq!( + contains, + BoolArray::from_iter([false, false, false]), + &mut ctx + ); + } + + #[test] + fn test_sorted_membership_fast_path_all_null_literal_list() { + // At `MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP` elements, so this reaches the fast path, but + // every element is null: `sorted_elements` ends up empty after filtering, exercising + // `SortedArray::try_new` and `sorted_membership_mask` on a zero-length member set. + let mut ctx = array_session().create_execution_ctx(); + let arr = PrimitiveArray::from_iter([1, 2, 3]).into_array(); + arr.statistics().compute_is_sorted(&mut ctx); + + let null_only_list = Scalar::list( + Arc::new(DType::Primitive(I32, Nullability::Nullable)), + vec![Scalar::null(DType::Primitive(I32, Nullability::Nullable)); 12], + Nullability::NonNullable, + ); + let contains = arr + .apply(&list_contains(lit(null_only_list), root())) + .unwrap(); + assert_arrays_eq!( + contains, + BoolArray::from_iter([false, false, false]), + &mut ctx + ); + } + + #[test] + fn test_sorted_membership_fast_path_utf8() { + let mut ctx = array_session().create_execution_ctx(); + + let arr = VarBinArray::from_iter( + ["ant", "bee", "cat", "dog", "eel"].map(Some), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + arr.statistics().compute_is_sorted(&mut ctx); + + let list_scalar = Scalar::list( + Arc::new(DType::Utf8(Nullability::NonNullable)), + vec![ + Scalar::from("dog"), + Scalar::from("ant"), + Scalar::from("dog"), + Scalar::from("bee"), + Scalar::from("xyz"), + Scalar::from("fff"), + Scalar::from("ggg"), + Scalar::from("hhh"), + Scalar::from("iii"), + Scalar::from("jjj"), + Scalar::from("kkk"), + Scalar::from("lll"), + ], + Nullability::NonNullable, + ); + let contains = arr.apply(&list_contains(lit(list_scalar), root())).unwrap(); + + let expected = BoolArray::from_iter([true, true, false, true, false]); + assert_arrays_eq!(contains, expected, &mut ctx); + } + + #[test] + fn test_sorted_membership_fast_path_skips_float_dtype() { + let mut ctx = array_session().create_execution_ctx(); + + // Sorted float column with a literal list past `MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP`: + // `Stat::IsSorted` is true, but the fast path deliberately excludes floats (see + // `try_sorted_membership_contains`), so this exercises the fan-out fallback specifically + // via the dtype check, not the length threshold. + let arr = PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0]).into_array(); + arr.statistics().compute_is_sorted(&mut ctx); + + let list_scalar = Scalar::list( + Arc::new(DType::Primitive( + crate::dtype::PType::F64, + Nullability::NonNullable, + )), + vec![ + Scalar::from(3.0f64), + Scalar::from(1.0f64), + Scalar::from(50.0f64), + Scalar::from(60.0f64), + Scalar::from(70.0f64), + Scalar::from(80.0f64), + Scalar::from(90.0f64), + Scalar::from(100.0f64), + Scalar::from(110.0f64), + Scalar::from(120.0f64), + Scalar::from(130.0f64), + Scalar::from(140.0f64), + ], + Nullability::NonNullable, + ); + let contains = arr.apply(&list_contains(lit(list_scalar), root())).unwrap(); + + let expected = BoolArray::from_iter([true, false, true, false]); + assert_arrays_eq!(contains, expected, &mut ctx); + } + + #[rstest] + // Below `MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP`: both sides take the fan-out. + #[case(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10], vec![10, 1, 1, 5])] + #[case(vec![-5, -3, -3, 0, 2, 2, 2, 9], vec![-3, 9, 9, 100])] + #[case(vec![1], vec![1])] + #[case(vec![1, 2, 3], vec![])] + #[case(vec![1, 2, 3], vec![100])] + // At or past the threshold: the sorted column genuinely exercises the fast path. + #[case(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10], vec![10, 1, 1, 5, 7, 7, 3, 2, 4, 6, 8, 9])] + #[case(vec![-5, -3, -3, 0, 2, 2, 2, 9], vec![-3, 9, 9, 100, 0, -5, -5, 2, 2, 50, 51, 52])] + fn test_sorted_membership_fast_path_matches_fanout( + #[case] data: Vec, + #[case] list_values: Vec, + ) { + let mut ctx = array_session().create_execution_ctx(); + let expr = list_contains(lit(int_list_scalar(list_values)), root()); + + // Same logical data, constructed independently so each copy owns its own stats: one + // exercises the sorted fast path, the other (stat never computed) exercises the + // pre-existing equality fan-out. They must agree. + let sorted_column = PrimitiveArray::from_iter(data.clone()).into_array(); + sorted_column.statistics().compute_is_sorted(&mut ctx); + let fast_path_result = sorted_column + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + + let plain_column = PrimitiveArray::from_iter(data).into_array(); + let fanout_result = plain_column + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + + assert_arrays_eq!(fast_path_result, fanout_result, &mut ctx); + } + + #[rstest] + // Below `MIN_ELEMENTS_FOR_SORTED_MEMBERSHIP`: both sides take the fan-out. + #[case(vec![None, None, Some(1), Some(3), Some(5), Some(5), Some(8)], vec![5, 100])] + #[case(vec![None, Some(-2), Some(-2), Some(0), Some(4)], vec![-2, 4, 4, 7])] + #[case(vec![None, None, None], vec![1, 2])] + // At or past the threshold: the sorted column genuinely exercises the fast path. + #[case( + vec![None, None, Some(1), Some(3), Some(5), Some(5), Some(8)], + vec![5, 100, 200, 300, 8, 301, 302, 303, 304, 305, 306, 307] + )] + #[case( + vec![None, Some(-2), Some(-2), Some(0), Some(4)], + vec![-2, 4, 4, 7, 99, 100, 101, 102, 103, 104, 105, 106] + )] + fn test_sorted_membership_fast_path_matches_fanout_with_nulls( + #[case] data: Vec>, + #[case] list_values: Vec, + ) { + let mut ctx = array_session().create_execution_ctx(); + let expr = list_contains(lit(int_list_scalar(list_values)), root()); + + let sorted_column = PrimitiveArray::from_option_iter::(data.clone()).into_array(); + sorted_column.statistics().compute_is_sorted(&mut ctx); + let fast_path_result = sorted_column + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + + let plain_column = PrimitiveArray::from_option_iter::(data).into_array(); + let fanout_result = plain_column + .apply(&expr) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + + assert_arrays_eq!(fast_path_result, fanout_result, &mut ctx); + } } diff --git a/vortex-array/src/search_sorted/membership.rs b/vortex-array/src/search_sorted/membership.rs new file mode 100644 index 00000000000..47683ff398c --- /dev/null +++ b/vortex-array/src/search_sorted/membership.rs @@ -0,0 +1,691 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cmp::Ordering; + +use vortex_buffer::BitBufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::Bool; +use crate::arrays::BoolArray; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::VarBinView; +use crate::arrays::VarBinViewArray; +use crate::arrays::bool::BoolArrayExt; +use crate::arrays::varbinview::BinaryView; +use crate::buffer::BufferHandle; +use crate::dtype::BigCast; +use crate::dtype::DType; +use crate::dtype::DecimalType; +use crate::dtype::NativePType; +use crate::dtype::i256; +use crate::match_each_native_ptype; + +/// Direction of a sorted array. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum SortedDirection { + Ascending, + Descending, +} + +/// Placement of nulls in a sorted array. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum SortedNulls { + First, + Last, +} + +/// Complete ordering contract for a [`SortedArray`]. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub struct SortedOrder { + pub direction: SortedDirection, + pub nulls: SortedNulls, +} + +/// Whether null is a member of a set that contains null. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum NullEquality { + Equal, + Unequal, +} + +/// A canonical array whose ordering has been validated once. +/// +/// The wrapper retains the canonical member representation and validity mask, +/// so repeated probe batches never canonicalize or validate the full member +/// set again. +#[derive(Clone, Debug)] +pub struct SortedArray { + array: ArrayRef, + validity: Mask, + order: SortedOrder, +} + +impl SortedArray { + pub fn try_new( + array: ArrayRef, + order: SortedOrder, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let array = canonicalize_supported(array, ctx)?; + let validity = array.validity()?.execute_mask(array.len(), ctx)?; + validate_array_order(&array, &validity, order)?; + Ok(Self { + array, + validity, + order, + }) + } + + pub fn array(&self) -> &ArrayRef { + &self.array + } + + pub fn order(&self) -> SortedOrder { + self.order + } + + pub fn len(&self) -> usize { + self.array.len() + } + + pub fn is_empty(&self) -> bool { + self.array.is_empty() + } +} + +/// Return a mask selecting every sorted `values` entry present in `members`. +/// +/// `values` and `members` must have the same logical Bool, Primitive, +/// Decimal, Binary, or UTF-8 dtype; outer nullability may differ. The probe +/// batch is canonicalized and its ordering is validated during the call. +/// Members are validated only by [`SortedArray::try_new`]. Two binary +/// searches narrow the member range to the probe batch's first and last +/// values before one linear merge. +pub fn sorted_membership_mask( + values: &ArrayRef, + members: &SortedArray, + null_equality: NullEquality, + ctx: &mut ExecutionCtx, +) -> VortexResult { + vortex_ensure!( + values.dtype().eq_ignore_nullability(members.array.dtype()), + "sorted membership requires matching logical dtypes, got {} and {}", + values.dtype(), + members.array.dtype() + ); + let values = canonicalize_supported(values.clone(), ctx)?; + let validity = values.validity()?.execute_mask(values.len(), ctx)?; + membership_dispatch( + &values, + &validity, + &members.array, + &members.validity, + members.order, + null_equality, + ) +} + +fn canonicalize_supported(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(match array.dtype() { + DType::Bool(_) => array.execute::(ctx)?.into_array(), + DType::Primitive(..) => array.execute::(ctx)?.into_array(), + DType::Decimal(..) => array.execute::(ctx)?.into_array(), + DType::Utf8(_) | DType::Binary(_) => array.execute::(ctx)?.into_array(), + dtype => { + vortex_ensure!( + false, + "sorted membership does not support dtype {dtype}; expected Bool, Primitive, Decimal, Binary, or UTF-8" + ); + unreachable!() + } + }) +} + +fn validate_array_order(array: &ArrayRef, validity: &Mask, order: SortedOrder) -> VortexResult<()> { + match array.dtype() { + DType::Bool(_) => { + let values = array + .as_typed::() + .vortex_expect("canonical boolean array"); + let bits = values.to_bit_buffer(); + validate_order(array.len(), validity, order, |left, right| { + bits.value(left).cmp(&bits.value(right)) + }) + } + DType::Primitive(ptype, _) => match_each_native_ptype!(*ptype, |T| { + let array = array + .as_typed::() + .vortex_expect("canonical primitive array"); + let values = array.as_slice::(); + validate_order(array.len(), validity, order, |left, right| { + values[left].total_compare(values[right]) + }) + }), + DType::Decimal(..) => { + let array = array + .as_typed::() + .vortex_expect("canonical decimal array"); + let values = DecimalValues::new(&array); + validate_order(array.len(), validity, order, |left, right| { + values.value(left).cmp(&values.value(right)) + }) + } + DType::Utf8(_) | DType::Binary(_) => { + let array = array + .as_typed::() + .vortex_expect("canonical variable-width array"); + let values = VarBinValues::new(&array); + validate_order(array.len(), validity, order, |left, right| { + values.value(left).cmp(values.value(right)) + }) + } + _ => unreachable!("canonicalize_supported rejects other dtypes"), + } +} + +fn membership_dispatch( + values: &ArrayRef, + value_validity: &Mask, + members: &ArrayRef, + member_validity: &Mask, + order: SortedOrder, + null_equality: NullEquality, +) -> VortexResult { + match values.dtype() { + DType::Bool(_) => { + let values = values + .as_typed::() + .vortex_expect("canonical boolean values") + .to_bit_buffer(); + let members = members + .as_typed::() + .vortex_expect("canonical boolean members") + .to_bit_buffer(); + membership_core( + values.len(), + value_validity, + members.len(), + member_validity, + order, + null_equality, + |left, right| values.value(left).cmp(&values.value(right)), + |member, value| members.value(member).cmp(&values.value(value)), + |member, value| members.value(member) == values.value(value), + ) + } + DType::Primitive(ptype, _) => match_each_native_ptype!(*ptype, |T| { + let value_array = values + .as_typed::() + .vortex_expect("canonical primitive values"); + let member_array = members + .as_typed::() + .vortex_expect("canonical primitive members"); + let values = value_array.as_slice::(); + let members = member_array.as_slice::(); + membership_core( + values.len(), + value_validity, + members.len(), + member_validity, + order, + null_equality, + |left, right| values[left].total_compare(values[right]), + |member, value| members[member].total_compare(values[value]), + |member, value| members[member].is_eq(values[value]), + ) + }), + DType::Decimal(..) => { + let value_array = values + .as_typed::() + .vortex_expect("canonical decimal values"); + let member_array = members + .as_typed::() + .vortex_expect("canonical decimal members"); + let values = DecimalValues::new(&value_array); + let members = DecimalValues::new(&member_array); + membership_core( + values.len(), + value_validity, + members.len(), + member_validity, + order, + null_equality, + |left, right| values.value(left).cmp(&values.value(right)), + |member, value| members.value(member).cmp(&values.value(value)), + |member, value| members.value(member) == values.value(value), + ) + } + DType::Utf8(_) | DType::Binary(_) => { + let value_array = values + .as_typed::() + .vortex_expect("canonical variable-width values"); + let member_array = members + .as_typed::() + .vortex_expect("canonical variable-width members"); + let values = VarBinValues::new(&value_array); + let members = VarBinValues::new(&member_array); + membership_core( + values.len(), + value_validity, + members.len(), + member_validity, + order, + null_equality, + |left, right| values.value(left).cmp(values.value(right)), + |member, value| members.value(member).cmp(values.value(value)), + |member, value| members.value(member) == values.value(value), + ) + } + _ => unreachable!("canonicalize_supported rejects other dtypes"), + } +} + +#[allow(clippy::too_many_arguments)] +fn membership_core( + values_len: usize, + value_validity: &Mask, + members_len: usize, + member_validity: &Mask, + order: SortedOrder, + null_equality: NullEquality, + value_cmp: impl Fn(usize, usize) -> Ordering, + member_value_cmp: impl Fn(usize, usize) -> Ordering, + member_value_eq: impl Fn(usize, usize) -> bool, +) -> VortexResult { + validate_order(values_len, value_validity, order, &value_cmp)?; + if values_len == 0 || members_len == 0 { + return Ok(Mask::new_false(values_len)); + } + + let compare = |member: usize, value: usize| { + entry_cmp( + member_validity.value(member), + value_validity.value(value), + order, + || member_value_cmp(member, value), + ) + }; + let lower = partition_point(members_len, |member| compare(member, 0).is_lt()); + let upper = lower + + partition_point(members_len - lower, |offset| { + !compare(lower + offset, values_len - 1).is_gt() + }); + + let mut selected = BitBufferMut::with_capacity(values_len); + let mut member = lower; + for value in 0..values_len { + while member < upper && compare(member, value).is_lt() { + member += 1; + } + let keep = if member == upper || !compare(member, value).is_eq() { + false + } else if value_validity.value(value) { + member_value_eq(member, value) + } else { + null_equality == NullEquality::Equal + }; + selected.append(keep); + } + Ok(Mask::from_buffer(selected.freeze())) +} + +fn validate_order( + len: usize, + validity: &Mask, + order: SortedOrder, + value_cmp: impl Fn(usize, usize) -> Ordering, +) -> VortexResult<()> { + for index in 1..len { + vortex_ensure!( + !entry_cmp( + validity.value(index - 1), + validity.value(index), + order, + || value_cmp(index - 1, index), + ) + .is_gt(), + "sorted membership input violates its {:?} order at index {index}", + order + ); + } + Ok(()) +} + +fn entry_cmp( + left_valid: bool, + right_valid: bool, + order: SortedOrder, + value_cmp: impl FnOnce() -> Ordering, +) -> Ordering { + match (left_valid, right_valid) { + (false, false) => Ordering::Equal, + (false, true) => match order.nulls { + SortedNulls::First => Ordering::Less, + SortedNulls::Last => Ordering::Greater, + }, + (true, false) => match order.nulls { + SortedNulls::First => Ordering::Greater, + SortedNulls::Last => Ordering::Less, + }, + (true, true) => match order.direction { + SortedDirection::Ascending => value_cmp(), + SortedDirection::Descending => value_cmp().reverse(), + }, + } +} + +fn partition_point(len: usize, mut predicate: impl FnMut(usize) -> bool) -> usize { + let mut left = 0; + let mut right = len; + while left < right { + let middle = left + (right - left) / 2; + if predicate(middle) { + left = middle + 1; + } else { + right = middle; + } + } + left +} + +struct VarBinValues<'a> { + views: &'a [BinaryView], + buffers: &'a [BufferHandle], +} + +impl<'a> VarBinValues<'a> { + fn new(array: &'a ArrayView<'a, VarBinView>) -> Self { + Self { + views: array.views(), + buffers: array.data_buffers(), + } + } + + fn len(&self) -> usize { + self.views.len() + } + + fn value(&self, index: usize) -> &[u8] { + let view = &self.views[index]; + if view.is_inlined() { + view.as_inlined().value() + } else { + let reference = view.as_view(); + &self.buffers[reference.buffer_index as usize].as_host()[reference.as_range()] + } + } +} + +enum DecimalValues { + I8(vortex_buffer::Buffer), + I16(vortex_buffer::Buffer), + I32(vortex_buffer::Buffer), + I64(vortex_buffer::Buffer), + I128(vortex_buffer::Buffer), + I256(vortex_buffer::Buffer), +} + +impl DecimalValues { + fn new(array: &ArrayView<'_, Decimal>) -> Self { + match array.values_type() { + DecimalType::I8 => Self::I8(array.buffer::()), + DecimalType::I16 => Self::I16(array.buffer::()), + DecimalType::I32 => Self::I32(array.buffer::()), + DecimalType::I64 => Self::I64(array.buffer::()), + DecimalType::I128 => Self::I128(array.buffer::()), + DecimalType::I256 => Self::I256(array.buffer::()), + } + } + + fn len(&self) -> usize { + match self { + Self::I8(values) => values.len(), + Self::I16(values) => values.len(), + Self::I32(values) => values.len(), + Self::I64(values) => values.len(), + Self::I128(values) => values.len(), + Self::I256(values) => values.len(), + } + } + + fn value(&self, index: usize) -> i256 { + match self { + Self::I8(values) => ::from(values[index]), + Self::I16(values) => ::from(values[index]), + Self::I32(values) => ::from(values[index]), + Self::I64(values) => ::from(values[index]), + Self::I128(values) => ::from(values[index]), + Self::I256(values) => Some(values[index]), + } + .vortex_expect("every decimal storage value widens to i256") + } +} + +#[cfg(test)] +mod tests { + use std::cell::Cell; + + use vortex_buffer::Buffer; + + use super::*; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::DecimalArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::VarBinViewArray; + use crate::dtype::DecimalDType; + use crate::validity::Validity; + + const ASC: SortedOrder = SortedOrder { + direction: SortedDirection::Ascending, + nulls: SortedNulls::First, + }; + + #[test] + fn primitive_membership_handles_duplicates_and_range_narrowing() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let members = PrimitiveArray::from_iter(0_i64..100_000).into_array(); + let members = SortedArray::try_new(members, ASC, &mut ctx)?; + let values = + PrimitiveArray::from_iter([49_999_i64, 50_000, 50_000, 75_000, 100_001]).into_array(); + let mask = sorted_membership_mask(&values, &members, NullEquality::Unequal, &mut ctx)?; + assert_eq!( + mask.iter().collect::>(), + [true, true, true, true, false] + ); + Ok(()) + } + + #[test] + fn boolean_and_empty_membership() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let members = SortedArray::try_new( + BoolArray::from_iter([false, true]).into_array(), + ASC, + &mut ctx, + )?; + let values = BoolArray::from_iter([false, false, true]).into_array(); + assert!( + sorted_membership_mask(&values, &members, NullEquality::Unequal, &mut ctx)?.all_true() + ); + + let empty = SortedArray::try_new( + PrimitiveArray::from_iter(Vec::::new()).into_array(), + ASC, + &mut ctx, + )?; + let values = PrimitiveArray::from_iter([1_i32, 2]).into_array(); + assert!( + sorted_membership_mask(&values, &empty, NullEquality::Unequal, &mut ctx)?.all_false() + ); + Ok(()) + } + + #[test] + fn float_membership_uses_vortex_bitwise_identity() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let members = SortedArray::try_new( + PrimitiveArray::from_iter([-0.0_f64, 1.0, f64::NAN]).into_array(), + ASC, + &mut ctx, + )?; + let values = PrimitiveArray::from_iter([-0.0_f64, 0.0, 1.0, f64::NAN]).into_array(); + assert_eq!( + sorted_membership_mask(&values, &members, NullEquality::Unequal, &mut ctx)? + .iter() + .collect::>(), + [true, false, true, true] + ); + Ok(()) + } + + #[test] + fn member_range_is_narrowed_before_linear_merge() -> VortexResult<()> { + let members = (0_i64..1_000_000).collect::>(); + let values = (500_000_i64..500_010).collect::>(); + let comparisons = Cell::new(0_usize); + let mask = membership_core( + values.len(), + &Mask::new_true(values.len()), + members.len(), + &Mask::new_true(members.len()), + ASC, + NullEquality::Unequal, + |left, right| values[left].cmp(&values[right]), + |member, value| { + comparisons.set(comparisons.get() + 1); + members[member].cmp(&values[value]) + }, + |member, value| members[member] == values[value], + )?; + assert!(mask.all_true()); + assert!( + comparisons.get() < 100, + "range narrowing performed {} comparisons", + comparisons.get() + ); + Ok(()) + } + + #[test] + fn descending_null_semantics_are_explicit() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let order = SortedOrder { + direction: SortedDirection::Descending, + nulls: SortedNulls::Last, + }; + let members = PrimitiveArray::from_option_iter([Some(9_i32), Some(3), None]).into_array(); + let members = SortedArray::try_new(members, order, &mut ctx)?; + let values = + PrimitiveArray::from_option_iter([Some(10_i32), Some(9), Some(3), None]).into_array(); + assert_eq!( + sorted_membership_mask(&values, &members, NullEquality::Equal, &mut ctx)? + .iter() + .collect::>(), + [false, true, true, true] + ); + assert_eq!( + sorted_membership_mask(&values, &members, NullEquality::Unequal, &mut ctx)? + .iter() + .collect::>(), + [false, true, true, false] + ); + Ok(()) + } + + #[test] + fn decimal_membership_accepts_mixed_physical_widths() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(20, 2); + let members = DecimalArray::new( + Buffer::from(vec![-100_i128, 0, 250]), + dtype, + Validity::NonNullable, + ) + .into_array(); + let members = SortedArray::try_new(members, ASC, &mut ctx)?; + let values = DecimalArray::new( + Buffer::from(vec![ + ::from(-100_i128).vortex_expect("test decimal fits"), + ::from(1_i128).vortex_expect("test decimal fits"), + ::from(250_i128).vortex_expect("test decimal fits"), + ]), + dtype, + Validity::NonNullable, + ) + .into_array(); + assert_eq!( + sorted_membership_mask(&values, &members, NullEquality::Unequal, &mut ctx)? + .iter() + .collect::>(), + [true, false, true] + ); + Ok(()) + } + + #[test] + fn varbin_membership_reads_inline_and_external_values() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let members = + VarBinViewArray::from_iter_bin([b"alpha".as_slice(), b"external-value-0001", b"omega"]) + .into_array(); + let members = SortedArray::try_new(members, ASC, &mut ctx)?; + let values = VarBinViewArray::from_iter_bin([ + b"alpha".as_slice(), + b"external-value-0000", + b"external-value-0001", + b"omega", + ]) + .into_array(); + assert_eq!( + sorted_membership_mask(&values, &members, NullEquality::Unequal, &mut ctx)? + .iter() + .collect::>(), + [true, false, true, true] + ); + Ok(()) + } + + #[test] + fn rejects_unsorted_members_and_values() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let unsorted = PrimitiveArray::from_iter([2_i32, 1]).into_array(); + assert!(SortedArray::try_new(unsorted, ASC, &mut ctx).is_err()); + + let members = SortedArray::try_new( + PrimitiveArray::from_iter([1_i32, 2]).into_array(), + ASC, + &mut ctx, + )?; + let values = PrimitiveArray::from_iter([2_i32, 1]).into_array(); + assert!( + sorted_membership_mask(&values, &members, NullEquality::Unequal, &mut ctx).is_err() + ); + Ok(()) + } + + #[test] + fn rejects_mismatched_and_unsupported_dtypes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let members = SortedArray::try_new( + PrimitiveArray::from_iter([1_i32]).into_array(), + ASC, + &mut ctx, + )?; + let wrong = PrimitiveArray::from_iter([1_i64]).into_array(); + assert!(sorted_membership_mask(&wrong, &members, NullEquality::Unequal, &mut ctx).is_err()); + Ok(()) + } +} diff --git a/vortex-array/src/search_sorted/mod.rs b/vortex-array/src/search_sorted/mod.rs index cda27cfc511..a100c7d4762 100644 --- a/vortex-array/src/search_sorted/mod.rs +++ b/vortex-array/src/search_sorted/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod membership; mod primitive; use std::cmp::Ordering; @@ -12,6 +13,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hint; +pub use membership::*; pub use primitive::*; use vortex_error::VortexResult;