Add arithmetic buffer compute

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/array/operator.rs

@@ -3,7 +3,8 @@
 
 use std::sync::Arc;
 
-use vortex_error::VortexResult;
+use vortex_dtype::DType;
+use vortex_error::{VortexResult, vortex_bail};
 use vortex_vector::Vector;
 
 use crate::execution::{BatchKernelRef, BindCtx};
@@ -33,6 +34,21 @@ pub trait ArrayOperator: 'static + Send + Sync {
 
 impl ArrayOperator for Arc<dyn Array> {
     fn execute_with_selection(&self, selection: Option<&ArrayRef>) -> VortexResult<Vector> {
+        if let Some(selection) = selection.as_ref() {
+            if !matches!(selection.dtype(), DType::Bool(_)) {
+                vortex_bail!(
+                    "Selection array must be of boolean type, got {}",
+                    selection.dtype()
+                );
+            }
+            if selection.len() != self.len() {
+                vortex_bail!(
+                    "Selection array length {} does not match array length {}",
+                    selection.len(),
+                    self.len()
+                );
+            }
+        }
         self.as_ref().execute_with_selection(selection)
     }
 

vortex-array/src/compute/arrays/arithmetic.rs

@@ -4,23 +4,23 @@
 use std::hash::{Hash, Hasher};
 use std::sync::LazyLock;
 
-use enum_map::{enum_map, Enum, EnumMap};
+use enum_map::{Enum, EnumMap, enum_map};
 use vortex_buffer::ByteBuffer;
 use vortex_compute::arithmetic::{Add, Checked, CheckedOperator, Div, Mul, Sub};
-use vortex_dtype::{match_each_native_ptype, DType, NativePType, PTypeDowncastExt};
-use vortex_error::{vortex_err, VortexExpect, VortexResult};
+use vortex_dtype::{DType, NativePType, PTypeDowncastExt, match_each_native_ptype};
+use vortex_error::{VortexExpect, VortexResult, vortex_err};
 use vortex_scalar::{PValue, Scalar};
 
 use crate::arrays::ConstantArray;
-use crate::execution::{kernel, BatchKernelRef, BindCtx};
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
 use crate::serde::ArrayChildren;
 use crate::stats::{ArrayStats, StatsSetRef};
 use crate::vtable::{
     ArrayVTable, NotSupported, OperatorVTable, SerdeVTable, VTable, VisitorVTable,
 };
 use crate::{
-    vtable, Array, ArrayBufferVisitor, ArrayChildVisitor, ArrayEq, ArrayHash,
-    ArrayRef, DeserializeMetadata, EmptyMetadata, EncodingId, EncodingRef, IntoArray, Precision,
+    Array, ArrayBufferVisitor, ArrayChildVisitor, ArrayEq, ArrayHash, ArrayRef,
+    DeserializeMetadata, EmptyMetadata, EncodingId, EncodingRef, IntoArray, Precision, vtable,
 };
 
 /// The set of operators supported by an arithmetic array.
@@ -223,11 +223,9 @@ impl OperatorVTable<ArithmeticVTable> for ArithmeticVTable {
         if let Some(rhs) = array.rhs.as_constant() {
             if rhs.is_null() {
                 // If the RHS is null, the result is always null.
-                return Ok(
-                    ConstantArray::new(Scalar::null(array.dtype().clone()), array.len())
-                        .into_array()
-                        .bind(selection, ctx)?,
-                );
+                return ConstantArray::new(Scalar::null(array.dtype().clone()), array.len())
+                    .into_array()
+                    .bind(selection, ctx);
             }
 
             let lhs = ctx.bind(&array.lhs, selection)?;
@@ -289,7 +287,7 @@ where
 
 #[cfg(test)]
 mod tests {
-    use vortex_buffer::buffer;
+    use vortex_buffer::{bitbuffer, buffer};
     use vortex_dtype::PTypeDowncastExt;
 
     use crate::arrays::PrimitiveArray;
@@ -365,7 +363,7 @@ mod tests {
         let lhs = PrimitiveArray::from_iter([1u32, 2, 3]).into_array();
         let rhs = PrimitiveArray::from_iter([10u32, 20, 30]).into_array();
 
-        let selection = PrimitiveArray::from_iter([0u64, 2]).into_array();
+        let selection = bitbuffer![1 0 1].into_array();
 
         let result = add(lhs, rhs)
             .execute_with_selection(Some(&selection))

vortex-compute/src/arithmetic/buffer.rs

@@ -145,9 +145,10 @@ where
 
 #[cfg(test)]
 mod tests {
-    use crate::arithmetic::{CheckedAdd, CheckedDiv, CheckedMul, CheckedSub};
     use vortex_buffer::buffer;
 
+    use crate::arithmetic::{CheckedAdd, CheckedDiv, CheckedMul, CheckedSub};
+
     #[test]
     fn test_add_buffers() {
         let left = buffer![1u32, 2, 3, 4];

vortex-compute/src/arithmetic/pvector.rs

@@ -1,12 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use crate::arithmetic::{Checked, CheckedOperator};
 use std::ops::BitAnd;
+
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_dtype::NativePType;
 use vortex_vector::{PVector, PVectorMut, VectorMutOps, VectorOps};
 
+use crate::arithmetic::{Checked, CheckedOperator};
+
 /// Implementation that attempts to downcast to a mutable vector and operates in-place.
 impl<Op, T> Checked<Op, &PVector<T>> for PVector<T>
 where
@@ -121,11 +123,11 @@ where
 
 #[cfg(test)]
 mod tests {
-    use super::*;
-    use crate::arithmetic::{CheckedAdd, CheckedDiv, CheckedMul, CheckedSub};
     use vortex_buffer::buffer;
     use vortex_mask::Mask;
-    use vortex_vector::PVector;
+    use vortex_vector::{PVector, VectorOps};
+
+    use crate::arithmetic::{CheckedAdd, CheckedDiv, CheckedMul, CheckedSub};
 
     #[test]
     fn test_add_pvectors() {

vortex-compute/src/logical/and.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::ops::BitAnd;
+
 use vortex_mask::Mask;
 use vortex_vector::{BoolVector, VectorOps};
 

vortex-compute/src/logical/and_not.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::ops::BitAnd;
+
 use vortex_vector::{BoolVector, VectorOps};
 
 use crate::logical::LogicalAndNot;

vortex-compute/src/logical/or.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::ops::{BitAnd, BitOr};
+
 use vortex_mask::Mask;
 use vortex_vector::{BoolVector, VectorOps};
 

vortex-dtype/src/ptype.rs

@@ -13,11 +13,11 @@ use num_traits::bounds::UpperBounded;
 use num_traits::{
     AsPrimitive, Bounded, FromPrimitive, Num, NumCast, PrimInt, ToPrimitive, Unsigned,
 };
-use vortex_error::{vortex_err, VortexError, VortexResult};
+use vortex_error::{VortexError, VortexResult, vortex_err};
 
+use crate::DType;
 use crate::half::f16;
 use crate::nullability::Nullability::NonNullable;
-use crate::DType;
 
 /// Physical type enum, represents the in-memory physical layout but might represent a different logical type.
 #[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Hash, prost::Enumeration)]

vortex-vector/src/primitive/generic.rs

@@ -5,7 +5,7 @@
 
 use vortex_buffer::Buffer;
 use vortex_dtype::NativePType;
-use vortex_error::{vortex_ensure, VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
 use vortex_mask::Mask;
 
 use crate::{PVectorMut, VectorOps};
@@ -111,18 +111,6 @@ impl<T: NativePType> AsRef<[T]> for PVector<T> {
     }
 }
 
-impl<T: NativePType> PVector<T> {
-    /// Returns a reference to the internal elements buffer.
-    pub fn elements(&self) -> &Buffer<T> {
-        &self.elements
-    }
-
-    /// Returns a reference to the internal validity mask.
-    pub fn validity(&self) -> &Mask {
-        &self.validity
-    }
-}
-
 impl<T: NativePType> VectorOps for PVector<T> {
     type Mutable = PVectorMut<T>;
 

vortex-vector/src/primitive/generic_mut.rs

@@ -5,10 +5,10 @@
 
 use vortex_buffer::BufferMut;
 use vortex_dtype::NativePType;
-use vortex_error::{vortex_ensure, VortexExpect, VortexResult};
+use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
 use vortex_mask::MaskMut;
 
-use crate::{PVector, VectorMutOps};
+use crate::{PVector, VectorMutOps, VectorOps};
 
 /// A mutable vector of generic primitive values.
 ///
@@ -166,14 +166,6 @@ impl<T> PVectorMut<T> {
     }
 }
 
-impl<T: NativePType> PVectorMut<T> {
-    /// Decomposes the mutable vector into its constituent parts: the elements buffer and the
-    /// validity mask.
-    pub fn into_parts(self) -> (BufferMut<T>, MaskMut) {
-        (self.elements, self.validity)
-    }
-}
-
 impl<T: NativePType> VectorMutOps for PVectorMut<T> {
     type Immutable = PVector<T>;