Skip to content

Commit f42310a

Browse files
authored
Port FillNullFn to FillNullKernel (#3208)
1 parent 2cfcfc0 commit f42310a

File tree

19 files changed

+172
-130
lines changed

19 files changed

+172
-130
lines changed

encodings/alp/src/alp_rd/compute/take.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ impl TakeFn<&ALPRDArray> for ALPRDEncoding {
2323
.transpose()?;
2424
let right_parts = fill_null(
2525
&take(array.right_parts(), indices)?,
26-
Scalar::new(array.right_parts().dtype().clone(), ScalarValue::from(0)),
26+
&Scalar::new(array.right_parts().dtype().clone(), ScalarValue::from(0)),
2727
)?;
2828

2929
Ok(ALPRDArray::try_new(

encodings/dict/src/compute/fill_null.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
use vortex_array::arrays::ConstantArray;
2-
use vortex_array::compute::{FillNullFn, Operator, compare, fill_null};
3-
use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
2+
use vortex_array::compute::{FillNullKernel, FillNullKernelAdapter, Operator, compare, fill_null};
3+
use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
44
use vortex_error::VortexResult;
55
use vortex_scalar::{Scalar, ScalarValue};
66

77
use crate::{DictArray, DictEncoding};
88

9-
impl FillNullFn<&DictArray> for DictEncoding {
10-
fn fill_null(&self, array: &DictArray, fill_value: Scalar) -> VortexResult<ArrayRef> {
9+
impl FillNullKernel for DictEncoding {
10+
fn fill_null(&self, array: &DictArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
1111
// If the fill value exists in the dictionary, we can simply rewrite the null codes to
1212
// point to the value.
1313
let found_fill_values = compare(
@@ -27,7 +27,7 @@ impl FillNullFn<&DictArray> for DictEncoding {
2727
// Now we rewrite the nullable codes to point at the fill value.
2828
let codes = fill_null(
2929
array.codes(),
30-
Scalar::new(
30+
&Scalar::new(
3131
array.codes().dtype().clone(),
3232
ScalarValue::from(first_fill_value),
3333
),
@@ -38,3 +38,5 @@ impl FillNullFn<&DictArray> for DictEncoding {
3838
Ok(DictArray::try_new(codes, values)?.into_array())
3939
}
4040
}
41+
42+
register_kernel!(FillNullKernelAdapter(DictEncoding).lift());

encodings/dict/src/compute/mod.rs

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,7 @@ mod like;
77
mod min_max;
88

99
use vortex_array::compute::{
10-
FillNullFn, FilterKernel, FilterKernelAdapter, MinMaxFn, ScalarAtFn, TakeFn, filter, scalar_at,
11-
take,
10+
FilterKernel, FilterKernelAdapter, MinMaxFn, ScalarAtFn, TakeFn, filter, scalar_at, take,
1211
};
1312
use vortex_array::vtable::ComputeVTable;
1413
use vortex_array::{Array, ArrayRef, register_kernel};
@@ -19,10 +18,6 @@ use vortex_scalar::Scalar;
1918
use crate::{DictArray, DictEncoding};
2019

2120
impl ComputeVTable for DictEncoding {
22-
fn fill_null_fn(&self) -> Option<&dyn FillNullFn<&dyn Array>> {
23-
Some(self)
24-
}
25-
2621
fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
2722
Some(self)
2823
}

encodings/fsst/src/compute/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ impl TakeFn<&FSSTArray> for FSSTEncoding {
3232
take(array.codes(), indices)?.as_::<VarBinArray>().clone(),
3333
fill_null(
3434
&take(array.uncompressed_lengths(), indices)?,
35-
Scalar::new(
35+
&Scalar::new(
3636
array.uncompressed_lengths_dtype().clone(),
3737
ScalarValue::from(0),
3838
),
Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1-
use vortex_array::compute::{FillNullFn, fill_null};
2-
use vortex_array::{Array, ArrayRef};
1+
use vortex_array::compute::{FillNullKernel, FillNullKernelAdapter, fill_null};
2+
use vortex_array::{Array, ArrayRef, register_kernel};
33
use vortex_error::VortexResult;
44
use vortex_scalar::Scalar;
55

66
use crate::{RunEndArray, RunEndEncoding};
77

8-
impl FillNullFn<&RunEndArray> for RunEndEncoding {
9-
fn fill_null(&self, array: &RunEndArray, fill_value: Scalar) -> VortexResult<ArrayRef> {
8+
impl FillNullKernel for RunEndEncoding {
9+
fn fill_null(&self, array: &RunEndArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
1010
Ok(RunEndArray::with_offset_and_length(
1111
array.ends().clone(),
1212
fill_null(array.values(), fill_value)?,
@@ -16,3 +16,5 @@ impl FillNullFn<&RunEndArray> for RunEndEncoding {
1616
.into_array())
1717
}
1818
}
19+
20+
register_kernel!(FillNullKernelAdapter(RunEndEncoding).lift());

encodings/runend/src/compute/mod.rs

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,16 +10,12 @@ pub(crate) mod take;
1010
mod take_from;
1111

1212
use vortex_array::Array;
13-
use vortex_array::compute::{FillNullFn, MinMaxFn, ScalarAtFn, TakeFn, TakeFromFn};
13+
use vortex_array::compute::{MinMaxFn, ScalarAtFn, TakeFn, TakeFromFn};
1414
use vortex_array::vtable::ComputeVTable;
1515

1616
use crate::RunEndEncoding;
1717

1818
impl ComputeVTable for RunEndEncoding {
19-
fn fill_null_fn(&self) -> Option<&dyn FillNullFn<&dyn Array>> {
20-
Some(self)
21-
}
22-
2319
fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
2420
Some(self)
2521
}

encodings/sparse/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ impl SparseArray {
178178
let non_top_mask = Mask::from_buffer(
179179
fill_null(
180180
&compare(array, &fill_array, Operator::NotEq)?,
181-
Scalar::bool(true, Nullability::NonNullable),
181+
&Scalar::bool(true, Nullability::NonNullable),
182182
)?
183183
.to_bool()?
184184
.boolean_buffer()

vortex-array/src/arrays/bool/compute/fill_null.rs

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,12 @@ use vortex_error::{VortexResult, vortex_err};
22
use vortex_scalar::Scalar;
33

44
use crate::arrays::{BoolArray, BoolEncoding, ConstantArray};
5-
use crate::compute::FillNullFn;
5+
use crate::compute::{FillNullKernel, FillNullKernelAdapter};
66
use crate::validity::Validity;
7-
use crate::{Array, ArrayRef, ToCanonical};
7+
use crate::{Array, ArrayRef, ToCanonical, register_kernel};
88

9-
impl FillNullFn<&BoolArray> for BoolEncoding {
10-
fn fill_null(&self, array: &BoolArray, fill_value: Scalar) -> VortexResult<ArrayRef> {
9+
impl FillNullKernel for BoolEncoding {
10+
fn fill_null(&self, array: &BoolArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
1111
let fill = fill_value
1212
.as_bool()
1313
.value()
@@ -19,7 +19,9 @@ impl FillNullFn<&BoolArray> for BoolEncoding {
1919
fill_value.dtype().nullability().into(),
2020
)
2121
.into_array(),
22-
Validity::AllInvalid => ConstantArray::new(fill_value, array.len()).into_array(),
22+
Validity::AllInvalid => {
23+
ConstantArray::new(fill_value.clone(), array.len()).into_array()
24+
}
2325
Validity::Array(v) => {
2426
let bool_buffer = if fill {
2527
array.boolean_buffer() | &!v.to_bool()?.boolean_buffer()
@@ -32,6 +34,8 @@ impl FillNullFn<&BoolArray> for BoolEncoding {
3234
}
3335
}
3436

37+
register_kernel!(FillNullKernelAdapter(BoolEncoding).lift());
38+
3539
#[cfg(test)]
3640
mod tests {
3741
use arrow_buffer::BooleanBuffer;
@@ -52,7 +56,7 @@ mod tests {
5256
BooleanBuffer::from_iter([true, true, false, false]),
5357
Validity::from_iter([true, false, true, false]),
5458
);
55-
let non_null_array = fill_null(&bool_array, fill_value.into())
59+
let non_null_array = fill_null(&bool_array, &fill_value.into())
5660
.unwrap()
5761
.to_bool()
5862
.unwrap();

vortex-array/src/arrays/bool/compute/mod.rs

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::Array;
22
use crate::arrays::BoolEncoding;
3-
use crate::compute::{FillNullFn, MinMaxFn, ScalarAtFn, TakeFn, UncompressedSizeFn};
3+
use crate::compute::{MinMaxFn, ScalarAtFn, TakeFn, UncompressedSizeFn};
44
use crate::vtable::ComputeVTable;
55

66
mod cast;
@@ -18,10 +18,6 @@ mod take;
1818
mod uncompressed_size;
1919

2020
impl ComputeVTable for BoolEncoding {
21-
fn fill_null_fn(&self) -> Option<&dyn FillNullFn<&dyn Array>> {
22-
Some(self)
23-
}
24-
2521
fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
2622
Some(self)
2723
}

vortex-array/src/arrays/bool/compute/take.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ impl TakeFn<&BoolArray> for BoolEncoding {
2323
)
2424
.into_array());
2525
}
26-
Mask::Values(_) => fill_null(indices, Scalar::from(0).cast(indices.dtype())?)?,
26+
Mask::Values(_) => fill_null(indices, &Scalar::from(0).cast(indices.dtype())?)?,
2727
};
2828
let indices_nulls_zeroed = indices_nulls_zeroed.to_primitive()?;
2929
let buffer = match_each_integer_ptype!(indices_nulls_zeroed.ptype(), |$I| {

0 commit comments

Comments
 (0)