Remove search sorted usize (#3240)

This PR impacts the performance of random access reads and should be
fixed up in the future. For now, it's blocking an incremental migration
of compute functions so I will merge and follow up.

Author: gatesn

encodings/fastlanes/src/bitpacking/compute/search_sorted.rs

@@ -6,7 +6,7 @@ use itertools::Itertools;
 use num_traits::AsPrimitive;
 use vortex_array::Array;
 use vortex_array::compute::{
-    IndexOrd, SearchResult, SearchSorted, SearchSortedFn, SearchSortedSide, SearchSortedUsizeFn,
+    IndexOrd, SearchResult, SearchSorted, SearchSortedFn, SearchSortedSide,
 };
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_dtype::{DType, NativePType, match_each_unsigned_integer_ptype};
@@ -50,46 +50,6 @@ impl SearchSortedFn<&BitPackedArray> for BitPackedEncoding {
     }
 }
 
-impl SearchSortedUsizeFn<&BitPackedArray> for BitPackedEncoding {
-    fn search_sorted_usize(
-        &self,
-        array: &BitPackedArray,
-        value: usize,
-        side: SearchSortedSide,
-    ) -> VortexResult<SearchResult> {
-        match_each_unsigned_integer_ptype!(array.ptype().to_unsigned(), |$P| {
-            // NOTE: conversion may truncate silently.
-            if let Some(pvalue) = num_traits::cast::<usize, $P>(value) {
-                search_sorted_native(array, pvalue, side)
-            } else {
-                // provided u64 is too large to fit in the provided PType, value must be off
-                // the right end of the array.
-                Ok(SearchResult::NotFound(array.len()))
-            }
-        })
-    }
-
-    fn search_sorted_usize_many(
-        &self,
-        array: &BitPackedArray,
-        values: &[usize],
-        side: SearchSortedSide,
-    ) -> VortexResult<Vec<SearchResult>> {
-        match_each_unsigned_integer_ptype!(array.ptype().to_unsigned(), |$P| {
-            let searcher = BitPackedSearch::<'_, $P>::try_new(array)?;
-
-            values
-                .iter()
-                .map(|&value| {
-                    // NOTE: truncating cast
-                    let cast_value: $P = value as $P;
-                    Ok(searcher.search_sorted(&cast_value, side))
-                })
-                .try_collect()
-        })
-    }
-}
-
 fn search_sorted_typed<T>(
     array: &BitPackedArray,
     value: &Scalar,
@@ -191,7 +151,7 @@ impl<T: BitPacking + NativePType> IndexOrd<T> for BitPackedSearch<'_, T> {
         Some(val.total_compare(*elem))
     }
 
-    fn len(&self) -> usize {
+    fn index_len(&self) -> usize {
         self.length
     }
 }
@@ -206,6 +166,7 @@ mod test {
         SearchResult, SearchSortedSide, search_sorted, search_sorted_many,
     };
     use vortex_array::validity::Validity;
+    use vortex_array::variants::PrimitiveArrayTrait;
     use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
     use vortex_buffer::buffer;
     use vortex_error::VortexUnwrap;
@@ -219,8 +180,6 @@ mod test {
         #[case] side: SearchSortedSide,
         #[case] expected: SearchResult,
     ) {
-        use vortex_array::variants::PrimitiveArrayTrait;
-
         let primitive_array = array.to_primitive().vortex_unwrap();
         // force patches
         let histogram = bit_width_histogram(&primitive_array).vortex_unwrap();

encodings/runend/src/array.rs

@@ -1,19 +1,19 @@
 use std::fmt::Debug;
 
 use vortex_array::arrays::PrimitiveArray;
-use vortex_array::compute::{SearchSortedSide, search_sorted_usize, search_sorted_usize_many};
+use vortex_array::compute::{SearchSorted, SearchSortedSide};
 use vortex_array::stats::{ArrayStats, StatsSetRef};
 use vortex_array::variants::{BoolArrayTrait, DecimalArrayTrait, PrimitiveArrayTrait};
 use vortex_array::vtable::VTableRef;
 use vortex_array::{
     Array, ArrayCanonicalImpl, ArrayImpl, ArrayRef, ArrayStatisticsImpl, ArrayValidityImpl,
-    ArrayVariantsImpl, Canonical, Encoding, IntoArray, ProstMetadata, ToCanonical,
+    ArrayVariants, ArrayVariantsImpl, Canonical, Encoding, IntoArray, ProstMetadata, ToCanonical,
     try_from_array_ref,
 };
-use vortex_buffer::Buffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect as _, VortexResult, vortex_bail};
 use vortex_mask::Mask;
+use vortex_scalar::PValue;
 
 use crate::compress::{runend_decode_bools, runend_decode_primitive, runend_encode};
 use crate::serde::RunEndMetadata;
