vortex_array/scalar_fn/fns/fill_null/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod kernel;
5
6pub use kernel::*;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_ensure;
10use vortex_error::vortex_err;
11use vortex_session::VortexSession;
12use vortex_session::registry::CachedId;
13
14use crate::AnyColumnar;
15use crate::ArrayRef;
16use crate::CanonicalView;
17use crate::ColumnarView;
18use crate::ExecutionCtx;
19use crate::arrays::Bool;
20use crate::arrays::Decimal;
21use crate::arrays::Primitive;
22use crate::arrays::ScalarFnArray;
23use crate::builtins::ArrayBuiltins;
24use crate::dtype::DType;
25use crate::expr::Expression;
26use crate::scalar::Scalar;
27use crate::scalar_fn::Arity;
28use crate::scalar_fn::ChildName;
29use crate::scalar_fn::EmptyOptions;
30use crate::scalar_fn::ExecutionArgs;
31use crate::scalar_fn::ScalarFnId;
32use crate::scalar_fn::ScalarFnVTable;
33use crate::scalar_fn::ScalarFnVTableExt;
34
35/// An expression that replaces null values in the input with a fill value.
36#[derive(Clone)]
37pub struct FillNull;
38
39impl FillNull {
40    /// Creates a lazy operation that replaces null input values with `fill_value`.
41    ///
42    /// # Errors
43    ///
44    /// Returns an error if the children have different lengths or incompatible dtypes.
45    pub fn try_new(input: ArrayRef, fill_value: ArrayRef) -> VortexResult<ScalarFnArray> {
46        ScalarFnArray::try_new(FillNull.bind(EmptyOptions), vec![input, fill_value])
47    }
48}
49
50impl ScalarFnVTable for FillNull {
51    type Options = EmptyOptions;
52
53    fn id(&self) -> ScalarFnId {
54        static ID: CachedId = CachedId::new("vortex.fill_null");
55        *ID
56    }
57
58    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
59        Ok(Some(vec![]))
60    }
61
62    fn deserialize(
63        &self,
64        _metadata: &[u8],
65        _session: &VortexSession,
66    ) -> VortexResult<Self::Options> {
67        Ok(EmptyOptions)
68    }
69
70    fn arity(&self, _options: &Self::Options) -> Arity {
71        Arity::Exact(2)
72    }
73
74    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
75        match child_idx {
76            0 => ChildName::from("input"),
77            1 => ChildName::from("fill_value"),
78            _ => unreachable!("Invalid child index {} for FillNull expression", child_idx),
79        }
80    }
81
82    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
83        vortex_ensure!(
84            arg_dtypes[0].eq_ignore_nullability(&arg_dtypes[1]),
85            "fill_null requires input and fill value to have the same base type, got {} and {}",
86            arg_dtypes[0],
87            arg_dtypes[1]
88        );
89        // The result dtype takes the nullability of the fill value.
90        Ok(arg_dtypes[0]
91            .clone()
92            .with_nullability(arg_dtypes[1].nullability()))
93    }
94
95    fn execute(
96        &self,
97        _options: &Self::Options,
98        args: &dyn ExecutionArgs,
99        ctx: &mut ExecutionCtx,
100    ) -> VortexResult<ArrayRef> {
101        let input = args.get(0)?;
102        let fill_value = args.get(1)?;
103
104        let fill_scalar = fill_value
105            .as_constant()
106            .ok_or_else(|| vortex_err!("fill_null fill_value must be a constant/scalar"))?;
107
108        vortex_ensure!(
109            !fill_scalar.is_null(),
110            "fill_null requires a non-null fill value"
111        );
112
113        let Some(columnar) = input.as_opt::<AnyColumnar>() else {
114            return input.execute::<ArrayRef>(ctx)?.fill_null(fill_scalar);
115        };
116
117        match columnar {
118            ColumnarView::Canonical(canonical) => fill_null_canonical(canonical, &fill_scalar, ctx),
119            ColumnarView::Constant(constant) => fill_null_constant(constant, &fill_scalar),
120        }
121    }
122
123    fn simplify(
124        &self,
125        _options: &Self::Options,
126        expr: &Expression,
127        ctx: &dyn crate::scalar_fn::SimplifyCtx,
128    ) -> VortexResult<Option<Expression>> {
129        let input_dtype = ctx.return_dtype(expr.child(0))?;
130
131        if !input_dtype.is_nullable() {
132            return Ok(Some(expr.child(0).clone()));
133        }
134
135        Ok(None)
136    }
137
138    fn validity(
139        &self,
140        _options: &Self::Options,
141        expression: &Expression,
142    ) -> VortexResult<Option<Expression>> {
143        // After fill_null, the result validity depends on the fill value's nullability.
144        // If fill_value is non-nullable, the result is always valid.
145        Ok(Some(expression.child(1).validity()?))
146    }
147
148    fn is_strict(&self, _options: &Self::Options) -> bool {
149        // This function replaces null input values instead of propagating them.
150        false
151    }
152
153    fn is_infallible(&self, _options: &Self::Options) -> bool {
154        true
155    }
156}
157
158/// Fill nulls on a canonical array by directly dispatching to the appropriate kernel.
159///
160/// Returns the filled array, or bails if no kernel is registered for the canonical type.
161fn fill_null_canonical(
162    canonical: CanonicalView<'_>,
163    fill_value: &Scalar,
164    ctx: &mut ExecutionCtx,
165) -> VortexResult<ArrayRef> {
166    let arr = canonical.to_array_ref();
167    if let Some(result) = short_circuit(&arr, fill_value)? {
168        // The short circuit can return a lazy `ScalarFn`, so this forces it for now.
169        // TODO(aduffy): Remove this once we have better driver check. We're also implicitly
170        //  relying on the fact that Cast execution will do an optimize on its result.
171        return result.execute::<ArrayRef>(ctx);
172    }
173    match canonical {
174        CanonicalView::Bool(a) => <Bool as FillNullKernel>::fill_null(a, fill_value, ctx)?
175            .ok_or_else(|| vortex_err!("FillNullKernel for BoolArray returned None")),
176        CanonicalView::Primitive(a) => {
177            <Primitive as FillNullKernel>::fill_null(a, fill_value, ctx)?
178                .ok_or_else(|| vortex_err!("FillNullKernel for PrimitiveArray returned None"))
179        }
180        CanonicalView::Decimal(a) => <Decimal as FillNullKernel>::fill_null(a, fill_value, ctx)?
181            .ok_or_else(|| vortex_err!("FillNullKernel for DecimalArray returned None")),
182        other => vortex_bail!(
183            "No FillNullKernel for canonical array {}",
184            other.to_array_ref().encoding_id()
185        ),
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use vortex_buffer::buffer;
192    use vortex_error::VortexExpect;
193
194    use crate::IntoArray;
195    use crate::VortexSessionExecute;
196    use crate::array_session;
197    use crate::arrays::PrimitiveArray;
198    use crate::arrays::StructArray;
199    use crate::assert_arrays_eq;
200    use crate::dtype::DType;
201    use crate::dtype::Nullability;
202    use crate::dtype::PType;
203    use crate::expr::fill_null;
204    use crate::expr::get_item;
205    use crate::expr::lit;
206    use crate::expr::root;
207
208    #[test]
209    fn dtype() {
210        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
211        assert_eq!(
212            fill_null(root(), lit(0i32)).return_dtype(&dtype).unwrap(),
213            DType::Primitive(PType::I32, Nullability::NonNullable)
214        );
215    }
216
217    #[test]
218    fn replace_children() {
219        let expr = fill_null(root(), lit(0i32));
220        expr.with_children(vec![root(), lit(0i32)])
221            .vortex_expect("operation should succeed in test");
222    }
223
224    #[test]
225    fn evaluate() {
226        let mut ctx = array_session().create_execution_ctx();
227        let test_array =
228            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)])
229                .into_array();
230
231        let expr = fill_null(root(), lit(42i32));
232        let result = test_array.apply(&expr).unwrap();
233
234        assert_eq!(
235            result.dtype(),
236            &DType::Primitive(PType::I32, Nullability::NonNullable)
237        );
238        assert_arrays_eq!(
239            result,
240            PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]),
241            &mut ctx
242        );
243    }
244
245    #[test]
246    fn evaluate_struct_field() {
247        let mut ctx = array_session().create_execution_ctx();
248        let test_array = StructArray::from_fields(&[(
249            "a",
250            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(),
251        )])
252        .unwrap()
253        .into_array();
254
255        let expr = fill_null(get_item("a", root()), lit(0i32));
256        let result = test_array.apply(&expr).unwrap();
257
258        assert_eq!(
259            result.dtype(),
260            &DType::Primitive(PType::I32, Nullability::NonNullable)
261        );
262        assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 0, 3]), &mut ctx);
263    }
264
265    #[test]
266    fn evaluate_non_nullable_input() {
267        let mut ctx = array_session().create_execution_ctx();
268        let test_array = buffer![1i32, 2, 3].into_array();
269        let expr = fill_null(root(), lit(0i32));
270        let result = test_array.apply(&expr).unwrap();
271        assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 2, 3]), &mut ctx);
272    }
273
274    #[test]
275    fn test_display() {
276        let expr = fill_null(get_item("value", root()), lit(0i32));
277        assert_eq!(expr.to_string(), "vortex.fill_null($.value, 0i32)");
278    }
279}
vortex_array/scalar_fn/fns/fill_null/mod.rs

vortex_array/scalar_fn/fns/fill_null/
mod.rs
