chore[fuzz]: added decimal array ops (#3571)

joseph-isaacs · web-flow · commit 776742aa90fa · 2025-06-19T19:58:26.000+01:00
Would be nice to be able to be generic over a `NativePType` and a
`NativeDecimalType`

---------

Signed-off-by: Joe Isaacs &lt;joe.isaacs@live.co.uk&gt;
diff --git a/fuzz/src/compare.rs b/fuzz/src/compare.rs
@@ -8,8 +8,8 @@ use vortex_array::compute::{Operator, scalar_cmp};
 use vortex_array::validity::Validity;
 use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
 use vortex_dtype::{DType, NativePType, match_each_native_ptype};
-use vortex_error::{VortexExpect, VortexResult};
-use vortex_scalar::Scalar;
+use vortex_error::{VortexExpect, VortexResult, vortex_err};
+use vortex_scalar::{NativeDecimalType, Scalar, match_each_decimal_value_type};
 
 pub fn compare_canonical_array(
     array: &dyn Array,
@@ -59,6 +59,27 @@ pub fn compare_canonical_array(
                 ))
             })
         }
+        DType::Decimal(..) => {
+            let decimal = value.as_decimal();
+            let decimal_array = array.to_decimal()?;
+            match_each_decimal_value_type!(decimal_array.values_type(), |D| {
+                let dval = decimal
+                    .decimal_value()
+                    .vortex_expect("nulls handled before")
+                    .cast::<D>()
+                    .ok_or_else(|| vortex_err!("todo: handle upcast of decimal array"))?;
+                let buf = decimal_array.buffer::<D>();
+                Ok(compare_native_decimal_type(
+                    buf.as_slice()
+                        .iter()
+                        .copied()
+                        .zip(array.validity_mask()?.to_boolean_buffer().iter())
+                        .map(|(b, v)| v.then_some(b)),
+                    dval,
+                    operator,
+                ))
+            })
+        }
         DType::Utf8(_) => array.to_varbinview()?.with_iterator(|iter| {
             let utf8_value = value
                 .as_utf8()
@@ -94,7 +115,9 @@ pub fn compare_canonical_array(
             )
             .into_array())
         }
-        d => unreachable!("DType {d} not supported for fuzzing"),
+        d @ (DType::Null | DType::Extension(_)) => {
+            unreachable!("DType {d} not supported for fuzzing")
+        }
     }
 }
 
@@ -133,3 +156,21 @@ fn compare_native_ptype<T: NativePType>(
     }))
     .into_array()
 }
+
+fn compare_native_decimal_type<D: NativeDecimalType>(
+    values: impl Iterator<Item = Option<D>>,
+    cmp_value: D,
+    operator: Operator,
+) -> ArrayRef {
+    BoolArray::from_iter(values.map(|val| {
+        val.map(|v| match operator {
+            Operator::Eq => v == cmp_value,
+            Operator::NotEq => v != cmp_value,
+            Operator::Gt => v > cmp_value,
+            Operator::Gte => v >= cmp_value,
+            Operator::Lt => v < cmp_value,
+            Operator::Lte => v <= cmp_value,
+        })
+    }))
+    .into_array()
+}
diff --git a/fuzz/src/filter.rs b/fuzz/src/filter.rs
@@ -1,12 +1,13 @@
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::{
-    BoolArray, BooleanBuffer, PrimitiveArray, StructArray, VarBinViewArray,
+    BoolArray, BooleanBuffer, DecimalArray, PrimitiveArray, StructArray, VarBinViewArray,
 };
 use vortex_array::validity::Validity;
 use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
 use vortex_buffer::Buffer;
 use vortex_dtype::{DType, match_each_native_ptype};
 use vortex_error::VortexResult;
+use vortex_scalar::match_each_decimal_value_type;
 
 use crate::take::take_canonical_array;
 
@@ -52,6 +53,23 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul
             )
             .into_array())
         }),
