Feature: add Decimal::bind (#5210)

Basically identical to `Primitive`

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

Author: connortsui20

vortex-array/src/arrays/decimal/vtable/mod.rs

@@ -8,6 +8,7 @@ use crate::{EncodingId, EncodingRef, vtable};
 mod array;
 mod canonical;
 mod operations;
+mod operator;
 mod serde;
 mod validity;
 mod visitor;

vortex-array/src/arrays/decimal/vtable/operator.rs

@@ -0,0 +1,66 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_compute::filter::Filter;
+use vortex_dtype::PrecisionScale;
+use vortex_error::VortexResult;
+use vortex_scalar::match_each_decimal_value_type;
+use vortex_vector::decimal::DVector;
+
+use crate::arrays::{DecimalArray, DecimalVTable, MaskedVTable};
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::vtable::{OperatorVTable, ValidityHelper};
+use crate::{ArrayRef, IntoArray};
+
+impl OperatorVTable<DecimalVTable> for DecimalVTable {
+    fn bind(
+        array: &DecimalArray,
+        selection: Option<&ArrayRef>,
+        ctx: &mut dyn BindCtx,
+    ) -> VortexResult<BatchKernelRef> {
+        let mask = ctx.bind_selection(array.len(), selection)?;
+        let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
+
+        match_each_decimal_value_type!(array.values_type(), |D| {
+            let elements = array.buffer::<D>();
+            let ps = PrecisionScale::<D>::try_from(&array.decimal_dtype())?;
+
+            Ok(kernel(move || {
+                let mask = mask.execute()?;
+                let validity = validity.execute()?;
+
+                // Note that validity already has the mask applied so we only need to apply it to
+                // the elements.
+                let elements = elements.filter(&mask);
+
+                Ok(DVector::<D>::try_new(ps, elements, validity)?.into())
+            }))
+        })
+    }
+
+    fn reduce_parent(
+        array: &DecimalArray,
+        parent: &ArrayRef,
+        _child_idx: usize,
+    ) -> VortexResult<Option<ArrayRef>> {
+        // Push-down masking of `validity` from the parent `MaskedArray`.
+        if let Some(masked) = parent.as_opt::<MaskedVTable>() {
+            let masked_array = match_each_decimal_value_type!(array.values_type(), |D| {
+                // SAFETY: Since we are only flipping some bits in the validity, all invariants that
+                // were upheld are still upheld.
+                unsafe {
+                    DecimalArray::new_unchecked(
+                        array.buffer::<D>(),
+                        array.decimal_dtype(),
+                        array.validity().clone().and(masked.validity().clone()),
+                    )
+                }
+                .into_array()
+            });
+
+            return Ok(Some(masked_array));
+        }
+
+        Ok(None)
+    }
+}

vortex-array/src/arrays/primitive/vtable/operator.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_buffer::Buffer;
 use vortex_compute::filter::Filter;
 use vortex_dtype::match_each_native_ptype;
 use vortex_error::VortexResult;
@@ -20,8 +21,8 @@ impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
         let mask = ctx.bind_selection(array.len(), selection)?;
         let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
 
-        match_each_native_ptype!(array.ptype(), |T| {
-            let elements = array.buffer::<T>();
+        match_each_native_ptype!(array.ptype(), |P| {
+            let elements = array.buffer::<P>();
             Ok(kernel(move || {
                 let mask = mask.execute()?;
                 let validity = validity.execute()?;
@@ -30,7 +31,7 @@ impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
                 // the elements.
                 let elements = elements.filter(&mask);
 
-                Ok(PVector::try_new(elements, validity)?.into())
+                Ok(PVector::<P>::try_new(elements, validity)?.into())
             }))
         })
     }
@@ -40,16 +41,21 @@ impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
         parent: &ArrayRef,
         _child_idx: usize,
     ) -> VortexResult<Option<ArrayRef>> {
-        // Push-down masking of validity from parent MaskedVTable.
+        // Push-down masking of `validity` from the parent `MaskedArray`.
         if let Some(masked) = parent.as_opt::<MaskedVTable>() {
-            return Ok(Some(
-                PrimitiveArray::from_byte_buffer(
-                    array.byte_buffer().clone(),
-                    array.ptype(),
-                    array.validity().clone().and(masked.validity().clone()),
-                )
-                .into_array(),
-            ));
+            let masked_array = match_each_native_ptype!(array.ptype(), |T| {
+                // SAFETY: Since we are only flipping some bits in the validity, all invariants that
+                // were upheld are still upheld.
+                unsafe {
+                    PrimitiveArray::new_unchecked(
+                        Buffer::<T>::from_byte_buffer(array.byte_buffer().clone()),
+                        array.validity().clone().and(masked.validity().clone()),
+                    )
+                }
+                .into_array()
+            });
+
+            return Ok(Some(masked_array));
         }
 
         Ok(None)

vortex-vector/src/decimal/vector.rs

@@ -3,7 +3,7 @@
 
 //! Definition and implementation of [`DecimalVector`].
 
-use vortex_dtype::{DecimalTypeDowncast, DecimalTypeUpcast, NativeDecimalType, i256};
+use vortex_dtype::{DecimalType, DecimalTypeDowncast, DecimalTypeUpcast, NativeDecimalType, i256};
 use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 
@@ -27,6 +27,20 @@ pub enum DecimalVector {
     D256(DVector<i256>),
 }
 
+impl DecimalVector {
+    /// Returns the [`DecimalType`] of the decimal vector.
+    pub fn decimal_type(&self) -> DecimalType {
+        match self {
+            Self::D8(_) => DecimalType::I8,
+            Self::D16(_) => DecimalType::I16,
+            Self::D32(_) => DecimalType::I32,
+            Self::D64(_) => DecimalType::I64,
+            Self::D128(_) => DecimalType::I128,
+            Self::D256(_) => DecimalType::I256,
+        }
+    }
+}
+
 impl VectorOps for DecimalVector {
     type Mutable = DecimalVectorMut;
 

vortex-vector/src/decimal/vector_mut.rs

@@ -31,6 +31,18 @@ pub enum DecimalVectorMut {
 }
 
 impl DecimalVectorMut {
+    /// Returns the [`DecimalType`] of the decimal vector.
+    pub fn decimal_type(&self) -> DecimalType {
+        match self {
+            Self::D8(_) => DecimalType::I8,
+            Self::D16(_) => DecimalType::I16,
+            Self::D32(_) => DecimalType::I32,
+            Self::D64(_) => DecimalType::I64,
+            Self::D128(_) => DecimalType::I128,
+            Self::D256(_) => DecimalType::I256,
+        }
+    }
+
     /// Create a new mutable decimal vector with the given primitive type and capacity.
     pub fn with_capacity(decimal_dtype: &DecimalDType, capacity: usize) -> Self {
         let decimal_type = DecimalType::smallest_decimal_value_type(decimal_dtype);