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

vortex_array/scalar_fn/fns/fill_null/
mod.rs
