Share fixed-width structural array operations

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

@@ -44,8 +44,6 @@ use crate::validity::Validity;
 
 /// The validity bitmap indicating which elements are non-null.
 pub(super) const VALIDITY_SLOT: usize = 0;
-pub(super) const NUM_SLOTS: usize = 1;
-pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"];
 
 /// A decimal array that stores fixed-precision decimal numbers with configurable scale.
 ///

vortex-array/src/arrays/decimal/compute/take.rs

@@ -1,33 +1,32 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::ptr;
-
-use itertools::Itertools as _;
-use vortex_buffer::Buffer;
-use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
-use vortex_error::vortex_ensure;
-use vortex_error::vortex_err;
 
 use crate::ArrayRef;
 use crate::Columnar;
 use crate::IntoArray;
 use crate::array::ArrayView;
+use crate::arrays::ConstantArray;
 use crate::arrays::Decimal;
 use crate::arrays::DecimalArray;
 use crate::arrays::PiecewiseSequence;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::dict::TakeExecute;
+use crate::arrays::fixed_width::take::PreparedTakeIndices;
+use crate::arrays::fixed_width::take::prepare_take_indices;
+use crate::arrays::fixed_width::take::take_slices as take_slices_to_buffer;
+use crate::arrays::fixed_width::take::take_slices_constant_length as take_slices_constant_length_to_buffer;
+use crate::arrays::fixed_width::take::take_values;
 use crate::arrays::piecewise_sequence::constant_unsigned_usize;
 use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
-use crate::dtype::IntegerPType;
 use crate::dtype::NativeDecimalType;
 use crate::dtype::UnsignedPType;
 use crate::executor::ExecutionCtx;
 use crate::match_each_decimal_value_type;
 use crate::match_each_integer_ptype;
 use crate::match_each_unsigned_integer_ptype;
