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::IntoArray;
20use crate::arrays::Bool;
21use crate::arrays::Constant;
22use crate::arrays::Decimal;
23use crate::arrays::Primitive;
24use crate::arrays::ScalarFnArray;
25use crate::arrays::VarBinView;
26use crate::dtype::DType;
27use crate::expr::BoundExpression;
28use crate::optimizer::ArrayOptimizer;
29use crate::scalar::Scalar;
30use crate::scalar_fn::Arity;
31use crate::scalar_fn::ChildName;
32use crate::scalar_fn::EmptyOptions;
33use crate::scalar_fn::ExecutionArgs;
34use crate::scalar_fn::ReduceNode;
35use crate::scalar_fn::ReduceNodeValidity;
36use crate::scalar_fn::ScalarFnId;
37use crate::scalar_fn::ScalarFnVTable;
38use crate::scalar_fn::ScalarFnVTableExt;
39use crate::scalar_fn::is_not_null_node;
40
41/// An expression that replaces null values in the input with a fill value.
42#[derive(Clone)]
43pub struct FillNull;
44
45impl FillNull {
46    /// Creates a lazy operation that replaces null input values with `fill_value`.
47    ///
48    /// # Errors
49    ///
50    /// Returns an error if the children have different lengths or incompatible dtypes.
51    pub fn try_new(input: ArrayRef, fill_value: ArrayRef) -> VortexResult<ScalarFnArray> {
52        ScalarFnArray::try_new(FillNull.bind(EmptyOptions), [input, fill_value])
53    }
54
55    /// Returns whether fill_null can execute inputs of `dtype`.
56    ///
57    /// These are the dtypes whose canonical arrays have a [`FillNullKernel`], so update this list
58    /// whenever `fill_null_canonical` gains or loses one. Rewrites into fill_null, such as the
59    /// CASE WHEN simplification, check it first because other dtypes fail at execution time.
60    pub(crate) fn supports_dtype(dtype: &DType) -> bool {
61        matches!(
62            dtype,
63            DType::Bool(_)
64                | DType::Primitive(..)
65                | DType::Decimal(..)
66                | DType::Utf8(_)
67                | DType::Binary(_)
68        )
69    }
70}
71
72impl ScalarFnVTable for FillNull {
73    type Options = EmptyOptions;
74
75    fn id(&self) -> ScalarFnId {
76        static ID: CachedId = CachedId::new("vortex.fill_null");
77        *ID
78    }
79
80    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
81        Ok(Some(vec![]))
82    }
83
84    fn deserialize(
85        &self,
86        _metadata: &[u8],
87        _session: &VortexSession,
88    ) -> VortexResult<Self::Options> {
89        Ok(EmptyOptions)
90    }
91
92    fn arity(&self, _options: &Self::Options) -> Arity {
93        Arity::Exact(2)
94    }
95
96    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
97        match child_idx {
98            0 => ChildName::from("input"),
99            1 => ChildName::from("fill_value"),
100            _ => unreachable!("Invalid child index {} for FillNull expression", child_idx),
101        }
102    }
103
104    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
105        vortex_ensure!(
106            arg_dtypes[0].eq_ignore_nullability(&arg_dtypes[1]),
107            "fill_null requires input and fill value to have the same base type, got {} and {}",
108            arg_dtypes[0],
109            arg_dtypes[1]
110        );
111        // The result dtype takes the nullability of the fill value.
112        Ok(arg_dtypes[0]
113            .clone()
114            .with_nullability(arg_dtypes[1].nullability()))
115    }
116
117    fn execute(
118        &self,
119        _options: &Self::Options,
120        args: &dyn ExecutionArgs,
121        ctx: &mut ExecutionCtx,
122    ) -> VortexResult<ArrayRef> {
123        let input = args.get(0)?;
124        let fill_value = args.get(1)?;
125        // Parent kernels decline a fill that is not already constant. Materialize it here so a
126        // lazy scalar, such as a literal `ScalarFnArray`, can still fill nulls.
127        let fill_was_lazy = !fill_value.is::<Constant>();
128        let fill_value = if fill_was_lazy {
129            fill_value.execute_until::<Constant>(ctx)?
130        } else {
131            fill_value
132        };
133
134        let fill_scalar = fill_value
135            .as_constant()
136            .ok_or_else(|| vortex_err!("fill_null fill_value must be a constant/scalar"))?;
137
138        vortex_ensure!(
139            !fill_scalar.is_null(),
140            "fill_null requires a non-null fill value"
141        );
142
143        let Some(columnar) = input.as_opt::<AnyColumnar>() else {
144            if fill_was_lazy {
145                // Hand the rebuilt node back with a constant fill so the input's FillNull
146                // kernels get a chance before we decode the input.
147                return FillNull::try_new(input, fill_value)?
148                    .into_array()
149                    .optimize_ctx(ctx.session());
150            }
151            let input = input.execute::<ArrayRef>(ctx)?;
152            return Ok(FillNull::try_new(input, fill_value)?.into_array());
153        };
154
155        match columnar {
156            ColumnarView::Canonical(canonical) => fill_null_canonical(canonical, &fill_scalar, ctx),
157            ColumnarView::Constant(constant) => fill_null_constant(constant, &fill_scalar),
158        }
159    }
160
161    fn simplify(
162        &self,
163        _options: &Self::Options,
164        expr: &BoundExpression,
165    ) -> VortexResult<Option<BoundExpression>> {
166        if !expr.child(0).dtype().is_nullable() {
167            return Ok(Some(expr.child(0).clone()));
168        }
169
170        Ok(None)
171    }
172
173    fn validity<T: ReduceNode>(
174        &self,
175        _options: &Self::Options,
176        node: &T,
177    ) -> VortexResult<ReduceNodeValidity<T>> {
178        // After fill_null, the result validity depends on the fill value's nullability.
179        // If fill_value is non-nullable, the result is always valid.
180        Ok(ReduceNodeValidity::Reduced(is_not_null_node(
181            &node.child(1),
182        )?))
183    }
184
185    fn is_strict(&self, _options: &Self::Options) -> bool {
186        // This function replaces null input values instead of propagating them.
187        false
188    }
189
190    fn is_infallible(&self, _options: &Self::Options) -> bool {
191        true
192    }
193}
194
195/// Fill nulls on a canonical array by directly dispatching to the appropriate kernel.
196///
197/// Returns the filled array, or bails if no kernel is registered for the canonical type.
198fn fill_null_canonical(
199    canonical: CanonicalView<'_>,
200    fill_value: &Scalar,
201    ctx: &mut ExecutionCtx,
202) -> VortexResult<ArrayRef> {
203    let arr = canonical.to_array_ref();
204    if let Some(result) = short_circuit(&arr, fill_value)? {
205        return Ok(result);
206    }
207    match canonical {
208        CanonicalView::Bool(a) => <Bool as FillNullKernel>::fill_null(a, fill_value, ctx)?
209            .ok_or_else(|| vortex_err!("FillNullKernel for BoolArray returned None")),
210        CanonicalView::Primitive(a) => {
211            <Primitive as FillNullKernel>::fill_null(a, fill_value, ctx)?
212                .ok_or_else(|| vortex_err!("FillNullKernel for PrimitiveArray returned None"))
213        }
214        CanonicalView::Decimal(a) => <Decimal as FillNullKernel>::fill_null(a, fill_value, ctx)?
215            .ok_or_else(|| vortex_err!("FillNullKernel for DecimalArray returned None")),
216        CanonicalView::VarBinView(a) => {
217            <VarBinView as FillNullKernel>::fill_null(a, fill_value, ctx)?
218                .ok_or_else(|| vortex_err!("FillNullKernel for VarBinViewArray returned None"))
219        }
220        other => vortex_bail!(
221            "No FillNullKernel for canonical array {}",
222            other.to_array_ref().encoding_id()
223        ),
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use std::sync::Arc;
230
231    use vortex_buffer::buffer;
232    use vortex_error::VortexExpect;
233
234    use super::*;
235    use crate::Canonical;
236    use crate::IntoArray;
237    use crate::VortexSessionExecute;
238    use crate::array_session;
239    use crate::arrays::Chunked;
240    use crate::arrays::ChunkedArray;
241    use crate::arrays::Constant;
242    use crate::arrays::ConstantArray;
243    use crate::arrays::Primitive;
244    use crate::arrays::PrimitiveArray;
245    use crate::arrays::ScalarFnArray;
246    use crate::arrays::StructArray;
247    use crate::arrays::scalar_fn::ExactScalarFn;
248    use crate::assert_arrays_eq;
249    use crate::builders::builder_with_capacity_in;
250    use crate::builtins::ArrayBuiltins;
251    use crate::dtype::DType;
252    use crate::dtype::DecimalDType;
253    use crate::dtype::MapDType;
254    use crate::dtype::Nullability;
255    use crate::dtype::PType;
256    use crate::dtype::StructFields;
257    use crate::expr::fill_null;
258    use crate::expr::get_item;
259    use crate::expr::lit;
260    use crate::expr::root;
261    use crate::extension::datetime::Date;
262    use crate::extension::datetime::TimeUnit;
263    use crate::kernel::ExecuteParentKernel;
264    use crate::matcher::Matcher;
265    use crate::optimizer::rules::ArrayParentReduceRule;
266    use crate::scalar::Scalar;
267    use crate::scalar_fn::fns::binary::Binary;
268    use crate::scalar_fn::fns::operators::Operator;
269
270    /// A fill of `40 + 2` that is not a [`ConstantArray`] until it executes.
271    fn lazy_fill_42(len: usize) -> VortexResult<ArrayRef> {
272        Ok(ScalarFnArray::try_new(
273            Binary.bind(Operator::Add),
274            vec![
275                ConstantArray::new(40i32, len).into_array(),
276                ConstantArray::new(2i32, len).into_array(),
277            ],
278        )?
279        .into_array())
280    }
281
282    #[test]
283    fn dtype() {
284        let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
285        assert_eq!(
286            fill_null(root(), lit(0i32)).return_dtype(&dtype).unwrap(),
287            DType::Primitive(PType::I32, Nullability::NonNullable)
288        );
289    }
290
291    #[test]
292    fn replace_children() {
293        let expr = fill_null(root(), lit(0i32));
294        let replaced = expr
295            .with_children(vec![get_item("a", root()), lit(42i32)])
296            .vortex_expect("operation should succeed in test");
297        assert_eq!(replaced, fill_null(get_item("a", root()), lit(42i32)));
298    }
299
300    #[test]
301    fn evaluate() {
302        let mut ctx = array_session().create_execution_ctx();
303        let test_array =
304            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)])
305                .into_array();
306
307        let expr = fill_null(root(), lit(42i32));
308        let result = test_array.apply(&expr).unwrap();
309
310        assert_eq!(
311            result.dtype(),
312            &DType::Primitive(PType::I32, Nullability::NonNullable)
313        );
314        assert_arrays_eq!(
315            result,
316            PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]),
317            &mut ctx
318        );
319    }
320
321    #[test]
322    fn evaluate_lazy_constant_fill() -> VortexResult<()> {
323        let mut ctx = array_session().create_execution_ctx();
324        let input = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();
325        let fill = lazy_fill_42(input.len())?;
326        assert!(fill.as_constant().is_none());
327
328        let result = FillNull::try_new(input, fill)?
329            .into_array()
330            .execute::<ArrayRef>(&mut ctx)?;
331
332        assert_eq!(
333            result.dtype(),
334            &DType::Primitive(PType::I32, Nullability::NonNullable)
335        );
336        assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 42, 3]), &mut ctx);
337        Ok(())
338    }
339
340    #[test]
341    fn evaluate_lazy_fill_preserves_chunked_input() -> VortexResult<()> {
342        let mut ctx = array_session().create_execution_ctx();
343        let input = ChunkedArray::try_new(
344            vec![
345                PrimitiveArray::from_option_iter([Some(1i32), None]).into_array(),
346                PrimitiveArray::from_option_iter([Some(3), None]).into_array(),
347            ],
348            DType::Primitive(PType::I32, Nullability::Nullable),
349        )?
350        .into_array();
351        let fill = lazy_fill_42(input.len())?;
352
353        let result = FillNull::try_new(input, fill)?
354            .into_array()
355            .execute::<ArrayRef>(&mut ctx)?;
356
357        assert!(result.is::<Chunked>());
358        assert_arrays_eq!(
359            result,
360            PrimitiveArray::from_iter([1i32, 42, 3, 42]),
361            &mut ctx
362        );
363        Ok(())
364    }
365
366    #[test]
367    fn non_constant_fill_declines_and_errors() -> VortexResult<()> {
368        let mut ctx = array_session().create_execution_ctx();
369        let varying = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
370
371        let constant_input = ConstantArray::new(Scalar::null_native::<i32>(), 3).into_array();
372        let constant_parent =
373            FillNull::try_new(constant_input.clone(), varying.clone())?.into_array();
374        let reduced = FillNullReduceAdaptor(Constant).reduce_parent(
375            constant_input
376                .as_opt::<Constant>()
377                .vortex_expect("constant input"),
378            ExactScalarFn::<FillNull>::try_match(&constant_parent)
379                .vortex_expect("fill_null parent"),
380            0,
381        )?;
382        assert!(reduced.is_none());
383
384        let primitive_input =
385            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();
386        let primitive_parent = FillNull::try_new(primitive_input.clone(), varying)?.into_array();
387        let executed = FillNullExecuteAdaptor(Primitive).execute_parent(
388            primitive_input
389                .as_opt::<Primitive>()
390                .vortex_expect("primitive input"),
391            ExactScalarFn::<FillNull>::try_match(&primitive_parent)
392                .vortex_expect("fill_null parent"),
393            0,
394            &mut ctx,
395        )?;
396        assert!(executed.is_none());
397
398        let err = primitive_parent.execute::<ArrayRef>(&mut ctx).unwrap_err();
399        assert!(
400            err.to_string()
401                .contains("fill_null fill_value must be a constant/scalar")
402        );
403        Ok(())
404    }
405
406    #[test]
407    fn evaluate_struct_field() {
408        let mut ctx = array_session().create_execution_ctx();
409        let test_array = StructArray::from_fields(&[(
410            "a",
411            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(),
412        )])
413        .unwrap()
414        .into_array();
415
416        let expr = fill_null(get_item("a", root()), lit(0i32));
417        let result = test_array.apply(&expr).unwrap();
418
419        assert_eq!(
420            result.dtype(),
421            &DType::Primitive(PType::I32, Nullability::NonNullable)
422        );
423        assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 0, 3]), &mut ctx);
424    }
425
426    #[test]
427    fn evaluate_non_nullable_input() {
428        let mut ctx = array_session().create_execution_ctx();
429        let test_array = buffer![1i32, 2, 3].into_array();
430        let expr = fill_null(root(), lit(0i32));
431        let result = test_array.apply(&expr).unwrap();
432        assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 2, 3]), &mut ctx);
433    }
434
435    /// Inputs with both valid and null rows execute exactly for the dtypes that `supports_dtype`
436    /// accepts, so adding a kernel without updating `supports_dtype` fails here.
437    #[test]
438    fn supports_dtype_matches_kernels() -> VortexResult<()> {
439        let mut ctx = array_session().create_execution_ctx();
440        let i32_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
441        let dtypes = [
442            DType::Bool(Nullability::Nullable),
443            DType::Primitive(PType::I32, Nullability::Nullable),
444            DType::Decimal(DecimalDType::new(19, 2), Nullability::Nullable),
445            DType::Utf8(Nullability::Nullable),
446            DType::Binary(Nullability::Nullable),
447            DType::List(Arc::new(i32_dtype.clone()), Nullability::Nullable),
448            DType::FixedSizeList(Arc::new(i32_dtype.clone()), 2, Nullability::Nullable),
449            DType::Struct(
450                StructFields::from_iter([("value", i32_dtype.clone())]),
451                Nullability::Nullable,
452            ),
453            DType::Map(
454                MapDType::try_new(i32_dtype.clone(), i32_dtype, false)?,
455                Nullability::Nullable,
456            ),
457            DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased()),
458        ];
459
460        let supported = [
461            true, true, true, true, true, false, false, false, false, false,
462        ];
463        for (dtype, supported) in dtypes.into_iter().zip(supported) {
464            assert_eq!(FillNull::supports_dtype(&dtype), supported, "{dtype}");
465            let fill = Scalar::default_value(&dtype.as_nonnullable());
466            let mut builder = builder_with_capacity_in(&dtype, 2, ctx.allocator());
467            builder.append_scalar(&fill.cast(&dtype)?)?;
468            builder.append_null();
469            let filled = builder
470                .finish()
471                .fill_null(fill.clone())?
472                .execute::<Canonical>(&mut ctx);
473
474            if supported {
475                let filled = filled?.into_array();
476                assert_eq!(filled.dtype(), fill.dtype(), "{dtype}");
477                assert_eq!(filled.execute_scalar(1, &mut ctx)?, fill, "{dtype}");
478                assert_arrays_eq!(filled, ConstantArray::new(fill, 2), &mut ctx);
479            } else {
480                assert!(filled.is_err(), "{dtype} executes but is not supported");
481            }
482        }
483        Ok(())
484    }
485
486    #[test]
487    fn test_display() {
488        let expr = fill_null(get_item("value", root()), lit(0i32));
489        assert_eq!(expr.to_string(), "vortex.fill_null($.value, 0i32)");
490    }
491}
vortex_array/scalar_fn/fns/fill_null/mod.rs

vortex_array/scalar_fn/fns/fill_null/
mod.rs
