Chore: reorganize compute function files in vortex-array (#4984)

Makes sure each compute function implementation has its own file.

Note that this is a purely cosmetic change.

Yes, this is verbose, but also it's better to be consistent (and it also
helps with discoverability, for example "which compute functions exist
for this array type?).

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

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

@@ -13,11 +13,13 @@ mod take;
 
 #[cfg(test)]
 mod test {
+    use rstest::rstest;
     use vortex_dtype::half::f16;
     use vortex_scalar::Scalar;
 
     use crate::IntoArray;
-    use crate::arrays::constant::ConstantArray;
+    use crate::arrays::ConstantArray;
+    use crate::compute::conformance::consistency::test_array_consistency;
     use crate::compute::conformance::filter::test_filter_conformance;
     use crate::compute::conformance::mask::test_mask_conformance;
 
@@ -40,15 +42,6 @@ mod test {
             &ConstantArray::new(Scalar::from(f16::from_f32(3.0f32)), 5).into_array(),
         );
     }
-}
-
-#[cfg(test)]
-mod tests {
-    use rstest::rstest;
-    use vortex_scalar::Scalar;
-
-    use crate::arrays::ConstantArray;
-    use crate::compute::conformance::consistency::test_array_consistency;
 
     #[rstest]
     // From test_all_consistency

vortex-array/src/arrays/extension/compute/cast.rs

@@ -4,7 +4,7 @@
 use vortex_dtype::DType;
 
 use crate::arrays::{ExtensionArray, ExtensionVTable};
-use crate::compute::{CastKernel, CastKernelAdapter, cast};
+use crate::compute::{self, CastKernel, CastKernelAdapter};
 use crate::{ArrayRef, IntoArray, register_kernel};
 
 impl CastKernel for ExtensionVTable {
@@ -21,7 +21,7 @@ impl CastKernel for ExtensionVTable {
             unreachable!("Already verified we have an extension dtype");
         };
 
-        let new_storage = match cast(array.storage(), ext_dtype.storage_dtype()) {
+        let new_storage = match compute::cast(array.storage(), ext_dtype.storage_dtype()) {
             Ok(arr) => arr,
             Err(e) => {
                 log::warn!("Failed to cast storage array: {e}");
@@ -49,6 +49,7 @@ mod tests {
     use super::*;
     use crate::IntoArray;
     use crate::arrays::PrimitiveArray;
+    use crate::compute::cast;
     use crate::compute::conformance::cast::test_cast_conformance;
 
     #[test]

vortex-array/src/arrays/extension/compute/compare.rs

@@ -4,7 +4,7 @@
 use vortex_error::VortexResult;
 
 use crate::arrays::{ConstantArray, ExtensionArray, ExtensionVTable};
-use crate::compute::{CompareKernel, CompareKernelAdapter, Operator, compare};
+use crate::compute::{self, CompareKernel, CompareKernelAdapter, Operator};
 use crate::{Array, ArrayRef, register_kernel};
 
 impl CompareKernel for ExtensionVTable {
@@ -17,7 +17,7 @@ impl CompareKernel for ExtensionVTable {
         // If the RHS is a constant, we can extract the storage scalar.
         if let Some(const_ext) = rhs.as_constant() {
             let storage_scalar = const_ext.as_extension().storage();
-            return compare(
+            return compute::compare(
                 lhs.storage(),
                 ConstantArray::new(storage_scalar, lhs.len()).as_ref(),
                 operator,
@@ -27,7 +27,7 @@ impl CompareKernel for ExtensionVTable {
 
         // If the RHS is an extension array matching ours, we can extract the storage.
         if let Some(rhs_ext) = rhs.as_opt::<ExtensionVTable>() {
-            return compare(lhs.storage(), rhs_ext.storage(), operator).map(Some);
+            return compute::compare(lhs.storage(), rhs_ext.storage(), operator).map(Some);
         }
 
         // Otherwise, we need the RHS to handle this comparison.

vortex-array/src/arrays/extension/compute/filter.rs

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use crate::arrays::{ExtensionArray, ExtensionVTable};
+use crate::compute::{self, FilterKernel, FilterKernelAdapter};
+use crate::{ArrayRef, IntoArray, register_kernel};
+
+impl FilterKernel for ExtensionVTable {
+    fn filter(&self, array: &ExtensionArray, mask: &Mask) -> VortexResult<ArrayRef> {
+        Ok(ExtensionArray::new(
+            array.ext_dtype().clone(),
+            compute::filter(array.storage(), mask)?,
+        )
+        .into_array())
+    }
+}
+
+register_kernel!(FilterKernelAdapter(ExtensionVTable).lift());

vortex-array/src/arrays/extension/compute/is_constant.rs

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::arrays::{ExtensionArray, ExtensionVTable};
+use crate::compute::{self, IsConstantKernel, IsConstantKernelAdapter, IsConstantOpts};
+use crate::register_kernel;
+
+impl IsConstantKernel for ExtensionVTable {
+    fn is_constant(
+        &self,
+        array: &ExtensionArray,
+        opts: &IsConstantOpts,
+    ) -> VortexResult<Option<bool>> {
+        compute::is_constant_opts(array.storage(), opts)
+    }
+}
+
+register_kernel!(IsConstantKernelAdapter(ExtensionVTable).lift());

vortex-array/src/arrays/extension/compute/is_sorted.rs

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::arrays::{ExtensionArray, ExtensionVTable};
+use crate::compute::{self, IsSortedKernel, IsSortedKernelAdapter};
+use crate::register_kernel;
+
+impl IsSortedKernel for ExtensionVTable {
+    fn is_sorted(&self, array: &ExtensionArray) -> VortexResult<Option<bool>> {
+        compute::is_sorted(array.storage())
+    }
+
+    fn is_strict_sorted(&self, array: &ExtensionArray) -> VortexResult<Option<bool>> {
+        compute::is_strict_sorted(array.storage())
+    }
+}
+
+register_kernel!(IsSortedKernelAdapter(ExtensionVTable).lift());

vortex-array/src/arrays/extension/compute/mask.rs

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use vortex_dtype::ExtDType;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use crate::arrays::{ExtensionArray, ExtensionVTable};
+use crate::compute::{self, MaskKernel, MaskKernelAdapter};
+use crate::{ArrayRef, IntoArray, register_kernel};
+
+impl MaskKernel for ExtensionVTable {
+    fn mask(&self, array: &ExtensionArray, mask_array: &Mask) -> VortexResult<ArrayRef> {
+        let masked_storage = compute::mask(array.storage(), mask_array)?;
+        if masked_storage.dtype().nullability() == array.ext_dtype().storage_dtype().nullability() {
+            Ok(ExtensionArray::new(array.ext_dtype().clone(), masked_storage).into_array())
+        } else {
+            // The storage dtype changed (i.e., became nullable due to masking)
+            let ext_dtype = Arc::new(ExtDType::new(
+                array.ext_dtype().id().clone(),
+                Arc::new(masked_storage.dtype().clone()),
+                array.ext_dtype().metadata().cloned(),
+            ));
+            Ok(ExtensionArray::new(ext_dtype, masked_storage).into_array())
+        }
+    }
+}
+
+register_kernel!(MaskKernelAdapter(ExtensionVTable).lift());

vortex-array/src/arrays/extension/compute/min_max.rs

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_scalar::Scalar;
+
+use crate::arrays::{ExtensionArray, ExtensionVTable};
+use crate::compute::{self, MinMaxKernel, MinMaxKernelAdapter, MinMaxResult};
+use crate::register_kernel;
+
+impl MinMaxKernel for ExtensionVTable {
+    fn min_max(&self, array: &ExtensionArray) -> VortexResult<Option<MinMaxResult>> {
+        Ok(
+            compute::min_max(array.storage())?.map(|MinMaxResult { min, max }| MinMaxResult {
+                min: Scalar::extension(array.ext_dtype().clone(), min),
+                max: Scalar::extension(array.ext_dtype().clone(), max),
+            }),
+        )
+    }
+}
+
+register_kernel!(MinMaxKernelAdapter(ExtensionVTable).lift());

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

@@ -3,117 +3,13 @@
 
 mod cast;
 mod compare;
-
-use std::sync::Arc;
-
-use vortex_dtype::ExtDType;
-use vortex_error::VortexResult;
-use vortex_mask::Mask;
-use vortex_scalar::Scalar;
-
-use crate::arrays::ExtensionVTable;
-use crate::arrays::extension::ExtensionArray;
-use crate::compute::{
-    FilterKernel, FilterKernelAdapter, IsConstantKernel, IsConstantKernelAdapter, IsConstantOpts,
-    IsSortedKernel, IsSortedKernelAdapter, MaskKernel, MaskKernelAdapter, MinMaxKernel,
-    MinMaxKernelAdapter, MinMaxResult, SumKernel, SumKernelAdapter, TakeKernel, TakeKernelAdapter,
-    filter, is_constant_opts, is_sorted, is_strict_sorted, mask, min_max, sum, take,
-};
-use crate::{Array, ArrayRef, IntoArray, register_kernel};
-
-impl FilterKernel for ExtensionVTable {
-    fn filter(&self, array: &ExtensionArray, mask: &Mask) -> VortexResult<ArrayRef> {
-        Ok(
-            ExtensionArray::new(array.ext_dtype().clone(), filter(array.storage(), mask)?)
-                .into_array(),
-        )
-    }
-}
-
-register_kernel!(FilterKernelAdapter(ExtensionVTable).lift());
-
-impl MaskKernel for ExtensionVTable {
-    fn mask(&self, array: &ExtensionArray, mask_array: &Mask) -> VortexResult<ArrayRef> {
-        let masked_storage = mask(array.storage(), mask_array)?;
-        if masked_storage.dtype().nullability() == array.ext_dtype().storage_dtype().nullability() {
-            Ok(ExtensionArray::new(array.ext_dtype().clone(), masked_storage).into_array())
-        } else {
-            // The storage dtype changed (i.e., became nullable due to masking)
-            let ext_dtype = Arc::new(ExtDType::new(
-                array.ext_dtype().id().clone(),
-                Arc::new(masked_storage.dtype().clone()),
-                array.ext_dtype().metadata().cloned(),
-            ));
-            Ok(ExtensionArray::new(ext_dtype, masked_storage).into_array())
-        }
-    }
-}
-
-register_kernel!(MaskKernelAdapter(ExtensionVTable).lift());
-
-impl SumKernel for ExtensionVTable {
-    fn sum(&self, array: &ExtensionArray) -> VortexResult<Scalar> {
-        sum(array.storage())
-    }
-}
-
-register_kernel!(SumKernelAdapter(ExtensionVTable).lift());
-
-impl TakeKernel for ExtensionVTable {
-    fn take(&self, array: &ExtensionArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
-        let taken_storage = take(array.storage(), indices)?;
-        if taken_storage.dtype().nullability() == array.ext_dtype().storage_dtype().nullability() {
-            Ok(ExtensionArray::new(array.ext_dtype().clone(), taken_storage).into_array())
-        } else {
-            // The storage dtype changed (i.e., became nullable due to nullable indices)
-            let ext_dtype = Arc::new(ExtDType::new(
-                array.ext_dtype().id().clone(),
-                Arc::new(taken_storage.dtype().clone()),
-                array.ext_dtype().metadata().cloned(),
-            ));
-            Ok(ExtensionArray::new(ext_dtype, taken_storage).into_array())
-        }
-    }
-}
-
-register_kernel!(TakeKernelAdapter(ExtensionVTable).lift());
-
-impl MinMaxKernel for ExtensionVTable {
-    fn min_max(&self, array: &ExtensionArray) -> VortexResult<Option<MinMaxResult>> {
-        Ok(
-            min_max(array.storage())?.map(|MinMaxResult { min, max }| MinMaxResult {
-                min: Scalar::extension(array.ext_dtype().clone(), min),
-                max: Scalar::extension(array.ext_dtype().clone(), max),
-            }),
-        )
-    }
-}
-
-register_kernel!(MinMaxKernelAdapter(ExtensionVTable).lift());
-
-impl IsConstantKernel for ExtensionVTable {
-    fn is_constant(
-        &self,
-        array: &ExtensionArray,
-        opts: &IsConstantOpts,
-    ) -> VortexResult<Option<bool>> {
-        is_constant_opts(array.storage(), opts)
-    }
-}
-
-register_kernel!(IsConstantKernelAdapter(ExtensionVTable).lift());
-
-impl IsSortedKernel for ExtensionVTable {
-    fn is_sorted(&self, array: &ExtensionArray) -> VortexResult<Option<bool>> {
-        is_sorted(array.storage())
-    }
-
-    fn is_strict_sorted(&self, array: &ExtensionArray) -> VortexResult<Option<bool>> {
-        is_strict_sorted(array.storage())
-    }
-}
-
-register_kernel!(IsSortedKernelAdapter(ExtensionVTable).lift());
+mod filter;
+mod is_constant;
+mod is_sorted;
+mod mask;
+mod min_max;
+mod sum;
+mod take;
 
 #[cfg(test)]
 mod test {

vortex-array/src/arrays/extension/compute/sum.rs

@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+use vortex_scalar::Scalar;
+
+use crate::arrays::{ExtensionArray, ExtensionVTable};
+use crate::compute::{self, SumKernel, SumKernelAdapter};
+use crate::register_kernel;
+
+impl SumKernel for ExtensionVTable {
+    fn sum(&self, array: &ExtensionArray) -> VortexResult<Scalar> {
+        compute::sum(array.storage())
+    }
+}
+
+register_kernel!(SumKernelAdapter(ExtensionVTable).lift());