Skip to content

Commit 083a5b5

Browse files
robert3005robert3005
committed
Deprecate the session-less optimize and keep constructors lazy
ArrayOptimizer::optimize uses only the static rules of the encodings, so it misses the kernels a session registers. It is now deprecated in favour of optimize_ctx, which takes the session or an ExecutionCtx created from it through the new AsOptimizerKernels trait. optimize_recursive takes the same argument. Constructors such as slice, filter, take, cast, fill_null, get_item and apply no longer optimize the arrays they build. Execution optimizes each node it reaches, and the callers that read an array's encoding directly optimize it with their session first: the scan tasks, the plan v2 operators, the dict layout writer, the FFI scan, and CUDA dispatch. The default append_to_builder optimizes the array first and appends through the encoding it reduces to, so a lazy slice of a constant appends as a constant. Signed-off-by: Robert Kruszewski <github@robertk.io>
1 parent 6f69042 commit 083a5b5

97 files changed

Lines changed: 735 additions & 361 deletions

File tree

Some content is hidden

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

‎encodings/datetime-parts/src/compute/rules.rs‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ use vortex_array::arrays::slice::SliceReduceAdaptor;
1616
use vortex_array::builtins::ArrayBuiltins;
1717
use vortex_array::dtype::DType;
1818
use vortex_array::extension::datetime::Timestamp;
19-
use vortex_array::optimizer::ArrayOptimizer;
2019
use vortex_array::optimizer::rules::ArrayParentReduceRule;
2120
use vortex_array::optimizer::rules::ParentRuleSet;
2221
use vortex_array::scalar_fn::fns::between::Between;
@@ -143,9 +142,7 @@ impl ArrayParentReduceRule<DateTimeParts> for DTPComparisonPushDownRule {
143142
}
144143
}
145144

146-
let result = ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)?
147-
.into_array()
148-
.optimize()?;
145+
let result = ScalarFnArray::try_new(parent.scalar_fn().clone(), new_children)?.into_array();
149146