+        DType::Decimal(d, _) => {
+            let decimal_array = array.to_decimal()?;
+            match_each_decimal_value_type!(decimal_array.values_type(), |D| {
+                let buf = decimal_array.buffer::<D>();
+                Ok(DecimalArray::new(
+                    filter
+                        .iter()
+                        .zip(buf.as_slice().iter().copied())
+                        .filter(|(f, _)| **f)
+                        .map(|(_, v)| v)
+                        .collect::<Buffer<_>>(),
+                    *d,
+                    validity,
+                )
+                .into_array())
+            })
+        }
         DType::Utf8(_) | DType::Binary(_) => {
             let utf8 = array.to_varbinview()?;
             let values = utf8.with_iterator(|iter| {
@@ -87,6 +105,8 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul
             }
             take_canonical_array(array, &indices)
         }
-        d => unreachable!("DType {d} not supported for fuzzing"),
+        d @ (DType::Null | DType::Extension(_)) => {
+            unreachable!("DType {d} not supported for fuzzing")
+        }
     }
 }
diff --git a/fuzz/src/search_sorted.rs b/fuzz/src/search_sorted.rs
@@ -6,8 +6,8 @@ use vortex_array::search_sorted::{IndexOrd, SearchResult, SearchSorted, SearchSo
 use vortex_array::{Array, ToCanonical};
 use vortex_buffer::{BufferString, ByteBuffer};
 use vortex_dtype::{DType, NativePType, match_each_native_ptype};
-use vortex_error::VortexResult;
-use vortex_scalar::Scalar;
+use vortex_error::{VortexResult, vortex_err};
+use vortex_scalar::{Scalar, match_each_decimal_value_type};
 
 struct SearchNullableSlice<T>(Vec<Option<T>>);
 
@@ -76,6 +76,31 @@ pub fn search_sorted_canonical_array(
                 Ok(SearchPrimitiveSlice(opt_values).search_sorted(&Some(to_find), side))
             })
         }
+        DType::Decimal(d, _) => {
+            let decimal_array = array.to_decimal()?;
+            let validity = decimal_array.validity_mask()?.to_boolean_buffer();
+            match_each_decimal_value_type!(decimal_array.values_type(), |D| {
+                let buf = decimal_array.buffer::<D>();
+                let opt_values = buf
+                    .as_slice()
+                    .iter()
+                    .copied()
+                    .zip(validity.iter())
+                    .map(|(b, v)| v.then_some(b))
+                    .collect::<Vec<_>>();
+                let to_find: D = scalar
+                    .as_decimal()
+                    .decimal_value()
+                    .map(|v| {
+                        v.cast::<D>().ok_or_else(|| {
+                            vortex_err!("cannot cast value {v} to decimal value type {d}")
+                        })
+                    })
+                    .transpose()?
+                    .ok_or_else(|| vortex_err!("unexpected null scalar"))?;
+                Ok(SearchNullableSlice(opt_values).search_sorted(&Some(to_find), side))
+            })
+        }
         DType::Utf8(_) | DType::Binary(_) => {
             let utf8 = array.to_varbinview()?;
             let opt_values =
@@ -99,6 +124,8 @@ pub fn search_sorted_canonical_array(
                 .collect::<VortexResult<Vec<_>>>()?;
             Ok(scalar_vals.search_sorted(&scalar.cast(array.dtype())?, side))
         }
-        d => unreachable!("DType {d} not supported for fuzzing"),
+        d @ (DType::Null | DType::Extension(_)) => {
+            unreachable!("DType {d} not supported for fuzzing")
+        }
     }
 }
diff --git a/fuzz/src/slice.rs b/fuzz/src/slice.rs
@@ -1,11 +1,14 @@
 use vortex_array::accessor::ArrayAccessor;
-use vortex_array::arrays::{BoolArray, ListArray, PrimitiveArray, StructArray, VarBinViewArray};
+use vortex_array::arrays::{
+    BoolArray, DecimalArray, ListArray, PrimitiveArray, StructArray, VarBinViewArray,
+};
 use vortex_array::validity::Validity;
 use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
 use vortex_dtype::{
     DType, NativePType, Nullability, match_each_integer_ptype, match_each_native_ptype,
 };
 use vortex_error::VortexResult;
+use vortex_scalar::match_each_decimal_value_type;
 
 #[allow(clippy::unnecessary_fallible_conversions)]
 pub fn slice_canonical_array(
@@ -80,7 +83,22 @@ pub fn slice_canonical_array(
             .into_array();
             ListArray::try_new(elements, offsets, validity).map(|a| a.into_array())
         }
-        d => unreachable!("DType {d} not supported for fuzzing"),
+        DType::Decimal(decimal_dtype, _) => {
+            let decimal_array = array.to_decimal()?;
+            Ok(
+                match_each_decimal_value_type!(decimal_array.values_type(), |D| {
+                    DecimalArray::new(
+                        decimal_array.buffer::<D>().slice(start..stop),
+                        *decimal_dtype,
+                        validity,
+                    )
+                })
+                .to_array(),
+            )
+        }
+        d @ (DType::Null | DType::Extension(_)) => {
+            unreachable!("DType {d} not supported for fuzzing")
+        }
     }
 }
 
