vortex_array/arrays/dict/compute/
fill_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use super::DictArray;
7use super::DictVTable;
8use crate::Array;
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::ToCanonical;
13use crate::arrays::ConstantArray;
14use crate::builtins::ArrayBuiltins;
15use crate::match_each_integer_ptype;
16use crate::scalar::Scalar;
17use crate::scalar::ScalarValue;
18use crate::scalar_fn::fns::fill_null::FillNullKernel;
19use crate::scalar_fn::fns::operators::Operator;
20
21impl FillNullKernel for DictVTable {
22    fn fill_null(
23        array: &DictArray,
24        fill_value: &Scalar,
25        _ctx: &mut ExecutionCtx,
26    ) -> VortexResult<Option<ArrayRef>> {
27        // If the fill value already exists in the dictionary, we can simply rewrite the null codes
28        // to point to the value.
29        let found_fill_values = array
30            .values()
31            .to_array()
32            .binary(
33                ConstantArray::new(fill_value.clone(), array.values().len()).to_array(),
34                Operator::Eq,
35            )?
36            .to_bool();
37
38        // We found the fill value already in the values at this given index.
39        let Some(existing_fill_value_index) =
40            found_fill_values.to_bit_buffer().set_indices().next()
41        else {
42            // No fill values found, so we must canonicalize and fill_null.
43            return Ok(Some(
44                array
45                    .to_canonical()?
46                    .into_array()
47                    .fill_null(fill_value.clone())?,
48            ));
49        };
50
51        // Now we rewrite the nullable codes to point at the fill value.
52        let codes = array.codes();
53
54        // Cast the index to the correct unsigned integer type matching the codes' ptype.
55        let codes_ptype = codes.dtype().as_ptype();
56
57        #[expect(
58            clippy::cast_possible_truncation,
59            reason = "The existing index must be representable by the existing ptype"
60        )]
61        let fill_scalar_value = match_each_integer_ptype!(codes_ptype, |P| {
62            ScalarValue::from(existing_fill_value_index as P)
63        });
64
65        // Fill nulls in both the codes and the values. Note that the precondition of this function
66        // states that the fill value is non-null, so we do not have to worry about the nullability.
67        let codes = codes.to_array().fill_null(Scalar::try_new(
68            codes.dtype().as_nonnullable(),
69            Some(fill_scalar_value),
70        )?)?;
71        let values = array.values().to_array().fill_null(fill_value.clone())?;
72
73        // SAFETY: invariants are still satisfied after patching nulls.
74        unsafe {
75            Ok(Some(
76                DictArray::new_unchecked(codes, values)
77                    .set_all_values_referenced(array.has_all_values_referenced())
78                    .into_array(),
79            ))
80        }
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use vortex_buffer::BitBuffer;
87    use vortex_buffer::buffer;
88    use vortex_error::VortexExpect;
89
90    use crate::IntoArray;
91    use crate::ToCanonical;
92    use crate::arrays::PrimitiveArray;
93    use crate::arrays::dict::DictArray;
94    use crate::assert_arrays_eq;
95    use crate::builtins::ArrayBuiltins;
96    use crate::dtype::Nullability;
97    use crate::scalar::Scalar;
98    use crate::validity::Validity;
99
100    #[test]
101    fn nullable_codes_fill_in_values() {
102        let dict = DictArray::try_new(
103            PrimitiveArray::new(
104                buffer![0u32, 1, 2],
105                Validity::from(BitBuffer::from(vec![true, false, true])),
106            )
107            .into_array(),
108            PrimitiveArray::new(buffer![10, 20, 20], Validity::AllValid).into_array(),
109        )
110        .vortex_expect("operation should succeed in test");
111
112        let filled = dict
113            .to_array()
114            .fill_null(Scalar::primitive(20, Nullability::NonNullable))
115            .vortex_expect("operation should succeed in test");
116        let filled_primitive = filled.to_primitive();
117        assert_arrays_eq!(filled_primitive, PrimitiveArray::from_iter([10, 20, 20]));
118        assert!(filled_primitive.all_valid().unwrap());
119    }
120}
vortex_array/arrays/dict/compute/fill_null.rs

vortex_array/arrays/dict/compute/
fill_null.rs