150147
Ok(Some(result))
151148
}
@@ -291,7 +288,7 @@ mod tests {
291288
.into_array();
292289

293290
// Optimize should push down to days
294-
let optimized = comparison.optimize().unwrap();
291+
let optimized = comparison.optimize_ctx(&SESSION).unwrap();
295292

296293
// Verify correctness: days [0, 1, 2] <= 1 should give [true, true, false]
297294
let mut ctx = SESSION.create_execution_ctx();
@@ -325,7 +322,7 @@ mod tests {
325322
.into_array();
326323

327324
// Optimize should push down to days
328-
let optimized = between.optimize().unwrap();
325+
let optimized = between.optimize_ctx(&SESSION).unwrap();
329326

330327
// Verify correctness: days [0, 1, 2, 3, 4] between 1 and 3 should give [false, true, true, true, false]
331328
let mut ctx = SESSION.create_execution_ctx();
@@ -363,7 +360,7 @@ mod tests {
363360
let constant = ConstantArray::new(scalar, dtp.len()).into_array();
364361
let comparison = Binary::try_new(dtp, constant, Operator::Lte)?.into_array();
365362

366-
let optimized = comparison.optimize()?;
363+
let optimized = comparison.optimize_ctx(&SESSION)?;
367364
assert_eq!(optimized.dtype(), comparison.dtype());
368365

369366
let mut ctx = SESSION.create_execution_ctx();
@@ -388,7 +385,7 @@ mod tests {
388385
.into_array();
389386

390387
// Optimize should NOT push down (constant has non-zero seconds)
391-
let optimized = comparison.optimize().unwrap();
388+
let optimized = comparison.optimize_ctx(&SESSION).unwrap();
392389

393390
// days [0, 1, 2] at midnight <= day 1 at noon: [true, true, false]
394391
let mut ctx = SESSION.create_execution_ctx();
@@ -424,7 +421,7 @@ mod tests {
424421
.into_array();
425422

426423
// Should still compute correctly (just not optimized via pushdown)
427-
let optimized = comparison.optimize().unwrap();
424+
let optimized = comparison.optimize_ctx(&SESSION).unwrap();
428425
// timestamps at 1am on days [0, 1, 2] <= day 1 midnight: [true, false, false]
429426
assert_arrays_eq!(
430427
optimized,

‎encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ mod tests {
4040
use vortex_array::assert_arrays_eq;
4141
use vortex_array::dtype::DecimalDType;
4242
use vortex_array::dtype::i256;
43+
use vortex_array::optimizer::ArrayOptimizer;
4344
use vortex_array::validity::Validity;
4445
use vortex_buffer::Buffer;
4546
use vortex_buffer::buffer;
@@ -55,8 +56,8 @@ mod tests {
5556
});
5657

5758
/// Taking pushes down into the parts during optimization, with no execution context in
58-
/// play: `ArrayRef::take` wraps the array in a `Dict` and optimizes, and the reduce rule
59-
/// must rewrite that into a `DecimalByteParts` of taken parts.
59+
/// play: `ArrayRef::take` wraps the array in a `Dict`, and optimizing it must rewrite that
60+
/// into a `DecimalByteParts` of taken parts.
6061
#[test]
6162
fn take_pushes_down_without_executing() -> VortexResult<()> {
6263
let decimal = DecimalArray::new(
@@ -67,7 +68,8 @@ mod tests {
6768
let indices = buffer![0u64, 2].into_array();
6869
let taken = DecimalByteParts::encode(&decimal, &mut SESSION.create_execution_ctx())?
6970
.into_array()
70-
.take(indices)?;
71+
.take(indices)?
72+
.optimize_ctx(&SESSION)?;
7173

7274
assert!(
7375
taken.is::<DecimalByteParts>(),

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,7 @@ mod tests {
205205
use vortex_array::arrays::slice::SliceKernel;
206206
use vortex_array::assert_arrays_eq;
207207
use vortex_array::builtins::ArrayBuiltins;
208+
use vortex_array::optimizer::ArrayOptimizer;
208209
use vortex_array::scalar_fn::fns::binary::CompareKernel;
209210
use vortex_array::scalar_fn::fns::operators::CompareOperator;
210211
use vortex_array::scalar_fn::fns::operators::Operator;
@@ -364,7 +365,10 @@ mod tests {
364365
let prim = PrimitiveArray::from_iter(values);
365366
let packed = BitPackedData::encode(&prim.clone().into_array(), 7, &mut ctx)?;
366367

367-
let sliced = packed.into_array().slice(start..start + slice_len)?;
368+
let sliced = packed
369+
.into_array()
370+
.slice(start..start + slice_len)?
371+
.optimize_ctx(&SESSION)?;
368372
let rhs = ConstantArray::new(50u32, slice_len).into_array();
369373
for op in [
370374
CompareOperator::Eq,

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@ mod tests {
150150
use vortex_array::assert_arrays_eq;
151151
use vortex_array::compute::conformance::filter::test_filter_conformance;
152152
use vortex_array::dtype::NativePType;
153+
use vortex_array::optimizer::ArrayOptimizer;
153154
use vortex_array::validity::Validity;
154155
use vortex_buffer::BitBuffer;
155156
use vortex_buffer::Buffer;
@@ -507,7 +508,8 @@ mod tests {
507508
PrimitiveArray::from_iter((0..5000u32).map(|i| <T as From<u8>>::from((i % 100) as u8)));
508509
let bitpacked = BitPackedData::encode(&unpacked.into_array(), 7, &mut ctx)?
509510
.into_array()
510-
.slice(range.clone())?;
511+
.slice(range.clone())?
512+
.optimize_ctx(&SESSION)?;
511513
let bitpacked = bitpacked.as_::<BitPacked>();
512514

513515
let bits = selection(pattern, range.len());

‎encodings/fastlanes/src/bitpacking/plugin/tests.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use vortex_array::builtins::ArrayBuiltins;
2020
use vortex_array::dtype::DType;
2121
use vortex_array::dtype::Nullability;
2222
use vortex_array::dtype::PType;
23+
use vortex_array::optimizer::ArrayOptimizer;
2324
use vortex_array::serde::SerializeOptions;
2425
use vortex_array::serde::SerializedArray;
2526
use vortex_array::session::ArraySessionExt;
@@ -224,7 +225,8 @@ fn serde_roundtrip(
224225
let mut ctx = PLUGIN_SESSION.create_execution_ctx();
225226
let array = BitPackedData::encode(&values, bit_width, &mut ctx)?
226227
.into_array()
227-
.slice(range.clone())?;
228+
.slice(range.clone())?
229+
.optimize_ctx(&PLUGIN_SESSION)?;
228230
let view = array.as_::<BitPacked>();
229231
let serialization = PLUGIN_SESSION
230232
.array_serialize(&array)?

‎encodings/fastlanes/src/delta/compute/cast.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ mod tests {
6363
use vortex_array::dtype::DType;
6464
use vortex_array::dtype::Nullability;
6565
use vortex_array::dtype::PType;
66+
use vortex_array::optimizer::ArrayOptimizer;
6667
use vortex_buffer::buffer;
6768
use vortex_error::VortexResult;
6869
use vortex_error::vortex_err;
@@ -129,7 +130,8 @@ mod tests {
129130

130131
let casted = array
131132
.into_array()
132-
.cast(DType::Primitive(PType::U32, Nullability::Nullable))?;
133+
.cast(DType::Primitive(PType::U32, Nullability::Nullable))?
134+
.optimize_ctx(&SESSION)?;
133135

134136
assert_eq!(
135137
casted.dtype(),

‎encodings/fastlanes/src/for_/plugin/tests.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ use vortex_array::IntoArray;
1010
use vortex_array::VortexSessionExecute;
1111
use vortex_array::arrays::PrimitiveArray;
1212
use vortex_array::assert_arrays_eq;
13+
use vortex_array::optimizer::ArrayOptimizer;
1314
use vortex_array::scalar::ScalarValue;
1415
use vortex_array::serde::SerializeOptions;
1516
use vortex_array::serde::SerializedArray;
@@ -81,7 +82,10 @@ fn chunked() -> VortexResult<FoRArray> {
8182
#[case::whole(0..5000)]
8283
#[case::sliced(1500..2600)]
8384
fn v2_roundtrip(#[case] range: std::ops::Range<usize>) -> VortexResult<()> {
84-
let array = chunked()?.into_array().slice(range.clone())?;
85+
let array = chunked()?
86+
.into_array()
87+
.slice(range.clone())?
88+
.optimize_ctx(&SESSION)?;
8589
let serialization = SESSION
8690
.array_serialize(&array)?
8791
.ok_or_else(|| vortex_err!("FoR must serialize"))?;

‎encodings/fastlanes/src/for_/tests.rs‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use vortex_array::assert_arrays_eq;
2525
use vortex_array::compute::conformance::consistency::test_array_consistency;
2626
use vortex_array::dtype::NativePType;
2727
use vortex_array::dtype::Nullability;
28+
use vortex_array::optimizer::ArrayOptimizer;
2829
use vortex_array::scalar::Scalar;
2930
use vortex_array::session::ArraySessionExt;
3031
use vortex_array::smallvec::smallvec;
@@ -168,7 +169,10 @@ fn global_reference_survives_slicing() -> VortexResult<()> {
168169
)?;
169170
assert_eq!(global(&array)?, 7u32.into());
170171
assert!(array.blocked_references().is_none());
171-
let sliced = array.into_array().slice(1500..1600)?;
172+
let sliced = array
173+
.into_array()
174+
.slice(1500..1600)?
175+
.optimize_ctx(&SESSION)?;
172176
assert_eq!(global(&sliced.as_::<FoR>().into_owned())?, 7u32.into());
173177
Ok(())
174178
}
@@ -335,7 +339,10 @@ fn fused_decode_sliced(
335339
) -> VortexResult<()> {
336340
let mut ctx = SESSION.create_execution_ctx();
337341
let (array, expected) = fused(bit_width, signed)?;
338-
let sliced = array.into_array().slice(start..end)?;
342+
let sliced = array
343+
.into_array()
344+
.slice(start..end)?
345+
.optimize_ctx(&SESSION)?;
339346
let sliced_for = sliced.as_::<FoR>();
340347
// A BitPacked child with patches slices lazily, which takes the unfused path.
341348
assert_eq!(sliced_for.encoded().is::<BitPacked>(), bit_width == 3);

‎encodings/fastlanes/tests/widen.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ use vortex_array::assert_arrays_eq;
3030
use vortex_array::builtins::ArrayBuiltins;
3131
use vortex_array::dtype::DType;
3232
use vortex_array::dtype::PType;
33+
use vortex_array::optimizer::ArrayOptimizer;
3334
use vortex_array::scalar::Scalar;
3435
use vortex_array::scalar_fn::fns::operators::Operator;
3536
use vortex_buffer::buffer;
@@ -58,7 +59,9 @@ fn test_operations() -> VortexResult<()> {
5859
let mut ctx = SESSION.create_execution_ctx();
5960
let array = widen_bitpacked(PrimitiveArray::from_iter([0u8, 1, 3, 7]), 3)?;
6061

61-
let selected = array.filter(Mask::from_iter([true, false, true, true]))?;
62+
let selected = array
63+
.filter(Mask::from_iter([true, false, true, true]))?
64+
.optimize_ctx(&SESSION)?;
6265
assert!(selected.is::<Widen>());
6366
assert_eq!(selected.as_::<Widen>().values().dtype(), &PType::U8.into());
6467
assert_arrays_eq!(selected, buffer![0u64, 3, 7].into_array(), &mut ctx);

‎encodings/fsst/src/array.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1072,6 +1072,7 @@ mod test {
10721072
use vortex_array::dtype::DType;
10731073
use vortex_array::dtype::Nullability;
10741074
use vortex_array::dtype::PType;
1075+
use vortex_array::optimizer::ArrayOptimizer;
10751076
use vortex_array::test_harness::check_metadata;
10761077
use vortex_array::vtable::VTable as _;
10771078
use vortex_buffer::Buffer;
@@ -1108,6 +1109,7 @@ mod test {
11081109
let compressor_ptr = fsst_array.compressor() as *const Compressor;
11091110
let sliced = fsst_array
11101111
.slice(1..3)?
1112+
.optimize_ctx(&SESSION)?
11111113
.try_downcast::<FSST>()
11121114
.map_err(|_| vortex_err!("slice must return an FSST array"))?;
11131115
let sliced_compressor_ptr = sliced.compressor() as *const Compressor;

0 commit comments

Comments
 (0)