diff --git a/fuzz/src/sort.rs b/fuzz/src/sort.rs
@@ -1,10 +1,11 @@
 use std::cmp::Ordering;
 
 use vortex_array::accessor::ArrayAccessor;
-use vortex_array::arrays::{BoolArray, PrimitiveArray, VarBinViewArray};
+use vortex_array::arrays::{BoolArray, DecimalArray, PrimitiveArray, VarBinViewArray};
 use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
 use vortex_dtype::{DType, NativePType, match_each_native_ptype};
 use vortex_error::{VortexExpect, VortexResult, VortexUnwrap};
+use vortex_scalar::match_each_decimal_value_type;
 
 use crate::take::take_canonical_array;
 
@@ -35,6 +36,21 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult<ArrayRef> {
                 Ok(PrimitiveArray::from_option_iter(opt_values).into_array())
             })
         }
+        DType::Decimal(d, _) => {
+            let decimal_array = array.to_decimal()?;
+            match_each_decimal_value_type!(decimal_array.values_type(), |D| {
+                let buf = decimal_array.buffer::<D>();
+                let mut opt_values = buf
+                    .as_slice()
+                    .iter()
+                    .copied()
+                    .zip(decimal_array.validity_mask()?.to_boolean_buffer().iter())
+                    .map(|(p, v)| v.then_some(p))
+                    .collect::<Vec<_>>();
+                opt_values.sort();
+                Ok(DecimalArray::from_option_iter(opt_values, *d).into_array())
+            })
+        }
         DType::Utf8(_) | DType::Binary(_) => {
             let utf8 = array.to_varbinview()?;
             let mut opt_values =
@@ -64,7 +80,9 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult<ArrayRef> {
             });
             take_canonical_array(array, &sort_indices)
         }
-        d => unreachable!("DType {d} not supported for fuzzing"),
+        d @ (DType::Null | DType::Extension(_)) => {
+            unreachable!("DType {d} not supported for fuzzing")
+        }
     }
 }
 
diff --git a/fuzz/src/take.rs b/fuzz/src/take.rs
@@ -1,12 +1,13 @@
 use arrow_buffer::ArrowNativeType;
 use vortex_array::accessor::ArrayAccessor;
-use vortex_array::arrays::{BoolArray, PrimitiveArray, StructArray, VarBinViewArray};
+use vortex_array::arrays::{BoolArray, DecimalArray, PrimitiveArray, StructArray, VarBinViewArray};
 use vortex_array::builders::{ArrayBuilderExt, builder_with_capacity};
 use vortex_array::validity::Validity;
 use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
 use vortex_buffer::Buffer;
-use vortex_dtype::{DType, NativePType, match_each_native_ptype};
+use vortex_dtype::{DType, DecimalDType, NativePType, match_each_native_ptype};
 use vortex_error::VortexResult;
+use vortex_scalar::{NativeDecimalType, match_each_decimal_value_type};
 
 pub fn take_canonical_array(array: &dyn Array, indices: &[usize]) -> VortexResult<ArrayRef> {
     let validity = if array.dtype().is_nullable() {
@@ -32,6 +33,13 @@ pub fn take_canonical_array(array: &dyn Array, indices: &[usize]) -> VortexResul
                 Ok(take_primitive::<P>(primitive_array, validity, indices))
             })
         }
+        DType::Decimal(d, _) => {
+            let decimal_array = array.to_decimal()?;
+
+            match_each_decimal_value_type!(decimal_array.values_type(), |D| {
+                Ok(take_decimal::<D>(decimal_array, d, validity, indices))
+            })
+        }
         DType::Utf8(_) | DType::Binary(_) => {
             let utf8 = array.to_varbinview()?;
             let values =
@@ -65,7 +73,9 @@ pub fn take_canonical_array(array: &dyn Array, indices: &[usize]) -> VortexResul
             }
             Ok(builder.finish())
         }
-        d => unreachable!("DType {d} not supported for fuzzing"),
+        d @ (DType::Null | DType::Extension(_)) => {
+            unreachable!("DType {d} not supported for fuzzing")
+        }
     }
 }
 
