From a8c4b0f793fd8945d5dffbe0bbcb61998f3e2535 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 16:03:00 +0100 Subject: [PATCH 01/11] Remove variants module Signed-off-by: Robert Kruszewski --- .../src/decimal_byte_parts/compute/compare.rs | 2 +- encodings/runend/src/array.rs | 18 +- encodings/runend/src/arrow.rs | 37 +-- encodings/runend/src/kernel.rs | 3 +- encodings/runend/src/ops.rs | 27 +- vortex-array/src/arrays/bool/compute/zip.rs | 2 +- .../src/arrays/listview/compute/zip.rs | 2 +- .../src/arrays/primitive/compute/zip.rs | 2 +- vortex-array/src/arrays/primitive/tests.rs | 19 +- .../src/arrays/struct_/compute/zip.rs | 6 +- .../src/arrays/varbinview/compute/zip.rs | 2 +- vortex-array/src/compute/conformance/cast.rs | 2 +- vortex-array/src/lib.rs | 1 - vortex-array/src/patches.rs | 69 ++--- vortex-array/src/scalar_fn/fns/merge.rs | 4 +- .../mod.rs} | 3 + vortex-array/src/search_sorted/primitive.rs | 92 +++++++ vortex-array/src/variants.rs | 241 ------------------ 18 files changed, 201 insertions(+), 331 deletions(-) rename vortex-array/src/{search_sorted.rs => search_sorted/mod.rs} (99%) create mode 100644 vortex-array/src/search_sorted/primitive.rs delete mode 100644 vortex-array/src/variants.rs diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index 04f321cf88f..602dbc26f77 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -49,7 +49,7 @@ impl CompareKernel for DecimalByteParts { match decimal_value_wrapper_to_primitive( rhs_decimal, - lhs.msp().as_primitive_typed().ptype(), + lhs.msp().dtype().as_ptype(), ) { Ok(value) => { let encoded_scalar = Scalar::try_new(scalar_type, Some(value))?; diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index 9dba648e248..d8f0648dd70 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -28,9 +28,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::legacy_session; -use vortex_array::scalar::PValue; -use vortex_array::search_sorted::SearchSorted; -use vortex_array::search_sorted::SearchSortedSide; use vortex_array::serde::ArrayChildren; use vortex_array::smallvec::smallvec; use vortex_array::validity::Validity; @@ -48,6 +45,7 @@ use crate::compress::runend_decode_primitive; use crate::compress::runend_decode_varbinview; use crate::compress::runend_encode; use crate::decompress_bool::runend_decode_bools; +use crate::ops::find_physical_index; use crate::rules::RULES; /// A [`RunEnd`]-encoded Vortex array. @@ -231,16 +229,14 @@ pub trait RunEndArrayExt: TypedArrayRef { } fn find_physical_index(&self, index: usize) -> VortexResult { - Ok(self - .ends() - .as_primitive_typed() - .search_sorted( - &PValue::from(index + self.offset()), - SearchSortedSide::Right, - )? - .to_ends_index(self.ends().len())) + find_physical_index( + self.ends(), + index + self.offset(), + &mut LEGACY_SESSION.create_execution_ctx(), + ) } } + impl> RunEndArrayExt for T {} #[derive(Clone, Debug)] diff --git a/encodings/runend/src/arrow.rs b/encodings/runend/src/arrow.rs index c927daaaffe..a24e69ef5c8 100644 --- a/encodings/runend/src/arrow.rs +++ b/encodings/runend/src/arrow.rs @@ -11,14 +11,13 @@ use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::NativePType; use vortex_array::legacy_session; -use vortex_array::scalar::PValue; use vortex_array::search_sorted::SearchSorted; -use vortex_array::search_sorted::SearchSortedSide; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; use crate::RunEndData; +use crate::ops::find_physical_index; use crate::ops::find_slice_end_index; impl FromArrowArray<&RunArray> for RunEndData @@ -41,14 +40,13 @@ where ends.validity()?, ) .into_array(); + let mut ctx = legacy_session().create_execution_ctx(); + let (ends_slice, values_slice) = if offset == 0 && len == array.run_ends().max_value() { (ends_array, values) } else { - let slice_begin = ends_array - .as_primitive_typed() - .search_sorted(&PValue::from(offset), SearchSortedSide::Right)? - .to_ends_index(ends_array.len()); - let slice_end = find_slice_end_index(&ends_array, offset + len)?; + let slice_begin = find_physical_index(&ends_array, offset, &mut ctx)?; + let slice_end = find_slice_end_index(&ends_array, offset + len, &mut ctx)?; ( ends_array.slice(slice_begin..slice_end)?, @@ -57,8 +55,6 @@ where }; // SAFETY: arrow-rs enforces the RunEndArray invariants, we inherit their guarantees. - // TODO(ctx): trait fixes - FromArrowArray::from_arrow has a fixed signature. - let mut ctx = legacy_session().create_execution_ctx(); RunEndData::validate_parts(&ends_slice, &values_slice, offset, len, &mut ctx)?; Ok(unsafe { RunEndData::new_unchecked(offset) }) } @@ -92,9 +88,6 @@ mod tests { use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::scalar::PValue; - use vortex_array::search_sorted::SearchSorted; - use vortex_array::search_sorted::SearchSortedSide; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -102,6 +95,8 @@ mod tests { use vortex_session::VortexSession; use crate::RunEnd; + use crate::RunEndArray; + use crate::ops::find_physical_index; use crate::ops::find_slice_end_index; static SESSION: LazyLock = LazyLock::new(|| { @@ -113,7 +108,7 @@ mod tests { fn decode_run_array( array: &RunArray, nullable: bool, - ) -> VortexResult + ) -> VortexResult where R::Native: NativePType, { @@ -131,14 +126,12 @@ mod tests { ends.validity()?, ) .into_array(); + let mut ctx = SESSION.create_execution_ctx(); let (ends_slice, values_slice) = if offset == 0 && len == array.run_ends().max_value() { (ends_array, values) } else { - let slice_begin = ends_array - .as_primitive_typed() - .search_sorted(&PValue::from(offset), SearchSortedSide::Right)? - .to_ends_index(ends_array.len()); - let slice_end = find_slice_end_index(&ends_array, offset + len)?; + let slice_begin = find_physical_index(&ends_array, offset, &mut ctx)?; + let slice_end = find_slice_end_index(&ends_array, offset + len, &mut ctx)?; ( ends_array.slice(slice_begin..slice_end)?, @@ -146,13 +139,7 @@ mod tests { ) }; - RunEnd::try_new_offset_length( - ends_slice, - values_slice, - offset, - array.len(), - &mut SESSION.create_execution_ctx(), - ) + RunEnd::try_new_offset_length(ends_slice, values_slice, offset, array.len(), &mut ctx) } #[test] diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 484a4233789..4f3f2a4fd78 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -25,6 +25,7 @@ use vortex_session::VortexSession; use crate::RunEnd; use crate::array::RunEndArrayExt; use crate::compute::take_from::RunEndTakeFrom; +use crate::ops::find_slice_end_index; pub(super) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -64,7 +65,7 @@ fn slice( let new_length = range.len(); let slice_begin = array.find_physical_index(range.start)?; - let slice_end = crate::ops::find_slice_end_index(array.ends(), range.end + array.offset())?; + let slice_end = find_slice_end_index(array.ends(), range.end + array.offset(), ctx)?; // If the sliced range contains only a single run, opt to return a ConstantArray. if slice_begin + 1 == slice_end { diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 56de024f988..2efb51b4cfc 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -4,11 +4,12 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; -use vortex_array::scalar::PValue; +use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar::Scalar; use vortex_array::search_sorted::SearchResult; use vortex_array::search_sorted::SearchSorted; use vortex_array::search_sorted::SearchSortedSide; +use vortex_array::search_sorted::TypedPrimitiveArray; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexResult; @@ -31,10 +32,14 @@ impl OperationsVTable for RunEnd { /// /// If the index exists in the array we want to take that position (as we are searching from the right) /// otherwise we want to take the next one -pub(crate) fn find_slice_end_index(array: &ArrayRef, index: usize) -> VortexResult { - let result = array - .as_primitive_typed() - .search_sorted(&PValue::from(index), SearchSortedSide::Right)?; +pub(crate) fn find_slice_end_index( + array: &ArrayRef, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let result = match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { + TypedPrimitiveArray::::new(array, ctx).search_sorted(&index, SearchSortedSide::Right)? + }); Ok(match result { SearchResult::Found(i) => i, SearchResult::NotFound(i) => { @@ -47,6 +52,18 @@ pub(crate) fn find_slice_end_index(array: &ArrayRef, index: usize) -> VortexResu }) } +pub(crate) fn find_physical_index( + array: &ArrayRef, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { + Ok(TypedPrimitiveArray::::new(array, ctx) + .search_sorted(&index, SearchSortedSide::Right)? + .to_ends_index(array.len())) + }) +} + #[cfg(test)] mod tests { use std::sync::LazyLock; diff --git a/vortex-array/src/arrays/bool/compute/zip.rs b/vortex-array/src/arrays/bool/compute/zip.rs index 80e7de80c82..ad311572395 100644 --- a/vortex-array/src/arrays/bool/compute/zip.rs +++ b/vortex-array/src/arrays/bool/compute/zip.rs @@ -34,7 +34,7 @@ impl ZipKernel for Bool { }; // Null mask entries select `if_false`, matching `Zip`'s SQL ELSE semantics. - let mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask = mask.clone().null_as_false().execute(ctx)?; let mask_values = match &mask { // Defer trivial masks to the generic zip, which just casts the surviving side. Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None), diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index 6007113300c..71640d4fb21 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -50,7 +50,7 @@ impl ZipKernel for ListView { }; // Null mask entries select `if_false`, matching `Zip`'s SQL ELSE semantics. - let mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask = mask.clone().null_as_false().execute(ctx)?; match &mask { // Defer the trivial masks to the generic zip, which just casts one side. Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None), diff --git a/vortex-array/src/arrays/primitive/compute/zip.rs b/vortex-array/src/arrays/primitive/compute/zip.rs index 83f03b10614..35e49831acc 100644 --- a/vortex-array/src/arrays/primitive/compute/zip.rs +++ b/vortex-array/src/arrays/primitive/compute/zip.rs @@ -46,7 +46,7 @@ impl ZipKernel for Primitive { } // Null mask entries select `if_false`, matching `Zip`'s SQL ELSE semantics. - let mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask = mask.clone().null_as_false().execute(ctx)?; match &mask { // Defer trivial masks to the generic zip, which just casts the surviving side. Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None), diff --git a/vortex-array/src/arrays/primitive/tests.rs b/vortex-array/src/arrays/primitive/tests.rs index f623479c7b0..9eda67ab287 100644 --- a/vortex-array/src/arrays/primitive/tests.rs +++ b/vortex-array/src/arrays/primitive/tests.rs @@ -1,10 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::LazyLock; + use vortex_buffer::buffer; +use vortex_error::VortexExpect; +use vortex_session::VortexSession; use crate::ArrayRef; use crate::IntoArray; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::filter::test_filter_conformance; @@ -12,24 +17,26 @@ use crate::compute::conformance::mask::test_mask_conformance; use crate::compute::conformance::search_sorted::rstest_reuse::apply; use crate::compute::conformance::search_sorted::search_sorted_conformance; use crate::compute::conformance::search_sorted::*; -use crate::scalar::PValue; +use crate::executor::VortexSessionExecute; use crate::search_sorted::SearchResult; use crate::search_sorted::SearchSorted; use crate::search_sorted::SearchSortedSide; +use crate::search_sorted::TypedPrimitiveArray; use crate::validity::Validity; +static SESSION: LazyLock = LazyLock::new(array_session); + #[apply(search_sorted_conformance)] fn test_search_sorted_primitive( #[case] array: ArrayRef, #[case] value: i32, #[case] side: SearchSortedSide, #[case] expected: SearchResult, -) -> vortex_error::VortexResult<()> { - let res = array - .as_primitive_typed() - .search_sorted(&Some(PValue::from(value)), side)?; +) { + let res = TypedPrimitiveArray::::new(&array, &mut SESSION.create_execution_ctx()) + .search_sorted(&value, side) + .vortex_expect("Failed to search sorted"); assert_eq!(res, expected); - Ok(()) } #[test] diff --git a/vortex-array/src/arrays/struct_/compute/zip.rs b/vortex-array/src/arrays/struct_/compute/zip.rs index cf7b7b91d6d..394cc5c8424 100644 --- a/vortex-array/src/arrays/struct_/compute/zip.rs +++ b/vortex-array/src/arrays/struct_/compute/zip.rs @@ -37,7 +37,7 @@ impl ZipKernel for Struct { let fields = if_true .iter_unmasked_fields() .zip(if_false.iter_unmasked_fields()) - .map(|(t, f)| ArrayBuiltins::zip(mask, t.clone(), f.clone())) + .map(|(t, f)| mask.zip(t.clone(), f.clone())) .collect::>>()?; let v1 = if_true.validity()?; @@ -48,11 +48,11 @@ impl ZipKernel for Struct { (Validity::AllInvalid, Validity::AllInvalid) => Validity::AllInvalid, (v1, v2) => { - let mask_mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask_mask = mask.clone().null_as_false().execute(ctx)?; let v1m = v1.execute_mask(if_true.len(), ctx)?; let v2m = v2.execute_mask(if_false.len(), ctx)?; - let combined = (v1m.bitand(&mask_mask)).bitor(&v2m.bitand(&mask_mask.not())); + let combined = v1m.bitand(&mask_mask).bitor(&v2m.bitand(&mask_mask.not())); Validity::from_mask( combined, if_true.dtype().nullability() | if_false.dtype().nullability(), diff --git a/vortex-array/src/arrays/varbinview/compute/zip.rs b/vortex-array/src/arrays/varbinview/compute/zip.rs index 75e30808821..193dc523823 100644 --- a/vortex-array/src/arrays/varbinview/compute/zip.rs +++ b/vortex-array/src/arrays/varbinview/compute/zip.rs @@ -57,7 +57,7 @@ impl ZipKernel for VarBinView { let true_validity = if_true.varbinview_validity().execute_mask(len, ctx)?; let false_validity = if_false.varbinview_validity().execute_mask(len, ctx)?; - let mask = mask.try_to_mask_fill_null_false(ctx)?; + let mask = mask.clone().null_as_false().execute(ctx)?; let if_false_view = if_false; match mask.slices() { AllOr::All => push_range( diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index c823b2a778a..922116f520a 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -216,7 +216,7 @@ fn test_cast_to_nullable(array: &ArrayRef) { } fn test_cast_from_floating_point_types(array: &ArrayRef) { - let ptype = array.as_primitive_typed().ptype(); + let ptype = array.dtype().as_ptype(); test_cast_to_primitive(array, PType::I8, false); test_cast_to_primitive(array, PType::U8, false); test_cast_to_primitive(array, PType::I16, false); diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index e219a929453..6456c412929 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -144,7 +144,6 @@ pub mod stream; #[cfg(any(test, feature = "_test-harness"))] pub mod test_harness; pub mod validity; -pub mod variants; pub mod flatbuffers { //! Re-exported autogenerated code from the core Vortex flatbuffer definitions. diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 32d220797a6..a59db47fd51 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -35,14 +35,12 @@ use crate::dtype::Nullability; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; use crate::dtype::UnsignedPType; -use crate::legacy_session; -use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; -use crate::scalar::PValue; use crate::scalar::Scalar; use crate::search_sorted::SearchResult; use crate::search_sorted::SearchSorted; use crate::search_sorted::SearchSortedSide; +use crate::search_sorted::TypedPrimitiveArray; use crate::validity::Validity; /// One patch index offset is stored for each chunk. @@ -476,32 +474,7 @@ impl Patches { return self.search_index_chunked(index); } - Self::search_index_binary_search(&self.indices, index + self.offset) - } - - /// Binary searches for `needle` in the indices array. - /// - /// # Returns - /// [`SearchResult::Found`] with the position if needle exists, or [`SearchResult::NotFound`] - /// with the insertion point if not found. - fn search_index_binary_search(indices: &ArrayRef, needle: usize) -> VortexResult { - if let Some(primitive) = indices.as_opt::() { - match_each_integer_ptype!(primitive.ptype(), |T| { - let Ok(needle) = T::try_from(needle) else { - // If the needle is not of type T, then it cannot possibly be in this array. - // - // The needle is a non-negative integer (a usize); therefore, it must be larger - // than all values in this array. - return Ok(SearchResult::NotFound(primitive.len())); - }; - return primitive - .as_slice::() - .search_sorted(&needle, SearchSortedSide::Left); - }); - } - indices - .as_primitive_typed() - .search_sorted(&PValue::U64(needle as u64), SearchSortedSide::Left) + search_index_binary_search(&self.indices, index + self.offset) } /// Constant time searches for `index` in the indices array. @@ -549,7 +522,7 @@ impl Patches { }; let chunk_indices = self.indices.slice(patches_start_idx..patches_end_idx)?; - let result = Self::search_index_binary_search(&chunk_indices, index + self.offset)?; + let result = search_index_binary_search(&chunk_indices, index + self.offset)?; Ok(match result { SearchResult::Found(idx) => SearchResult::Found(patches_start_idx + idx), @@ -1010,6 +983,42 @@ impl Patches { } } +/// Binary searches for `needle` in the indices array. +/// +/// # Returns +/// [`SearchResult::Found`] with the position if needle exists, or [`SearchResult::NotFound`] +/// with the insertion point if not found. +fn search_index_binary_search(indices: &ArrayRef, needle: usize) -> VortexResult { + if let Some(primitive) = indices.as_opt::() { + match_each_unsigned_integer_ptype!(primitive.ptype(), |T| { + let Ok(needle) = T::try_from(needle) else { + // If the needle is not of type T, then it cannot possibly be in this array. + // + // The needle is a non-negative integer (a usize); therefore, it must be larger + // than all values in this array. + return Ok(SearchResult::NotFound(primitive.len())); + }; + return primitive + .as_slice::() + .search_sorted(&needle, SearchSortedSide::Left); + }); + } + + search_index_binary_search_scalar(indices, needle) +} + +fn search_index_binary_search_scalar( + indices: &ArrayRef, + needle: usize, +) -> VortexResult { + match_each_unsigned_integer_ptype!(indices.dtype().as_ptype(), |T| { + TypedPrimitiveArray::::new(indices, &mut LEGACY_SESSION.create_execution_ctx()) + .search_sorted(&needle, SearchSortedSide::Left) + .map_err(|_| vortex_err!("indices must be a primitive array")) + }) +} + + #[expect(clippy::too_many_arguments)] // private function, can clean up one day fn take_map, T: NativePType>( indices: &[I], diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index c9f5783de27..23fedb5009b 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -364,7 +364,7 @@ mod tests { let actual_array = test_array.apply(&expr).unwrap(); assert_eq!( - actual_array.as_struct_typed().names(), + actual_array.dtype().as_struct_fields().names(), ["a", "b", "c", "d", "e"] ); @@ -450,7 +450,7 @@ mod tests { .into_array(); let actual_array = test_array.clone().apply(&expr).unwrap(); assert_eq!(actual_array.len(), test_array.len()); - assert_eq!(actual_array.as_struct_typed().nfields(), 0); + assert_eq!(actual_array.nchildren(), 0); } #[test] diff --git a/vortex-array/src/search_sorted.rs b/vortex-array/src/search_sorted/mod.rs similarity index 99% rename from vortex-array/src/search_sorted.rs rename to vortex-array/src/search_sorted/mod.rs index f2ceb5dc24d..cda27cfc511 100644 --- a/vortex-array/src/search_sorted.rs +++ b/vortex-array/src/search_sorted/mod.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod primitive; + use std::cmp::Ordering; use std::cmp::Ordering::Equal; use std::cmp::Ordering::Greater; @@ -10,6 +12,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hint; +pub use primitive::*; use vortex_error::VortexResult; use crate::ArrayRef; diff --git a/vortex-array/src/search_sorted/primitive.rs b/vortex-array/src/search_sorted/primitive.rs new file mode 100644 index 00000000000..94dced9c294 --- /dev/null +++ b/vortex-array/src/search_sorted/primitive.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cell::RefCell; +use std::cmp::Ordering; +use std::marker::PhantomData; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::NativePType; +use crate::search_sorted::IndexOrd; + +pub struct TypedPrimitiveArray<'a, T>(&'a ArrayRef, RefCell<&'a mut ExecutionCtx>, PhantomData); + +impl<'a, T: NativePType> TypedPrimitiveArray<'a, T> { + pub fn new(array: &'a ArrayRef, ctx: &'a mut ExecutionCtx) -> Self { + assert_eq!( + array.dtype().as_ptype(), + T::PTYPE, + "Array PType must match primitive type" + ); + Self(array, RefCell::new(ctx), PhantomData) + } +} + +impl IndexOrd for TypedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &T) -> VortexResult> { + let value = self + .0 + .execute_scalar(idx, &mut self.1.borrow_mut())? + .as_primitive() + .typed_value::() + .unwrap_or_else(|| T::zero()); + + Ok(Some(value.total_compare(*elem))) + } + + fn index_len(&self) -> usize { + self.0.len() + } +} + +impl IndexOrd> for TypedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { + let ctx_mut = &mut self.1.borrow_mut(); + let value = self + .0 + .is_valid(idx, ctx_mut)? + .then(|| { + self.0.execute_scalar(idx, ctx_mut).map(|s| { + s.as_primitive() + .typed_value::() + .unwrap_or_else(|| T::zero()) + }) + }) + .transpose()?; + + Ok(match (value, elem.as_ref()) { + (Some(l), Some(r)) => Some(l.total_compare(*r)), + (Some(_), None) => Some(Ordering::Greater), + (None, Some(_)) => Some(Ordering::Less), + (None, None) => Some(Ordering::Equal), + }) + } + + fn index_len(&self) -> usize { + self.0.len() + } +} + +impl IndexOrd for TypedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &usize) -> VortexResult> { + let value = self + .0 + .execute_scalar(idx, &mut self.1.borrow_mut())? + .as_primitive() + .typed_value::() + .unwrap_or_else(|| T::zero()); + + let Some(elem_t) = T::from_usize(*elem) else { + return Ok(Some(Ordering::Less)); + }; + + Ok(Some(value.total_compare(elem_t))) + } + + fn index_len(&self) -> usize { + self.0.len() + } +} diff --git a/vortex-array/src/variants.rs b/vortex-array/src/variants.rs deleted file mode 100644 index 13c7a242c7a..00000000000 --- a/vortex-array/src/variants.rs +++ /dev/null @@ -1,241 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! This module defines extension functionality specific to each Vortex DType. -use std::cmp::Ordering; - -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_panic; -use vortex_mask::Mask; - -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::VortexSessionExecute; -use crate::aggregate_fn::fns::sum::sum; -use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::FieldNames; -use crate::dtype::PType; -use crate::dtype::extension::ExtDTypeRef; -use crate::legacy_session; -use crate::scalar::PValue; -use crate::scalar::Scalar; -use crate::search_sorted::IndexOrd; - -impl ArrayRef { - /// Downcasts the array for null-specific behavior. - pub fn as_null_typed(&self) -> NullTyped<'_> { - matches!(self.dtype(), DType::Null) - .then(|| NullTyped(self)) - .vortex_expect("Array does not have DType::Null") - } - - /// Downcasts the array for bool-specific behavior. - pub fn as_bool_typed(&self) -> BoolTyped<'_> { - matches!(self.dtype(), DType::Bool(..)) - .then(|| BoolTyped(self)) - .vortex_expect("Array does not have DType::Bool") - } - - /// Downcasts the array for primitive-specific behavior. - pub fn as_primitive_typed(&self) -> PrimitiveTyped<'_> { - matches!(self.dtype(), DType::Primitive(..)) - .then(|| PrimitiveTyped(self)) - .vortex_expect("Array does not have DType::Primitive") - } - - /// Downcasts the array for decimal-specific behavior. - pub fn as_decimal_typed(&self) -> DecimalTyped<'_> { - matches!(self.dtype(), DType::Decimal(..)) - .then(|| DecimalTyped(self)) - .vortex_expect("Array does not have DType::Decimal") - } - - /// Downcasts the array for utf8-specific behavior. - pub fn as_utf8_typed(&self) -> Utf8Typed<'_> { - matches!(self.dtype(), DType::Utf8(..)) - .then(|| Utf8Typed(self)) - .vortex_expect("Array does not have DType::Utf8") - } - - /// Downcasts the array for binary-specific behavior. - pub fn as_binary_typed(&self) -> BinaryTyped<'_> { - matches!(self.dtype(), DType::Binary(..)) - .then(|| BinaryTyped(self)) - .vortex_expect("Array does not have DType::Binary") - } - - /// Downcasts the array for struct-specific behavior. - pub fn as_struct_typed(&self) -> StructTyped<'_> { - matches!(self.dtype(), DType::Struct(..)) - .then(|| StructTyped(self)) - .vortex_expect("Array does not have DType::Struct") - } - - /// Downcasts the array for list-specific behavior. - pub fn as_list_typed(&self) -> ListTyped<'_> { - matches!(self.dtype(), DType::List(..)) - .then(|| ListTyped(self)) - .vortex_expect("Array does not have DType::List") - } - - /// Downcasts the array for extension-specific behavior. - pub fn as_extension_typed(&self) -> ExtensionTyped<'_> { - matches!(self.dtype(), DType::Extension(..)) - .then(|| ExtensionTyped(self)) - .vortex_expect("Array does not have DType::Extension") - } - - pub fn try_to_mask_fill_null_false(&self, ctx: &mut ExecutionCtx) -> VortexResult { - if !matches!(self.dtype(), DType::Bool(_)) { - vortex_bail!("mask must be bool array, has dtype {}", self.dtype()); - } - - // Convert nulls to false first in case this can be done cheaply by the encoding. - let array = self - .clone() - .fill_null(Scalar::bool(false, self.dtype().nullability()))?; - - Ok(array - .execute::(ctx)? - .to_mask_fill_null_false(ctx)) - } -} - -#[expect(dead_code)] -pub struct NullTyped<'a>(&'a ArrayRef); - -pub struct BoolTyped<'a>(&'a ArrayRef); - -impl BoolTyped<'_> { - #[deprecated( - note = "Relies on the hidden global `legacy_session()`; use `sum(array, ctx)` with an explicit `ExecutionCtx` instead" - )] - #[allow(clippy::disallowed_methods)] - pub fn true_count(&self) -> VortexResult { - let mut ctx = legacy_session().create_execution_ctx(); - let true_count = sum(self.0, &mut ctx)?; - Ok(true_count - .as_primitive() - .as_::() - .vortex_expect("true count should never be null")) - } -} - -pub struct PrimitiveTyped<'a>(&'a ArrayRef); - -impl PrimitiveTyped<'_> { - pub fn ptype(&self) -> PType { - let DType::Primitive(ptype, _) = self.0.dtype() else { - vortex_panic!("Expected Primitive DType") - }; - *ptype - } - - /// Return the primitive value at the given index. - #[deprecated( - note = "Relies on the hidden global `legacy_session()`; use `is_valid`/`execute_scalar` with an explicit `ExecutionCtx` instead" - )] - #[allow(deprecated)] - #[allow(clippy::disallowed_methods)] - pub fn value(&self, idx: usize) -> VortexResult> { - self.0 - .is_valid(idx, &mut legacy_session().create_execution_ctx())? - .then(|| self.value_unchecked(idx)) - .transpose() - } - - /// Return the primitive value at the given index, ignoring nullability. - #[deprecated( - note = "Relies on the hidden global `legacy_session()`; use `execute_scalar` with an explicit `ExecutionCtx` instead" - )] - #[allow(clippy::disallowed_methods)] - pub fn value_unchecked(&self, idx: usize) -> VortexResult { - Ok(self - .0 - .execute_scalar(idx, &mut legacy_session().create_execution_ctx())? - .as_primitive() - .pvalue() - .unwrap_or_else(|| PValue::zero(&self.ptype()))) - } -} - -impl IndexOrd> for PrimitiveTyped<'_> { - #[allow(deprecated)] - fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { - let value = self.value(idx)?; - Ok(value.partial_cmp(elem)) - } - - fn index_len(&self) -> usize { - self.0.len() - } -} - -// TODO(ngates): add generics to the `value` function and implement this over T. -impl IndexOrd for PrimitiveTyped<'_> { - #[allow(deprecated)] - #[allow(clippy::disallowed_methods)] - fn index_cmp(&self, idx: usize, elem: &PValue) -> VortexResult> { - assert!( - self.0 - .all_valid(&mut legacy_session().create_execution_ctx())? - ); - let value = self.value_unchecked(idx)?; - Ok(value.partial_cmp(elem)) - } - - fn index_len(&self) -> usize { - self.0.len() - } -} - -#[expect(dead_code)] -pub struct Utf8Typed<'a>(&'a ArrayRef); - -#[expect(dead_code)] -pub struct BinaryTyped<'a>(&'a ArrayRef); - -#[expect(dead_code)] -pub struct DecimalTyped<'a>(&'a ArrayRef); - -pub struct StructTyped<'a>(&'a ArrayRef); - -impl StructTyped<'_> { - pub fn names(&self) -> &FieldNames { - let DType::Struct(st, _) = self.0.dtype() else { - unreachable!() - }; - st.names() - } - - pub fn dtypes(&self) -> Vec { - let DType::Struct(st, _) = self.0.dtype() else { - unreachable!() - }; - st.fields().collect() - } - - pub fn nfields(&self) -> usize { - self.names().len() - } -} - -#[expect(dead_code)] -pub struct ListTyped<'a>(&'a ArrayRef); - -pub struct ExtensionTyped<'a>(&'a ArrayRef); - -impl ExtensionTyped<'_> { - /// Returns the extension logical [`DType`]. - pub fn ext_dtype(&self) -> &ExtDTypeRef { - let DType::Extension(ext_dtype) = self.0.dtype() else { - vortex_panic!("Expected ExtDType") - }; - ext_dtype - } -} From 1bdb43913eccd99c7e8495dabb3675a8c1ab2014 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 16:14:15 +0100 Subject: [PATCH 02/11] fixes Signed-off-by: Robert Kruszewski --- encodings/runend/src/ops.rs | 6 +-- vortex-array/src/arrays/primitive/tests.rs | 4 +- vortex-array/src/patches.rs | 4 +- vortex-array/src/search_sorted/primitive.rs | 43 ++++++++++----------- 4 files changed, 27 insertions(+), 30 deletions(-) diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 2efb51b4cfc..3bd686c9935 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -8,8 +8,8 @@ use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::scalar::Scalar; use vortex_array::search_sorted::SearchResult; use vortex_array::search_sorted::SearchSorted; +use vortex_array::search_sorted::SearchSortedPrimitiveArray; use vortex_array::search_sorted::SearchSortedSide; -use vortex_array::search_sorted::TypedPrimitiveArray; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexResult; @@ -38,7 +38,7 @@ pub(crate) fn find_slice_end_index( ctx: &mut ExecutionCtx, ) -> VortexResult { let result = match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { - TypedPrimitiveArray::::new(array, ctx).search_sorted(&index, SearchSortedSide::Right)? + SearchSortedPrimitiveArray::::new(array, ctx).search_sorted(&index, SearchSortedSide::Right)? }); Ok(match result { SearchResult::Found(i) => i, @@ -58,7 +58,7 @@ pub(crate) fn find_physical_index( ctx: &mut ExecutionCtx, ) -> VortexResult { match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { - Ok(TypedPrimitiveArray::::new(array, ctx) + Ok(SearchSortedPrimitiveArray::::new(array, ctx) .search_sorted(&index, SearchSortedSide::Right)? .to_ends_index(array.len())) }) diff --git a/vortex-array/src/arrays/primitive/tests.rs b/vortex-array/src/arrays/primitive/tests.rs index 9eda67ab287..cf2968637e9 100644 --- a/vortex-array/src/arrays/primitive/tests.rs +++ b/vortex-array/src/arrays/primitive/tests.rs @@ -20,8 +20,8 @@ use crate::compute::conformance::search_sorted::*; use crate::executor::VortexSessionExecute; use crate::search_sorted::SearchResult; use crate::search_sorted::SearchSorted; +use crate::search_sorted::SearchSortedPrimitiveArray; use crate::search_sorted::SearchSortedSide; -use crate::search_sorted::TypedPrimitiveArray; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(array_session); @@ -33,7 +33,7 @@ fn test_search_sorted_primitive( #[case] side: SearchSortedSide, #[case] expected: SearchResult, ) { - let res = TypedPrimitiveArray::::new(&array, &mut SESSION.create_execution_ctx()) + let res = SearchSortedPrimitiveArray::::new(&array, &mut SESSION.create_execution_ctx()) .search_sorted(&value, side) .vortex_expect("Failed to search sorted"); assert_eq!(res, expected); diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index a59db47fd51..6ab01fce495 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -39,8 +39,8 @@ use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; use crate::search_sorted::SearchResult; use crate::search_sorted::SearchSorted; +use crate::search_sorted::SearchSortedPrimitiveArray; use crate::search_sorted::SearchSortedSide; -use crate::search_sorted::TypedPrimitiveArray; use crate::validity::Validity; /// One patch index offset is stored for each chunk. @@ -1012,7 +1012,7 @@ fn search_index_binary_search_scalar( needle: usize, ) -> VortexResult { match_each_unsigned_integer_ptype!(indices.dtype().as_ptype(), |T| { - TypedPrimitiveArray::::new(indices, &mut LEGACY_SESSION.create_execution_ctx()) + SearchSortedPrimitiveArray::::new(indices, &mut LEGACY_SESSION.create_execution_ctx()) .search_sorted(&needle, SearchSortedSide::Left) .map_err(|_| vortex_err!("indices must be a primitive array")) }) diff --git a/vortex-array/src/search_sorted/primitive.rs b/vortex-array/src/search_sorted/primitive.rs index 94dced9c294..8c3ec0cf40f 100644 --- a/vortex-array/src/search_sorted/primitive.rs +++ b/vortex-array/src/search_sorted/primitive.rs @@ -12,9 +12,13 @@ use crate::ExecutionCtx; use crate::dtype::NativePType; use crate::search_sorted::IndexOrd; -pub struct TypedPrimitiveArray<'a, T>(&'a ArrayRef, RefCell<&'a mut ExecutionCtx>, PhantomData); +pub struct SearchSortedPrimitiveArray<'a, T>( + &'a ArrayRef, + RefCell<&'a mut ExecutionCtx>, + PhantomData, +); -impl<'a, T: NativePType> TypedPrimitiveArray<'a, T> { +impl<'a, T: NativePType> SearchSortedPrimitiveArray<'a, T> { pub fn new(array: &'a ArrayRef, ctx: &'a mut ExecutionCtx) -> Self { assert_eq!( array.dtype().as_ptype(), @@ -23,17 +27,21 @@ impl<'a, T: NativePType> TypedPrimitiveArray<'a, T> { ); Self(array, RefCell::new(ctx), PhantomData) } -} -impl IndexOrd for TypedPrimitiveArray<'_, T> { - fn index_cmp(&self, idx: usize, elem: &T) -> VortexResult> { - let value = self + fn value(&self, idx: usize) -> VortexResult { + let ctx_mut = &mut self.1.borrow_mut(); + Ok(self .0 - .execute_scalar(idx, &mut self.1.borrow_mut())? + .execute_scalar(idx, ctx_mut)? .as_primitive() .typed_value::() - .unwrap_or_else(|| T::zero()); + .unwrap_or_else(|| T::zero())) + } +} +impl IndexOrd for SearchSortedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &T) -> VortexResult> { + let value = self.value(idx)?; Ok(Some(value.total_compare(*elem))) } @@ -42,19 +50,13 @@ impl IndexOrd for TypedPrimitiveArray<'_, T> { } } -impl IndexOrd> for TypedPrimitiveArray<'_, T> { +impl IndexOrd> for SearchSortedPrimitiveArray<'_, T> { fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { let ctx_mut = &mut self.1.borrow_mut(); let value = self .0 .is_valid(idx, ctx_mut)? - .then(|| { - self.0.execute_scalar(idx, ctx_mut).map(|s| { - s.as_primitive() - .typed_value::() - .unwrap_or_else(|| T::zero()) - }) - }) + .then(|| self.value(idx)) .transpose()?; Ok(match (value, elem.as_ref()) { @@ -70,14 +72,9 @@ impl IndexOrd> for TypedPrimitiveArray<'_, T> { } } -impl IndexOrd for TypedPrimitiveArray<'_, T> { +impl IndexOrd for SearchSortedPrimitiveArray<'_, T> { fn index_cmp(&self, idx: usize, elem: &usize) -> VortexResult> { - let value = self - .0 - .execute_scalar(idx, &mut self.1.borrow_mut())? - .as_primitive() - .typed_value::() - .unwrap_or_else(|| T::zero()); + let value = self.value(idx)?; let Some(elem_t) = T::from_usize(*elem) else { return Ok(Some(Ordering::Less)); From 9eaf3841ec45157c1d12077da51fccf809bb8b6b Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 16:19:19 +0100 Subject: [PATCH 03/11] format Signed-off-by: Robert Kruszewski --- .../src/decimal_byte_parts/compute/compare.rs | 5 +-- encodings/runend/src/ops.rs | 3 +- vortex-array/src/patches.rs | 31 +++++++++---------- 3 files changed, 18 insertions(+), 21 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index 602dbc26f77..af6e98bddc2 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -47,10 +47,7 @@ impl CompareKernel for DecimalByteParts { .decimal_value() .vortex_expect("checked for null in entry func"); - match decimal_value_wrapper_to_primitive( - rhs_decimal, - lhs.msp().dtype().as_ptype(), - ) { + match decimal_value_wrapper_to_primitive(rhs_decimal, lhs.msp().dtype().as_ptype()) { Ok(value) => { let encoded_scalar = Scalar::try_new(scalar_type, Some(value))?; let encoded_const = ConstantArray::new(encoded_scalar, rhs.len()); diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 3bd686c9935..70dcee4a0a4 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -38,7 +38,8 @@ pub(crate) fn find_slice_end_index( ctx: &mut ExecutionCtx, ) -> VortexResult { let result = match_each_unsigned_integer_ptype!(array.dtype().as_ptype(), |T| { - SearchSortedPrimitiveArray::::new(array, ctx).search_sorted(&index, SearchSortedSide::Right)? + SearchSortedPrimitiveArray::::new(array, ctx) + .search_sorted(&index, SearchSortedSide::Right)? }); Ok(match result { SearchResult::Found(i) => i, diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 6ab01fce495..ea1e92b3f3c 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -991,17 +991,17 @@ impl Patches { fn search_index_binary_search(indices: &ArrayRef, needle: usize) -> VortexResult { if let Some(primitive) = indices.as_opt::() { match_each_unsigned_integer_ptype!(primitive.ptype(), |T| { - let Ok(needle) = T::try_from(needle) else { - // If the needle is not of type T, then it cannot possibly be in this array. - // - // The needle is a non-negative integer (a usize); therefore, it must be larger - // than all values in this array. - return Ok(SearchResult::NotFound(primitive.len())); - }; - return primitive - .as_slice::() - .search_sorted(&needle, SearchSortedSide::Left); - }); + let Ok(needle) = T::try_from(needle) else { + // If the needle is not of type T, then it cannot possibly be in this array. + // + // The needle is a non-negative integer (a usize); therefore, it must be larger + // than all values in this array. + return Ok(SearchResult::NotFound(primitive.len())); + }; + return primitive + .as_slice::() + .search_sorted(&needle, SearchSortedSide::Left); + }); } search_index_binary_search_scalar(indices, needle) @@ -1012,13 +1012,12 @@ fn search_index_binary_search_scalar( needle: usize, ) -> VortexResult { match_each_unsigned_integer_ptype!(indices.dtype().as_ptype(), |T| { - SearchSortedPrimitiveArray::::new(indices, &mut LEGACY_SESSION.create_execution_ctx()) - .search_sorted(&needle, SearchSortedSide::Left) - .map_err(|_| vortex_err!("indices must be a primitive array")) - }) + SearchSortedPrimitiveArray::::new(indices, &mut LEGACY_SESSION.create_execution_ctx()) + .search_sorted(&needle, SearchSortedSide::Left) + .map_err(|_| vortex_err!("indices must be a primitive array")) + }) } - #[expect(clippy::too_many_arguments)] // private function, can clean up one day fn take_map, T: NativePType>( indices: &[I], From b59f8106689959e18d15cec2f6e57df2edf175e5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 16:29:14 +0100 Subject: [PATCH 04/11] Address review findings on SearchSortedPrimitiveArray - Fix RefCell double-borrow panic in IndexOrd>::index_cmp: the RefMut guard from is_valid was still alive when value() re-borrowed the ctx, so searching any valid element with an optional needle panicked with BorrowMutError. Add regression tests covering Option searches with and without nulls. - Stop swallowing the underlying error in Patches search_index_binary_search_scalar; the map_err replaced real execution errors with a misleading "indices must be a primitive array" message. - Pass the in-scope ExecutionCtx to find_physical_index in RunEnd scalar_at and slice instead of spinning up a fresh LEGACY_SESSION ctx per call. - Add doc comments to the new public SearchSortedPrimitiveArray API and document the null-as-zero semantics of plain T searches. - Restore VortexResult<()> + ? style in test_search_sorted_primitive. Signed-off-by: "Robert Kruszewski" Co-Authored-By: Claude Fable 5 --- encodings/runend/src/kernel.rs | 3 +- encodings/runend/src/ops.rs | 5 +- vortex-array/src/arrays/primitive/tests.rs | 7 +-- vortex-array/src/patches.rs | 1 - vortex-array/src/search_sorted/primitive.rs | 64 ++++++++++++++++++--- 5 files changed, 63 insertions(+), 17 deletions(-) diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 4f3f2a4fd78..309a0140679 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -25,6 +25,7 @@ use vortex_session::VortexSession; use crate::RunEnd; use crate::array::RunEndArrayExt; use crate::compute::take_from::RunEndTakeFrom; +use crate::ops::find_physical_index; use crate::ops::find_slice_end_index; pub(super) fn initialize(session: &VortexSession) { @@ -64,7 +65,7 @@ fn slice( ) -> VortexResult { let new_length = range.len(); - let slice_begin = array.find_physical_index(range.start)?; + let slice_begin = find_physical_index(array.ends(), range.start + array.offset(), ctx)?; let slice_end = find_slice_end_index(array.ends(), range.end + array.offset(), ctx)?; // If the sliced range contains only a single run, opt to return a ConstantArray. diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 70dcee4a0a4..d31bb061011 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -22,9 +22,8 @@ impl OperationsVTable for RunEnd { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - array - .values() - .execute_scalar(array.find_physical_index(index)?, ctx) + let physical_index = find_physical_index(array.ends(), index + array.offset(), ctx)?; + array.values().execute_scalar(physical_index, ctx) } } diff --git a/vortex-array/src/arrays/primitive/tests.rs b/vortex-array/src/arrays/primitive/tests.rs index cf2968637e9..1f50c8edb1c 100644 --- a/vortex-array/src/arrays/primitive/tests.rs +++ b/vortex-array/src/arrays/primitive/tests.rs @@ -4,7 +4,6 @@ use std::sync::LazyLock; use vortex_buffer::buffer; -use vortex_error::VortexExpect; use vortex_session::VortexSession; use crate::ArrayRef; @@ -32,11 +31,11 @@ fn test_search_sorted_primitive( #[case] value: i32, #[case] side: SearchSortedSide, #[case] expected: SearchResult, -) { +) -> vortex_error::VortexResult<()> { let res = SearchSortedPrimitiveArray::::new(&array, &mut SESSION.create_execution_ctx()) - .search_sorted(&value, side) - .vortex_expect("Failed to search sorted"); + .search_sorted(&value, side)?; assert_eq!(res, expected); + Ok(()) } #[test] diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index ea1e92b3f3c..05116dc213d 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -1014,7 +1014,6 @@ fn search_index_binary_search_scalar( match_each_unsigned_integer_ptype!(indices.dtype().as_ptype(), |T| { SearchSortedPrimitiveArray::::new(indices, &mut LEGACY_SESSION.create_execution_ctx()) .search_sorted(&needle, SearchSortedSide::Left) - .map_err(|_| vortex_err!("indices must be a primitive array")) }) } diff --git a/vortex-array/src/search_sorted/primitive.rs b/vortex-array/src/search_sorted/primitive.rs index 8c3ec0cf40f..3da28c7be2b 100644 --- a/vortex-array/src/search_sorted/primitive.rs +++ b/vortex-array/src/search_sorted/primitive.rs @@ -12,6 +12,12 @@ use crate::ExecutionCtx; use crate::dtype::NativePType; use crate::search_sorted::IndexOrd; +/// A [`SearchSorted`](crate::search_sorted::SearchSorted) adapter over a sorted primitive-typed +/// array, comparing elements as the native type `T`. +/// +/// Values can be searched as `T`, `Option`, or `usize`. Searching as `T` or `usize` treats +/// null elements as `T::zero()`; use `Option` when the array may contain nulls, in which case +/// nulls sort before all non-null values. pub struct SearchSortedPrimitiveArray<'a, T>( &'a ArrayRef, RefCell<&'a mut ExecutionCtx>, @@ -19,6 +25,8 @@ pub struct SearchSortedPrimitiveArray<'a, T>( ); impl<'a, T: NativePType> SearchSortedPrimitiveArray<'a, T> { + /// Wraps `array` for searching, panicking if the array's [`PType`](crate::dtype::PType) is + /// not `T::PTYPE`. pub fn new(array: &'a ArrayRef, ctx: &'a mut ExecutionCtx) -> Self { assert_eq!( array.dtype().as_ptype(), @@ -28,11 +36,11 @@ impl<'a, T: NativePType> SearchSortedPrimitiveArray<'a, T> { Self(array, RefCell::new(ctx), PhantomData) } + /// Returns the value at `idx`, with nulls mapped to `T::zero()`. fn value(&self, idx: usize) -> VortexResult { - let ctx_mut = &mut self.1.borrow_mut(); Ok(self .0 - .execute_scalar(idx, ctx_mut)? + .execute_scalar(idx, &mut self.1.borrow_mut())? .as_primitive() .typed_value::() .unwrap_or_else(|| T::zero())) @@ -52,12 +60,9 @@ impl IndexOrd for SearchSortedPrimitiveArray<'_, T> { impl IndexOrd> for SearchSortedPrimitiveArray<'_, T> { fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { - let ctx_mut = &mut self.1.borrow_mut(); - let value = self - .0 - .is_valid(idx, ctx_mut)? - .then(|| self.value(idx)) - .transpose()?; + // The borrow must end before `self.value` re-borrows the ctx. + let valid = self.0.is_valid(idx, &mut self.1.borrow_mut())?; + let value = valid.then(|| self.value(idx)).transpose()?; Ok(match (value, elem.as_ref()) { (Some(l), Some(r)) => Some(l.total_compare(*r)), @@ -87,3 +92,46 @@ impl IndexOrd for SearchSortedPrimitiveArray<'_, T> { self.0.len() } } + +#[cfg(test)] +mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::IntoArray; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::executor::VortexSessionExecute; + use crate::search_sorted::SearchResult; + use crate::search_sorted::SearchSorted; + use crate::search_sorted::SearchSortedPrimitiveArray; + use crate::search_sorted::SearchSortedSide; + use crate::validity::Validity; + + #[test] + fn search_sorted_optional_value() -> VortexResult<()> { + let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::AllValid).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let res = SearchSortedPrimitiveArray::::new(&array, &mut ctx) + .search_sorted(&Some(3i32), SearchSortedSide::Left)?; + assert_eq!(res, SearchResult::Found(2)); + Ok(()) + } + + #[test] + fn search_sorted_optional_value_with_nulls() -> VortexResult<()> { + let array = + PrimitiveArray::from_option_iter([None, None, Some(2i32), Some(3)]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + let searcher = SearchSortedPrimitiveArray::::new(&array, &mut ctx); + assert_eq!( + searcher.search_sorted(&Some(2i32), SearchSortedSide::Left)?, + SearchResult::Found(2) + ); + assert_eq!( + searcher.search_sorted(&None, SearchSortedSide::Right)?, + SearchResult::Found(2) + ); + Ok(()) + } +} From e7069dfd8164b06e2914a9724f135d5c108d3963 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 18:20:07 +0100 Subject: [PATCH 05/11] fixes Signed-off-by: Robert Kruszewski --- encodings/runend/src/array.rs | 13 +++++++------ encodings/runend/src/kernel.rs | 6 ++---- encodings/runend/src/ops.rs | 2 +- 3 files changed, 10 insertions(+), 11 deletions(-) diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index d8f0648dd70..637947bba53 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -46,6 +46,7 @@ use crate::compress::runend_decode_varbinview; use crate::compress::runend_encode; use crate::decompress_bool::runend_decode_bools; use crate::ops::find_physical_index; +use crate::ops::find_slice_end_index; use crate::rules::RULES; /// A [`RunEnd`]-encoded Vortex array. @@ -228,12 +229,12 @@ pub trait RunEndArrayExt: TypedArrayRef { self.values().dtype() } - fn find_physical_index(&self, index: usize) -> VortexResult { - find_physical_index( - self.ends(), - index + self.offset(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) + fn find_physical_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + find_physical_index(self.ends(), index + self.offset(), ctx) + } + + fn find_slice_end_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult { + find_slice_end_index(self.ends(), index + self.offset(), ctx) } } diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 309a0140679..13f32df07a1 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -25,8 +25,6 @@ use vortex_session::VortexSession; use crate::RunEnd; use crate::array::RunEndArrayExt; use crate::compute::take_from::RunEndTakeFrom; -use crate::ops::find_physical_index; -use crate::ops::find_slice_end_index; pub(super) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -65,8 +63,8 @@ fn slice( ) -> VortexResult { let new_length = range.len(); - let slice_begin = find_physical_index(array.ends(), range.start + array.offset(), ctx)?; - let slice_end = find_slice_end_index(array.ends(), range.end + array.offset(), ctx)?; + let slice_begin = array.find_physical_index(range.start, ctx)?; + let slice_end = array.find_slice_end_index(range.end, ctx)?; // If the sliced range contains only a single run, opt to return a ConstantArray. if slice_begin + 1 == slice_end { diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index d31bb061011..c929f3d4251 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -22,7 +22,7 @@ impl OperationsVTable for RunEnd { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - let physical_index = find_physical_index(array.ends(), index + array.offset(), ctx)?; + let physical_index = array.find_physical_index(index, ctx)?; array.values().execute_scalar(physical_index, ctx) } } From 6d8da7da14ede5ce4687ad9b45e96a76e7617491 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 18:25:58 +0100 Subject: [PATCH 06/11] DCO Remediation Commit for Robert Kruszewski I, Robert Kruszewski , hereby add my Signed-off-by to this commit: 12082ae4649735da9e465e71f29e0668c2ab6ef8 Signed-off-by: Robert Kruszewski From a6cc3f2b858a68c626a8098244c29ebb8646ebd0 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 22:22:15 +0100 Subject: [PATCH 07/11] fixes Signed-off-by: Robert Kruszewski --- encodings/runend/src/arrow.rs | 1 - vortex-array/src/patches.rs | 5 ++--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/encodings/runend/src/arrow.rs b/encodings/runend/src/arrow.rs index a24e69ef5c8..1f58e21bef2 100644 --- a/encodings/runend/src/arrow.rs +++ b/encodings/runend/src/arrow.rs @@ -11,7 +11,6 @@ use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::NativePType; use vortex_array::legacy_session; -use vortex_array::search_sorted::SearchSorted; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 05116dc213d..fec9b826e19 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -19,7 +19,7 @@ use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_utils::aliases::hash_map::HashMap; -use crate::ArrayRef; +use crate::{legacy_session, ArrayRef}; use crate::ArraySlots; use crate::ExecutionCtx; use crate::IntoArray; @@ -275,7 +275,6 @@ impl Patches { #[cfg(debug_assertions)] { - use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_sorted::is_sorted; let mut ctx = legacy_session().create_execution_ctx(); assert!( @@ -1012,7 +1011,7 @@ fn search_index_binary_search_scalar( needle: usize, ) -> VortexResult { match_each_unsigned_integer_ptype!(indices.dtype().as_ptype(), |T| { - SearchSortedPrimitiveArray::::new(indices, &mut LEGACY_SESSION.create_execution_ctx()) + SearchSortedPrimitiveArray::::new(indices, &mut legacy_session().create_execution_ctx()) .search_sorted(&needle, SearchSortedSide::Left) }) } From 47721dbd268d699305620b8321baa846d40dd812 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 22:27:07 +0100 Subject: [PATCH 08/11] fix Signed-off-by: Robert Kruszewski --- vortex-array/src/patches.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index fec9b826e19..f4596a5ea3e 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -19,7 +19,7 @@ use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_utils::aliases::hash_map::HashMap; -use crate::{legacy_session, ArrayRef}; +use crate::ArrayRef; use crate::ArraySlots; use crate::ExecutionCtx; use crate::IntoArray; @@ -35,6 +35,7 @@ use crate::dtype::Nullability; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; use crate::dtype::UnsignedPType; +use crate::legacy_session; use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; use crate::search_sorted::SearchResult; @@ -987,7 +988,10 @@ impl Patches { /// # Returns /// [`SearchResult::Found`] with the position if needle exists, or [`SearchResult::NotFound`] /// with the insertion point if not found. -fn search_index_binary_search(indices: &ArrayRef, needle: usize) -> VortexResult { +fn search_index_binary_search( + indices: &ArrayRef, + needle: usize, +) -> VortexResult { if let Some(primitive) = indices.as_opt::() { match_each_unsigned_integer_ptype!(primitive.ptype(), |T| { let Ok(needle) = T::try_from(needle) else { @@ -1006,6 +1010,7 @@ fn search_index_binary_search(indices: &ArrayRef, needle: usize) -> VortexResult search_index_binary_search_scalar(indices, needle) } +#[allow(clippy::disallowed_methods)] fn search_index_binary_search_scalar( indices: &ArrayRef, needle: usize, From 6dd1f0ece15ae6292e322c62e4f49f5393a1ab6b Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 22:51:32 +0100 Subject: [PATCH 09/11] DCO Remediation Commit for Robert Kruszewski I, Robert Kruszewski , hereby add my Signed-off-by to this commit: 1d76252036d102816593c20c186f2464d4ba9ffa Signed-off-by: Robert Kruszewski From ddd570daaac92ae2112a9ff443bc075de689896f Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 22:57:23 +0100 Subject: [PATCH 10/11] format Signed-off-by: Robert Kruszewski --- vortex-array/src/patches.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index f4596a5ea3e..f5d7b0e2d2a 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -988,10 +988,7 @@ impl Patches { /// # Returns /// [`SearchResult::Found`] with the position if needle exists, or [`SearchResult::NotFound`] /// with the insertion point if not found. -fn search_index_binary_search( - indices: &ArrayRef, - needle: usize, -) -> VortexResult { +fn search_index_binary_search(indices: &ArrayRef, needle: usize) -> VortexResult { if let Some(primitive) = indices.as_opt::() { match_each_unsigned_integer_ptype!(primitive.ptype(), |T| { let Ok(needle) = T::try_from(needle) else { From ccfbcdb7752eafa6cb0056f683491aa5e11d6458 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 2 Jul 2026 23:07:03 +0100 Subject: [PATCH 11/11] fixes Signed-off-by: Robert Kruszewski --- vortex-array/benches/listview_zip.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/vortex-array/benches/listview_zip.rs b/vortex-array/benches/listview_zip.rs index 28db66f04a9..c50180a0e1a 100644 --- a/vortex-array/benches/listview_zip.rs +++ b/vortex-array/benches/listview_zip.rs @@ -3,23 +3,28 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ListViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_mask::Mask; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = LazyLock::new(array_session); + // Smaller than the value-path benches: listview zip cost is dominated by element concatenation and // per-list canonicalization. A few thousand lists already exercise the select while keeping each // case well under a few hundred microseconds under CodSpeed's instruction-count simulation, which @@ -50,7 +55,7 @@ fn run(bencher: Bencher, if_true: ArrayRef, if_false: ArrayRef) { if_true.clone(), if_false.clone(), mask.clone().into_array(), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(t, f, m, ctx)| {