vortex_array/expr/exprs/fill_null/
kernel.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexExpect;
5use vortex_error::VortexResult;
6
7use crate::Array;
8use crate::ArrayRef;
9use crate::ExecutionCtx;
10use crate::IntoArray;
11use crate::arrays::ConstantArray;
12use crate::arrays::ExactScalarFn;
13use crate::arrays::ScalarFnArrayView;
14use crate::arrays::ScalarFnVTable;
15use crate::builtins::ArrayBuiltins;
16use crate::expr::FillNull as FillNullExpr;
17use crate::kernel::ExecuteParentKernel;
18use crate::optimizer::rules::ArrayParentReduceRule;
19use crate::scalar::Scalar;
20use crate::vtable::VTable;
21
22/// Fill nulls in an array with a scalar value without reading buffers.
23///
24/// This trait is for fill_null implementations that can operate purely on array metadata
25/// and structure without needing to read or execute on the underlying buffers.
26/// Implementations should return `None` if the operation requires buffer access.
27///
28/// # Preconditions
29///
30/// The fill value is guaranteed to be non-null. The array is guaranteed to have mixed
31/// validity (neither all-valid nor all-invalid).
32pub trait FillNullReduce: VTable {
33    fn fill_null(array: &Self::Array, fill_value: &Scalar) -> VortexResult<Option<ArrayRef>>;
34}
35
36/// Fill nulls in an array with a scalar value, potentially reading buffers.
37///
38/// Unlike [`FillNullReduce`], this trait is for fill_null implementations that may need
39/// to read and execute on the underlying buffers to produce the result.
40///
41/// # Preconditions
42///
43/// The fill value is guaranteed to be non-null. The array is guaranteed to have mixed
44/// validity (neither all-valid nor all-invalid).
45pub trait FillNullKernel: VTable {
46    fn fill_null(
47        array: &Self::Array,
48        fill_value: &Scalar,
49        ctx: &mut ExecutionCtx,
50    ) -> VortexResult<Option<ArrayRef>>;
51}
52
53/// Common preconditions for fill_null operations that apply to all arrays.
54///
55/// Returns `Some(result)` if the precondition short-circuits the fill_null operation,
56/// or `None` if fill_null should proceed with the encoding-specific implementation.
57pub(super) fn precondition(
58    array: &dyn Array,
59    fill_value: &Scalar,
60) -> VortexResult<Option<ArrayRef>> {
61    // If the array has no nulls, fill_null is a no-op (just cast for nullability).
62    if !array.dtype().is_nullable() || array.all_valid()? {
63        return array.to_array().cast(fill_value.dtype().clone()).map(Some);
64    }
65
66    // If all values are null, replace the entire array with the fill value.
67    if array.all_invalid()? {
68        return Ok(Some(
69            ConstantArray::new(fill_value.clone(), array.len()).into_array(),
70        ));
71    }
72
73    Ok(None)
74}
75
76/// Fill null on a [`ConstantArray`] by replacing null scalars with the fill value,
77/// or casting non-null scalars to the fill value's dtype.
78pub(crate) fn fill_null_constant(
79    array: &ConstantArray,
80    fill_value: &Scalar,
81) -> VortexResult<ArrayRef> {
82    let scalar = if array.scalar().is_null() {
83        fill_value.clone()
84    } else {
85        array.scalar().cast(fill_value.dtype())?
86    };
87    Ok(ConstantArray::new(scalar, array.len()).into_array())
88}
89
90/// Adaptor that wraps a [`FillNullReduce`] impl as an [`ArrayParentReduceRule`].
91#[derive(Default, Debug)]
92pub struct FillNullReduceAdaptor<V>(pub V);
93
94impl<V> ArrayParentReduceRule<V> for FillNullReduceAdaptor<V>
95where
96    V: FillNullReduce,
97{
98    type Parent = ExactScalarFn<FillNullExpr>;
99
100    fn reduce_parent(
101        &self,
102        array: &V::Array,
103        parent: ScalarFnArrayView<'_, FillNullExpr>,
104        child_idx: usize,
105    ) -> VortexResult<Option<ArrayRef>> {
106        // Only process the input child (index 0), not the fill_value child (index 1).
107        if child_idx != 0 {
108            return Ok(None);
109        }
110        let scalar_fn_array = parent
111            .as_opt::<ScalarFnVTable>()
112            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
113        let fill_value = scalar_fn_array.children()[1]
114            .as_constant()
115            .vortex_expect("fill_null fill_value must be constant");
116        if let Some(result) = precondition(&**array, &fill_value)? {
117            return Ok(Some(result));
118        }
119        <V as FillNullReduce>::fill_null(array, &fill_value)
120    }
121}
122
123/// Adaptor that wraps a [`FillNullKernel`] impl as an [`ExecuteParentKernel`].
124#[derive(Default, Debug)]
125pub struct FillNullExecuteAdaptor<V>(pub V);
126
127impl<V> ExecuteParentKernel<V> for FillNullExecuteAdaptor<V>
128where
129    V: FillNullKernel,
130{
131    type Parent = ExactScalarFn<FillNullExpr>;
132
133    fn execute_parent(
134        &self,
135        array: &V::Array,
136        parent: ScalarFnArrayView<'_, FillNullExpr>,
137        child_idx: usize,
138        ctx: &mut ExecutionCtx,
139    ) -> VortexResult<Option<ArrayRef>> {
140        // Only process the input child (index 0), not the fill_value child (index 1).
141        if child_idx != 0 {
142            return Ok(None);
143        }
144        let scalar_fn_array = parent
145            .as_opt::<ScalarFnVTable>()
146            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
147        let fill_value = scalar_fn_array.children()[1]
148            .as_constant()
149            .vortex_expect("fill_null fill_value must be constant");
150        if let Some(result) = precondition(&**array, &fill_value)? {
151            return Ok(Some(result));
152        }
153        <V as FillNullKernel>::fill_null(array, &fill_value, ctx)
154    }
155}
vortex_array/expr/exprs/fill_null/kernel.rs

vortex_array/expr/exprs/fill_null/
kernel.rs
