vortex_array/arrays/constant/compute/
fill_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_scalar::Scalar;
6
7use crate::ArrayRef;
8use crate::IntoArray;
9use crate::arrays::ConstantArray;
10use crate::arrays::ConstantVTable;
11use crate::compute::FillNullKernel;
12use crate::compute::FillNullKernelAdapter;
13use crate::compute::cast;
14use crate::register_kernel;
15
16impl FillNullKernel for ConstantVTable {
17    fn fill_null(&self, array: &ConstantArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
18        if array.scalar().is_null() {
19            Ok(ConstantArray::new(fill_value.clone(), array.len()).into_array())
20        } else {
21            cast(array.as_ref(), fill_value.dtype())
22        }
23    }
24}
25
26register_kernel!(FillNullKernelAdapter(ConstantVTable).lift());
27
28#[cfg(test)]
29mod test {
30    use vortex_scalar::Scalar;
31
32    use crate::IntoArray as _;
33    use crate::arrays::ConstantArray;
34    use crate::arrow::IntoArrowArray as _;
35    use crate::compute::fill_null;
36
37    #[test]
38    fn test_null() {
39        let actual = fill_null(
40            &ConstantArray::new(Scalar::from(None::<i32>), 3).into_array(),
41            &Scalar::from(1),
42        )
43        .unwrap();
44        let expected = ConstantArray::new(Scalar::from(1), 3).into_array();
45
46        assert!(!actual.dtype().is_nullable());
47
48        let actual_arrow = actual.clone().into_arrow_preferred().unwrap();
49        let expected_arrow = expected.clone().into_arrow_preferred().unwrap();
50        assert_eq!(
51            &actual_arrow,
52            &expected_arrow,
53            "{}, {}",
54            actual.display_values(),
55            expected.display_values()
56        );
57    }
58
59    #[test]
60    fn test_non_null() {
61        let actual = fill_null(
62            &ConstantArray::new(Scalar::from(Some(1)), 3).into_array(),
63            &Scalar::from(1),
64        )
65        .unwrap();
66        let expected = ConstantArray::new(Scalar::from(1), 3).into_array();
67
68        assert!(!actual.dtype().is_nullable());
69
70        let actual_arrow = actual.clone().into_arrow_preferred().unwrap();
71        let expected_arrow = expected.clone().into_arrow_preferred().unwrap();
72        assert_eq!(
73            &actual_arrow,
74            &expected_arrow,
75            "{}, {}",
76            actual.display_values(),
77            expected.display_values()
78        );
79    }
80
81    #[test]
82    fn test_non_nullable_with_nullable() {
83        let actual = fill_null(
84            &ConstantArray::new(Scalar::from(1), 3).into_array(),
85            &Scalar::from(Some(1)),
86        )
87        .unwrap();
88        let expected = ConstantArray::new(Scalar::from(1), 3).into_array();
89
90        assert!(!Scalar::from(1).dtype().is_nullable());
91
92        assert!(actual.dtype().is_nullable());
93
94        let actual_arrow = actual.clone().into_arrow_preferred().unwrap();
95        let expected_arrow = expected.clone().into_arrow_preferred().unwrap();
96        assert_eq!(
97            &actual_arrow,
98            &expected_arrow,
99            "{}, {}",
100            actual.display_values(),
101            expected.display_values()
102        );
103    }
104}