Skip to content

Commit 2b1337c

Browse files
gatesngatesn
authored
Numeric kernel (#3148)
1 parent f353817 commit 2b1337c

File tree

24 files changed

+472
-442
lines changed

24 files changed

+472
-442
lines changed

‎bench-vortex/src/bin/notimplemented.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,6 @@ fn compute_funcs(encodings: &[ArrayRef]) {
162162
"fill_forward",
163163
"fill_null",
164164
"scalar_at",
165-
"binary_numeric",
166165
"search_sorted",
167166
"slice",
168167
"take",
@@ -177,7 +176,6 @@ fn compute_funcs(encodings: &[ArrayRef]) {
177176
encoding.fill_forward_fn().is_some(),
178177
encoding.fill_null_fn().is_some(),
179178
encoding.scalar_at_fn().is_some(),
180-
encoding.binary_numeric_fn().is_some(),
181179
encoding.search_sorted_fn().is_some(),
182180
encoding.slice_fn().is_some(),
183181
encoding.take_fn().is_some(),

‎encodings/dict/src/compute/binary_numeric.rs‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
use vortex_array::arrays::ConstantArray;
2-
use vortex_array::compute::{BinaryNumericFn, binary_numeric};
3-
use vortex_array::{Array, ArrayRef};
2+
use vortex_array::compute::{NumericKernel, NumericKernelAdapter, numeric};
3+
use vortex_array::{Array, ArrayRef, register_kernel};
44
use vortex_error::VortexResult;
5-
use vortex_scalar::BinaryNumericOperator;
5+
use vortex_scalar::NumericOperator;
66

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

9-
impl BinaryNumericFn<&DictArray> for DictEncoding {
10-
fn binary_numeric(
9+
impl NumericKernel for DictEncoding {
10+
fn numeric(
1111
&self,
1212
array: &DictArray,
1313
rhs: &dyn Array,
14-
op: BinaryNumericOperator,
14+
op: NumericOperator,
1515
) -> VortexResult<Option<ArrayRef>> {
1616
// if we have more values than codes, it is faster to canonicalise first.
1717
if array.values().len() > array.codes().len() {
@@ -26,18 +26,20 @@ impl BinaryNumericFn<&DictArray> for DictEncoding {
2626
Ok(Some(
2727
DictArray::try_new(
2828
array.codes().clone(),
29-
binary_numeric(array.values(), &rhs_const_array, op)?,
29+
numeric(array.values(), &rhs_const_array, op)?,
3030
)?
3131
.into_array(),
3232
))
3333
}
3434
}
3535

36+
register_kernel!(NumericKernelAdapter(DictEncoding).lift());
37+
3638
#[cfg(test)]
3739
mod tests {
3840
use vortex_array::ArrayRef;
3941
use vortex_array::arrays::PrimitiveArray;
40-
use vortex_array::compute::conformance::binary_numeric::test_binary_numeric;
42+
use vortex_array::compute::conformance::binary_numeric::test_numeric;
4143
use vortex_array::compute::slice;
4244

4345
use crate::builders::dict_encode;
@@ -58,6 +60,6 @@ mod tests {
5860
#[test]
5961
fn test_dict_binary_numeric() {
6062
let array = sliced_dict_array();
61-
test_binary_numeric::<i32>(array)
63+
test_numeric::<i32>(array)
6264
}
6365
}

‎encodings/dict/src/compute/mod.rs‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ mod min_max;
88
mod optimize;
99

1010
use vortex_array::compute::{
11-
BinaryNumericFn, FillNullFn, FilterKernel, FilterKernelAdapter, IsConstantFn, IsSortedFn,
12-
LikeFn, MinMaxFn, OptimizeFn, ScalarAtFn, SliceFn, TakeFn, filter, scalar_at, slice, take,
11+
FillNullFn, FilterKernel, FilterKernelAdapter, IsConstantFn, IsSortedFn, LikeFn, MinMaxFn,
12+
OptimizeFn, ScalarAtFn, SliceFn, TakeFn, filter, scalar_at, slice, take,
1313
};
1414
use vortex_array::vtable::ComputeVTable;
1515
use vortex_array::{Array, ArrayRef, register_kernel};
@@ -20,10 +20,6 @@ use vortex_scalar::Scalar;
2020
use crate::{DictArray, DictEncoding};
2121

2222
impl ComputeVTable for DictEncoding {
23-
fn binary_numeric_fn(&self) -> Option<&dyn BinaryNumericFn<&dyn Array>> {
24-
Some(self)
25-
}
26-
2723
fn fill_null_fn(&self) -> Option<&dyn FillNullFn<&dyn Array>> {
2824
Some(self)
2925
}

‎encodings/fastlanes/benches/compute_between.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ mod primitive {
6363
use vortex_array::arrays::ConstantArray;
6464
use vortex_array::compute::StrictComparison::NonStrict;
6565
use vortex_array::compute::{
66-
BetweenOptions, BooleanOperator, Operator, between, binary_boolean, compare,
66+
BetweenOptions, BooleanOperator, Operator, between, boolean, compare,
6767
};
6868
use vortex_dtype::NativePType;
6969
use vortex_error::VortexExpect;
@@ -85,7 +85,7 @@ mod primitive {
8585
let arr = generate_primitive_array::<T>(&mut rng, len);
8686

8787
bencher.with_inputs(|| arr.clone()).bench_values(|arr| {
88-
binary_boolean(
88+
boolean(
8989
&compare(&arr, &ConstantArray::new(min, arr.len()), Operator::Gte)
9090
.vortex_expect(""),
9191
&compare(&arr, &ConstantArray::new(max, arr.len()), Operator::Lt).vortex_expect(""),
@@ -132,7 +132,7 @@ mod bitpack {
132132
use vortex_array::arrays::ConstantArray;
133133
use vortex_array::compute::StrictComparison::NonStrict;
134134
use vortex_array::compute::{
135-
BetweenOptions, BooleanOperator, Operator, between, binary_boolean, compare,
135+
BetweenOptions, BooleanOperator, Operator, between, boolean, compare,
136136
};
137137
use vortex_dtype::NativePType;
138138
use vortex_error::VortexExpect;
@@ -154,7 +154,7 @@ mod bitpack {
154154
let arr = generate_bit_pack_primitive_array::<T>(&mut rng, len);
155155

156156
bencher.with_inputs(|| arr.clone()).bench_values(|arr| {
157-
binary_boolean(
157+
boolean(
158158
&compare(&arr, &ConstantArray::new(min, arr.len()), Operator::Gte)
159159
.vortex_expect(""),
160160
&compare(&arr, &ConstantArray::new(max, arr.len()), Operator::Lt).vortex_expect(""),
@@ -200,7 +200,7 @@ mod alp {
200200
use vortex_array::arrays::ConstantArray;
201201
use vortex_array::compute::StrictComparison::NonStrict;
202202
use vortex_array::compute::{
203-
BetweenOptions, BooleanOperator, Operator, between, binary_boolean, compare,
203+
BetweenOptions, BooleanOperator, Operator, between, boolean, compare,
204204
};
205205
use vortex_dtype::NativePType;
206206
use vortex_error::VortexExpect;
@@ -222,7 +222,7 @@ mod alp {
222222
let arr = generate_alp_bit_pack_primitive_array::<T>(&mut rng, len);
223223

224224
bencher.with_inputs(|| arr.clone()).bench_values(|arr| {
225-
binary_boolean(
225+
boolean(
226226
&compare(&arr, &ConstantArray::new(min, arr.len()), Operator::Gte)
227227
.vortex_expect(""),
228228
&compare(&arr, &ConstantArray::new(max, arr.len()), Operator::Lt).vortex_expect(""),

‎encodings/runend/src/compute/binary_numeric.rs‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
use vortex_array::arrays::ConstantArray;
2-
use vortex_array::compute::{BinaryNumericFn, binary_numeric};
3-
use vortex_array::{Array, ArrayRef};
2+
use vortex_array::compute::{NumericKernel, NumericKernelAdapter, numeric};
3+
use vortex_array::{Array, ArrayRef, register_kernel};
44
use vortex_error::VortexResult;
5-
use vortex_scalar::BinaryNumericOperator;
5+
use vortex_scalar::NumericOperator;
66

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

9-
impl BinaryNumericFn<&RunEndArray> for RunEndEncoding {
10-
fn binary_numeric(
9+
impl NumericKernel for RunEndEncoding {
10+
fn numeric(
1111
&self,
1212
array: &RunEndArray,
1313
rhs: &dyn Array,
14-
op: BinaryNumericOperator,
14+
op: NumericOperator,
1515
) -> VortexResult<Option<ArrayRef>> {
1616
let Some(rhs_scalar) = rhs.as_constant() else {
1717
return Ok(None);
@@ -22,7 +22,7 @@ impl BinaryNumericFn<&RunEndArray> for RunEndEncoding {
2222
Ok(Some(
2323
RunEndArray::with_offset_and_length(
2424
array.ends().clone(),
25-
binary_numeric(array.values(), &rhs_const_array, op)?,
25+
numeric(array.values(), &rhs_const_array, op)?,
2626
array.offset(),
2727
array.len(),
2828
)?
@@ -31,11 +31,13 @@ impl BinaryNumericFn<&RunEndArray> for RunEndEncoding {
3131
}
3232
}
3333

34+
register_kernel!(NumericKernelAdapter(RunEndEncoding).lift());
35+
3436
#[cfg(test)]
3537
mod tests {
3638
use vortex_array::Array;
3739
use vortex_array::arrays::PrimitiveArray;
38-
use vortex_array::compute::conformance::binary_numeric::test_binary_numeric;
40+
use vortex_array::compute::conformance::binary_numeric::test_numeric;
3941

4042
use crate::RunEndArray;
4143

@@ -49,6 +51,6 @@ mod tests {
4951
#[test]
5052
fn test_runend_binary_numeric() {
5153
let array = ree_array().into_array();
52-
test_binary_numeric::<i32>(array)
54+
test_numeric::<i32>(array)
5355
}
5456
}

‎encodings/runend/src/compute/mod.rs‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,18 +12,13 @@ mod take_from;
1212

1313
use vortex_array::Array;
1414
use vortex_array::compute::{
15-
BinaryNumericFn, FillNullFn, InvertFn, IsSortedFn, MinMaxFn, ScalarAtFn, SliceFn, TakeFn,
16-
TakeFromFn,
15+
FillNullFn, InvertFn, IsSortedFn, MinMaxFn, ScalarAtFn, SliceFn, TakeFn, TakeFromFn,
1716
};
1817
use vortex_array::vtable::ComputeVTable;
1918

2019
use crate::RunEndEncoding;
2120

2221
impl ComputeVTable for RunEndEncoding {
23-
fn binary_numeric_fn(&self) -> Option<&dyn BinaryNumericFn<&dyn Array>> {
24-
Some(self)
25-
}
26-
2722
fn fill_null_fn(&self) -> Option<&dyn FillNullFn<&dyn Array>> {
2823
Some(self)
2924
}
@@ -61,7 +56,7 @@ impl ComputeVTable for RunEndEncoding {
6156
mod test {
6257
use vortex_array::Array;
6358
use vortex_array::arrays::PrimitiveArray;
64-
use vortex_array::compute::conformance::binary_numeric::test_binary_numeric;
59+
use vortex_array::compute::conformance::binary_numeric::test_numeric;
6560

6661
use crate::RunEndArray;
6762

@@ -75,6 +70,6 @@ mod test {
7570
#[test]
7671
fn test_runend_binary_numeric() {
7772
let array = ree_array().into_array();
78-
test_binary_numeric::<i32>(array)
73+
test_numeric::<i32>(array)
7974
}
8075
}

‎encodings/sparse/src/compute/binary_numeric.rs‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
use vortex_array::arrays::ConstantArray;
2-
use vortex_array::compute::{BinaryNumericFn, binary_numeric};
3-
use vortex_array::{Array, ArrayRef};
2+
use vortex_array::compute::{NumericKernel, NumericKernelAdapter, numeric};
3+
use vortex_array::{Array, ArrayRef, register_kernel};
44
use vortex_error::{VortexResult, vortex_err};
5-
use vortex_scalar::BinaryNumericOperator;
5+
use vortex_scalar::NumericOperator;
66

77
use crate::{SparseArray, SparseEncoding};
88

9-
impl BinaryNumericFn<&SparseArray> for SparseEncoding {
10-
fn binary_numeric(
9+
impl NumericKernel for SparseEncoding {
10+
fn numeric(
1111
&self,
1212
array: &SparseArray,
1313
rhs: &dyn Array,
14-
op: BinaryNumericOperator,
14+
op: NumericOperator,
1515
) -> VortexResult<Option<ArrayRef>> {
1616
let Some(rhs_scalar) = rhs.as_constant() else {
1717
return Ok(None);
@@ -20,7 +20,7 @@ impl BinaryNumericFn<&SparseArray> for SparseEncoding {
2020
let new_patches = array.patches().clone().map_values(|values| {
2121
let rhs_const_array = ConstantArray::new(rhs_scalar.clone(), values.len()).into_array();
2222

23-
binary_numeric(&values, &rhs_const_array, op)
23+
numeric(&values, &rhs_const_array, op)
2424
})?;
2525
let new_fill_value = array
2626
.fill_scalar()
@@ -33,3 +33,5 @@ impl BinaryNumericFn<&SparseArray> for SparseEncoding {
3333
))
3434
}
3535
}
36+
37+
register_kernel!(NumericKernelAdapter(SparseEncoding).lift());

‎encodings/sparse/src/compute/mod.rs‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use vortex_array::arrays::ConstantArray;
22
use vortex_array::compute::{
3-
BinaryNumericFn, FilterKernel, FilterKernelAdapter, InvertFn, ScalarAtFn, SearchSortedFn,
4-
SearchSortedUsizeFn, SliceFn, TakeFn,
3+
FilterKernel, FilterKernelAdapter, InvertFn, ScalarAtFn, SearchSortedFn, SearchSortedUsizeFn,
4+
SliceFn, TakeFn,
55
};
66
use vortex_array::vtable::ComputeVTable;
77
use vortex_array::{Array, ArrayRef, register_kernel};
@@ -18,10 +18,6 @@ mod slice;
1818
mod take;
1919

2020
impl ComputeVTable for SparseEncoding {
21-
fn binary_numeric_fn(&self) -> Option<&dyn BinaryNumericFn<&dyn Array>> {
22-
Some(self)
23-
}
24-
2521
fn invert_fn(&self) -> Option<&dyn InvertFn<&dyn Array>> {
2622
Some(self)
2723
}
@@ -77,7 +73,7 @@ register_kernel!(FilterKernelAdapter(SparseEncoding).lift());
7773
mod test {
7874
use rstest::{fixture, rstest};
7975
use vortex_array::arrays::PrimitiveArray;
80-
use vortex_array::compute::conformance::binary_numeric::test_binary_numeric;
76+
use vortex_array::compute::conformance::binary_numeric::test_numeric;
8177
use vortex_array::compute::conformance::mask::test_mask;
8278
use vortex_array::compute::{filter, try_cast};
8379
use vortex_array::validity::Validity;
@@ -138,7 +134,7 @@ mod test {
138134

139135
#[rstest]
140136
fn test_sparse_binary_numeric(array: ArrayRef) {
141-
test_binary_numeric::<i32>(array)
137+
test_numeric::<i32>(array)
142138
}
143139

144140
#[test]

‎vortex-array/src/arrays/chunked/compute/binary_numeric.rs‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,23 @@
11
use vortex_error::VortexResult;
2-
use vortex_scalar::BinaryNumericOperator;
2+
use vortex_scalar::NumericOperator;
33

44
use crate::arrays::{ChunkedArray, ChunkedEncoding};
5-
use crate::compute::{BinaryNumericFn, binary_numeric, slice};
6-
use crate::{Array, ArrayRef};
5+
use crate::compute::{NumericKernel, NumericKernelAdapter, numeric, slice};
6+
use crate::{Array, ArrayRef, register_kernel};
77

8-
impl BinaryNumericFn<&ChunkedArray> for ChunkedEncoding {
9-
fn binary_numeric(
8+
impl NumericKernel for ChunkedEncoding {
9+
fn numeric(
1010
&self,
1111
array: &ChunkedArray,
1212
rhs: &dyn Array,
13-
op: BinaryNumericOperator,
13+
op: NumericOperator,
1414
) -> VortexResult<Option<ArrayRef>> {
1515
let mut start = 0;
1616

1717
let mut new_chunks = Vec::with_capacity(array.nchunks());
1818
for chunk in array.non_empty_chunks() {
1919
let end = start + chunk.len();
20-
new_chunks.push(binary_numeric(chunk, &slice(rhs, start, end)?, op)?);
20+
new_chunks.push(numeric(chunk, &slice(rhs, start, end)?, op)?);
2121
start = end;
2222
}
2323

@@ -26,3 +26,5 @@ impl BinaryNumericFn<&ChunkedArray> for ChunkedEncoding {
2626
.map(Some)
2727
}
2828
}
29+
30+
register_kernel!(NumericKernelAdapter(ChunkedEncoding).lift());

‎vortex-array/src/arrays/chunked/compute/boolean.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@ use vortex_dtype::DType;
22
use vortex_error::VortexResult;
33

44
use crate::arrays::{ChunkedArray, ChunkedEncoding};
5-
use crate::compute::{BooleanKernel, BooleanKernelAdapter, BooleanOperator, binary_boolean, slice};
5+
use crate::compute::{BooleanKernel, BooleanKernelAdapter, BooleanOperator, boolean, slice};
66
use crate::{Array, ArrayRef, register_kernel};
77

88
impl BooleanKernel for ChunkedEncoding {
9-
fn binary_boolean(
9+
fn boolean(
1010
&self,
1111
lhs: &ChunkedArray,
1212
rhs: &dyn Array,
@@ -17,7 +17,7 @@ impl BooleanKernel for ChunkedEncoding {
1717

1818
for chunk in lhs.non_empty_chunks() {
1919
let sliced = slice(rhs, idx, idx + chunk.len())?;
20-
let result = binary_boolean(chunk, &sliced, op)?;
20+
let result = boolean(chunk, &sliced, op)?;
2121
chunks.push(result);
2222
idx += chunk.len();
2323
}
@@ -37,7 +37,7 @@ mod tests {
3737
use crate::array::Array;
3838
use crate::arrays::{BoolArray, BooleanBuffer, ChunkedArray};
3939
use crate::canonical::ToCanonical;
40-
use crate::compute::{BooleanOperator, binary_boolean};
40+
use crate::compute::{BooleanOperator, boolean};
4141

4242
#[test]
4343
fn test_bin_bool_chunked() {
@@ -51,7 +51,7 @@ mod tests {
5151
let chunked2 =
5252
ChunkedArray::try_new(vec![arr2, arr3], DType::Bool(Nullability::Nullable)).unwrap();
5353

54-
let result = binary_boolean(&chunked1, &chunked2, BooleanOperator::Or)
54+
let result = boolean(&chunked1, &chunked2, BooleanOperator::Or)
5555
.unwrap()
5656
.to_bool()
5757
.unwrap();

0 commit comments

Comments
 (0)