@@ -84,21 +84,33 @@ impl RunEndArray {
 
     /// Convert the given logical index to an index into the `values` array
     pub fn find_physical_index(&self, index: usize) -> VortexResult<usize> {
-        search_sorted_usize(self.ends(), index + self.offset(), SearchSortedSide::Right)
-            .map(|s| s.to_ends_index(self.ends().len()))
+        Ok(self
+            .ends()
+            .as_primitive_typed()
+            .vortex_expect("ends array must be primitive")
+            .search_sorted(
+                &Some(PValue::from(index + self.offset())),
+                SearchSortedSide::Right,
+            )
+            .to_ends_index(self.ends().len()))
     }
 
     /// Convert a batch of logical indices into an index for the values. Expects indices to be adjusted by offset unlike
     /// [Self::find_physical_index]
     ///
     /// See: [find_physical_index][Self::find_physical_index].
-    pub fn find_physical_indices(&self, indices: &[usize]) -> VortexResult<Buffer<u64>> {
-        search_sorted_usize_many(self.ends(), indices, SearchSortedSide::Right).map(|results| {
-            results
-                .into_iter()
-                .map(|result| result.to_ends_index(self.ends().len()) as u64)
-                .collect()
-        })
+    pub fn find_physical_indices<I: IntoIterator<Item = usize>>(
+        &self,
+        indices: I,
+    ) -> impl Iterator<Item = usize> {
+        self.ends()
+            .as_primitive_typed()
+            .vortex_expect("ends array must be primitive")
+            .search_sorted_many(
+                indices.into_iter().map(|i| Some(PValue::from(i))),
+                SearchSortedSide::Right,
+            )
+            .map(|result| result.to_ends_index(self.ends().len()))
     }
 
     /// Run the array through run-end encoding.

encodings/runend/src/compute/take.rs

@@ -2,6 +2,7 @@ use num_traits::AsPrimitive;
 use vortex_array::compute::{TakeKernel, TakeKernelAdapter, take};
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
+use vortex_buffer::Buffer;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::{VortexResult, vortex_bail};
 
@@ -37,11 +38,11 @@ pub fn take_indices_unchecked<T: AsPrimitive<usize>>(
     array: &RunEndArray,
     indices: &[T],
 ) -> VortexResult<ArrayRef> {
-    let adjusted_indices = indices
-        .iter()
-        .map(|idx| idx.as_() + array.offset())
-        .collect::<Vec<_>>();
-    let physical_indices = array.find_physical_indices(&adjusted_indices)?.into_array();
+    let physical_indices = array
+        .find_physical_indices(indices.iter().map(|idx| idx.as_() + array.offset()))
+        .map(|idx| idx as u64)
+        .collect::<Buffer<u64>>()
+        .into_array();
     take(array.values(), &physical_indices)
 }
 

encodings/sparse/src/compute/mod.rs

@@ -1,7 +1,5 @@
 use vortex_array::arrays::ConstantArray;
-use vortex_array::compute::{
-    FilterKernel, FilterKernelAdapter, SearchSortedFn, SearchSortedUsizeFn,
-};
+use vortex_array::compute::{FilterKernel, FilterKernelAdapter, SearchSortedFn};
 use vortex_array::vtable::ComputeVTable;
 use vortex_array::{Array, ArrayRef, register_kernel};
 use vortex_error::VortexResult;
@@ -18,10 +16,6 @@ impl ComputeVTable for SparseEncoding {
     fn search_sorted_fn(&self) -> Option<&dyn SearchSortedFn<&dyn Array>> {
         Some(self)
     }
-
-    fn search_sorted_usize_fn(&self) -> Option<&dyn SearchSortedUsizeFn<&dyn Array>> {
-        Some(self)
-    }
 }
 
 impl FilterKernel for SparseEncoding {

encodings/sparse/src/compute/search_sorted.rs

@@ -1,7 +1,7 @@
 use std::cmp::Ordering;
 
 use vortex_array::Array;
-use vortex_array::compute::{SearchResult, SearchSortedFn, SearchSortedSide, SearchSortedUsizeFn};
+use vortex_array::compute::{SearchResult, SearchSortedFn, SearchSortedSide};
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_scalar::Scalar;
 
@@ -99,21 +99,6 @@ fn fill_position(array: &SparseArray, side: SearchSortedSide) -> VortexResult<us
     })
 }
 
