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_buffer::buffer;
59    use vortex_session::VortexSession;
60
61    use crate::Sparse;
62    use crate::initialize;
63
64    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
65        let session = vortex_array::array_session();
66        initialize(&session);
67        session
68    });
69
70    fn nullable_i32() -> DType {
71        DType::Primitive(PType::I32, Nullability::Nullable)
72    }
73
74    #[rstest]
75    // null fill, some null patches
76    #[case(Sparse::try_new(
77        buffer![1u64, 3, 5].into_array(),
78        PrimitiveArray::from_option_iter([Some(10i32), None, Some(30)]).into_array().cast(nullable_i32()).unwrap(),
79        8,
80        Scalar::null(nullable_i32()),
81    ).unwrap().into_array())]
82    // non-null fill, nullable patches with a null
83    #[case(Sparse::try_new(
84        buffer![0u64, 2].into_array(),
85        PrimitiveArray::from_option_iter([Some(7i32), None]).into_array().cast(nullable_i32()).unwrap(),
86        4,
87        Scalar::from(1i32).cast(&nullable_i32()).unwrap(),
88    ).unwrap().into_array())]
89    fn fill_null_matches_canonical(#[case] array: vortex_array::ArrayRef) {
90        let mut ctx = SESSION.create_execution_ctx();
91        let fill = Scalar::from(0i32);
92
93        // Kernel path: fill_null pushes through the Sparse encoding.
94        let kernel = array
95            .fill_null(fill.clone())
96            .unwrap()
97            .execute::<Canonical>(&mut ctx)
98            .unwrap();
99
100        // Baseline: canonicalize first, then fill_null on the PrimitiveArray.
101        let canonical_input = array.execute::<Canonical>(&mut ctx).unwrap().into_array();
102        let baseline = canonical_input
103            .fill_null(fill)
104            .unwrap()
105            .execute::<Canonical>(&mut ctx)
106            .unwrap();
107
108        assert_arrays_eq!(kernel, baseline, &mut ctx);
109    }
110}
vortex_sparse/compute/fill_null.rs

vortex_sparse/compute/
fill_null.rs