+use crate::scalar::Scalar;
 use crate::validity::Validity;
 
 impl TakeExecute for Decimal {
@@ -42,15 +41,27 @@ impl TakeExecute for Decimal {
             return Ok(Some(taken));
         }
 
-        let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
-        let validity = array.validity()?.take(&indices.clone().into_array())?;
+        let PreparedTakeIndices::Values {
+            indices,
+            validity: indices_validity,
+        } = prepare_take_indices(indices, ctx)?
+        else {
+            return Ok(Some(
+                ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len())
+                    .into_array(),
+            ));
+        };
+        let validity = array
+            .validity()?
+            .take(&indices.clone().into_array())?
+            .and(indices_validity)?;
 
         // TODO(joe): if the true count of take indices validity is low, only take array values with
         // valid indices.
         let decimal = match_each_decimal_value_type!(array.values_type(), |D| {
             match_each_integer_ptype!(indices.ptype(), |I| {
                 let buffer =
-                    take_to_buffer::<I, D>(indices.as_slice::<I>(), array.buffer::<D>().as_slice());
+                    take_values::<D, I>(array.buffer::<D>().as_slice(), indices.as_slice::<I>())?;
                 // SAFETY: Take operation preserves decimal dtype and creates valid buffer.
                 // Validity is computed correctly from the parent array and indices.
                 unsafe { DecimalArray::new_unchecked(buffer, array.decimal_dtype(), validity) }
@@ -108,10 +119,10 @@ where
     D: NativeDecimalType,
 {
     match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
-        let values = take_slices_constant_length_to_buffer::<S, D>(
+        let values = take_slices_constant_length_to_buffer::<D, S>(
+            array.buffer::<D>().as_slice(),
             starts.as_slice::<S>(),
             length,
-            array.buffer::<D>().as_slice(),
             output_len,
         )?;
 
@@ -162,10 +173,10 @@ where
     S: UnsignedPType,
 {
     match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
-        let values = take_slices_to_buffer::<S, L, D>(
+        let values = take_slices_to_buffer::<D, S, L>(
+            array.buffer::<D>().as_slice(),
             starts.as_slice::<S>(),
             lengths.as_slice::<L>(),
-            array.buffer::<D>().as_slice(),
             output_len,
         )?;
 
@@ -177,93 +188,6 @@ where
     })
 }
 
-fn take_to_buffer<I: IntegerPType, T: NativeDecimalType>(indices: &[I], values: &[T]) -> Buffer<T> {
-    indices.iter().map(|idx| values[idx.as_()]).collect()
-}
-
-fn take_slices_constant_length_to_buffer<S, T>(
-    starts: &[S],
-    length: usize,
-    values: &[T],
-    output_len: usize,
-) -> VortexResult<Buffer<T>>
-where
-    S: UnsignedPType,
-    T: NativeDecimalType,
-{
-    let computed_len = starts
-        .len()
-        .checked_mul(length)
-        .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
-    vortex_ensure!(
-        computed_len == output_len,
-        "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}"
-    );
-
-    let mut result = BufferMut::<T>::with_capacity(output_len);
-    let spare = &mut result.spare_capacity_mut()[..output_len];
-    let mut cursor = 0usize;
-    for &start in starts {
-        let start = start.as_();
-        let src = &values[start..][..length];
-        // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap.
-        unsafe {
-            ptr::copy_nonoverlapping(
-                src.as_ptr(),
-                spare[cursor..][..src.len()].as_mut_ptr().cast::<T>(),
-                src.len(),
-            );
-        }
-        cursor += src.len();
-    }
-    // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity.
-    unsafe { result.set_len(cursor) };
-    vortex_ensure!(
-        result.len() == output_len,
-        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
-        result.len()
-    );
-    Ok(result.freeze())
-}
-
-fn take_slices_to_buffer<S, L, T>(
-    starts: &[S],
-    lengths: &[L],
-    values: &[T],
-    output_len: usize,
-) -> VortexResult<Buffer<T>>
-where
-    S: UnsignedPType,
-    L: UnsignedPType,
-    T: NativeDecimalType,
-{
-    let mut result = BufferMut::<T>::with_capacity(output_len);
-    let spare = &mut result.spare_capacity_mut()[..output_len];
-    let mut cursor = 0usize;
-    for (&start, &length) in starts.iter().zip_eq(lengths) {
-        let start = start.as_();
-        let length = length.as_();
-        let src = &values[start..][..length];
-        // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap.
-        unsafe {
-            ptr::copy_nonoverlapping(
-                src.as_ptr(),
-                spare[cursor..][..src.len()].as_mut_ptr().cast::<T>(),
-                src.len(),
-            );
-        }
-        cursor += src.len();
-    }
-    // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity.
-    unsafe { result.set_len(cursor) };
-    vortex_ensure!(
-        result.len() == output_len,
-        "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}",
-        result.len()
-    );
-    Ok(result.freeze())
-}
-
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
@@ -314,6 +238,27 @@ mod tests {
         assert_arrays_eq!(expected, taken, &mut ctx);
     }
 
+    #[test]
+    fn test_take_null_indices_ignore_physical_value() {
+        let mut ctx = array_session().create_execution_ctx();
+        let ddtype = DecimalDType::new(19, 1);
+        let array = DecimalArray::new(
+            buffer![10i128, 11i128, 12i128, 13i128],
+            ddtype,
+            Validity::NonNullable,
+        );
+
+        let indices = PrimitiveArray::new(
+            buffer![u32::MAX, 2, 3],
+            Validity::from_iter([false, true, true]),
+        )
+        .into_array();
+        let taken = array.take(indices).unwrap();
+
+        let expected = DecimalArray::from_option_iter([None, Some(12i128), Some(13)], ddtype);
+        assert_arrays_eq!(expected, taken, &mut ctx);
+    }
+
     #[rstest]
     #[case(DecimalArray::new(
         buffer![100i128, 200i128, 300i128, 400i128, 500i128],

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

@@ -8,7 +8,6 @@ use vortex_buffer::Alignment;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
-use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
 
 use crate::ArrayParts;
@@ -19,6 +18,7 @@ use crate::array::Array;
 use crate::array::ArrayView;
 use crate::array::VTable;
 use crate::arrays::decimal::DecimalData;
+use crate::arrays::fixed_width::vtable as fixed_width;
 use crate::buffer::BufferHandle;
 use crate::builders::ArrayBuilder;
 use crate::builders::DecimalBuilder;
@@ -27,7 +27,6 @@ use crate::dtype::DecimalType;
 use crate::dtype::NativeDecimalType;
 use crate::match_each_decimal_value_type;
 use crate::serde::ArrayChildren;
-use crate::validity::Validity;
 mod kernel;
 mod operations;
 mod validity;
@@ -38,7 +37,6 @@ use vortex_session::registry::CachedId;
 
 use crate::EqMode;
 use crate::array::ArrayId;
-use crate::arrays::decimal::array::SLOT_NAMES;
 use crate::arrays::decimal::compute::rules::RULES;
 use crate::hash::ArrayEq;
 use crate::hash::ArrayHash;
@@ -85,31 +83,20 @@ impl VTable for Decimal {
     }
 
     fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
-        match idx {
-            0 => array.values.clone(),
-            _ => vortex_panic!("DecimalArray buffer index {idx} out of bounds"),
-        }
+        fixed_width::buffer("DecimalArray", &array.values, idx)
     }
 
     fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
-        match idx {
-            0 => Some("values".to_string()),
-            _ => None,
-        }
+        fixed_width::buffer_name(idx)
     }
 
     fn with_buffers(
         &self,
         array: ArrayView<'_, Self>,
         buffers: &[BufferHandle],
     ) -> VortexResult<ArrayParts<Self>> {
-        vortex_ensure!(
-            buffers.len() == 1,
-            "Expected 1 buffer, got {}",
-            buffers.len()
-        );
         let mut data = array.data().clone();
-        data.values = buffers[0].clone();
+        data.values = fixed_width::replacement_buffer(buffers)?;
         Ok(
             ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
                 .with_slots(array.slots().iter().cloned().collect()),
@@ -138,25 +125,7 @@ impl VTable for Decimal {
         let DType::Decimal(_, nullability) = dtype else {
             vortex_bail!("Expected decimal dtype, got {dtype:?}");
         };
-        vortex_ensure!(
-            data.len() == len,
-            InvalidArgument:
-            "DecimalArray length {} does not match outer length {}",
-            data.len(),
-            len
-        );
-        let validity = crate::array::child_to_validity(slots[0].as_ref(), *nullability);
-        if let Some(validity_len) = validity.maybe_len() {
-            vortex_ensure!(
-                validity_len == len,
-                InvalidArgument:
-                "DecimalArray validity len {} does not match outer length {}",
-                validity_len,
-                len
-            );
-        }
-
-        Ok(())
+        fixed_width::validate_layout("DecimalArray", data.len(), *nullability, len, slots)
     }
 
     fn deserialize(
@@ -174,14 +143,7 @@ impl VTable for Decimal {
         }
         let values = buffers[0].clone();
 
-        let validity = if children.is_empty() {
-            Validity::from(dtype.nullability())
-        } else if children.len() == 1 {
-            let validity = children.get(0, &Validity::DTYPE, len)?;
-            Validity::Array(validity)
-        } else {
-            vortex_bail!("Expected 0 or 1 child, got {}", children.len());
-        };
+        let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?;
 
         let Some(decimal_dtype) = dtype.as_decimal_opt() else {
             vortex_bail!("Expected Decimal dtype, got {:?}", dtype)
@@ -201,7 +163,7 @@ impl VTable for Decimal {
     }
 
     fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
-        SLOT_NAMES[idx].to_string()
+        fixed_width::slot_name(idx)
     }
 
     fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {

vortex-array/src/arrays/fixed_width/mod.rs

@@ -0,0 +1,7 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Shared structural operations for fixed-width canonical arrays.
+
+pub(crate) mod take;
+pub(crate) mod vtable;