chore: add an unsafe take_unchecked to TakeFn for bounds check ellisi…

…on (#1611)
spiraldb · Dec 9, 2024 · 173d499 · 173d499
1 parent 69ef100
commit 173d499
Show file tree

Hide file tree

Showing 39 changed files with 298 additions and 519 deletions.
diff --git a/encodings/alp/src/alp/compute/mod.rs b/encodings/alp/src/alp/compute/mod.rs
@@ -1,6 +1,6 @@
 use vortex_array::compute::{
     filter, scalar_at, slice, take, ComputeVTable, FilterFn, FilterMask, ScalarAtFn, SliceFn,
-    TakeFn, TakeOptions,
+    TakeFn,
 };
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData};
@@ -48,15 +48,10 @@ impl ScalarAtFn<ALPArray> for ALPEncoding {
 }
 
 impl TakeFn<ALPArray> for ALPEncoding {
-    fn take(
-        &self,
-        array: &ALPArray,
-        indices: &ArrayData,
-        options: TakeOptions,
-    ) -> VortexResult<ArrayData> {
+    fn take(&self, array: &ALPArray, indices: &ArrayData) -> VortexResult<ArrayData> {
         // TODO(ngates): wrap up indices in an array that caches decompression?
         Ok(ALPArray::try_new(
-            take(array.encoded(), indices, options)?,
+            take(array.encoded(), indices)?,
             array.exponents(),
             array
                 .patches()

diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs
@@ -1,26 +1,21 @@
-use vortex_array::compute::{take, TakeFn, TakeOptions};
+use vortex_array::compute::{take, TakeFn};
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData};
 use vortex_error::VortexResult;
 
 use crate::{ALPRDArray, ALPRDEncoding};
 
 impl TakeFn<ALPRDArray> for ALPRDEncoding {
-    fn take(
-        &self,
-        array: &ALPRDArray,
-        indices: &ArrayData,
-        options: TakeOptions,
-    ) -> VortexResult<ArrayData> {
+    fn take(&self, array: &ALPRDArray, indices: &ArrayData) -> VortexResult<ArrayData> {
         let left_parts_exceptions = array
             .left_parts_exceptions()
-            .map(|array| take(&array, indices, options))
+            .map(|array| take(&array, indices))
             .transpose()?;
 
         Ok(ALPRDArray::try_new(
             array.dtype().clone(),
-            take(array.left_parts(), indices, options)?,
+            take(array.left_parts(), indices)?,
             array.left_parts_dict(),
-            take(array.right_parts(), indices, options)?,
+            take(array.right_parts(), indices)?,
             array.right_bit_width(),
             left_parts_exceptions,
         )?
@@ -32,7 +27,7 @@ impl TakeFn<ALPRDArray> for ALPRDEncoding {
 mod test {
     use rstest::rstest;
     use vortex_array::array::PrimitiveArray;
-    use vortex_array::compute::{take, TakeOptions};
+    use vortex_array::compute::take;
     use vortex_array::IntoArrayVariant;
 
     use crate::{ALPRDFloat, RDEncoder};
@@ -46,14 +41,10 @@ mod test {
 
         assert!(encoded.left_parts_exceptions().is_some());
 
-        let taken = take(
-            encoded.as_ref(),
-            PrimitiveArray::from(vec![0, 2]).as_ref(),
-            TakeOptions::default(),
-        )
-        .unwrap()
-        .into_primitive()
-        .unwrap();
+        let taken = take(encoded.as_ref(), PrimitiveArray::from(vec![0, 2]).as_ref())
+            .unwrap()
+            .into_primitive()
+            .unwrap();
 
         assert_eq!(taken.maybe_null_slice::<T>(), &[a, outlier]);
     }

diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs
@@ -1,7 +1,5 @@
 use num_traits::AsPrimitive;
-use vortex_array::compute::{
-    ComputeVTable, FillForwardFn, ScalarAtFn, SliceFn, TakeFn, TakeOptions,
-};
+use vortex_array::compute::{ComputeVTable, FillForwardFn, ScalarAtFn, SliceFn, TakeFn};
 use vortex_array::validity::{ArrayValidity, Validity};
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayDType, ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant, ToArrayData};
@@ -49,12 +47,7 @@ impl SliceFn<ByteBoolArray> for ByteBoolEncoding {
 }
 
 impl TakeFn<ByteBoolArray> for ByteBoolEncoding {
-    fn take(
-        &self,
-        array: &ByteBoolArray,
-        indices: &ArrayData,
-        _options: TakeOptions,
-    ) -> VortexResult<ArrayData> {
+    fn take(&self, array: &ByteBoolArray, indices: &ArrayData) -> VortexResult<ArrayData> {
         let validity = array.validity();
         let indices = indices.clone().into_primitive()?;
         let bools = array.maybe_null_slice();

diff --git a/encodings/datetime-parts/src/compute/take.rs b/encodings/datetime-parts/src/compute/take.rs
@@ -1,21 +1,16 @@
-use vortex_array::compute::{take, TakeFn, TakeOptions};
+use vortex_array::compute::{take, TakeFn};
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData};
 use vortex_error::VortexResult;
 
 use crate::{DateTimePartsArray, DateTimePartsEncoding};
 
 impl TakeFn<DateTimePartsArray> for DateTimePartsEncoding {
-    fn take(
-        &self,
-        array: &DateTimePartsArray,
-        indices: &ArrayData,
-        options: TakeOptions,
-    ) -> VortexResult<ArrayData> {
+    fn take(&self, array: &DateTimePartsArray, indices: &ArrayData) -> VortexResult<ArrayData> {
         Ok(DateTimePartsArray::try_new(
             array.dtype().clone(),
-            take(array.days(), indices, options)?,
-            take(array.seconds(), indices, options)?,
-            take(array.subsecond(), indices, options)?,
+            take(array.days(), indices)?,
+            take(array.seconds(), indices)?,
+            take(array.subsecond(), indices)?,
         )?
         .into_array())
     }

diff --git a/encodings/dict/src/array.rs b/encodings/dict/src/array.rs
@@ -3,7 +3,7 @@ use std::fmt::{Debug, Display};
 use arrow_buffer::BooleanBuffer;
 use serde::{Deserialize, Serialize};
 use vortex_array::array::BoolArray;
-use vortex_array::compute::{scalar_at, take, TakeOptions};
+use vortex_array::compute::{scalar_at, take};
 use vortex_array::encoding::ids;
 use vortex_array::stats::StatsSet;
 use vortex_array::validity::{ArrayValidity, LogicalValidity, ValidityVTable};
@@ -74,10 +74,10 @@ impl IntoCanonical for DictArray {
             // copies of the view pointers.
             DType::Utf8(_) | DType::Binary(_) => {
                 let canonical_values: ArrayData = self.values().into_canonical()?.into();
-                take(canonical_values, self.codes(), TakeOptions::default())?.into_canonical()
+                take(canonical_values, self.codes())?.into_canonical()
             }
             // Non-string case: take and then canonicalize
-            _ => take(self.values(), self.codes(), TakeOptions::default())?.into_canonical(),
+            _ => take(self.values(), self.codes())?.into_canonical(),
         }
     }
 }

diff --git a/encodings/dict/src/compute/compare.rs b/encodings/dict/src/compute/compare.rs
@@ -1,5 +1,5 @@
 use vortex_array::array::ConstantArray;
-use vortex_array::compute::{compare, take, CompareFn, Operator, TakeOptions};
+use vortex_array::compute::{compare, take, CompareFn, Operator};
 use vortex_array::ArrayData;
 use vortex_error::VortexResult;
 
@@ -20,7 +20,7 @@ impl CompareFn<DictArray> for DictEncoding {
                 ConstantArray::new(const_scalar, lhs.values().len()),
                 operator,
             )?;
-            return take(compare_result, lhs.codes(), TakeOptions::default()).map(Some);
+            return take(compare_result, lhs.codes()).map(Some);
         }
 
         // It's a little more complex, but we could perform a comparison against the dictionary

diff --git a/encodings/dict/src/compute/mod.rs b/encodings/dict/src/compute/mod.rs
@@ -3,7 +3,7 @@ mod like;
 
 use vortex_array::compute::{
     filter, scalar_at, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, LikeFn,
-    ScalarAtFn, SliceFn, TakeFn, TakeOptions,
+    ScalarAtFn, SliceFn, TakeFn,
 };
 use vortex_array::{ArrayData, IntoArrayData};
 use vortex_error::VortexResult;
@@ -45,16 +45,11 @@ impl ScalarAtFn<DictArray> for DictEncoding {
 }
 
 impl TakeFn<DictArray> for DictEncoding {
-    fn take(
-        &self,
-        array: &DictArray,
-        indices: &ArrayData,
-        options: TakeOptions,
-    ) -> VortexResult<ArrayData> {
+    fn take(&self, array: &DictArray, indices: &ArrayData) -> VortexResult<ArrayData> {
         // Dict
         //   codes: 0 0 1
         //   dict: a b c d e f g h
-        let codes = take(array.codes(), indices, options)?;
+        let codes = take(array.codes(), indices)?;
         DictArray::try_new(codes, array.values()).map(|a| a.into_array())
     }
 }