chore[array]: fill_null cleanup + decimal impl (#5175)

joseph-isaacs · web-flow · commit 42a65ff83456 · 2025-11-04T15:27:32.000Z
Signed-off-by: Joe Isaacs &lt;joe.isaacs@live.co.uk&gt;
diff --git a/vortex-array/src/arrays/bool/compute/fill_null.rs b/vortex-array/src/arrays/bool/compute/fill_null.rs
@@ -4,7 +4,7 @@
 use vortex_error::{VortexResult, vortex_err};
 use vortex_scalar::Scalar;
 
-use crate::arrays::{BoolArray, BoolVTable, ConstantArray};
+use crate::arrays::{BoolArray, BoolVTable};
 use crate::compute::{FillNullKernel, FillNullKernelAdapter};
 use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
@@ -18,14 +18,6 @@ impl FillNullKernel for BoolVTable {
             .ok_or_else(|| vortex_err!("Fill value must be non null"))?;
 
         Ok(match array.validity() {
-            Validity::NonNullable | Validity::AllValid => BoolArray::from_bit_buffer(
-                array.bit_buffer().clone(),
-                fill_value.dtype().nullability().into(),
-            )
-            .into_array(),
-            Validity::AllInvalid => {
-                ConstantArray::new(fill_value.clone(), array.len()).into_array()
-            }
             Validity::Array(v) => {
                 let bool_buffer = if fill {
                     array.bit_buffer() | &!v.to_bool().bit_buffer()
@@ -35,6 +27,7 @@ impl FillNullKernel for BoolVTable {
                 BoolArray::from_bit_buffer(bool_buffer, fill_value.dtype().nullability().into())
                     .into_array()
             }
+            _ => unreachable!("checked in entry point"),
         })
     }
 }
diff --git a/vortex-array/src/arrays/decimal/array.rs b/vortex-array/src/arrays/decimal/array.rs
@@ -204,6 +204,19 @@ impl DecimalArray {
         self.decimal_dtype().scale()
     }
 
+    pub fn from_iter<T: NativeDecimalType, I: IntoIterator<Item = T>>(
+        iter: I,
+        decimal_dtype: DecimalDType,
+    ) -> Self {
+        let iter = iter.into_iter();
+
+        Self::new(
+            BufferMut::from_iter(iter).freeze(),
+            decimal_dtype,
+            Validity::NonNullable,
+        )
+    }
+
     pub fn from_option_iter<T: NativeDecimalType, I: IntoIterator<Item = Option<T>>>(
         iter: I,
         decimal_dtype: DecimalDType,
diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs
@@ -0,0 +1,133 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ops::Not;
+
+use vortex_error::{VortexExpect, VortexResult};
+use vortex_scalar::{Scalar, match_each_decimal_value_type};
+
+use crate::arrays::DecimalVTable;
+use crate::arrays::decimal::DecimalArray;
+use crate::compute::{FillNullKernel, FillNullKernelAdapter};
+use crate::validity::Validity;
+use crate::vtable::ValidityHelper;
+use crate::{ArrayRef, IntoArray, ToCanonical, register_kernel};
+
+impl FillNullKernel for DecimalVTable {
+    fn fill_null(&self, array: &DecimalArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
+        let result_validity = Validity::from(fill_value.dtype().nullability());
+
+        Ok(match array.validity() {
+            Validity::Array(is_valid) => {
+                let is_invalid = is_valid.to_bool().bit_buffer().not();
+                match_each_decimal_value_type!(array.values_type(), |T| {
+                    let mut buffer = array.buffer::<T>().into_mut();
+                    let fill_value = fill_value
+                        .as_decimal()
+                        .decimal_value()
+                        .and_then(|v| v.cast::<T>())
+                        .vortex_expect("top-level fill_null ensure non-null fill value");
+                    for invalid_index in is_invalid.set_indices() {
+                        buffer[invalid_index] = fill_value;
+                    }
+                    DecimalArray::new(buffer.freeze(), array.decimal_dtype(), result_validity)
+                        .into_array()
+                })
+            }
+            _ => unreachable!("checked in entry point"),
+        })
+    }
+}
+
+register_kernel!(FillNullKernelAdapter(DecimalVTable).lift());
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::buffer;
+    use vortex_dtype::{DecimalDType, Nullability};
+    use vortex_scalar::{DecimalValue, Scalar};
+
+    use crate::arrays::decimal::DecimalArray;
+    use crate::assert_arrays_eq;
+    use crate::canonical::ToCanonical;
+    use crate::compute::fill_null;
+    use crate::validity::Validity;
+
+    #[test]
+    fn fill_null_leading_none() {
+        let decimal_dtype = DecimalDType::new(19, 2);
+        let arr = DecimalArray::from_option_iter(
+            [None, Some(800i128), None, Some(1000i128), None],
+            decimal_dtype,
+        );
+        let p = fill_null(
+            arr.as_ref(),
+            &Scalar::decimal(
+                DecimalValue::I128(4200i128),
+                DecimalDType::new(19, 2),
+                Nullability::NonNullable,
+            ),
+        )
+        .unwrap()
+        .to_decimal();
+        assert_arrays_eq!(
+            p,
+            DecimalArray::from_iter([4200, 800, 4200, 1000, 4200], decimal_dtype)
+        );
+        assert_eq!(
+            p.buffer::<i128>().as_slice(),
+            vec![4200, 800, 4200, 1000, 4200]
+        );
+        assert!(p.validity_mask().all_true());
+    }
+
+    #[test]
+    fn fill_null_all_none() {
+        let decimal_dtype = DecimalDType::new(19, 2);
+
+        let arr = DecimalArray::from_option_iter(
+            [Option::<i128>::None, None, None, None, None],
+            decimal_dtype,
+        );
+
+        let p = fill_null(
+            arr.as_ref(),
+            &Scalar::decimal(
+                DecimalValue::I128(25500i128),
+                DecimalDType::new(19, 2),
+                Nullability::NonNullable,
+            ),
+        )
+        .unwrap()
+        .to_decimal();
+        assert_arrays_eq!(
+            p,
+            DecimalArray::from_iter([25500, 25500, 25500, 25500, 25500], decimal_dtype)
+        );
+    }
+
+    #[test]
+    fn fill_null_non_nullable() {
+        let decimal_dtype = DecimalDType::new(19, 2);
+
+        let arr = DecimalArray::new(
+            buffer![800i128, 1000i128, 1200i128, 1400i128, 1600i128],
+            decimal_dtype,
+            Validity::NonNullable,
+        );
+        let p = fill_null(
+            arr.as_ref(),
+            &Scalar::decimal(
+                DecimalValue::I128(25500i128),
+                DecimalDType::new(19, 2),
+                Nullability::NonNullable,
+            ),
+        )
+        .unwrap()
+        .to_decimal();
+        assert_arrays_eq!(
+            p,
+            DecimalArray::from_iter([800i128, 1000, 1200, 1400, 1600], decimal_dtype)
+        );
+    }
+}
diff --git a/vortex-array/src/arrays/decimal/compute/mod.rs b/vortex-array/src/arrays/decimal/compute/mod.rs
@@ -3,6 +3,7 @@
 
 mod between;
 mod cast;
+mod fill_null;
 mod filter;
 mod is_constant;
 mod is_sorted;
diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs
@@ -3,39 +3,26 @@
 
 use std::ops::Not;
 
-use vortex_buffer::BufferMut;
-use vortex_dtype::{Nullability, match_each_native_ptype};
+use vortex_dtype::match_each_native_ptype;
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_scalar::Scalar;
 
+use crate::arrays::PrimitiveVTable;
 use crate::arrays::primitive::PrimitiveArray;
-use crate::arrays::{ConstantArray, PrimitiveVTable};
 use crate::compute::{FillNullKernel, FillNullKernelAdapter};
 use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
 use crate::{ArrayRef, IntoArray, ToCanonical, register_kernel};
 
 impl FillNullKernel for PrimitiveVTable {
     fn fill_null(&self, array: &PrimitiveArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
-        let result_validity = match fill_value.dtype().nullability() {
-            Nullability::NonNullable => Validity::NonNullable,
-            Nullability::Nullable => Validity::AllValid,
-        };
+        let result_validity = Validity::from(fill_value.dtype().nullability());
 
         Ok(match array.validity() {
-            Validity::NonNullable | Validity::AllValid => {
-                match_each_native_ptype!(array.ptype(), |T| {
-                    PrimitiveArray::new::<T>(array.buffer(), result_validity).into_array()
-                })
-            }
-            Validity::AllInvalid => {
-                ConstantArray::new(fill_value.clone(), array.len()).into_array()
-            }
             Validity::Array(is_valid) => {
-                // TODO(danking): when we take PrimitiveArray by value, we should mutate in-place
                 let is_invalid = is_valid.to_bool().bit_buffer().not();
                 match_each_native_ptype!(array.ptype(), |T| {
-                    let mut buffer = BufferMut::copy_from(array.as_slice::<T>());
+                    let mut buffer = array.buffer::<T>().into_mut();
                     let fill_value = fill_value
                         .as_primitive()
                         .typed_value::<T>()
@@ -46,6 +33,7 @@ impl FillNullKernel for PrimitiveVTable {
                     PrimitiveArray::new(buffer.freeze(), result_validity).into_array()
                 })
             }
+            _ => unreachable!("checked in entry point"),
         })
     }
 }
diff --git a/vortex-array/src/compute/fill_null.rs b/vortex-array/src/compute/fill_null.rs
@@ -8,6 +8,7 @@ use vortex_dtype::DType;
 use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
 use vortex_scalar::Scalar;
 
+use crate::arrays::ConstantArray;
 use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output, cast};
 use crate::vtable::VTable;
 use crate::{Array, ArrayRef, IntoArray};
@@ -89,6 +90,12 @@ impl ComputeFnVTable for FillNull {
             return Ok(cast(array, fill_value.dtype())?.into());
         }
 
+        if array.all_invalid() {
+            return Ok(ConstantArray::new(fill_value.clone(), array.len())
+                .into_array()
+                .into());
+        }
+
         if fill_value.is_null() {
             vortex_bail!("Cannot fill_null with a null value")
         }
