Skip to content

Commit 5ec8d2f

Browse files
gatesngatesn
authored
Use execute where possible (#6835)
When the ExecutionCtx is in scope, prefer to execute rather than use the legacy ToCanonical trait. --------- Signed-off-by: Nicholas Gates <nick@nickgates.com>
1 parent 6a34208 commit 5ec8d2f

11 files changed

Lines changed: 27 additions & 39 deletions

File tree

‎encodings/alp/src/alp/decompress.rs‎

Lines changed: 5 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
use std::mem::transmute;
55

66
use vortex_array::ExecutionCtx;
7-
use vortex_array::ToCanonical;
87
use vortex_array::arrays::PrimitiveArray;
98
use vortex_array::arrays::primitive::chunk_range;
109
use vortex_array::arrays::primitive::patch_chunk;
@@ -33,14 +32,10 @@ pub fn decompress_into_array(
3332
if let Some(ref patches) = patches
3433
&& let Some(chunk_offsets) = patches.chunk_offsets()
3534
{
36-
let prim_encoded = encoded.to_primitive();
37-
// We need to drop ALPArray here in case converting encoded buffer into
38-
// primitive didn't create a copy. In that case both alp_encoded and array
39-
// will hold a reference to the buffer we want to mutate.
40-
drop(encoded);
41-
let patches_chunk_offsets = chunk_offsets.as_ref().to_primitive();
42-
let patches_indices = patches.indices().to_primitive();
43-
let patches_values = patches.values().to_primitive();
35+
let prim_encoded = encoded.execute::<PrimitiveArray>(ctx)?;
36+
let patches_chunk_offsets = chunk_offsets.clone().execute::<PrimitiveArray>(ctx)?;
37+
let patches_indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
38+
let patches_values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
4439
Ok(decompress_chunked_core(
4540
prim_encoded,
4641
exponents,
@@ -51,11 +46,7 @@ pub fn decompress_into_array(
5146
dtype,
5247
))
5348
} else {
54-
let encoded_prim = encoded.to_primitive();
55-
// We need to drop ALPArray here in case converting encoded buffer into
56-
// primitive didn't create a copy. In that case both alp_encoded and array
57-
// will hold a reference to the buffer we want to mutate.
58-
drop(encoded);
49+
let encoded_prim = encoded.execute::<PrimitiveArray>(ctx)?;
5950
decompress_unchunked_core(encoded_prim, exponents, patches, dtype, ctx)
6051
}
6152
}

‎encodings/datetime-parts/src/canonical.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,6 @@ mod test {
110110
use rstest::rstest;
111111
use vortex_array::ExecutionCtx;
112112
use vortex_array::IntoArray;
113-
use vortex_array::ToCanonical;
114113
use vortex_array::arrays::PrimitiveArray;
115114
use vortex_array::arrays::TemporalArray;
116115
use vortex_array::assert_arrays_eq;
@@ -156,7 +155,8 @@ mod test {
156155
let mut ctx = ExecutionCtx::new(VortexSession::empty());
157156
let primitive_values = decode_to_temporal(&date_times, &mut ctx)?
158157
.temporal_values()
159-
.to_primitive();
158+
.clone()
159+
.execute::<PrimitiveArray>(&mut ctx)?;
160160

161161
assert_arrays_eq!(primitive_values, milliseconds);
162162
assert_eq!(primitive_values.validity(), &validity);

‎encodings/fastlanes/src/bitpacking/compute/take.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ use vortex_array::ArrayRef;
99
use vortex_array::DynArray;
1010
use vortex_array::ExecutionCtx;
1111
use vortex_array::IntoArray;
12-
use vortex_array::ToCanonical;
1312
use vortex_array::arrays::PrimitiveArray;
1413
use vortex_array::arrays::dict::TakeExecute;
1514
use vortex_array::dtype::IntegerPType;
@@ -43,7 +42,8 @@ impl TakeExecute for BitPackedVTable {
4342
) -> VortexResult<Option<ArrayRef>> {
4443
// If the indices are large enough, it's faster to flatten and take the primitive array.
4544
if indices.len() * UNPACK_CHUNK_THRESHOLD > array.len() {
46-
return array.to_primitive().take(indices.to_array()).map(Some);
45+
let prim = array.clone().into_array().execute::<PrimitiveArray>(ctx)?;
46+
return prim.take(indices.to_array()).map(Some);
4747
}
4848

4949
// NOTE: we use the unsigned PType because all values in the BitPackedArray must
@@ -52,7 +52,7 @@ impl TakeExecute for BitPackedVTable {
5252
let validity = array.validity();
5353
let taken_validity = validity.take(indices)?;
5454

55-
let indices = indices.to_primitive();
55+
let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
5656
let taken = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| {
5757
match_each_integer_ptype!(indices.ptype(), |I| {
5858
take_primitive::<T, I>(array, &indices, taken_validity, ctx)?

‎encodings/runend/public-api.lock‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,11 +122,11 @@ pub fn vortex_runend::RunEndVTable::fmt(&self, f: &mut core::fmt::Formatter<'_>)
122122

123123
impl vortex_array::arrays::dict::take::TakeExecute for vortex_runend::RunEndVTable
124124

125-
pub fn vortex_runend::RunEndVTable::take(array: &vortex_runend::RunEndArray, indices: &vortex_array::array::ArrayRef, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
125+
pub fn vortex_runend::RunEndVTable::take(array: &vortex_runend::RunEndArray, indices: &vortex_array::array::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
126126

127127
impl vortex_array::arrays::filter::kernel::FilterKernel for vortex_runend::RunEndVTable
128128

129-
pub fn vortex_runend::RunEndVTable::filter(array: &vortex_runend::RunEndArray, mask: &vortex_mask::Mask, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
129+
pub fn vortex_runend::RunEndVTable::filter(array: &vortex_runend::RunEndArray, mask: &vortex_mask::Mask, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
130130

131131
impl vortex_array::compute::is_constant::IsConstantKernel for vortex_runend::RunEndVTable
132132

@@ -144,7 +144,7 @@ pub fn vortex_runend::RunEndVTable::min_max(&self, array: &vortex_runend::RunEnd
144144

145145
impl vortex_array::scalar_fn::fns::binary::compare::CompareKernel for vortex_runend::RunEndVTable
146146

147-
pub fn vortex_runend::RunEndVTable::compare(lhs: &vortex_runend::RunEndArray, 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>>
147+
pub fn vortex_runend::RunEndVTable::compare(lhs: &vortex_runend::RunEndArray, 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>>
148148

149149
impl vortex_array::scalar_fn::fns::cast::kernel::CastReduce for vortex_runend::RunEndVTable
150150

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,9 @@ use vortex_array::ArrayRef;
55
use vortex_array::DynArray;
66
use vortex_array::ExecutionCtx;
77
use vortex_array::IntoArray;
8-
use vortex_array::ToCanonical;
8+
use vortex_array::arrays::BoolArray;
99
use vortex_array::arrays::ConstantArray;
10+
use vortex_array::arrays::PrimitiveArray;
1011
use vortex_array::builtins::ArrayBuiltins;
1112
use vortex_array::scalar_fn::fns::binary::CompareKernel;
1213
use vortex_array::scalar_fn::fns::operators::CompareOperator;
@@ -22,7 +23,7 @@ impl CompareKernel for RunEndVTable {
2223
lhs: &RunEndArray,
2324
rhs: &ArrayRef,
2425
operator: CompareOperator,
25-
_ctx: &mut ExecutionCtx,
26+
ctx: &mut ExecutionCtx,
2627
) -> VortexResult<Option<ArrayRef>> {
2728
// If the RHS is constant, then we just need to compare against our encoded values.
2829
if let Some(const_scalar) = rhs.as_constant() {
@@ -31,8 +32,8 @@ impl CompareKernel for RunEndVTable {
3132
Operator::from(operator),
3233
)?;
3334
let decoded = runend_decode_bools(
34-
lhs.ends().to_primitive(),
35-
values.to_bool(),
35+
lhs.ends().clone().execute::<PrimitiveArray>(ctx)?,
36+
values.execute::<BoolArray>(ctx)?,
3637
lhs.offset(),
3738
lhs.len(),
3839
)?;

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ use num_traits::AsPrimitive;
88
use vortex_array::ArrayRef;
99
use vortex_array::ExecutionCtx;
1010
use vortex_array::IntoArray;
11-
use vortex_array::ToCanonical;
1211
use vortex_array::arrays::PrimitiveArray;
1312
use vortex_array::arrays::filter::FilterKernel;
1413
use vortex_array::dtype::NativePType;
@@ -30,7 +29,7 @@ impl FilterKernel for RunEndVTable {
3029
fn filter(
3130
array: &RunEndArray,
3231
mask: &Mask,
33-
_ctx: &mut ExecutionCtx,
32+
ctx: &mut ExecutionCtx,
3433
) -> VortexResult<Option<ArrayRef>> {
3534
let mask_values = mask
3635
.values()
@@ -45,7 +44,7 @@ impl FilterKernel for RunEndVTable {
4544
&Validity::NonNullable,
4645
)?))
4746
} else {
48-
let primitive_run_ends = array.ends().to_primitive();
47+
let primitive_run_ends = array.ends().clone().execute::<PrimitiveArray>(ctx)?;
4948
let (run_ends, values_mask) =
5049
match_each_unsigned_integer_ptype!(primitive_run_ends.ptype(), |P| {
5150
filter_run_end_primitive(

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,9 @@ impl TakeExecute for RunEndVTable {
3131
fn take(
3232
array: &RunEndArray,
3333
indices: &ArrayRef,
34-
_ctx: &mut ExecutionCtx,
34+
ctx: &mut ExecutionCtx,
3535
) -> VortexResult<Option<ArrayRef>> {
36-
let primitive_indices = indices.to_primitive();
36+
let primitive_indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
3737

3838
let checked_indices = match_each_integer_ptype!(primitive_indices.ptype(), |P| {
3939
primitive_indices

‎encodings/sequence/public-api.lock‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ pub fn vortex_sequence::SequenceVTable::fmt(&self, f: &mut core::fmt::Formatter<
7272

7373
impl vortex_array::arrays::dict::take::TakeExecute for vortex_sequence::SequenceVTable
7474

75-
pub fn vortex_sequence::SequenceVTable::take(array: &vortex_sequence::SequenceArray, indices: &vortex_array::array::ArrayRef, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
75+
pub fn vortex_sequence::SequenceVTable::take(array: &vortex_sequence::SequenceArray, indices: &vortex_array::array::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
7676

7777
impl vortex_array::arrays::filter::kernel::FilterKernel for vortex_sequence::SequenceVTable
7878

‎encodings/sequence/src/compute/take.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ use vortex_array::ArrayRef;
66
use vortex_array::DynArray;
77
use vortex_array::ExecutionCtx;
88
use vortex_array::IntoArray;
9-
use vortex_array::ToCanonical;
109
use vortex_array::arrays::ConstantArray;
1110
use vortex_array::arrays::PrimitiveArray;
1211
use vortex_array::arrays::dict::TakeExecute;
@@ -77,10 +76,10 @@ impl TakeExecute for SequenceVTable {
7776
fn take(
7877
array: &SequenceArray,
7978
indices: &ArrayRef,
80-
_ctx: &mut ExecutionCtx,
79+
ctx: &mut ExecutionCtx,
8180
) -> VortexResult<Option<ArrayRef>> {
8281
let mask = indices.validity_mask()?;
83-
let indices = indices.to_primitive();
82+
let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
8483
let result_nullability = array.dtype().nullability() | indices.dtype().nullability();
8584

8685
match_each_integer_ptype!(indices.ptype(), |T| {

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ use crate::arrays::ChunkedVTable;
1414
use crate::arrays::PrimitiveArray;
1515
use crate::arrays::dict::TakeExecute;
1616
use crate::builtins::ArrayBuiltins;
17-
use crate::canonical::ToCanonical;
1817
use crate::dtype::DType;
1918
use crate::dtype::PType;
2019
use crate::executor::ExecutionCtx;
@@ -30,7 +29,7 @@ fn take_chunked(
3029
let indices = indices
3130
.to_array()
3231
.cast(DType::Primitive(PType::U64, indices.dtype().nullability()))?
33-
.to_primitive();
32+
.execute::<PrimitiveArray>(ctx)?;
3433

3534
let indices_mask = indices.validity_mask()?;
3635
let indices_values = indices.as_slice::<u64>();

0 commit comments

Comments
 (0)