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

vortex_array/arrays/dict/compute/
fill_null.rs
