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

vortex_array/scalar_fn/fns/fill_null/
kernel.rs
