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