SearchSorted Many Side (#1427)

Fixes #1413

Author: gatesn

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

@@ -53,18 +53,16 @@ impl SearchSortedFn<BitPackedArray> for BitPackedEncoding {
         &self,
         array: &BitPackedArray,
         values: &[Scalar],
-        sides: &[SearchSortedSide],
+        side: SearchSortedSide,
     ) -> VortexResult<Vec<SearchResult>> {
         match_each_unsigned_integer_ptype!(array.ptype(), |$P| {
             let searcher = BitPackedSearch::<'_, $P>::new(array);
 
             values
                 .iter()
-                .zip(sides.iter().copied())
-                .map(|(value, side)| {
+                .map(|value| {
                     // Unwrap to native value
                     let unwrapped_value: $P = value.cast(array.dtype())?.try_into()?;
-
                     Ok(searcher.search_sorted(&unwrapped_value, side))
                 })
                 .try_collect()
@@ -75,16 +73,14 @@ impl SearchSortedFn<BitPackedArray> for BitPackedEncoding {
         &self,
         array: &BitPackedArray,
         values: &[usize],
-        sides: &[SearchSortedSide],
+        side: SearchSortedSide,
     ) -> VortexResult<Vec<SearchResult>> {
         match_each_unsigned_integer_ptype!(array.ptype(), |$P| {
             let searcher = BitPackedSearch::<'_, $P>::new(array);
 
             values
                 .iter()
-                .copied()
-                .zip(sides.iter().copied())
-                .map(|(value, side)| {
+                .map(|&value| {
                     // NOTE: truncating cast
                     let cast_value: $P = value as $P;
                     Ok(searcher.search_sorted(&cast_value, side))
@@ -299,11 +295,7 @@ mod test {
         let results = search_sorted_many(
             bitpacked.as_ref(),
             &[3u64, 2u64, 1u64],
-            &[
-                SearchSortedSide::Left,
-                SearchSortedSide::Left,
-                SearchSortedSide::Left,
-            ],
+            SearchSortedSide::Left,
         )
         .unwrap();
 

encodings/runend/src/array.rs

@@ -123,12 +123,7 @@ impl RunEndArray {
     ///
     /// See: [find_physical_index][Self::find_physical_index].
     pub fn find_physical_indices(&self, indices: &[usize]) -> VortexResult<Vec<usize>> {
-        search_sorted_usize_many(
-            &self.ends(),
-            indices,
-            &vec![SearchSortedSide::Right; indices.len()],
-        )
-        .map(|results| {
+        search_sorted_usize_many(&self.ends(), indices, SearchSortedSide::Right).map(|results| {
             results
                 .iter()
                 .map(|result| result.to_ends_index(self.ends().len()))

vortex-array/src/compute/search_sorted.rs

@@ -121,26 +121,23 @@ pub trait SearchSortedFn<Array> {
         &self,
         array: &Array,
         values: &[Scalar],
-        sides: &[SearchSortedSide],
+        side: SearchSortedSide,
     ) -> VortexResult<Vec<SearchResult>> {
         values
             .iter()
-            .zip(sides.iter())
-            .map(|(value, side)| self.search_sorted(array, value, *side))
+            .map(|value| self.search_sorted(array, value, side))
             .try_collect()
     }
 
     fn search_sorted_usize_many(
         &self,
         array: &Array,
         values: &[usize],
-        sides: &[SearchSortedSide],
+        side: SearchSortedSide,
     ) -> VortexResult<Vec<SearchResult>> {
         values
             .iter()
-            .copied()
-            .zip(sides.iter().copied())
-            .map(|(value, side)| self.search_sorted_usize(array, value, side))
+            .map(|&value| self.search_sorted_usize(array, value, side))
             .try_collect()
     }
 }
@@ -184,30 +181,30 @@ where
         &self,
         array: &ArrayData,
         values: &[Scalar],
-        sides: &[SearchSortedSide],
+        side: SearchSortedSide,
     ) -> VortexResult<Vec<SearchResult>> {
         let array_ref = <&E::Array>::try_from(array)?;
         let encoding = array
             .encoding()
             .as_any()
             .downcast_ref::<E>()
             .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
-        SearchSortedFn::search_sorted_many(encoding, array_ref, values, sides)
+        SearchSortedFn::search_sorted_many(encoding, array_ref, values, side)
     }
 
     fn search_sorted_usize_many(
         &self,
         array: &ArrayData,
         values: &[usize],
-        sides: &[SearchSortedSide],
+        side: SearchSortedSide,
     ) -> VortexResult<Vec<SearchResult>> {
         let array_ref = <&E::Array>::try_from(array)?;
         let encoding = array
             .encoding()
             .as_any()
             .downcast_ref::<E>()
             .ok_or_else(|| vortex_err!("Mismatched encoding"))?;
-        SearchSortedFn::search_sorted_usize_many(encoding, array_ref, values, sides)
+        SearchSortedFn::search_sorted_usize_many(encoding, array_ref, values, side)
     }
 }
 
@@ -261,41 +258,38 @@ pub fn search_sorted_usize(
 pub fn search_sorted_many<T: Into<Scalar> + Clone>(
     array: &ArrayData,
     targets: &[T],
-    sides: &[SearchSortedSide],
+    side: SearchSortedSide,
 ) -> VortexResult<Vec<SearchResult>> {
     if let Some(f) = array.encoding().search_sorted_fn() {
         let values: Vec<Scalar> = targets
             .iter()
             .map(|t| t.clone().into().cast(array.dtype()))
             .try_collect()?;
 
-        return f.search_sorted_many(array, &values, sides);
+        return f.search_sorted_many(array, &values, side);
     }
 
     // Call in loop and collect
     targets
         .iter()
-        .zip(sides.iter().copied())
-        .map(|(target, side)| search_sorted(array, target.clone(), side))
+        .map(|target| search_sorted(array, target.clone(), side))
         .try_collect()
 }
 
 // Native functions for each of the values, cast up to u64 or down to something lower.
 pub fn search_sorted_usize_many(
     array: &ArrayData,
     targets: &[usize],
-    sides: &[SearchSortedSide],
+    side: SearchSortedSide,
 ) -> VortexResult<Vec<SearchResult>> {
     if let Some(f) = array.encoding().search_sorted_fn() {
-        return f.search_sorted_usize_many(array, targets, sides);
+        return f.search_sorted_usize_many(array, targets, side);
     }
 
     // Call in loop and collect
     targets
         .iter()
-        .copied()
-        .zip(sides.iter().copied())
-        .map(|(target, side)| search_sorted_usize(array, target, side))
+        .map(|&target| search_sorted_usize(array, target, side))
         .try_collect()
 }