-impl SearchSortedUsizeFn<&SparseArray> for SparseEncoding {
-    fn search_sorted_usize(
-        &self,
-        array: &SparseArray,
-        value: usize,
-        side: SearchSortedSide,
-    ) -> VortexResult<SearchResult> {
-        let Ok(target) = Scalar::from(value).cast(array.dtype()) else {
-            // If the downcast fails, then the target is too large for the dtype.
-            return Ok(SearchResult::NotFound(array.len()));
-        };
-        SearchSortedFn::search_sorted(self, array, &target, side)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use vortex_array::compute::conformance::search_sorted::rstest_reuse::apply;

fuzz/src/search_sorted.rs

@@ -17,7 +17,7 @@ impl<T: PartialOrd + Debug> IndexOrd<Option<T>> for SearchNullableSlice<T> {
         unsafe { self.0.get_unchecked(idx) }.partial_cmp(elem)
     }
 
-    fn len(&self) -> usize {
+    fn index_len(&self) -> usize {
         self.0.len()
     }
 }
@@ -38,7 +38,7 @@ impl<T: NativePType> IndexOrd<Option<T>> for SearchPrimitiveSlice<T> {
         }
     }
 
-    fn len(&self) -> usize {
+    fn index_len(&self) -> usize {
         self.0.len()
     }
 }

vortex-array/src/arrays/chunked/compute/take.rs

@@ -7,7 +7,7 @@ use vortex_scalar::Scalar;
 use crate::arrays::ChunkedEncoding;
 use crate::arrays::chunked::ChunkedArray;
 use crate::compute::{
-    SearchSortedSide, TakeKernel, TakeKernelAdapter, cast, search_sorted_usize, sub_scalar, take,
+    SearchSortedSide, TakeKernel, TakeKernelAdapter, cast, search_sorted, sub_scalar, take,
 };
 use crate::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
 
@@ -74,8 +74,7 @@ fn take_strict_sorted(chunked: &ChunkedArray, indices: &dyn Array) -> VortexResu
         // Find the end of this chunk, and locate that position in the indices array.
         let chunk_begin = usize::try_from(chunked.chunk_offsets()[chunk_idx])?;
         let chunk_end = usize::try_from(chunked.chunk_offsets()[chunk_idx + 1])?;
-        let chunk_end_pos =
-            search_sorted_usize(indices, chunk_end, SearchSortedSide::Left)?.to_index();
+        let chunk_end_pos = search_sorted(indices, chunk_end, SearchSortedSide::Left)?.to_index();
 
         // Now we can say the slice of indices belonging to this chunk is [pos, chunk_end_pos)
         let chunk_indices = indices.slice(pos, chunk_end_pos)?;

vortex-array/src/arrays/chunked/variants.rs

@@ -2,13 +2,14 @@ use std::sync::Arc;
 
 use vortex_dtype::{DType, FieldName};
 use vortex_error::{VortexExpect, VortexResult, vortex_err, vortex_panic};
+use vortex_scalar::PValue;
 
 use crate::arrays::chunked::ChunkedArray;
 use crate::variants::{
     BinaryArrayTrait, BoolArrayTrait, DecimalArrayTrait, ExtensionArrayTrait, ListArrayTrait,
     NullArrayTrait, PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait,
 };
-use crate::{Array, ArrayRef, ArrayVariantsImpl};
+use crate::{Array, ArrayRef, ArrayVariants, ArrayVariantsImpl};
 
 /// Chunked arrays support all DTypes
 impl ArrayVariantsImpl for ChunkedArray {
@@ -53,7 +54,19 @@ impl NullArrayTrait for ChunkedArray {}
 
 impl BoolArrayTrait for ChunkedArray {}
 
-impl PrimitiveArrayTrait for ChunkedArray {}
+impl PrimitiveArrayTrait for ChunkedArray {
+    fn value(&self, idx: usize) -> Option<PValue> {
+        let (chunk, offset_in_chunk) = self.find_chunk_idx(idx);
+        let chunk = self
+            .chunks()
+            .get(chunk)
+            .vortex_expect("Chunk index out of bounds");
+        chunk
+            .as_primitive_typed()
+            .vortex_expect("Chunk was not a PrimitiveArray")
+            .value(offset_in_chunk)
+    }
+}
 
 impl DecimalArrayTrait for ChunkedArray {}
 

vortex-array/src/arrays/constant/variants.rs

@@ -1,5 +1,6 @@
 use vortex_dtype::FieldName;
 use vortex_error::VortexResult;
+use vortex_scalar::PValue;
 
 use crate::arrays::constant::ConstantArray;
 use crate::variants::{
@@ -51,7 +52,11 @@ impl NullArrayTrait for ConstantArray {}
 
 impl BoolArrayTrait for ConstantArray {}
 
-impl PrimitiveArrayTrait for ConstantArray {}
+impl PrimitiveArrayTrait for ConstantArray {
+    fn value(&self, _idx: usize) -> Option<PValue> {
+        self.scalar().as_primitive().pvalue()
+    }
+}
 
 impl Utf8ArrayTrait for ConstantArray {}
 

vortex-array/src/arrays/primitive/compute/mod.rs

@@ -1,6 +1,6 @@
 use crate::Array;
 use crate::arrays::PrimitiveEncoding;
-use crate::compute::{SearchSortedFn, SearchSortedUsizeFn};
+use crate::compute::SearchSortedFn;
 use crate::vtable::ComputeVTable;
 
 mod between;
@@ -22,8 +22,4 @@ impl ComputeVTable for PrimitiveEncoding {
     fn search_sorted_fn(&self) -> Option<&dyn SearchSortedFn<&dyn Array>> {
         Some(self)
     }
-
-    fn search_sorted_usize_fn(&self) -> Option<&dyn SearchSortedUsizeFn<&dyn Array>> {
-        Some(self)
-    }
 }