Flatten unary compute mod (#1489)

We probably ought to group functions at some point, e.g. by logical, arithmetic, select, etc. But grouping by arity is weird...
spiraldb · Nov 27, 2024 · 7507c99 · 7507c99
1 parent 630835b
commit 7507c99
Show file tree

Hide file tree

Showing 85 changed files with 143 additions and 185 deletions.
diff --git a/bench-vortex/benches/compressor_throughput.rs b/bench-vortex/benches/compressor_throughput.rs
@@ -6,8 +6,7 @@ use rand::seq::SliceRandom as _;
 use rand::{thread_rng, Rng, SeedableRng as _};
 use vortex::aliases::hash_set::HashSet;
 use vortex::array::{ConstantArray, PrimitiveArray, VarBinViewArray};
-use vortex::compute::unary::try_cast;
-use vortex::compute::{compare, Operator};
+use vortex::compute::{compare, try_cast, Operator};
 use vortex::dict::{dict_encode_varbinview, DictArray};
 use vortex::dtype::PType;
 use vortex::fsst::{fsst_compress, fsst_train_compressor};

diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs
@@ -107,7 +107,7 @@ fn decompress_primitive<T: NativePType + ALPFloat>(
 mod tests {
     use core::f64;
 
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::scalar_at;
 
     use super::*;
 

diff --git a/encodings/alp/src/alp/compute/mod.rs b/encodings/alp/src/alp/compute/mod.rs
@@ -1,9 +1,8 @@
 mod compare;
 
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
 use vortex_array::compute::{
-    filter, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, SliceFn, TakeFn,
-    TakeOptions,
+    filter, scalar_at, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, ScalarAtFn,
+    SliceFn, TakeFn, TakeOptions,
 };
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData};

diff --git a/encodings/alp/src/alp_rd/compute/mod.rs b/encodings/alp/src/alp_rd/compute/mod.rs
@@ -1,5 +1,4 @@
-use vortex_array::compute::unary::ScalarAtFn;
-use vortex_array::compute::{ComputeVTable, FilterFn, SliceFn, TakeFn};
+use vortex_array::compute::{ComputeVTable, FilterFn, ScalarAtFn, SliceFn, TakeFn};
 use vortex_array::ArrayData;
 
 use crate::ALPRDEncoding;

diff --git a/encodings/alp/src/alp_rd/compute/scalar_at.rs b/encodings/alp/src/alp_rd/compute/scalar_at.rs
@@ -1,4 +1,4 @@
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
+use vortex_array::compute::{scalar_at, ScalarAtFn};
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
 
@@ -36,7 +36,7 @@ impl ScalarAtFn<ALPRDArray> for ALPRDEncoding {
 mod test {
     use rstest::rstest;
     use vortex_array::array::PrimitiveArray;
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::scalar_at;
     use vortex_scalar::Scalar;
 
     use crate::{ALPRDFloat, RDEncoder};

diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs
@@ -1,6 +1,7 @@
 use num_traits::AsPrimitive;
-use vortex_array::compute::unary::{FillForwardFn, ScalarAtFn};
-use vortex_array::compute::{ComputeVTable, SliceFn, TakeFn, TakeOptions};
+use vortex_array::compute::{
+    ComputeVTable, FillForwardFn, ScalarAtFn, SliceFn, TakeFn, TakeOptions,
+};
 use vortex_array::validity::{ArrayValidity, Validity};
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayDType, ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant, ToArrayData};
@@ -142,8 +143,7 @@ impl FillForwardFn<ByteBoolArray> for ByteBoolEncoding {
 
 #[cfg(test)]
 mod tests {
-    use vortex_array::compute::unary::scalar_at;
-    use vortex_array::compute::{compare, slice, Operator};
+    use vortex_array::compute::{compare, scalar_at, slice, Operator};
 
     use super::*;
 

diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs
@@ -2,7 +2,7 @@ use std::fmt::{Debug, Display};
 
 use serde::{Deserialize, Serialize};
 use vortex_array::array::StructArray;
-use vortex_array::compute::unary::try_cast;
+use vortex_array::compute::try_cast;
 use vortex_array::encoding::ids;
 use vortex_array::stats::{Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{LogicalValidity, Validity, ValidityVTable};

diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs
@@ -1,5 +1,5 @@
 use vortex_array::array::{PrimitiveArray, TemporalArray};
-use vortex_array::compute::unary::try_cast;
+use vortex_array::compute::try_cast;
 use vortex_array::{ArrayDType as _, ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant};
 use vortex_datetime_dtype::TimeUnit;
 use vortex_dtype::{DType, PType};

diff --git a/encodings/datetime-parts/src/compute.rs b/encodings/datetime-parts/src/compute.rs
@@ -1,7 +1,8 @@
 use itertools::Itertools as _;
 use vortex_array::array::{PrimitiveArray, TemporalArray};
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
-use vortex_array::compute::{slice, take, ComputeVTable, SliceFn, TakeFn, TakeOptions};
+use vortex_array::compute::{
+    scalar_at, slice, take, ComputeVTable, ScalarAtFn, SliceFn, TakeFn, TakeOptions,
+};
 use vortex_array::validity::ArrayValidity;
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData, IntoArrayVariant};
 use vortex_datetime_dtype::{TemporalMetadata, TimeUnit};

diff --git a/encodings/dict/src/array.rs b/encodings/dict/src/array.rs
@@ -3,8 +3,7 @@ use std::fmt::{Debug, Display};
 use arrow_buffer::BooleanBuffer;
 use serde::{Deserialize, Serialize};
 use vortex_array::array::BoolArray;
-use vortex_array::compute::unary::scalar_at;
-use vortex_array::compute::{take, TakeOptions};
+use vortex_array::compute::{scalar_at, take, TakeOptions};
 use vortex_array::encoding::ids;
 use vortex_array::stats::StatsSet;
 use vortex_array::validity::{LogicalValidity, ValidityVTable};

diff --git a/encodings/dict/src/compress.rs b/encodings/dict/src/compress.rs
@@ -195,7 +195,7 @@ mod test {
 
     use vortex_array::accessor::ArrayAccessor;
     use vortex_array::array::{PrimitiveArray, VarBinArray};
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::scalar_at;
     use vortex_dtype::Nullability::Nullable;
     use vortex_dtype::{DType, PType};
     use vortex_scalar::Scalar;

diff --git a/encodings/dict/src/compute/mod.rs b/encodings/dict/src/compute/mod.rs
@@ -1,9 +1,8 @@
 mod compare;
 
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
 use vortex_array::compute::{
-    filter, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, SliceFn, TakeFn,
-    TakeOptions,
+    filter, scalar_at, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, ScalarAtFn,
+    SliceFn, TakeFn, TakeOptions,
 };
 use vortex_array::{ArrayData, IntoArrayData};
 use vortex_error::VortexResult;
@@ -74,8 +73,7 @@ impl SliceFn<DictArray> for DictEncoding {
 mod test {
     use vortex_array::accessor::ArrayAccessor;
     use vortex_array::array::{ConstantArray, PrimitiveArray, VarBinViewArray};
-    use vortex_array::compute::unary::scalar_at;
-    use vortex_array::compute::{compare, slice, Operator};
+    use vortex_array::compute::{compare, scalar_at, slice, Operator};
     use vortex_array::{ArrayLen, IntoArrayData, IntoArrayVariant, ToArrayData};
     use vortex_dtype::{DType, Nullability};
     use vortex_scalar::Scalar;

diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs
@@ -1,5 +1,4 @@
-use vortex_array::compute::unary::ScalarAtFn;
-use vortex_array::compute::{ComputeVTable, FilterFn, SearchSortedFn, SliceFn, TakeFn};
+use vortex_array::compute::{ComputeVTable, FilterFn, ScalarAtFn, SearchSortedFn, SliceFn, TakeFn};
 use vortex_array::ArrayData;
 
 use crate::BitPackedEncoding;

diff --git a/encodings/fastlanes/src/bitpacking/compute/scalar_at.rs b/encodings/fastlanes/src/bitpacking/compute/scalar_at.rs
@@ -1,4 +1,4 @@
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
+use vortex_array::compute::{scalar_at, ScalarAtFn};
 use vortex_array::ArrayDType;
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
@@ -21,7 +21,7 @@ impl ScalarAtFn<BitPackedArray> for BitPackedEncoding {
 #[cfg(test)]
 mod test {
     use vortex_array::array::{PrimitiveArray, SparseArray};
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::scalar_at;
     use vortex_array::validity::Validity;
     use vortex_array::IntoArrayData;
     use vortex_buffer::Buffer;

diff --git a/encodings/fastlanes/src/bitpacking/compute/slice.rs b/encodings/fastlanes/src/bitpacking/compute/slice.rs
@@ -47,8 +47,7 @@ impl SliceFn<BitPackedArray> for BitPackedEncoding {
 mod test {
     use itertools::Itertools;
     use vortex_array::array::{PrimitiveArray, SparseArray};
-    use vortex_array::compute::unary::scalar_at;
-    use vortex_array::compute::{slice, take, TakeOptions};
+    use vortex_array::compute::{scalar_at, slice, take, TakeOptions};
     use vortex_array::{ArrayLen, IntoArrayData};
 
     use crate::BitPackedArray;

diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs
@@ -241,8 +241,7 @@ mod test {
     use rand::distributions::Uniform;
     use rand::{thread_rng, Rng};
     use vortex_array::array::{PrimitiveArray, SparseArray};
-    use vortex_array::compute::unary::scalar_at;
-    use vortex_array::compute::{slice, take, TakeOptions};
+    use vortex_array::compute::{scalar_at, slice, take, TakeOptions};
     use vortex_array::{IntoArrayData, IntoArrayVariant};
 
     use crate::BitPackedArray;

diff --git a/encodings/fastlanes/src/delta/compress.rs b/encodings/fastlanes/src/delta/compress.rs
@@ -2,8 +2,7 @@ use arrayref::{array_mut_ref, array_ref};
 use fastlanes::{Delta, Transpose};
 use num_traits::{WrappingAdd, WrappingSub};
 use vortex_array::array::PrimitiveArray;
-use vortex_array::compute::slice;
-use vortex_array::compute::unary::fill_forward;
+use vortex_array::compute::{fill_forward, slice};
 use vortex_array::validity::Validity;
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayLen, IntoArrayVariant};

diff --git a/encodings/fastlanes/src/delta/compute.rs b/encodings/fastlanes/src/delta/compute.rs
@@ -1,7 +1,6 @@
 use std::cmp::min;
 
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
-use vortex_array::compute::{slice, ComputeVTable, SliceFn};
+use vortex_array::compute::{scalar_at, slice, ComputeVTable, ScalarAtFn, SliceFn};
 use vortex_array::{ArrayData, IntoArrayData, IntoArrayVariant};
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
@@ -68,8 +67,7 @@ impl SliceFn<DeltaArray> for DeltaEncoding {
 
 #[cfg(test)]
 mod test {
-    use vortex_array::compute::slice;
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::{scalar_at, slice};
     use vortex_array::IntoArrayVariant;
     use vortex_error::VortexError;
 

diff --git a/encodings/fastlanes/src/for/compress.rs b/encodings/fastlanes/src/for/compress.rs
@@ -144,7 +144,7 @@ fn decompress_primitive<T: NativePType + WrappingAdd + PrimInt>(
 
 #[cfg(test)]
 mod test {
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::scalar_at;
     use vortex_array::IntoArrayVariant;
     use vortex_dtype::Nullability;
 

diff --git a/encodings/fastlanes/src/for/compute.rs b/encodings/fastlanes/src/for/compute.rs
@@ -1,10 +1,9 @@
 use std::ops::AddAssign;
 
 use num_traits::{CheckedShl, CheckedShr, WrappingAdd, WrappingSub};
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
 use vortex_array::compute::{
-    filter, search_sorted, slice, take, ComputeVTable, FilterFn, FilterMask, SearchResult,
-    SearchSortedFn, SearchSortedSide, SliceFn, TakeFn, TakeOptions,
+    filter, scalar_at, search_sorted, slice, take, ComputeVTable, FilterFn, FilterMask, ScalarAtFn,
+    SearchResult, SearchSortedFn, SearchSortedSide, SliceFn, TakeFn, TakeOptions,
 };
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData};
@@ -178,8 +177,7 @@ where
 #[cfg(test)]
 mod test {
     use vortex_array::array::PrimitiveArray;
-    use vortex_array::compute::unary::scalar_at;
-    use vortex_array::compute::{search_sorted, SearchResult, SearchSortedSide};
+    use vortex_array::compute::{scalar_at, search_sorted, SearchResult, SearchSortedSide};
     use vortex_array::IntoArrayData;
 
     use crate::{for_compress, FoRArray};

diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs
@@ -82,8 +82,7 @@ fn compare_fsst_constant(
 #[cfg(test)]
 mod tests {
     use vortex_array::array::{ConstantArray, VarBinArray};
-    use vortex_array::compute::unary::scalar_at;
-    use vortex_array::compute::{compare, Operator};
+    use vortex_array::compute::{compare, scalar_at, Operator};
     use vortex_array::{ArrayLen, IntoArrayData, IntoArrayVariant};
     use vortex_dtype::{DType, Nullability};
     use vortex_scalar::Scalar;

diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs
@@ -1,10 +1,9 @@
 mod compare;
 
 use vortex_array::array::varbin_scalar;
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
 use vortex_array::compute::{
-    filter, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, SliceFn, TakeFn,
-    TakeOptions,
+    filter, scalar_at, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, ScalarAtFn,
+    SliceFn, TakeFn, TakeOptions,
 };
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData};
 use vortex_buffer::Buffer;

diff --git a/encodings/fsst/tests/fsst_tests.rs b/encodings/fsst/tests/fsst_tests.rs
@@ -2,8 +2,7 @@
 
 use vortex_array::array::builder::VarBinBuilder;
 use vortex_array::array::PrimitiveArray;
-use vortex_array::compute::unary::scalar_at;
-use vortex_array::compute::{filter, slice, take, FilterMask, TakeOptions};
+use vortex_array::compute::{filter, scalar_at, slice, take, FilterMask, TakeOptions};
 use vortex_array::validity::Validity;
 use vortex_array::{ArrayData, ArrayDef, IntoArrayData, IntoCanonical};
 use vortex_dtype::{DType, Nullability};

diff --git a/encodings/roaring/src/boolean/compute.rs b/encodings/roaring/src/boolean/compute.rs
@@ -1,6 +1,5 @@
 use croaring::Bitmap;
-use vortex_array::compute::unary::ScalarAtFn;
-use vortex_array::compute::{ComputeVTable, SliceFn};
+use vortex_array::compute::{ComputeVTable, ScalarAtFn, SliceFn};
 use vortex_array::{ArrayData, IntoArrayData};
 use vortex_error::VortexResult;
 use vortex_scalar::Scalar;
@@ -43,8 +42,7 @@ impl SliceFn<RoaringBoolArray> for RoaringBoolEncoding {
 #[cfg(test)]
 mod tests {
     use vortex_array::array::BoolArray;
-    use vortex_array::compute::slice;
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::{scalar_at, slice};
     use vortex_array::{IntoArrayData, IntoArrayVariant};
     use vortex_scalar::Scalar;
 

diff --git a/encodings/roaring/src/integer/compute.rs b/encodings/roaring/src/integer/compute.rs
@@ -1,6 +1,5 @@
 use croaring::Bitmap;
-use vortex_array::compute::unary::ScalarAtFn;
-use vortex_array::compute::{ComputeVTable, SliceFn};
+use vortex_array::compute::{ComputeVTable, ScalarAtFn, SliceFn};
 use vortex_array::{ArrayData, ArrayLen, IntoArrayData};
 use vortex_dtype::PType;
 use vortex_error::{vortex_err, VortexResult};
@@ -57,8 +56,7 @@ impl SliceFn<RoaringIntArray> for RoaringIntEncoding {
 #[cfg(test)]
 mod tests {
     use vortex_array::array::PrimitiveArray;
-    use vortex_array::compute::slice;
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::{scalar_at, slice};
 
     use super::*;
 

diff --git a/encodings/roaring/src/integer/mod.rs b/encodings/roaring/src/integer/mod.rs
@@ -5,7 +5,7 @@ pub use compress::*;
 use croaring::{Bitmap, Portable};
 use serde::{Deserialize, Serialize};
 use vortex_array::array::PrimitiveArray;
-use vortex_array::compute::unary::try_cast;
+use vortex_array::compute::try_cast;
 use vortex_array::encoding::ids;
 use vortex_array::stats::{ArrayStatistics, Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{LogicalValidity, Validity, ValidityVTable};

diff --git a/encodings/runend-bool/src/array.rs b/encodings/runend-bool/src/array.rs
@@ -2,8 +2,7 @@ use std::fmt::{Debug, Display};
 
 use serde::{Deserialize, Serialize};
 use vortex_array::array::{BoolArray, PrimitiveArray};
-use vortex_array::compute::unary::scalar_at;
-use vortex_array::compute::{search_sorted, SearchSortedSide};
+use vortex_array::compute::{scalar_at, search_sorted, SearchSortedSide};
 use vortex_array::encoding::ids;
 use vortex_array::stats::{ArrayStatistics, Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{LogicalValidity, Validity, ValidityMetadata, ValidityVTable};
@@ -266,8 +265,7 @@ mod test {
     use itertools::Itertools as _;
     use rstest::rstest;
     use vortex_array::array::{BoolArray, PrimitiveArray};
-    use vortex_array::compute::unary::scalar_at;
-    use vortex_array::compute::{slice, take, TakeOptions};
+    use vortex_array::compute::{scalar_at, slice, take, TakeOptions};
     use vortex_array::stats::ArrayStatistics;
     use vortex_array::validity::Validity;
     use vortex_array::{

diff --git a/encodings/runend-bool/src/compute.rs b/encodings/runend-bool/src/compute.rs
@@ -1,6 +1,5 @@
 use vortex_array::array::BoolArray;
-use vortex_array::compute::unary::ScalarAtFn;
-use vortex_array::compute::{slice, ComputeVTable, SliceFn, TakeFn, TakeOptions};
+use vortex_array::compute::{slice, ComputeVTable, ScalarAtFn, SliceFn, TakeFn, TakeOptions};
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant, ToArrayData};
 use vortex_dtype::match_each_integer_ptype;

diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs
@@ -2,8 +2,7 @@ use std::fmt::{Debug, Display};
 
 use serde::{Deserialize, Serialize};
 use vortex_array::array::PrimitiveArray;
-use vortex_array::compute::unary::scalar_at;
-use vortex_array::compute::{search_sorted, search_sorted_usize_many, SearchSortedSide};
+use vortex_array::compute::{scalar_at, search_sorted, search_sorted_usize_many, SearchSortedSide};
 use vortex_array::encoding::ids;
 use vortex_array::stats::{ArrayStatistics, Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{
@@ -276,7 +275,7 @@ impl StatisticsVTable<RunEndArray> for RunEndEncoding {
 
 #[cfg(test)]
 mod tests {
-    use vortex_array::compute::unary::scalar_at;
+    use vortex_array::compute::scalar_at;
     use vortex_array::validity::Validity;
     use vortex_array::{ArrayDType, ArrayLen, IntoArrayData};
     use vortex_dtype::{DType, Nullability, PType};

diff --git a/encodings/runend/src/compute/mod.rs b/encodings/runend/src/compute/mod.rs
@@ -5,10 +5,9 @@ use std::ops::AddAssign;
 
 use num_traits::AsPrimitive;
 use vortex_array::array::{BooleanBuffer, ConstantArray, PrimitiveArray, SparseArray};
-use vortex_array::compute::unary::{scalar_at, ScalarAtFn};
 use vortex_array::compute::{
-    filter, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, SliceFn, TakeFn,
-    TakeOptions,
+    filter, scalar_at, slice, take, CompareFn, ComputeVTable, FilterFn, FilterMask, ScalarAtFn,
+    SliceFn, TakeFn, TakeOptions,
 };
 use vortex_array::validity::Validity;
 use vortex_array::variants::PrimitiveArrayTrait;
@@ -178,8 +177,9 @@ fn filter_run_ends<R: NativePType + AddAssign + From<bool> + AsPrimitive<u64>>(
 #[cfg(test)]
 mod test {
     use vortex_array::array::{BoolArray, PrimitiveArray};
-    use vortex_array::compute::unary::{scalar_at, try_cast};
-    use vortex_array::compute::{filter, slice, take, FilterMask, TakeOptions};
+    use vortex_array::compute::{
+        filter, scalar_at, slice, take, try_cast, FilterMask, TakeOptions,
+    };
     use vortex_array::validity::{ArrayValidity, Validity};
     use vortex_array::{ArrayDType, ArrayLen, IntoArrayData, IntoArrayVariant, ToArrayData};
     use vortex_dtype::{DType, Nullability, PType};