vortex_sparse/compute/
fill_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::builtins::ArrayBuiltins;
9use vortex_array::scalar::Scalar;
10use vortex_array::scalar_fn::fns::fill_null::FillNullKernel;
11use vortex_error::VortexResult;
12
13use crate::Sparse;
14use crate::SparseExt as _;
15
16/// Sparse-specific fill_null kernel.
17///
18/// `fill_null(Sparse{ F, patches }, v)` replaces nulls in the fill and in each patch value
19/// with the (non-null) `v`, staying sparse: the new fill is `v` if `F` was null, else `F`
20/// cast to the non-nullable result dtype. The work is `O(P)`.
21impl FillNullKernel for Sparse {
22    fn fill_null(
23        array: ArrayView<'_, Self>,
24        fill_value: &Scalar,
25        _ctx: &mut ExecutionCtx,
26    ) -> VortexResult<Option<ArrayRef>> {
27        let new_fill = if array.fill_scalar().is_null() {
28            fill_value.clone()
29        } else {
30            array.fill_scalar().cast(fill_value.dtype())?
31        };
32
33        let new_patches = array
34            .patches()
35            .map_values(|values| values.fill_null(fill_value.clone()))?;
36
37        Ok(Some(
38            Sparse::try_new_from_patches(new_patches, new_fill)?.into_array(),
39        ))
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use std::sync::LazyLock;
46
47    use rstest::rstest;
48    use vortex_array::Canonical;
49    use vortex_array::IntoArray;
50    use vortex_array::VortexSessionExecute;
51    use vortex_array::arrays::PrimitiveArray;
52    use vortex_array::assert_arrays_eq;
53    use vortex_array::builtins::ArrayBuiltins;
54    use vortex_array::dtype::DType;
55    use vortex_array::dtype::Nullability;
56    use vortex_array::dtype::PType;
57    use vortex_array::scalar::Scalar;
58    use vortex_array::session::ArraySession;
59    use vortex_buffer::buffer;
60    use vortex_session::VortexSession;
61
62    use crate::Sparse;
63    use crate::initialize;
64
65    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
66        let session = VortexSession::empty().with::<ArraySession>();
67        initialize(&session);
68        session
69    });
70
71    fn nullable_i32() -> DType {
72        DType::Primitive(PType::I32, Nullability::Nullable)
73    }
74
75    #[rstest]
76    // null fill, some null patches
77    #[case(Sparse::try_new(
78        buffer![1u64, 3, 5].into_array(),
79        PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array().cast(nullable_i32()).unwrap(),
80        8,
81        Scalar::null(nullable_i32()),
82    ).unwrap().into_array())]
83    // non-null fill, nullable patches with a null
84    #[case(Sparse::try_new(
85        buffer![0u64, 2].into_array(),
86        PrimitiveArray::from_option_iter([Some(7i32), None]).into_array().cast(nullable_i32()).unwrap(),
87        4,
88        Scalar::from(1i32).cast(&nullable_i32()).unwrap(),
89    ).unwrap().into_array())]
90    fn fill_null_matches_canonical(#[case] array: vortex_array::ArrayRef) {
91        let mut ctx = SESSION.create_execution_ctx();
92        let fill = Scalar::from(0i32);
93
94        // Kernel path: fill_null pushes through the Sparse encoding.
95        let kernel = array
96            .fill_null(fill.clone())
97            .unwrap()
98            .execute::<Canonical>(&mut ctx)
99            .unwrap();
100
101        // Baseline: canonicalize first, then fill_null on the PrimitiveArray.
102        let canonical_input = array.execute::<Canonical>(&mut ctx).unwrap().into_array();
103        let baseline = canonical_input
104            .fill_null(fill)
105            .unwrap()
106            .execute::<Canonical>(&mut ctx)
107            .unwrap();
108
109        assert_arrays_eq!(kernel, baseline);
110    }
111}
vortex_sparse/compute/fill_null.rs

vortex_sparse/compute/
fill_null.rs
