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

vortex_array/arrays/dict/compute/
fill_null.rs
