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;
6use vortex_error::vortex_ensure;
7
8use crate::ArrayRef;
9use crate::ExecutionCtx;
10use crate::IntoArray;
11use crate::arrays::ConstantArray;
12use crate::arrays::ScalarFnVTable;
13use crate::arrays::scalar_fn::ExactScalarFn;
14use crate::arrays::scalar_fn::ScalarFnArrayView;
15use crate::builtins::ArrayBuiltins;
16use crate::kernel::ExecuteParentKernel;
17use crate::optimizer::rules::ArrayParentReduceRule;
18use crate::scalar::Scalar;
19use crate::scalar_fn::fns::fill_null::FillNull as FillNullExpr;
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: &ArrayRef,
59    fill_value: &Scalar,
60) -> VortexResult<Option<ArrayRef>> {
61    vortex_ensure!(
62        !fill_value.is_null(),
63        "fill_null requires a non-null fill value"
64    );
65
66    // If the array has no nulls, fill_null is a no-op (just cast for nullability).
67    if !array.dtype().is_nullable() || array.all_valid()? {
68        return array.to_array().cast(fill_value.dtype().clone()).map(Some);
69    }
70
71    // If all values are null, replace the entire array with the fill value.
72    if array.all_invalid()? {
73        return Ok(Some(
74            ConstantArray::new(fill_value.clone(), array.len()).into_array(),
75        ));
76    }
77
78    Ok(None)
79}
80
81/// Fill null on a [`ConstantArray`] by replacing null scalars with the fill value,
82/// or casting non-null scalars to the fill value's dtype.
83pub(crate) fn fill_null_constant(
84    array: &ConstantArray,
85    fill_value: &Scalar,
86) -> VortexResult<ArrayRef> {
87    let scalar = if array.scalar().is_null() {
88        fill_value.clone()
89    } else {
90        array.scalar().cast(fill_value.dtype())?
91    };
92    Ok(ConstantArray::new(scalar, array.len()).into_array())
93}
94
95/// Adaptor that wraps a [`FillNullReduce`] impl as an [`ArrayParentReduceRule`].
96#[derive(Default, Debug)]
97pub struct FillNullReduceAdaptor<V>(pub V);
98
99impl<V> ArrayParentReduceRule<V> for FillNullReduceAdaptor<V>
100where
101    V: FillNullReduce,
102{
103    type Parent = ExactScalarFn<FillNullExpr>;
104
105    fn reduce_parent(
106        &self,
107        array: &V::Array,
108        parent: ScalarFnArrayView<'_, FillNullExpr>,
109        child_idx: usize,
110    ) -> VortexResult<Option<ArrayRef>> {
111        // Only process the input child (index 0), not the fill_value child (index 1).
112        if child_idx != 0 {
113            return Ok(None);
114        }
115        let scalar_fn_array = parent
116            .as_opt::<ScalarFnVTable>()
117            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
118        let fill_value = scalar_fn_array.children()[1]
119            .as_constant()
120            .vortex_expect("fill_null fill_value must be constant");
121        let arr = array.to_array();
122        if let Some(result) = precondition(&arr, &fill_value)? {
123            return Ok(Some(result));
124        }
125        <V as FillNullReduce>::fill_null(array, &fill_value)
126    }
127}
128
129/// Adaptor that wraps a [`FillNullKernel`] impl as an [`ExecuteParentKernel`].
130#[derive(Default, Debug)]
131pub struct FillNullExecuteAdaptor<V>(pub V);
132
133impl<V> ExecuteParentKernel<V> for FillNullExecuteAdaptor<V>
134where
135    V: FillNullKernel,
136{
137    type Parent = ExactScalarFn<FillNullExpr>;
138
139    fn execute_parent(
140        &self,
141        array: &V::Array,
142        parent: ScalarFnArrayView<'_, FillNullExpr>,
143        child_idx: usize,
144        ctx: &mut ExecutionCtx,
145    ) -> VortexResult<Option<ArrayRef>> {
146        // Only process the input child (index 0), not the fill_value child (index 1).
147        if child_idx != 0 {
148            return Ok(None);
149        }
150        let scalar_fn_array = parent
151            .as_opt::<ScalarFnVTable>()
152            .vortex_expect("ExactScalarFn matcher confirmed ScalarFnArray");
153        let fill_value = scalar_fn_array.children()[1]
154            .as_constant()
155            .vortex_expect("fill_null fill_value must be constant");
156        let arr = array.to_array();
157        if let Some(result) = precondition(&arr, &fill_value)? {
158            return Ok(Some(result));
159        }
160        <V as FillNullKernel>::fill_null(array, &fill_value, ctx)
161    }
162}
vortex_array/scalar_fn/fns/fill_null/kernel.rs

vortex_array/scalar_fn/fns/fill_null/
kernel.rs