@@ -84,3 +94,22 @@ fn take_primitive<T: NativePType + ArrowNativeType>(
     )
     .into_array()
 }
+
+fn take_decimal<D: NativeDecimalType>(
+    array: DecimalArray,
+    decimal_type: &DecimalDType,
+    validity: Validity,
+    indices: &[usize],
+) -> ArrayRef {
+    let buf = array.buffer::<D>();
+    let vec_values = buf.as_slice();
+    DecimalArray::new(
+        indices
+            .iter()
+            .map(|i| vec_values[*i])
+            .collect::<Buffer<D>>(),
+        *decimal_type,
+        validity,
+    )
+    .into_array()
+}
diff --git a/vortex-array/src/arrays/arbitrary.rs b/vortex-array/src/arrays/arbitrary.rs
@@ -70,7 +70,7 @@ fn random_array(u: &mut Unstructured, dtype: &DType, len: Option<usize>) -> Resu
                     PType::F64 => random_primitive::<f64>(u, *n, chunk_len),
                 },
                 DType::Decimal(decimal, n) => {
-                    let elem_len = u.int_in_range(0..=20)?;
+                    let elem_len = chunk_len.unwrap_or(u.int_in_range(0..=20)?);
                     match_each_decimal_value_type!(smallest_storage_type(decimal), |DVT| {
                         let mut builder =
                             DecimalBuilder::new::<DVT>(decimal.precision(), decimal.scale(), *n);
diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs
@@ -3,7 +3,8 @@ mod macros;
 mod ops;
 mod serde;
 
-use vortex_buffer::{Buffer, ByteBuffer};
+use arrow_buffer::BooleanBufferBuilder;
+use vortex_buffer::{Buffer, BufferMut, ByteBuffer};
 use vortex_dtype::{DType, DecimalDType};
 use vortex_error::{VortexResult, vortex_panic};
 use vortex_scalar::{DecimalValueType, NativeDecimalType};
@@ -130,6 +131,33 @@ impl DecimalArray {
     pub fn scale(&self) -> i8 {
         self.decimal_dtype().scale()
     }
+
+    pub fn from_option_iter<T: NativeDecimalType, I: IntoIterator<Item = Option<T>>>(
+        iter: I,
+        decimal_dtype: DecimalDType,
+    ) -> Self {
+        let iter = iter.into_iter();
+        let mut values = BufferMut::with_capacity(iter.size_hint().0);
+        let mut validity = BooleanBufferBuilder::new(values.capacity());
+
+        for i in iter {
+            match i {
+                None => {
+                    validity.append(false);
+                    values.push(T::default());
+                }
+                Some(e) => {
+                    validity.append(true);
+                    values.push(e);
+                }
+            }
+        }
+        Self::new(
+            values.freeze(),
+            decimal_dtype,
+            Validity::from(validity.finish()),
+        )
+    }
 }
 
 impl ArrayVTable<DecimalVTable> for DecimalVTable {
diff --git a/vortex-array/src/arrays/decimal/ops.rs b/vortex-array/src/arrays/decimal/ops.rs
@@ -10,7 +10,7 @@ use crate::{ArrayRef, IntoArray};
 
 impl OperationsVTable<DecimalVTable> for DecimalVTable {
     fn slice(array: &DecimalArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
-        match_each_decimal_value_type!(array.values_type, |D| {
+        match_each_decimal_value_type!(array.values_type(), |D| {
             slice_typed(
                 array.buffer::<D>(),
                 start,
diff --git a/vortex-array/src/arrays/struct_/mod.rs b/vortex-array/src/arrays/struct_/mod.rs
@@ -105,7 +105,7 @@ impl StructArray {
             if field.len() != length {
                 vortex_bail!(
                     "Expected all struct fields to have length {length}, found {}",
-                    field.len()
+                    fields.iter().map(|f| f.len()).format(","),
                 );
             }
         }

Original file line number	Diff line number	Diff line change
`@@ -105,7 +105,7 @@ impl StructArray {`
`105`	`105`	`if field.len() != length {`
`106`	`106`	`vortex_bail!(`
`107`	`107`	`"Expected all struct fields to have length {length}, found {}",`
`108`		`- field.len()`
	`108`	`+ fields.iter().map(\|f\| f.len()).format(","),`
`109`	`109`	`);`
`110`	`110`	`}`
`111`	`111`	`}`