Skip to content

Commit ec85ffe

Browse files
gatesngatesn
authored
IsSorted AggregateFn (#7029)
Port IsSorted compute function over to an AggregateFn --------- Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent 3cf3a17 commit ec85ffe

70 files changed

Lines changed: 1323 additions & 1385 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎encodings/alp/benches/alp_compress.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,12 @@ use vortex_alp::decompress_into_array;
1515
use vortex_array::LEGACY_SESSION;
1616
use vortex_array::VortexSessionExecute;
1717
use vortex_array::arrays::PrimitiveArray;
18-
use vortex_array::compute::warm_up_vtables;
1918
use vortex_array::dtype::NativePType;
2019
use vortex_array::validity::Validity;
2120
use vortex_buffer::Buffer;
2221
use vortex_buffer::buffer;
2322

2423
fn main() {
25-
warm_up_vtables();
2624
divan::main();
2725
}
2826

‎encodings/fastlanes/benches/bitpacking_take.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,12 @@ use vortex_array::LEGACY_SESSION;
1515
use vortex_array::RecursiveCanonical;
1616
use vortex_array::VortexSessionExecute;
1717
use vortex_array::arrays::PrimitiveArray;
18-
use vortex_array::compute::warm_up_vtables;
1918
use vortex_array::validity::Validity;
2019
use vortex_buffer::Buffer;
2120
use vortex_buffer::buffer;
2221
use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width;
2322

2423
fn main() {
25-
warm_up_vtables();
2624
divan::main();
2725
}
2826

‎encodings/fastlanes/benches/canonicalize_bench.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,12 @@ use vortex_array::VortexSessionExecute;
1313
use vortex_array::arrays::ChunkedArray;
1414
use vortex_array::builders::ArrayBuilder;
1515
use vortex_array::builders::PrimitiveBuilder;
16-
use vortex_array::compute::warm_up_vtables;
1716
use vortex_array::session::ArraySession;
1817
use vortex_error::VortexExpect;
1918
use vortex_fastlanes::bitpack_compress::test_harness::make_array;
2019
use vortex_session::VortexSession;
2120

2221
fn main() {
23-
warm_up_vtables();
2422
divan::main();
2523
}
2624

‎encodings/fastlanes/benches/compute_between.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,11 @@ use vortex_array::ArrayRef;
1111
use vortex_array::IntoArray;
1212
use vortex_array::ToCanonical;
1313
use vortex_array::arrays::PrimitiveArray;
14-
use vortex_array::compute::warm_up_vtables;
1514
use vortex_array::dtype::NativePType;
1615
use vortex_error::VortexExpect;
1716
use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width;
1817

1918
fn main() {
20-
warm_up_vtables();
2119
divan::main();
2220
}
2321

‎encodings/fastlanes/public-api.lock‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -420,12 +420,6 @@ impl vortex_array::arrays::slice::SliceReduce for vortex_fastlanes::FoR
420420

421421
pub fn vortex_fastlanes::FoR::slice(array: &Self::Array, range: core::ops::range::Range<usize>) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
422422

423-
impl vortex_array::compute::is_sorted::IsSortedKernel for vortex_fastlanes::FoR
424-
425-
pub fn vortex_fastlanes::FoR::is_sorted(&self, array: &vortex_fastlanes::FoRArray) -> vortex_error::VortexResult<core::option::Option<bool>>
426-
427-
pub fn vortex_fastlanes::FoR::is_strict_sorted(&self, array: &vortex_fastlanes::FoRArray) -> vortex_error::VortexResult<core::option::Option<bool>>
428-
429423
impl vortex_array::scalar_fn::fns::binary::compare::CompareKernel for vortex_fastlanes::FoR
430424

431425
pub fn vortex_fastlanes::FoR::compare(lhs: &vortex_fastlanes::FoRArray, rhs: &vortex_array::array::ArrayRef, operator: vortex_array::scalar_fn::fns::operators::CompareOperator, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>

‎encodings/fastlanes/src/for/compute/is_sorted.rs‎

Lines changed: 48 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,19 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4+
use vortex_array::ArrayRef;
5+
use vortex_array::ExecutionCtx;
46
use vortex_array::IntoArray;
57
use vortex_array::ToCanonical;
6-
use vortex_array::compute::IsSortedKernel;
7-
use vortex_array::compute::IsSortedKernelAdapter;
8-
use vortex_array::compute::is_sorted;
9-
use vortex_array::compute::is_strict_sorted;
10-
use vortex_array::register_kernel;
8+
use vortex_array::aggregate_fn::AggregateFnRef;
9+
use vortex_array::aggregate_fn::fns::is_sorted::IsSorted;
10+
use vortex_array::aggregate_fn::fns::is_sorted::is_sorted;
11+
use vortex_array::aggregate_fn::fns::is_sorted::is_strict_sorted;
12+
use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
13+
use vortex_array::scalar::Scalar;
1114
use vortex_error::VortexResult;
1215

1316
use crate::FoR;
14-
use crate::FoRArray;
1517

1618
/// FoR can express sortedness directly on its encoded form.
1719
///
@@ -71,92 +73,107 @@ use crate::FoRArray;
7173
/// Addition is order-preserving, so all the wrapped values preserve their order and they're all
7274
/// represented as unsigned values larger than 127 so they also preserve their order with the
7375
/// unwrapped values.
74-
impl IsSortedKernel for FoR {
75-
fn is_sorted(&self, array: &FoRArray) -> VortexResult<Option<bool>> {
76-
let encoded = array.encoded().to_primitive();
77-
is_sorted(
78-
&encoded
79-
.reinterpret_cast(encoded.ptype().to_unsigned())
80-
.into_array(),
81-
)
82-
}
76+
#[derive(Debug)]
77+
pub(crate) struct FoRIsSortedKernel;
78+
79+
impl DynAggregateKernel for FoRIsSortedKernel {
80+
fn aggregate(
81+
&self,
82+
aggregate_fn: &AggregateFnRef,
83+
batch: &ArrayRef,
84+
ctx: &mut ExecutionCtx,
85+
) -> VortexResult<Option<Scalar>> {
86+
let Some(options) = aggregate_fn.as_opt::<IsSorted>() else {
87+
return Ok(None);
88+
};
89+
90+
let Some(array) = batch.as_opt::<FoR>() else {
91+
return Ok(None);
92+
};
8393

84-
fn is_strict_sorted(&self, array: &FoRArray) -> VortexResult<Option<bool>> {
8594
let encoded = array.encoded().to_primitive();
86-
is_strict_sorted(
87-
&encoded
88-
.reinterpret_cast(encoded.ptype().to_unsigned())
89-
.into_array(),
90-
)
95+
let unsigned_array = encoded
96+
.reinterpret_cast(encoded.ptype().to_unsigned())
97+
.into_array();
98+
99+
let result = if options.strict {
100+
is_strict_sorted(&unsigned_array, ctx)?
101+
} else {
102+
is_sorted(&unsigned_array, ctx)?
103+
};
104+
105+
Ok(Some(IsSorted::make_partial(batch, result, options.strict)?))
91106
}
92107
}
93108

94-
register_kernel!(IsSortedKernelAdapter(FoR).lift());
95-
96109
#[cfg(test)]
97110
mod test {
98111
use vortex_array::IntoArray;
112+
use vortex_array::LEGACY_SESSION;
113+
use vortex_array::VortexSessionExecute;
114+
use vortex_array::aggregate_fn::fns::is_sorted::is_sorted;
99115
use vortex_array::arrays::PrimitiveArray;
100-
use vortex_array::compute::is_sorted;
101116
use vortex_array::validity::Validity;
102117
use vortex_buffer::buffer;
103118

104119
use crate::FoRArray;
105120

106121
#[test]
107122
fn test_sorted() {
123+
let mut ctx = LEGACY_SESSION.create_execution_ctx();
124+
108125
let a = PrimitiveArray::new(buffer![-1, 0, i8::MAX], Validity::NonNullable);
109126
let b = FoRArray::encode(a).unwrap();
110127
assert!(
111-
is_sorted(&b.clone().into_array()).unwrap().unwrap(),
128+
is_sorted(&b.clone().into_array(), &mut ctx).unwrap(),
112129
"{}",
113130
b.encoded().display_values()
114131
);
115132

116133
let a = PrimitiveArray::new(buffer![i8::MIN, 0, i8::MAX], Validity::NonNullable);
117134
let b = FoRArray::encode(a).unwrap();
118135
assert!(
119-
is_sorted(&b.clone().into_array()).unwrap().unwrap(),
136+
is_sorted(&b.clone().into_array(), &mut ctx).unwrap(),
120137
"{}",
121138
b.encoded().display_values()
122139
);
123140

124141
let a = PrimitiveArray::new(buffer![i8::MIN, 0, 30, 127], Validity::NonNullable);
125142
let b = FoRArray::encode(a).unwrap();
126143
assert!(
127-
is_sorted(&b.clone().into_array()).unwrap().unwrap(),
144+
is_sorted(&b.clone().into_array(), &mut ctx).unwrap(),
128145
"{}",
129146
b.encoded().display_values()
130147
);
131148

132149
let a = PrimitiveArray::new(buffer![i8::MIN, -3, -1], Validity::NonNullable);
133150
let b = FoRArray::encode(a).unwrap();
134151
assert!(
135-
is_sorted(&b.clone().into_array()).unwrap().unwrap(),
152+
is_sorted(&b.clone().into_array(), &mut ctx).unwrap(),
136153
"{}",
137154
b.encoded().display_values()
138155
);
139156

140157
let a = PrimitiveArray::new(buffer![-10, -3, -1], Validity::NonNullable);
141158
let b = FoRArray::encode(a).unwrap();
142159
assert!(
143-
is_sorted(&b.clone().into_array()).unwrap().unwrap(),
160+
is_sorted(&b.clone().into_array(), &mut ctx).unwrap(),
144161
"{}",
145162
b.encoded().display_values()
146163
);
147164

148165
let a = PrimitiveArray::new(buffer![-10, -11, -1], Validity::NonNullable);
149166
let b = FoRArray::encode(a).unwrap();
150167
assert!(
151-
!is_sorted(&b.clone().into_array()).unwrap().unwrap(),
168+
!is_sorted(&b.clone().into_array(), &mut ctx).unwrap(),
152169
"{}",
153170
b.encoded().display_values()
154171
);
155172

156173
let a = PrimitiveArray::new(buffer![-10, i8::MIN, -1], Validity::NonNullable);
157174
let b = FoRArray::encode(a).unwrap();
158175
assert!(
159-
!is_sorted(&b.clone().into_array()).unwrap().unwrap(),
176+
!is_sorted(&b.clone().into_array(), &mut ctx).unwrap(),
160177
"{}",
161178
b.encoded().display_values()
162179
);

‎encodings/fastlanes/src/for/compute/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
mod cast;
55
mod compare;
66
pub(crate) mod is_constant;
7-
mod is_sorted;
7+
pub(crate) mod is_sorted;
88

99
use vortex_array::ArrayRef;
1010
use vortex_array::DynArray;

‎encodings/fastlanes/src/lib.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@ pub(crate) const FL_CHUNK_SIZE: usize = 1024;
1818

1919
use bitpacking::compute::is_constant::BitPackedIsConstantKernel;
2020
use r#for::compute::is_constant::FoRIsConstantKernel;
21+
use r#for::compute::is_sorted::FoRIsSortedKernel;
2122
use vortex_array::aggregate_fn::AggregateFnVTable;
2223
use vortex_array::aggregate_fn::fns::is_constant::IsConstant;
24+
use vortex_array::aggregate_fn::fns::is_sorted::IsSorted;
2325
use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
2426
use vortex_array::session::ArraySessionExt;
2527
use vortex_session::VortexSession;
@@ -42,6 +44,11 @@ pub fn initialize(session: &mut VortexSession) {
4244
Some(IsConstant.id()),
4345
&FoRIsConstantKernel,
4446
);
47+
session.aggregate_fns().register_aggregate_kernel(
48+
FoR::ID,
49+
Some(IsSorted.id()),
50+
&FoRIsSortedKernel,
51+
);
4552
}
4653

4754
#[cfg(test)]

‎encodings/fsst/benches/chunked_dict_fsst_builder.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,13 @@ use vortex_array::IntoArray;
1010
use vortex_array::VortexSessionExecute;
1111
use vortex_array::arrays::ChunkedArray;
1212
use vortex_array::builders::builder_with_capacity;
13-
use vortex_array::compute::warm_up_vtables;
1413
use vortex_array::dtype::NativePType;
1514
use vortex_array::session::ArraySession;
1615
use vortex_error::VortexExpect;
1716
use vortex_fsst::test_utils::gen_dict_fsst_test_data;
1817
use vortex_session::VortexSession;
1918

2019
fn main() {
21-
warm_up_vtables();
2220
divan::main();
2321
}
2422

‎encodings/fsst/benches/fsst_compress.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ use vortex_array::arrays::VarBinArray;
1919
use vortex_array::builders::ArrayBuilder;
2020
use vortex_array::builders::VarBinViewBuilder;
2121
use vortex_array::builtins::ArrayBuiltins;
22-
use vortex_array::compute::warm_up_vtables;
2322
use vortex_array::dtype::DType;
2423
use vortex_array::dtype::Nullability;
2524
use vortex_array::scalar::Scalar;
@@ -30,7 +29,6 @@ use vortex_fsst::fsst_train_compressor;
3029
use vortex_session::VortexSession;
3130

3231
fn main() {
33-
warm_up_vtables();
3432
divan::main();
3533
}
3634

0 commit comments

Comments
 (0)