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..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().as_primitive_typed().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/array.rs b/encodings/runend/src/array.rs index 9dba648e248..637947bba53 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,8 @@ 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::ops::find_slice_end_index; use crate::rules::RULES; /// A [`RunEnd`]-encoded Vortex array. @@ -230,17 +229,15 @@ pub trait RunEndArrayExt: TypedArrayRef { self.values().dtype() } - 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())) + 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) } } + impl> RunEndArrayExt for T {} #[derive(Clone, Debug)] diff --git a/encodings/runend/src/arrow.rs b/encodings/runend/src/arrow.rs index c927daaaffe..1f58e21bef2 100644 --- a/encodings/runend/src/arrow.rs +++ b/encodings/runend/src/arrow.rs @@ -11,14 +11,12 @@ 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 +39,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 +54,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 +87,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 +94,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 +107,7 @@ mod tests { fn decode_run_array( array: &RunArray, nullable: bool, - ) -> VortexResult + ) -> VortexResult where R::Native: NativePType, { @@ -131,14 +125,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 +138,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..13f32df07a1 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -63,8 +63,8 @@ fn slice( ) -> VortexResult { 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_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 56de024f988..c929f3d4251 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -4,10 +4,11 @@ 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::SearchSortedPrimitiveArray; use vortex_array::search_sorted::SearchSortedSide; use vortex_array::vtable::OperationsVTable; use vortex_error::VortexResult; @@ -21,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 = array.find_physical_index(index, ctx)?; + array.values().execute_scalar(physical_index, ctx) } } @@ -31,10 +31,15 @@ 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| { + SearchSortedPrimitiveArray::::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(SearchSortedPrimitiveArray::::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/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)| { 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..1f50c8edb1c 100644 --- a/vortex-array/src/arrays/primitive/tests.rs +++ b/vortex-array/src/arrays/primitive/tests.rs @@ -1,10 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::LazyLock; + use vortex_buffer::buffer; +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,12 +16,15 @@ 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::SearchSortedPrimitiveArray; use crate::search_sorted::SearchSortedSide; use crate::validity::Validity; +static SESSION: LazyLock = LazyLock::new(array_session); + #[apply(search_sorted_conformance)] fn test_search_sorted_primitive( #[case] array: ArrayRef, @@ -25,9 +32,8 @@ fn test_search_sorted_primitive( #[case] side: SearchSortedSide, #[case] expected: SearchResult, ) -> vortex_error::VortexResult<()> { - let res = array - .as_primitive_typed() - .search_sorted(&Some(PValue::from(value)), side)?; + let res = SearchSortedPrimitiveArray::::new(&array, &mut SESSION.create_execution_ctx()) + .search_sorted(&value, side)?; assert_eq!(res, expected); Ok(()) } 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..f5d7b0e2d2a 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -36,12 +36,11 @@ 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::SearchSortedPrimitiveArray; use crate::search_sorted::SearchSortedSide; use crate::validity::Validity; @@ -277,7 +276,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!( @@ -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,41 @@ 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) +} + +#[allow(clippy::disallowed_methods)] +fn search_index_binary_search_scalar( + indices: &ArrayRef, + 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) + }) +} + #[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..3da28c7be2b --- /dev/null +++ b/vortex-array/src/search_sorted/primitive.rs @@ -0,0 +1,137 @@ +// 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; + +/// 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>, + PhantomData, +); + +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(), + T::PTYPE, + "Array PType must match primitive type" + ); + Self(array, RefCell::new(ctx), PhantomData) + } + + /// Returns the value at `idx`, with nulls mapped to `T::zero()`. + fn value(&self, idx: usize) -> VortexResult { + Ok(self + .0 + .execute_scalar(idx, &mut self.1.borrow_mut())? + .as_primitive() + .typed_value::() + .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))) + } + + fn index_len(&self) -> usize { + self.0.len() + } +} + +impl IndexOrd> for SearchSortedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &Option) -> VortexResult> { + // 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)), + (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 SearchSortedPrimitiveArray<'_, T> { + fn index_cmp(&self, idx: usize, elem: &usize) -> VortexResult> { + let value = self.value(idx)?; + + 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() + } +} + +#[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(()) + } +} 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 - } -}