vortex_array/arrays/varbinview/compute/
fill_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Not;
5use std::sync::Arc;
6
7use vortex_buffer::BufferMut;
8use vortex_buffer::BufferString;
9use vortex_buffer::ByteBuffer;
10use vortex_error::VortexResult;
11use vortex_error::vortex_err;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::IntoArray;
16use crate::array::ArrayView;
17use crate::arrays::BoolArray;
18use crate::arrays::VarBinView;
19use crate::arrays::VarBinViewArray;
20use crate::arrays::varbinview::BinaryView;
21use crate::dtype::DType;
22use crate::scalar::Scalar;
23use crate::scalar_fn::fns::fill_null::FillNullKernel;
24use crate::validity::Validity;
25
26// Points every null row at one view of the fill value. The data buffers are reused as they are,
27// so only the views are copied, plus one buffer for a fill value too long to inline.
28impl FillNullKernel for VarBinView {
29    fn fill_null(
30        array: ArrayView<'_, VarBinView>,
31        fill_value: &Scalar,
32        ctx: &mut ExecutionCtx,
33    ) -> VortexResult<Option<ArrayRef>> {
34        let fill_bytes = match fill_value.dtype() {
35            DType::Utf8(_) => fill_value
36                .as_utf8()
37                .value()
38                .cloned()
39                .map(BufferString::into_inner),
40            DType::Binary(_) => fill_value.as_binary().value().cloned(),
41            _ => None,
42        }
43        .ok_or_else(|| vortex_err!("Fill value must be a non-null utf8 or binary scalar"))?;
44
45        let Validity::Array(is_valid) = array.validity()? else {
46            unreachable!("checked in entry point");
47        };
48        let is_invalid = is_valid.execute::<BoolArray>(ctx)?.into_bit_buffer().not();
49
50        let mut buffers: Vec<ByteBuffer> = array
51            .data_buffers()
52            .iter()
53            .map(|buffer| buffer.as_host().clone())
54            .collect();
55        let fill_view =
56            BinaryView::make_view(fill_bytes.as_slice(), u32::try_from(buffers.len())?, 0);
57        if !fill_view.is_inlined() {
58            buffers.push(fill_bytes);
59        }
60
61        let mut views = BufferMut::copy_from_in(array.views(), ctx.allocator().clone());
62        for (start, end) in is_invalid.set_slices() {
63            views[start..end].fill(fill_view);
64        }
65
66        let nullability = fill_value.dtype().nullability();
67        // SAFETY: the views of valid rows are unchanged. Null rows now hold the fill value's view,
68        // which is either inlined or points at the buffer pushed above. The fill value has the
69        // array's dtype, so it is valid UTF-8 whenever the array has to be.
70        let filled = unsafe {
71            VarBinViewArray::new_unchecked(
72                views.freeze(),
73                Arc::from(buffers),
74                array.dtype().with_nullability(nullability),
75                Validity::from(nullability),
76            )
77        };
78        Ok(Some(filled.into_array()))
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use rstest::rstest;
85    use vortex_error::VortexResult;
86
87    use crate::IntoArray;
88    use crate::VortexSessionExecute;
89    use crate::array_session;
90    use crate::arrays::VarBinArray;
91    use crate::arrays::VarBinViewArray;
92    use crate::assert_arrays_eq;
93    use crate::builtins::ArrayBuiltins;
94    use crate::dtype::DType;
95    use crate::dtype::Nullability;
96    use crate::scalar::Scalar;
97
98    const LONG: &str = "a value that is too long to inline";
99
100    #[rstest]
101    #[case::inlined_fill("fill", 0)]
102    #[case::outlined_fill("a fill value that is too long to inline", 1)]
103    fn fill_null_utf8(#[case] fill: &str, #[case] added_buffers: usize) -> VortexResult<()> {
104        let mut ctx = array_session().create_execution_ctx();
105        let array =
106            VarBinViewArray::from_iter_nullable_str([None, Some("short"), None, Some(LONG), None]);
107        let input_buffers = array.data_buffers().len();
108
109        let filled = array
110            .into_array()
111            .fill_null(Scalar::utf8(fill, Nullability::NonNullable))?
112            .execute::<VarBinViewArray>(&mut ctx)?;
113
114        assert_eq!(filled.dtype(), &DType::Utf8(Nullability::NonNullable));
115        assert_eq!(filled.data_buffers().len(), input_buffers + added_buffers);
116        assert_arrays_eq!(
117            filled,
118            VarBinViewArray::from_iter_str([fill, "short", fill, LONG, fill]),
119            &mut ctx
120        );
121        Ok(())
122    }
123
124    #[test]
125    fn fill_null_binary_with_nullable_fill() -> VortexResult<()> {
126        let mut ctx = array_session().create_execution_ctx();
127        let array = VarBinViewArray::from_iter_nullable_bin([Some(b"a".as_slice()), None]);
128
129        let filled = array
130            .into_array()
131            .fill_null(Scalar::binary(
132                LONG.as_bytes().to_vec(),
133                Nullability::Nullable,
134            ))?
135            .execute::<VarBinViewArray>(&mut ctx)?;
136
137        assert_eq!(filled.dtype(), &DType::Binary(Nullability::Nullable));
138        assert!(filled.validity()?.definitely_no_nulls());
139        assert_arrays_eq!(
140            filled,
141            VarBinViewArray::from_iter_nullable_bin([Some(b"a".as_slice()), Some(LONG.as_bytes())]),
142            &mut ctx
143        );
144        Ok(())
145    }
146
147    #[test]
148    fn fill_null_varbin_input() -> VortexResult<()> {
149        let mut ctx = array_session().create_execution_ctx();
150        let array = VarBinArray::from_nullable_strs(vec![Some("a"), None, Some(LONG)]);
151
152        let filled = array
153            .into_array()
154            .fill_null(Scalar::utf8("fill", Nullability::NonNullable))?;
155
156        assert_arrays_eq!(
157            filled,
158            VarBinViewArray::from_iter_str(["a", "fill", LONG]),
159            &mut ctx
160        );
161        Ok(())
162    }
163}
vortex_array/arrays/varbinview/compute/fill_null.rs

vortex_array/arrays/varbinview/compute/
fill_null.rs
