From 262c53b09f822d014f7e9e918ec059d619c3092a Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 19 Jun 2026 15:43:53 +0100 Subject: [PATCH 01/12] assert_arrays_eq and assert_nth_scalar require ExecutionCtx Signed-off-by: Robert Kruszewski --- encodings/alp/src/alp/array.rs | 126 ++++----- encodings/alp/src/alp/compress.rs | 262 +++++++++--------- encodings/alp/src/alp/compute/between.rs | 23 +- encodings/alp/src/alp/compute/cast.rs | 26 +- encodings/alp/src/alp/compute/compare.rs | 85 +++--- encodings/alp/src/alp_rd/array.rs | 14 +- encodings/alp/src/alp_rd/compute/filter.rs | 18 +- encodings/alp/src/alp_rd/compute/take.rs | 23 +- encodings/alp/src/alp_rd/ops.rs | 23 +- encodings/bytebool/src/array.rs | 13 +- encodings/bytebool/src/compute.rs | 42 ++- encodings/datetime-parts/src/canonical.rs | 15 +- .../src/decimal_byte_parts/compute/compare.rs | 26 +- .../onpair/src/compute/compare.rs | 11 +- .../src/bitpacking/array/bitpack_compress.rs | 30 +- .../bitpacking/array/bitpack_decompress.rs | 44 ++- .../fastlanes/src/bitpacking/array/mod.rs | 11 +- .../src/bitpacking/compute/between.rs | 22 +- .../fastlanes/src/bitpacking/compute/cast.rs | 21 +- .../src/bitpacking/compute/compare.rs | 18 +- .../src/bitpacking/compute/filter.rs | 47 +++- .../fastlanes/src/bitpacking/compute/take.rs | 53 ++-- .../src/bitpacking/vtable/operations.rs | 86 ++++-- .../src/delta/array/delta_compress.rs | 10 +- encodings/fastlanes/src/delta/compute/cast.rs | 27 +- .../fastlanes/src/delta/vtable/operations.rs | 46 +-- .../fastlanes/src/for/array/for_compress.rs | 16 +- encodings/fastlanes/src/for/compute/cast.rs | 13 +- .../fastlanes/src/for/compute/compare.rs | 43 ++- .../fastlanes/src/for/vtable/operations.rs | 12 +- encodings/fastlanes/src/rle/array/mod.rs | 26 +- .../fastlanes/src/rle/array/rle_compress.rs | 82 +++--- encodings/fastlanes/src/rle/compute/cast.rs | 15 +- .../fastlanes/src/rle/vtable/operations.rs | 61 ++-- encodings/fsst/src/compute/compare.rs | 24 +- encodings/fsst/src/compute/like.rs | 84 ++++-- encodings/fsst/src/dfa/tests.rs | 9 +- encodings/fsst/src/kernel.rs | 18 +- encodings/fsst/src/tests.rs | 51 ++-- encodings/parquet-variant/src/array.rs | 30 +- encodings/parquet-variant/src/arrow.rs | 4 +- encodings/parquet-variant/src/kernel.rs | 56 ++-- encodings/parquet-variant/src/vtable.rs | 17 +- encodings/pco/src/array.rs | 5 +- encodings/pco/src/compute/cast.rs | 18 +- encodings/pco/src/tests.rs | 40 +-- encodings/runend/src/array.rs | 18 +- encodings/runend/src/arrow.rs | 24 +- encodings/runend/src/compress.rs | 32 ++- encodings/runend/src/compute/cast.rs | 15 +- encodings/runend/src/compute/compare.rs | 14 +- encodings/runend/src/compute/filter.rs | 5 +- encodings/runend/src/compute/mod.rs | 22 +- encodings/runend/src/compute/take.rs | 66 +++-- encodings/runend/src/compute/take_from.rs | 28 +- encodings/runend/src/decompress_bool.rs | 40 +-- encodings/runend/src/ops.rs | 43 +-- encodings/sequence/src/array.rs | 16 +- encodings/sequence/src/compress.rs | 26 +- encodings/sequence/src/compute/cast.rs | 18 +- encodings/sequence/src/compute/compare.rs | 16 +- .../sequence/src/compute/list_contains.rs | 13 +- encodings/sparse/src/canonical.rs | 102 +++---- encodings/sparse/src/compute/between.rs | 2 +- encodings/sparse/src/compute/cast.rs | 10 +- encodings/sparse/src/compute/compare.rs | 2 +- encodings/sparse/src/compute/fill_null.rs | 2 +- encodings/sparse/src/compute/filter.rs | 22 +- encodings/sparse/src/compute/mod.rs | 22 +- encodings/sparse/src/compute/take.rs | 42 ++- encodings/sparse/src/lib.rs | 43 +-- encodings/sparse/src/ops.rs | 14 +- encodings/zigzag/src/compress.rs | 14 +- encodings/zigzag/src/compute/cast.rs | 22 +- encodings/zigzag/src/compute/mod.rs | 18 +- encodings/zstd/src/compute/cast.rs | 14 +- encodings/zstd/src/test.rs | 74 ++--- encodings/zstd/src/zstd_buffers.rs | 2 +- fuzz/src/array/fill_null.rs | 20 +- fuzz/src/array/mask.rs | 14 +- .../fns/all_non_distinct/primitive.rs | 7 +- .../src/aggregate_fn/fns/count/grouped.rs | 22 +- .../src/aggregate_fn/fns/sum/grouped.rs | 30 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 28 +- vortex-array/src/array/typed.rs | 9 +- vortex-array/src/arrays/assertions.rs | 67 ++--- vortex-array/src/arrays/bool/array.rs | 34 ++- .../src/arrays/bool/compute/filter.rs | 33 ++- vortex-array/src/arrays/bool/compute/take.rs | 38 ++- vortex-array/src/arrays/bool/compute/zip.rs | 4 +- vortex-array/src/arrays/bool/patch.rs | 9 +- vortex-array/src/arrays/bool/vtable/mod.rs | 5 +- .../src/arrays/bool/vtable/operations.rs | 6 +- vortex-array/src/arrays/chunked/array.rs | 16 +- .../src/arrays/chunked/compute/cast.rs | 9 +- .../src/arrays/chunked/compute/take.rs | 38 ++- vortex-array/src/arrays/chunked/tests.rs | 83 ++++-- .../src/arrays/chunked/vtable/canonical.rs | 9 +- .../src/arrays/chunked/vtable/operations.rs | 30 +- .../src/arrays/constant/compute/take.rs | 10 +- .../src/arrays/constant/vtable/canonical.rs | 36 ++- .../src/arrays/constant/vtable/mod.rs | 2 +- vortex-array/src/arrays/datetime/test.rs | 11 +- .../src/arrays/decimal/compute/fill_null.rs | 20 +- .../src/arrays/decimal/compute/take.rs | 9 +- vortex-array/src/arrays/decimal/vtable/mod.rs | 5 +- vortex-array/src/arrays/dict/array.rs | 4 +- vortex-array/src/arrays/dict/compute/cast.rs | 20 +- .../src/arrays/dict/compute/fill_null.rs | 8 +- vortex-array/src/arrays/dict/compute/like.rs | 9 +- vortex-array/src/arrays/dict/compute/mod.rs | 16 +- vortex-array/src/arrays/dict/compute/slice.rs | 9 +- vortex-array/src/arrays/dict/tests.rs | 5 +- .../src/arrays/extension/compute/cast.rs | 4 +- .../arrays/filter/execute/fixed_size_list.rs | 13 +- .../src/arrays/filter/execute/listview.rs | 17 +- .../src/arrays/filter/execute/struct_.rs | 13 +- .../src/arrays/filter/execute/take/tests.rs | 55 ++-- .../arrays/fixed_size_list/tests/filter.rs | 45 ++- .../src/arrays/fixed_size_list/tests/take.rs | 16 +- vortex-array/src/arrays/interleave/mod.rs | 2 +- vortex-array/src/arrays/list/tests.rs | 52 +++- .../src/arrays/listview/compute/zip.rs | 8 +- .../src/arrays/listview/conversion.rs | 72 +++-- vortex-array/src/arrays/listview/rebuild.rs | 38 ++- .../src/arrays/listview/tests/basic.rs | 33 ++- .../src/arrays/listview/tests/filter.rs | 13 +- .../src/arrays/listview/tests/operations.rs | 19 +- .../src/arrays/listview/tests/take.rs | 13 +- vortex-array/src/arrays/masked/tests.rs | 5 +- .../src/arrays/patched/compute/compare.rs | 19 +- .../src/arrays/patched/compute/filter.rs | 31 +-- .../src/arrays/patched/compute/take.rs | 24 +- vortex-array/src/arrays/patched/vtable/mod.rs | 28 +- .../src/arrays/patched/vtable/operations.rs | 7 +- .../src/arrays/patched/vtable/slice.rs | 13 +- .../src/arrays/primitive/array/patch.rs | 6 +- .../src/arrays/primitive/compute/cast.rs | 57 +++- .../src/arrays/primitive/compute/fill_null.rs | 32 ++- .../src/arrays/primitive/compute/take/mod.rs | 6 +- .../src/arrays/primitive/compute/zip.rs | 4 +- .../src/arrays/primitive/vtable/mod.rs | 5 +- .../src/arrays/scalar_fn/vtable/operations.rs | 16 +- vortex-array/src/arrays/shared/tests.rs | 4 +- vortex-array/src/arrays/slice/vtable.rs | 9 +- .../src/arrays/struct_/compute/cast.rs | 3 +- .../src/arrays/struct_/compute/mod.rs | 10 +- .../src/arrays/struct_/compute/rules.rs | 24 +- vortex-array/src/arrays/struct_/tests.rs | 36 ++- .../src/arrays/varbin/compute/compare.rs | 9 +- .../src/arrays/varbin/compute/filter.rs | 37 ++- .../src/arrays/varbin/compute/take.rs | 5 +- vortex-array/src/arrays/varbin/tests.rs | 11 +- .../src/arrays/varbin/vtable/canonical.rs | 4 +- vortex-array/src/arrays/varbinview/compact.rs | 38 ++- vortex-array/src/arrays/varbinview/tests.rs | 16 +- .../src/arrays/varbinview/vtable/mod.rs | 5 +- vortex-array/src/arrays/variant/mod.rs | 15 +- vortex-array/src/builders/bool.rs | 4 +- vortex-array/src/builders/decimal.rs | 7 +- vortex-array/src/builders/dict/primitive.rs | 12 +- vortex-array/src/builders/extension.rs | 5 +- vortex-array/src/builders/list.rs | 10 +- vortex-array/src/builders/listview.rs | 44 ++- vortex-array/src/builders/primitive.rs | 17 +- vortex-array/src/builders/struct_.rs | 5 +- vortex-array/src/builders/varbinview.rs | 18 +- .../src/compute/conformance/filter.rs | 4 +- vortex-array/src/executor.rs | 5 +- vortex-array/src/normalize.rs | 17 +- vortex-array/src/patches.rs | 237 +++++++++++++--- vortex-array/src/scalar_fn/fns/between/mod.rs | 8 +- .../src/scalar_fn/fns/binary/boolean.rs | 16 +- .../src/scalar_fn/fns/binary/compare.rs | 30 +- vortex-array/src/scalar_fn/fns/binary/mod.rs | 37 ++- .../src/scalar_fn/fns/binary/numeric.rs | 55 +++- vortex-array/src/scalar_fn/fns/byte_length.rs | 8 +- vortex-array/src/scalar_fn/fns/case_when.rs | 158 +++++++++-- vortex-array/src/scalar_fn/fns/dynamic.rs | 39 ++- .../src/scalar_fn/fns/fill_null/mod.rs | 25 +- vortex-array/src/scalar_fn/fns/like/mod.rs | 6 +- .../src/scalar_fn/fns/list_contains/mod.rs | 38 ++- vortex-array/src/scalar_fn/fns/merge.rs | 18 +- vortex-array/src/scalar_fn/fns/pack.rs | 26 +- .../src/scalar_fn/fns/variant_get/mod.rs | 16 +- vortex-array/src/scalar_fn/fns/zip/mod.rs | 18 +- vortex-array/src/stats/expr.rs | 66 +++-- vortex-btrblocks/src/canonical_compressor.rs | 48 ++-- vortex-btrblocks/src/schemes/float/tests.rs | 5 +- .../schemes/integer/scheme_selection_tests.rs | 5 +- vortex-btrblocks/src/schemes/integer/tests.rs | 13 +- vortex-compressor/src/builtins/dict/float.rs | 2 +- .../src/builtins/dict/integer.rs | 2 +- vortex-compressor/src/sample.rs | 5 +- vortex-cuda/src/dynamic_dispatch/mod.rs | 104 +++++-- vortex-cuda/src/hybrid_dispatch/mod.rs | 28 +- vortex-cuda/src/kernel/arrays/constant.rs | 13 +- vortex-cuda/src/kernel/arrays/dict.rs | 177 ++++++++++-- vortex-cuda/src/kernel/encodings/alp.rs | 28 +- vortex-cuda/src/kernel/encodings/bitpacked.rs | 40 ++- .../src/kernel/encodings/date_time_parts.rs | 13 +- .../kernel/encodings/decimal_byte_parts.rs | 9 +- vortex-cuda/src/kernel/encodings/for_.rs | 8 +- vortex-cuda/src/kernel/encodings/fsst.rs | 9 +- vortex-cuda/src/kernel/encodings/runend.rs | 42 +-- vortex-cuda/src/kernel/encodings/sequence.rs | 5 +- vortex-cuda/src/kernel/encodings/zigzag.rs | 5 +- vortex-cuda/src/kernel/encodings/zstd.rs | 29 +- .../src/kernel/encodings/zstd_buffers.rs | 17 +- vortex-cuda/src/kernel/filter/decimal.rs | 9 +- vortex-cuda/src/kernel/filter/primitive.rs | 9 +- vortex-cuda/src/kernel/filter/varbinview.rs | 5 +- vortex-cuda/src/kernel/patches/mod.rs | 2 +- vortex-cuda/src/kernel/patches/types.rs | 5 +- vortex-duckdb/src/convert/vector.rs | 35 ++- vortex-duckdb/src/exporter/dict.rs | 3 +- vortex-ffi/src/scan.rs | 21 +- vortex-ffi/src/struct_array.rs | 2 +- vortex-file/src/tests.rs | 60 ++-- vortex-file/tests/test_write_table.rs | 5 +- vortex-ipc/src/iterator.rs | 5 +- vortex-ipc/src/stream.rs | 13 +- vortex-layout/src/layouts/chunked/reader.rs | 5 +- vortex-layout/src/layouts/dict/reader.rs | 24 +- vortex-layout/src/layouts/flat/reader.rs | 13 +- vortex-layout/src/layouts/repartition.rs | 1 + vortex-layout/src/layouts/row_idx/mod.rs | 16 +- vortex-layout/src/layouts/struct_/reader.rs | 20 +- vortex-layout/src/layouts/zoned/reader.rs | 5 +- vortex-layout/src/layouts/zoned/zone_map.rs | 102 +++++-- vortex-layout/src/scan/arrow.rs | 1 + vortex-python/src/arrays/range_to_sequence.rs | 29 +- vortex-tensor/src/types/vector/arrow.rs | 4 +- vortex-test/compat-gen/src/check.rs | 6 +- .../arrays/synthetic/encodings/patched.rs | 4 +- .../src/fixtures/arrays/synthetic/mod.rs | 5 +- .../common_encoding_tree_throughput.rs | 1 + 237 files changed, 4297 insertions(+), 1783 deletions(-) diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index f0f87e2229d..8e581fe51fc 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -471,7 +471,6 @@ mod tests { use rstest::rstest; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -482,7 +481,11 @@ mod tests { use crate::alp_encode; use crate::decompress_into_array; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[rstest] #[case(0)] @@ -507,10 +510,9 @@ mod tests { .unwrap() }; // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -525,26 +527,19 @@ mod tests { #[case(2049)] fn test_execute_f64(#[case] size: usize) { let values = PrimitiveArray::from_iter((0..size).map(|i| i as f64)); - let encoded = alp_encode( - values.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(values.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); - let result_canonical = { - let mut ctx = SESSION.create_execution_ctx(); - encoded - .clone() - .into_array() - .execute::(&mut ctx) - .unwrap() - }; + let mut ctx = SESSION.create_execution_ctx(); + let result_canonical = encoded + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap(); // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -564,27 +559,20 @@ mod tests { .collect(); let array = PrimitiveArray::from_iter(values); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().unwrap().array_len() > 0); - let result_canonical = { - let mut ctx = SESSION.create_execution_ctx(); - encoded - .clone() - .into_array() - .execute::(&mut ctx) - .unwrap() - }; + let mut ctx = SESSION.create_execution_ctx(); + let result_canonical = encoded + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap(); // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -603,26 +591,19 @@ mod tests { .collect(); let array = PrimitiveArray::from_option_iter(values); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); - let result_canonical = { - let mut ctx = SESSION.create_execution_ctx(); - encoded - .clone() - .into_array() - .execute::(&mut ctx) - .unwrap() - }; + let mut ctx = SESSION.create_execution_ctx(); + let result_canonical = encoded + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap(); // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -643,27 +624,20 @@ mod tests { .collect(); let array = PrimitiveArray::from_option_iter(values); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().unwrap().array_len() > 0); - let result_canonical = { - let mut ctx = SESSION.create_execution_ctx(); - encoded - .clone() - .into_array() - .execute::(&mut ctx) - .unwrap() - }; + let mut ctx = SESSION.create_execution_ctx(); + let result_canonical = encoded + .clone() + .into_array() + .execute::(&mut ctx) + .unwrap(); // Compare against the traditional array-based decompress path - let expected = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let expected = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(result_canonical.into_array(), expected); + assert_arrays_eq!(result_canonical.into_array(), expected, &mut ctx); } #[rstest] @@ -720,7 +694,7 @@ mod tests { #[case(1000, 200)] #[case(2048, 512)] fn test_sliced_to_primitive(#[case] size: usize, #[case] slice_start: usize) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec> = (0..size) .map(|i| { if i % 5 == 0 { @@ -784,7 +758,7 @@ mod tests { let normally_encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); assert!( @@ -818,7 +792,7 @@ mod tests { // The legacy decompress_into_array path should work correctly. let result_legacy = decompress_into_array( alp_without_chunk_offsets.clone(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let legacy_slice = result_legacy.as_slice::(); diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index dff1c3dda2d..819fc4b3dc4 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -133,59 +133,64 @@ where mod tests { use core::f32; use core::f64; + use std::sync::LazyLock; use f64::consts::E; use f64::consts::PI; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_session::VortexSession; use super::*; use crate::alp::array::ALPArrayExt; use crate::alp::array::ALPArraySlotsExt; use crate::decompress_into_array; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_compress() { let array = PrimitiveArray::new(buffer![1.234f32; 1025], Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); let expected_encoded = PrimitiveArray::from_iter(vec![1234i32; 1025]); - assert_arrays_eq!(encoded.encoded(), expected_encoded); + assert_arrays_eq!( + encoded.encoded(), + expected_encoded, + &mut SESSION.create_execution_ctx() + ); assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 }); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] fn test_nullable_compress() { let array = PrimitiveArray::from_option_iter([None, Some(1.234f32), None]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); let expected_encoded = PrimitiveArray::from_option_iter([None, Some(1234i32), None]); - assert_arrays_eq!(encoded.encoded(), expected_encoded); + assert_arrays_eq!( + encoded.encoded(), + expected_encoded, + &mut SESSION.create_execution_ctx() + ); assert_eq!(encoded.exponents(), Exponents { e: 9, f: 6 }); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); let expected = PrimitiveArray::from_option_iter(vec![None, Some(1.234f32), None]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -193,21 +198,24 @@ mod tests { fn test_patched_compress() { let values = buffer![1.234f64, 2.718, PI, 4.0]; let array = PrimitiveArray::new(values.clone(), Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); let expected_encoded = PrimitiveArray::from_iter(vec![1234i64, 2718, 1234, 4000]); - assert_arrays_eq!(encoded.encoded(), expected_encoded); + assert_arrays_eq!( + encoded.encoded(), + expected_encoded, + &mut SESSION.create_execution_ctx() + ); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); let expected_decoded = PrimitiveArray::new(values, Validity::NonNullable); - assert_arrays_eq!(decoded, expected_decoded); + assert_arrays_eq!( + decoded, + expected_decoded, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -215,21 +223,20 @@ mod tests { fn test_compress_ignores_invalid_exceptional_values() { let values = buffer![1.234f64, 2.718, PI, 4.0]; let array = PrimitiveArray::new(values, Validity::from_iter([true, true, false, true])); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); let expected_encoded = PrimitiveArray::from_option_iter(buffer![Some(1234i64), Some(2718), None, Some(4000)]); - assert_arrays_eq!(encoded.encoded(), expected_encoded); + assert_arrays_eq!( + encoded.encoded(), + expected_encoded, + &mut SESSION.create_execution_ctx() + ); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -242,20 +249,15 @@ mod tests { Some(4.0), None, ]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); assert_eq!(encoded.exponents(), Exponents { e: 16, f: 13 }); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut SESSION.create_execution_ctx()); - let _decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let _decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); } #[test] @@ -264,15 +266,15 @@ mod tests { let alp_arr = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); - assert_arrays_eq!(alp_arr, original); + assert_arrays_eq!(alp_arr, original, &mut SESSION.create_execution_ctx()); } #[test] fn roundtrips_all_null() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let original = PrimitiveArray::new(buffer![195.26274f64, PI, -48.815685], Validity::AllInvalid); let alp_arr = alp_encode(original.as_view(), None, &mut ctx).unwrap(); @@ -287,12 +289,12 @@ mod tests { decompressed.as_slice::() ); - assert_arrays_eq!(decompressed, original); + assert_arrays_eq!(decompressed, original, &mut ctx); } #[test] fn non_finite_numbers() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let original = PrimitiveArray::new( buffer![0.0f32, -0.0, f32::NAN, f32::NEG_INFINITY, f32::INFINITY], Validity::NonNullable, @@ -315,7 +317,7 @@ mod tests { #[test] fn test_chunk_offsets() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = vec![1.0f64; 3072]; values[1023] = PI; @@ -333,7 +335,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 3]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx); let patch_indices = patches .indices() @@ -341,7 +343,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![1023u64, 1024, 1025]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices, expected_indices, &mut ctx); let patch_values = patches .values() @@ -349,12 +351,12 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E, PI]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values, expected_values, &mut ctx); } #[test] fn test_chunk_offsets_no_patches_in_middle() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = vec![1.0f64; 3072]; values[0] = PI; values[2048] = E; @@ -370,7 +372,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx); let patch_indices = patches .indices() @@ -378,7 +380,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64, 2048]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices, expected_indices, &mut ctx); let patch_values = patches .values() @@ -386,12 +388,12 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values, expected_values, &mut ctx); } #[test] fn test_chunk_offsets_trailing_empty_chunks() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = vec![1.0f64; 3072]; values[0] = PI; @@ -406,7 +408,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64, 1, 1]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx); let patch_indices = patches .indices() @@ -414,7 +416,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices, expected_indices, &mut ctx); let patch_values = patches .values() @@ -422,12 +424,12 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values, expected_values, &mut ctx); } #[test] fn test_chunk_offsets_single_chunk() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = vec![1.0f64; 512]; values[0] = PI; values[100] = E; @@ -443,7 +445,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_offsets = PrimitiveArray::from_iter(vec![0u64]); - assert_arrays_eq!(chunk_offsets, expected_offsets); + assert_arrays_eq!(chunk_offsets, expected_offsets, &mut ctx); let patch_indices = patches .indices() @@ -451,7 +453,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_indices = PrimitiveArray::from_iter(vec![0u64, 100]); - assert_arrays_eq!(patch_indices, expected_indices); + assert_arrays_eq!(patch_indices, expected_indices, &mut ctx); let patch_values = patches .values() @@ -459,7 +461,7 @@ mod tests { .execute::(&mut ctx) .unwrap(); let expected_values = PrimitiveArray::from_iter(vec![PI, E]); - assert_arrays_eq!(patch_values, expected_values); + assert_arrays_eq!(patch_values, expected_values, &mut ctx); } #[test] @@ -470,14 +472,18 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let sliced_alp = encoded.slice(512..1024).unwrap(); let expected_slice = original.slice(512..1024).unwrap(); - assert_arrays_eq!(sliced_alp, expected_slice); + assert_arrays_eq!( + sliced_alp, + expected_slice, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -487,14 +493,18 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let sliced_alp = encoded.slice(512..1024).unwrap(); let expected_slice = original.slice(512..1024).unwrap(); - assert_arrays_eq!(sliced_alp, expected_slice); + assert_arrays_eq!( + sliced_alp, + expected_slice, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -508,14 +518,18 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let sliced_alp = encoded.slice(512..1024).unwrap(); let expected_slice = original.slice(512..1024).unwrap(); - assert_arrays_eq!(sliced_alp, expected_slice); + assert_arrays_eq!( + sliced_alp, + expected_slice, + &mut SESSION.create_execution_ctx() + ); assert!(encoded.patches().is_some()); } @@ -533,20 +547,24 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let sliced_alp = encoded.slice(1023..1025).unwrap(); let expected_slice = original.slice(1023..1025).unwrap(); - assert_arrays_eq!(sliced_alp, expected_slice); + assert_arrays_eq!( + sliced_alp, + expected_slice, + &mut SESSION.create_execution_ctx() + ); assert!(encoded.patches().is_some()); } #[test] fn test_slice_half_chunk_nullable_roundtrip() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = (0..1024) .map(|i| if i % 3 == 0 { None } else { Some(2.5f32) }) .collect::>(); @@ -558,41 +576,31 @@ mod tests { let decoded = sliced_alp.execute::(&mut ctx).unwrap(); let expected_slice = original.slice(512..1024).unwrap(); - assert_arrays_eq!(decoded, expected_slice); + assert_arrays_eq!(decoded, expected_slice, &mut ctx); } #[test] fn test_large_f32_array_uniform_values() { let size = 10_000; let array = PrimitiveArray::new(buffer![42.125f32; size], Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] fn test_large_f64_array_uniform_values() { let size = 50_000; let array = PrimitiveArray::new(buffer![123.456789f64; size], Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -605,17 +613,12 @@ mod tests { values[4500] = f32::INFINITY; let array = PrimitiveArray::new(Buffer::from(values), Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -632,16 +635,11 @@ mod tests { values[7000] = 999.999999999; let array = PrimitiveArray::new(Buffer::from(values.clone()), Validity::NonNullable); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); for idx in 0..size { let decoded_val = decoded.as_slice::()[idx]; @@ -667,16 +665,11 @@ mod tests { .collect(); let array = PrimitiveArray::from_option_iter(values); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -693,16 +686,11 @@ mod tests { let validity = Validity::from_iter((0..size).map(|i| !matches!(i, 500 | 2500))); let array = PrimitiveArray::new(Buffer::from(values), validity); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); - let decoded = - decompress_into_array(encoded, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); + let decoded = decompress_into_array(encoded, &mut SESSION.create_execution_ctx()).unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } /// Regression test for patch_chunk index-out-of-bounds when slicing a multi-chunk @@ -732,7 +720,7 @@ mod tests { let encoded = alp_encode( original.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); assert!(encoded.patches().is_some()); @@ -742,11 +730,11 @@ mod tests { // but patches_indices only includes patches up to element 2500. let sliced_alp = encoded.slice(0..2500).unwrap(); let expected = original.slice(0..2500).unwrap(); - assert_arrays_eq!(sliced_alp, expected); + assert_arrays_eq!(sliced_alp, expected, &mut SESSION.create_execution_ctx()); // Also test slicing that starts mid-chunk (both start and end mid-chunk). let sliced_alp = encoded.slice(500..2500).unwrap(); let expected = original.slice(500..2500).unwrap(); - assert_arrays_eq!(sliced_alp, expected); + assert_arrays_eq!(sliced_alp, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/alp/src/alp/compute/between.rs b/encodings/alp/src/alp/compute/between.rs index c1390288e84..dbb6fec9a22 100644 --- a/encodings/alp/src/alp/compute/between.rs +++ b/encodings/alp/src/alp/compute/between.rs @@ -135,8 +135,9 @@ fn encode_upper_bound( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -145,12 +146,19 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::between::BetweenOptions; use vortex_array::scalar_fn::fns::between::StrictComparison; + use vortex_session::VortexSession; use crate::ALPArray; use crate::alp::array::ALPArrayExt; use crate::alp::compute::between::between_impl; use crate::alp_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn assert_between( arr: &ALPArray, lower: f32, @@ -163,7 +171,8 @@ mod tests { assert_arrays_eq!( res, ConstantArray::new(Scalar::bool(expected, res.dtype().nullability()), arr.len()) - .into_array() + .into_array(), + &mut SESSION.create_execution_ctx() ); } @@ -171,12 +180,8 @@ mod tests { fn comparison_range() { let value = 0.0605_f32; let array = PrimitiveArray::from_iter([value; 1]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_none()); assert_between( @@ -237,7 +242,7 @@ mod tests { #[test] fn non_finite_bounds_use_total_order() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([1.234f32; 10]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); diff --git a/encodings/alp/src/alp/compute/cast.rs b/encodings/alp/src/alp/compute/cast.rs index 71b51ecf9af..d77a7257544 100644 --- a/encodings/alp/src/alp/compute/cast.rs +++ b/encodings/alp/src/alp/compute/cast.rs @@ -45,9 +45,10 @@ impl CastReduce for ALP { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -59,13 +60,20 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::alp::array::ALPArrayExt; use crate::alp_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn issue_5766_test_cast_alp_with_patches_to_nullable() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![1.234f32, f32::NAN, 2.345, f32::INFINITY, 3.456].into_array(); let values_primitive = values.clone().execute::(&mut ctx)?; let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?; @@ -81,14 +89,14 @@ mod tests { let expected = values.cast(nullable_dtype)?; let casted_prim = casted.execute::(&mut ctx)?; - assert_arrays_eq!(casted_prim, expected); + assert_arrays_eq!(casted_prim, expected, &mut ctx); Ok(()) } #[test] fn test_cast_alp_f32_to_f64() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![1.5f32, 2.5, 3.5, 4.5].into_array(); let values_primitive = values.execute::(&mut ctx)?; let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?; @@ -112,7 +120,7 @@ mod tests { #[test] fn test_cast_alp_to_int() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![1.0f32, 2.0, 3.0, 4.0].into_array(); let values_primitive = values.execute::(&mut ctx)?; let alp = alp_encode(values_primitive.as_view(), None, &mut ctx)?; @@ -126,7 +134,11 @@ mod tests { ); let decoded = casted.execute::(&mut ctx)?; - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i32, 2, 3, 4])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([1i32, 2, 3, 4]), + &mut ctx + ); Ok(()) } @@ -138,7 +150,7 @@ mod tests { #[case(buffer![42.42f64].into_array())] #[case(buffer![0.0f32, -1.5, 2.5, -3.5, 4.5].into_array())] fn test_cast_alp_conformance(#[case] array: vortex_array::ArrayRef) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array_primitive = array.execute::(&mut ctx)?; let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).vortex_expect("cannot fail"); diff --git a/encodings/alp/src/alp/compute/compare.rs b/encodings/alp/src/alp/compute/compare.rs index 6cd918f8be3..83e38e19bda 100644 --- a/encodings/alp/src/alp/compute/compare.rs +++ b/encodings/alp/src/alp/compute/compare.rs @@ -151,10 +151,10 @@ where #[cfg(test)] mod tests { use std::f32; + use std::sync::LazyLock; use rstest::rstest; use vortex_array::ArrayRef; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; @@ -167,10 +167,17 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::CompareOperator; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_session::VortexSession; use super::*; use crate::alp_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn test_alp_compare>( alp: ArrayView, value: F, @@ -185,7 +192,7 @@ mod tests { #[test] fn basic_comparison_test() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([1.234f32; 1025]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); @@ -200,18 +207,18 @@ mod tests { .unwrap() .unwrap(); let expected = BoolArray::from_iter([false; 1025]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut ctx); let r = alp_scalar_compare(encoded.as_view(), 1.234f32, CompareOperator::Eq) .unwrap() .unwrap(); let expected = BoolArray::from_iter([true; 1025]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut ctx); } #[test] fn comparison_with_unencodable_value() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([1.234f32; 1025]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); @@ -226,18 +233,18 @@ mod tests { .unwrap() .unwrap(); let expected = BoolArray::from_iter([false; 1025]); - assert_arrays_eq!(r_eq, expected); + assert_arrays_eq!(r_eq, expected, &mut ctx); let r_neq = alp_scalar_compare(encoded.as_view(), 1.234444f32, CompareOperator::NotEq) .unwrap() .unwrap(); let expected = BoolArray::from_iter([true; 1025]); - assert_arrays_eq!(r_neq, expected); + assert_arrays_eq!(r_neq, expected, &mut ctx); } #[test] fn comparison_range() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([0.0605_f32; 10]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); @@ -253,33 +260,33 @@ mod tests { .unwrap() .unwrap(); let expected = BoolArray::from_iter([false; 10]); - assert_arrays_eq!(r_gte, expected); + assert_arrays_eq!(r_gte, expected, &mut ctx); // (0.0605_f32 > 0.06051_f32); let r_gt = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Gt) .unwrap() .unwrap(); let expected = BoolArray::from_iter([false; 10]); - assert_arrays_eq!(r_gt, expected); + assert_arrays_eq!(r_gt, expected, &mut ctx); // 0.0605_f32 <= 0.06051_f32; let r_lte = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lte) .unwrap() .unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_lte, expected); + assert_arrays_eq!(r_lte, expected, &mut ctx); // 0.0605_f32 < 0.06051_f32; let r_lt = alp_scalar_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lt) .unwrap() .unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_lt, expected); + assert_arrays_eq!(r_lt, expected, &mut ctx); } #[test] fn comparison_zeroes() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([0.0_f32; 10]); let encoded = alp_encode(array.as_view(), None, &mut ctx).unwrap(); assert!(encoded.patches().is_none()); @@ -293,43 +300,39 @@ mod tests { let r_gte = test_alp_compare(encoded.as_view(), -0.00000001_f32, CompareOperator::Gte).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_gte, expected); + assert_arrays_eq!(r_gte, expected, &mut ctx); let r_gte = test_alp_compare(encoded.as_view(), -0.0_f32, CompareOperator::Gte).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_gte, expected); + assert_arrays_eq!(r_gte, expected, &mut ctx); let r_gt = test_alp_compare(encoded.as_view(), -0.0000000001f32, CompareOperator::Gt).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_gt, expected); + assert_arrays_eq!(r_gt, expected, &mut ctx); let r_gte = test_alp_compare(encoded.as_view(), -0.0_f32, CompareOperator::Gt).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_gte, expected); + assert_arrays_eq!(r_gte, expected, &mut ctx); let r_lte = test_alp_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lte).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_lte, expected); + assert_arrays_eq!(r_lte, expected, &mut ctx); let r_lt = test_alp_compare(encoded.as_view(), 0.06051_f32, CompareOperator::Lt).unwrap(); let expected = BoolArray::from_iter([true; 10]); - assert_arrays_eq!(r_lt, expected); + assert_arrays_eq!(r_lt, expected, &mut ctx); let r_lt = test_alp_compare(encoded.as_view(), -0.00001_f32, CompareOperator::Lt).unwrap(); let expected = BoolArray::from_iter([false; 10]); - assert_arrays_eq!(r_lt, expected); + assert_arrays_eq!(r_lt, expected, &mut ctx); } #[test] fn compare_with_patches() { let array = PrimitiveArray::from_iter([1.234f32, 1.5, 19.0, f32::consts::E, 1_000_000.9]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); assert!(encoded.patches().is_some()); // Not supported! @@ -343,12 +346,8 @@ mod tests { #[test] fn compare_to_null() { let array = PrimitiveArray::from_iter([1.234f32; 10]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); let other = ConstantArray::new( Scalar::null(DType::Primitive(PType::F32, Nullability::Nullable)), @@ -361,7 +360,7 @@ mod tests { .unwrap(); // Comparing to null yields null results let expected = BoolArray::from_iter([None::; 10]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx()); } #[rstest] @@ -371,16 +370,12 @@ mod tests { #[case(f32::NEG_INFINITY, true)] fn compare_to_non_finite_gt(#[case] value: f32, #[case] result: bool) { let array = PrimitiveArray::from_iter([1.234f32; 10]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); let r = test_alp_compare(encoded.as_view(), value, CompareOperator::Gt).unwrap(); let expected = BoolArray::from_iter([result; 10]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx()); } #[rstest] @@ -390,15 +385,11 @@ mod tests { #[case(f32::NEG_INFINITY, false)] fn compare_to_non_finite_lt(#[case] value: f32, #[case] result: bool) { let array = PrimitiveArray::from_iter([1.234f32; 10]); - let encoded = alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let encoded = + alp_encode(array.as_view(), None, &mut SESSION.create_execution_ctx()).unwrap(); let r = test_alp_compare(encoded.as_view(), value, CompareOperator::Lt).unwrap(); let expected = BoolArray::from_iter([result; 10]); - assert_arrays_eq!(r, expected); + assert_arrays_eq!(r, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index 0a812fbc691..160723e83b2 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -621,20 +621,28 @@ impl ValidityChild for ALPRD { #[cfg(test)] mod test { + use std::sync::LazyLock; + use prost::Message; use rstest::rstest; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::PType; use vortex_array::patches::PatchesMetadata; use vortex_array::test_harness::check_metadata; + use vortex_session::VortexSession; use super::ALPRDMetadata; use crate::ALPRDFloat; use crate::alp_rd; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(vec![0.1f32.next_up(); 1024], 1.123_848_f32)] #[case(vec![0.1f64.next_up(); 1024], 1.123_848_591_110_992_f64)] @@ -642,7 +650,7 @@ mod test { #[case] reals: Vec, #[case] seed: T, ) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_eq!(reals.len(), 1024, "test expects 1024-length fixture"); // Null out some of the values. let mut reals: Vec> = reals.into_iter().map(Some).collect(); @@ -664,7 +672,7 @@ mod test { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decoded, PrimitiveArray::from_option_iter(reals)); + assert_arrays_eq!(decoded, PrimitiveArray::from_option_iter(reals), &mut ctx); } #[cfg_attr(miri, ignore)] diff --git a/encodings/alp/src/alp_rd/compute/filter.rs b/encodings/alp/src/alp_rd/compute/filter.rs index 5deaa7c466b..44d6d39c00a 100644 --- a/encodings/alp/src/alp_rd/compute/filter.rs +++ b/encodings/alp/src/alp_rd/compute/filter.rs @@ -41,9 +41,10 @@ impl FilterKernel for ALPRD { #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -51,16 +52,23 @@ mod test { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] fn test_filter(#[case] a: T, #[case] b: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::new(buffer![a, b, outlier], Validity::NonNullable); let encoded = RDEncoder::new(&[a, b]).encode(array.as_view(), &mut ctx); @@ -71,14 +79,14 @@ mod test { let filtered = encoded .filter(Mask::from_iter([true, false, true])) .unwrap(); - assert_arrays_eq!(filtered, PrimitiveArray::from_iter([a, outlier])); + assert_arrays_eq!(filtered, PrimitiveArray::from_iter([a, outlier]), &mut ctx); } #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] fn test_filter_simple(#[case] a: T, #[case] b: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); test_filter_conformance( &RDEncoder::new(&[a, b]) .encode( @@ -93,7 +101,7 @@ mod test { #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] fn test_filter_with_nulls(#[case] a: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); test_filter_conformance( &RDEncoder::new(&[a]) .encode( diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs index 0d980e0abfc..827d2ecbeaa 100644 --- a/encodings/alp/src/alp_rd/compute/take.rs +++ b/encodings/alp/src/alp_rd/compute/take.rs @@ -57,18 +57,26 @@ impl TakeExecute for ALPRD { #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_session::VortexSession; use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] @@ -76,7 +84,7 @@ mod test { use vortex_array::IntoArray as _; use vortex_buffer::buffer; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([a, b, outlier]); let encoded = RDEncoder::new(&[a, b]).encode(array.as_view(), &mut ctx); @@ -95,14 +103,14 @@ mod test { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(taken, PrimitiveArray::from_iter([a, outlier])); + assert_arrays_eq!(taken, PrimitiveArray::from_iter([a, outlier]), &mut ctx); } #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] fn take_with_nulls(#[case] a: T, #[case] b: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([a, b, outlier]); let encoded = RDEncoder::new(&[a, b]).encode(array.as_view(), &mut ctx); @@ -123,7 +131,8 @@ mod test { assert_arrays_eq!( taken, - PrimitiveArray::from_option_iter([Some(a), Some(outlier), None]) + PrimitiveArray::from_option_iter([Some(a), Some(outlier), None]), + &mut ctx ); } @@ -131,7 +140,7 @@ mod test { #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] fn test_take_conformance_alprd(#[case] a: T, #[case] b: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); test_take_conformance( &RDEncoder::new(&[a, b]) .encode( @@ -146,7 +155,7 @@ mod test { #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] fn test_take_with_nulls_conformance(#[case] a: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); test_take_conformance( &RDEncoder::new(&[a]) .encode( diff --git a/encodings/alp/src/alp_rd/ops.rs b/encodings/alp/src/alp_rd/ops.rs index 0351d446424..be8d8f88948 100644 --- a/encodings/alp/src/alp_rd/ops.rs +++ b/encodings/alp/src/alp_rd/ops.rs @@ -65,27 +65,35 @@ impl OperationsVTable for ALPRD { #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::scalar::Scalar; + use vortex_session::VortexSession; use crate::ALPRDArrayExt; use crate::ALPRDFloat; use crate::RDEncoder; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] fn test_slice(#[case] a: T, #[case] b: T, #[case] outlier: T) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([a, b, outlier]); let encoded = RDEncoder::new(&[a, b]).encode(array.as_view(), &mut ctx); assert!(encoded.left_parts_patches().is_some()); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut ctx); } #[rstest] @@ -96,16 +104,16 @@ mod test { #[case] b: T, #[case] outlier: T, ) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_iter([a, b, outlier]); let encoded = RDEncoder::new(&[a, b]).encode(array.as_view(), &mut ctx); assert!(encoded.left_parts_patches().is_some()); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut ctx); } #[test] fn nullable_scalar_at() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let a = 0.1f64; let b = 0.2f64; let outlier = 3e100f64; @@ -114,7 +122,8 @@ mod test { assert!(encoded.left_parts_patches().is_some()); assert_arrays_eq!( encoded, - PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)]) + PrimitiveArray::from_option_iter([Some(a), Some(b), Some(outlier)]), + &mut ctx ); } } diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index 6b92bead37f..dd5db041c15 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -304,9 +304,10 @@ impl OperationsVTable for ByteBool { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::ArrayContext; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::serde::SerializeOptions; @@ -317,6 +318,12 @@ mod tests { use super::*; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_validity_construction() { let v = vec![true, false]; @@ -325,7 +332,7 @@ mod tests { let arr = ByteBool::from_vec(v, Validity::AllValid); assert_eq!(v_len, arr.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for idx in 0..arr.len() { assert!(arr.is_valid(idx, &mut ctx).unwrap()); } @@ -374,6 +381,6 @@ mod tests { .decode(&dtype, len, &ReadContext::new(ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index fb72fc4e549..1870a0b56b5 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -163,8 +163,9 @@ fn truthy_bit_buffer(array: ArrayView<'_, ByteBool>) -> BitBuffer { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; @@ -178,10 +179,17 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::vortex_err; + use vortex_session::VortexSession; use super::*; use crate::ByteBoolArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn bb(v: Vec) -> ByteBoolArray { ByteBool::from_vec(v, Validity::AllValid) } @@ -198,7 +206,11 @@ mod tests { let sliced_arr = vortex_arr.slice(1..4).unwrap(); let expected = bb_opt(vec![Some(true), None, Some(false)]); - assert_arrays_eq!(sliced_arr, expected.into_array()); + assert_arrays_eq!( + sliced_arr, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -212,7 +224,11 @@ mod tests { .unwrap(); let expected = bb(vec![true; 5]); - assert_arrays_eq!(arr, expected.into_array()); + assert_arrays_eq!( + arr, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -226,7 +242,11 @@ mod tests { .unwrap(); let expected = bb(vec![false; 5]); - assert_arrays_eq!(arr, expected.into_array()); + assert_arrays_eq!( + arr, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -240,21 +260,26 @@ mod tests { .unwrap(); let expected = bb_opt(vec![Some(true), Some(true), Some(true), Some(false), None]); - assert_arrays_eq!(arr, expected.into_array()); + assert_arrays_eq!( + arr, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] fn test_boolean_kernel_kleene() -> VortexResult<()> { let lhs = bb_opt(vec![Some(false), Some(true), None, Some(false), None]); let rhs = bb_opt(vec![None, None, Some(true), Some(false), None]).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let and_result = ::boolean(lhs.as_view(), &rhs, Operator::And, &mut ctx)? .ok_or_else(|| vortex_err!("ByteBool should handle ByteBool boolean AND"))?; assert_arrays_eq!( and_result, - BoolArray::from_iter([Some(false), None, None, Some(false), None]) + BoolArray::from_iter([Some(false), None, None, Some(false), None]), + &mut ctx ); let or_result = @@ -262,7 +287,8 @@ mod tests { .ok_or_else(|| vortex_err!("ByteBool should handle ByteBool boolean OR"))?; assert_arrays_eq!( or_result, - BoolArray::from_iter([None, Some(true), Some(true), Some(false), None]) + BoolArray::from_iter([None, Some(true), Some(true), Some(false), None]), + &mut ctx ); Ok(()) diff --git a/encodings/datetime-parts/src/canonical.rs b/encodings/datetime-parts/src/canonical.rs index eee60bbe913..b08b76b148c 100644 --- a/encodings/datetime-parts/src/canonical.rs +++ b/encodings/datetime-parts/src/canonical.rs @@ -99,9 +99,11 @@ pub fn decode_to_temporal( #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; - use vortex_array::ExecutionCtx; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::assert_arrays_eq; @@ -109,12 +111,19 @@ mod test { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::DateTimeParts; use crate::array::DateTimePartsArraySlotsExt; use crate::array::DateTimePartsParts; use crate::canonical::decode_to_temporal; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(Validity::NonNullable)] #[case(Validity::AllValid)] @@ -132,7 +141,7 @@ mod test { ], validity.clone(), ); - let mut ctx = ExecutionCtx::new(vortex_array::array_session()); + let mut ctx = SESSION.create_execution_ctx(); let date_times = DateTimeParts::try_from_temporal( TemporalArray::new_timestamp( milliseconds.clone().into_array(), @@ -160,7 +169,7 @@ mod test { .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(primitive_values, milliseconds); + assert_arrays_eq!(primitive_values, milliseconds, &mut ctx); Ok(()) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index e6cc8d4d72d..04f321cf88f 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -144,7 +144,10 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -159,9 +162,16 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::DecimalByteParts; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn compare_decimal_const() { let decimal_dtype = DecimalDType::new(8, 2); @@ -180,7 +190,7 @@ mod tests { let res = lhs.binary(rhs.into_array(), Operator::Eq).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(true)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -208,7 +218,7 @@ mod tests { let res = lhs.into_array().binary(rhs, Operator::Lte)?; let expected = BoolArray::from_iter([None, Some(true), Some(true), Some(true)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -235,15 +245,15 @@ mod tests { let res = lhs.binary(rhs.clone().into_array(), Operator::Eq).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(false)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); let res = lhs.binary(rhs.clone().into_array(), Operator::Gt).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(true), Some(true)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); let res = lhs.binary(rhs.into_array(), Operator::Lt).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(false)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); // This cannot be converted to a i32. let rhs = ConstantArray::new( @@ -253,14 +263,14 @@ mod tests { let res = lhs.binary(rhs.clone().into_array(), Operator::Eq).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(false)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); let res = lhs.binary(rhs.clone().into_array(), Operator::Gt).unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(false)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); let res = lhs.binary(rhs.into_array(), Operator::Lt).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(true), Some(true)]).into_array(); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs index 2e84411bb58..8c01a36977c 100644 --- a/encodings/experimental/onpair/src/compute/compare.rs +++ b/encodings/experimental/onpair/src/compute/compare.rs @@ -106,7 +106,7 @@ mod tests { let result = arr .binary(ConstantArray::new("", input.len()).into_array(), op)? .execute::(&mut ctx)?; - assert_arrays_eq!(&result, &BoolArray::from_iter(expected)); + assert_arrays_eq!(&result, &BoolArray::from_iter(expected), &mut ctx); Ok(()) } @@ -127,7 +127,8 @@ mod tests { .execute::(&mut ctx)?; assert_arrays_eq!( &eq_empty, - &BoolArray::from_iter([Some(true), None, Some(false)]) + &BoolArray::from_iter([Some(true), None, Some(false)]), + &mut ctx ); let null_rhs = @@ -135,7 +136,11 @@ mod tests { let eq_null = arr .binary(null_rhs.into_array(), Operator::Eq)? .execute::(&mut ctx)?; - assert_arrays_eq!(&eq_null, &BoolArray::from_iter([None::, None, None])); + assert_arrays_eq!( + &eq_null, + &BoolArray::from_iter([None::, None, None]), + &mut ctx + ); Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index 890bf72f4a5..f7a3485c113 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -447,7 +447,11 @@ mod test { use crate::bitpack_compress::test_harness::make_array; use crate::bitpacking::array::BitPackedArrayExt; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_best_bit_width() { @@ -513,14 +517,14 @@ mod test { chunked.append_to_builder(&mut primitive_builder, &mut ctx)?; let ca_into = primitive_builder.finish(); - assert_arrays_eq!(into_ca, ca_into); + assert_arrays_eq!(into_ca, ca_into, &mut ctx); let mut primitive_builder = PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); primitive_builder.extend_from_array(&chunked); let ca_into = primitive_builder.finish(); - assert_arrays_eq!(into_ca, ca_into); + assert_arrays_eq!(into_ca, ca_into, &mut ctx); Ok(()) } @@ -553,7 +557,11 @@ mod test { // chunk 1 (1024-2047): no patches -> points to patch index 2 // chunk 2 (2048-3071): patch at 3000 -> starts at patch index 2 // chunk 3 (3072-4095): patch at 3100 -> starts at patch index 3 - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64, 2, 2, 3])); + assert_arrays_eq!( + chunk_offsets, + PrimitiveArray::from_iter([0u64, 2, 2, 3]), + &mut ctx + ); Ok(()) } @@ -581,7 +589,11 @@ mod test { .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64, 2, 2])); + assert_arrays_eq!( + chunk_offsets, + PrimitiveArray::from_iter([0u64, 2, 2]), + &mut ctx + ); Ok(()) } @@ -614,7 +626,11 @@ mod test { // chunk 2 (2048-3071): no patches -> points to patch index 3 // chunk 3 (3072-4095): no patches -> points to patch index 3 (remaining chunks filled) // chunk 4 (4096-5119): no patches -> points to patch index 3 (remaining chunks filled) - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64, 2, 3, 3, 3])); + assert_arrays_eq!( + chunk_offsets, + PrimitiveArray::from_iter([0u64, 2, 3, 3, 3]), + &mut ctx + ); Ok(()) } @@ -643,7 +659,7 @@ mod test { .execute::(&mut ctx)?; // Single chunk starting at patch index 0. - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64])); + assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64]), &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 71e25ec8197..78aac1aa86e 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -225,7 +225,11 @@ mod tests { bitpack_encode(array, bit_width, None, &mut SESSION.create_execution_ctx()).unwrap() } - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn unpack(bitpacked: &BitPackedArray) -> VortexResult { unpack_array(bitpacked.as_view(), &mut SESSION.create_execution_ctx()) @@ -235,7 +239,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let values = PrimitiveArray::from_iter((0..n).map(|i| (i % 2047) as u16)); let compressed = BitPackedData::encode(&values.clone().into_array(), 11, &mut ctx).unwrap(); - assert_arrays_eq!(compressed, values); + assert_arrays_eq!(compressed, values, &mut ctx); values .as_slice::() @@ -270,7 +274,7 @@ mod tests { .execute::(&mut ctx)?; let bitpacked = encode(&zeros, 0); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 0, 0, 0])); + assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 0, 0, 0]), &mut ctx); Ok(()) } @@ -282,7 +286,7 @@ mod tests { .execute::(&mut ctx)?; let bitpacked = encode(&zeros, 0); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 1, 0, 1])); + assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 1, 0, 1]), &mut ctx); Ok(()) } @@ -294,7 +298,7 @@ mod tests { .execute::(&mut ctx)?; let bitpacked = encode(&zeros, 10); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1024)); + assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1024), &mut ctx); Ok(()) } @@ -309,7 +313,8 @@ mod tests { let actual = unpack(&bitpacked)?; assert_arrays_eq!( actual, - PrimitiveArray::from_iter((5u16..1029).chain(5u16..1029).chain(5u16..1029)) + PrimitiveArray::from_iter((5u16..1029).chain(5u16..1029).chain(5u16..1029)), + &mut ctx ); Ok(()) } @@ -323,7 +328,7 @@ mod tests { let bitpacked = encode(&zeros, 11); assert!(bitpacked.patches().is_none()); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1025)); + assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1025), &mut ctx); Ok(()) } @@ -337,7 +342,7 @@ mod tests { assert_eq!(bitpacked.len(), 1025); assert!(bitpacked.patches().is_some()); let actual = unpack(&bitpacked)?; - assert_arrays_eq!(actual, PrimitiveArray::from_iter(512u16..1537)); + assert_arrays_eq!(actual, PrimitiveArray::from_iter(512u16..1537), &mut ctx); Ok(()) } @@ -352,7 +357,7 @@ mod tests { assert!(bitpacked.patches().is_some()); let slice_ref = bitpacked.into_array().slice(1023..1025)?; let actual = slice_ref.execute::(&mut ctx)?.into_primitive(); - assert_arrays_eq!(actual, PrimitiveArray::from_iter([1535u16, 1536])); + assert_arrays_eq!(actual, PrimitiveArray::from_iter([1535u16, 1536]), &mut ctx); Ok(()) } @@ -369,7 +374,8 @@ mod tests { let actual = slice_ref.execute::(&mut ctx)?.into_primitive(); assert_arrays_eq!( actual, - PrimitiveArray::from_iter((1023u16..2049).map(|x| x + 512)) + PrimitiveArray::from_iter((1023u16..2049).map(|x| x + 512)), + &mut ctx ); Ok(()) } @@ -457,7 +463,11 @@ mod tests { let result = builder.finish_into_primitive(); // Verify all values were correctly unpacked including patches. - assert_arrays_eq!(result, PrimitiveArray::from_iter(values)); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter(values), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -652,7 +662,11 @@ mod tests { // Verify consistency with unpack_array. let zeros_array = unpack(&zeros_bp)?; assert_eq!(zeros_result.len(), zeros_array.len()); - assert_arrays_eq!(zeros_result, zeros_array); + assert_arrays_eq!( + zeros_result, + zeros_array, + &mut SESSION.create_execution_ctx() + ); // Maximum bit width for u16 (15 bits, since bitpacking requires bit_width < type bit width). let max_values = PrimitiveArray::from_iter([32767u16; 50]); // 2^15 - 1 @@ -679,7 +693,11 @@ mod tests { // Verify consistency. let boundary_unpacked = unpack(&boundary_bp)?; assert_eq!(boundary_result.len(), boundary_unpacked.len()); - assert_arrays_eq!(boundary_result, boundary_unpacked); + assert_arrays_eq!( + boundary_result, + boundary_unpacked, + &mut SESSION.create_execution_ctx() + ); // Single element. let single = PrimitiveArray::from_iter([42u8]); diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 03c0973ce52..772de2ac641 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -333,7 +333,11 @@ mod test { use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_encode() { @@ -355,7 +359,7 @@ mod test { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(packed_primitive, expected); + assert_arrays_eq!(packed_primitive, expected, &mut ctx); } #[test] @@ -384,7 +388,8 @@ mod test { .unwrap(); assert_arrays_eq!( packed_primitive, - PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable) + PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable), + &mut ctx ); } } diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index dcc372ad091..1dd3b61dbf4 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -142,9 +142,10 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; @@ -156,10 +157,17 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::BitPackedArrayExt; use crate::BitPackedData; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn opts(lower: StrictComparison, upper: StrictComparison) -> BetweenOptions { BetweenOptions { lower_strict: lower, @@ -176,7 +184,7 @@ mod tests { #[case] lower_strict: StrictComparison, #[case] upper_strict: StrictComparison, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: BufferMut = (0..3000u32).map(|i| i % 257).collect(); let prim = PrimitiveArray::new(values.freeze(), Validity::NonNullable); let packed = BitPackedData::encode(&prim.clone().into_array(), 9, &mut ctx)?; @@ -194,13 +202,13 @@ mod tests { .between(lower, upper, options)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } #[test] fn signed_with_patches_against_primitive_baseline() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec = (0..1500) .map(|i| if i % 73 == 0 { 100_000 + i } else { i % 100 }) .collect(); @@ -221,13 +229,13 @@ mod tests { .between(lower, upper, options)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } #[test] fn nullable_propagates_validity() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let prim = PrimitiveArray::from_option_iter([Some(1u32), None, Some(3), Some(4), None, Some(6)]); let packed = BitPackedData::encode(&prim.clone().into_array(), 3, &mut ctx)?; @@ -244,7 +252,7 @@ mod tests { .into_array() .between(lower, upper, options)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 10060eb57e2..3b917aa36af 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -112,10 +112,11 @@ impl CastKernel for BitPacked { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -128,12 +129,19 @@ mod tests { use vortex_array::match_each_integer_ptype; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::BitPackedArray; use crate::BitPackedData; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray { - BitPackedData::encode(array, bit_width, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + BitPackedData::encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() } #[test] @@ -151,7 +159,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_iter([10u32, 20, 30, 40, 50, 60]) + PrimitiveArray::from_iter([10u32, 20, 30, 40, 50, 60]), + &mut SESSION.create_execution_ctx() ); } @@ -213,7 +222,7 @@ mod tests { let source = match_each_integer_ptype!(src, |S| { values::(len) }); let source_ref = source.into_array(); let target = DType::Primitive(tgt, Nullability::NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Reference: plain primitive cast of the same values. let reference = source_ref @@ -227,7 +236,7 @@ mod tests { let casted = packed .cast(target.clone())? .execute::(&mut ctx)?; - assert_arrays_eq!(casted, reference); + assert_arrays_eq!(casted, reference, &mut ctx); // Also exercise the sliced/offset path (offset > 0, trailer present). if len >= 4 { @@ -242,7 +251,7 @@ mod tests { .slice(lo..hi)? .cast(target.clone())? .execute::(&mut ctx)?; - assert_arrays_eq!(casted, reference); + assert_arrays_eq!(casted, reference, &mut ctx); } } } diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index d5c50751bae..51aca62a9aa 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -125,7 +125,11 @@ mod tests { use crate::BitPackedArrayExt; use crate::BitPackedData; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); /// All six operators on a small in-range input. #[rstest] @@ -146,7 +150,7 @@ mod tests { .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter(expected)); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); } /// Sweep every native int type across several bit-widths. 2048 elements spans two @@ -177,7 +181,7 @@ mod tests { .into_array() .binary(rhs.clone(), Operator::from(op))? .execute::(&mut ctx)?; - assert_arrays_eq!(got, want); + assert_arrays_eq!(got, want, &mut SESSION.create_execution_ctx()); } } Ok(()) @@ -215,7 +219,7 @@ mod tests { .into_array() .binary(rhs, Operator::Eq)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } @@ -259,7 +263,7 @@ mod tests { .slice(start..start + slice_len)? .binary(rhs.clone(), Operator::from(op))? .execute::(&mut ctx)?; - assert_arrays_eq!(got, want); + assert_arrays_eq!(got, want, &mut ctx); } Ok(()) } @@ -296,7 +300,7 @@ mod tests { .slice(start..end)? .binary(rhs, Operator::Eq)? .execute::(&mut ctx)?; - assert_arrays_eq!(got, want); + assert_arrays_eq!(got, want, &mut ctx); Ok(()) } @@ -315,7 +319,7 @@ mod tests { .into_array() .binary(rhs, Operator::Eq)? .execute::(&mut ctx)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 284ae5bce4c..21184d785a5 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -177,8 +177,9 @@ fn filter_with_indices( #[cfg(test)] mod test { + use std::sync::LazyLock; + use vortex_array::IntoArray as _; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -187,13 +188,20 @@ mod test { use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn take_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a u8 array modulo 63. let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); @@ -203,13 +211,14 @@ mod test { let primitive_result = bitpacked.filter(mask).unwrap(); assert_arrays_eq!( primitive_result, - PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]) + PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]), + &mut ctx ); } #[test] fn take_sliced_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a u8 array modulo 63. let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); @@ -218,25 +227,30 @@ mod test { let mask = Mask::from_indices(sliced.len(), vec![1919, 1921]); let primitive_result = sliced.filter(mask).unwrap(); - assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([31u8, 33])); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([31u8, 33]), + &mut ctx + ); } #[test] fn filter_bitpacked() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); let filtered = bitpacked.filter(Mask::from_indices(4096, 0..1024)).unwrap(); let filtered_prim = filtered.execute::(&mut ctx).unwrap(); assert_arrays_eq!( filtered_prim, - PrimitiveArray::from_iter((0..1024).map(|i| (i % 63) as u8)) + PrimitiveArray::from_iter((0..1024).map(|i| (i % 63) as u8)), + &mut ctx ); } #[test] fn filter_bitpacked_signed() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = (0..500).collect(); let unpacked = PrimitiveArray::new(values.clone(), Validity::NonNullable); let bitpacked = BitPackedData::encode(&unpacked.into_array(), 9, &mut ctx).unwrap(); @@ -248,13 +262,14 @@ mod test { assert_arrays_eq!( filtered, - PrimitiveArray::from_iter(values[0..250].iter().copied()) + PrimitiveArray::from_iter(values[0..250].iter().copied()), + &mut ctx ); } #[test] fn test_filter_bitpacked_conformance() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test with u8 values let unpacked = buffer![1u8, 2, 3, 4, 5].into_array(); let bitpacked = BitPackedData::encode(&unpacked, 3, &mut ctx).unwrap(); @@ -278,7 +293,7 @@ mod test { /// This test ensures that the type handling is correct. #[test] fn filter_bitpacked_signed_with_patches() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create signed integer values where some exceed the bit width (causing patches). // Values 0-127 fit in 7 bits, but 1000 and 2000 do not. let values: Vec = vec![0, 10, 1000, 20, 30, 2000, 40, 50, 60, 70]; @@ -296,7 +311,11 @@ mod test { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(filtered, PrimitiveArray::from_iter([0i32, 1000, 2000, 70])); + assert_arrays_eq!( + filtered, + PrimitiveArray::from_iter([0i32, 1000, 2000, 70]), + &mut ctx + ); } /// Regression test for signed integers with patches using low selectivity. @@ -305,7 +324,7 @@ mod test { /// that doesn't fully decompress the array first. #[test] fn filter_bitpacked_signed_with_patches_low_selectivity() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a larger array with signed integers and some patches. let values: Vec = (0..1000) .map(|i| { @@ -332,6 +351,6 @@ mod test { .unwrap(); let expected: Vec = values[0..20].to_vec(); - assert_arrays_eq!(filtered, PrimitiveArray::from_iter(expected)); + assert_arrays_eq!(filtered, PrimitiveArray::from_iter(expected), &mut ctx); } } diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index bbf5bfaf2c8..7ffe051ca40 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -160,27 +160,35 @@ fn take_primitive( #[cfg(test)] #[expect(clippy::cast_possible_truncation)] mod test { + use std::sync::LazyLock; + use rand::RngExt; use rand::distr::Uniform; use rand::rng; use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; use crate::bitpacking::compute::take::take_primitive; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn take_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0, 125, 2047, 2049, 2151, 2790].into_array(); // Create a u8 array modulo 63. @@ -190,25 +198,30 @@ mod test { let primitive_result = bitpacked.take(indices).unwrap(); assert_arrays_eq!( primitive_result, - PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]) + PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]), + &mut ctx ); } #[test] fn take_with_patches() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let unpacked = Buffer::from_iter(0u32..1024).into_array(); let bitpacked = BitPackedData::encode(&unpacked, 2, &mut ctx).unwrap(); let indices = buffer![0, 2, 4, 6].into_array(); let primitive_result = bitpacked.take(indices).unwrap(); - assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([0u32, 2, 4, 6])); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([0u32, 2, 4, 6]), + &mut ctx + ); } #[test] fn take_sliced_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![1919, 1921].into_array(); // Create a u8 array modulo 63. @@ -217,13 +230,17 @@ mod test { let sliced = bitpacked.slice(128..2050).unwrap(); let primitive_result = sliced.take(indices).unwrap(); - assert_arrays_eq!(primitive_result, PrimitiveArray::from_iter([31u8, 33])); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([31u8, 33]), + &mut ctx + ); } #[test] #[cfg_attr(miri, ignore)] // This test is too slow on miri fn take_random_indices() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let num_patches: usize = 128; let values = (0..u16::MAX as u32 + num_patches as u32).collect::>(); let uncompressed = PrimitiveArray::new(values.clone(), Validity::NonNullable); @@ -256,7 +273,7 @@ mod test { #[test] #[cfg_attr(miri, ignore)] fn take_signed_with_patches() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let start = BitPackedData::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx) .unwrap(); @@ -268,12 +285,16 @@ mod test { &mut ctx, ) .unwrap(); - assert_arrays_eq!(taken_primitive, PrimitiveArray::from_iter([1i32, 2, 3, 4])); + assert_arrays_eq!( + taken_primitive, + PrimitiveArray::from_iter([1i32, 2, 3, 4]), + &mut ctx + ); } #[test] fn take_nullable_with_nullables() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let start = BitPackedData::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx) .unwrap(); @@ -285,19 +306,15 @@ mod test { .unwrap(); assert_arrays_eq!( taken_primitive, - PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]) + PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]), + &mut ctx ); let taken_primitive_prim = taken_primitive.execute::(&mut ctx).unwrap(); assert_eq!(taken_primitive_prim.invalid_count(&mut ctx).unwrap(), 1); } fn bp(array: vortex_array::ArrayRef, bit_width: u8) -> BitPackedArray { - BitPackedData::encode( - &array, - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap() + BitPackedData::encode(&array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() } #[rstest] diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index 92ff63f4408..e14b27323c1 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -34,7 +34,6 @@ mod test { use vortex_array::ArrayRef; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SliceArray; @@ -56,10 +55,10 @@ mod test { use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; - use crate::test::SESSION as FASTLANES_SESSION; + use crate::test::SESSION; fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray { - BitPackedData::encode(array, bit_width, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() + BitPackedData::encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() } fn slice_via_reduce(array: &BitPackedArray, range: Range) -> BitPackedArray { @@ -79,8 +78,13 @@ mod test { 6, ); let sliced = slice_via_reduce(&arr, 1024..2048); - assert_nth_scalar!(sliced, 0, 1024u32 % 64); - assert_nth_scalar!(sliced, 1023, 2047u32 % 64); + assert_nth_scalar!(sliced, 0, 1024u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 1023, + 2047u32 % 64, + &mut SESSION.create_execution_ctx() + ); assert_eq!(sliced.offset(), 0); assert_eq!(sliced.len(), 1024); } @@ -92,8 +96,13 @@ mod test { 6, ); let sliced = slice_via_reduce(&arr, 512..1434); - assert_nth_scalar!(sliced, 0, 512u32 % 64); - assert_nth_scalar!(sliced, 921, 1433u32 % 64); + assert_nth_scalar!(sliced, 0, 512u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 921, + 1433u32 % 64, + &mut SESSION.create_execution_ctx() + ); assert_eq!(sliced.offset(), 512); assert_eq!(sliced.len(), 922); } @@ -106,8 +115,18 @@ mod test { ); let compressed = packed.slice(768..9999).unwrap(); - assert_nth_scalar!(compressed, 0, (768 % 63) as u8); - assert_nth_scalar!(compressed, compressed.len() - 1, (9998 % 63) as u8); + assert_nth_scalar!( + compressed, + 0, + (768 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + compressed, + compressed.len() - 1, + (9998 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -118,8 +137,18 @@ mod test { ); let compressed = packed.slice(7168..9216).unwrap(); - assert_nth_scalar!(compressed, 0, (7168 % 63) as u8); - assert_nth_scalar!(compressed, compressed.len() - 1, (9215 % 63) as u8); + assert_nth_scalar!( + compressed, + 0, + (7168 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + compressed, + compressed.len() - 1, + (9215 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -129,20 +158,35 @@ mod test { 6, ); let sliced = slice_via_reduce(&arr, 512..1434); - assert_nth_scalar!(sliced, 0, 512u32 % 64); - assert_nth_scalar!(sliced, 921, 1433u32 % 64); + assert_nth_scalar!(sliced, 0, 512u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 921, + 1433u32 % 64, + &mut SESSION.create_execution_ctx() + ); assert_eq!(sliced.offset(), 512); assert_eq!(sliced.len(), 922); let doubly_sliced = slice_via_reduce(&sliced, 127..911); - assert_nth_scalar!(doubly_sliced, 0, (512u32 + 127) % 64); - assert_nth_scalar!(doubly_sliced, 783, (512u32 + 910) % 64); + assert_nth_scalar!( + doubly_sliced, + 0, + (512u32 + 127) % 64, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + doubly_sliced, + 783, + (512u32 + 910) % 64, + &mut SESSION.create_execution_ctx() + ); assert_eq!(doubly_sliced.offset(), 639); assert_eq!(doubly_sliced.len(), 784); } #[test] fn slice_empty_patches() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // We create an array that has 1 element that does not fit in the 6-bit range. let array = BitPackedData::encode(&buffer![0u32..=64].into_array(), 6, &mut ctx).unwrap(); @@ -156,7 +200,7 @@ mod test { // resulting array should have no patches. let array_ref = array.into_array(); let slice_array = SliceArray::new(array_ref, 0..64); - let mut ctx = FASTLANES_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let sliced_bp = slice_array .into_array() .execute::(&mut ctx) @@ -217,7 +261,7 @@ mod test { .into_array(); assert_eq!( packed_array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap(), Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)) ); @@ -225,7 +269,7 @@ mod test { #[test] fn scalar_at() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = (0u32..257).collect::>(); let uncompressed = values.clone().into_array(); let packed = BitPackedData::encode(&uncompressed, 8, &mut ctx).unwrap(); @@ -235,7 +279,7 @@ mod test { assert_eq!( usize::try_from( &patches - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() ) .unwrap(), @@ -243,6 +287,6 @@ mod test { ); let expected = PrimitiveArray::from_iter(values.iter().copied()); - assert_arrays_eq!(packed, expected); + assert_arrays_eq!(packed, expected, &mut ctx); } } diff --git a/encodings/fastlanes/src/delta/array/delta_compress.rs b/encodings/fastlanes/src/delta/array/delta_compress.rs index fd17d481f43..7f00eb172dc 100644 --- a/encodings/fastlanes/src/delta/array/delta_compress.rs +++ b/encodings/fastlanes/src/delta/array/delta_compress.rs @@ -115,7 +115,11 @@ mod tests { use crate::delta::array::delta_decompress::delta_decompress; use crate::delta_compress; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[rstest] #[case::u32((0u32..10_000).collect())] @@ -147,7 +151,7 @@ mod tests { let delta = Delta::try_from_primitive_array(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(delta.len(), array.len()); let decompressed = delta_decompress(&delta, &mut SESSION.create_execution_ctx())?; - assert_arrays_eq!(decompressed, array); + assert_arrays_eq!(decompressed, array, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -173,7 +177,7 @@ mod tests { .as_array() .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(packed_delta_prim, array); + assert_arrays_eq!(packed_delta_prim, array, &mut ctx); Ok(()) } diff --git a/encodings/fastlanes/src/delta/compute/cast.rs b/encodings/fastlanes/src/delta/compute/cast.rs index 5a9167f126f..2ee10e7619a 100644 --- a/encodings/fastlanes/src/delta/compute/cast.rs +++ b/encodings/fastlanes/src/delta/compute/cast.rs @@ -56,7 +56,11 @@ mod tests { use vortex_session::VortexSession; use crate::Delta; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_cast_delta_unsigned_widening_wraps() { @@ -75,7 +79,11 @@ mod tests { .cast(DType::Primitive(PType::U32, Nullability::NonNullable)) .unwrap(); - assert_arrays_eq!(casted, PrimitiveArray::from_iter([200u32, 50, 75, 10, 255])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([200u32, 50, 75, 10, 255]), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -95,7 +103,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_iter([10f32, 20.0, 30.0, 40.0, 50.0]) + PrimitiveArray::from_iter([10f32, 20.0, 30.0, 40.0, 50.0]), + &mut SESSION.create_execution_ctx() ); } @@ -115,7 +124,8 @@ mod tests { ); assert_arrays_eq!( casted, - PrimitiveArray::from_option_iter([Some(10u32), Some(20), Some(5), Some(30), Some(15)]) + PrimitiveArray::from_option_iter([Some(10u32), Some(20), Some(5), Some(30), Some(15)]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -134,7 +144,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_option_iter([Some(10u32), None, Some(30), Some(15), None]) + PrimitiveArray::from_option_iter([Some(10u32), None, Some(30), Some(15), None]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -177,7 +188,11 @@ mod tests { ); // Verify by decoding - assert_arrays_eq!(casted, PrimitiveArray::from_iter([10u32, 20, 30, 40, 50])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([10u32, 20, 30, 40, 50]), + &mut SESSION.create_execution_ctx() + ); } #[test] diff --git a/encodings/fastlanes/src/delta/vtable/operations.rs b/encodings/fastlanes/src/delta/vtable/operations.rs index 53c5d726158..60bf86db97e 100644 --- a/encodings/fastlanes/src/delta/vtable/operations.rs +++ b/encodings/fastlanes/src/delta/vtable/operations.rs @@ -43,7 +43,11 @@ mod tests { use crate::Delta; use crate::DeltaArray; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn da(array: &PrimitiveArray) -> DeltaArray { Delta::try_from_primitive_array(array, &mut SESSION.create_execution_ctx()) @@ -56,7 +60,7 @@ mod tests { let actual = delta.slice(10..250).unwrap(); let expected = PrimitiveArray::from_iter(10u32..250).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -65,7 +69,7 @@ mod tests { let actual = delta.slice(1024 + 10..1024 + 250).unwrap(); let expected = PrimitiveArray::from_iter((1024 + 10u32)..(1024 + 250)).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -74,7 +78,7 @@ mod tests { let actual = delta.slice(1000..1048).unwrap(); let expected = PrimitiveArray::from_iter(1000u32..1048).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -83,7 +87,7 @@ mod tests { let actual = delta.slice(2040..2050).unwrap(); let expected = PrimitiveArray::from_iter(2040u32..2050).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -92,7 +96,7 @@ mod tests { let actual = delta.slice(0..4096).unwrap(); let expected = PrimitiveArray::from_iter(0u32..4096).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -101,15 +105,15 @@ mod tests { let actual = delta.slice(0..0).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); let actual = delta.slice(4096..4096).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); let actual = delta.slice(1024..1024).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -118,7 +122,7 @@ mod tests { let actual = delta.slice(1024 + 10..1024 + 250).unwrap(); let expected = PrimitiveArray::from_iter((1024 + 10u32)..(1024 + 250)).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -127,15 +131,15 @@ mod tests { let actual = delta.slice(0..0).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); let actual = delta.slice(4000..4000).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); let actual = delta.slice(1024..1024).unwrap(); let expected = PrimitiveArray::from_iter(Vec::::new()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -146,7 +150,7 @@ mod tests { let sliced_again = sliced.slice(0..2).unwrap(); let expected = PrimitiveArray::from_iter(vec![10u32, 11]).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -157,7 +161,7 @@ mod tests { let sliced_again = sliced.slice(0..2).unwrap(); let expected = PrimitiveArray::from_iter(vec![10u32, 11]).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -168,7 +172,7 @@ mod tests { let sliced_again = sliced.slice(0..2).unwrap(); let expected = PrimitiveArray::from_iter(vec![1034u32, 1035]).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -179,7 +183,7 @@ mod tests { let sliced_again = sliced.slice(0..2).unwrap(); let expected = PrimitiveArray::from_iter(vec![1034u32, 1035]).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -190,7 +194,7 @@ mod tests { let sliced_again = sliced.slice(5..20).unwrap(); let expected = PrimitiveArray::from_iter(1015u32..1030).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -201,7 +205,7 @@ mod tests { let sliced_again = sliced.slice(5..20).unwrap(); let expected = PrimitiveArray::from_iter(1015u32..1030).into_array(); - assert_arrays_eq!(sliced_again, expected); + assert_arrays_eq!(sliced_again, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -209,7 +213,7 @@ mod tests { let delta = da(&(0u32..2048).collect()).into_array(); let expected = PrimitiveArray::from_iter(0u32..2048).into_array(); - assert_arrays_eq!(delta, expected); + assert_arrays_eq!(delta, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -225,7 +229,7 @@ mod tests { let delta = da(&(0u32..2000).collect()).into_array(); let expected = PrimitiveArray::from_iter(0u32..2000).into_array(); - assert_arrays_eq!(delta, expected); + assert_arrays_eq!(delta, expected, &mut SESSION.create_execution_ctx()); } #[test] diff --git a/encodings/fastlanes/src/for/array/for_compress.rs b/encodings/fastlanes/src/for/array/for_compress.rs index 48efe87e834..e39b4924e78 100644 --- a/encodings/fastlanes/src/for/array/for_compress.rs +++ b/encodings/fastlanes/src/for/array/for_compress.rs @@ -68,7 +68,11 @@ mod test { use crate::r#for::array::for_decompress::decompress; use crate::r#for::array::for_decompress::fused_decompress; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_compress_round_trip_small() { @@ -79,7 +83,7 @@ mod test { let compressed = FoRData::encode(array.clone()).unwrap(); assert_eq!(i32::try_from(compressed.reference_scalar()).unwrap(), 1); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -121,7 +125,7 @@ mod test { // Create a range offset by a million. let array = PrimitiveArray::from_iter((0u32..100_000).step_by(1024).map(|v| v + 1_000_000)); let compressed = FoRData::encode(array.clone()).unwrap(); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut SESSION.create_execution_ctx()); } #[test] @@ -132,7 +136,7 @@ mod test { let array = PrimitiveArray::from_iter((0u32..1024).map(|x| x % 7)); let bp = BitPackedData::encode(&array.into_array(), 3, &mut ctx).unwrap(); let compressed = FoR::try_new(bp.into_array(), 10u32.into()).unwrap(); - assert_arrays_eq!(compressed, expect); + assert_arrays_eq!(compressed, expect, &mut ctx); } #[test] @@ -144,7 +148,7 @@ mod test { let bp = BitPackedData::encode(&array.into_array(), 2, &mut ctx)?; let compressed = FoR::try_new(bp.clone().into_array(), 10u32.into())?; let decompressed = fused_decompress::(&compressed, bp.as_view(), &mut ctx)?; - assert_arrays_eq!(decompressed, expect); + assert_arrays_eq!(decompressed, expect, &mut ctx); Ok(()) } @@ -182,7 +186,7 @@ mod test { i8::try_from(&compressed.execute_scalar(i, &mut ctx).unwrap()).unwrap() ); }); - assert_arrays_eq!(decompressed, array); + assert_arrays_eq!(decompressed, array, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/for/compute/cast.rs b/encodings/fastlanes/src/for/compute/cast.rs index e534f3c76fa..b865af43d12 100644 --- a/encodings/fastlanes/src/for/compute/cast.rs +++ b/encodings/fastlanes/src/for/compute/cast.rs @@ -30,9 +30,12 @@ impl CastReduce for FoR { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; @@ -43,10 +46,17 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexExpect; + use vortex_session::VortexSession; use crate::FoR; use crate::FoRArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn for_arr(encoded: ArrayRef, reference: Scalar) -> FoRArray { FoR::try_new(encoded, reference).vortex_expect("FoR array construction should succeed") } @@ -70,7 +80,8 @@ mod tests { // Verify the values after decoding assert_arrays_eq!( casted, - PrimitiveArray::from_iter([100i64, 110, 120, 130, 140]) + PrimitiveArray::from_iter([100i64, 110, 120, 130, 140]), + &mut SESSION.create_execution_ctx() ); } diff --git a/encodings/fastlanes/src/for/compute/compare.rs b/encodings/fastlanes/src/for/compute/compare.rs index 8cd2fb9c011..384bdee8037 100644 --- a/encodings/fastlanes/src/for/compute/compare.rs +++ b/encodings/fastlanes/src/for/compute/compare.rs @@ -90,18 +90,28 @@ where #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_session::VortexSession; use super::*; use crate::FoR; use crate::FoRArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn for_arr(encoded: ArrayRef, reference: Scalar) -> FoRArray { FoR::try_new(encoded, reference).vortex_expect("FoR array construction should succeed") } @@ -123,7 +133,11 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([false, true, false].map(Some))); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, true, false].map(Some)), + &mut SESSION.create_execution_ctx() + ); let result = compare_constant( lhs.as_view(), @@ -133,7 +147,11 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([true, true, false].map(Some))); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, true, false].map(Some)), + &mut SESSION.create_execution_ctx() + ); for op in [ CompareOperator::Lt, @@ -203,7 +221,8 @@ mod tests { .unwrap(); assert_arrays_eq!( result, - BoolArray::from_iter([false, false, false].map(Some)) + BoolArray::from_iter([false, false, false].map(Some)), + &mut SESSION.create_execution_ctx() ); let result = compare_constant( @@ -214,7 +233,11 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([true, true, true].map(Some))); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, true, true].map(Some)), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -237,7 +260,11 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([Some(false), Some(true)])); + assert_arrays_eq!( + result, + BoolArray::from_iter([Some(false), Some(true)]), + &mut SESSION.create_execution_ctx() + ); let result = compare_constant( lhs.as_view(), @@ -247,6 +274,10 @@ mod tests { ) .unwrap() .unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([Some(true), Some(false)])); + assert_arrays_eq!( + result, + BoolArray::from_iter([Some(true), Some(false)]), + &mut SESSION.create_execution_ctx() + ); } } diff --git a/encodings/fastlanes/src/for/vtable/operations.rs b/encodings/fastlanes/src/for/vtable/operations.rs index 5fea6f4bdaa..b6deefbe038 100644 --- a/encodings/fastlanes/src/for/vtable/operations.rs +++ b/encodings/fastlanes/src/for/vtable/operations.rs @@ -40,15 +40,25 @@ impl OperationsVTable for FoR { #[cfg(test)] mod test { + use std::sync::LazyLock; + + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_session::VortexSession; use crate::FoRData; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn for_scalar_at() { let for_arr = FoRData::encode(PrimitiveArray::from_iter([-100, 1100, 1500, 1900])).unwrap(); let expected = PrimitiveArray::from_iter([-100, 1100, 1500, 1900]); - assert_arrays_eq!(for_arr, expected); + assert_arrays_eq!(for_arr, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/fastlanes/src/rle/array/mod.rs b/encodings/fastlanes/src/rle/array/mod.rs index 8793179c7d3..7ddff474836 100644 --- a/encodings/fastlanes/src/rle/array/mod.rs +++ b/encodings/fastlanes/src/rle/array/mod.rs @@ -140,7 +140,6 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -274,7 +273,7 @@ mod tests { let invalid_slice = rle_array .slice(2..5) .unwrap() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut SESSION.create_execution_ctx()) .unwrap() .into_primitive(); let mut ctx = SESSION.create_execution_ctx(); @@ -311,13 +310,10 @@ mod tests { let validity_mask = sliced_array .validity() .unwrap() - .execute_mask( - sliced_array.len(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) + .execute_mask(sliced_array.len(), &mut SESSION.create_execution_ctx()) .unwrap(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let expected_mask = Validity::from_iter([false, true, false]) .execute_mask(3, &mut ctx) .unwrap(); @@ -347,7 +343,7 @@ mod tests { #[test] fn test_multi_chunk_two_chunks() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = PrimitiveArray::from_iter([10u32, 20, 30, 40]).into_array(); let indices = PrimitiveArray::from_iter([0u16, 1].repeat(1024)).into_array(); let values_idx_offsets = PrimitiveArray::from_iter([0u64, 2]).into_array(); @@ -395,7 +391,11 @@ mod tests { let decoded_data = decoded.execute::(&mut exec_ctx)?; - assert_arrays_eq!(original_data, decoded_data); + assert_arrays_eq!( + original_data, + decoded_data, + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -442,7 +442,11 @@ mod tests { .execute::(&mut exec_ctx)?; let decoded_data = decoded.execute::(&mut exec_ctx)?; - assert_arrays_eq!(original_data, decoded_data); + assert_arrays_eq!( + original_data, + decoded_data, + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -493,7 +497,7 @@ mod tests { .as_array() .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/rle/array/rle_compress.rs b/encodings/fastlanes/src/rle/array/rle_compress.rs index 107c912c23f..b7e04a13699 100644 --- a/encodings/fastlanes/src/rle/array/rle_compress.rs +++ b/encodings/fastlanes/src/rle/array/rle_compress.rs @@ -153,10 +153,11 @@ fn padded_validity(array: &PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResu #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::MaskedArray; @@ -166,13 +167,20 @@ mod tests { use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use super::*; use crate::rle::array::RLEArrayExt; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_encode_decode() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // u8 let array_u8: Buffer = buffer![1, 1, 2, 2, 3, 3]; let encoded_u8 = RLEData::encode( @@ -184,7 +192,7 @@ mod tests { .clone() .execute::(&mut ctx)?; let expected_u8 = PrimitiveArray::from_iter(vec![1u8, 1, 2, 2, 3, 3]); - assert_arrays_eq!(decoded_u8, expected_u8); + assert_arrays_eq!(decoded_u8, expected_u8, &mut ctx); // u16 let array_u16: Buffer = buffer![100, 100, 200, 200]; @@ -197,7 +205,7 @@ mod tests { .clone() .execute::(&mut ctx)?; let expected_u16 = PrimitiveArray::from_iter(vec![100u16, 100, 200, 200]); - assert_arrays_eq!(decoded_u16, expected_u16); + assert_arrays_eq!(decoded_u16, expected_u16, &mut ctx); // u64 let array_u64: Buffer = buffer![1000, 1000, 2000]; @@ -210,13 +218,13 @@ mod tests { .clone() .execute::(&mut ctx)?; let expected_u64 = PrimitiveArray::from_iter(vec![1000u64, 1000, 2000]); - assert_arrays_eq!(decoded_u64, expected_u64); + assert_arrays_eq!(decoded_u64, expected_u64, &mut ctx); Ok(()) } #[test] fn test_length() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = buffer![1, 1, 2, 2, 2, 3]; let encoded = RLEData::encode( PrimitiveArray::new(values, Validity::NonNullable).as_view(), @@ -228,7 +236,7 @@ mod tests { #[test] fn test_empty_length() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = Buffer::empty(); let encoded = RLEData::encode( PrimitiveArray::new(values, Validity::NonNullable).as_view(), @@ -242,7 +250,7 @@ mod tests { #[test] fn test_single_value() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = vec![42; 2000].into_iter().collect(); let encoded = RLEData::encode( @@ -256,13 +264,13 @@ mod tests { .clone() .execute::(&mut ctx)?; // Verify round-trip let expected = PrimitiveArray::from_iter(vec![42u16; 2000]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn test_all_different() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = (0u8..=255).collect(); let encoded = RLEData::encode( @@ -276,13 +284,13 @@ mod tests { .clone() .execute::(&mut ctx)?; // Verify round-trip let expected = PrimitiveArray::from_iter((0u8..=255).collect::>()); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn test_partial_last_chunk() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test array with partial last chunk (not divisible by 1024) let values: Buffer = (0..1500).map(|i| (i / 100) as u32).collect(); let array = PrimitiveArray::new(values, Validity::NonNullable); @@ -290,14 +298,14 @@ mod tests { let encoded = RLEData::encode(array.as_view(), &mut ctx).unwrap(); assert_eq!(encoded.len(), 1500); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut ctx); // 2 chunks: 1024 + 476 elements assert_eq!(encoded.values_idx_offsets().len(), 2); } #[test] fn test_two_full_chunks() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Array that spans exactly 2 chunks (2048 elements) let values: Buffer = (0..2048).map(|i| (i / 100) as u32).collect(); let array = PrimitiveArray::new(values, Validity::NonNullable); @@ -305,7 +313,7 @@ mod tests { let encoded = RLEData::encode(array.as_view(), &mut ctx).unwrap(); assert_eq!(encoded.len(), 2048); - assert_arrays_eq!(encoded, array); + assert_arrays_eq!(encoded, array, &mut ctx); assert_eq!(encoded.values_idx_offsets().len(), 2); } @@ -324,7 +332,7 @@ mod tests { fn test_roundtrip_primitive_types( #[case] values: Buffer, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive = values .clone() .into_array() @@ -335,7 +343,7 @@ mod tests { .clone() .execute::(&mut ctx)?; let expected = PrimitiveArray::new(values, primitive.validity()?); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } @@ -366,44 +374,44 @@ mod tests { #[test] fn test_encode_all_null_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec> = vec![None; FL_CHUNK_SIZE]; let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let decoded = with_masked_constant_indices(&rle, &mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } #[test] fn test_encode_all_null_chunk_then_value_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // First chunk is entirely null, second chunk has a value preceded by nulls. let mut values: Vec> = vec![None; 2 * FL_CHUNK_SIZE]; values[FL_CHUNK_SIZE + 100] = Some(42); let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let decoded = with_masked_constant_indices(&rle, &mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } #[test] fn test_encode_one_value_near_end() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Single distinct value near the end of the chunk. let mut values: Vec> = vec![None; FL_CHUNK_SIZE]; values[1000] = Some(42); let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let decoded = with_masked_constant_indices(&rle, &mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } #[test] fn test_encode_value_chunk_then_all_null_remainder() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // 1085 elements (2 chunks: 1024 + 61 padded to 1024). // Chunk 0 has -1i16 at scattered positions (273..=366), rest null. // Chunk 1 (the remainder) is entirely null. @@ -420,7 +428,7 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let decoded = with_masked_constant_indices(&rle, &mut ctx)?; - assert_arrays_eq!(decoded, original); + assert_arrays_eq!(decoded, original, &mut ctx); Ok(()) } @@ -463,18 +471,18 @@ mod tests { #[test] fn test_random_invalid_indices_all_null_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec> = vec![None; FL_CHUNK_SIZE]; let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } #[test] fn test_random_invalid_indices_sparse_values() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values: Vec> = vec![None; FL_CHUNK_SIZE]; values[0] = Some(10); values[500] = Some(20); @@ -482,13 +490,13 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } #[test] fn test_random_invalid_indices_multi_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Two chunks: first has scattered values, second is all null. let mut values: Vec> = vec![None; 2 * FL_CHUNK_SIZE]; values[0] = Some(10); @@ -497,13 +505,13 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } #[test] fn test_random_invalid_indices_partial_last_chunk() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // 1085 elements: chunk 0 has values at scattered positions, chunk 1 is // a partial (61 elements padded to 1024) that is entirely null. let mut values: Vec> = vec![None; 1085]; @@ -513,13 +521,13 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } #[test] fn test_random_invalid_indices_mostly_valid() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Most positions are valid, only a few are null with garbage indices. let mut values: Vec> = (0..FL_CHUNK_SIZE).map(|i| Some((i / 100) as u64)).collect(); @@ -530,7 +538,7 @@ mod tests { let original = PrimitiveArray::from_option_iter(values); let rle = RLEData::encode(original.as_view(), &mut ctx)?; let clobbered = with_random_invalid_indices(&rle, &mut ctx)?; - assert_arrays_eq!(clobbered, original); + assert_arrays_eq!(clobbered, original, &mut ctx); Ok(()) } @@ -543,11 +551,11 @@ mod tests { fn test_float_zeros( #[case] values: Vec, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive = PrimitiveArray::from_iter(values); let rle = RLEData::encode(primitive.as_view(), &mut ctx)?; let decoded = rle.as_array().clone().execute::(&mut ctx)?; - assert_arrays_eq!(primitive, decoded); + assert_arrays_eq!(primitive, decoded, &mut ctx); Ok(()) } } diff --git a/encodings/fastlanes/src/rle/compute/cast.rs b/encodings/fastlanes/src/rle/compute/cast.rs index 4f5fec886f8..92947d927bc 100644 --- a/encodings/fastlanes/src/rle/compute/cast.rs +++ b/encodings/fastlanes/src/rle/compute/cast.rs @@ -49,7 +49,6 @@ mod tests { use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -65,7 +64,11 @@ mod tests { use crate::RLEData; use crate::rle::RLEArray; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn rle(primitive: &PrimitiveArray, ctx: &mut ExecutionCtx) -> RLEArray { RLEData::encode(primitive.as_view(), ctx).unwrap() @@ -84,7 +87,11 @@ mod tests { .into_array() .cast(DType::Primitive(PType::U16, Nullability::NonNullable)) .unwrap(); - assert_arrays_eq!(casted, PrimitiveArray::from_iter([10u16, 20, 30, 40, 50])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([10u16, 20, 30, 40, 50]), + &mut ctx + ); } #[test] @@ -153,7 +160,7 @@ mod tests { ) )] fn test_cast_rle_conformance(#[case] primitive: PrimitiveArray) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let rle_array = rle(&primitive, &mut ctx); test_cast_conformance(&rle_array.into_array()); } diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index 0441a1e61b9..b8f79db95a6 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -41,20 +41,29 @@ impl OperationsVTable for RLE { #[cfg(test)] mod tests { + + use std::sync::LazyLock; + use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_session::VortexSession; use super::*; use crate::RLE; use crate::RLEArray; use crate::RLEData; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + mod fixture { use super::*; @@ -122,7 +131,11 @@ mod tests { let array = fixture::rle_array(); let expected = PrimitiveArray::from_iter([10u32, 10, 20, 20, 20, 30, 10]); - assert_arrays_eq!(array.slice(0..7).unwrap(), expected); + assert_arrays_eq!( + array.slice(0..7).unwrap(), + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -139,7 +152,11 @@ mod tests { Some(30), Some(10), ]); - assert_arrays_eq!(array.slice(0..7).unwrap(), expected); + assert_arrays_eq!( + array.slice(0..7).unwrap(), + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -151,7 +168,7 @@ mod tests { assert_eq!(sliced.len(), 4); let expected = PrimitiveArray::from_iter([20u32, 20, 20, 30]); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -163,12 +180,12 @@ mod tests { assert_eq!(sliced.len(), 4); let expected = PrimitiveArray::from_option_iter([Some(20u32), Some(20), None, Some(30)]); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); } #[test] fn test_scalar_at_multiple_chunks() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test accessing elements around chunk boundaries let values: Buffer = (0..3000).map(|i| (i / 50) as u16).collect(); let expected: Vec = (0..3000).map(|i| (i / 50) as u16).collect(); @@ -197,7 +214,7 @@ mod tests { fn test_scalar_at_out_of_bounds() { let array = fixture::rle_array(); array - .execute_scalar(1025, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1025, &mut SESSION.create_execution_ctx()) .unwrap(); } @@ -206,7 +223,7 @@ mod tests { fn test_scalar_at_slice_out_of_bounds() { let array = fixture::rle_array().slice(0..1).unwrap(); array - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap(); } @@ -224,7 +241,11 @@ mod tests { Some(30), Some(10), ]); - assert_arrays_eq!(sliced.into_array(), expected.into_array()); + assert_arrays_eq!( + sliced.into_array(), + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -233,7 +254,7 @@ mod tests { let sliced = array.slice(4..6).unwrap(); // [20, 30] let expected = buffer![20u32, 30].into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -242,7 +263,7 @@ mod tests { let sliced = array.slice(5..6).unwrap(); // [30] let expected = buffer![30u32].into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -259,12 +280,16 @@ mod tests { let sliced = array.slice(1..4).unwrap(); // [null, 20, 20] let expected = PrimitiveArray::from_option_iter([Option::::None, Some(20), Some(20)]); - assert_arrays_eq!(sliced.into_array(), expected.into_array()); + assert_arrays_eq!( + sliced.into_array(), + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] fn test_slice_decode_with_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = fixture::rle_array_with_nulls(); let sliced = array .slice(1..4) @@ -273,7 +298,7 @@ mod tests { .unwrap(); // [null, 20, 20] let expected = PrimitiveArray::from_option_iter([Option::::None, Some(20), Some(20)]); - assert_arrays_eq!(sliced.into_array(), expected.into_array()); + assert_arrays_eq!(sliced.into_array(), expected.into_array(), &mut ctx); } #[test] @@ -286,7 +311,7 @@ mod tests { #[test] fn test_slice_across_chunk_boundaries() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Buffer = (0..2100).map(|i| (i / 100) as u32).collect(); let expected: Vec = (0..2100).map(|i| (i / 100) as u32).collect(); let array = values.into_array(); @@ -298,14 +323,16 @@ mod tests { let slice = encoded.slice(500..1500).unwrap(); assert_arrays_eq!( slice, - PrimitiveArray::from_iter(expected[500..1500].iter().copied()) + PrimitiveArray::from_iter(expected[500..1500].iter().copied()), + &mut ctx ); // Slice across second and third chunk. let slice = encoded.slice(1000..2000).unwrap(); assert_arrays_eq!( slice, - PrimitiveArray::from_iter(expected[1000..2000].iter().copied()) + PrimitiveArray::from_iter(expected[1000..2000].iter().copied()), + &mut ctx ); } } diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index 732e1a446dd..ec63af3735f 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -122,8 +122,9 @@ fn compare_fsst_constant( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; @@ -135,14 +136,21 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::fsst_compress; use crate::fsst_train_compressor; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] #[cfg_attr(miri, ignore)] fn test_compare_fsst() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let lhs = VarBinArray::from_iter( [ Some("hello"), @@ -170,7 +178,8 @@ mod tests { assert_arrays_eq!( &equals, - &BoolArray::from_iter([Some(false), None, Some(true), None, Some(false)]) + &BoolArray::from_iter([Some(false), None, Some(true), None, Some(false)]), + &mut ctx ); // Ensure fastpath for Eq exists, and returns correct answer @@ -183,7 +192,8 @@ mod tests { assert_eq!(not_equals.dtype(), &DType::Bool(Nullability::Nullable)); assert_arrays_eq!( ¬_equals, - &BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]) + &BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]), + &mut ctx ); // Ensure null constants are handled correctly. @@ -195,7 +205,8 @@ mod tests { .binary(null_rhs.clone().into_array(), Operator::Eq)?; assert_arrays_eq!( &equals_null, - &BoolArray::from_iter([None::, None, None, None, None]) + &BoolArray::from_iter([None::, None, None, None, None]), + &mut ctx ); let noteq_null = lhs @@ -203,7 +214,8 @@ mod tests { .binary(null_rhs.into_array(), Operator::NotEq)?; assert_arrays_eq!( ¬eq_null, - &BoolArray::from_iter([None::, None, None, None, None]) + &BoolArray::from_iter([None::, None, None, None, None]), + &mut ctx ); Ok(()) } diff --git a/encodings/fsst/src/compute/like.rs b/encodings/fsst/src/compute/like.rs index 1b0695ecc7f..c922fba088e 100644 --- a/encodings/fsst/src/compute/like.rs +++ b/encodings/fsst/src/compute/like.rs @@ -104,7 +104,11 @@ mod tests { use crate::fsst_compress; use crate::fsst_train_compressor; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn make_fsst(strings: &[Option<&str>], nullability: Nullability) -> FSSTArray { let array = @@ -144,7 +148,8 @@ mod tests { let result = like(fsst, "http%")?; assert_arrays_eq!( &result, - &BoolArray::from_iter([true, true, false, true, false]) + &BoolArray::from_iter([true, true, false, true, false]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -158,7 +163,8 @@ mod tests { let result = like(fsst, "hel%")?; // spellchecker:disable-line assert_arrays_eq!( &result, - &BoolArray::from_iter([Some(true), None, Some(true), None, Some(false)]) + &BoolArray::from_iter([Some(true), None, Some(true), None, Some(false)]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -175,7 +181,11 @@ mod tests { Nullability::NonNullable, ); let result = like(fsst, "%hello%")?; - assert_arrays_eq!(&result, &BoolArray::from_iter([true, true, false, true])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, true, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -191,7 +201,11 @@ mod tests { Nullability::NonNullable, ); let result = like(fsst, "%lazy dog%")?; - assert_arrays_eq!(&result, &BoolArray::from_iter([true, false, true, false])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, false, true, false]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -206,7 +220,11 @@ mod tests { case_insensitive: false, }; let result = run_like(fsst, "%sdf%", opts)?; - assert_arrays_eq!(&result, &BoolArray::from_iter([false, false, true])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([false, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -217,7 +235,11 @@ mod tests { Nullability::NonNullable, ); let result = like(fsst, "%")?; - assert_arrays_eq!(&result, &BoolArray::from_iter([true, true, true])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, true, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -236,7 +258,11 @@ mod tests { let fsst = fsst.as_view(); let result = ::like(fsst, &pattern, LikeOptions::default(), &mut ctx)?; assert!(result.is_some(), "FSST LikeKernel should handle prefix%"); - assert_arrays_eq!(result.unwrap(), BoolArray::from_iter([true, false])); + assert_arrays_eq!( + result.unwrap(), + BoolArray::from_iter([true, false]), + &mut ctx + ); Ok(()) } @@ -253,7 +279,11 @@ mod tests { let fsst = fsst.as_view(); let result = ::like(fsst, &pattern, LikeOptions::default(), &mut ctx)?; assert!(result.is_some(), "FSST LikeKernel should handle %needle%"); - assert_arrays_eq!(result.unwrap(), BoolArray::from_iter([true, false])); + assert_arrays_eq!( + result.unwrap(), + BoolArray::from_iter([true, false]), + &mut ctx + ); Ok(()) } @@ -311,7 +341,8 @@ mod tests { assert!(result.is_some(), "escaped percent prefix should use FSST"); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([true, false, false, false, false, false, false]) + BoolArray::from_iter([true, false, false, false, false, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"\_%", fsst.len()).into_array(); @@ -323,7 +354,8 @@ mod tests { ); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([false, true, false, false, false, false, false]) + BoolArray::from_iter([false, true, false, false, false, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"\\%", fsst.len()).into_array(); @@ -332,7 +364,8 @@ mod tests { assert!(result.is_some(), "escaped backslash prefix should use FSST"); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([false, false, true, false, false, false, false]) + BoolArray::from_iter([false, false, true, false, false, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"%\%%", fsst.len()).into_array(); @@ -341,7 +374,8 @@ mod tests { assert!(result.is_some(), "escaped percent contains should use FSST"); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([true, false, false, true, false, false, false]) + BoolArray::from_iter([true, false, false, true, false, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"%\_%", fsst.len()).into_array(); @@ -353,7 +387,8 @@ mod tests { ); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([false, true, false, false, true, false, false]) + BoolArray::from_iter([false, true, false, false, true, false, false]), + &mut ctx ); let pattern = ConstantArray::new(r"%\\%", fsst.len()).into_array(); @@ -365,7 +400,8 @@ mod tests { ); assert_arrays_eq!( result.unwrap(), - BoolArray::from_iter([false, false, true, false, false, true, false]) + BoolArray::from_iter([false, false, true, false, false, true, false]), + &mut ctx ); Ok(()) @@ -394,7 +430,11 @@ mod tests { direct.is_some(), "14-byte prefixes are now handled by the flat prefix DFA" ); - assert_arrays_eq!(direct.unwrap(), BoolArray::from_iter([true, false, true])); + assert_arrays_eq!( + direct.unwrap(), + BoolArray::from_iter([true, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -424,7 +464,11 @@ mod tests { ); let result = like(fsst, &pattern)?; - assert_arrays_eq!(&result, &BoolArray::from_iter([true, false, true])); + assert_arrays_eq!( + &result, + &BoolArray::from_iter([true, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -451,7 +495,11 @@ mod tests { direct.is_some(), "254-byte contains needle should stay on the DFA path" ); - assert_arrays_eq!(direct.unwrap(), BoolArray::from_iter([true, false, true])); + assert_arrays_eq!( + direct.unwrap(), + BoolArray::from_iter([true, false, true]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } } diff --git a/encodings/fsst/src/dfa/tests.rs b/encodings/fsst/src/dfa/tests.rs index 101eb3989d4..3ea9a507c20 100644 --- a/encodings/fsst/src/dfa/tests.rs +++ b/encodings/fsst/src/dfa/tests.rs @@ -31,7 +31,11 @@ use crate::FSSTArray; use crate::fsst_compress; use crate::fsst_train_compressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); /// Helper: make a Symbol from a byte string (up to 8 bytes, zero-padded). fn sym(bytes: &[u8]) -> Symbol { @@ -356,6 +360,7 @@ fn test_like_edge_cases( ConstantArray::new(pattern, opts.len()).into_array(), )?; let expected_arr = BoolArray::from_iter(expected.iter().copied()); - assert_arrays_eq!(&result, &expected_arr); + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(&result, &expected_arr, &mut ctx); Ok(()) } diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index f137ea95025..30f25006195 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -54,7 +54,11 @@ mod tests { use crate::fsst_compress; use crate::fsst_train_compressor; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); fn build_test_fsst_array() -> ArrayRef { let mut builder = VarBinBuilder::::with_capacity(10); @@ -98,7 +102,7 @@ mod tests { let expected = fsst_array.filter(mask)?; assert_eq!(result.len(), 2); - assert_arrays_eq!(result.into_array(), expected); + assert_arrays_eq!(result.into_array(), expected, &mut ctx); Ok(()) } @@ -118,7 +122,7 @@ mod tests { let expected = fsst_array.filter(mask)?; assert_eq!(result.len(), 5); - assert_arrays_eq!(result.into_array(), expected); + assert_arrays_eq!(result.into_array(), expected, &mut ctx); Ok(()) } @@ -162,7 +166,7 @@ mod tests { let expected = input.filter(mask)?; assert_eq!(result.len(), 1); - assert_arrays_eq!(result.into_array(), expected); + assert_arrays_eq!(result.into_array(), expected, &mut ctx); Ok(()) } @@ -188,7 +192,7 @@ mod tests { let expected = input.filter(mask)?; assert_eq!(result.len(), 2); - assert_arrays_eq!(result.into_array(), expected); + assert_arrays_eq!(result.into_array(), expected, &mut ctx); Ok(()) } @@ -203,7 +207,7 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let result = filter_array.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(result, fsst_array); + assert_arrays_eq!(result, fsst_array, &mut ctx); Ok(()) } @@ -223,7 +227,7 @@ mod tests { let fsst = fsst_compress(&varbin, &compressor, &mut ctx)?.into_array(); let result = fsst.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 7, 18, 0]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } } diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 7cf7e9e9f6c..26e0975da22 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -1,10 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::LazyLock; + use fsst::CompressorBuilder; -use vortex_array::ArrayRef; +use vortex_array::{ArrayRef, ExecutionCtx}; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::builder::VarBinBuilder; @@ -14,13 +15,20 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_buffer::buffer; use vortex_mask::Mask; +use vortex_session::VortexSession; use crate::FSST; use crate::fsst_compress; use crate::fsst_train_compressor; +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + /// this function is VERY slow on miri, so we only want to run it once -pub(crate) fn build_fsst_array() -> ArrayRef { +pub(crate) fn build_fsst_array(ctx: &mut ExecutionCtx) -> ArrayRef { let mut input_array = VarBinBuilder::::with_capacity(3); input_array.append_value(b"The Greeks never said that the limit could not be overstepped"); input_array.append_value( @@ -31,32 +39,34 @@ pub(crate) fn build_fsst_array() -> ArrayRef { .finish(DType::Utf8(Nullability::NonNullable)) .into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let compressor = fsst_train_compressor(&input_array, &mut ctx).unwrap(); - fsst_compress(&input_array, &compressor, &mut ctx) + let compressor = fsst_train_compressor(&input_array, ctx).unwrap(); + fsst_compress(&input_array, &compressor, ctx) .unwrap() .into_array() } #[test] fn test_fsst_array_ops() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // first test the scalar_at values - let fsst_array = build_fsst_array(); + let fsst_array = build_fsst_array(&mut ctx); assert_nth_scalar!( fsst_array, 0, - "The Greeks never said that the limit could not be overstepped" + "The Greeks never said that the limit could not be overstepped", + &mut ctx ); assert_nth_scalar!( fsst_array, 1, - "They said it existed and that whoever dared to exceed it was mercilessly struck down" + "They said it existed and that whoever dared to exceed it was mercilessly struck down", + &mut ctx ); assert_nth_scalar!( fsst_array, 2, - "Nothing in present history can contradict them" + "Nothing in present history can contradict them", + &mut ctx ); // test slice @@ -66,12 +76,14 @@ fn test_fsst_array_ops() { assert_nth_scalar!( fsst_sliced, 0, - "They said it existed and that whoever dared to exceed it was mercilessly struck down" + "They said it existed and that whoever dared to exceed it was mercilessly struck down", + &mut ctx ); assert_nth_scalar!( fsst_sliced, 1, - "Nothing in present history can contradict them" + "Nothing in present history can contradict them", + &mut ctx ); // test take @@ -81,12 +93,14 @@ fn test_fsst_array_ops() { assert_nth_scalar!( fsst_taken, 0, - "The Greeks never said that the limit could not be overstepped" + "The Greeks never said that the limit could not be overstepped", + &mut ctx ); assert_nth_scalar!( fsst_taken, 1, - "Nothing in present history can contradict them" + "Nothing in present history can contradict them", + &mut ctx ); // test filter @@ -98,7 +112,8 @@ fn test_fsst_array_ops() { assert_nth_scalar!( fsst_filtered, 0, - "They said it existed and that whoever dared to exceed it was mercilessly struck down" + "They said it existed and that whoever dared to exceed it was mercilessly struck down", + &mut ctx ); // test to_canonical @@ -108,7 +123,7 @@ fn test_fsst_array_ops() { .unwrap() .into_array(); - assert_arrays_eq!(fsst_array, canonical_array); + assert_arrays_eq!(fsst_array, canonical_array, &mut ctx); } // TODO(someone): ideally CI would run this in release mode as well since debug builds make the @@ -156,7 +171,7 @@ fn fsst_compress_offsets_overflow_i32() { let compressor = CompressorBuilder::default().build(); let len = array.len(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); println!("compressing to FSST"); let compressed = fsst_compress(&array, &compressor, &mut ctx).unwrap(); diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 37247430c37..7be0265fc96 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -385,6 +385,7 @@ impl> ParquetVariantArrayExt for T {} #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::Array as _; use arrow_array::ArrayRef as ArrowArrayRef; @@ -399,7 +400,6 @@ mod tests { use parquet_variant_compute::VariantArray as ArrowVariantArray; use parquet_variant_compute::VariantArrayBuilder; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -410,10 +410,17 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; + use vortex_session::VortexSession; use crate::ParquetVariant; use crate::array::ParquetVariantArrayExt; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn assert_arrow_variant_storage_roundtrip(struct_array: StructArray) -> VortexResult<()> { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; @@ -421,7 +428,7 @@ mod tests { .as_opt::() .ok_or_else(|| vortex_err!("expected parquet variant child"))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let roundtripped = inner.to_arrow(&mut ctx)?; let roundtripped = roundtripped.inner(); @@ -513,10 +520,11 @@ mod tests { .typed_value_array() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut SESSION.create_execution_ctx())?; assert_arrays_eq!( typed_value, - PrimitiveArray::from_option_iter([Some(10), None, Some(30)]) + PrimitiveArray::from_option_iter([Some(10), None, Some(30)]), + &mut SESSION.create_execution_ctx() ); Ok(()) @@ -528,7 +536,7 @@ mod tests { let value = VarBinViewArray::from_iter_bin([b"\x10", b"\x11"]).into_array(); let pv_array = ParquetVariant::try_new(Validity::NonNullable, metadata, Some(value), None)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let variant_arr = pv_array.to_arrow(&mut ctx)?; let struct_arr = variant_arr.inner(); @@ -550,7 +558,7 @@ mod tests { Some(typed_value), )?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let variant_arr = pv_array.to_arrow(&mut ctx)?; let struct_arr = variant_arr.inner(); @@ -599,8 +607,12 @@ mod tests { .typed_value_array() .ok_or_else(|| vortex_err!("expected typed_value child"))? .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; - assert_arrays_eq!(typed_value, PrimitiveArray::from_iter([10i32, 20, 30])); + .execute::(&mut SESSION.create_execution_ctx())?; + assert_arrays_eq!( + typed_value, + PrimitiveArray::from_iter([10i32, 20, 30]), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -629,7 +641,7 @@ mod tests { assert!(parquet_array.value_array().is_some()); assert!(parquet_array.typed_value_array().is_some()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let roundtripped = parquet_array.to_arrow(&mut ctx)?; let roundtripped = roundtripped.inner(); assert_eq!( diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index 760cd1fb1aa..326e7f9d7ee 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -499,7 +499,7 @@ mod tests { .arrow() .from_arrow_array(Arc::clone(&exported), &field)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } @@ -700,7 +700,7 @@ mod tests { let actual = session.arrow().from_arrow_array(exported, &field)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } } diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 5624437d461..684ce0217f4 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -207,6 +207,7 @@ impl TakeExecute for ParquetVariant { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; @@ -226,7 +227,6 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray as VortexStructArray; @@ -256,11 +256,11 @@ mod tests { use crate::ParquetVariant; use crate::ParquetVariantArrayExt; - fn session() -> VortexSession { + static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); crate::initialize(&session); session - } + }); fn make_unshredded_array() -> VortexResult { let mut builder = VariantArrayBuilder::new(4); @@ -382,10 +382,9 @@ mod tests { dtype: Option, ) -> VortexResult { let expr = variant_get(root(), parse_path(path)?, dtype); - let session = session(); array .apply(&expr)? - .execute::(&mut session.create_execution_ctx()) + .execute::(&mut SESSION.create_execution_ctx()) } macro_rules! assert_rows_eq { @@ -395,7 +394,7 @@ mod tests { let expected_rows = [$($expected_idx),*]; assert_eq!(actual.len(), expected_rows.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for (actual_idx, expected_idx) in expected_rows.into_iter().enumerate() { assert_eq!( actual.execute_scalar(actual_idx, &mut ctx)?, @@ -412,7 +411,7 @@ mod tests { let expected = [$($is_null),*]; assert_eq!(array.len(), expected.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for (idx, is_null) in expected.into_iter().enumerate() { assert_eq!( array.execute_scalar(idx, &mut ctx)?.is_null(), @@ -520,7 +519,8 @@ mod tests { result.dtype(), &VortexDType::Primitive(PType::I32, Nullability::Nullable) ); - assert_arrays_eq!(result, PrimitiveArray::from_option_iter(expected)); + let mut ctx = SESSION.create_execution_ctx(); + assert_arrays_eq!(result, PrimitiveArray::from_option_iter(expected), &mut ctx); Ok(()) } @@ -536,7 +536,7 @@ mod tests { let result = execute_variant_get(arr, "$.a", None)?; assert_eq!(result.dtype(), &VortexDType::Variant(Nullability::Nullable)); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let row0 = result.execute_scalar(0, &mut ctx)?; assert_eq!( row0.as_variant() @@ -545,7 +545,7 @@ mod tests { .map(|value| value.as_str()), Some("ok") ); - assert_nth_scalar_is_null!(result, 1); + assert_nth_scalar_is_null!(result, 1, &mut ctx); assert_eq!( result .execute_scalar(2, &mut ctx)? @@ -553,7 +553,7 @@ mod tests { .is_variant_null(), Some(true) ); - assert_nth_scalar_is_null!(result, 3); + assert_nth_scalar_is_null!(result, 3, &mut ctx); Ok(()) } @@ -711,7 +711,7 @@ mod tests { fn make_partially_shredded_object_array() -> VortexResult { let arrow_variant = make_partially_shredded_arrow_variant()?; let parquet_array = ParquetVariant::from_arrow_variant(&arrow_variant)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(canonical) = parquet_array.execute::(&mut ctx)? else { return Err(vortex_err!("expected canonical variant array")); }; @@ -749,7 +749,7 @@ mod tests { fn assert_variant_i32_scalars(array: &ArrayRef, expected: &[Option]) -> VortexResult<()> { assert_eq!(array.len(), expected.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for (idx, expected) in expected.iter().enumerate() { let scalar = array.execute_scalar(idx, &mut ctx)?; let variant = scalar.as_variant(); @@ -775,7 +775,7 @@ mod tests { ) -> VortexResult<()> { assert_eq!(array.len(), expected_a.len()); assert_eq!(array.len(), expected_b.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); for idx in 0..array.len() { let scalar = array.execute_scalar(idx, &mut ctx)?; let object = scalar @@ -834,8 +834,7 @@ mod tests { array: &ArrayRef, expected: impl IntoIterator>, ) -> VortexResult<()> { - let session = session(); - let mut ctx = session.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let executed = array.clone().execute::(&mut ctx)?; let typed_value = executed .as_::() @@ -844,7 +843,11 @@ mod tests { .clone() .execute::(&mut ctx)?; - assert_arrays_eq!(typed_value, PrimitiveArray::from_option_iter(expected)); + assert_arrays_eq!( + typed_value, + PrimitiveArray::from_option_iter(expected), + &mut ctx + ); Ok(()) } @@ -884,7 +887,8 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10), Some(30), None]) + PrimitiveArray::from_option_iter([Some(10), Some(30), None]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -901,7 +905,8 @@ mod tests { assert_arrays_eq!( result, - VarBinArray::from(vec![Some("left"), Some("right"), Some("missing_a")]) + VarBinArray::from(vec![Some("left"), Some("right"), Some("missing_a")]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -922,7 +927,7 @@ mod tests { .shredded() .ok_or_else(|| vortex_err!("expected canonical shredded child"))? .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + .execute::(&mut SESSION.create_execution_ctx())?; assert_eq!( shredded.unmasked_field_by_name("a")?.dtype(), &VortexDType::Variant(Nullability::Nullable) @@ -937,7 +942,8 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10), Some(30), None]) + PrimitiveArray::from_option_iter([Some(10), Some(30), None]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -954,7 +960,8 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10), Some(30), None]) + PrimitiveArray::from_option_iter([Some(10), Some(30), None]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -971,7 +978,8 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(100), Some(20), Some(30)]) + PrimitiveArray::from_option_iter([Some(100), Some(20), Some(30)]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -1015,7 +1023,7 @@ mod tests { .is_some() ); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(canonical) = parquet_array.execute::(&mut ctx)? else { return Err(vortex_err!("expected canonical variant array")); }; diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index d5b98ce03aa..5aae70c7d24 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -284,6 +284,7 @@ impl VTable for ParquetVariant { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_array::Int32Array; @@ -300,7 +301,6 @@ mod tests { use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; @@ -333,6 +333,12 @@ mod tests { use crate::ParquetVariant; use crate::array::ParquetVariantArrayExt; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn roundtrip(array: ArrayRef) -> VortexResult { let dtype = array.dtype().clone(); let len = array.len(); @@ -403,7 +409,7 @@ mod tests { let parquet_variant = ParquetVariant::try_new(Validity::NonNullable, metadata, None, Some(typed_value))?; assert!(parquet_variant.typed_value_array().is_some()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(variant) = parquet_variant .into_array() @@ -427,7 +433,8 @@ mod tests { let shredded = shredded.clone().execute::(&mut ctx)?; assert_arrays_eq!( shredded, - PrimitiveArray::from_option_iter([Some(10), None, Some(30)]) + PrimitiveArray::from_option_iter([Some(10), None, Some(30)]), + &mut ctx ); Ok(()) @@ -456,7 +463,7 @@ mod tests { .read_all() .await?; - assert_arrays_eq!(expected, actual); + assert_arrays_eq!(expected, actual, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -484,7 +491,7 @@ mod tests { .read_all() .await?; - assert_arrays_eq!(expected, actual); + assert_arrays_eq!(expected, actual, &mut SESSION.create_execution_ctx()); Ok(()) } diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index df2d79f63f1..5d073b56fbf 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -696,7 +696,8 @@ mod tests { Some(40), Some(50), None - ]) + ]), + &mut ctx ); // Slice to get only the non-null values in the middle @@ -704,6 +705,6 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(20u32), Some(30), Some(40), Some(50)]) .into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut ctx); } } diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index 89100279a3f..5f562fafca6 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -90,7 +90,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]) + PrimitiveArray::from_iter([1.0f64, 2.0, 3.0, 4.0, 5.0]), + &mut ctx ); } @@ -107,7 +108,8 @@ mod tests { .unwrap(); assert_arrays_eq!( casted, - PrimitiveArray::new(buffer![10u32, 20, 30, 40], Validity::AllValid,) + PrimitiveArray::new(buffer![10u32, 20, 30, 40], Validity::AllValid,), + &mut ctx ); } @@ -128,7 +130,11 @@ mod tests { &DType::Primitive(PType::U32, Nullability::NonNullable) ); // Verify the values are correct - assert_arrays_eq!(casted, PrimitiveArray::from_iter([20u32, 30, 40, 50])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([20u32, 30, 40, 50]), + &mut ctx + ); } #[test] @@ -151,7 +157,11 @@ mod tests { casted.dtype(), &DType::Primitive(PType::U32, Nullability::NonNullable) ); - assert_arrays_eq!(casted, PrimitiveArray::from_iter([20u32, 30, 40, 50])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([20u32, 30, 40, 50]), + &mut ctx + ); } #[rstest] diff --git a/encodings/pco/src/tests.rs b/encodings/pco/src/tests.rs index 32656c9b7fa..692e3f370db 100644 --- a/encodings/pco/src/tests.rs +++ b/encodings/pco/src/tests.rs @@ -53,17 +53,25 @@ fn test_compress_decompress() { compressed.dtype().nullability(), ); let decompressed = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); - assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data), &mut ctx); // check slicing works let slice = compressed.slice(100..105).unwrap(); for i in 0_i32..5 { - assert_nth_scalar!(slice, i as usize, 100 + i); + assert_nth_scalar!(slice, i as usize, 100 + i, &mut ctx); } - assert_arrays_eq!(slice, PrimitiveArray::from_iter([100, 101, 102, 103, 104])); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter([100, 101, 102, 103, 104]), + &mut ctx + ); let slice = compressed.slice(200..200).unwrap(); - assert_arrays_eq!(slice, PrimitiveArray::from_iter(Vec::::new())); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter(Vec::::new()), + &mut ctx + ); } #[test] @@ -73,14 +81,14 @@ fn test_compress_decompress_small() { let compressed = Pco::from_primitive(array.as_view(), 3, 0, &mut ctx).unwrap(); let expected = array.into_array(); - assert_arrays_eq!(compressed, expected); + assert_arrays_eq!(compressed, expected, &mut ctx); let unsliced_validity = child_to_validity( compressed.as_ref().slots()[0].as_ref(), compressed.dtype().nullability(), ); let decompressed = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); - assert_arrays_eq!(decompressed, expected); + assert_arrays_eq!(decompressed, expected, &mut ctx); } #[test] @@ -94,7 +102,7 @@ fn test_empty() { compressed.dtype().nullability(), ); let primitive = compressed.decompress(&unsliced_validity, &mut ctx).unwrap(); - assert_arrays_eq!(primitive, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(primitive, PrimitiveArray::from_iter(data), &mut ctx); } #[test] @@ -128,18 +136,18 @@ fn test_validity_and_multiple_chunks_and_pages() { assert_eq!(compressed.metadata.chunks.len(), 6); // 191 values / 33 rounds up to 6 assert_eq!(compressed.metadata.chunks[0].pages.len(), 4); // 33 / 10 rounds up to 4 - assert_nth_scalar!(compressed, 0, 0); - assert_nth_scalar!(compressed, 3, 3); - assert_nth_scalar!(compressed, 7, None::); - assert_nth_scalar!(compressed, 14, None::); - assert_nth_scalar!(compressed, 15, 15); - assert_nth_scalar!(compressed, 101, None::); - assert_nth_scalar!(compressed, 199, 199); + assert_nth_scalar!(compressed, 0, 0, &mut ctx); + assert_nth_scalar!(compressed, 3, 3, &mut ctx); + assert_nth_scalar!(compressed, 7, None::, &mut ctx); + assert_nth_scalar!(compressed, 14, None::, &mut ctx); + assert_nth_scalar!(compressed, 15, 15, &mut ctx); + assert_nth_scalar!(compressed, 101, None::, &mut ctx); + assert_nth_scalar!(compressed, 199, 199, &mut ctx); // check slicing works let slice = compressed.slice(100..103).unwrap(); - assert_nth_scalar!(slice, 0, 100); - assert_nth_scalar!(slice, 2, 102); + assert_nth_scalar!(slice, 0, 100, &mut ctx); + assert_nth_scalar!(slice, 2, 102, &mut ctx); let primitive = slice.execute::(&mut ctx).unwrap(); assert!( diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index f69e9175222..fde553a2fd4 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -401,12 +401,14 @@ impl RunEndData { /// ``` /// # use vortex_array::arrays::BoolArray; /// # use vortex_array::IntoArray; - /// # use vortex_array::{LEGACY_SESSION, VortexSessionExecute}; + /// # use vortex_array::VortexSessionExecute; /// # use vortex_buffer::buffer; /// # use vortex_error::VortexResult; /// # use vortex_runend::RunEnd; /// # fn main() -> VortexResult<()> { - /// let mut ctx = LEGACY_SESSION.create_execution_ctx(); + /// let session = vortex_array::array_session(); + /// vortex_runend::initialize(&session); + /// let mut ctx = session.create_execution_ctx(); /// let ends = buffer![2u8, 3u8].into_array(); /// let values = BoolArray::from_iter([false, true]).into_array(); /// let run_end = RunEnd::new(ends, values, &mut ctx); @@ -516,7 +518,11 @@ mod tests { use crate::RunEnd; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_runend_constructor() { @@ -536,7 +542,7 @@ mod tests { // 2, 3, 4 => 2 // 5, 6, 7, 8, 9 => 3 let expected = buffer![1, 1, 2, 2, 2, 3, 3, 3, 3, 3].into_array(); - assert_arrays_eq!(arr.into_array(), expected); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } #[test] @@ -550,7 +556,7 @@ mod tests { let expected = VarBinViewArray::from_iter_str(["a", "a", "b", "b", "b", "c", "c", "c", "c", "c"]) .into_array(); - assert_arrays_eq!(arr.into_array(), expected); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } #[test] @@ -571,6 +577,6 @@ mod tests { let expected = VarBinViewArray::from_iter_str(["x", "x", "y", "y", "y", "z", "z", "z", "z", "z"]) .into_array(); - assert_arrays_eq!(arr.into_array(), expected); + assert_arrays_eq!(arr.into_array(), expected, &mut ctx); } } diff --git a/encodings/runend/src/arrow.rs b/encodings/runend/src/arrow.rs index aa03e3be5dc..193937122ef 100644 --- a/encodings/runend/src/arrow.rs +++ b/encodings/runend/src/arrow.rs @@ -81,7 +81,6 @@ mod tests { use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray as _; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute as _; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -95,7 +94,6 @@ mod tests { use vortex_array::scalar::PValue; use vortex_array::search_sorted::SearchSorted; use vortex_array::search_sorted::SearchSortedSide; - use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -107,7 +105,7 @@ mod tests { static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); - session.arrays().register(RunEnd); + crate::initialize(&session); session }); @@ -170,7 +168,8 @@ mod tests { assert_arrays_eq!( vortex_array.into_array(), - buffer![10i32, 10, 10, 20, 20, 30, 30, 30].into_array() + buffer![10i32, 10, 10, 20, 20, 30, 30, 30].into_array(), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -194,7 +193,8 @@ mod tests { None, Some(300i32), Some(300i32) - ]) + ]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -209,7 +209,11 @@ mod tests { // Convert to Vortex let vortex_array = decode_run_array(&arrow_run_array, false)?; - assert_arrays_eq!(vortex_array, buffer![1.5f64, 2.5, 2.5, 3.5].into_array()); + assert_arrays_eq!( + vortex_array, + buffer![1.5f64, 2.5, 2.5, 3.5].into_array(), + &mut SESSION.create_execution_ctx() + ); Ok(()) } @@ -230,7 +234,8 @@ mod tests { let vortex_array = decode_run_array(&sliced_array, false)?; assert_arrays_eq!( vortex_array, - buffer![100, 200, 200, 200, 300, 300].into_array() + buffer![100, 200, 200, 200, 300, 300].into_array(), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -261,7 +266,8 @@ mod tests { Some(30), Some(30), Some(40), - ]) + ]), + &mut SESSION.create_execution_ctx() ); Ok(()) } @@ -345,7 +351,7 @@ mod tests { #[case] expected_ends: &[i32], #[case] expected_values: &[i32], ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = RunEnd::encode( PrimitiveArray::from_iter(input.iter().copied()).into_array(), &mut ctx, diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index ba46588d88b..544b677dfa3 100644 --- a/encodings/runend/src/compress.rs +++ b/encodings/runend/src/compress.rs @@ -320,7 +320,8 @@ pub fn runend_decode_varbinview( #[cfg(test)] mod tests { - use vortex_array::LEGACY_SESSION; + use std::sync::LazyLock; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -328,27 +329,34 @@ mod tests { use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::compress::runend_decode_primitive; use crate::compress::runend_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn encode() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = PrimitiveArray::from_iter([1i32, 1, 2, 2, 2, 3, 3, 3, 3, 3]); let (ends, values) = runend_encode(arr.as_view(), &mut ctx); let values = values.execute::(&mut ctx)?; let expected_ends = PrimitiveArray::from_iter(vec![2u8, 5, 10]); - assert_arrays_eq!(ends, expected_ends); + assert_arrays_eq!(ends, expected_ends, &mut ctx); let expected_values = PrimitiveArray::from_iter(vec![1i32, 2, 3]); - assert_arrays_eq!(values, expected_values); + assert_arrays_eq!(values, expected_values, &mut ctx); Ok(()) } #[test] fn encode_nullable() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1i32, 1, 2, 2, 2, 3, 3, 3, 3, 3], Validity::from(BitBuffer::from(vec![ @@ -359,16 +367,16 @@ mod tests { let values = values.execute::(&mut ctx)?; let expected_ends = PrimitiveArray::from_iter(vec![2u8, 4, 5, 8, 10]); - assert_arrays_eq!(ends, expected_ends); + assert_arrays_eq!(ends, expected_ends, &mut ctx); let expected_values = PrimitiveArray::from_option_iter(vec![Some(1i32), None, Some(2), Some(3), None]); - assert_arrays_eq!(values, expected_values); + assert_arrays_eq!(values, expected_values, &mut ctx); Ok(()) } #[test] fn encode_all_null() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = PrimitiveArray::new( buffer![0, 0, 0, 0, 0], Validity::from(BitBuffer::new_unset(5)), @@ -377,21 +385,21 @@ mod tests { let values = values.execute::(&mut ctx)?; let expected_ends = PrimitiveArray::from_iter(vec![5u64]); - assert_arrays_eq!(ends, expected_ends); + assert_arrays_eq!(ends, expected_ends, &mut ctx); let expected_values = PrimitiveArray::from_option_iter(vec![Option::::None]); - assert_arrays_eq!(values, expected_values); + assert_arrays_eq!(values, expected_values, &mut ctx); Ok(()) } #[test] fn decode() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = PrimitiveArray::from_iter([1i32, 2, 3]); let decoded = runend_decode_primitive(ends, values, 0, 10, &mut ctx)?; let expected = PrimitiveArray::from_iter(vec![1i32, 1, 2, 2, 2, 3, 3, 3, 3, 3]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } } diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index ff28baeb2d3..fe740739ca9 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -52,7 +52,11 @@ mod tests { use crate::RunEnd; use crate::RunEndArray; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_cast_runend_i32_to_i64() { @@ -130,7 +134,11 @@ mod tests { let sliced = runend.slice(3..8).unwrap(); // Verify the slice is correct before casting - assert_arrays_eq!(sliced, PrimitiveArray::from_iter([200, 200, 300, 300, 300])); + assert_arrays_eq!( + sliced, + PrimitiveArray::from_iter([200, 200, 300, 300, 300]), + &mut ctx + ); // Cast the sliced array let casted = sliced @@ -140,7 +148,8 @@ mod tests { // Verify the cast preserved the offset assert_arrays_eq!( casted, - PrimitiveArray::from_iter([200i64, 200, 300, 300, 300]) + PrimitiveArray::from_iter([200i64, 200, 300, 300, 300]), + &mut ctx ); } diff --git a/encodings/runend/src/compute/compare.rs b/encodings/runend/src/compute/compare.rs index 4da7a453000..ecdd03e3686 100644 --- a/encodings/runend/src/compute/compare.rs +++ b/encodings/runend/src/compute/compare.rs @@ -48,9 +48,10 @@ impl CompareKernel for RunEnd { #[cfg(test)] mod test { + use std::sync::LazyLock; + use vortex_array::ExecutionCtx; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; @@ -58,10 +59,17 @@ mod test { use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_session::VortexSession; use crate::RunEnd; use crate::RunEndArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn ree_array(ctx: &mut ExecutionCtx) -> RunEndArray { RunEnd::encode( PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(), @@ -72,7 +80,7 @@ mod test { #[test] fn compare_run_end() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = ree_array(&mut ctx); let res = arr .into_array() @@ -81,6 +89,6 @@ mod test { let expected = BoolArray::from_iter([ false, false, false, false, false, false, false, false, true, true, true, true, ]); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut ctx); } } diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index 83eb811e682..ca99fdf1912 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -146,7 +146,8 @@ mod tests { PrimitiveArray::from_iter([1u8, 2, 3]).into_array(), PrimitiveArray::from_iter([1i32, 4, 2]).into_array(), &mut ctx, - ) + ), + &mut ctx ); Ok(()) } @@ -186,7 +187,7 @@ mod tests { .chain(std::iter::repeat_n(30, 16)) .chain(std::iter::repeat_n(40, 16)) .collect(); - assert_arrays_eq!(executed, PrimitiveArray::from_iter(expected)); + assert_arrays_eq!(executed, PrimitiveArray::from_iter(expected), &mut ctx); Ok(()) } diff --git a/encodings/runend/src/compute/mod.rs b/encodings/runend/src/compute/mod.rs index 9bdb6d67c00..32926f41c45 100644 --- a/encodings/runend/src/compute/mod.rs +++ b/encodings/runend/src/compute/mod.rs @@ -13,44 +13,52 @@ pub(crate) mod take_from; #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::RunEnd; use crate::RunEndArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] // Simple run-end arrays #[case::runend_i32(RunEnd::encode( buffer![1i32, 1, 1, 2, 2, 3, 3, 3, 3].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::runend_single_run(RunEnd::encode( buffer![5i32, 5, 5, 5, 5].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::runend_alternating(RunEnd::encode( buffer![1i32, 2, 1, 2, 1, 2].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] // Different types #[case::runend_u64(RunEnd::encode( buffer![100u64, 100, 200, 200, 200].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] // Edge cases #[case::runend_single(RunEnd::encode( buffer![42i32].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case::runend_large(RunEnd::encode( PrimitiveArray::from_iter((0..1000).map(|i| i / 10)).into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] fn test_runend_consistency(#[case] array: RunEndArray) { diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index 46c8c2a35fc..c61854635e3 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -484,11 +484,12 @@ fn advance_run(ends: &[I], run_idx: &mut usize, logical_idx: u #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; @@ -497,6 +498,7 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use super::physical_indices_binary; use super::physical_indices_linear_sorted; @@ -505,10 +507,16 @@ mod tests { use crate::RunEnd; use crate::RunEndArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn ree_array() -> RunEndArray { RunEnd::encode( buffer![1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() } @@ -517,14 +525,14 @@ mod tests { fn ree_take() { let taken = ree_array().take(buffer![9, 8, 1, 3].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![5i32, 5, 1, 4]).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] fn ree_take_end() { let taken = ree_array().take(buffer![11].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![5i32]).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -533,7 +541,7 @@ mod tests { .take(buffer![0, 2, 3, 6, 8, 11].into_array()) .unwrap(); let expected = PrimitiveArray::from_iter(vec![1i32, 1, 4, 2, 5, 5]).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -542,7 +550,7 @@ mod tests { let _array = ree_array() .take(buffer![12].into_array()) .unwrap() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut SESSION.create_execution_ctx()) .unwrap(); } @@ -552,7 +560,7 @@ mod tests { let taken = sliced.take(buffer![1, 3, 4].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![4i32, 2, 5]).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -561,7 +569,7 @@ mod tests { let taken = sliced.take(buffer![4, 0, 2, 1].into_array()).unwrap(); let expected = PrimitiveArray::from_iter(vec![5i32, 4, 2, 4]).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -571,7 +579,11 @@ mod tests { .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -581,7 +593,11 @@ mod tests { .unwrap(); let expected = PrimitiveArray::from_option_iter([None::, None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -593,7 +609,11 @@ mod tests { let taken = ree_array().take(indices.into_array()).unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -605,7 +625,11 @@ mod tests { let taken = ree_array().take(indices.into_array()).unwrap(); let expected = PrimitiveArray::from_option_iter([Some(4i32), None, Some(1)]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -635,7 +659,11 @@ mod tests { Some(5), None, ]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[rstest] @@ -687,7 +715,7 @@ mod tests { #[case(ree_array())] #[case(RunEnd::encode( buffer![1u8, 1, 2, 2, 2, 3, 3, 3, 3, 4].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case(RunEnd::encode( PrimitiveArray::from_option_iter([ @@ -700,14 +728,14 @@ mod tests { Some(20), ]) .into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case(RunEnd::encode(buffer![42i32, 42, 42, 42, 42].into_array(), - &mut LEGACY_SESSION.create_execution_ctx()) + &mut SESSION.create_execution_ctx()) .unwrap())] #[case(RunEnd::encode( buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ).unwrap())] #[case({ let mut values = Vec::new(); @@ -718,7 +746,7 @@ mod tests { } RunEnd::encode( PrimitiveArray::from_iter(values).into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() })] @@ -731,7 +759,7 @@ mod tests { #[case({ let array = RunEnd::encode( buffer![1i32, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3].into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); array.slice(2..8).unwrap() diff --git a/encodings/runend/src/compute/take_from.rs b/encodings/runend/src/compute/take_from.rs index 8c3417459a0..0987cd3fc55 100644 --- a/encodings/runend/src/compute/take_from.rs +++ b/encodings/runend/src/compute/take_from.rs @@ -52,21 +52,31 @@ impl ExecuteParentKernel for RunEndTakeFrom { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::kernel::ExecuteParentKernel; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::RunEnd; use crate::RunEndArray; use crate::array::RunEndArrayExt; use crate::compute::take_from::RunEndTakeFrom; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + /// Build a DictArray whose codes are run-end encoded. /// /// Input: `[2, 2, 2, 3, 3, 2, 2]` @@ -82,7 +92,7 @@ mod tests { #[test] fn test_execute_parent_no_offset() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(vortex_array::array_session()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); let result = RunEndTakeFrom @@ -91,13 +101,13 @@ mod tests { let expected = PrimitiveArray::from_iter([2i32, 2, 2, 3, 3, 2, 2]); let canonical = result.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); Ok(()) } #[test] fn test_execute_parent_with_offset() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(vortex_array::array_session()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); // Slice codes to positions 2..5 → logical codes [0, 1, 1] → values [2, 3, 3] let sliced_codes = unsafe { @@ -115,13 +125,13 @@ mod tests { let expected = PrimitiveArray::from_iter([2i32, 3, 3]); let canonical = result.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); Ok(()) } #[test] fn test_execute_parent_offset_at_run_boundary() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(vortex_array::array_session()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); // Slice codes to positions 3..7 → logical codes [1, 1, 0, 0] → values [3, 3, 2, 2] let sliced_codes = unsafe { @@ -139,13 +149,13 @@ mod tests { let expected = PrimitiveArray::from_iter([3i32, 3, 2, 2]); let canonical = result.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); Ok(()) } #[test] fn test_execute_parent_single_element_offset() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(vortex_array::array_session()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); // Slice to single element at position 4 → code=1 → value=3 let sliced_codes = unsafe { @@ -163,13 +173,13 @@ mod tests { let expected = PrimitiveArray::from_iter([3i32]); let canonical = result.execute::(&mut ctx)?.into_array(); - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut ctx); Ok(()) } #[test] fn test_execute_parent_returns_none_for_non_codes_child() -> VortexResult<()> { - let mut ctx = ExecutionCtx::new(vortex_array::array_session()); + let mut ctx = SESSION.create_execution_ctx(); let (codes, dict) = make_dict_with_runend_codes(&mut ctx); let result = RunEndTakeFrom.execute_parent(codes.as_view(), dict.as_view(), 1, &mut ctx)?; diff --git a/encodings/runend/src/decompress_bool.rs b/encodings/runend/src/decompress_bool.rs index 68b6d3d5981..d2b778fdbbf 100644 --- a/encodings/runend/src/decompress_bool.rs +++ b/encodings/runend/src/decompress_bool.rs @@ -246,7 +246,8 @@ fn decode_nullable_sequential( #[cfg(test)] mod tests { - use vortex_array::LEGACY_SESSION; + use std::sync::LazyLock; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; @@ -255,12 +256,19 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use super::runend_decode_bools; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn decode_bools_alternating() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Alternating true/false: [T, T, F, F, F, T, T, T, T, T] let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = BoolArray::from(BitBuffer::from(vec![true, false, true])); @@ -269,13 +277,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ true, true, false, false, false, true, true, true, true, true, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_mostly_true() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Mostly true: [T, T, T, T, T, F, T, T, T, T] let ends = PrimitiveArray::from_iter([5u32, 6, 10]); let values = BoolArray::from(BitBuffer::from(vec![true, false, true])); @@ -284,13 +292,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ true, true, true, true, true, false, true, true, true, true, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_mostly_false() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Mostly false: [F, F, F, F, F, T, F, F, F, F] let ends = PrimitiveArray::from_iter([5u32, 6, 10]); let values = BoolArray::from(BitBuffer::from(vec![false, true, false])); @@ -299,13 +307,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ false, false, false, false, false, true, false, false, false, false, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_all_true_single_run() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let ends = PrimitiveArray::from_iter([10u32]); let values = BoolArray::from(BitBuffer::from(vec![true])); let decoded = runend_decode_bools(ends, values, 0, 10, &mut ctx)?; @@ -313,13 +321,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ true, true, true, true, true, true, true, true, true, true, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_all_false_single_run() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let ends = PrimitiveArray::from_iter([10u32]); let values = BoolArray::from(BitBuffer::from(vec![false])); let decoded = runend_decode_bools(ends, values, 0, 10, &mut ctx)?; @@ -327,13 +335,13 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![ false, false, false, false, false, false, false, false, false, false, ])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_with_offset() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test with offset: [T, T, F, F, F, T, T, T, T, T] -> slice [2..8] = [F, F, F, T, T, T] let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = BoolArray::from(BitBuffer::from(vec![true, false, true])); @@ -341,7 +349,7 @@ mod tests { let expected = BoolArray::from(BitBuffer::from(vec![false, false, false, true, true, true])); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } @@ -349,7 +357,7 @@ mod tests { fn decode_bools_nullable() -> VortexResult<()> { use vortex_array::validity::Validity; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // 3 runs: T (valid), F (null), T (valid) -> [T, T, null, null, null, T, T, T, T, T] let ends = PrimitiveArray::from_iter([2u32, 5, 10]); let values = BoolArray::new( @@ -367,13 +375,13 @@ mod tests { true, true, false, false, false, true, true, true, true, true, ])), ); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn decode_bools_nullable_few_runs() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test few runs (uses fast path): 5 runs of length 2000 each let ends = PrimitiveArray::from_iter([2000u32, 4000, 6000, 8000, 10000]); let values = BoolArray::new( diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 41a9e69419c..56de024f988 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -49,9 +49,9 @@ pub(crate) fn find_slice_end_index(array: &ArrayRef, index: usize) -> VortexResu #[cfg(test)] mod tests { + use std::sync::LazyLock; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::is_constant::is_constant; use vortex_array::arrays::PrimitiveArray; @@ -60,12 +60,19 @@ mod tests { use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::RunEnd; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn slice_array() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = RunEnd::try_new( buffer![2u32, 5, 10].into_array(), buffer![1i32, 2, 3].into_array(), @@ -81,12 +88,12 @@ mod tests { assert_eq!(arr.len(), 5); let expected = PrimitiveArray::from_iter(vec![2i32, 2, 3, 3, 3]).into_array(); - assert_arrays_eq!(arr, expected); + assert_arrays_eq!(arr, expected, &mut ctx); } #[test] fn double_slice() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = RunEnd::try_new( buffer![2u32, 5, 10].into_array(), buffer![1i32, 2, 3].into_array(), @@ -100,12 +107,12 @@ mod tests { let doubly_sliced = arr.slice(0..3).unwrap(); let expected = PrimitiveArray::from_iter(vec![2i32, 2, 3]).into_array(); - assert_arrays_eq!(doubly_sliced, expected); + assert_arrays_eq!(doubly_sliced, expected, &mut ctx); } #[test] fn slice_end_inclusive() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = RunEnd::try_new( buffer![2u32, 5, 10].into_array(), buffer![1i32, 2, 3].into_array(), @@ -121,12 +128,12 @@ mod tests { assert_eq!(arr.len(), 6); let expected = PrimitiveArray::from_iter(vec![2i32, 3, 3, 3, 3, 3]).into_array(); - assert_arrays_eq!(arr, expected); + assert_arrays_eq!(arr, expected, &mut ctx); } #[test] fn slice_at_end() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let re_array = RunEnd::try_new( buffer![7_u64, 10].into_array(), buffer![2_u64, 3].into_array(), @@ -142,7 +149,7 @@ mod tests { #[test] fn slice_single_end() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let re_array = RunEnd::try_new( buffer![7_u64, 10].into_array(), buffer![2_u64, 3].into_array(), @@ -159,7 +166,7 @@ mod tests { #[test] fn ree_scalar_at_end() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let scalar = RunEnd::encode( buffer![1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5].into_array(), &mut ctx, @@ -172,7 +179,7 @@ mod tests { #[test] fn slice_along_run_boundaries() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a runend array with runs: [1, 1, 1] [4, 4, 4] [2, 2] [5, 5, 5, 5] // Run ends at indices: 3, 6, 8, 12 let arr = RunEnd::try_new( @@ -186,42 +193,42 @@ mod tests { let slice1 = arr.slice(0..3).unwrap(); assert_eq!(slice1.len(), 3); let expected = PrimitiveArray::from_iter(vec![1i32, 1, 1]).into_array(); - assert_arrays_eq!(slice1, expected); + assert_arrays_eq!(slice1, expected, &mut ctx); // Slice from start of second run to end of second run (indices 3..6) let slice2 = arr.slice(3..6).unwrap(); assert_eq!(slice2.len(), 3); let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4]).into_array(); - assert_arrays_eq!(slice2, expected); + assert_arrays_eq!(slice2, expected, &mut ctx); // Slice from start of third run to end of third run (indices 6..8) let slice3 = arr.slice(6..8).unwrap(); assert_eq!(slice3.len(), 2); let expected = PrimitiveArray::from_iter(vec![2i32, 2]).into_array(); - assert_arrays_eq!(slice3, expected); + assert_arrays_eq!(slice3, expected, &mut ctx); // Slice from start of last run to end of last run (indices 8..12) let slice4 = arr.slice(8..12).unwrap(); assert_eq!(slice4.len(), 4); let expected = PrimitiveArray::from_iter(vec![5i32, 5, 5, 5]).into_array(); - assert_arrays_eq!(slice4, expected); + assert_arrays_eq!(slice4, expected, &mut ctx); // Slice spanning exactly two runs (indices 3..8) let slice5 = arr.slice(3..8).unwrap(); assert_eq!(slice5.len(), 5); let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4, 2, 2]).into_array(); - assert_arrays_eq!(slice5, expected); + assert_arrays_eq!(slice5, expected, &mut ctx); // Slice from middle of first run to end of second run (indices 1..6) let slice6 = arr.slice(1..6).unwrap(); assert_eq!(slice6.len(), 5); let expected = PrimitiveArray::from_iter(vec![1i32, 1, 4, 4, 4]).into_array(); - assert_arrays_eq!(slice6, expected); + assert_arrays_eq!(slice6, expected, &mut ctx); // Slice from start of second run to middle of third run (indices 3..7) let slice7 = arr.slice(3..7).unwrap(); assert_eq!(slice7.len(), 4); let expected = PrimitiveArray::from_iter(vec![4i32, 4, 4, 2]).into_array(); - assert_arrays_eq!(slice7, expected); + assert_arrays_eq!(slice7, expected, &mut ctx); } } diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index 15365b04464..56e62764ce4 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -442,7 +442,8 @@ impl Sequence { #[cfg(test)] mod tests { - use vortex_array::LEGACY_SESSION; + use std::sync::LazyLock; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -453,16 +454,23 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::Sequence; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_sequence_canonical() { let arr = Sequence::try_new_typed(2i64, 3, Nullability::NonNullable, 4).unwrap(); let canon = PrimitiveArray::from_iter((0..4).map(|i| 2i64 + i * 3)); - assert_arrays_eq!(arr, canon); + assert_arrays_eq!(arr, canon, &mut SESSION.create_execution_ctx()); } #[test] @@ -474,14 +482,14 @@ mod tests { let canon = PrimitiveArray::from_iter((2..3).map(|i| 2i64 + i * 3)); - assert_arrays_eq!(arr, canon); + assert_arrays_eq!(arr, canon, &mut SESSION.create_execution_ctx()); } #[test] fn test_sequence_scalar_at() { let scalar = Sequence::try_new_typed(2i64, 3, Nullability::NonNullable, 4) .unwrap() - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut SESSION.create_execution_ctx()) .unwrap(); assert_eq!( diff --git a/encodings/sequence/src/compress.rs b/encodings/sequence/src/compress.rs index 9f5e62a5217..f34cf4ffe34 100644 --- a/encodings/sequence/src/compress.rs +++ b/encodings/sequence/src/compress.rs @@ -148,18 +148,26 @@ fn encode_primitive_array + CheckedAdd + CheckedSu #[cfg(test)] mod tests { + use std::sync::LazyLock; + #[expect(unused_imports)] use itertools::Itertools; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_session::VortexSession; use crate::sequence_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_encode_array_success() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); assert!(encoded.is_some()); @@ -167,12 +175,12 @@ mod tests { .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decoded, primitive_array); + assert_arrays_eq!(decoded, primitive_array, &mut ctx); } #[test] fn test_encode_array_1_success() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter([0]); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); assert!(encoded.is_some()); @@ -180,12 +188,12 @@ mod tests { .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decoded, primitive_array); + assert_arrays_eq!(decoded, primitive_array, &mut ctx); } #[test] fn test_encode_array_fail() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0]); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); @@ -194,7 +202,7 @@ mod tests { #[test] fn test_encode_array_fail_oob() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter(vec![100i8; 1000]); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); @@ -203,7 +211,7 @@ mod tests { #[test] fn test_encode_all_u8_values() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let primitive_array = PrimitiveArray::from_iter(0u8..=255); let encoded = sequence_encode(primitive_array.as_view(), &mut ctx).unwrap(); assert!(encoded.is_some()); @@ -211,6 +219,6 @@ mod tests { .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decoded, primitive_array); + assert_arrays_eq!(decoded, primitive_array, &mut ctx); } } diff --git a/encodings/sequence/src/compute/cast.rs b/encodings/sequence/src/compute/cast.rs index 3750fe20fd7..e0a63d0c922 100644 --- a/encodings/sequence/src/compute/cast.rs +++ b/encodings/sequence/src/compute/cast.rs @@ -104,7 +104,11 @@ mod tests { use crate::Sequence; use crate::SequenceArray; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_cast_sequence_nullability() { @@ -137,7 +141,11 @@ mod tests { // Verify the values let decoded = casted.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([100i64, 110, 120, 130])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([100i64, 110, 120, 130]), + &mut ctx + ); } #[test] @@ -159,7 +167,8 @@ mod tests { let decoded = casted.execute::(&mut ctx).unwrap(); assert_arrays_eq!( decoded, - PrimitiveArray::from_option_iter([Some(5i32), Some(8), Some(11)]) + PrimitiveArray::from_option_iter([Some(5i32), Some(8), Some(11)]), + &mut ctx ); } @@ -183,7 +192,8 @@ mod tests { let decoded = casted.execute::(&mut ctx).unwrap(); assert_arrays_eq!( decoded, - PrimitiveArray::from_iter([0.0f32, 1.0, 2.0, 3.0, 4.0]) + PrimitiveArray::from_iter([0.0f32, 1.0, 2.0, 3.0, 4.0]), + &mut ctx ); } diff --git a/encodings/sequence/src/compute/compare.rs b/encodings/sequence/src/compute/compare.rs index c0c9d1367d6..5895330ea81 100644 --- a/encodings/sequence/src/compute/compare.rs +++ b/encodings/sequence/src/compute/compare.rs @@ -135,7 +135,10 @@ fn find_intersection( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::assert_arrays_eq; @@ -143,9 +146,16 @@ mod tests { use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::Nullability::Nullable; use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_session::VortexSession; use crate::Sequence; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_compare_match() { let lhs = Sequence::try_new_typed(2i64, 1, NonNullable, 4).unwrap(); @@ -155,7 +165,7 @@ mod tests { .binary(rhs.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([false, false, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -167,7 +177,7 @@ mod tests { .binary(rhs.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([Some(false), Some(false), Some(true), Some(false)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -179,6 +189,6 @@ mod tests { .binary(rhs.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([false, false, false, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index fa2cacb5e19..93bd823a5f2 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -52,8 +52,10 @@ impl ListContainsElementReduce for Sequence { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::Nullability; @@ -62,9 +64,16 @@ mod tests { use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar::Scalar; + use vortex_session::VortexSession; use crate::Sequence; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_list_contains_seq() { let list_scalar = Scalar::list( @@ -84,7 +93,7 @@ mod tests { let expr = list_contains(lit(list_scalar.clone()), root()); let result = array.into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(false), Some(true)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } { @@ -98,7 +107,7 @@ mod tests { let expr = list_contains(lit(list_scalar), root()); let result = array.into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([Some(true), Some(true), Some(false)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } } } diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 17b4241f293..1829d4fffd1 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -597,11 +597,11 @@ fn execute_varbin_inner( #[cfg(test)] mod test { use std::sync::Arc; + use std::sync::LazyLock; use rstest::rstest; use vortex_array::Canonical; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; @@ -631,15 +631,22 @@ mod test { use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case(Some(true))] #[case(Some(false))] #[case(None)] fn test_sparse_bool(#[case] fill_value: Option) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u64, 1, 7].into_array(); let values = BoolArray::from_iter([Some(true), None, Some(false)]).into_array(); let sparse_bools = Sparse::try_new(indices, values, 10, Scalar::from(fill_value)).unwrap(); @@ -662,7 +669,7 @@ mod test { fill_value, ]); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[rstest] @@ -670,7 +677,7 @@ mod test { #[case(Some(-1i32))] #[case(None)] fn test_sparse_primitive(#[case] fill_value: Option) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u64, 1, 7].into_array(); let values = PrimitiveArray::from_option_iter([Some(0i32), None, Some(1)]).into_array(); let sparse_ints = Sparse::try_new(indices, values, 10, Scalar::from(fill_value)).unwrap(); @@ -694,12 +701,12 @@ mod test { fill_value, ]); - assert_arrays_eq!(&flat_ints, &expected); + assert_arrays_eq!(&flat_ints, &expected, &mut ctx); } #[test] fn test_sparse_struct_valid_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let field_names = FieldNames::from_iter(["a", "b"]); let field_types = vec![ DType::Primitive(PType::I32, Nullable), @@ -769,12 +776,12 @@ mod test { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_struct_invalid_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let field_names = FieldNames::from_iter(["a", "b"]); let field_types = vec![ DType::Primitive(PType::I32, Nullable), @@ -841,12 +848,12 @@ mod test { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_decimal() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let indices = buffer![0u32, 1u32, 7u32, 8u32].into_array(); let decimal_dtype = DecimalDType::new(3, 2); let patch_values = DecimalArray::new( @@ -859,7 +866,7 @@ mod test { let fill_scalar = Scalar::decimal(DecimalValue::I32(123), decimal_dtype, Nullable); let sparse_struct = Sparse::try_new(indices, patch_values, len, fill_scalar).unwrap(); - let expected = LEGACY_SESSION + let expected = SESSION .arrow() .execute_arrow( DecimalArray::new( @@ -874,7 +881,7 @@ mod test { ) .unwrap(); - let actual = LEGACY_SESSION + let actual = SESSION .arrow() .execute_arrow( sparse_struct @@ -894,7 +901,7 @@ mod test { #[test] fn test_sparse_utf8_varbinview_non_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = >::from_iter([ Some("hello"), Some("goodbye"), @@ -936,12 +943,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_utf8_varbinview_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = >::from_iter([ Some("hello"), Some("goodbye"), @@ -983,12 +990,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_utf8_varbinview_non_nullable() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = VarBinViewArray::from_iter_str(["hello", "goodbye", "hello", "bonjour", "你好"]) .into_array(); @@ -1012,12 +1019,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_utf8_varbin_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = >::from_iter([ Some("hello"), Some("goodbye"), @@ -1059,12 +1066,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_binary_varbinview_non_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let binaries = VarBinViewArray::from_iter_nullable_bin([ Some(b"hello" as &[u8]), Some(b"goodbye"), @@ -1106,12 +1113,12 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_sparse_list_null_fill() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Use ListViewArray consistently let elements = buffer![1i32, 2, 1, 2].into_array(); // Create ListView with offsets and sizes @@ -1169,7 +1176,7 @@ mod test { #[test] fn test_sparse_list_null_fill_sliced_sparse_values() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create ListViewArray with 8 elements forming 8 single-element lists let elements = buffer![1i32, 2, 1, 2, 1, 2, 1, 2].into_array(); let offsets = buffer![0u32, 1, 2, 3, 4, 5, 6, 7].into_array(); @@ -1225,7 +1232,7 @@ mod test { #[test] fn test_sparse_list_non_null_fill() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create ListViewArray with 4 single-element lists let elements = buffer![1i32, 2, 1, 2].into_array(); let offsets = buffer![0u32, 1, 2, 3].into_array(); @@ -1297,7 +1304,7 @@ mod test { #[test] fn test_sparse_binary_varbin_null_fill() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let strings = >::from_iter([ Some(b"hello" as &[u8]), Some(b"goodbye"), @@ -1339,7 +1346,7 @@ mod test { ]) .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] @@ -1357,7 +1364,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 5, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Expected: [1,2,3], null, [4,5,6], [7,8,9], null. @@ -1371,7 +1378,7 @@ mod test { )? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -1392,7 +1399,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 6, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Expected: [1,2], [99,88], [3,4], [99,88], [5,6], [99,88]. @@ -1400,7 +1407,7 @@ mod test { let expected = FixedSizeListArray::try_new(expected_elements, 2, Validity::NonNullable, 6)? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -1428,7 +1435,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 6, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Expected validity: [true, true, true, false, true, true]. @@ -1444,7 +1451,7 @@ mod test { )? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -1473,7 +1480,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 100, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Build expected: 97 copies of [99,99] with patches at positions 5, 50, 95. @@ -1505,7 +1512,7 @@ mod test { FixedSizeListArray::try_new(expected_elements, 2, Validity::NonNullable, 100)? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -1527,7 +1534,7 @@ mod test { let sparse = Sparse::try_new(indices, fsl, 1, fill_value)?.into_array(); let actual = sparse - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut SESSION.create_execution_ctx())? .into_array(); // Expected: just [42, 43]. @@ -1535,13 +1542,13 @@ mod test { let expected = FixedSizeListArray::try_new(expected_elements, 2, Validity::NonNullable, 1)? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn test_sparse_list_grows_offset_type() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let elements = buffer![1i32, 2, 1, 2].into_array(); let offsets = buffer![0u8, 1, 2, 3, 4].into_array(); let lists = ListArray::try_new(elements, offsets, Validity::AllValid)?.into_array(); @@ -1565,19 +1572,16 @@ mod test { actual_listview.offsets().dtype(), &DType::Primitive(PType::U16, NonNullable) ); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &expected, &mut ctx); // Note that the preferred arrow list representation is `List` (not `ListView`). - let arrow_dtype = LEGACY_SESSION - .arrow() - .to_arrow_field("", expected.dtype())?; - let actual = LEGACY_SESSION + let arrow_dtype = SESSION.arrow().to_arrow_field("", expected.dtype())?; + let actual = SESSION .arrow() .execute_arrow(actual, Some(&arrow_dtype), &mut ctx)?; - let expected = - LEGACY_SESSION - .arrow() - .execute_arrow(expected, Some(&arrow_dtype), &mut ctx)?; + let expected = SESSION + .arrow() + .execute_arrow(expected, Some(&arrow_dtype), &mut ctx)?; assert_eq!(actual.data_type(), expected.data_type()); Ok(()) @@ -1585,7 +1589,7 @@ mod test { #[test] fn test_sparse_listview_null_fill_with_gaps() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // This test specifically catches the bug where the old implementation // incorrectly tracked `last_valid_offset` as the START of the last list // instead of properly handling ListView's offset/size pairs. @@ -1667,7 +1671,7 @@ mod test { #[test] fn test_sparse_listview_sliced_values_null_fill() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // This test uses sliced ListView values to ensure proper handling // of non-zero starting offsets in the source data. diff --git a/encodings/sparse/src/compute/between.rs b/encodings/sparse/src/compute/between.rs index ba50a120174..46fcd7f1e59 100644 --- a/encodings/sparse/src/compute/between.rs +++ b/encodings/sparse/src/compute/between.rs @@ -131,6 +131,6 @@ mod tests { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(kernel, baseline); + assert_arrays_eq!(kernel, baseline, &mut ctx); } } diff --git a/encodings/sparse/src/compute/cast.rs b/encodings/sparse/src/compute/cast.rs index 1068987b9e4..39c21805426 100644 --- a/encodings/sparse/src/compute/cast.rs +++ b/encodings/sparse/src/compute/cast.rs @@ -54,7 +54,11 @@ mod tests { use crate::Sparse; use crate::SparseArray; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[test] fn test_cast_sparse_i32_to_i64() { @@ -78,7 +82,7 @@ mod tests { let expected = PrimitiveArray::from_iter([0i64, 0, 100, 0, 0, 200, 0, 0, 300, 0]); let casted_primitive = casted.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(casted_primitive, expected); + assert_arrays_eq!(casted_primitive, expected, &mut ctx); } #[test] @@ -156,7 +160,7 @@ mod tests { let expected = PrimitiveArray::from_iter([10u64, 20, 30, 40, 50]); let casted_primitive = casted.execute::(&mut ctx)?; - assert_arrays_eq!(casted_primitive, expected); + assert_arrays_eq!(casted_primitive, expected, &mut ctx); Ok(()) } diff --git a/encodings/sparse/src/compute/compare.rs b/encodings/sparse/src/compute/compare.rs index f9637da516a..75bb78622df 100644 --- a/encodings/sparse/src/compute/compare.rs +++ b/encodings/sparse/src/compute/compare.rs @@ -111,6 +111,6 @@ mod tests { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(kernel_bool, canonical_bool); + assert_arrays_eq!(kernel_bool, canonical_bool, &mut ctx); } } diff --git a/encodings/sparse/src/compute/fill_null.rs b/encodings/sparse/src/compute/fill_null.rs index 8f7d2b75a27..f36032265bc 100644 --- a/encodings/sparse/src/compute/fill_null.rs +++ b/encodings/sparse/src/compute/fill_null.rs @@ -105,6 +105,6 @@ mod tests { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(kernel, baseline); + assert_arrays_eq!(kernel, baseline, &mut ctx); } } diff --git a/encodings/sparse/src/compute/filter.rs b/encodings/sparse/src/compute/filter.rs index f7871f124a6..91e67297fd2 100644 --- a/encodings/sparse/src/compute/filter.rs +++ b/encodings/sparse/src/compute/filter.rs @@ -34,10 +34,13 @@ impl FilterKernel for Sparse { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::fixture; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; @@ -49,9 +52,16 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[fixture] fn array() -> ArrayRef { Sparse::try_new( @@ -82,7 +92,11 @@ mod tests { ) .unwrap(); - assert_arrays_eq!(filtered_array, expected); + assert_arrays_eq!( + filtered_array, + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -111,7 +125,11 @@ mod tests { ) .unwrap(); - assert_arrays_eq!(filtered_array, expected); + assert_arrays_eq!( + filtered_array, + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] diff --git a/encodings/sparse/src/compute/mod.rs b/encodings/sparse/src/compute/mod.rs index fac33edf47e..a64ccc869e2 100644 --- a/encodings/sparse/src/compute/mod.rs +++ b/encodings/sparse/src/compute/mod.rs @@ -15,10 +15,13 @@ mod take; #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::fixture; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; @@ -31,9 +34,16 @@ mod test { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; + use vortex_session::VortexSession; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[fixture] fn array() -> ArrayRef { Sparse::try_new( @@ -64,7 +74,11 @@ mod test { ) .unwrap(); - assert_arrays_eq!(filtered_array, expected); + assert_arrays_eq!( + filtered_array, + expected, + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -93,7 +107,11 @@ mod test { ) .unwrap(); - assert_arrays_eq!(filtered_array, expected); + assert_arrays_eq!( + filtered_array, + expected, + &mut SESSION.create_execution_ctx() + ); } #[rstest] diff --git a/encodings/sparse/src/compute/take.rs b/encodings/sparse/src/compute/take.rs index ddf9000fac1..e260eb26e9e 100644 --- a/encodings/sparse/src/compute/take.rs +++ b/encodings/sparse/src/compute/take.rs @@ -55,9 +55,12 @@ impl TakeExecute for Sparse { #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -65,10 +68,17 @@ mod test { use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::Sparse; use crate::SparseArray; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + fn test_array_fill_value() -> Scalar { // making this const is annoying Scalar::null_native::() @@ -91,7 +101,11 @@ mod test { let sparse = sparse.slice(30..40).unwrap(); let taken = sparse.take(buffer![6, 7, 8].into_array()).unwrap(); let expected = PrimitiveArray::from_option_iter([Option::::None, Some(0.47), None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -105,7 +119,11 @@ mod test { Some(1.23), Some(3.5), ]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -113,7 +131,7 @@ mod test { let sparse = sparse_array(); let taken = sparse.take(buffer![69].into_array()).unwrap(); let expected = ConstantArray::new(test_array_fill_value(), 1).into_array(); - assert_arrays_eq!(taken, expected); + assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx()); } #[test] @@ -123,7 +141,11 @@ mod test { let taken = sparse.take(buffer![69, 37].into_array()).unwrap(); // Index 69 is not in sparse array (fill value is null), index 37 has value 0.47 let expected = PrimitiveArray::from_option_iter([Option::::None, Some(0.47f64)]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -144,7 +166,11 @@ mod test { .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1), Some(10), Option::::None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -165,7 +191,11 @@ mod test { .unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1), Some(10), Option::::None]); - assert_arrays_eq!(taken, expected.into_array()); + assert_arrays_eq!( + taken, + expected.into_array(), + &mut SESSION.create_execution_ctx() + ); } #[rstest] diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index d47cf36563e..8ae38b56052 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -688,9 +688,10 @@ impl ValidityVTable for Sparse { #[cfg(test)] mod test { + use std::sync::LazyLock; + use itertools::Itertools; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; @@ -707,6 +708,12 @@ mod test { use super::*; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + initialize(&session); + session + }); + fn nullable_fill() -> Scalar { Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)) } @@ -731,19 +738,19 @@ mod test { assert_eq!( array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap(), nullable_fill() ); assert_eq!( array - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(2, &mut SESSION.create_execution_ctx()) .unwrap(), Scalar::from(Some(100_i32)) ); assert_eq!( array - .execute_scalar(5, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(5, &mut SESSION.create_execution_ctx()) .unwrap(), Scalar::from(Some(200_i32)) ); @@ -754,7 +761,7 @@ mod test { fn test_scalar_at_oob() { let array = sparse_array(nullable_fill()); array - .execute_scalar(10, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(10, &mut SESSION.create_execution_ctx()) .unwrap(); } @@ -769,19 +776,19 @@ mod test { .unwrap(); assert_eq!( - arr.execute_scalar(10, &mut LEGACY_SESSION.create_execution_ctx()) + arr.execute_scalar(10, &mut SESSION.create_execution_ctx()) .unwrap() .as_primitive() .typed_value::(), Some(1234) ); assert!( - arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + arr.execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() .is_null() ); assert!( - arr.execute_scalar(99, &mut LEGACY_SESSION.create_execution_ctx()) + arr.execute_scalar(99, &mut SESSION.create_execution_ctx()) .unwrap() .is_null() ); @@ -793,7 +800,7 @@ mod test { assert_eq!( usize::try_from( &sliced - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap() ) .unwrap(), @@ -808,7 +815,7 @@ mod test { sliced .validity() .unwrap() - .execute_mask(sliced.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(sliced.len(), &mut SESSION.create_execution_ctx()) .unwrap(), Mask::from_iter(vec![true, false, false, true, false]) ); @@ -834,7 +841,7 @@ mod test { sliced .validity() .unwrap() - .execute_mask(sliced.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(sliced.len(), &mut SESSION.create_execution_ctx()) .unwrap(), Mask::from_iter(vec![false, true, true, false, true]) ); @@ -846,7 +853,7 @@ mod test { assert_eq!( usize::try_from( &sliced_once - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap() ) .unwrap(), @@ -857,7 +864,7 @@ mod test { assert_eq!( usize::try_from( &sliced_twice - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(3, &mut SESSION.create_execution_ctx()) .unwrap() ) .unwrap(), @@ -872,7 +879,7 @@ mod test { array .validity() .unwrap() - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(array.len(), &mut SESSION.create_execution_ctx()) .unwrap() .to_bit_buffer() .iter() @@ -890,7 +897,7 @@ mod test { array .validity() .unwrap() - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(array.len(), &mut SESSION.create_execution_ctx()) .unwrap() .all_true() ); @@ -915,7 +922,7 @@ mod test { #[test] fn encode_with_nulls() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let original = PrimitiveArray::new( buffer![0i32, 1, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4], Validity::from_iter(vec![ @@ -935,7 +942,7 @@ mod test { ]) ); let sparse_primitive = sparse.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(sparse_primitive, original); + assert_arrays_eq!(sparse_primitive, original, &mut ctx); } #[test] @@ -947,7 +954,7 @@ mod test { let actual = array .validity() .unwrap() - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(array.len(), &mut SESSION.create_execution_ctx()) .unwrap(); let expected = Mask::from_iter([ true, false, true, false, false, false, false, false, true, false, diff --git a/encodings/sparse/src/ops.rs b/encodings/sparse/src/ops.rs index 172114b7dfb..568d8d377d1 100644 --- a/encodings/sparse/src/ops.rs +++ b/encodings/sparse/src/ops.rs @@ -25,18 +25,26 @@ impl OperationsVTable for Sparse { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_buffer::buffer; + use vortex_session::VortexSession; use crate::Sparse; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn slice_partially_invalid() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![0u64].into_array(); let indices = buffer![0u8].into_array(); @@ -46,6 +54,6 @@ mod tests { expected[0] = 0; let values = sliced.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(values, PrimitiveArray::from_iter(expected)); + assert_arrays_eq!(values, PrimitiveArray::from_iter(expected), &mut ctx); } } diff --git a/encodings/zigzag/src/compress.rs b/encodings/zigzag/src/compress.rs index 260e2e375b5..e70650a56e1 100644 --- a/encodings/zigzag/src/compress.rs +++ b/encodings/zigzag/src/compress.rs @@ -77,25 +77,33 @@ where #[cfg(test)] mod test { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; + use vortex_session::VortexSession; use super::*; use crate::ZigZag; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[rstest] #[case::i8(PrimitiveArray::from_iter(-100_i8..100))] #[case::i16(PrimitiveArray::from_iter(-100_i16..100))] #[case::i32(PrimitiveArray::from_iter(-100_i32..100))] #[case::i64(PrimitiveArray::from_iter(-100_i64..100))] fn test_compress(#[case] input: PrimitiveArray) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let compressed = zigzag_encode(input.as_view()).unwrap().into_array(); assert!(compressed.is::()); let decompressed = compressed.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(decompressed, input); + assert_arrays_eq!(decompressed, input, &mut ctx); } } diff --git a/encodings/zigzag/src/compute/cast.rs b/encodings/zigzag/src/compute/cast.rs index a71eee7a838..c3fae84e346 100644 --- a/encodings/zigzag/src/compute/cast.rs +++ b/encodings/zigzag/src/compute/cast.rs @@ -26,8 +26,11 @@ impl CastReduce for ZigZag { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; @@ -35,10 +38,17 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_session::VortexSession; use crate::ZigZagArray; use crate::zigzag_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] fn test_cast_zigzag_i32_to_i64() { let values = PrimitiveArray::from_iter([-100i32, -1, 0, 1, 100]); @@ -61,7 +71,11 @@ mod tests { "Cast should preserve ZigZag encoding" ); - assert_arrays_eq!(casted, PrimitiveArray::from_iter([-100i64, -1, 0, 1, 100])); + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([-100i64, -1, 0, 1, 100]), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -82,7 +96,8 @@ mod tests { assert_arrays_eq!( casted, - PrimitiveArray::from_iter([100i16, -50, 0, 25, -100]) + PrimitiveArray::from_iter([100i16, -50, 0, 25, -100]), + &mut SESSION.create_execution_ctx() ); // Test i16 to i64 (widening) @@ -101,7 +116,8 @@ mod tests { assert_arrays_eq!( casted64, - PrimitiveArray::from_iter([1000i64, -500, 0, 250, -1000]) + PrimitiveArray::from_iter([1000i64, -500, 0, 250, -1000]), + &mut SESSION.create_execution_ctx() ); } diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index a90db7eacc7..41e83d187e2 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -70,10 +70,11 @@ impl ZigZagEncoded for u64 { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -85,17 +86,24 @@ mod tests { use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::ZigZagArray; use crate::zigzag_encode; + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + #[test] pub fn nullable_scalar_at() -> VortexResult<()> { let zigzag = zigzag_encode( PrimitiveArray::new(buffer![-189, -160, 1], Validity::AllValid).as_view(), )?; assert_eq!( - zigzag.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())?, + zigzag.execute_scalar(1, &mut SESSION.create_execution_ctx())?, Scalar::primitive(-160, Nullability::Nullable) ); Ok(()) @@ -112,7 +120,7 @@ mod tests { let expected = zigzag_encode(PrimitiveArray::new(buffer![-189, 1], Validity::AllValid).as_view())? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -127,7 +135,7 @@ mod tests { let expected = zigzag_encode(PrimitiveArray::new(buffer![-189, 1], Validity::AllValid).as_view())? .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -187,7 +195,7 @@ mod tests { fn test_take_zigzag_conformance(#[case] array: ArrayRef) -> VortexResult<()> { use vortex_array::compute::conformance::take::test_take_conformance; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array_primitive = array.execute::(&mut ctx)?; let zigzag = zigzag_encode(array_primitive.as_view())?; test_take_conformance(&zigzag.into_array()); diff --git a/encodings/zstd/src/compute/cast.rs b/encodings/zstd/src/compute/cast.rs index 52f6f8020c0..c5f15eba284 100644 --- a/encodings/zstd/src/compute/cast.rs +++ b/encodings/zstd/src/compute/cast.rs @@ -107,7 +107,11 @@ mod tests { ); let decoded = casted.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i64, 2, 3, 4, 5])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([1i64, 2, 3, 4, 5]), + &mut ctx + ); } #[test] @@ -144,7 +148,11 @@ mod tests { ); // Verify the values are correct let decoded = casted.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([20u32, 30, 40, 50])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([20u32, 30, 40, 50]), + &mut ctx + ); } #[test] @@ -169,7 +177,7 @@ mod tests { ); let decoded = casted.execute::(&mut ctx).unwrap(); let expected = PrimitiveArray::from_iter([20u32, 30, 40, 50]); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut ctx); } #[rstest] diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index 094c5ef60cc..e52cea12afa 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -32,17 +32,25 @@ fn test_zstd_compress_decompress() { // check full decompression works let decompressed = Zstd::decompress(&compressed, &mut ctx).unwrap(); - assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data), &mut ctx); // check slicing works let slice = compressed.slice(100..105).unwrap(); for i in 0_i32..5 { - assert_nth_scalar!(slice, i as usize, 100 + i); + assert_nth_scalar!(slice, i as usize, 100 + i, &mut ctx); } - assert_arrays_eq!(slice, PrimitiveArray::from_iter([100, 101, 102, 103, 104])); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter([100, 101, 102, 103, 104]), + &mut ctx + ); let slice = compressed.slice(200..200).unwrap(); - assert_arrays_eq!(slice, PrimitiveArray::from_iter(Vec::::new())); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter(Vec::::new()), + &mut ctx + ); } #[test] @@ -56,7 +64,7 @@ fn test_zstd_empty() { let compressed = Zstd::from_primitive(&array, 3, 100, &mut ctx).unwrap(); - assert_arrays_eq!(compressed, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(compressed, PrimitiveArray::from_iter(data), &mut ctx); } #[test] @@ -73,10 +81,10 @@ fn test_zstd_with_validity_and_multi_frame() { let compressed = Zstd::from_primitive(&array, 0, 30, &mut ctx).unwrap(); assert!(compressed.dictionary.is_none()); - assert_nth_scalar!(compressed, 0, None::); - assert_nth_scalar!(compressed, 3, 3); - assert_nth_scalar!(compressed, 10, None::); - assert_nth_scalar!(compressed, 177, 177); + assert_nth_scalar!(compressed, 0, None::, &mut ctx); + assert_nth_scalar!(compressed, 3, 3, &mut ctx); + assert_nth_scalar!(compressed, 10, None::, &mut ctx); + assert_nth_scalar!(compressed, 177, 177, &mut ctx); let decompressed = Zstd::decompress(&compressed, &mut ctx) .unwrap() @@ -124,19 +132,23 @@ fn test_zstd_with_dict() { let compressed = Zstd::from_primitive(&array, 0, 16, &mut ctx).unwrap(); assert!(compressed.dictionary.is_some()); - assert_nth_scalar!(compressed, 0, 0); - assert_nth_scalar!(compressed, 199, 199); + assert_nth_scalar!(compressed, 0, 0, &mut ctx); + assert_nth_scalar!(compressed, 199, 199, &mut ctx); let decompressed = Zstd::decompress(&compressed, &mut ctx) .unwrap() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data)); + assert_arrays_eq!(decompressed, PrimitiveArray::from_iter(data), &mut ctx); // check slicing works let slice = compressed.slice(176..179).unwrap(); let primitive = slice.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(primitive, PrimitiveArray::from_iter([176, 177, 178])); + assert_arrays_eq!( + primitive, + PrimitiveArray::from_iter([176, 177, 178]), + &mut ctx + ); } #[test] @@ -181,16 +193,16 @@ fn test_zstd_var_bin_view() { let compressed = Zstd::from_var_bin_view(&array, 0, 3, &mut ctx).unwrap(); assert!(compressed.dictionary.is_none()); - assert_nth_scalar!(compressed, 0, "foo"); - assert_nth_scalar!(compressed, 1, "bar"); - assert_nth_scalar!(compressed, 2, None::); - assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet"); - assert_nth_scalar!(compressed, 4, "baz"); + assert_nth_scalar!(compressed, 0, "foo", &mut ctx); + assert_nth_scalar!(compressed, 1, "bar", &mut ctx); + assert_nth_scalar!(compressed, 2, None::, &mut ctx); + assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet", &mut ctx); + assert_nth_scalar!(compressed, 4, "baz", &mut ctx); let sliced = compressed.slice(1..4).unwrap(); - assert_nth_scalar!(sliced, 0, "bar"); - assert_nth_scalar!(sliced, 1, None::); - assert_nth_scalar!(sliced, 2, "Lorem ipsum dolor sit amet"); + assert_nth_scalar!(sliced, 0, "bar", &mut ctx); + assert_nth_scalar!(sliced, 1, None::, &mut ctx); + assert_nth_scalar!(sliced, 2, "Lorem ipsum dolor sit amet", &mut ctx); } #[test] @@ -207,21 +219,21 @@ fn test_zstd_decompress_var_bin_view() { let compressed = Zstd::from_var_bin_view(&array, 0, 3, &mut ctx).unwrap(); assert!(compressed.dictionary.is_none()); - assert_nth_scalar!(compressed, 0, "foo"); - assert_nth_scalar!(compressed, 1, "bar"); - assert_nth_scalar!(compressed, 2, None::); - assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet"); - assert_nth_scalar!(compressed, 4, "baz"); + assert_nth_scalar!(compressed, 0, "foo", &mut ctx); + assert_nth_scalar!(compressed, 1, "bar", &mut ctx); + assert_nth_scalar!(compressed, 2, None::, &mut ctx); + assert_nth_scalar!(compressed, 3, "Lorem ipsum dolor sit amet", &mut ctx); + assert_nth_scalar!(compressed, 4, "baz", &mut ctx); let decompressed = Zstd::decompress(&compressed, &mut ctx) .unwrap() .execute::(&mut ctx) .unwrap(); - assert_nth_scalar!(decompressed, 0, "foo"); - assert_nth_scalar!(decompressed, 1, "bar"); - assert_nth_scalar!(decompressed, 2, None::); - assert_nth_scalar!(decompressed, 3, "Lorem ipsum dolor sit amet"); - assert_nth_scalar!(decompressed, 4, "baz"); + assert_nth_scalar!(decompressed, 0, "foo", &mut ctx); + assert_nth_scalar!(decompressed, 1, "bar", &mut ctx); + assert_nth_scalar!(decompressed, 2, None::, &mut ctx); + assert_nth_scalar!(decompressed, 3, "Lorem ipsum dolor sit amet", &mut ctx); + assert_nth_scalar!(decompressed, 4, "baz", &mut ctx); } #[test] diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index dd689fd76bd..df6ee882d6f 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -552,7 +552,7 @@ mod tests { let mut ctx = LEGACY_SESSION.create_execution_ctx(); let decompressed = compressed.into_array().execute::(&mut ctx)?; - assert_arrays_eq!(input, decompressed); + assert_arrays_eq!(input, decompressed, &mut ctx); Ok(()) } diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index d19a9f2b360..7b3ebdbdf22 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -335,7 +335,7 @@ mod tests { let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -349,7 +349,7 @@ mod tests { let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = BoolArray::from(BitBuffer::from(vec![true, true, false, true])); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -364,7 +364,7 @@ mod tests { let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = VarBinViewArray::from_iter_str(["hello", "default", "world"]); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -376,7 +376,7 @@ mod tests { let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = PrimitiveArray::from_iter([100i32, 100, 100]); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -388,7 +388,7 @@ mod tests { let result = fill_null_canonical_array(canonical(array), &fill_value, &mut ctx).unwrap(); let expected = PrimitiveArray::from_iter([1i32, 2, 3]); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -428,7 +428,7 @@ mod tests { [100i32, 999i32, 300i32, 999i32, 500i32], DecimalDType::new(10, 2), ); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -448,7 +448,7 @@ mod tests { let expected = DecimalArray::from_iter([1000i64, 9999i64, 3000i64], DecimalDType::new(15, 3)); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -470,7 +470,7 @@ mod tests { [10000i128, 99999i128, 30000i128, 99999i128], DecimalDType::new(20, 4), ); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -493,7 +493,7 @@ mod tests { .into_array() .cast(result.dtype().clone()) .unwrap(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -518,6 +518,6 @@ mod tests { .into_array() .cast(result.dtype().clone()) .unwrap(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } } diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index 49fc2ad40f9..537421d4bd0 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -206,7 +206,7 @@ mod tests { let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); let expected = BoolArray::from_iter([None, Some(false), Some(true), None, Some(true)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] @@ -218,7 +218,7 @@ mod tests { let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] @@ -230,7 +230,7 @@ mod tests { let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); let expected = PrimitiveArray::from_option_iter([None, None, Some(3i32), None, None]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] @@ -247,7 +247,7 @@ mod tests { let expected = DecimalArray::from_option_iter([Some(1i128), Some(2), None, Some(4), Some(5)], dtype); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] @@ -260,7 +260,7 @@ mod tests { let expected = VarBinViewArray::from_iter_nullable_str([None, Some("two"), None, Some("four"), None]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] @@ -335,7 +335,7 @@ mod tests { let result = mask_canonical_array(canonical(array), &mask, &mut ctx).unwrap(); let expected = PrimitiveArray::from_option_iter([None, None, None, None, None::]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] @@ -348,7 +348,7 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4), Some(5)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs index f11a9a2bc52..1ad0ca6d646 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::dtype::NativePType; use vortex_error::VortexResult; use crate::arrays::primitive::PrimitiveArrayExt; @@ -12,6 +13,10 @@ where R: PrimitiveArrayExt, { match_each_native_ptype!(lhs.ptype(), |P| { - Ok(lhs.as_slice::

() == rhs.as_slice::

()) + Ok(lhs + .as_slice::

() + .iter() + .zip(rhs.as_slice::

()) + .all(|(l, r)| l.is_eq(*r))) }) } diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index 3cc75a79b1f..e5e780e42f0 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -151,6 +151,8 @@ mod tests { #[test] fn listview_counts_all_valid() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I32, NonNullable); @@ -162,13 +164,15 @@ mod tests { let direct = PrimitiveArray::new(buffer![2u64, 1, 3, 0], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct, &mut assertion_ctx); + assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); Ok(()) } #[test] fn listview_counts_with_nulls() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) .into_array(); @@ -181,13 +185,15 @@ mod tests { // Group 0: {1, null, 3} -> 2. Group 1: {null, null} -> 0. Group 2: {9} -> 1. let direct = PrimitiveArray::new(buffer![2u64, 0, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct, &mut assertion_ctx); + assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); Ok(()) } #[test] fn listview_counts_varbinview_with_nulls() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = VarBinViewArray::from_iter_nullable_str([ Some("a"), None, @@ -204,8 +210,8 @@ mod tests { let expected = grouped_count_reference(&elements, &ranges)?; let direct = PrimitiveArray::new(buffer![1u64, 1, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct, &mut assertion_ctx); + assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); Ok(()) } @@ -234,6 +240,8 @@ mod tests { #[test] fn fixed_size_counts_with_nulls() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(); let elem_dtype = DType::Primitive(PType::I32, Nullable); @@ -242,7 +250,7 @@ mod tests { let actual = grouped_count_actual(&groups, &elem_dtype)?; let direct = PrimitiveArray::new(buffer![1u64, 2], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct); + assert_arrays_eq!(&actual, &direct, &mut assertion_ctx); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 1bd17df14cd..51394873ad5 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -236,6 +236,8 @@ mod tests { #[test] fn listview_matches_reference_unsigned() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1u32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::U32, NonNullable); @@ -248,13 +250,15 @@ mod tests { // Unsigned input sums to U64. let direct = PrimitiveArray::from_option_iter([Some(3u64), Some(3u64), Some(15u64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); + assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); Ok(()) } #[test] fn listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Offsets are not in group order and a group is null: the group validity must be indexed by // group index, not by element offset. let elements = @@ -269,13 +273,15 @@ mod tests { let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; let direct = PrimitiveArray::from_option_iter([Some(110i64), None, Some(70i64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); + assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); Ok(()) } #[test] fn listview_interior_and_full_nulls() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Group 1 has an interior null, group 2 is entirely null, group 3 is empty. let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) @@ -290,13 +296,15 @@ mod tests { let direct = PrimitiveArray::from_option_iter([Some(4i64), Some(0i64), Some(0i64), Some(9i64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); + assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); Ok(()) } #[test] fn listview_overflow_group_is_null() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I64, NonNullable); @@ -309,8 +317,8 @@ mod tests { // First group overflows -> null sum; second group sums normally. let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); + assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); Ok(()) } @@ -380,6 +388,8 @@ mod tests { #[test] fn fixed_size_overflow_and_nan() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // FixedSize path: first group overflows -> null sum, second sums normally. let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); @@ -391,8 +401,8 @@ mod tests { let expected = grouped_sum_reference(&elements, &[(0, 2), (2, 2)], &[true, true], &elem_dtype)?; let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); - assert_arrays_eq!(&actual, &direct.into_array()); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); + assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 26f618a5ffa..5262cc72a34 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -586,6 +586,8 @@ mod tests { #[test] fn grouped_sum_fixed_size_list() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; @@ -594,12 +596,14 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(6i64), Some(15i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut assertion_ctx); Ok(()) } #[test] fn grouped_sum_with_null_elements() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) .into_array(); @@ -609,12 +613,14 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(4i64), Some(11i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut assertion_ctx); Ok(()) } #[test] fn grouped_sum_with_null_group() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) .into_array(); @@ -626,12 +632,14 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(6i64), None, Some(24i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut assertion_ctx); Ok(()) } #[test] fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; @@ -640,12 +648,14 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut assertion_ctx); Ok(()) } #[test] fn grouped_sum_bool() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); let groups = FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; @@ -654,7 +664,7 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut assertion_ctx); Ok(()) } @@ -672,7 +682,7 @@ mod tests { let result1 = acc.finish()?; let expected1 = PrimitiveArray::from_option_iter([Some(3i64), Some(7i64)]).into_array(); - assert_arrays_eq!(&result1, &expected1); + assert_arrays_eq!(&result1, &expected1, &mut ctx); let elements2 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); let groups2 = FixedSizeListArray::try_new(elements2, 2, Validity::NonNullable, 1)?; @@ -680,12 +690,14 @@ mod tests { let result2 = acc.finish()?; let expected2 = PrimitiveArray::from_option_iter([Some(30i64)]).into_array(); - assert_arrays_eq!(&result2, &expected2); + assert_arrays_eq!(&result2, &expected2, &mut ctx); Ok(()) } #[test] fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); @@ -699,7 +711,7 @@ mod tests { // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200. let expected = PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut assertion_ctx); Ok(()) } diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 180f771bc9c..06a7ca538fb 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -489,6 +489,7 @@ mod tests { use vortex_buffer::buffer; use super::Array; + use crate::VortexSessionExecute; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -496,17 +497,21 @@ mod tests { #[test] fn typed_array_into_parts_roundtrips() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let expected = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let parts = array.try_into_parts().unwrap(); let rebuilt = Array::::try_from_parts(parts).unwrap(); - assert_arrays_eq!(rebuilt, expected); + assert_arrays_eq!(rebuilt, expected, &mut assertion_ctx); } #[test] fn typed_array_try_into_parts_requires_unique_owner() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let alias = array.clone(); @@ -515,6 +520,6 @@ mod tests { Err(array) => array, }; - assert_arrays_eq!(array, alias); + assert_arrays_eq!(array, alias, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/assertions.rs b/vortex-array/src/arrays/assertions.rs index 389baa23857..7d2948e5d79 100644 --- a/vortex-array/src/arrays/assertions.rs +++ b/vortex-array/src/arrays/assertions.rs @@ -9,9 +9,8 @@ use vortex_error::VortexExpect; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::RecursiveCanonical; -use crate::VortexSessionExecute; +use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct; fn format_indices>(indices: I) -> impl Display { indices.into_iter().format(",") @@ -28,14 +27,14 @@ fn execute_to_canonical(array: ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { /// Finds indices where two arrays differ based on `scalar_at` comparison. #[expect(clippy::unwrap_used)] -fn find_mismatched_indices(left: &ArrayRef, right: &ArrayRef) -> Vec { +fn find_mismatched_indices( + left: &ArrayRef, + right: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> Vec { assert_eq!(left.len(), right.len()); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); (0..left.len()) - .filter(|i| { - left.execute_scalar(*i, &mut ctx).unwrap() - != right.execute_scalar(*i, &mut ctx).unwrap() - }) + .filter(|i| left.execute_scalar(*i, ctx).unwrap() != right.execute_scalar(*i, ctx).unwrap()) .collect() } @@ -46,22 +45,16 @@ fn find_mismatched_indices(left: &ArrayRef, right: &ArrayRef) -> Vec { /// # Example /// ```ignore /// let arr = PrimitiveArray::from_iter([1, 2, 3]); -/// assert_nth_scalar!(arr, 0, 1); -/// assert_nth_scalar!(arr, 1, 2); +/// assert_nth_scalar!(arr, 0, 1, &mut ctx); +/// assert_nth_scalar!(arr, 1, 2, &mut ctx); /// ``` #[macro_export] macro_rules! assert_nth_scalar { - ($arr:expr, $n:expr, $expected:expr) => {{ + ($arr:expr, $n:expr, $expected:expr, $ctx:expr) => {{ use $crate::IntoArray as _; - use $crate::LEGACY_SESSION; - use $crate::VortexSessionExecute as _; let arr_ref: $crate::ArrayRef = $crate::IntoArray::into_array($arr.clone()); - assert_eq!( - arr_ref - .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - $expected.try_into().unwrap() - ); + let expected = $expected.try_into().unwrap(); + assert_eq!(arr_ref.execute_scalar($n, $ctx).unwrap(), expected); }}; } @@ -71,33 +64,25 @@ macro_rules! assert_nth_scalar { /// /// ```ignore /// let arr = PrimitiveArray::from_option_iter([Some(1), None, Some(3)]); -/// assert_nth_scalar_null!(arr, 1); +/// assert_nth_scalar_null!(arr, 1, &mut ctx); /// ``` #[macro_export] macro_rules! assert_nth_scalar_is_null { - ($arr:expr, $n:expr) => {{ - use $crate::LEGACY_SESSION; - use $crate::VortexSessionExecute as _; + ($arr:expr, $n:expr, $ctx:expr) => {{ let arr_ref: $crate::ArrayRef = $crate::IntoArray::into_array($arr.clone()); + let scalar = arr_ref.execute_scalar($n, $ctx).unwrap(); assert!( - arr_ref - .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_null(), + scalar.is_null(), "expected scalar at index {} to be null, but was {:?}", $n, - arr_ref - .execute_scalar($n, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + scalar ); }}; } #[macro_export] macro_rules! assert_arrays_eq { - ($left:expr, $right:expr) => {{ - - + ($left:expr, $right:expr, $ctx:expr) => {{ let left: $crate::ArrayRef = $crate::IntoArray::into_array($left.clone()); let right: $crate::ArrayRef = $crate::IntoArray::into_array($right.clone()); if left.dtype() != right.dtype() { @@ -122,7 +107,7 @@ macro_rules! assert_arrays_eq { let left = left.clone(); let right = right.clone(); - $crate::arrays::assert_arrays_eq_impl(&left, &right); + $crate::arrays::assert_arrays_eq_impl(&left, &right, $ctx); }}; } @@ -130,13 +115,17 @@ macro_rules! assert_arrays_eq { /// `ArrayRef`. #[track_caller] #[expect(clippy::panic)] -pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef) { - let executed = execute_to_canonical(left.clone(), &mut LEGACY_SESSION.create_execution_ctx()); +pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef, ctx: &mut ExecutionCtx) { + let executed = execute_to_canonical(left.clone(), ctx); - let left_right = find_mismatched_indices(left, right); - let executed_right = find_mismatched_indices(&executed, right); + let left_right_the_same = + all_non_distinct(left, right, ctx).vortex_expect("failed to compare left and right"); + let executed_right_the_same = all_non_distinct(&executed, right, ctx) + .vortex_expect("failed to compare executed left and right"); - if !left_right.is_empty() || !executed_right.is_empty() { + if !left_right_the_same || !executed_right_the_same { + let left_right = find_mismatched_indices(left, right, ctx); + let executed_right = find_mismatched_indices(&executed, right, ctx); let mut msg = String::new(); if !left_right.is_empty() { msg.push_str(&format!( diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 82df2f9c167..100487a607f 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -449,9 +449,11 @@ mod tests { #[test] fn patch_sliced_bools() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.slice(4..12).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8])); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut assertion_ctx); let arr = { let mut builder = BitBufferMut::new_unset(12); @@ -460,7 +462,11 @@ mod tests { }; let sliced = arr.slice(4..12).unwrap(); let expected_slice: Vec = (4..12).map(|i| (1..12).contains(&i)).collect(); - assert_arrays_eq!(sliced, BoolArray::from_iter(expected_slice.clone())); + assert_arrays_eq!( + sliced, + BoolArray::from_iter(expected_slice.clone()), + &mut assertion_ctx + ); // patch the underlying array at index 4 to false let patches = Patches::new( @@ -476,21 +482,33 @@ mod tests { .unwrap(); // After patching index 4 to false: indices 1-3 and 5-11 are true, index 0 and 4 are false let expected_patched: Vec = (0..12).map(|i| (1..12).contains(&i) && i != 4).collect(); - assert_arrays_eq!(arr, BoolArray::from_iter(expected_patched)); + assert_arrays_eq!( + arr, + BoolArray::from_iter(expected_patched), + &mut assertion_ctx + ); // the slice should be unchanged (still has original values before patch) - assert_arrays_eq!(sliced, BoolArray::from_iter(expected_slice)); + assert_arrays_eq!( + sliced, + BoolArray::from_iter(expected_slice), + &mut assertion_ctx + ); } #[test] fn slice_array_in_middle() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let sliced = arr.slice(4..12).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8])); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut assertion_ctx); } #[test] fn patch_bools_owned() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let buf_ptr = arr.to_bit_buffer().inner().as_ptr(); @@ -510,13 +528,15 @@ mod tests { // After patching index 0 to false: [false, true, true, ..., true] (16 values) let expected: BoolArray = once(false).chain(repeat_n(true, 15)).collect(); - assert_arrays_eq!(arr, expected); + assert_arrays_eq!(arr, expected, &mut assertion_ctx); } #[test] fn patch_sliced_bools_offset() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.slice(4..15).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11])); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11]), &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index 315c1db2474..6c1b692006e 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -303,44 +303,69 @@ mod tests { use super::*; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; #[test] fn filter_bool_test() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_iter([true, false, true]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!(filtered, BoolArray::from_iter([true, false])); + assert_arrays_eq!( + filtered, + BoolArray::from_iter([true, false]), + &mut assertion_ctx + ); } #[test] fn filter_bool_sparse_index_mask() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_indices(3, [0, 2]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!(filtered, BoolArray::from_iter([true, false])); + assert_arrays_eq!( + filtered, + BoolArray::from_iter([true, false]), + &mut assertion_ctx + ); } #[test] fn filter_bool_sparse_slice_mask() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_slices(3, vec![(0, 1), (2, 3)]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!(filtered, BoolArray::from_iter([true, false])); + assert_arrays_eq!( + filtered, + BoolArray::from_iter([true, false]), + &mut assertion_ctx + ); } #[test] fn filter_bool_sparse_buffer_mask() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_buffer(BitBuffer::from_iter([true, false, true])); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!(filtered, BoolArray::from_iter([true, false])); + assert_arrays_eq!( + filtered, + BoolArray::from_iter([true, false]), + &mut assertion_ctx + ); } #[test] diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index d70579136a4..1981e3aad85 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -90,6 +90,7 @@ mod test { use crate::IntoArray as _; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; @@ -99,6 +100,8 @@ mod test { #[test] fn take_nullable() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let reference = BoolArray::from_iter(vec![ Some(false), Some(true), @@ -119,11 +122,17 @@ mod test { let all_invalid_indices = PrimitiveArray::from_option_iter([None::, None, None]); let b = reference.take(all_invalid_indices.into_array()).unwrap(); - assert_arrays_eq!(b, BoolArray::from_iter([None, None, None])); + assert_arrays_eq!( + b, + BoolArray::from_iter([None, None, None]), + &mut assertion_ctx + ); } #[test] fn test_bool_array_take_with_null_out_of_bounds_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -132,11 +141,17 @@ mod test { let actual = values.take(indices.into_array()).unwrap(); // position 3 is null, the third index is null - assert_arrays_eq!(actual, BoolArray::from_iter([Some(false), None, None])); + assert_arrays_eq!( + actual, + BoolArray::from_iter([Some(false), None, None]), + &mut assertion_ctx + ); } #[test] fn test_non_null_bool_array_take_with_null_out_of_bounds_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -146,30 +161,43 @@ mod test { // the third index is null assert_arrays_eq!( actual, - BoolArray::from_iter([Some(false), Some(true), None]) + BoolArray::from_iter([Some(false), Some(true), None]), + &mut assertion_ctx ); } #[test] fn test_bool_array_take_all_null_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([false, false, false]).into_array()), ); let actual = values.take(indices.into_array()).unwrap(); - assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None])); + assert_arrays_eq!( + actual, + BoolArray::from_iter([None, None, None]), + &mut assertion_ctx + ); } #[test] fn test_non_null_bool_array_take_all_null_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([false, false, false]).into_array()), ); let actual = values.take(indices.into_array()).unwrap(); - assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None])); + assert_arrays_eq!( + actual, + BoolArray::from_iter([None, None, None]), + &mut assertion_ctx + ); } #[rstest] diff --git a/vortex-array/src/arrays/bool/compute/zip.rs b/vortex-array/src/arrays/bool/compute/zip.rs index 2628b457a85..bad41eb19c4 100644 --- a/vortex-array/src/arrays/bool/compute/zip.rs +++ b/vortex-array/src/arrays/bool/compute/zip.rs @@ -147,7 +147,7 @@ mod tests { } })) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -182,7 +182,7 @@ mod tests { } })) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/bool/patch.rs b/vortex-array/src/arrays/bool/patch.rs index 5efdac30231..840e8388b07 100644 --- a/vortex-array/src/arrays/bool/patch.rs +++ b/vortex-array/src/arrays/bool/patch.rs @@ -47,22 +47,27 @@ mod tests { use vortex_buffer::BitBuffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; #[test] fn patch_sliced_bools() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.into_array().slice(4..12).unwrap(); let expected = BoolArray::from_iter([true; 8]); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut assertion_ctx); } #[test] fn patch_sliced_bools_offset() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.into_array().slice(4..15).unwrap(); let expected = BoolArray::from_iter([true; 11]); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index eaa39462894..8d390c4b7ff 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -212,6 +212,7 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -219,6 +220,8 @@ mod tests { #[test] fn test_nullable_bool_serde_roundtrip() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = BoolArray::from_iter([Some(true), None, Some(false), None]); let dtype = array.dtype().clone(); let len = array.len(); @@ -244,6 +247,6 @@ mod tests { ) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index b4ac4682d23..5a7f0334ae6 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -30,6 +30,7 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; @@ -46,13 +47,16 @@ mod tests { #[test] fn test_slice() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); #[expect(deprecated)] let sliced_arr = arr.into_array().slice(1..4).unwrap().to_bool(); assert_arrays_eq!( sliced_arr, - BoolArray::from_iter([Some(true), None, Some(false)]) + BoolArray::from_iter([Some(true), None, Some(false)]), + &mut assertion_ctx ); } } diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 9faa1486e64..4b828c7262b 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -291,6 +291,8 @@ mod test { #[test] fn test_rechunk_one_chunk() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -299,11 +301,13 @@ mod test { let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap(); - assert_arrays_eq!(chunked, rechunked); + assert_arrays_eq!(chunked, rechunked, &mut assertion_ctx); } #[test] fn test_rechunk_two_chunks() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array(), buffer![5u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -313,11 +317,13 @@ mod test { let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap(); assert_eq!(rechunked.nchunks(), 1); - assert_arrays_eq!(chunked, rechunked); + assert_arrays_eq!(chunked, rechunked, &mut assertion_ctx); } #[test] fn test_rechunk_tiny_target_chunks() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2, 3].into_array(), @@ -331,11 +337,13 @@ mod test { assert_eq!(rechunked.nchunks(), 2); assert!(rechunked.iter_chunks().all(|c| c.len() < 5)); - assert_arrays_eq!(chunked, rechunked); + assert_arrays_eq!(chunked, rechunked, &mut assertion_ctx); } #[test] fn test_rechunk_with_too_big_chunk() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2].into_array(), @@ -352,7 +360,7 @@ mod test { // greedy so should be: [0, 1, 2] [42, 42, 42, 42, 42, 42] [4, 5, 6, 7] [8, 9] assert_eq!(rechunked.nchunks(), 4); - assert_arrays_eq!(chunked, rechunked); + assert_arrays_eq!(chunked, rechunked, &mut assertion_ctx); } #[test] diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index 9eed54d4300..b74a04f8ba2 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -35,6 +35,7 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -46,6 +47,8 @@ mod test { #[test] fn test_cast_chunked() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr0 = buffer![0u32, 1].into_array(); let arr1 = buffer![2u32, 3].into_array(); @@ -67,7 +70,11 @@ mod test { let result = root .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) .unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([0u64, 1, 2, 3])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([0u64, 1, 2, 3]), + &mut assertion_ctx + ); } #[rstest] diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 7024c2b3875..3ab6ed147db 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -128,6 +128,7 @@ mod test { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; @@ -141,6 +142,8 @@ mod test { #[test] fn test_take() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -149,11 +152,17 @@ mod test { let indices = buffer![0u64, 0, 6, 4].into_array(); let result = arr.take(indices).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 1, 1, 2])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 1, 1, 2]), + &mut assertion_ctx + ); } #[test] fn test_take_nullable_values() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let a = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -164,12 +173,15 @@ mod test { let result = arr.take(indices.into_array()).unwrap(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([1i32, 1, 1, 2].map(Some)) + PrimitiveArray::from_option_iter([1i32, 1, 1, 2].map(Some)), + &mut assertion_ctx ); } #[test] fn test_take_nullable_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -183,12 +195,15 @@ mod test { let result = arr.take(indices.into_array()).unwrap(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(1i32), None, None, Some(2)]) + PrimitiveArray::from_option_iter([Some(1i32), None, None, Some(2)]), + &mut assertion_ctx ); } #[test] fn test_take_nullable_struct() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let struct_array = StructArray::try_new(FieldNames::default(), vec![], 100, Validity::NonNullable) .unwrap(); @@ -209,11 +224,13 @@ mod test { Validity::Array(BoolArray::from_iter(vec![true, false, true]).into_array()), ) .unwrap(); - assert_arrays_eq!(result, expect); + assert_arrays_eq!(result, expect, &mut assertion_ctx); } #[test] fn test_empty_take() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -227,12 +244,15 @@ mod test { assert_eq!(result.dtype(), arr.dtype()); assert_arrays_eq!( result, - PrimitiveArray::empty::(Nullability::NonNullable) + PrimitiveArray::empty::(Nullability::NonNullable), + &mut assertion_ctx ); } #[test] fn test_take_shuffled_indices() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let c0 = buffer![0i32, 1, 2].into_array(); let c1 = buffer![3i32, 4, 5].into_array(); let c2 = buffer![6i32, 7, 8].into_array(); @@ -249,7 +269,8 @@ mod test { assert_arrays_eq!( result, - PrimitiveArray::from_iter([8i32, 0, 5, 3, 2, 7, 1, 6, 4]) + PrimitiveArray::from_iter([8i32, 0, 5, 3, 2, 7, 1, 6, 4]), + &mut assertion_ctx ); Ok(()) } @@ -300,6 +321,8 @@ mod test { #[test] fn test_take_null_indices() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let c0 = buffer![10i32, 20, 30].into_array(); let c1 = buffer![40i32, 50, 60].into_array(); let arr = ChunkedArray::try_new( @@ -323,7 +346,8 @@ mod test { Some(40), None, Some(30) - ]) + ]), + &mut assertion_ctx ); Ok(()) } diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index 970656491ee..bd1a21cf74e 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -63,7 +63,8 @@ fn builder_kernel_path_canonicalizes_primitive_chunks() { assert_arrays_eq!( output, - PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut ctx ); } @@ -96,7 +97,11 @@ fn builder_kernel_nested_chunked_of_chunked() { let mut builder = execute_into_builder(outer, builder, &mut ctx).unwrap(); let output = builder.finish(); - assert_arrays_eq!(output, PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6])); + assert_arrays_eq!( + output, + PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]), + &mut ctx + ); } #[test] @@ -128,8 +133,8 @@ fn builder_kernel_path_repeated_shared_chunked_dict_execution() { drop(keep_alive); - assert_arrays_eq!(first, expected); - assert_arrays_eq!(second, expected); + assert_arrays_eq!(first, expected, &mut ctx); + assert_arrays_eq!(second, expected, &mut ctx); } #[test] @@ -154,8 +159,8 @@ fn execute_path_repeated_shared_chunked_dict_execution() { drop(keep_alive); - assert_arrays_eq!(first, expected); - assert_arrays_eq!(second, expected); + assert_arrays_eq!(first, expected, &mut ctx); + assert_arrays_eq!(second, expected, &mut ctx); } #[test] @@ -192,12 +197,14 @@ fn execute_path_nested_chunked_dict_of_dict_into_canonical() { drop(keep_alive); - assert_arrays_eq!(first, expected); - assert_arrays_eq!(second, expected); + assert_arrays_eq!(first, expected, &mut ctx); + assert_arrays_eq!(second, expected, &mut ctx); } #[test] fn with_slot_rewrites_chunk_and_offsets() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = chunked_array().into_array(); let replacement = buffer![10u64, 11, 12].into_array(); @@ -208,11 +215,13 @@ fn with_slot_rewrites_chunk_and_offsets() { assert_eq!(array.chunk_offsets(), [0, 3, 6, 9]); assert_arrays_eq!( array.chunk(0).clone(), - PrimitiveArray::from_iter([10u64, 11, 12]) + PrimitiveArray::from_iter([10u64, 11, 12]), + &mut assertion_ctx ); assert_arrays_eq!( array.array().clone(), - PrimitiveArray::from_iter([10u64, 11, 12, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([10u64, 11, 12, 4, 5, 6, 7, 8, 9]), + &mut assertion_ctx ); } @@ -228,49 +237,67 @@ fn with_slot_rejects_len_mismatch() { #[test] fn slice_middle() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(2..5).unwrap(), - PrimitiveArray::from_iter([3u64, 4, 5]) + PrimitiveArray::from_iter([3u64, 4, 5]), + &mut assertion_ctx ); } #[test] fn slice_begin() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(1..3).unwrap(), - PrimitiveArray::from_iter([2u64, 3]) + PrimitiveArray::from_iter([2u64, 3]), + &mut assertion_ctx ); } #[test] fn slice_aligned() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(3..6).unwrap(), - PrimitiveArray::from_iter([4u64, 5, 6]) + PrimitiveArray::from_iter([4u64, 5, 6]), + &mut assertion_ctx ); } #[test] fn slice_many_aligned() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(0..6).unwrap(), - PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]) + PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]), + &mut assertion_ctx ); } #[test] fn slice_end() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(7..8).unwrap(), - PrimitiveArray::from_iter([8u64]) + PrimitiveArray::from_iter([8u64]), + &mut assertion_ctx ); } #[test] fn slice_exactly_end() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(6..9).unwrap(), - PrimitiveArray::from_iter([7u64, 8, 9]) + PrimitiveArray::from_iter([7u64, 8, 9]), + &mut assertion_ctx ); } @@ -284,6 +311,8 @@ fn slice_empty() { #[test] fn scalar_at_empty_children_both_sides() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -295,11 +324,17 @@ fn scalar_at_empty_children_both_sides() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1u64, 2]), + &mut assertion_ctx + ); } #[test] fn scalar_at_empty_children_trailing() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -310,11 +345,17 @@ fn scalar_at_empty_children_trailing() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1u64, 2, 3, 4]), + &mut assertion_ctx + ); } #[test] fn scalar_at_empty_children_leading() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -325,7 +366,11 @@ fn scalar_at_empty_children_leading() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1u64, 2, 3, 4]), + &mut assertion_ctx + ); } #[test] diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 8c004f59ccb..3cc9505d16b 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -297,7 +297,6 @@ mod tests { use crate::ArrayRef; use crate::Canonical; - use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::accessor::ArrayAccessor; @@ -437,7 +436,11 @@ mod tests { assert_variant_values(&variant, &[10, 20, 30])?; let shredded = shredded.clone().execute::(&mut ctx)?; - assert_arrays_eq!(shredded, PrimitiveArray::from_iter([10i32, 20, 30])); + assert_arrays_eq!( + shredded, + PrimitiveArray::from_iter([10i32, 20, 30]), + &mut ctx + ); Ok(()) } @@ -641,7 +644,7 @@ mod tests { let session = crate::array_session().with_allocator(Arc::new(CountingAllocator { allocations: Arc::clone(&allocations), })); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let l1 = ListArray::try_new( buffer![1, 2, 3, 4].into_array(), diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index ce44b8231f8..1df291aa5dc 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -30,6 +30,7 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -57,9 +58,12 @@ mod tests { #[case::end(7..8, &[8u64])] #[case::exactly_end(6..9, &[7u64, 8, 9])] fn slice(#[case] range: Range, #[case] expected: &[u64]) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(range).unwrap(), - PrimitiveArray::from_iter(expected.iter().copied()) + PrimitiveArray::from_iter(expected.iter().copied()), + &mut assertion_ctx ); } @@ -73,6 +77,8 @@ mod tests { #[test] fn scalar_at_empty_children_both_sides() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -84,11 +90,17 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1u64, 2]), + &mut assertion_ctx + ); } #[test] fn scalar_at_empty_children_trailing() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -99,11 +111,17 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1u64, 2, 3, 4]), + &mut assertion_ctx + ); } #[test] fn scalar_at_empty_children_leading() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -114,6 +132,10 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1u64, 2, 3, 4]), + &mut assertion_ctx + ); } } diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index eb3d7918207..eb2debc0594 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -87,6 +87,8 @@ mod tests { #[test] fn take_nullable_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take( @@ -108,7 +110,8 @@ mod tests { PrimitiveArray::new( buffer![42i32, 42, 42], Validity::from_iter([false, true, false]) - ) + ), + &mut assertion_ctx ); assert_eq!( taken @@ -123,6 +126,8 @@ mod tests { #[test] fn take_all_valid_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take(PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array()) @@ -134,7 +139,8 @@ mod tests { assert_arrays_eq!( #[expect(deprecated)] taken.to_primitive(), - PrimitiveArray::new(buffer![42i32, 42, 42], Validity::AllValid) + PrimitiveArray::new(buffer![42i32, 42, 42], Validity::AllValid), + &mut assertion_ctx ); assert_eq!( taken diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index e12db00c118..8e9b78c9584 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -384,10 +384,12 @@ mod tests { #[test] fn test_canonicalize_const_str() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let const_array = ConstantArray::new("four".to_string(), 4); let expected = VarBinArray::from(vec!["four", "four", "four", "four"]); - assert_arrays_eq!(const_array, expected); + assert_arrays_eq!(const_array, expected, &mut assertion_ctx); } #[test] @@ -450,21 +452,24 @@ mod tests { .elements() .clone() .execute::(&mut ctx)?, - PrimitiveArray::from_iter([1u64, 2, 1, 2]) + PrimitiveArray::from_iter([1u64, 2, 1, 2]), + &mut ctx ); assert_arrays_eq!( list_array .offsets() .clone() .execute::(&mut ctx)?, - PrimitiveArray::from_iter([0u64, 2]) + PrimitiveArray::from_iter([0u64, 2]), + &mut ctx ); assert_arrays_eq!( list_array .sizes() .clone() .execute::(&mut ctx)?, - PrimitiveArray::from_iter([2u64, 2]) + PrimitiveArray::from_iter([2u64, 2]), + &mut ctx ); Ok(()) } @@ -491,7 +496,8 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(), - PrimitiveArray::from_iter([0u64, 0]) + PrimitiveArray::from_iter([0u64, 0]), + &mut ctx ); assert_arrays_eq!( canonical_const @@ -499,7 +505,8 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(), - PrimitiveArray::from_iter([0u64, 0]) + PrimitiveArray::from_iter([0u64, 0]), + &mut ctx ); } @@ -524,7 +531,8 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(), - PrimitiveArray::from_iter([0u64, 0]) + PrimitiveArray::from_iter([0u64, 0]), + &mut ctx ); assert_arrays_eq!( canonical_const @@ -532,7 +540,8 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(), - PrimitiveArray::from_iter([0u64, 0]) + PrimitiveArray::from_iter([0u64, 0]), + &mut ctx ); } @@ -593,7 +602,11 @@ mod tests { for i in 0..4 { let list = canonical.fixed_size_list_elements_at(i).unwrap(); let list_primitive = list.execute::(&mut ctx).unwrap(); - assert_arrays_eq!(list_primitive, PrimitiveArray::from_iter([10i32, 20, 30])); + assert_arrays_eq!( + list_primitive, + PrimitiveArray::from_iter([10i32, 20, 30]), + &mut ctx + ); } } @@ -625,7 +638,8 @@ mod tests { .unwrap(); assert_arrays_eq!( elements, - PrimitiveArray::from_iter([1.5f64, 2.5, 1.5, 2.5, 1.5, 2.5]) + PrimitiveArray::from_iter([1.5f64, 2.5, 1.5, 2.5, 1.5, 2.5]), + &mut ctx ); } @@ -738,7 +752,7 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(elements, PrimitiveArray::from_iter([42i16])); + assert_arrays_eq!(elements, PrimitiveArray::from_iter([42i16]), &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 89c571eb827..6eb599153a7 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -294,7 +294,7 @@ mod tests { .into_array() .append_to_builder(builder.as_mut(), &mut ctx)?; let result = builder.finish(); - assert_arrays_eq!(&result, &expected); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index ed2adc7cfc8..0ca0a8abfbf 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -10,6 +10,7 @@ use crate::EqMode; use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; +use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::datetime::TemporalData; use crate::assert_arrays_eq; @@ -26,12 +27,15 @@ use crate::validity::Validity; macro_rules! test_temporal_roundtrip { ($prim:ty, $constructor:expr, $unit:expr) => {{ + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![100 as $prim].into_array(); let temporal: TemporalData = $constructor(array, $unit); assert_arrays_eq!( temporal.temporal_values(), - PrimitiveArray::from_iter([100 as $prim]) + PrimitiveArray::from_iter([100 as $prim]), + &mut assertion_ctx ); assert_eq!(temporal.temporal_metadata().time_unit(), $unit); }}; @@ -145,6 +149,8 @@ test_fail_case!( // We test Timestamp explicitly to avoid the macro getting too complex. #[test] fn test_timestamp() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let ts = buffer![100i64].into_array(); let ts_array = ts.into_array(); @@ -159,7 +165,8 @@ fn test_timestamp() { assert_arrays_eq!( temporal_array.temporal_values(), - PrimitiveArray::from_iter([100i64]) + PrimitiveArray::from_iter([100i64]), + &mut assertion_ctx ); assert_eq!( temporal_array.temporal_metadata(), diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index 144c90c74f9..b28a97933c0 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -104,6 +104,8 @@ mod tests { #[test] fn fill_null_leading_none() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( [None, Some(800i128), None, Some(1000i128), None], @@ -121,7 +123,8 @@ mod tests { .to_decimal(); assert_arrays_eq!( p, - DecimalArray::from_iter([4200, 800, 4200, 1000, 4200], decimal_dtype) + DecimalArray::from_iter([4200, 800, 4200, 1000, 4200], decimal_dtype), + &mut assertion_ctx ); assert_eq!( p.buffer::().as_slice(), @@ -139,6 +142,8 @@ mod tests { #[test] fn fill_null_all_none() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( @@ -158,13 +163,16 @@ mod tests { .to_decimal(); assert_arrays_eq!( p, - DecimalArray::from_iter([25500, 25500, 25500, 25500, 25500], decimal_dtype) + DecimalArray::from_iter([25500, 25500, 25500, 25500, 25500], decimal_dtype), + &mut assertion_ctx ); } /// fill_null with a value that overflows the array's storage type should upcast the array. #[test] fn fill_null_overflow_upcasts() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let decimal_dtype = DecimalDType::new(3, 0); let arr = DecimalArray::from_option_iter([None, Some(10i8), None], decimal_dtype); // i8 max is 127, so 200 doesn't fit — the array should be widened to i16. @@ -180,12 +188,15 @@ mod tests { .to_decimal(); assert_arrays_eq!( result, - DecimalArray::from_iter([200i16, 10, 200], decimal_dtype) + DecimalArray::from_iter([200i16, 10, 200], decimal_dtype), + &mut assertion_ctx ); } #[test] fn fill_null_non_nullable() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::new( @@ -205,7 +216,8 @@ mod tests { .to_decimal(); assert_arrays_eq!( p, - DecimalArray::from_iter([800i128, 1000, 1200, 1400, 1600], decimal_dtype) + DecimalArray::from_iter([800i128, 1000, 1200, 1400, 1600], decimal_dtype), + &mut assertion_ctx ); } } diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 444fe313d15..8ddd8c84b38 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -53,6 +53,7 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -62,6 +63,8 @@ mod tests { #[test] fn test_take() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![10i128, 11i128, 12i128, 13i128], @@ -73,11 +76,13 @@ mod tests { let taken = array.take(indices).unwrap(); let expected = DecimalArray::from_iter([10i128, 12, 13], ddtype); - assert_arrays_eq!(expected, taken); + assert_arrays_eq!(expected, taken, &mut assertion_ctx); } #[test] fn test_take_null_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![i128::MAX, 11i128, 12i128, 13i128], @@ -89,7 +94,7 @@ mod tests { let taken = array.take(indices).unwrap(); let expected = DecimalArray::from_option_iter([None, Some(12i128), Some(13)], ddtype); - assert_arrays_eq!(expected, taken); + assert_arrays_eq!(expected, taken, &mut assertion_ctx); } #[rstest] diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index a8dce4a1254..1e83d022f4c 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -209,6 +209,7 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; @@ -248,6 +249,8 @@ mod tests { #[test] fn test_nullable_decimal_serde_roundtrip() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = DecimalArray::new( buffer![1234567i32, 0i32, -9999999i32], DecimalDType::new(7, 3), @@ -277,6 +280,6 @@ mod tests { ) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 59758924a19..c97b137c514 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -462,6 +462,8 @@ mod test { #[test] fn test_dict_array_from_primitive_chunks() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let len = 2; let chunk_count = 2; let array = make_dict_primitive_chunks::(len, 2, chunk_count); @@ -476,7 +478,7 @@ mod test { let into_prim = array.to_primitive(); let prim_into = builder.finish_into_canonical().into_primitive(); - assert_arrays_eq!(into_prim, prim_into); + assert_arrays_eq!(into_prim, prim_into, &mut assertion_ctx); Ok(()) } diff --git a/vortex-array/src/arrays/dict/compute/cast.rs b/vortex-array/src/arrays/dict/compute/cast.rs index 63b20989f4c..ed881d5db53 100644 --- a/vortex-array/src/arrays/dict/compute/cast.rs +++ b/vortex-array/src/arrays/dict/compute/cast.rs @@ -71,6 +71,8 @@ mod tests { #[test] fn test_cast_dict_to_wider_type() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![1i32, 2, 3, 2, 1].into_array(); let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); @@ -85,7 +87,11 @@ mod tests { #[expect(deprecated)] let decoded = casted.to_primitive(); - assert_arrays_eq!(decoded, PrimitiveArray::from_iter([1i64, 2, 3, 2, 1])); + assert_arrays_eq!( + decoded, + PrimitiveArray::from_iter([1i64, 2, 3, 2, 1]), + &mut assertion_ctx + ); } #[test] @@ -106,6 +112,8 @@ mod tests { #[test] fn test_cast_dict_allvalid_to_nonnullable_and_back() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create an AllValid dict array (no nulls) let values = buffer![10i32, 20, 30, 40].into_array(); let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); @@ -173,7 +181,7 @@ mod tests { let original_values = dict.as_array().to_primitive(); #[expect(deprecated)] let final_values = back_to_non_nullable.to_primitive(); - assert_arrays_eq!(original_values, final_values); + assert_arrays_eq!(original_values, final_values, &mut assertion_ctx); } #[rstest] @@ -187,6 +195,8 @@ mod tests { #[test] fn test_cast_dict_with_unreferenced_null_values_to_nonnullable() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use crate::arrays::DictArray; use crate::validity::Validity; @@ -222,6 +232,10 @@ mod tests { ); #[expect(deprecated)] let casted_prim = casted.to_primitive(); - assert_arrays_eq!(casted_prim, PrimitiveArray::from_iter([1.0f64, 3.0, 1.0])); + assert_arrays_eq!( + casted_prim, + PrimitiveArray::from_iter([1.0f64, 3.0, 1.0]), + &mut assertion_ctx + ); } } diff --git a/vortex-array/src/arrays/dict/compute/fill_null.rs b/vortex-array/src/arrays/dict/compute/fill_null.rs index 8f146d728a3..57f5ab8c34b 100644 --- a/vortex-array/src/arrays/dict/compute/fill_null.rs +++ b/vortex-array/src/arrays/dict/compute/fill_null.rs @@ -108,6 +108,8 @@ mod tests { #[test] fn nullable_codes_fill_in_values() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::new( buffer![0u32, 1, 2], @@ -124,7 +126,11 @@ mod tests { .vortex_expect("operation should succeed in test"); #[expect(deprecated)] let filled_primitive = filled.to_primitive(); - assert_arrays_eq!(filled_primitive, PrimitiveArray::from_iter([10, 20, 20])); + assert_arrays_eq!( + filled_primitive, + PrimitiveArray::from_iter([10, 20, 20]), + &mut assertion_ctx + ); assert!( filled_primitive .all_valid(&mut LEGACY_SESSION.create_execution_ctx()) diff --git a/vortex-array/src/arrays/dict/compute/like.rs b/vortex-array/src/arrays/dict/compute/like.rs index 40bcea8fcf1..dfc4e5e8803 100644 --- a/vortex-array/src/arrays/dict/compute/like.rs +++ b/vortex-array/src/arrays/dict/compute/like.rs @@ -56,6 +56,7 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::DictArray; use crate::arrays::VarBinArray; @@ -68,6 +69,8 @@ mod tests { #[test] fn like_reduce_dict() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict = DictArray::try_new( buffer![0u8, 1, 0, 2].into_array(), VarBinArray::from(vec!["hello", "world", "help"]).into_array(), @@ -79,7 +82,11 @@ mod tests { .try_new_array(4, LikeOptions::default(), [dict, pattern])? .optimize()?; - assert_arrays_eq!(result, BoolArray::from_iter([true, false, true, false])); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, true, false]), + &mut assertion_ctx + ); Ok(()) } } diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index 5ed65c49bbc..adc3470a3dc 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -86,6 +86,8 @@ mod test { #[test] fn canonicalise_nullable_primitive() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values: Vec> = (0..65) .map(|i| match i % 3 { 0 => Some(42), @@ -105,11 +107,13 @@ mod test { let expected = PrimitiveArray::from_option_iter(values); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut assertion_ctx); } #[test] fn canonicalise_non_nullable_primitive_32_unique_values() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let unique_values: Vec = (0..32).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 32])); @@ -121,11 +125,13 @@ mod test { #[expect(deprecated)] let actual = dict.as_array().to_primitive(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut assertion_ctx); } #[test] fn canonicalise_non_nullable_primitive_100_unique_values() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let unique_values: Vec = (0..100).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 100])); @@ -137,7 +143,7 @@ mod test { #[expect(deprecated)] let actual = dict.as_array().to_primitive(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut assertion_ctx); } #[test] @@ -180,6 +186,8 @@ mod test { #[test] fn compare_sliced_dict() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use crate::arrays::BoolArray; let sliced = sliced_dict_array(); let compared = sliced @@ -187,7 +195,7 @@ mod test { .unwrap(); let expected = BoolArray::from_iter([Some(false), None, Some(true)]); - assert_arrays_eq!(compared, expected.into_array()); + assert_arrays_eq!(compared, expected.into_array(), &mut assertion_ctx); } #[test] diff --git a/vortex-array/src/arrays/dict/compute/slice.rs b/vortex-array/src/arrays/dict/compute/slice.rs index 1552c9c9fea..ba92370f16c 100644 --- a/vortex-array/src/arrays/dict/compute/slice.rs +++ b/vortex-array/src/arrays/dict/compute/slice.rs @@ -91,6 +91,7 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::dict::compute::slice::ConstantArray; @@ -102,18 +103,22 @@ mod tests { #[test] fn slice_constant_valid_code() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(1u8, 5).into_array(), buffer![10i32, 20, 30].into_array(), ); let sliced = dict.slice(1..4)?; let expected = PrimitiveArray::from_iter([20i32, 20, 20]).into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut assertion_ctx); Ok(()) } #[test] fn slice_constant_null_code() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(Scalar::null(DType::Primitive(PType::U8, Nullable)), 5).into_array(), buffer![10i32, 20, 30].into_array(), @@ -121,7 +126,7 @@ mod tests { let sliced = dict.slice(1..4)?; let expected = PrimitiveArray::from_option_iter([Option::::None, None, None]).into_array(); - assert_arrays_eq!(sliced, expected); + assert_arrays_eq!(sliced, expected, &mut assertion_ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/dict/tests.rs b/vortex-array/src/arrays/dict/tests.rs index bc69998b0de..ae0454db757 100644 --- a/vortex-array/src/arrays/dict/tests.rs +++ b/vortex-array/src/arrays/dict/tests.rs @@ -5,6 +5,7 @@ use vortex_buffer::buffer; use super::DictArray; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; @@ -14,6 +15,8 @@ use crate::validity::Validity; #[test] fn test_scalar_at_null_code() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::from_option_iter(vec![None, Some(0u32), None]).into_array(), buffer![1i32].into_array(), @@ -21,7 +24,7 @@ fn test_scalar_at_null_code() { .unwrap(); let expected = PrimitiveArray::from_option_iter(vec![None, Some(1i32), None]).into_array(); - assert_arrays_eq!(dict, expected); + assert_arrays_eq!(dict, expected, &mut assertion_ctx); } #[test] diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index 9558eb35308..5b50dbacc85 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -121,6 +121,8 @@ mod tests { #[test] fn cast_timestamp_to_i64() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let ext_dtype = Timestamp::new_with_tz( TimeUnit::Nanoseconds, Some("UTC".into()), @@ -135,7 +137,7 @@ mod tests { result.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_arrays_eq!(result, buffer![1i64, 2, 3].into_array()); + assert_arrays_eq!(result, buffer![1i64, 2, 3].into_array(), &mut assertion_ctx); Ok(()) } diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index 3085debf9e1..0c30352b564 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -124,6 +124,7 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -149,6 +150,8 @@ mod test { #[test] fn filter_fixed_size_list_selects_correct_lists() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let array = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -160,11 +163,13 @@ mod test { let expected = FixedSizeListArray::new(expected_elements.into_array(), 2, Validity::NonNullable, 2); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn filter_degenerate_list_size_zero() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::empty::(Nullability::NonNullable); let array = FixedSizeListArray::new(elements.into_array(), 0, Validity::NonNullable, 5); @@ -175,11 +180,13 @@ mod test { let expected = FixedSizeListArray::new(expected_elements.into_array(), 0, Validity::NonNullable, 3); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn filter_nested_fixed_size_lists() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Inner lists of size 2, outer lists of size 2 (so 2 outer lists, each with 2 inner lists). let inner_elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); let inner_fsl = FixedSizeListArray::new(inner_elements, 2, Validity::NonNullable, 4); @@ -196,6 +203,6 @@ mod test { let expected_outer = FixedSizeListArray::new(expected_inner.into_array(), 2, Validity::NonNullable, 1); - assert_arrays_eq!(filtered, expected_outer); + assert_arrays_eq!(filtered, expected_outer, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index d6d55a72837..ce0fb1cd1d8 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -135,6 +135,8 @@ mod test { #[test] fn filter_listview_selects_correct_lists() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // 3 lists: [10,20], [30,40], [50,60] let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let offsets = buffer![0u32, 2, 4].into_array(); @@ -157,11 +159,13 @@ mod test { Validity::NonNullable, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn test_filter_preserves_unreferenced_elements() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -185,7 +189,8 @@ mod test { // Verify the entire elements array is preserved. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut assertion_ctx ); // Verify offsets are unchanged. @@ -195,6 +200,8 @@ mod test { #[test] fn test_filter_with_gaps() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] @@ -218,7 +225,8 @@ mod test { // Verify the entire elements array is preserved including gaps. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]) + PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), + &mut assertion_ctx ); // Verify offsets are unchanged. @@ -229,7 +237,8 @@ mod test { // Verify the lists still read correctly. assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([7i32, 8, 9]) + PrimitiveArray::from_iter([7i32, 8, 9]), + &mut assertion_ctx ); } diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index ee209b90fc3..06ac484bafa 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -49,6 +49,7 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -86,6 +87,8 @@ mod test { #[test] fn filter_struct_selects_correct_rows() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = StructArray::try_new( ["x", "y"].into(), vec![ @@ -111,11 +114,13 @@ mod test { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn filter_empty_struct() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let mask = Mask::from_iter([ @@ -125,18 +130,20 @@ mod test { let expected = StructArray::try_new(FieldNames::empty(), vec![], 5, Validity::NonNullable).unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn filter_empty_struct_with_empty_filter() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); let filtered = struct_arr.filter(Mask::from_iter::<[bool; 0]>([])).unwrap(); let expected = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] diff --git a/vortex-array/src/arrays/filter/execute/take/tests.rs b/vortex-array/src/arrays/filter/execute/take/tests.rs index ecfebe46679..36ab1a63dcb 100644 --- a/vortex-array/src/arrays/filter/execute/take/tests.rs +++ b/vortex-array/src/arrays/filter/execute/take/tests.rs @@ -7,6 +7,7 @@ use vortex_mask::Mask; use crate::IntoArray; use crate::RecursiveCanonical; +use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::DecimalArray; use crate::arrays::Dict; @@ -58,14 +59,15 @@ fn test_take_execute_kernel_maps_indices_through_filter() -> VortexResult<()> { filter.clone(), )? .into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - PrimitiveArray::from_option_iter([Some(40i32), None, Some(10)]).into_array() + PrimitiveArray::from_option_iter([Some(40i32), None, Some(10)]).into_array(), + &mut ctx ); Ok(()) } @@ -86,7 +88,7 @@ fn test_take_execute_kernel_nullable_fast_path_maps_indices_through_filter() -> filter.clone(), )? .into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -94,7 +96,8 @@ fn test_take_execute_kernel_nullable_fast_path_maps_indices_through_filter() -> assert!(result.as_opt::().is_some()); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - PrimitiveArray::from_option_iter([Some(40i32), None, Some(20)]).into_array() + PrimitiveArray::from_option_iter([Some(40i32), None, Some(20)]).into_array(), + &mut ctx ); Ok(()) } @@ -107,7 +110,7 @@ fn test_take_execute_kernel_fast_path_maps_indices_through_filter() -> VortexRes ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 3].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -115,7 +118,8 @@ fn test_take_execute_kernel_fast_path_maps_indices_through_filter() -> VortexRes assert!(result.as_opt::().is_some()); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - PrimitiveArray::from_iter([50i32, 20, 60]).into_array() + PrimitiveArray::from_iter([50i32, 20, 60]).into_array(), + &mut ctx ); Ok(()) } @@ -127,7 +131,7 @@ fn assert_take_execute_rejects_out_of_bounds_rank( ) -> VortexResult<()> { let filter = FilterArray::new(child, filter_mask).into_array(); let parent = DictArray::try_new(codes, filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); if let Err(err) = execute_parent(&filter, &parent, 1, &mut ctx) { assert!( @@ -205,7 +209,7 @@ fn test_take_execute_kernel_handles_empty_sequential_take() -> VortexResult<()> filter.clone(), )? .into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -217,7 +221,8 @@ fn test_take_execute_kernel_handles_empty_sequential_take() -> VortexResult<()> buffer![0u32].into_array(), Validity::NonNullable, )? - .into_array() + .into_array(), + &mut ctx ); Ok(()) } @@ -229,12 +234,16 @@ fn assert_take_execute_maps_child_dtype( let filter = FilterArray::new(child, Mask::from_iter([true, false, true, true, false])).into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); - assert_arrays_eq!(result.execute::(&mut ctx)?.0, expected); + assert_arrays_eq!( + result.execute::(&mut ctx)?.0, + expected, + &mut ctx + ); Ok(()) } @@ -246,7 +255,7 @@ fn test_take_execute_kernel_skips_bool_filter_child() -> VortexResult<()> { ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)?; @@ -268,7 +277,7 @@ fn execute_primitive_take( .into_array(); let indices = PrimitiveArray::from_iter((0..take_len).map(|idx| (idx % filtered_len) as u64)); let parent = DictArray::try_new(indices.into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); execute_parent(&filter, &parent, 1, &mut ctx) } @@ -335,7 +344,7 @@ fn test_take_execute_kernel_handles_nullable_primitive_filter_child() -> VortexR ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)?; @@ -344,7 +353,8 @@ fn test_take_execute_kernel_handles_nullable_primitive_filter_child() -> VortexR .expect("filter child should execute its take parent") .execute::(&mut ctx)? .0, - PrimitiveArray::from_option_iter([Some(40i32), Some(10), None]).into_array() + PrimitiveArray::from_option_iter([Some(40i32), Some(10), None]).into_array(), + &mut ctx ); Ok(()) } @@ -357,7 +367,7 @@ fn test_take_execute_kernel_preserves_nullable_all_valid_fixed_width_child() -> ) .into_array(); let parent = DictArray::try_new(buffer![0u64, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -365,7 +375,8 @@ fn test_take_execute_kernel_preserves_nullable_all_valid_fixed_width_child() -> assert_eq!(result.dtype(), parent.dtype()); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - PrimitiveArray::new(buffer![10i32, 20], Validity::AllValid).into_array() + PrimitiveArray::new(buffer![10i32, 20], Validity::AllValid).into_array(), + &mut ctx ); Ok(()) } @@ -383,7 +394,7 @@ fn test_take_execute_kernel_handles_nullable_decimal_filter_child() -> VortexRes ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)?; @@ -393,7 +404,8 @@ fn test_take_execute_kernel_handles_nullable_decimal_filter_child() -> VortexRes .execute::(&mut ctx)? .0, DecimalArray::from_option_iter([Some(400i128), Some(100), None], decimal_dtype) - .into_array() + .into_array(), + &mut ctx ); Ok(()) } @@ -476,7 +488,7 @@ fn test_take_execute_kernel_preserves_nullable_indices_dtype_fast_path() -> Vort filter.clone(), )? .into_array(); - let mut ctx = ExecutionCtx::new(crate::array_session()); + let mut ctx = crate::array_session().create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its nullable take parent"); @@ -484,7 +496,8 @@ fn test_take_execute_kernel_preserves_nullable_indices_dtype_fast_path() -> Vort assert_eq!(result.dtype(), parent.dtype()); assert_arrays_eq!( result.execute::(&mut ctx)?.0, - VarBinViewArray::from_iter_nullable_str([Some("a"), Some("b")]).into_array() + VarBinViewArray::from_iter_nullable_str([Some("a"), Some("b")]).into_array(), + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs index 0228790f0fd..5615753e90c 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs @@ -8,6 +8,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -51,6 +52,8 @@ fn test_filter_degenerate_list_size_zero( #[case] mask_values: Vec, #[case] expected_len: usize, ) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let new_validity = if matches!(validity, Validity::NonNullable) { Validity::NonNullable } else { @@ -74,11 +77,13 @@ fn test_filter_degenerate_list_size_zero( new_validity, expected_len, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn test_filter_with_nulls() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4), Some(5), Some(6)]); let validity = Validity::from_iter([true, false, true]); @@ -101,11 +106,13 @@ fn test_filter_with_nulls() { 2, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn test_filter_all_null_array() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create an array where all elements are null. let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let validity = Validity::AllInvalid; @@ -116,12 +123,14 @@ fn test_filter_all_null_array() { // Verify the result is an array of nulls. assert_eq!(filtered.len(), 2, "All-null FSL should produce 2 elements"); - assert_nth_scalar_is_null!(filtered, 0); - assert_nth_scalar_is_null!(filtered, 1); + assert_nth_scalar_is_null!(filtered, 0, &mut assertion_ctx); + assert_nth_scalar_is_null!(filtered, 1, &mut assertion_ctx); } #[test] fn test_filter_nested_fixed_size_lists() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create nested fixed-size lists: FSL>. // Inner lists are of size 2, outer lists are of size 3. // So we have 2 outer lists, each containing 3 inner lists, each containing 2 i32s. @@ -160,7 +169,7 @@ fn test_filter_nested_fixed_size_lists() { let expected_outer = FixedSizeListArray::new(expected_inner.into_array(), 3, Validity::NonNullable, 1); - assert_arrays_eq!(filtered, expected_outer); + assert_arrays_eq!(filtered, expected_outer, &mut assertion_ctx); } // Conformance tests using rstest for various array configurations. @@ -247,6 +256,8 @@ fn create_fsl_empty() -> ArrayRef { #[test] fn test_filter_all_null_various_list_sizes() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test filtering with all-null arrays of different list sizes. // The implementation returns ConstantArray only when validity_mask() is Mask::AllFalse. @@ -257,8 +268,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered0 = fsl0.filter(mask0).unwrap(); assert_eq!(filtered0.len(), 2); // Check that all elements are null (might be ConstantArray or FixedSizeListArray). - assert_nth_scalar_is_null!(filtered0, 0); - assert_nth_scalar_is_null!(filtered0, 1); + assert_nth_scalar_is_null!(filtered0, 0, &mut assertion_ctx); + assert_nth_scalar_is_null!(filtered0, 1, &mut assertion_ctx); // Case 2: list_size == 1. let elements1 = buffer![1i32, 2, 3].into_array(); @@ -267,8 +278,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered1 = fsl1.filter(mask1).unwrap(); assert_eq!(filtered1.len(), 2); // Check that all elements are null. - assert_nth_scalar_is_null!(filtered1, 0); - assert_nth_scalar_is_null!(filtered1, 1); + assert_nth_scalar_is_null!(filtered1, 0, &mut assertion_ctx); + assert_nth_scalar_is_null!(filtered1, 1, &mut assertion_ctx); // Case 3: list_size == 10 (large). let elements10 = buffer![0..50i32].into_array(); @@ -277,14 +288,16 @@ fn test_filter_all_null_various_list_sizes() { let filtered10 = fsl10.filter(mask10).unwrap(); assert_eq!(filtered10.len(), 5); // Check that all elements are null. - assert_nth_scalar_is_null!(filtered10, 0); - assert_nth_scalar_is_null!(filtered10, 4); + assert_nth_scalar_is_null!(filtered10, 0, &mut assertion_ctx); + assert_nth_scalar_is_null!(filtered10, 4, &mut assertion_ctx); } // Note: test_filter_to_empty_degenerate has been consolidated into test_filter_degenerate_list_size_zero above. #[test] fn test_mask_expansion_threshold_boundary() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test with list_size == 8 (the FSL_SPARSE_MASK_LIST_SIZE_THRESHOLD). let list_size = 8u32; let num_lists = 100; @@ -323,7 +336,7 @@ fn test_mask_expansion_threshold_boundary() { 3, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); // Test with list_size == 7 (just below threshold). let list_size_7 = 7u32; @@ -355,12 +368,14 @@ fn test_mask_expansion_threshold_boundary() { 3, ); - assert_arrays_eq!(filtered7, expected7); + assert_arrays_eq!(filtered7, expected7, &mut assertion_ctx); } // Test FSL-specific behavior with very large list sizes. #[test] fn test_filter_large_list_size() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test with list_size=100, which is significantly larger than typical use cases. let list_size = 100u32; let num_lists = 5; @@ -394,7 +409,7 @@ fn test_filter_large_list_size() { 3, ); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); // Test edge case: filter out all but one large list. let mask_single = Mask::from_iter([false, false, true, false, false]); @@ -414,5 +429,5 @@ fn test_filter_large_list_size() { 1, ); - assert_arrays_eq!(filtered_single, expected_single); + assert_arrays_eq!(filtered_single, expected_single, &mut assertion_ctx); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 8bccb80c8cb..5875b380267 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -40,6 +40,8 @@ fn test_take_fsl_conformance(#[case] fsl: FixedSizeListArray) { #[test] fn test_take_basic_smoke_test() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -53,7 +55,7 @@ fn test_take_basic_smoke_test() { Validity::NonNullable, 3, ); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut assertion_ctx); } // Parameterized test for FSL-specific degenerate (list_size=0) cases. @@ -107,6 +109,8 @@ fn test_take_degenerate_lists( #[test] fn test_take_large_list_size() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = buffer![0i32..300].into_array(); let fsl = FixedSizeListArray::new(elements, 100, Validity::NonNullable, 3); @@ -116,11 +120,13 @@ fn test_take_large_list_size() { // Expected: [[200..300], [0..100]] let expected_elems = PrimitiveArray::from_iter((200i32..300).chain(0..100)).into_array(); let expected = FixedSizeListArray::new(expected_elems, 100, Validity::NonNullable, 2); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut assertion_ctx); } #[test] fn test_take_fsl_with_null_indices_preserves_elements() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -135,7 +141,7 @@ fn test_take_fsl_with_null_indices_preserves_elements() { Validity::from_iter([true, false, true]), 3, ); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut assertion_ctx); } // Element index overflow: with u8 indices and list_size=16, data_idx=16 produces element index @@ -167,8 +173,10 @@ fn test_element_index_overflow( #[case] indices: ArrayRef, #[case] expected: FixedSizeListArray, ) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let result = fsl.take(indices).unwrap(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } // Parameterized test for nullable array scenarios that are specific to FSL's implementation. diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index dd5619559dc..b085d9571e1 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -499,7 +499,7 @@ mod tests { let mut ctx = LEGACY_SESSION.create_execution_ctx(); let reference = interleave_reference(&values, &array_indices, &row_indices, &mut ctx)?; - assert_arrays_eq!(interleaved, reference); + assert_arrays_eq!(interleaved, reference, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index b58b5e0a365..26e2a03dbbf 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -113,6 +113,8 @@ fn test_simple_list_filter() { #[test] fn test_list_filter_dense_mask() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test filtering with a dense mask (high density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -139,11 +141,13 @@ fn test_list_filter_dense_mask() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn test_list_filter_sparse_mask() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test filtering with a sparse mask (low density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -172,11 +176,13 @@ fn test_list_filter_sparse_mask() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn test_list_filter_empty_lists() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test filtering arrays that contain empty lists. let elements = buffer![0..10].into_array(); let offsets = buffer![0, 0, 3, 3, 7, 10, 10].into_array(); // Lists at indices 0, 2, 5 are empty. @@ -200,7 +206,7 @@ fn test_list_filter_empty_lists() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] @@ -233,6 +239,8 @@ fn test_list_filter_with_nulls() { #[test] fn test_list_filter_all_true() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test filtering with an all-true mask. let elements = buffer![0..20].into_array(); let offsets = buffer![0, 5, 10, 15, 20].into_array(); @@ -250,7 +258,7 @@ fn test_list_filter_all_true() { assert_eq!(filtered.len(), 4); let expected = ListArray::try_new(elements, offsets, validity).unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] @@ -275,6 +283,8 @@ fn test_list_filter_all_false() { #[test] fn test_list_filter_single_element() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test filtering to keep only one element. let elements = buffer![0..50].into_array(); let offsets = buffer![0, 10, 20, 30, 40, 50].into_array(); @@ -298,11 +308,13 @@ fn test_list_filter_single_element() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn test_list_filter_alternating_pattern() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test filtering with an alternating pattern. let elements = buffer![0..60].into_array(); let offsets = buffer![0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60].into_array(); @@ -333,11 +345,13 @@ fn test_list_filter_alternating_pattern() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] fn test_list_filter_variable_sizes() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test filtering lists with highly variable sizes. let elements = buffer![0..100].into_array(); let offsets = buffer![0, 1, 2, 5, 10, 20, 35, 60, 100].into_array(); @@ -373,7 +387,7 @@ fn test_list_filter_variable_sizes() { ) .unwrap(); - assert_arrays_eq!(filtered, expected); + assert_arrays_eq!(filtered, expected, &mut assertion_ctx); } #[test] @@ -574,6 +588,8 @@ fn create_list_of_lists_nullable(data: OptVec>>) -> ListArray #[test] #[expect(clippy::cognitive_complexity)] fn test_list_of_lists() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let data = vec![ Some(vec![Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]), Some(vec![Some(vec![Some(4), Some(5), Some(6)])]), @@ -604,11 +620,19 @@ fn test_list_of_lists() { // Check first inner list [1, 2]. let first_inner = first_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(first_inner, PrimitiveArray::from_iter([1, 2])); + assert_arrays_eq!( + first_inner, + PrimitiveArray::from_iter([1, 2]), + &mut assertion_ctx + ); // Check second inner list [3]. let second_inner = first_outer_list.list_elements_at(1).unwrap(); - assert_arrays_eq!(second_inner, PrimitiveArray::from_iter([3])); + assert_arrays_eq!( + second_inner, + PrimitiveArray::from_iter([3]), + &mut assertion_ctx + ); // Check the second list of lists [[4, 5, 6]]. let second_outer = list_of_lists.list_elements_at(1).unwrap(); @@ -616,7 +640,11 @@ fn test_list_of_lists() { assert_eq!(second_outer_list.len(), 1); let inner = second_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(inner, PrimitiveArray::from_iter([4, 5, 6])); + assert_arrays_eq!( + inner, + PrimitiveArray::from_iter([4, 5, 6]), + &mut assertion_ctx + ); // Check the third list of lists (empty). let third_outer = list_of_lists.list_elements_at(2).unwrap(); @@ -629,7 +657,7 @@ fn test_list_of_lists() { assert_eq!(fourth_outer_list.len(), 1); let inner = fourth_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(inner, PrimitiveArray::from_iter([7])); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([7]), &mut assertion_ctx); // Test scalar conversion. let scalar = list_of_lists @@ -792,7 +820,7 @@ fn test_list_of_lists_both_nullable() { let third_list = third_outer.as_::(); assert_eq!(third_list.len(), 1); let inner = third_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(inner, PrimitiveArray::from_iter([3])); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([3]), &mut ctx); // Fourth outer list should have a null inner list. let fourth_outer = list_of_lists.list_elements_at(3).unwrap(); diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index 1423a14804c..f682dbe8714 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -232,7 +232,7 @@ mod tests { buffer![2u32, 2, 3].into_array(), Validity::NonNullable, ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -269,7 +269,7 @@ mod tests { buffer![1u32, 0, 1].into_array(), Validity::Array(BoolArray::from_iter([true, false, true]).into_array()), ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -307,7 +307,7 @@ mod tests { buffer![2u32, 1, 2].into_array(), Validity::AllValid, ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -362,7 +362,7 @@ mod tests { buffer![2u32, 1].into_array(), Validity::NonNullable, ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index 7b553041c68..ae85225c440 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -331,37 +331,47 @@ mod tests { // Verify structure. assert_eq!(list_view.len(), 4); - assert_arrays_eq!(elements, list_view.elements().clone()); + assert_arrays_eq!(elements, list_view.elements().clone(), &mut ctx); // Verify offsets (should be same but without last element). let expected_offsets = buffer![0u32, 3, 5, 7].into_array(); - assert_arrays_eq!(expected_offsets, list_view.offsets().clone()); + assert_arrays_eq!(expected_offsets, list_view.offsets().clone(), &mut ctx); // Verify sizes. let expected_sizes = buffer![3u32, 2, 2, 3].into_array(); - assert_arrays_eq!(expected_sizes, list_view.sizes().clone()); + assert_arrays_eq!(expected_sizes, list_view.sizes().clone(), &mut ctx); // Verify data integrity. - assert_arrays_eq!(list_array, list_view); + assert_arrays_eq!(list_array, list_view, &mut ctx); Ok(()) } #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let list_view = create_basic_listview(); let list_array = list_from_list_view(list_view.clone(), &mut SESSION.create_execution_ctx())?; // Should have same elements. - assert_arrays_eq!(list_view.elements().clone(), list_array.elements().clone()); + assert_arrays_eq!( + list_view.elements().clone(), + list_array.elements().clone(), + &mut assertion_ctx + ); // ListArray offsets should have n+1 elements for n lists (add the final offset). // Check that the first n offsets match. let list_array_offsets_without_last = list_array.offsets().slice(0..list_view.len())?; - assert_arrays_eq!(list_view.offsets().clone(), list_array_offsets_without_last); + assert_arrays_eq!( + list_view.offsets().clone(), + list_array_offsets_without_last, + &mut assertion_ctx + ); // Verify data integrity. - assert_arrays_eq!(list_view, list_array); + assert_arrays_eq!(list_view, list_array, &mut assertion_ctx); Ok(()) } @@ -409,12 +419,14 @@ mod tests { // Round-trip conversion. let converted_back = list_from_list_view(nullable_list_view, &mut ctx)?; - assert_arrays_eq!(nullable_list, converted_back); + assert_arrays_eq!(nullable_list, converted_back, &mut ctx); Ok(()) } #[test] fn test_non_zero_copy_listview_to_list() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create ListViewArray with overlapping lists (not zero-copyable). let list_view = create_overlapping_listview(); let list_array = @@ -428,7 +440,7 @@ mod tests { } // The data should still be correct even though it required a rebuild. - assert_arrays_eq!(list_view, list_array); + assert_arrays_eq!(list_view, list_array, &mut assertion_ctx); Ok(()) } @@ -448,7 +460,7 @@ mod tests { // Round-trip. let converted_back = list_view_from_list(list_array, &mut ctx)?; - assert_arrays_eq!(empty_lists_view, converted_back); + assert_arrays_eq!(empty_lists_view, converted_back, &mut ctx); Ok(()) } @@ -474,8 +486,8 @@ mod tests { assert_eq!(list_view_i64.sizes().dtype(), i64_offsets.dtype()); // Verify data integrity. - assert_arrays_eq!(list_i32, list_view_i32); - assert_arrays_eq!(list_i64, list_view_i64); + assert_arrays_eq!(list_i32, list_view_i32, &mut ctx); + assert_arrays_eq!(list_i64, list_view_i64, &mut ctx); Ok(()) } @@ -487,20 +499,20 @@ mod tests { let original = create_basic_listview(); let to_list = list_from_list_view(original.clone(), &mut ctx)?; let back_to_view = list_view_from_list(to_list, &mut ctx)?; - assert_arrays_eq!(original, back_to_view); + assert_arrays_eq!(original, back_to_view, &mut ctx); // Test 2: Nullable round-trip. let nullable = create_nullable_listview(); let nullable_to_list = list_from_list_view(nullable.clone(), &mut ctx)?; let nullable_back = list_view_from_list(nullable_to_list, &mut ctx)?; - assert_arrays_eq!(nullable, nullable_back); + assert_arrays_eq!(nullable, nullable_back, &mut ctx); // Test 3: Non-zero-copyable round-trip. let overlapping = create_overlapping_listview(); let overlapping_to_list = list_from_list_view(overlapping.clone(), &mut ctx)?; let overlapping_back = list_view_from_list(overlapping_to_list, &mut ctx)?; - assert_arrays_eq!(overlapping, overlapping_back); + assert_arrays_eq!(overlapping, overlapping_back, &mut ctx); Ok(()) } @@ -517,11 +529,11 @@ mod tests { // Verify sizes are all 1. let expected_sizes = buffer![1u32, 1, 1].into_array(); - assert_arrays_eq!(expected_sizes, list_view.sizes().clone()); + assert_arrays_eq!(expected_sizes, list_view.sizes().clone(), &mut ctx); // Round-trip. let converted_back = list_from_list_view(list_view, &mut ctx)?; - assert_arrays_eq!(single_elem_list, converted_back); + assert_arrays_eq!(single_elem_list, converted_back, &mut ctx); Ok(()) } @@ -538,16 +550,18 @@ mod tests { // Verify sizes. let expected_sizes = buffer![2u32, 0, 1, 0, 3].into_array(); - assert_arrays_eq!(expected_sizes, list_view.sizes().clone()); + assert_arrays_eq!(expected_sizes, list_view.sizes().clone(), &mut ctx); // Round-trip. let converted_back = list_from_list_view(list_view, &mut ctx)?; - assert_arrays_eq!(mixed_list, converted_back); + assert_arrays_eq!(mixed_list, converted_back, &mut ctx); Ok(()) } #[test] fn test_recursive_simple_listview() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let list_view = create_basic_listview(); let result = recursive_list_from_list_view( list_view.clone().into_array(), @@ -555,12 +569,14 @@ mod tests { )?; assert_eq!(result.len(), list_view.len()); - assert_arrays_eq!(list_view.into_array(), result); + assert_arrays_eq!(list_view.into_array(), result, &mut assertion_ctx); Ok(()) } #[test] fn test_recursive_nested_listview() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let inner_elements = buffer![1i32, 2, 3].into_array(); let inner_offsets = buffer![0u32, 2].into_array(); let inner_sizes = buffer![2u32, 1].into_array(); @@ -592,12 +608,14 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(outer_listview.into_array(), result); + assert_arrays_eq!(outer_listview.into_array(), result, &mut assertion_ctx); Ok(()) } #[test] fn test_recursive_struct_with_listview_fields() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let listview_field = create_basic_listview().into_array(); let primitive_field = buffer![10i32, 20, 30, 40].into_array(); @@ -614,12 +632,14 @@ mod tests { )?; assert_eq!(result.len(), 4); - assert_arrays_eq!(struct_array.into_array(), result); + assert_arrays_eq!(struct_array.into_array(), result, &mut assertion_ctx); Ok(()) } #[test] fn test_recursive_fixed_size_list_with_listview_elements() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let lv1_elements = buffer![1i32, 2].into_array(); let lv1_offsets = buffer![0u32].into_array(); let lv1_sizes = buffer![2u32].into_array(); @@ -659,12 +679,14 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(fixed_list.into_array(), result); + assert_arrays_eq!(fixed_list.into_array(), result, &mut assertion_ctx); Ok(()) } #[test] fn test_recursive_deep_nesting() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let innermost_elements = buffer![1i32, 2, 3].into_array(); let innermost_offsets = buffer![0u32, 2].into_array(); let innermost_sizes = buffer![2u32, 1].into_array(); @@ -703,7 +725,7 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(outer_listview.into_array(), result); + assert_arrays_eq!(outer_listview.into_array(), result, &mut assertion_ctx); Ok(()) } @@ -733,7 +755,7 @@ mod tests { let result = recursive_list_from_list_view(struct_array.clone().into_array(), &mut ctx)?; assert_eq!(result.len(), 4); - assert_arrays_eq!(struct_array.into_array(), result); + assert_arrays_eq!(struct_array.into_array(), result, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 7ef752e4dcb..bae2c4f3fb9 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -455,12 +455,14 @@ mod tests { // Verify the data is correct assert_arrays_eq!( flattened.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); assert_arrays_eq!( flattened.list_elements_at(1)?, - PrimitiveArray::from_iter([2i32, 3]) + PrimitiveArray::from_iter([2i32, 3]), + &mut ctx ); Ok(()) } @@ -495,12 +497,14 @@ mod tests { // Verify valid lists contain correct data assert_arrays_eq!( flattened.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); assert_arrays_eq!( flattened.list_elements_at(2)?, - PrimitiveArray::from_iter([3i32]) + PrimitiveArray::from_iter([3i32]), + &mut ctx ); Ok(()) } @@ -536,12 +540,14 @@ mod tests { // Verify the data is correct. assert_arrays_eq!( trimmed.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); assert_arrays_eq!( trimmed.list_elements_at(1)?, - PrimitiveArray::from_iter([3i32, 4]) + PrimitiveArray::from_iter([3i32, 4]), + &mut ctx ); // Note that element at index 2 (97) is preserved as a gap. @@ -598,12 +604,14 @@ mod tests { // Verify data is preserved assert_arrays_eq!( exact.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); assert_arrays_eq!( exact.list_elements_at(1)?, - PrimitiveArray::from_iter([3i32, 4]) + PrimitiveArray::from_iter([3i32, 4]), + &mut ctx ); Ok(()) } @@ -626,7 +634,8 @@ mod tests { let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; assert_arrays_eq!( trimmed.list_elements_at(1)?, - PrimitiveArray::from_iter([30i32, 40]) + PrimitiveArray::from_iter([30i32, 40]), + &mut ctx ); Ok(()) } @@ -647,7 +656,8 @@ mod tests { let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; assert_arrays_eq!( trimmed.list_elements_at(1)?, - PrimitiveArray::from_iter([30i32, 40]) + PrimitiveArray::from_iter([30i32, 40]), + &mut ctx ); Ok(()) } @@ -670,11 +680,13 @@ mod tests { // Values: [1,2,3] and [2,3]. assert_arrays_eq!( rebuilt.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); assert_arrays_eq!( rebuilt.list_elements_at(1)?, - PrimitiveArray::from_iter([2i32, 3]) + PrimitiveArray::from_iter([2i32, 3]), + &mut ctx ); // Signed input -> signed result (offsets widened to i32, sizes kept i16). @@ -711,7 +723,7 @@ mod tests { let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; // min(offsets) = 0, so nothing to trim; output should equal input. - assert_arrays_eq!(trimmed, listview); + assert_arrays_eq!(trimmed, listview, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index 61664c2f608..c596b108f07 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -26,6 +26,8 @@ use crate::validity::Validity; #[test] fn test_basic_listview_comprehensive() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Comprehensive test for basic ListView functionality including scalar_at. // Logical lists: [[1,2,3], [4,5], [6,7,8,9]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -50,7 +52,8 @@ fn test_basic_listview_comprehensive() { // Check individual list elements. assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut assertion_ctx ); // Test scalar_at which returns entire lists as Scalar values. @@ -68,17 +71,21 @@ fn test_basic_listview_comprehensive() { assert_arrays_eq!( listview.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([4i32, 5]) + PrimitiveArray::from_iter([4i32, 5]), + &mut assertion_ctx ); assert_arrays_eq!( listview.list_elements_at(2).unwrap(), - PrimitiveArray::from_iter([6i32, 7, 8, 9]) + PrimitiveArray::from_iter([6i32, 7, 8, 9]), + &mut assertion_ctx ); } #[test] fn test_out_of_order_offsets() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // ListView-specific: Tests that offsets can be non-sequential and out-of-order. // Logical lists: [[7,8,9], [1,2,3], [4,5,6]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -92,13 +99,15 @@ fn test_out_of_order_offsets() { // First list starts at offset 6: [7, 8, 9]. assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([7i32, 8, 9]) + PrimitiveArray::from_iter([7i32, 8, 9]), + &mut assertion_ctx ); // Second list starts at offset 0: [1, 2, 3]. assert_arrays_eq!( listview.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut assertion_ctx ); } @@ -136,7 +145,8 @@ fn test_from_list_array() -> VortexResult<()> { // Check first list. assert_arrays_eq!( list_view.list_elements_at(0)?, - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut ctx ); // Check validity is preserved. @@ -147,7 +157,8 @@ fn test_from_list_array() -> VortexResult<()> { // Check third list. assert_arrays_eq!( list_view.list_elements_at(2)?, - PrimitiveArray::from_iter([5i32, 6, 7]) + PrimitiveArray::from_iter([5i32, 6, 7]), + &mut ctx ); Ok(()) } @@ -158,6 +169,8 @@ fn test_from_list_array() -> VortexResult<()> { #[case::constant_offsets(false, true)] // Varying sizes, constant offsets #[case::both_constant(true, true)] // Both constant fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_offsets: bool) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Logical lists vary by case: // - constant_sizes: [[1,2,3], [4,5,6], [7,8,9]] (size 3 each, varying offsets) // - constant_offsets: [[1,2,3], [1,2], [1]] (all start at 0, varying sizes) @@ -190,7 +203,11 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o // All lists are identical [1, 2, 3] (overlapping). let expected = PrimitiveArray::from_iter([1i32, 2, 3]); for i in 0..3 { - assert_arrays_eq!(listview.list_elements_at(i).unwrap(), expected); + assert_arrays_eq!( + listview.list_elements_at(i).unwrap(), + expected, + &mut assertion_ctx + ); } } else if const_sizes { // All lists have size 3, different offsets (no overlap). diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index 756b1192aef..83504a1a225 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -39,6 +39,8 @@ fn test_filter_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_filter_preserves_unreferenced_elements() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -61,7 +63,8 @@ fn test_filter_preserves_unreferenced_elements() { // Verify the entire elements array is preserved. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut assertion_ctx ); // Verify offsets are unchanged. @@ -71,6 +74,8 @@ fn test_filter_preserves_unreferenced_elements() { #[test] fn test_filter_with_gaps() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] @@ -93,7 +98,8 @@ fn test_filter_with_gaps() { // Verify the entire elements array is preserved including gaps. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]) + PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), + &mut assertion_ctx ); // Verify offsets are unchanged. @@ -104,7 +110,8 @@ fn test_filter_with_gaps() { // Verify the lists still read correctly. assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([7i32, 8, 9]) + PrimitiveArray::from_iter([7i32, 8, 9]), + &mut assertion_ctx ); } diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index c911b9ba7a5..ba28419e3de 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -84,6 +84,8 @@ fn test_slice_comprehensive() { #[test] fn test_slice_out_of_order() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // ListView-specific: Test slicing with out-of-order offsets. // Logical lists: [[70,80], [10,20,30], [40,50,60], [90], [30]] let elements = buffer![10i32, 20, 30, 40, 50, 60, 70, 80, 90].into_array(); @@ -123,15 +125,18 @@ fn test_slice_out_of_order() { // Verify the actual list contents are correct. assert_arrays_eq!( sliced_list.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([10i32, 20, 30]) + PrimitiveArray::from_iter([10i32, 20, 30]), + &mut assertion_ctx ); assert_arrays_eq!( sliced_list.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([40i32, 50, 60]) + PrimitiveArray::from_iter([40i32, 50, 60]), + &mut assertion_ctx ); assert_arrays_eq!( sliced_list.list_elements_at(2).unwrap(), - PrimitiveArray::from_iter([90i32]) + PrimitiveArray::from_iter([90i32]), + &mut assertion_ctx ); } @@ -390,6 +395,8 @@ fn test_cast_large_dataset() { #[test] fn test_zip_widens_false_element_nullability() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // [[1, 2], [3], [4]] let if_true = ListViewArray::new( buffer![1i32, 2, 3, 4].into_array(), @@ -429,12 +436,14 @@ fn test_zip_widens_false_element_nullability() -> VortexResult<()> { Validity::NonNullable, ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn test_zip_widens_true_element_nullability() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // [[1, null], [3], [4]] let if_true = ListViewArray::new( PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(), @@ -474,7 +483,7 @@ fn test_zip_widens_true_element_nullability() -> VortexResult<()> { Validity::NonNullable, ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 4391af98ee2..7c86cb360fb 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -40,6 +40,8 @@ fn test_take_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_take_preserves_unreferenced_elements() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // ListView-specific: Test that take preserves the entire elements array // even when taking only a subset of lists. let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -60,7 +62,8 @@ fn test_take_preserves_unreferenced_elements() { // Verify the entire elements array is preserved. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), + &mut assertion_ctx ); // Verify offsets are preserved. @@ -70,6 +73,8 @@ fn test_take_preserves_unreferenced_elements() { #[test] fn test_take_with_gaps() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // ListView-specific: Test with gaps in elements array. // Elements with gaps (999 values are "gaps" between used ranges). let elements = buffer![1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12].into_array(); @@ -87,13 +92,15 @@ fn test_take_with_gaps() { // Verify the entire elements array is preserved including gaps. assert_arrays_eq!( result_list.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]) + PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), + &mut assertion_ctx ); // Verify the lists still read correctly despite gaps. assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([7i32, 8, 9]) + PrimitiveArray::from_iter([7i32, 8, 9]), + &mut assertion_ctx ); } diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index b26d101eb33..fb5f2e2727e 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -95,6 +95,8 @@ fn test_masked_child_with_validity() { #[test] fn test_masked_child_all_valid() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // When validity is AllValid, masked_child should invert to AllInvalid. let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); @@ -108,7 +110,8 @@ fn test_masked_child_all_valid() { ); assert_arrays_eq!( PrimitiveArray::from_option_iter([10i32, 20, 30].map(Some)), - array + array, + &mut assertion_ctx ); } diff --git a/vortex-array/src/arrays/patched/compute/compare.rs b/vortex-array/src/arrays/patched/compute/compare.rs index 0f1e6add00f..a9262b6621e 100644 --- a/vortex-array/src/arrays/patched/compute/compare.rs +++ b/vortex-array/src/arrays/patched/compute/compare.rs @@ -163,9 +163,8 @@ mod tests { use vortex_error::VortexResult; use vortex_error::vortex_err; - use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::Patched; @@ -189,7 +188,7 @@ mod tests { ) .unwrap(); - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx) .unwrap() @@ -207,7 +206,7 @@ mod tests { let expected = BoolArray::from_indices(512, [509, 510, 511], Validity::NonNullable).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -222,7 +221,7 @@ mod tests { ) .unwrap(); - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx).unwrap(); // Slice the array so that the first patch should be skipped. @@ -240,7 +239,7 @@ mod tests { let expected = BoolArray::from_indices(502, [500, 501], Validity::NonNullable).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -259,7 +258,7 @@ mod tests { None, )?; - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx)? .into_array() .try_downcast::() @@ -277,7 +276,7 @@ mod tests { let expected = BoolArray::from_indices(512, [510], Validity::NonNullable).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } @@ -293,7 +292,7 @@ mod tests { None, )?; - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx)? .into_array() .try_downcast::() @@ -311,7 +310,7 @@ mod tests { let expected = BoolArray::from_indices(10, [7], Validity::NonNullable).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/patched/compute/filter.rs b/vortex-array/src/arrays/patched/compute/filter.rs index d2c60941059..8cf0c186fb6 100644 --- a/vortex-array/src/arrays/patched/compute/filter.rs +++ b/vortex-array/src/arrays/patched/compute/filter.rs @@ -67,9 +67,8 @@ mod tests { use vortex_error::VortexResult; use vortex_mask::Mask; - use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -79,7 +78,7 @@ mod tests { #[test] fn test_filter_noop() -> VortexResult<()> { // Filter that doesn't prune any chunks (all data fits in one chunk). - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let array = buffer![u16::MIN; 5].into_array(); let patched_indices = buffer![3u8, 4].into_array(); @@ -98,7 +97,7 @@ mod tests { // Values at indices 0 and 4: MIN and MAX. let expected = PrimitiveArray::from_iter([u16::MIN, u16::MAX]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -106,7 +105,7 @@ mod tests { #[test] fn test_filter_with_offset() -> VortexResult<()> { // Test filtering where offset > 0. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![5u16, 1030].into_array(); @@ -126,7 +125,7 @@ mod tests { .execute::(&mut ctx)?; let expected = PrimitiveArray::from_iter([u16::MAX, u16::MIN, u16::MIN, u16::MAX]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -134,7 +133,7 @@ mod tests { #[test] fn test_filter_basic() -> VortexResult<()> { // Basic test: filter with mask that crosses boundaries. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![1024u16, 1025].into_array(); @@ -153,7 +152,7 @@ mod tests { let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -161,7 +160,7 @@ mod tests { #[test] fn test_filter_complex() -> VortexResult<()> { // Filter with mask that crosses boundaries, with patches offset. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![1024u16, 1025].into_array(); @@ -180,7 +179,7 @@ mod tests { let expected = PrimitiveArray::from_iter([u16::MAX, u16::MIN, u16::MIN]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -188,7 +187,7 @@ mod tests { #[test] fn test_filter_sliced() -> VortexResult<()> { // Test filter on a sliced PatchedArray to exercise codepath where offset > 0. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); // Create a larger array (6 chunks) so we can slice and still have room // for the filter to prune chunks. @@ -217,7 +216,7 @@ mod tests { let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -226,7 +225,7 @@ mod tests { fn test_filter_with_offset_nonuniform() -> VortexResult<()> { // Test filtering with offset > 0 using non-uniform base values. // This catches slice_chunks bugs where inner coordinates are miscalculated. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); // Use non-uniform values so that incorrect slicing is detectable. let base_values: Vec = (0u16..4096).collect(); @@ -254,7 +253,7 @@ mod tests { // Expected: 9999 (patched at logical 0), 6 (original at logical 1), 8888 (patched at logical 1025). let expected = PrimitiveArray::from_iter([9999u16, 6, 8888]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } @@ -263,7 +262,7 @@ mod tests { fn test_filter_with_offset_last_chunk() -> VortexResult<()> { // Test filtering with offset > 0 where the mask touches the last chunk. // This ensures we don't accidentally slice past the end of the array or mask. - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); // Create a 6-chunk array (6144 elements). let array = buffer![u16::MIN; 6144].into_array(); @@ -294,7 +293,7 @@ mod tests { // Expected: patch at 3976 (was 5000), patch at 4976 (was 6000), and MIN at 5119. let expected = PrimitiveArray::from_iter([u16::MAX, u16::MAX, u16::MIN]); - assert_arrays_eq!(expected, filtered); + assert_arrays_eq!(expected, filtered, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/patched/compute/take.rs b/vortex-array/src/arrays/patched/compute/take.rs index 893d18db1ac..248ec8dc616 100644 --- a/vortex-array/src/arrays/patched/compute/take.rs +++ b/vortex-array/src/arrays/patched/compute/take.rs @@ -136,8 +136,8 @@ mod tests { use vortex_session::VortexSession; use crate::ArrayRef; - use crate::ExecutionCtx; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -159,7 +159,7 @@ mod tests { )?; let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); Patched::from_array_and_patches(values, &patches, &mut ctx)? .into_array() @@ -168,6 +168,8 @@ mod tests { #[test] fn test_take_basic() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -177,13 +179,15 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 10, 0, 30, 0]).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut assertion_ctx); Ok(()) } #[test] fn test_take_sliced() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = make_patched_array(&[0; 10], &[1, 3], &[100, 200], 2..10)?; let indices = buffer![0u32, 1, 2, 3, 7].into_array(); @@ -191,13 +195,15 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 200, 0, 0, 0]).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut assertion_ctx); Ok(()) } #[test] fn test_take_out_of_order() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -207,13 +213,15 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 30, 0, 10, 0]).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut assertion_ctx); Ok(()) } #[test] fn test_take_duplicates() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at index [2] with value [99] let array = make_patched_array(&[0; 5], &[2], &[99], 0..5)?; @@ -227,13 +235,15 @@ mod tests { let _canonical = result.to_canonical()?.into_primitive(); let expected = PrimitiveArray::from_iter([99u16, 99, 0, 99]).into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut assertion_ctx); Ok(()) } #[test] fn test_take_with_null_indices() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use crate::arrays::BoolArray; use crate::validity::Validity; @@ -271,7 +281,7 @@ mod tests { .into_array(), ), ); - assert_arrays_eq!(expected.into_array(), result); + assert_arrays_eq!(expected.into_array(), result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index 7790d734ed6..aff9f426457 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -354,9 +354,9 @@ mod tests { use crate::ArrayContext; use crate::ArraySlots; use crate::Canonical; - use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::Patched; use crate::arrays::PatchedArray; use crate::arrays::PrimitiveArray; @@ -384,7 +384,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -417,7 +417,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -450,7 +450,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -467,7 +467,7 @@ mod tests { expected[3] = 30; let expected = expected.into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -483,7 +483,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -500,7 +500,7 @@ mod tests { expected[0] = 30; let expected = expected.into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] @@ -520,7 +520,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -546,7 +546,7 @@ mod tests { ]) .into_array(); - assert_arrays_eq!(expected, result); + assert_arrays_eq!(expected, result, &mut ctx); } fn make_patched_array( @@ -564,7 +564,7 @@ mod tests { let patches = Patches::new(len, 0, indices, patch_vals, None)?; let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); Patched::from_array_and_patches(array, &patches, &mut ctx) } @@ -640,11 +640,11 @@ mod tests { assert_eq!(array_ref.dtype(), new_array.dtype()); // Execute both and compare results - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = VortexSession::empty().create_execution_ctx(); let original_executed = array_ref.execute::(&mut ctx)?.into_primitive(); let new_executed = new_array.execute::(&mut ctx)?.into_primitive(); - assert_arrays_eq!(original_executed, new_executed); + assert_arrays_eq!(original_executed, new_executed, &mut ctx); Ok(()) } @@ -666,12 +666,12 @@ mod tests { let new_array = array_ref.with_slots(slots.into_slots())?; // Execute and verify the inner values changed (except at patch positions) - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = VortexSession::empty().create_execution_ctx(); let executed = new_array.execute::(&mut ctx)?.into_primitive(); // Expected: all 5s except indices 1, 2, 3 which are patched to 10, 20, 30 let expected = PrimitiveArray::from_iter([5u16, 10, 20, 30, 5, 5, 5, 5, 5, 5]); - assert_arrays_eq!(expected, executed); + assert_arrays_eq!(expected, executed, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/patched/vtable/operations.rs b/vortex-array/src/arrays/patched/vtable/operations.rs index f0491666e56..a6468ef03aa 100644 --- a/vortex-array/src/arrays/patched/vtable/operations.rs +++ b/vortex-array/src/arrays/patched/vtable/operations.rs @@ -60,7 +60,6 @@ mod tests { use vortex_buffer::buffer; use vortex_session::VortexSession; - use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -83,7 +82,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -128,7 +127,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -160,7 +159,7 @@ mod tests { .unwrap(); let session = VortexSession::empty(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() diff --git a/vortex-array/src/arrays/patched/vtable/slice.rs b/vortex-array/src/arrays/patched/vtable/slice.rs index 8655e0c7d84..1f59c2a854d 100644 --- a/vortex-array/src/arrays/patched/vtable/slice.rs +++ b/vortex-array/src/arrays/patched/vtable/slice.rs @@ -65,9 +65,8 @@ mod tests { use vortex_error::VortexResult; use crate::Canonical; - use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -81,7 +80,7 @@ mod tests { let patch_values = buffer![u16::MAX; 3].into_array(); let patches = Patches::new(512, 0, patch_indices, patch_values, None)?; - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)?; @@ -127,7 +126,7 @@ mod tests { ) .unwrap(); - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let patched_array = Patched::from_array_and_patches(inner.into_array(), &patches, &mut ctx) .unwrap() @@ -148,7 +147,7 @@ mod tests { .slice(range) .unwrap(); - assert_arrays_eq!(slice_first, slice_last); + assert_arrays_eq!(slice_first, slice_last, &mut ctx); } #[test] @@ -159,7 +158,7 @@ mod tests { let patched_values = buffer![0u64, 1, 2, 3, 4, 5].into_array(); let patches = Patches::new(10_000, 0, patched_indices, patched_values, None).unwrap(); - let mut ctx = ExecutionCtx::new(LEGACY_SESSION.clone()); + let mut ctx = crate::array_session().create_execution_ctx(); let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() @@ -181,6 +180,6 @@ mod tests { let expected = expected.into_array(); - assert_arrays_eq!(expected, sliced); + assert_arrays_eq!(expected, sliced, &mut ctx); } } diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index 197db4d3854..3161b2562e5 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -137,6 +137,7 @@ mod tests { use super::*; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::assert_arrays_eq; use crate::validity::Validity; @@ -173,13 +174,16 @@ mod tests { #[test] fn patch_sliced() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid); let sliced = input.slice(2..8).unwrap(); #[expect(deprecated)] let sliced_primitive = sliced.to_primitive(); assert_arrays_eq!( sliced_primitive, - PrimitiveArray::new(buffer![2u32; 6], Validity::AllValid) + PrimitiveArray::new(buffer![2u32; 6], Validity::AllValid), + &mut assertion_ctx ); } } diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 5c5ab43401f..64ec5542a7c 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -286,12 +286,18 @@ mod test { #[test] fn cast_u32_u8() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); // cast from u32 to u8 #[expect(deprecated)] let p = arr.cast(PType::U8.into()).unwrap().to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([0u8, 10, 200]), + &mut assertion_ctx + ); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); // to nullable @@ -303,7 +309,8 @@ mod test { .to_primitive(); assert_arrays_eq!( p, - PrimitiveArray::new(buffer![0u8, 10, 200], Validity::AllValid) + PrimitiveArray::new(buffer![0u8, 10, 200], Validity::AllValid), + &mut assertion_ctx ); assert!(matches!(p.validity(), Ok(Validity::AllValid))); @@ -314,7 +321,11 @@ mod test { .cast(DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([0u8, 10, 200]), + &mut assertion_ctx + ); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); // to nullable u32 @@ -326,7 +337,8 @@ mod test { .to_primitive(); assert_arrays_eq!( p, - PrimitiveArray::new(buffer![0u32, 10, 200], Validity::AllValid) + PrimitiveArray::new(buffer![0u32, 10, 200], Validity::AllValid), + &mut assertion_ctx ); assert!(matches!(p.validity(), Ok(Validity::AllValid))); @@ -337,16 +349,26 @@ mod test { .cast(DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([0u8, 10, 200]), + &mut assertion_ctx + ); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); } #[test] fn cast_u32_f32() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); #[expect(deprecated)] let u8arr = arr.cast(PType::F32.into()).unwrap().to_primitive(); - assert_arrays_eq!(u8arr, PrimitiveArray::from_iter([0.0f32, 10., 200.])); + assert_arrays_eq!( + u8arr, + PrimitiveArray::from_iter([0.0f32, 10., 200.]), + &mut assertion_ctx + ); } #[test] @@ -380,6 +402,8 @@ mod test { #[test] fn cast_with_invalid_nulls() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = PrimitiveArray::new( buffer![-1i32, 0, 10], Validity::from_iter([false, true, true]), @@ -392,7 +416,8 @@ mod test { .to_primitive(); assert_arrays_eq!( p, - PrimitiveArray::from_option_iter([None, Some(0u32), Some(10)]) + PrimitiveArray::from_option_iter([None, Some(0u32), Some(10)]), + &mut assertion_ctx ); assert_eq!( p.as_ref() @@ -408,6 +433,8 @@ mod test { /// buffer without allocation (pointer identity). #[test] fn cast_same_width_int_reinterprets_buffer() -> vortex_error::VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let src = PrimitiveArray::from_iter([0u32, 10, 100]); let src_ptr = src.as_slice::().as_ptr(); @@ -417,7 +444,11 @@ mod test { // Zero-copy: the data pointer should be identical. assert_eq!(src_ptr as usize, dst_ptr as usize); - assert_arrays_eq!(dst, PrimitiveArray::from_iter([0i32, 10, 100])); + assert_arrays_eq!( + dst, + PrimitiveArray::from_iter([0i32, 10, 100]), + &mut assertion_ctx + ); Ok(()) } @@ -453,6 +484,8 @@ mod test { /// not prevent the cast from succeeding. #[test] fn cast_same_width_int_nullable_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // The null position holds u32::MAX which doesn't fit in i32, but it's // masked as invalid so the cast should still succeed via reinterpret. let arr = PrimitiveArray::new( @@ -466,13 +499,16 @@ mod test { .to_primitive(); assert_arrays_eq!( casted, - PrimitiveArray::from_option_iter([None, Some(0i32), Some(42)]) + PrimitiveArray::from_option_iter([None, Some(0i32), Some(42)]), + &mut assertion_ctx ); Ok(()) } #[test] fn cast_u32_to_u8_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1000u32, 10u32, 42u32], Validity::from_iter([false, true, true]), @@ -484,7 +520,8 @@ mod test { .to_primitive(); assert_arrays_eq!( casted, - PrimitiveArray::from_option_iter([None, Some(10u8), Some(42)]) + PrimitiveArray::from_option_iter([None, Some(10u8), Some(42)]), + &mut assertion_ctx ); Ok(()) } diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index 67089bd534e..93d815a9365 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -64,6 +64,8 @@ mod test { #[test] fn fill_null_leading_none() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([None, Some(8u8), None, Some(10), None]); #[expect(deprecated)] let p = arr @@ -71,7 +73,11 @@ mod test { .fill_null(Scalar::from(42u8)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([42u8, 8, 42, 10, 42])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([42u8, 8, 42, 10, 42]), + &mut assertion_ctx + ); assert!( p.as_ref() .validity() @@ -84,6 +90,8 @@ mod test { #[test] fn fill_null_all_none() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Option::::None, None, None, None, None]); #[expect(deprecated)] @@ -92,7 +100,11 @@ mod test { .fill_null(Scalar::from(255u8)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([255u8, 255, 255, 255, 255])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([255u8, 255, 255, 255, 255]), + &mut assertion_ctx + ); assert!( p.as_ref() .validity() @@ -105,6 +117,8 @@ mod test { #[test] fn fill_null_nullable_non_null() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = PrimitiveArray::new( buffer![8u8, 10, 12, 14, 16], Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()), @@ -115,7 +129,11 @@ mod test { .fill_null(Scalar::from(255u8)) .unwrap() .to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([8u8, 10, 12, 14, 16]), + &mut assertion_ctx + ); assert!( p.as_ref() .validity() @@ -128,10 +146,16 @@ mod test { #[test] fn fill_null_non_nullable() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = buffer![8u8, 10, 12, 14, 16].into_array(); #[expect(deprecated)] let p = arr.fill_null(Scalar::from(255u8)).unwrap().to_primitive(); - assert_arrays_eq!(p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16])); + assert_arrays_eq!( + p, + PrimitiveArray::from_iter([8u8, 10, 12, 14, 16]), + &mut assertion_ctx + ); assert!( p.as_ref() .validity() diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index c253ea8e2f7..ee0547244b2 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -222,6 +222,7 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -229,6 +230,8 @@ mod tests { #[test] fn take_null_index_skips_out_of_bounds_value() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = PrimitiveArray::from_iter([10i32, 20, 30]); let indices = PrimitiveArray::new( buffer![1u64, 3], @@ -239,7 +242,8 @@ mod tests { assert_arrays_eq!( taken, - PrimitiveArray::from_option_iter([Some(20i32), None]).into_array() + PrimitiveArray::from_option_iter([Some(20i32), None]).into_array(), + &mut assertion_ctx ); } } diff --git a/vortex-array/src/arrays/primitive/compute/zip.rs b/vortex-array/src/arrays/primitive/compute/zip.rs index 0dcb7298460..99b92c78703 100644 --- a/vortex-array/src/arrays/primitive/compute/zip.rs +++ b/vortex-array/src/arrays/primitive/compute/zip.rs @@ -181,7 +181,7 @@ mod tests { (0..len).map(|i| if bits[i] { i as i64 } else { 1_000 + i as i64 }), ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -217,7 +217,7 @@ mod tests { } })) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index d9d778c34bb..d669f861fcf 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -226,6 +226,7 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -234,6 +235,8 @@ mod tests { #[test] fn test_nullable_primitive_serde_roundtrip() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::new( buffer![1i32, 2, 3, 4], Validity::from_iter([true, false, true, false]), @@ -262,6 +265,6 @@ mod tests { ) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index 67de19d936f..eba022d1e5d 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -76,6 +76,8 @@ mod tests { #[test] fn test_scalar_fn_add() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let lhs = buffer![1i32, 2, 3].into_array(); let rhs = buffer![10i32, 20, 30].into_array(); @@ -89,7 +91,7 @@ mod tests { .execute::(&mut LEGACY_SESSION.create_execution_ctx())? .into_array(); let expected = buffer![11i32, 22, 33].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } @@ -130,6 +132,8 @@ mod tests { #[test] fn test_scalar_fn_mul() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let lhs = buffer![2i32, 3, 4].into_array(); let rhs = buffer![5i32, 6, 7].into_array(); @@ -141,13 +145,15 @@ mod tests { .execute::(&mut LEGACY_SESSION.create_execution_ctx())? .into_array(); let expected = buffer![10i32, 18, 28].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn test_scalar_fn_with_nullable() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let rhs = PrimitiveArray::new( buffer![10i32, 20, 30], @@ -167,13 +173,15 @@ mod tests { Validity::from_iter([true, false, true]), ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn test_scalar_fn_comparison() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let lhs = buffer![1i32, 5, 3].into_array(); let rhs = buffer![2i32, 5, 1].into_array(); @@ -185,7 +193,7 @@ mod tests { .execute::(&mut LEGACY_SESSION.create_execution_ctx())? .into_array(); let expected = BoolArray::from_iter([false, true, false]).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } diff --git a/vortex-array/src/arrays/shared/tests.rs b/vortex-array/src/arrays/shared/tests.rs index fd454a8521e..d1a33ed3e20 100644 --- a/vortex-array/src/arrays/shared/tests.rs +++ b/vortex-array/src/arrays/shared/tests.rs @@ -5,8 +5,8 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use crate::Canonical; -use crate::ExecutionCtx; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::SharedArray; use crate::arrays::shared::SharedArrayExt; @@ -20,7 +20,7 @@ fn shared_array_caches_on_canonicalize() -> VortexResult<()> { let shared = SharedArray::new(array); let session = crate::array_session(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); let first = shared.get_or_compute(|source| source.clone().execute::(&mut ctx))?; diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 1ca91254123..efbfb029734 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -188,18 +188,25 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::assert_arrays_eq; #[test] fn test_slice_slice() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Slice(1..4, Slice(2..8, base)) combines to Slice(3..6, base) let arr = PrimitiveArray::from_iter(0i32..10).into_array(); let inner_slice = SliceArray::new(arr, 2..8).into_array(); let slice = inner_slice.slice(1..4)?; - assert_arrays_eq!(slice, PrimitiveArray::from_iter([3i32, 4, 5])); + assert_arrays_eq!( + slice, + PrimitiveArray::from_iter([3i32, 4, 5]), + &mut assertion_ctx + ); Ok(()) } diff --git a/vortex-array/src/arrays/struct_/compute/cast.rs b/vortex-array/src/arrays/struct_/compute/cast.rs index 2cadb386e38..ac14983f0e2 100644 --- a/vortex-array/src/arrays/struct_/compute/cast.rs +++ b/vortex-array/src/arrays/struct_/compute/cast.rs @@ -227,7 +227,8 @@ mod tests { assert_eq!(result.dtype(), &target); assert_arrays_eq!( result.unmasked_field_by_name("b").unwrap(), - ConstantArray::new(Scalar::null(utf8_null), 1) + ConstantArray::new(Scalar::null(utf8_null), 1), + &mut ctx ); } diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index c30ef532008..9614d601616 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -40,6 +40,8 @@ mod tests { #[test] fn take_empty_struct() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let indices = PrimitiveArray::from_option_iter([Some(1), None]); @@ -52,7 +54,8 @@ mod tests { vec![], 2, Validity::from_iter([true, false]) - ) + ), + &mut assertion_ctx ); } @@ -92,6 +95,8 @@ mod tests { #[test] fn take_field_struct() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let struct_arr = StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).into_array())]) .unwrap(); @@ -103,7 +108,8 @@ mod tests { [("a", buffer![1, 0])], Validity::from_iter([true, false]) ) - .unwrap() + .unwrap(), + &mut assertion_ctx ); } diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 8443c2dd25e..6a4b803096a 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -165,6 +165,8 @@ mod tests { #[test] fn test_struct_cast_field_reorder() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Source: {a, b}, Target: {c, b, a} - reordered + new null field let source = StructArray::try_new( FieldNames::from(["a", "b"]), @@ -196,15 +198,18 @@ mod tests { .unwrap(); assert_arrays_eq!( result.unmasked_field_by_name("a").unwrap(), - VarBinViewArray::from_iter_nullable_str([Some("A")]) + VarBinViewArray::from_iter_nullable_str([Some("A")]), + &mut assertion_ctx ); assert_arrays_eq!( result.unmasked_field_by_name("b").unwrap(), - VarBinViewArray::from_iter_nullable_str([Some("B")]) + VarBinViewArray::from_iter_nullable_str([Some("B")]), + &mut assertion_ctx ); assert_arrays_eq!( result.unmasked_field_by_name("c").unwrap(), - ConstantArray::new(Scalar::null(utf8_null), 1) + ConstantArray::new(Scalar::null(utf8_null), 1), + &mut assertion_ctx ); } @@ -298,6 +303,8 @@ mod tests { #[test] fn cast_struct_drop_field() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Casting to a struct with a subset of fields should succeed. let source = StructArray::try_new( FieldNames::from(["a", "b", "c"]), @@ -331,16 +338,20 @@ mod tests { assert_eq!(result.unmasked_fields().len(), 2); assert_arrays_eq!( result.unmasked_field_by_name("a").unwrap(), - buffer![1i32, 2, 3].into_array() + buffer![1i32, 2, 3].into_array(), + &mut assertion_ctx ); assert_arrays_eq!( result.unmasked_field_by_name("c").unwrap(), - buffer![100u8, 200, 255].into_array() + buffer![100u8, 200, 255].into_array(), + &mut assertion_ctx ); } #[test] fn cast_struct_field_type_widening() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Casting struct fields to wider types (i32 -> i64). let source = StructArray::try_new( FieldNames::from(["val"]), @@ -370,7 +381,8 @@ mod tests { ); assert_arrays_eq!( result.unmasked_field_by_name("val").unwrap(), - buffer![1i64, 2, 3].into_array() + buffer![1i64, 2, 3].into_array(), + &mut assertion_ctx ); } diff --git a/vortex-array/src/arrays/struct_/tests.rs b/vortex-array/src/arrays/struct_/tests.rs index 2db9e62f4f5..da198c7c1e2 100644 --- a/vortex-array/src/arrays/struct_/tests.rs +++ b/vortex-array/src/arrays/struct_/tests.rs @@ -24,6 +24,8 @@ use crate::validity::Validity; #[test] fn test_project() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = VarBinArray::from_vec( vec!["a", "b", "c", "d", "e"], @@ -52,15 +54,22 @@ fn test_project() { let bools = struct_b.unmasked_field(0); assert_arrays_eq!( bools, - BoolArray::from_iter([true, true, true, false, false]) + BoolArray::from_iter([true, true, true, false, false]), + &mut assertion_ctx ); let prims = struct_b.unmasked_field(1); - assert_arrays_eq!(prims, PrimitiveArray::from_iter([0i64, 1, 2, 3, 4])); + assert_arrays_eq!( + prims, + PrimitiveArray::from_iter([0i64, 1, 2, 3, 4]), + &mut assertion_ctx + ); } #[test] fn test_remove_column() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = PrimitiveArray::new(buffer![4u64, 5, 6, 7, 8], Validity::NonNullable); @@ -77,7 +86,11 @@ fn test_remove_column() { removed.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_arrays_eq!(removed, PrimitiveArray::from_iter([0i64, 1, 2, 3, 4])); + assert_arrays_eq!( + removed, + PrimitiveArray::from_iter([0i64, 1, 2, 3, 4]), + &mut assertion_ctx + ); assert_eq!(data.names(), &["ys"]); assert_eq!(data.struct_fields().nfields(), 1); @@ -88,7 +101,8 @@ fn test_remove_column() { ); assert_arrays_eq!( data.unmasked_field(0), - PrimitiveArray::from_iter([4u64, 5, 6, 7, 8]) + PrimitiveArray::from_iter([4u64, 5, 6, 7, 8]), + &mut assertion_ctx ); let empty = data.remove_column("non_existent"); @@ -101,6 +115,8 @@ fn test_remove_column() { #[test] fn test_duplicate_field_names() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test that StructArray allows duplicate field names and returns the first match let field1 = buffer![1i32, 2, 3].into_array(); let field2 = buffer![10i32, 20, 30].into_array(); @@ -119,19 +135,25 @@ fn test_duplicate_field_names() { let first_value_field = struct_array.unmasked_field_by_name("value").unwrap(); assert_arrays_eq!( first_value_field, - PrimitiveArray::from_iter([1i32, 2, 3]) // This is field1, not field3 + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut assertion_ctx ); // Verify field_by_name_opt also returns the first match let opt_field = struct_array.unmasked_field_by_name_opt("value").unwrap(); assert_arrays_eq!( opt_field, - PrimitiveArray::from_iter([1i32, 2, 3]) // First "value" field + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut assertion_ctx ); // Verify the third field (second "value") can be accessed by index let third_field = struct_array.unmasked_field(2); - assert_arrays_eq!(third_field, PrimitiveArray::from_iter([100i32, 200, 300])); + assert_arrays_eq!( + third_field, + PrimitiveArray::from_iter([100i32, 200, 300]), + &mut assertion_ctx + ); } #[test] diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index 9aef7d0f162..b6055cfb144 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -254,6 +254,7 @@ mod tests { use vortex_buffer::ByteBuffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; @@ -289,6 +290,8 @@ mod tests { /// [`CompareKernel`]: super::CompareKernel #[test] fn varbin_i64_offsets_compare_constant() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -304,11 +307,13 @@ mod tests { .unwrap(); let expected = BoolArray::from_iter([true, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[test] fn varbin_i64_offsets_compare_constant_binary() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -328,6 +333,6 @@ mod tests { .unwrap(); let expected = BoolArray::from_iter([true, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index a0c78a64363..ddddc3fbd37 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -240,6 +240,8 @@ mod test { #[test] fn filter_var_bin_test() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -257,11 +259,17 @@ mod test { ) .unwrap(); - assert_arrays_eq!(buf, VarBinArray::from(vec!["hello", "filter"])); + assert_arrays_eq!( + buf, + VarBinArray::from(vec!["hello", "filter"]), + &mut assertion_ctx + ); } #[test] fn filter_var_bin_slice_test() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -282,11 +290,17 @@ mod test { ) .unwrap(); - assert_arrays_eq!(buf, VarBinArray::from(vec!["hello", "filter", "filter3"])); + assert_arrays_eq!( + buf, + VarBinArray::from(vec!["hello", "filter", "filter3"]), + &mut assertion_ctx + ); } #[test] fn filter_var_bin_slice_null() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bytes = [ b"one".as_slice(), b"two".as_slice(), @@ -322,12 +336,15 @@ mod test { Some("three"), Some("five"), Some("six") - ]) + ]), + &mut assertion_ctx ); } #[test] fn filter_varbin_nulls() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bytes = [b"".as_slice(), b"two".as_slice(), b"two".as_slice()] .into_iter() .flat_map(|x| x.iter().cloned()) @@ -346,11 +363,17 @@ mod test { ) .unwrap(); - assert_arrays_eq!(buf, VarBinArray::from(vec![None, Some("two")])); + assert_arrays_eq!( + buf, + VarBinArray::from(vec![None, Some("two")]), + &mut assertion_ctx + ); } #[test] fn filter_varbin_all_null() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let offsets = buffer![0, 0, 0, 0].into_array(); let validity = Validity::Array(BoolArray::from_iter([false, false, false]).into_array()); let arr = VarBinArray::try_new( @@ -370,7 +393,11 @@ mod test { ) .unwrap(); - assert_arrays_eq!(buf, VarBinArray::from(vec![None::<&str>, None])); + assert_arrays_eq!( + buf, + VarBinArray::from(vec![None::<&str>, None]), + &mut assertion_ctx + ); } #[test] diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 5826633aea3..3345eb6d952 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -256,6 +256,7 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::compute::take::PrimitiveArray; @@ -304,6 +305,8 @@ mod tests { #[test] fn test_take_overflow() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let scream = std::iter::once("a").cycle().take(128).collect::(); let bytes = ByteBuffer::copy_from(scream.as_bytes()); let offsets = buffer![0u8, 128u8].into_array(); @@ -322,6 +325,6 @@ mod tests { [Some(scream.clone()), Some(scream.clone()), Some(scream)], DType::Utf8(Nullability::NonNullable), ); - assert_arrays_eq!(expected, taken); + assert_arrays_eq!(expected, taken, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/varbin/tests.rs b/vortex-array/src/arrays/varbin/tests.rs index 4711c69102e..eb5bb97e5de 100644 --- a/vortex-array/src/arrays/varbin/tests.rs +++ b/vortex-array/src/arrays/varbin/tests.rs @@ -8,6 +8,7 @@ use vortex_buffer::buffer; use crate::ArrayRef; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; @@ -32,17 +33,23 @@ fn binary_array() -> ArrayRef { #[rstest] pub fn test_scalar_at(binary_array: ArrayRef) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); assert_arrays_eq!( binary_array, - VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) + VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), + &mut assertion_ctx ); } #[rstest] pub fn slice_array(binary_array: ArrayRef) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let binary_arr = binary_array.slice(1..2).unwrap(); assert_arrays_eq!( binary_arr, - VarBinViewArray::from_iter_str(["hello world this is a long string"]) + VarBinViewArray::from_iter_str(["hello world this is a long string"]), + &mut assertion_ctx ); } diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index fc30ee2c94e..f4739bf70a1 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -96,6 +96,8 @@ mod tests { #[case(DType::Utf8(Nullability::NonNullable))] #[case(DType::Binary(Nullability::NonNullable))] fn test_canonical_varbin_unsliced(#[case] dtype: DType) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let varbin = VarBinArray::from_iter_nonnull(["foo", "bar", "baz"], dtype.clone()); #[expect(deprecated)] let canonical = varbin.as_array().to_varbinview(); @@ -103,7 +105,7 @@ mod tests { DType::Utf8(_) => VarBinViewArray::from_iter_str(["foo", "bar", "baz"]), _ => VarBinViewArray::from_iter_bin(["foo", "bar", "baz"]), }; - assert_arrays_eq!(canonical, expected); + assert_arrays_eq!(canonical, expected, &mut assertion_ctx); } // Empty array: offsets has exactly one element; no elements to canonicalize. diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 6effc7c656a..ebcceb55652 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -202,6 +202,8 @@ mod tests { use crate::dtype::Nullability; #[test] fn test_optimize_compacts_buffers() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create a VarBinViewArray with some long strings that will create multiple buffers let original = VarBinViewArray::from_iter_nullable_str([ Some("short"), @@ -235,12 +237,15 @@ mod tests { // Verify the data is still correct assert_arrays_eq!( optimized_array, - >::from_iter([Some("short"), Some("tiny")]) + >::from_iter([Some("short"), Some("tiny")]), + &mut assertion_ctx ); } #[test] fn test_optimize_with_long_strings() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create strings that are definitely longer than 12 bytes let long_string_1 = "this is definitely a very long string that exceeds the inline limit"; let long_string_2 = "another extremely long string that also needs external buffer storage"; @@ -270,12 +275,15 @@ mod tests { // Verify the data is still correct assert_arrays_eq!( optimized_array, - VarBinArray::from(vec![long_string_1, long_string_3]) + VarBinArray::from(vec![long_string_1, long_string_3]), + &mut assertion_ctx ); } #[test] fn test_optimize_no_buffers() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create an array with only short strings (all inlined) let original = VarBinViewArray::from_iter_str(["a", "bb", "ccc", "dddd"]); @@ -287,11 +295,13 @@ mod tests { assert_eq!(optimized_array.data_buffers().len(), 0); - assert_arrays_eq!(optimized_array, original); + assert_arrays_eq!(optimized_array, original, &mut assertion_ctx); } #[test] fn test_optimize_single_buffer() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create an array that naturally has only one buffer let str1 = "this is a long string that goes into a buffer"; let str2 = "another long string in the same buffer"; @@ -306,11 +316,13 @@ mod tests { assert_eq!(optimized_array.data_buffers().len(), 1); - assert_arrays_eq!(optimized_array, original); + assert_arrays_eq!(optimized_array, original, &mut assertion_ctx); } #[test] fn test_selective_compaction_with_threshold_zero() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // threshold=0 should keep all buffers (no compaction) let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -333,11 +345,13 @@ mod tests { assert_eq!(compacted.data_buffers().len(), taken.data_buffers().len()); // Verify correctness - assert_arrays_eq!(compacted, taken); + assert_arrays_eq!(compacted, taken, &mut assertion_ctx); } #[test] fn test_selective_compaction_with_high_threshold() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // threshold=1.0 should compact any buffer with waste let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -361,11 +375,13 @@ mod tests { assert!(compacted.data_buffers().len() <= original_buffers); // Verify correctness - assert_arrays_eq!(compacted, taken); + assert_arrays_eq!(compacted, taken, &mut assertion_ctx); } #[test] fn test_selective_compaction_preserves_well_utilized_buffers() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create an array with multiple strings in one buffer (well-utilized) let str1 = "first long string that needs external buffer storage"; let str2 = "second long string also in buffer"; @@ -383,11 +399,13 @@ mod tests { assert_eq!(compacted.data_buffers().len(), 1); // Verify all data is correct - assert_arrays_eq!(compacted, original); + assert_arrays_eq!(compacted, original, &mut assertion_ctx); } #[test] fn test_selective_compaction_with_mixed_utilization() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create array with some long strings let strings: Vec = (0..10) .map(|i| { @@ -414,11 +432,13 @@ mod tests { [0, 2, 4, 6, 8].map(|i| Some(strings[i].as_str())), DType::Utf8(Nullability::NonNullable), ); - assert_arrays_eq!(expected, compacted); + assert_arrays_eq!(expected, compacted, &mut assertion_ctx); } #[test] fn test_slice_strategy_with_contiguous_range() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create array with strings that will be in one buffer let strings: Vec = (0..20) .map(|i| format!("this is a long string number {} for slice test", i)) @@ -447,7 +467,7 @@ mod tests { ); // Verify correctness - assert_arrays_eq!(&compacted, taken); + assert_arrays_eq!(&compacted, taken, &mut assertion_ctx); // Verify that if there was only one buffer, the compacted version also has one // (it was sliced, not rewritten into multiple buffers) diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index 2892ff16283..82bcb01b939 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -3,40 +3,50 @@ #[expect(deprecated)] use crate::ToCanonical as _; +use crate::VortexSessionExecute; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; #[test] pub fn varbin_view() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]); assert_arrays_eq!( binary_arr, - VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) + VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), + &mut assertion_ctx ); } #[test] pub fn slice_array() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) .slice(1..2) .unwrap(); assert_arrays_eq!( binary_arr, - VarBinViewArray::from_iter_str(["hello world this is a long string"]) + VarBinViewArray::from_iter_str(["hello world this is a long string"]), + &mut assertion_ctx ); } #[test] pub fn flatten_array() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["string1", "string2"]); #[expect(deprecated)] let var_bin = binary_arr.as_array().to_varbinview(); assert_arrays_eq!( var_bin, - VarBinViewArray::from_iter_str(["string1", "string2"]) + VarBinViewArray::from_iter_str(["string1", "string2"]), + &mut assertion_ctx ); } diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index a9ffe5a5859..9545aef2c9c 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -249,6 +249,7 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -266,6 +267,8 @@ mod tests { let len = array.len(); let ctx = ArrayContext::empty(); + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let serialized = array .clone() .into_array() @@ -286,6 +289,6 @@ mod tests { ) .unwrap(); - assert_arrays_eq!(decoded, array); + assert_arrays_eq!(decoded, array, &mut assertion_ctx); } } diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index a04894f2981..ed89b75f4ca 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -170,7 +170,7 @@ mod tests { } else { PrimitiveArray::from_option_iter(expected_shredded.iter().copied()) }; - assert_arrays_eq!(shredded, expected_shredded_array); + assert_arrays_eq!(shredded, expected_shredded_array, &mut ctx); Ok(()) } @@ -379,6 +379,8 @@ mod tests { #[test] fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([( "a", @@ -398,13 +400,16 @@ mod tests { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]) + PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), + &mut assertion_ctx ); Ok(()) } #[test] fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([ ( @@ -430,7 +435,8 @@ mod tests { .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; assert_arrays_eq!( value_result, - PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]) + PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), + &mut assertion_ctx ); let typed_value_expr = variant_get( @@ -444,7 +450,8 @@ mod tests { .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; assert_arrays_eq!( typed_value_result, - PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]) + PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]), + &mut assertion_ctx ); Ok(()) } diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index 221ceb75538..51fcfa2a288 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -221,6 +221,8 @@ mod tests { #[test] fn test_append_scalar() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10); // Test appending true value. @@ -237,7 +239,7 @@ mod tests { let array = builder.finish_into_bool(); let expected = BoolArray::from_iter([Some(true), Some(false), None]); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut assertion_ctx); // Test wrong dtype error. let mut builder = BoolBuilder::with_capacity(Nullability::NonNullable, 10); diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index f8d438d477b..cd9f6bbc93a 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -309,6 +309,7 @@ impl Default for DecimalBuffer { #[cfg(test)] mod tests { + use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::assert_arrays_eq; @@ -344,6 +345,8 @@ mod tests { #[test] fn test_append_scalar() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use crate::scalar::Scalar; // Simply test that the builder accepts its own finish output via scalar. @@ -357,7 +360,7 @@ mod tests { [Some(1234i64), Some(5678), None], DecimalDType::new(10, 2), ); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut assertion_ctx); // Test by taking a scalar from the array and appending it to a new builder. let mut builder2 = DecimalBuilder::new::(DecimalDType::new(10, 2), true.into()); @@ -369,7 +372,7 @@ mod tests { } let array2 = builder2.finish(); - assert_arrays_eq!(&array2, &array); + assert_arrays_eq!(&array2, &array, &mut assertion_ctx); // Test wrong dtype error. let mut builder = DecimalBuilder::new::(DecimalDType::new(10, 2), false.into()); diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index 8d7defd6a50..a97f0327742 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -173,18 +173,22 @@ mod test { #[test] fn encode_primitive() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = buffer![1, 1, 3, 3, 3].into_array(); let dict = dict_encode(&arr, &mut SESSION.create_execution_ctx()).unwrap(); let expected_codes = buffer![0u8, 0, 1, 1, 1].into_array(); - assert_arrays_eq!(dict.codes(), expected_codes); + assert_arrays_eq!(dict.codes(), expected_codes, &mut assertion_ctx); let expected_values = buffer![1i32, 3].into_array(); - assert_arrays_eq!(dict.values(), expected_values); + assert_arrays_eq!(dict.values(), expected_values, &mut assertion_ctx); } #[test] fn encode_primitive_nulls() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([ Some(1), Some(1), @@ -198,10 +202,10 @@ mod test { let dict = dict_encode(&arr.into_array(), &mut SESSION.create_execution_ctx()).unwrap(); let expected_codes = buffer![0u8, 0, 1, 2, 2, 1, 2, 1].into_array(); - assert_arrays_eq!(dict.codes(), expected_codes); + assert_arrays_eq!(dict.codes(), expected_codes, &mut assertion_ctx); let expected_values = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); - assert_arrays_eq!(dict.values(), expected_values); + assert_arrays_eq!(dict.values(), expected_values, &mut assertion_ctx); } } diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index 7e1da3b8f3e..ba844e3f178 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -127,6 +127,7 @@ impl ArrayBuilder for ExtensionBuilder { #[cfg(test)] mod tests { use super::*; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -137,6 +138,8 @@ mod tests { #[test] fn test_append_scalar() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); let mut builder = ExtensionBuilder::new(ext_dtype.clone()); @@ -165,7 +168,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(42i32), Some(84), None]).into_array(), ); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut assertion_ctx); assert_eq!(array.len(), 3); // Test wrong dtype error. diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index c23e35cccff..663f191e605 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -483,9 +483,9 @@ mod tests { let actual = builder.finish_into_canonical().into_listview(); - assert_arrays_eq!(actual.elements(), expected.elements()); + assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx); - assert_arrays_eq!(actual.offsets(), expected.offsets()); + assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx); assert!( actual @@ -681,13 +681,15 @@ mod tests { // Verify elements array: [1, 2, 3, 10, 11, 4, 5]. assert_arrays_eq!( list.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]) + PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]), + &mut ctx ); // Verify offsets array. assert_arrays_eq!( list.offsets(), - PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]) + PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]), + &mut ctx ); // Test dtype mismatch error. diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index fef03662181..49b4338c466 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -503,7 +503,8 @@ mod tests { // Check first list: [1, 2, 3]. assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); // Check empty list. @@ -521,12 +522,15 @@ mod tests { // Check last list: [4, 5]. assert_arrays_eq!( listview.list_elements_at(3).unwrap(), - PrimitiveArray::from_iter([4i32, 5]) + PrimitiveArray::from_iter([4i32, 5]), + &mut ctx ); } #[test] fn test_different_offset_size_types() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test u32 offsets with u8 sizes. let dtype: Arc = Arc::new(I32.into()); let mut builder = @@ -560,13 +564,15 @@ mod tests { // Verify first list: [1, 2]. assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([1i32, 2]) + PrimitiveArray::from_iter([1i32, 2]), + &mut assertion_ctx ); // Verify second list: [3, 4, 5]. assert_arrays_eq!( listview.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([3i32, 4, 5]) + PrimitiveArray::from_iter([3i32, 4, 5]), + &mut assertion_ctx ); // Test u64 offsets with u16 sizes. @@ -589,7 +595,8 @@ mod tests { for i in 0..5i32 { assert_arrays_eq!( listview2.list_elements_at(i as usize).unwrap(), - PrimitiveArray::from_iter([i * 10]) + PrimitiveArray::from_iter([i * 10]), + &mut assertion_ctx ); } } @@ -642,7 +649,8 @@ mod tests { // Last is the regular list: [10, 20]. assert_arrays_eq!( listview.list_elements_at(4).unwrap(), - PrimitiveArray::from_iter([10i32, 20]) + PrimitiveArray::from_iter([10i32, 20]), + &mut ctx ); } @@ -684,13 +692,15 @@ mod tests { // First list: [0] (initial data). assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([0i32]) + PrimitiveArray::from_iter([0i32]), + &mut ctx ); // Second list: [1, 2, 3] (from source). assert_arrays_eq!( listview.list_elements_at(1).unwrap(), - PrimitiveArray::from_iter([1i32, 2, 3]) + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut ctx ); // Third list: null (from source). @@ -705,7 +715,8 @@ mod tests { // Fourth list: [4, 5] (from source). assert_arrays_eq!( listview.list_elements_at(3).unwrap(), - PrimitiveArray::from_iter([4i32, 5]) + PrimitiveArray::from_iter([4i32, 5]), + &mut ctx ); } @@ -738,7 +749,8 @@ mod tests { assert_arrays_eq!( listview.list_elements_at(0).unwrap(), - PrimitiveArray::from_iter([10i32, 20]) + PrimitiveArray::from_iter([10i32, 20]), + &mut ctx ); assert!( !listview @@ -750,7 +762,8 @@ mod tests { assert_eq!(listview.list_elements_at(1).unwrap().len(), 0); assert_arrays_eq!( listview.list_elements_at(2).unwrap(), - PrimitiveArray::from_iter([10i32]) + PrimitiveArray::from_iter([10i32]), + &mut ctx ); } @@ -817,19 +830,22 @@ mod tests { // Verify elements array: [1, 2, 3, 10, 11, 4, 5]. assert_arrays_eq!( listview.elements(), - PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]) + PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]), + &mut ctx ); // Verify offsets array. assert_arrays_eq!( listview.offsets(), - PrimitiveArray::from_iter([0u32, 3, 5, 7, 7]) + PrimitiveArray::from_iter([0u32, 3, 5, 7, 7]), + &mut ctx ); // Verify sizes array. assert_arrays_eq!( listview.sizes(), - PrimitiveArray::from_iter([3u32, 2, 2, 0, 0]) + PrimitiveArray::from_iter([3u32, 2, 2, 0, 0]), + &mut ctx ); // Test dtype mismatch error. diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 870ad2275bd..79b96a7a762 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -376,6 +376,7 @@ mod tests { use vortex_error::VortexExpect; use super::*; + use crate::VortexSessionExecute; use crate::assert_arrays_eq; /// REGRESSION TEST: This test verifies that multiple sequential ranges have correct offsets. @@ -384,6 +385,8 @@ mod tests { /// buffer. #[test] fn test_multiple_uninit_ranges_correct_offsets() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); // First range. @@ -413,7 +416,11 @@ mod tests { assert_eq!(builder.values(), &[1, 2, 3, 4, 5]); let array = builder.finish_into_primitive(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]), + &mut assertion_ctx + ); } /// REGRESSION TEST: This test verifies that `append_mask` was correctly moved from @@ -490,6 +497,8 @@ mod tests { /// This verifies the new simplified API without the redundant `len` parameter. #[test] fn test_copy_from_slice_with_offsets() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); let mut range = builder.uninit_range(6); @@ -504,7 +513,11 @@ mod tests { } let array = builder.finish_into_primitive(); - assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6])); + assert_arrays_eq!( + array, + PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]), + &mut assertion_ctx + ); } /// Test that `set_bit` uses relative indexing within the range. diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index d396cf186a0..fb35bf2fdcf 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -210,6 +210,7 @@ impl ArrayBuilder for StructBuilder { #[cfg(test)] mod tests { + use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -266,6 +267,8 @@ mod tests { #[test] fn test_append_scalar() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use crate::scalar::Scalar; let dtype = DType::Struct( @@ -328,7 +331,7 @@ mod tests { Validity::from_iter([true, true, false]), ) .unwrap(); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut assertion_ctx); // Test wrong dtype error. let struct_fields = match &dtype { diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index c517972fb11..6c653b48c3b 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -884,6 +884,8 @@ mod tests { #[test] fn test_utf8_builder() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); builder.append_value("Hello"); @@ -906,11 +908,13 @@ mod tests { Some(""), Some("test"), ]); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut assertion_ctx); } #[test] fn test_utf8_builder_with_extend() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = { let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); @@ -934,7 +938,11 @@ mod tests { None, Some("Hello3"), ]); - assert_arrays_eq!(actual.into_array(), expected.into_array()); + assert_arrays_eq!( + actual.into_array(), + expected.into_array(), + &mut assertion_ctx + ); } #[test] @@ -986,6 +994,8 @@ mod tests { #[test] fn test_append_scalar() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use crate::scalar::Scalar; // Test with Utf8 builder. @@ -1007,7 +1017,7 @@ mod tests { let array = utf8_builder.finish(); let expected = >::from_iter([Some("hello"), Some("world"), None]); - assert_arrays_eq!(&array, &expected); + assert_arrays_eq!(&array, &expected, &mut assertion_ctx); // Test with Binary builder. let mut binary_builder = @@ -1022,7 +1032,7 @@ mod tests { let binary_array = binary_builder.finish(); let expected = >::from_iter([Some(vec![1u8, 2, 3]), None]); - assert_arrays_eq!(&binary_array, &expected); + assert_arrays_eq!(&binary_array, &expected, &mut assertion_ctx); // Test wrong dtype error. let mut builder = diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index 7d5f4abcad8..bcbae669675 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -62,12 +62,14 @@ pub fn create_runs_pattern(len: usize, run_length: usize) -> Vec { /// Tests that filtering with an all-true mask returns all elements unchanged fn test_all_filter(array: &ArrayRef) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let len = array.len(); let mask = Mask::new_true(len); let filtered = array .filter(mask) .vortex_expect("filter should succeed in conformance test"); - assert_arrays_eq!(filtered, array); + assert_arrays_eq!(filtered, array, &mut assertion_ctx); } /// Tests that filtering with an all-false mask returns an empty array with the same dtype diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 515e6f7398e..0c083f18ee5 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -858,6 +858,7 @@ mod tests { use super::*; use crate::VTable as _; + use crate::VortexSessionExecute; use crate::arrays::Bool; use crate::arrays::Primitive; use crate::optimizer::kernels::ExecuteParentFn; @@ -878,7 +879,7 @@ mod tests { let session = VortexSession::empty().with_some(KernelSession::empty()); let key = execute_parent_key(Bool.id(), Primitive.id()); - let before_registration = ExecutionCtx::new(session.clone()); + let before_registration = session.create_execution_ctx(); assert!( !before_registration .execute_parent_kernels @@ -897,7 +898,7 @@ mod tests { .contains_key(&key) ); - let after_registration = ExecutionCtx::new(session); + let after_registration = session.create_execution_ctx(); assert!(after_registration.execute_parent_kernels.contains_key(&key)); } } diff --git a/vortex-array/src/normalize.rs b/vortex-array/src/normalize.rs index 24a8f50d8a1..9de1603fd3f 100644 --- a/vortex-array/src/normalize.rs +++ b/vortex-array/src/normalize.rs @@ -104,8 +104,8 @@ mod tests { use super::NormalizeOptions; use super::Operation; use crate::ArrayRef; - use crate::ExecutionCtx; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::array::VTable; use crate::arrays::Dict; use crate::arrays::DictArray; @@ -128,7 +128,7 @@ mod tests { )? .into_array(); let allowed = HashSet::from_iter([array.encoding_id(), field.encoding_id()]); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = VortexSession::empty().create_execution_ctx(); let normalized = array.clone().normalize(&mut NormalizeOptions { allowed: &allowed, @@ -173,7 +173,7 @@ mod tests { )? .into_array(); let allowed = HashSet::from_iter([array.encoding_id(), unchanged.encoding_id()]); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = VortexSession::empty().create_execution_ctx(); let normalized = array.clone().normalize(&mut NormalizeOptions { allowed: &allowed, @@ -192,7 +192,11 @@ mod tests { &original_children[1], &normalized_children[1] )); - assert_arrays_eq!(normalized_children[1], PrimitiveArray::from_iter(12i32..16)); + assert_arrays_eq!( + normalized_children[1], + PrimitiveArray::from_iter(12i32..16), + &mut ctx + ); Ok(()) } @@ -208,7 +212,7 @@ mod tests { assert_eq!(sliced.encoding_id(), Slice.id()); let allowed = HashSet::from_iter([Dict.id(), Primitive.id()]); - let mut ctx = ExecutionCtx::new(VortexSession::empty()); + let mut ctx = VortexSession::empty().create_execution_ctx(); let normalized = sliced.normalize(&mut NormalizeOptions { allowed: &allowed, @@ -224,7 +228,8 @@ mod tests { let normalized_canonical = normalized.to_canonical()?; assert_arrays_eq!( normalized_canonical, - PrimitiveArray::from_iter(vec![20i32, 10, 20]) + PrimitiveArray::from_iter(vec![20i32, 10, 20]), + &mut ctx ); Ok(()) diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 76f1c8ad99c..558b3329ed8 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -1222,6 +1222,8 @@ mod test { #[test] fn test_filter() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 100, 0, @@ -1239,12 +1241,22 @@ mod test { .unwrap() .unwrap(); - assert_arrays_eq!(filtered.indices(), PrimitiveArray::from_iter([0u64, 1])); - assert_arrays_eq!(filtered.values(), PrimitiveArray::from_iter([100i32, 200])); + assert_arrays_eq!( + filtered.indices(), + PrimitiveArray::from_iter([0u64, 1]), + &mut assertion_ctx + ); + assert_arrays_eq!( + filtered.values(), + PrimitiveArray::from_iter([100i32, 200]), + &mut assertion_ctx + ); } #[test] fn take_with_nulls() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1269,9 +1281,14 @@ mod test { assert_eq!(taken.array_len(), 2); assert_arrays_eq!( primitive_values, - PrimitiveArray::from_option_iter([Some(44i32)]) + PrimitiveArray::from_option_iter([Some(44i32)]), + &mut assertion_ctx + ); + assert_arrays_eq!( + primitive_indices, + PrimitiveArray::from_iter([0u64]), + &mut assertion_ctx ); - assert_arrays_eq!(primitive_indices, PrimitiveArray::from_iter([0u64])); assert_eq!( primitive_values .as_ref() @@ -1288,6 +1305,8 @@ mod test { #[test] fn take_search_with_nulls_chunked() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1311,9 +1330,14 @@ mod test { assert_eq!(taken.array_len(), 2); assert_arrays_eq!( primitive_values, - PrimitiveArray::from_option_iter([Some(44i32), None]) + PrimitiveArray::from_option_iter([Some(44i32), None]), + &mut assertion_ctx + ); + assert_arrays_eq!( + taken.indices(), + PrimitiveArray::from_iter([0u64, 1]), + &mut assertion_ctx ); - assert_arrays_eq!(taken.indices(), PrimitiveArray::from_iter([0u64, 1])); assert_eq!( primitive_values @@ -1331,6 +1355,8 @@ mod test { #[test] fn take_search_chunked_multiple_chunks() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 2048, 0, @@ -1352,7 +1378,8 @@ mod test { assert_eq!(taken.array_len(), 3); assert_arrays_eq!( taken.values(), - PrimitiveArray::from_option_iter([Some(20i32), Some(30)]) + PrimitiveArray::from_option_iter([Some(20i32), Some(30)]), + &mut assertion_ctx ); } @@ -1380,6 +1407,8 @@ mod test { #[test] fn take_search_chunked_interleaved() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 30, 0, @@ -1401,7 +1430,8 @@ mod test { assert_eq!(taken.array_len(), 4); assert_arrays_eq!( taken.values(), - PrimitiveArray::from_option_iter([Some(200i32), Some(300)]) + PrimitiveArray::from_option_iter([Some(200i32), Some(300)]), + &mut assertion_ctx ); } @@ -1430,6 +1460,8 @@ mod test { #[test] fn test_slice() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1437,11 +1469,17 @@ mod test { let sliced = patches.slice(15..100).unwrap().unwrap(); assert_eq!(sliced.array_len(), 100 - 15); - assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([13531u32])); + assert_arrays_eq!( + sliced.values(), + PrimitiveArray::from_iter([13531u32]), + &mut assertion_ctx + ); } #[test] fn doubly_sliced() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1449,12 +1487,17 @@ mod test { let sliced = patches.slice(15..100).unwrap().unwrap(); assert_eq!(sliced.array_len(), 100 - 15); - assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([13531u32])); + assert_arrays_eq!( + sliced.values(), + PrimitiveArray::from_iter([13531u32]), + &mut assertion_ctx + ); let doubly_sliced = sliced.slice(35..36).unwrap().unwrap(); assert_arrays_eq!( doubly_sliced.values(), - PrimitiveArray::from_iter([13531u32]) + PrimitiveArray::from_iter([13531u32]), + &mut assertion_ctx ); } @@ -1478,6 +1521,8 @@ mod test { #[test] fn test_mask_all_false() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1496,7 +1541,8 @@ mod test { // No patch values should be masked assert_arrays_eq!( masked.values(), - PrimitiveArray::from_iter([100i32, 200, 300]) + PrimitiveArray::from_iter([100i32, 200, 300]), + &mut assertion_ctx ); assert!( masked @@ -1518,11 +1564,17 @@ mod test { ); // Indices should remain unchanged - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([2u64, 5, 8]), + &mut assertion_ctx + ); } #[test] fn test_mask_partial() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1543,14 +1595,24 @@ mod test { // Only the patch at index 5 should remain assert_eq!(masked.values().len(), 1); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([200i32])); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_iter([200i32]), + &mut assertion_ctx + ); // Only index 5 should remain - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([5u64])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([5u64]), + &mut assertion_ctx + ); } #[test] fn test_mask_with_offset() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1571,12 +1633,22 @@ mod test { .unwrap(); assert_eq!(masked.array_len(), 10); assert_eq!(masked.offset(), 5); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([10u64, 13])); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([200i32, 300])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([10u64, 13]), + &mut assertion_ctx + ); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_iter([200i32, 300]), + &mut assertion_ctx + ); } #[test] fn test_mask_nullable_values() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1596,7 +1668,11 @@ mod test { .unwrap(); // Patches at indices 5 and 8 should remain - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([5u64, 8])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([5u64, 8]), + &mut assertion_ctx + ); // Values should be the null and 300 #[expect(deprecated)] @@ -1625,6 +1701,8 @@ mod test { #[test] fn test_filter_keep_all() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1641,10 +1719,15 @@ mod test { .unwrap() .unwrap(); - assert_arrays_eq!(filtered.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); + assert_arrays_eq!( + filtered.indices(), + PrimitiveArray::from_iter([2u64, 5, 8]), + &mut assertion_ctx + ); assert_arrays_eq!( filtered.values(), - PrimitiveArray::from_iter([100i32, 200, 300]) + PrimitiveArray::from_iter([100i32, 200, 300]), + &mut assertion_ctx ); } @@ -1669,6 +1752,8 @@ mod test { #[test] fn test_filter_with_indices() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1685,12 +1770,22 @@ mod test { .unwrap() .unwrap(); - assert_arrays_eq!(filtered.indices(), PrimitiveArray::from_iter([0u64, 1])); // Adjusted indices - assert_arrays_eq!(filtered.values(), PrimitiveArray::from_iter([100i32, 200])); + assert_arrays_eq!( + filtered.indices(), + PrimitiveArray::from_iter([0u64, 1]), + &mut assertion_ctx + ); // Adjusted indices + assert_arrays_eq!( + filtered.values(), + PrimitiveArray::from_iter([100i32, 200]), + &mut assertion_ctx + ); } #[test] fn test_slice_full_range() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1702,15 +1797,22 @@ mod test { let sliced = patches.slice(0..10).unwrap().unwrap(); - assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); + assert_arrays_eq!( + sliced.indices(), + PrimitiveArray::from_iter([2u64, 5, 8]), + &mut assertion_ctx + ); assert_arrays_eq!( sliced.values(), - PrimitiveArray::from_iter([100i32, 200, 300]) + PrimitiveArray::from_iter([100i32, 200, 300]), + &mut assertion_ctx ); } #[test] fn test_slice_partial() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1723,8 +1825,16 @@ mod test { // Slice from 3 to 8 (includes patch at 5) let sliced = patches.slice(3..8).unwrap().unwrap(); - assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([5u64])); // Index stays the same - assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([200i32])); + assert_arrays_eq!( + sliced.indices(), + PrimitiveArray::from_iter([5u64]), + &mut assertion_ctx + ); // Index stays the same + assert_arrays_eq!( + sliced.values(), + PrimitiveArray::from_iter([200i32]), + &mut assertion_ctx + ); assert_eq!(sliced.array_len(), 5); // 8 - 3 = 5 assert_eq!(sliced.offset(), 3); // New offset } @@ -1747,6 +1857,8 @@ mod test { #[test] fn test_slice_with_offset() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1759,8 +1871,16 @@ mod test { // Slice from 3 to 8 (includes patch at actual index 5) let sliced = patches.slice(3..8).unwrap().unwrap(); - assert_arrays_eq!(sliced.indices(), PrimitiveArray::from_iter([10u64])); // Index stays the same (offset + 5 = 10) - assert_arrays_eq!(sliced.values(), PrimitiveArray::from_iter([200i32])); + assert_arrays_eq!( + sliced.indices(), + PrimitiveArray::from_iter([10u64]), + &mut assertion_ctx + ); // Index stays the same (offset + 5 = 10) + assert_arrays_eq!( + sliced.values(), + PrimitiveArray::from_iter([200i32]), + &mut assertion_ctx + ); assert_eq!(sliced.offset(), 8); // New offset = 5 + 3 } @@ -1846,6 +1966,8 @@ mod test { #[test] fn test_mask_boundary_patches() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test masking patches at array boundaries let patches = Patches::new( 10, @@ -1864,8 +1986,16 @@ mod test { .unwrap(); assert!(masked.is_some()); let masked = masked.unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([9u64])); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([200i32])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([9u64]), + &mut assertion_ctx + ); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_iter([200i32]), + &mut assertion_ctx + ); } #[test] @@ -1892,6 +2022,8 @@ mod test { #[test] fn test_mask_no_patches_removed() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test when no patches are masked let patches = Patches::new( 10, @@ -1911,15 +2043,22 @@ mod test { .unwrap() .unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([2u64, 5, 8])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([2u64, 5, 8]), + &mut assertion_ctx + ); assert_arrays_eq!( masked.values(), - PrimitiveArray::from_iter([100i32, 200, 300]) + PrimitiveArray::from_iter([100i32, 200, 300]), + &mut assertion_ctx ); } #[test] fn test_mask_single_patch() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test with a single patch let patches = Patches::new( 5, @@ -1943,11 +2082,17 @@ mod test { .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) .unwrap() .unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([2u64])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([2u64]), + &mut assertion_ctx + ); } #[test] fn test_mask_contiguous_patches() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test with contiguous patches let patches = Patches::new( 10, @@ -1967,12 +2112,22 @@ mod test { .unwrap() .unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([3u64, 6])); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([100i32, 400])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([3u64, 6]), + &mut assertion_ctx + ); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_iter([100i32, 400]), + &mut assertion_ctx + ); } #[test] fn test_mask_with_large_offset() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test with a large offset that shifts all indices let patches = Patches::new( 20, @@ -1993,8 +2148,16 @@ mod test { .unwrap() .unwrap(); - assert_arrays_eq!(masked.indices(), PrimitiveArray::from_iter([16u64, 19])); - assert_arrays_eq!(masked.values(), PrimitiveArray::from_iter([100i32, 300])); + assert_arrays_eq!( + masked.indices(), + PrimitiveArray::from_iter([16u64, 19]), + &mut assertion_ctx + ); + assert_arrays_eq!( + masked.values(), + PrimitiveArray::from_iter([100i32, 300]), + &mut assertion_ctx + ); } #[test] diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index 2cbe02183c8..bdb3f1a5276 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -469,6 +469,8 @@ mod tests { #[test] fn test_between_decimal() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![100i128, 200i128, 300i128, 400i128]; let decimal_type = DecimalDType::new(3, 2); let array = DecimalArray::new(values, decimal_type, Validity::NonNullable).into_array(); @@ -506,7 +508,8 @@ mod tests { .unwrap(); assert_arrays_eq!( between_strict, - BoolArray::from_iter([false, true, true, true]) + BoolArray::from_iter([false, true, true, true]), + &mut assertion_ctx ); // Non-strict lower bound, strict upper bound @@ -523,7 +526,8 @@ mod tests { .unwrap(); assert_arrays_eq!( between_strict, - BoolArray::from_iter([true, true, true, false]) + BoolArray::from_iter([true, true, true, false]), + &mut assertion_ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index 40e86deb839..a5ad61bb973 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -772,6 +772,8 @@ mod tests { #[test] fn test_kleene_truth_table() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let lhs = BoolArray::from_iter([ Some(true), Some(true), @@ -809,7 +811,8 @@ mod tests { None, Some(false), None, - ]) + ]), + &mut assertion_ctx ); assert_arrays_eq!( @@ -824,7 +827,8 @@ mod tests { Some(true), None, None, - ]) + ]), + &mut assertion_ctx ); Ok(()) @@ -832,17 +836,21 @@ mod tests { #[test] fn test_null_constant_kleene() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let lhs = BoolArray::from_iter([Some(false), Some(true), None]).into_array(); let null = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len()) .into_array(); assert_arrays_eq!( lhs.binary(null.clone(), Operator::And)?, - BoolArray::from_iter([Some(false), None, None]) + BoolArray::from_iter([Some(false), None, None]), + &mut assertion_ctx ); assert_arrays_eq!( lhs.binary(null, Operator::Or)?, - BoolArray::from_iter([None, Some(true), None]) + BoolArray::from_iter([None, Some(true), None]), + &mut assertion_ctx ); Ok(()) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare.rs b/vortex-array/src/scalar_fn/fns/binary/compare.rs index 5d4509fe406..02dc2d2ac19 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare.rs @@ -369,13 +369,17 @@ mod tests { #[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())] #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())] fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let res = left.binary(right, Operator::Eq).unwrap(); let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(res, expected); + assert_arrays_eq!(res, expected, &mut assertion_ctx); } #[test] fn test_list_array_comparison() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values1 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets1 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list1 = ListArray::try_new( @@ -400,7 +404,7 @@ mod tests { .binary(list2.clone().into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); let result = list1 .clone() @@ -408,18 +412,20 @@ mod tests { .binary(list2.clone().into_array(), Operator::NotEq) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); let result = list1 .into_array() .binary(list2.into_array(), Operator::Lt) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[test] fn test_list_array_constant_comparison() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list = ListArray::try_new( @@ -441,11 +447,13 @@ mod tests { .binary(constant.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([false, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[test] fn test_struct_array_comparison() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bool_field1 = BoolArray::from_iter([Some(true), Some(false), Some(true)]); let int_field1 = PrimitiveArray::from_iter([1i32, 2, 3]); @@ -470,18 +478,20 @@ mod tests { .binary(struct2.clone().into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); let result = struct1 .into_array() .binary(struct2.into_array(), Operator::Gt) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[test] fn test_empty_struct_compare() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let empty1 = StructArray::try_new( FieldNames::from(Vec::::new()), Vec::new(), @@ -503,13 +513,15 @@ mod tests { .binary(empty2.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, true, true, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } /// Regression test: comparing struct arrays where the same logical field is backed by /// different Vortex encodings (VarBinArray vs VarBinViewArray) must not panic. #[test] fn struct_compare_mixed_binary_encodings() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // LHS: struct with a VarBinArray (offset-based) binary field let bin_field1 = VarBinArray::from(vec![ "apple".as_bytes(), @@ -531,7 +543,7 @@ mod tests { .binary(struct2.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } /// Regression test: `scalar_cmp` must error when comparing scalars with incompatible diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index bda7e61ffd7..09938e13cf1 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -519,6 +519,8 @@ mod tests { #[test] fn test_or_kleene_validity() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::StructArray; @@ -537,11 +539,17 @@ mod tests { let expr = or(col("a"), col("b")); let result = struct_arr.apply(&expr).unwrap(); - assert_arrays_eq!(result, BoolArray::from_iter([Some(true)]).into_array()) + assert_arrays_eq!( + result, + BoolArray::from_iter([Some(true)]).into_array(), + &mut assertion_ctx + ) } #[test] fn test_scalar_subtract_unsigned() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -551,11 +559,17 @@ mod tests { let values = buffer![1u16, 2, 3].into_array(); let rhs = ConstantArray::new(Scalar::from(1u16), 3).into_array(); let result = values.binary(rhs, Operator::Sub).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([0u16, 1, 2]), + &mut assertion_ctx + ); } #[test] fn test_scalar_subtract_signed() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -565,11 +579,17 @@ mod tests { let values = buffer![1i64, 2, 3].into_array(); let rhs = ConstantArray::new(Scalar::from(-1i64), 3).into_array(); let result = values.binary(rhs, Operator::Sub).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([2i64, 3, 4]), + &mut assertion_ctx + ); } #[test] fn test_scalar_subtract_nullable() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -579,12 +599,15 @@ mod tests { let result = values.into_array().binary(rhs, Operator::Sub).unwrap(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]) + PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]), + &mut assertion_ctx ); } #[test] fn test_scalar_subtract_float() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -594,7 +617,11 @@ mod tests { let values = buffer![1.0f64, 2.0, 3.0].into_array(); let rhs = ConstantArray::new(Scalar::from(-1f64), 3).into_array(); let result = values.binary(rhs, Operator::Sub).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2.0f64, 3.0, 4.0])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]), + &mut assertion_ctx + ); } #[test] diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index cb85b727503..17af9dd8289 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -951,33 +951,54 @@ mod test { #[test] fn test_scalar_subtract_unsigned() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![1u16, 2, 3].into_array(); let result = sub_scalar(&values, 1u16).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([0u16, 1, 2]), + &mut assertion_ctx + ); } #[test] fn test_scalar_subtract_signed() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![1i64, 2, 3].into_array(); let result = sub_scalar(&values, -1i64).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([2i64, 3, 4]), + &mut assertion_ctx + ); } #[test] fn test_scalar_subtract_nullable() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); let result = sub_scalar(&values.into_array(), Some(1u16)).unwrap(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]) + PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]), + &mut assertion_ctx ); } #[test] fn test_scalar_subtract_float() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![1.0f64, 2.0, 3.0].into_array(); let result = sub_scalar(&values, -1f64).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([2.0f64, 3.0, 4.0])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]), + &mut assertion_ctx + ); } #[test] @@ -989,6 +1010,8 @@ mod test { #[test] fn test_float_divide_by_zero_is_ok() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![1.0f64, -1.0].into_array(); let result = values .binary( @@ -1000,7 +1023,8 @@ mod test { assert_arrays_eq!( result, - PrimitiveArray::from_iter([f64::INFINITY, f64::NEG_INFINITY]) + PrimitiveArray::from_iter([f64::INFINITY, f64::NEG_INFINITY]), + &mut assertion_ctx ); } @@ -1045,6 +1069,8 @@ mod test { #[test] fn test_integer_divide_errors_ignore_null_lanes() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![10i32, 10], Validity::from_iter([false, true])) .into_array(); let rhs = buffer![0i32, 2].into_array(); @@ -1056,11 +1082,17 @@ mod test { .map(|a| a.0.into_array()) .unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_option_iter([None, Some(5i32)])); + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([None, Some(5i32)]), + &mut assertion_ctx + ); } #[test] fn test_integer_errors_ignore_null_lanes() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true])) .into_array(); let result = values @@ -1074,7 +1106,11 @@ mod test { .map(|a| a.0.into_array()) .unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_option_iter([None, Some(2u8)])); + assert_arrays_eq!( + result, + PrimitiveArray::from_option_iter([None, Some(2u8)]), + &mut assertion_ctx + ); } #[test] @@ -1094,6 +1130,8 @@ mod test { #[test] fn test_present_nullable_constant_preserves_nullable_output() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = buffer![1u8, 2].into_array(); let result = values .binary( @@ -1105,7 +1143,8 @@ mod test { assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(2u8), Some(3)]) + PrimitiveArray::from_option_iter([Some(2u8), Some(3)]), + &mut assertion_ctx ); } } diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-array/src/scalar_fn/fns/byte_length.rs index aa9c508ea89..9e0f491bf21 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-array/src/scalar_fn/fns/byte_length.rs @@ -202,18 +202,22 @@ mod tests { #[case] array: ArrayRef, #[case] expected_lens: Vec, ) -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(expected_lens); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn test_varbinview_byte_length() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["short", "a longer string here"]).into_array(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 20]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index 0a5f6b4ee14..2e41daa1069 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -783,6 +783,8 @@ mod tests { #[test] fn test_evaluate_simple_condition() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -795,11 +797,17 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 100, 100, 100].into_array()); + assert_arrays_eq!( + result, + buffer![0i32, 0, 100, 100, 100].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_nary_multiple_conditions() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Test n-ary via nested_case_when let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) @@ -815,11 +823,17 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![10i32, 0, 30, 0, 0].into_array()); + assert_arrays_eq!( + result, + buffer![10i32, 0, 30, 0, 0].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_nary_first_match_wins() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -835,11 +849,17 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 100, 100, 100].into_array()); + assert_arrays_eq!( + result, + buffer![0i32, 0, 100, 100, 100].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_no_else_returns_null() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -852,12 +872,15 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([None::, None, None, Some(100), Some(100)]) - .into_array() + .into_array(), + &mut assertion_ctx ); } #[test] fn test_evaluate_all_conditions_false() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -870,11 +893,17 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 0, 0, 0].into_array()); + assert_arrays_eq!( + result, + buffer![0i32, 0, 0, 0, 0].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_all_conditions_true() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -887,11 +916,17 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![100i32, 100, 100, 100, 100].into_array()); + assert_arrays_eq!( + result, + buffer![100i32, 100, 100, 100, 100].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_all_true_no_else_returns_correct_dtype() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // CASE WHEN value > 0 THEN 100 END — condition is always true, no ELSE. // Result must be Nullable because the implicit ELSE is NULL. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -908,12 +943,15 @@ mod tests { ); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(100i32), Some(100), Some(100)]).into_array() + PrimitiveArray::from_option_iter([Some(100i32), Some(100), Some(100)]).into_array(), + &mut assertion_ctx ); } #[test] fn test_merge_case_branches_widens_nullability_of_later_branch() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // When a later THEN branch is Nullable and branches[0] and ELSE are NonNullable, // the result dtype must still be Nullable. // @@ -943,22 +981,31 @@ mod tests { ); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10), Some(20), Some(0)]).into_array() + PrimitiveArray::from_option_iter([Some(10), Some(20), Some(0)]).into_array(), + &mut assertion_ctx ); Ok(()) } #[test] fn test_evaluate_with_literal_condition() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = case_when(lit(true), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![100i32, 100, 100].into_array()); + assert_arrays_eq!( + result, + buffer![100i32, 100, 100].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_with_bool_column_result() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -973,12 +1020,15 @@ mod tests { let result = evaluate_expr(&expr, &test_array); assert_arrays_eq!( result, - BoolArray::from_iter([false, false, true, true, true]).into_array() + BoolArray::from_iter([false, false, true, true, true]).into_array(), + &mut assertion_ctx ); } #[test] fn test_evaluate_with_nullable_condition() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]).into_array(), @@ -989,11 +1039,17 @@ mod tests { let expr = case_when(get_item("cond", root()), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![100i32, 0, 0, 0, 100].into_array()); + assert_arrays_eq!( + result, + buffer![100i32, 0, 0, 0, 100].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_with_nullable_result_values() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[ ("value", buffer![1i32, 2, 3, 4, 5].into_array()), ( @@ -1015,12 +1071,15 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(0i32), Some(0), Some(30), Some(40), Some(50)]) - .into_array() + .into_array(), + &mut assertion_ctx ); } #[test] fn test_evaluate_with_all_null_condition() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([None, None, None]).into_array(), @@ -1031,13 +1090,15 @@ mod tests { let expr = case_when(get_item("cond", root()), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 0].into_array()); + assert_arrays_eq!(result, buffer![0i32, 0, 0].into_array(), &mut assertion_ctx); } // ==================== N-ary Evaluate Tests ==================== #[test] fn test_evaluate_nary_no_else_returns_null() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1057,12 +1118,15 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None, None]) - .into_array() + .into_array(), + &mut assertion_ctx ); } #[test] fn test_evaluate_nary_many_conditions() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1081,11 +1145,17 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![10i32, 20, 30, 40, 50].into_array()); + assert_arrays_eq!( + result, + buffer![10i32, 20, 30, 40, 50].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_nary_all_false_no_else() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) .unwrap() .into_array(); @@ -1103,12 +1173,15 @@ mod tests { assert!(result.dtype().is_nullable()); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([None::, None, None]).into_array() + PrimitiveArray::from_option_iter([None::, None, None]).into_array(), + &mut assertion_ctx ); } #[test] fn test_evaluate_nary_overlapping_conditions_first_wins() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![10i32, 20, 30].into_array())]) .unwrap() @@ -1128,11 +1201,13 @@ mod tests { let result = evaluate_expr(&expr, &test_array); // First matching condition always wins - assert_arrays_eq!(result, buffer![1i32, 1, 1].into_array()); + assert_arrays_eq!(result, buffer![1i32, 1, 1].into_array(), &mut assertion_ctx); } #[test] fn test_evaluate_nary_early_exit_when_remaining_empty() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // After branch 0 claims all rows, remaining becomes all_false. // The loop breaks before evaluating branch 1's condition. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1149,11 +1224,17 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![100i32, 100, 100].into_array()); + assert_arrays_eq!( + result, + buffer![100i32, 100, 100].into_array(), + &mut assertion_ctx + ); } #[test] fn test_evaluate_nary_skips_branch_with_empty_effective_mask() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Branch 0 claims value=1. Branch 1 targets the same rows but they are already // matched → effective_mask is all_false → branch 1 is skipped (THEN not used). let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1172,7 +1253,11 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![10i32, 20, 0].into_array()); + assert_arrays_eq!( + result, + buffer![10i32, 20, 0].into_array(), + &mut assertion_ctx + ); } #[test] @@ -1215,6 +1300,8 @@ mod tests { #[test] fn test_evaluate_nary_with_nullable_conditions() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[ ( "cond1", @@ -1240,7 +1327,11 @@ mod tests { // row 0: cond1=true → 10 // row 1: cond1=NULL(→false), cond2=true → 20 // row 2: cond1=false, cond2=NULL(→false) → else=0 - assert_arrays_eq!(result, buffer![10i32, 20, 0].into_array()); + assert_arrays_eq!( + result, + buffer![10i32, 20, 0].into_array(), + &mut assertion_ctx + ); } // ==================== Simplify: COALESCE -> fill_null ==================== @@ -1331,6 +1422,8 @@ mod tests { #[test] fn test_simplify_null_fill_semantic_equivalence() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // The collapse-to-input rewrite must preserve values (and `x`'s nullability). let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1348,8 +1441,16 @@ mod tests { ); let expected = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); - assert_arrays_eq!(evaluate_expr(&original, &array), expected); - assert_arrays_eq!(evaluate_expr(&optimized, &array), expected); + assert_arrays_eq!( + evaluate_expr(&original, &array), + expected, + &mut assertion_ctx + ); + assert_arrays_eq!( + evaluate_expr(&optimized, &array), + expected, + &mut assertion_ctx + ); Ok(()) } @@ -1383,6 +1484,8 @@ mod tests { #[test] fn test_simplify_semantic_equivalence() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // The optimized expression must produce the same values as the original CASE WHEN. let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1398,17 +1501,21 @@ mod tests { // NonNullable because a non-null fill cannot leave any nulls behind. Values match. assert_arrays_eq!( evaluate_expr(&original, &array), - PrimitiveArray::from_option_iter([Some(1i64), Some(0), Some(3)]).into_array() + PrimitiveArray::from_option_iter([Some(1i64), Some(0), Some(3)]).into_array(), + &mut assertion_ctx ); assert_arrays_eq!( evaluate_expr(&optimized, &array), - buffer![1i64, 0, 3].into_array() + buffer![1i64, 0, 3].into_array(), + &mut assertion_ctx ); Ok(()) } #[test] fn test_merge_case_branches_alternating_mask() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Exercises the scalar path: alternating rows produce one slice per row (no runs), // triggering the per-row cursor path in merge_case_branches. let n = 100usize; @@ -1438,7 +1545,8 @@ mod tests { .collect(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter(expected).into_array() + PrimitiveArray::from_option_iter(expected).into_array(), + &mut assertion_ctx ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index f6e6619282a..9ea3a17de51 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -277,6 +277,7 @@ mod tests { use super::*; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -303,6 +304,8 @@ mod tests { #[test] fn execute_with_value() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -312,12 +315,18 @@ mod tests { root(), ); let result = input.apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([true, false, false])); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, false]), + &mut assertion_ctx + ); Ok(()) } #[test] fn execute_without_value_default_true() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -327,12 +336,18 @@ mod tests { root(), ); let result = input.apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([true, true, true])); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, true, true]), + &mut assertion_ctx + ); Ok(()) } #[test] fn execute_without_value_default_false() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -342,12 +357,18 @@ mod tests { root(), ); let result = input.apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([false, false, false])); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, false]), + &mut assertion_ctx + ); Ok(()) } #[test] fn execute_value_flips() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let threshold = Arc::new(AtomicI32::new(5)); let threshold_clone = Arc::clone(&threshold); let expr = dynamic( @@ -360,11 +381,19 @@ mod tests { let input = buffer![1i32, 5, 10].into_array(); let result = input.clone().apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([true, false, false])); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, false]), + &mut assertion_ctx + ); threshold.store(10, Ordering::SeqCst); let result = input.apply(&expr)?; - assert_arrays_eq!(result, BoolArray::from_iter([true, true, false])); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, true, false]), + &mut assertion_ctx + ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index b2d07b41255..9069d82cab4 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -179,6 +179,7 @@ mod tests { use vortex_error::VortexExpect; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::assert_arrays_eq; @@ -208,6 +209,8 @@ mod tests { #[test] fn evaluate() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]) .into_array(); @@ -219,11 +222,17 @@ mod tests { result.dtype(), &DType::Primitive(PType::I32, Nullability::NonNullable) ); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 42, 3, 42, 5])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]), + &mut assertion_ctx + ); } #[test] fn evaluate_struct_field() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = StructArray::from_fields(&[( "a", PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), @@ -238,15 +247,25 @@ mod tests { result.dtype(), &DType::Primitive(PType::I32, Nullability::NonNullable) ); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 0, 3])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 0, 3]), + &mut assertion_ctx + ); } #[test] fn evaluate_non_nullable_input() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = fill_null(root(), lit(0i32)); let result = test_array.apply(&expr).unwrap(); - assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 2, 3])); + assert_arrays_eq!( + result, + PrimitiveArray::from_iter([1i32, 2, 3]), + &mut assertion_ctx + ); } #[test] diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index f1c53b146cb..a2da7735c6b 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -244,6 +244,8 @@ mod tests { use std::borrow::Cow; use crate::IntoArray; + use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -260,9 +262,11 @@ mod tests { fn invert_booleans() { let not_expr = not(root()); let bools = BoolArray::from_iter([false, true, false, false, true, true]); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); assert_arrays_eq!( bools.into_array().apply(¬_expr).unwrap(), - BoolArray::from_iter([true, false, true, true, false, false]) + BoolArray::from_iter([true, false, true, true, false, false]), + &mut ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 4b39f51a7f2..c2afb16a458 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -752,6 +752,8 @@ mod tests { #[case] value: Option<&str>, #[case] expected: BoolArray, ) { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let element_nullability = list_array .dtype() .as_list_element_opt() @@ -764,11 +766,13 @@ mod tests { let elem = ConstantArray::new(scalar, list_array.len()); let expr = list_contains(root(), lit(elem.scalar().clone())); let result = list_array.apply(&expr).unwrap(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[test] fn test_constant_list() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let list_array = ConstantArray::new( Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -782,11 +786,13 @@ mod tests { let expr = list_contains(root(), lit(2i32)); let contains = list_array.apply(&expr).unwrap(); let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(contains, expected); + assert_arrays_eq!(contains, expected, &mut assertion_ctx); } #[test] fn test_all_nulls() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let list_array = ConstantArray::new( Scalar::null(DType::List( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -803,11 +809,13 @@ mod tests { [false, false, false, false, false].into_iter().collect(), Validity::AllInvalid, ); - assert_arrays_eq!(contains, expected); + assert_arrays_eq!(contains, expected, &mut assertion_ctx); } #[test] fn test_list_array_element() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let list_scalar = Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), vec![1.into(), 3.into(), 6.into()], @@ -819,11 +827,13 @@ mod tests { let contains = arr.apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, true, false, true, false, false, true]); - assert_arrays_eq!(contains, expected); + assert_arrays_eq!(contains, expected, &mut assertion_ctx); } #[test] fn test_list_contains_empty_listview() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let empty_elements = PrimitiveArray::empty::(Nullability::NonNullable); let offsets = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); let sizes = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); @@ -842,11 +852,13 @@ mod tests { let result = list_array.into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, false, false, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[test] fn test_list_contains_all_null_elements() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::from_option_iter::([None, None, None, None, None]); let offsets = Buffer::from_iter([0u32, 2, 4]).into_array(); let sizes = Buffer::from_iter([2u32, 2, 1]).into_array(); @@ -870,18 +882,20 @@ mod tests { [false, false, false].into_iter().collect(), Validity::AllInvalid, ); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); // Searching for non-null let expr2 = list_contains(root(), lit(42i32)); let result2 = list_array.into_array().apply(&expr2).unwrap(); let expected2 = BoolArray::from_iter([false, false, false]); - assert_arrays_eq!(result2, expected2); + assert_arrays_eq!(result2, expected2, &mut assertion_ctx); } #[test] fn test_list_contains_large_offsets() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = Buffer::from_iter([1i32, 2, 3, 4, 5]).into_array(); let offsets = Buffer::from_iter([0u32, 1, 4, 0]).into_array(); @@ -894,17 +908,19 @@ mod tests { let result = list_array.clone().into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, true, false, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); let expr5 = list_contains(root(), lit(5i32)); let result5 = list_array.into_array().apply(&expr5).unwrap(); let expected5 = BoolArray::from_iter([false, false, true, false]); - assert_arrays_eq!(result5, expected5); + assert_arrays_eq!(result5, expected5, &mut assertion_ctx); } #[test] fn test_list_contains_offset_size_boundary() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = Buffer::from_iter(0..256).into_array(); let offsets = Buffer::from_iter([0u8, 100, 200, 254]).into_array(); let sizes = Buffer::from_iter([50u8, 50, 54, 2]).into_array(); @@ -916,12 +932,12 @@ mod tests { let result = list_array.clone().into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, false, false, true]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); let expr_zero = list_contains(root(), lit(0i32)); let result_zero = list_array.into_array().apply(&expr_zero).unwrap(); let expected_zero = BoolArray::from_iter([true, false, false, false]); - assert_arrays_eq!(result_zero, expected_zero); + assert_arrays_eq!(result_zero, expected_zero, &mut assertion_ctx); } } diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 390d3a2213e..d152f55ad3e 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -279,6 +279,7 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -318,6 +319,8 @@ mod tests { #[test] pub fn test_merge_right_most() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let expr = merge_opts( vec![ get_item("0", root()), @@ -367,23 +370,28 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array, &["a"]).unwrap(), - PrimitiveArray::from_iter([0i32, 0, 0]) + PrimitiveArray::from_iter([0i32, 0, 0]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["b"]).unwrap(), - PrimitiveArray::from_iter([2i32, 2, 2]) + PrimitiveArray::from_iter([2i32, 2, 2]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["c"]).unwrap(), - PrimitiveArray::from_iter([3i32, 3, 3]) + PrimitiveArray::from_iter([3i32, 3, 3]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["d"]).unwrap(), - PrimitiveArray::from_iter([4i32, 4, 4]) + PrimitiveArray::from_iter([4i32, 4, 4]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["e"]).unwrap(), - PrimitiveArray::from_iter([5i32, 5, 5]) + PrimitiveArray::from_iter([5i32, 5, 5]), + &mut assertion_ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 2a87bf069bc..209621109fe 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -172,6 +172,7 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -230,6 +231,8 @@ mod tests { #[test] pub fn test_simple_pack() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -246,20 +249,25 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]) + PrimitiveArray::from_iter([0i32, 1, 2]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]) + PrimitiveArray::from_iter([4i32, 5, 6]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array.into_array(), &["three"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]) + PrimitiveArray::from_iter([0i32, 1, 2]), + &mut assertion_ctx ); } #[test] pub fn test_nested_pack() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -285,19 +293,23 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]) + PrimitiveArray::from_iter([0i32, 1, 2]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two", "two_one"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]) + PrimitiveArray::from_iter([4i32, 5, 6]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two", "two_two"]).unwrap(), - PrimitiveArray::from_iter([4i32, 5, 6]) + PrimitiveArray::from_iter([4i32, 5, 6]), + &mut assertion_ctx ); assert_arrays_eq!( primitive_field(&actual_array.into_array(), &["three"]).unwrap(), - PrimitiveArray::from_iter([0i32, 1, 2]) + PrimitiveArray::from_iter([0i32, 1, 2]), + &mut assertion_ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index ac6aa562aa6..179264ba033 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -672,10 +672,12 @@ mod tests { "$.items[1]", Some(DType::Primitive(PType::I32, Nullability::NonNullable)), )?; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(20i32), None, None]) + PrimitiveArray::from_option_iter([Some(20i32), None, None]), + &mut ctx ); Ok(()) } @@ -704,10 +706,12 @@ mod tests { "$.a", Some(DType::Primitive(PType::I32, Nullability::NonNullable)), )?; + let mut ctx = LEGACY_SESSION.create_execution_ctx(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None]) + PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None]), + &mut ctx ); Ok(()) } @@ -736,9 +740,11 @@ mod tests { )?; assert!(!result.is::()); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]) + PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]), + &mut ctx ); Ok(()) } @@ -769,7 +775,7 @@ mod tests { .map(|value| value.as_str()), Some("ok") ); - assert_nth_scalar_is_null!(result, 1); + assert_nth_scalar_is_null!(result, 1, &mut ctx); assert_eq!( result .execute_scalar(2, &mut ctx)? @@ -777,7 +783,7 @@ mod tests { .is_variant_null(), Some(true) ); - assert_nth_scalar_is_null!(result, 3); + assert_nth_scalar_is_null!(result, 3, &mut ctx); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 8a3e1b5bbc6..4fcff8c728c 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -317,6 +317,8 @@ mod tests { #[test] fn test_zip_basic() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mask = Mask::from_iter([true, false, false, true, false]); let if_true = buffer![10, 20, 30, 40, 50].into_array(); let if_false = buffer![1, 2, 3, 4, 5].into_array(); @@ -324,11 +326,13 @@ mod tests { let result = mask.into_array().zip(if_true, if_false).unwrap(); let expected = buffer![10, 2, 3, 40, 5].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[test] fn test_zip_all_true() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mask = Mask::new_true(4); let if_true = buffer![10, 20, 30, 40].into_array(); let if_false = @@ -338,12 +342,14 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), Some(40)]).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); assert_eq!(result.dtype(), if_false.dtype()) } #[test] fn test_zip_all_false_widens_nullability() { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mask = Mask::new_false(4); let if_true = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array(); @@ -353,7 +359,7 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), Some(4)]).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); assert_eq!(result.dtype(), if_true.dtype()); } @@ -369,7 +375,8 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30), Some(40)]) - .into_array() + .into_array(), + &mut ctx ); assert_eq!(result.dtype(), if_false.dtype()); Ok(()) @@ -386,7 +393,8 @@ mod tests { let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?; assert_arrays_eq!( result, - PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]).into_array() + PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]).into_array(), + &mut ctx ); assert_eq!(result.dtype(), if_true.dtype()); Ok(()) diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index d792050321d..036d72abe30 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -109,6 +109,8 @@ mod tests { #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let sum_scalar = Scalar::primitive(6i64, Nullability::Nullable); array.statistics().set( @@ -123,13 +125,15 @@ mod tests { let expected = ConstantArray::new(Scalar::primitive(6i64, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_returns_null_when_sum_is_missing() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let result = array @@ -142,13 +146,15 @@ mod tests { 3, ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_sum_per_chunk() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let chunk0 = buffer![1i32, 2].into_array(); let sum_scalar = Scalar::primitive(3i64, Nullability::Nullable); chunk0.statistics().set( @@ -177,13 +183,15 @@ mod tests { Validity::from_iter([true, true, false, false, false]), ) .into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_null_count() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); let null_count_scalar = Scalar::primitive(2u64, Nullability::NonNullable); @@ -203,13 +211,15 @@ mod tests { let expected = ConstantArray::new(Scalar::primitive(2u64, Nullability::Nullable), 4).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_all_null_from_null_count() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); array .statistics() @@ -222,13 +232,15 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_all_null_false_from_inexact_low_null_count() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, Some(2), None]).into_array(); array .statistics() @@ -241,13 +253,15 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_returns_null_for_inexact_full_null_count_as_all_null() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, Some(2), None]).into_array(); array .statistics() @@ -260,13 +274,15 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_all_non_null_from_null_count() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); array .statistics() @@ -279,13 +295,15 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_all_non_null_true_from_inexact_zero_null_count() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); array .statistics() @@ -298,13 +316,15 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_returns_null_for_inexact_nonzero_null_count_as_all_non_null() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); array @@ -318,7 +338,7 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 4).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } @@ -338,6 +358,8 @@ mod tests { #[test] fn stat_expr_reads_cached_all_nan_from_nan_count() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(f32::NAN), Some(f32::NAN)]) .into_array(); @@ -352,13 +374,15 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_all_nan_false_from_inexact_low_nan_count() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32), Some(f32::NAN)]) .into_array(); @@ -373,13 +397,15 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_returns_null_for_inexact_full_nan_count_as_all_nan() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32), Some(f32::NAN)]) .into_array(); @@ -394,13 +420,15 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_all_non_nan_true_from_inexact_zero_nan_count() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![1.0f32, 2.0, 3.0].into_array(); array .statistics() @@ -413,13 +441,15 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_returns_null_for_inexact_nonzero_nan_count_as_all_non_nan() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1.0f32), Some(f32::NAN), Some(3.0)]) .into_array(); array @@ -433,13 +463,15 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } #[test] fn stat_expr_reads_cached_min_and_max() -> VortexResult<()> { + let assertion_session = crate::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![3i32, 1, 2].into_array(); array .statistics() @@ -460,7 +492,7 @@ mod tests { .into_array(); let expected_min = ConstantArray::new(Scalar::primitive(1i32, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(min_result, expected_min); + assert_arrays_eq!(min_result, expected_min, &mut assertion_ctx); let max_result = array .apply(&stat( @@ -473,7 +505,7 @@ mod tests { .into_array(); let expected_max = ConstantArray::new(Scalar::primitive(3i32, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(max_result, expected_max); + assert_arrays_eq!(max_result, expected_max, &mut assertion_ctx); Ok(()) } diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index ca77bc24ca8..37fab4f7c07 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -112,6 +112,8 @@ mod tests { #[case] input: ListViewArray, #[case] expect_list: bool, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array_ref = input.clone().into_array(); let result = BtrBlocksCompressor::default() .compress(&array_ref, &mut SESSION.create_execution_ctx())?; @@ -120,12 +122,14 @@ mod tests { } else { assert!(result.as_opt::().is_some()); } - assert_arrays_eq!(result, input); + assert_arrays_eq!(result, input, &mut assertion_ctx); Ok(()) } #[test] fn test_constant_all_true() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -133,12 +137,14 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut assertion_ctx); Ok(()) } #[test] fn test_constant_all_false() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -146,12 +152,14 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut assertion_ctx); Ok(()) } #[test] fn test_nullable_all_valid_compressed() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true; 100]), Validity::from(BitBuffer::from(vec![true; 100])), @@ -162,12 +170,14 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut assertion_ctx); Ok(()) } #[test] fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0))); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity); let btr = BtrBlocksCompressor::default(); @@ -176,12 +186,14 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(!compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut assertion_ctx); Ok(()) } #[test] fn test_mixed_not_constant() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true, false, true]), Validity::NonNullable, @@ -192,12 +204,14 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(!compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut assertion_ctx); Ok(()) } #[test] fn test_binary_constant_compressed() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = vec![Some(b"constant-bytes".as_slice()); 100]; let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); let btr = BtrBlocksCompressor::default(); @@ -206,12 +220,14 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut assertion_ctx); Ok(()) } #[test] fn test_binary_dict_compressed() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"]; let values = (0..1000) .map(|idx| Some(distinct_values[idx % distinct_values.len()])) @@ -223,7 +239,7 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut assertion_ctx); Ok(()) } @@ -244,17 +260,15 @@ mod tests { ); let compressor = BtrBlocksCompressorBuilder::default().with_compact().build(); - let compressed = compressor.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; assert!( compressed.is::(), "expected Zstd, got {}", compressed.encoding_id() ); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } @@ -277,17 +291,15 @@ mod tests { let compressor = BtrBlocksCompressorBuilder::default() .only_cuda_compatible() .build(); - let compressed = compressor.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; assert!( compressed.is::(), "expected ZstdBuffers, got {}", compressed.encoding_id() ); - assert_arrays_eq!(compressed, array); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } } diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index da84831fcc0..7752989ea08 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -22,7 +22,6 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::float::FloatRLEScheme; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] @@ -58,6 +57,8 @@ fn test_compress() -> VortexResult<()> { #[test] fn test_rle_compression() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(1.5f32, 100)); values.extend(iter::repeat_n(2.7f32, 200)); @@ -71,7 +72,7 @@ fn test_rle_compression() -> VortexResult<()> { assert!(compressed.is::()); let expected = Buffer::copy_from(&values).into_array(); - assert_arrays_eq!(compressed, expected); + assert_arrays_eq!(compressed, expected, &mut assertion_ctx); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index 94d2c0760e0..48e7bbd6bb8 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -28,7 +28,6 @@ use vortex_session::VortexSession; use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] @@ -162,6 +161,8 @@ fn test_rle_compressed() -> VortexResult<()> { #[cfg(feature = "unstable_encodings")] #[test] fn test_delta_compressed() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex_array::assert_arrays_eq; use vortex_fastlanes::Delta; @@ -191,7 +192,7 @@ fn test_delta_compressed() -> VortexResult<()> { "Delta was applied more than once in the tree:\n{}", compressed.display_tree() ); - assert_arrays_eq!(compressed, array.into_array()); + assert_arrays_eq!(compressed, array.into_array(), &mut assertion_ctx); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index 5a3a0f33a40..27c9afb4150 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -27,7 +27,6 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] @@ -72,6 +71,8 @@ fn test_dict_encodable() -> VortexResult<()> { #[test] fn constant_mostly_nulls() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = PrimitiveArray::new( buffer![189u8, 189, 189, 189, 189, 189, 189, 189, 189, 0, 46], Validity::from_iter(vec![ @@ -89,12 +90,14 @@ fn constant_mostly_nulls() -> VortexResult<()> { let decoded = compressed; let expected = PrimitiveArray::new(buffer![0u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 46], validity).into_array(); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut assertion_ctx); Ok(()) } #[test] fn nullable_sequence() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = (0i32..20).step_by(7).collect_vec(); let array = PrimitiveArray::from_option_iter(values.clone().into_iter().map(Some)); @@ -104,12 +107,14 @@ fn nullable_sequence() -> VortexResult<()> { let decoded = compressed; let expected = PrimitiveArray::from_option_iter(values.into_iter().map(Some)).into_array(); - assert_arrays_eq!(decoded, expected); + assert_arrays_eq!(decoded, expected, &mut assertion_ctx); Ok(()) } #[test] fn test_rle_compression() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(42i32, 100)); values.extend(iter::repeat_n(123i32, 200)); @@ -122,7 +127,7 @@ fn test_rle_compression() -> VortexResult<()> { assert!(compressed.is::()); let expected = Buffer::copy_from(&values).into_array(); - assert_arrays_eq!(compressed, expected); + assert_arrays_eq!(compressed, expected, &mut assertion_ctx); Ok(()) } diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index 2ab34d52e0c..5288637a609 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -290,7 +290,7 @@ mod tests { .clone() .execute::(&mut ctx)? .into_array(); - assert_arrays_eq!(undict, expected); + assert_arrays_eq!(undict, expected, &mut ctx); Ok(()) } } diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 1bb9ede306c..8ec3bf53345 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -296,7 +296,7 @@ mod tests { .clone() .execute::(&mut ctx)? .into_array(); - assert_arrays_eq!(undict, expected); + assert_arrays_eq!(undict, expected, &mut ctx); Ok(()) } } diff --git a/vortex-compressor/src/sample.rs b/vortex-compressor/src/sample.rs index 4821c7c6d95..4b117cb1ae7 100644 --- a/vortex-compressor/src/sample.rs +++ b/vortex-compressor/src/sample.rs @@ -133,6 +133,7 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> #[cfg(test)] mod tests { use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; @@ -143,6 +144,8 @@ mod tests { #[test] fn sample_is_deterministic() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Create a deterministic array with linear-with-noise pattern let values: Vec = (0i64..100_000).map(|i| i + (i * 7 + 3) % 11).collect(); @@ -153,7 +156,7 @@ mod tests { for _ in 0..10 { let again = sample(&array, SAMPLE_SIZE, SAMPLE_COUNT); assert_eq!(first.nbytes(), again.nbytes()); - assert_arrays_eq!(&first, &again); + assert_arrays_eq!(&first, &again, &mut assertion_ctx); } Ok(()) } diff --git a/vortex-cuda/src/dynamic_dispatch/mod.rs b/vortex-cuda/src/dynamic_dispatch/mod.rs index 0439436d505..0d8519dfd3d 100644 --- a/vortex-cuda/src/dynamic_dispatch/mod.rs +++ b/vortex-cuda/src/dynamic_dispatch/mod.rs @@ -1145,6 +1145,8 @@ mod tests { #[crate::test] async fn test_dict_mixed_width_u8_codes_u32_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let len = 3000; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u8).collect(); @@ -1170,13 +1172,15 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_dict_mixed_width_u16_codes_u32_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict_values: Vec = vec![1000, 2000, 3000, 4000, 5000]; let len = 2048; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u16).collect(); @@ -1202,13 +1206,15 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_runend_mixed_width_u64_ends_u32_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let ends: Vec = vec![1000, 2000, 3000]; let values: Vec = vec![10, 20, 30]; let len = 3000; @@ -1243,7 +1249,7 @@ mod tests { }) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -1631,6 +1637,8 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u8() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bit_width: u8 = 4; let len = 3000; let reference = 100u8; @@ -1659,12 +1667,14 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_for_bitpacked_u16() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let reference = 1000u16; @@ -1695,12 +1705,14 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_for_bitpacked_u64() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bit_width: u8 = 20; let len = 3000; let reference = 100_000u64; @@ -1729,7 +1741,7 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -1746,6 +1758,8 @@ mod tests { #[crate::test] async fn test_single_element() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values: Vec = vec![42]; let primitive = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); let bp = BitPacked::encode( @@ -1765,12 +1779,14 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected, result); + vortex::array::assert_arrays_eq!(expected, result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_exactly_elements_per_block() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Exactly 2048 elements — one full block, no remainder let bit_width: u8 = 6; let len = 2048; @@ -1799,7 +1815,7 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -1929,6 +1945,8 @@ mod tests { #[crate::test] async fn test_alp_f64_for_bitpacked() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut ctx = LEGACY_SESSION.create_execution_ctx(); // ALP(FoR(BitPacked)) with f64: same structure as the f32 test. let len = 3000; @@ -1965,7 +1983,7 @@ mod tests { let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let gpu = CanonicalCudaExt::into_host(canonical).await?.into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -2029,6 +2047,8 @@ mod tests { #[crate::test] async fn alp_slice_device_patches() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Regression test for https://github.com/vortex-data/vortex/issues/7838#issuecomment-4452796116. let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 4096; @@ -2084,13 +2104,15 @@ mod tests { expected[2048 - 100] = std::f64::consts::LN_2; let expected = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected, gpu); + vortex::array::assert_arrays_eq!(expected, gpu, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_runend_u32_ends_u16_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // RunEnd with u32 ends, u16 values. Output type = u16. // Ends (u32) differ from output (u16) → pending subtree. let ends: Vec = vec![500, 1000, 1500, 2000]; @@ -2128,13 +2150,15 @@ mod tests { }) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_dict_bitpacked_u8_codes_u32_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Dict with BitPacked u8 codes (narrower than u32 output) and u32 values. // The kernel's bitunpack_typed decodes at the source's native width and // widens to T, so this fuses into a single kernel launch. @@ -2166,7 +2190,7 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -2184,6 +2208,8 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); @@ -2210,7 +2236,7 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -2227,6 +2253,8 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); @@ -2254,7 +2282,7 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -2280,6 +2308,8 @@ mod tests { #[case] ends: Vec, #[case] values: Vec, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let ends_u64: Vec = ends.iter().map(|e| (*e).into()).collect(); let len = *ends_u64.last().unwrap() as usize; let bit_width = 64 - ends_u64.iter().max().unwrap().leading_zeros() as u8; @@ -2314,7 +2344,7 @@ mod tests { prev = *end; } let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -2324,6 +2354,8 @@ mod tests { /// at native width and widens to T, fusing everything. #[crate::test] async fn test_runend_mixed_width_for_bp_u16_ends_u32_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let ends: Vec = vec![500, 1000, 1500, 2000]; let values: Vec = vec![100, 200, 300, 400]; let len = 2000usize; @@ -2359,13 +2391,15 @@ mod tests { prev = end; } let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_sliced_dict_mixed_width() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Sliced Dict with u8 codes and u32 values — combines PartiallyFused + slice handling. let dict_values: Vec = vec![100, 200, 300, 400]; let full_len = 4096; @@ -2389,7 +2423,7 @@ mod tests { .map(|&c| dict_values[c as usize]) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -2476,6 +2510,8 @@ mod tests { /// Nullable Primitive array — LOAD source with validity propagated. #[crate::test] async fn test_nullable_primitive() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let array = PrimitiveArray::from_option_iter( @@ -2489,7 +2525,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -2499,6 +2535,8 @@ mod tests { /// validity, so this produces a real nullable FoR(BitPacked) tree. #[crate::test] async fn test_nullable_for_bitpacked() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 2048; @@ -2546,7 +2584,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -2571,6 +2609,8 @@ mod tests { /// AllValid nullable array — should fuse and produce AllValid output. #[crate::test] async fn test_all_valid_nullable() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let values: Vec = (0..2048).collect(); @@ -2583,7 +2623,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -2607,6 +2647,8 @@ mod tests { /// Dict with non-nullable codes but nullable values should still fuse. #[crate::test] async fn test_dict_nullable_values_fuses() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex::buffer::buffer; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; @@ -2624,7 +2666,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -2632,6 +2674,8 @@ mod tests { /// Validity must survive through fused dispatch and into the filter. #[crate::test] async fn test_nullable_fused_then_filter() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex::array::arrays::FilterArray; use vortex::mask::Mask; @@ -2662,7 +2706,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -2912,6 +2956,8 @@ mod tests { /// u8 BitPacked with patches (bit_width=3, patch values > 7). #[crate::test] async fn test_bitpacked_with_patches_u8() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bit_width: u8 = 3; let len = 3000usize; let max_val = (1u8 << bit_width) - 1; @@ -2938,13 +2984,15 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } /// u16 BitPacked with patches (bit_width=6, patch values > 63). #[crate::test] async fn test_bitpacked_with_patches_u16() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bit_width: u8 = 6; let len = 3000usize; let max_val = (1u16 << bit_width) - 1; @@ -2971,13 +3019,15 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } /// u64 BitPacked with patches (bit_width=4, patch values > 15). #[crate::test] async fn test_bitpacked_with_patches_u64() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bit_width: u8 = 4; let len = 3000usize; let max_val = (1u64 << bit_width) - 1; @@ -3004,7 +3054,7 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); Ok(()) } @@ -3111,6 +3161,8 @@ mod tests { /// dispatch alongside patch application. #[crate::test] async fn test_nullable_bitpacked_with_patches() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 3000usize; @@ -3146,7 +3198,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu); + vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index a46cf517c16..5cd4dede718 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -194,6 +194,8 @@ mod tests { /// FoR(BitPacked) u32 — entire tree compiles into a single fused plan. #[crate::test] async fn test_fused() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let values: Vec = (0..2048).map(|i| (i % 128) as u32).collect(); @@ -213,7 +215,7 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -221,6 +223,8 @@ mod tests { /// Exercises the unsigned type reinterpretation in CudaDispatchPlan::execute. #[crate::test] async fn test_fused_f32() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex::encodings::alp::ALP; use vortex::encodings::alp::Exponents; @@ -249,13 +253,15 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } /// ALP with patches — plan builder rejects it, falls back to ALPExecutor. #[crate::test] async fn test_fallback() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex::array::patches::Patches; use vortex::array::validity::Validity::NonNullable as NN; use vortex::buffer::buffer; @@ -287,7 +293,7 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -297,6 +303,8 @@ mod tests { #[cfg(feature = "unstable_encodings")] #[crate::test] async fn test_partial_fusion() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex::array::arrays::DictArray; use vortex::array::session::ArraySessionExt; use vortex::encodings::fastlanes; @@ -353,13 +361,15 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } /// Filter(FoR(BP), mask) — FoR+BP fuses via dyn dispatch, then CUB filters the result. #[crate::test] async fn test_filter_fused_child() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -385,7 +395,7 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu); + assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); Ok(()) } @@ -403,6 +413,8 @@ mod tests { ))] #[crate::test] async fn test_ext_storage_gpu_decode(#[case] ext: ExtensionArray) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -419,13 +431,15 @@ mod tests { assert!(!storage.is_host(), "storage was not decoded on the device"); let actual = actual.into_host().await?.into_array(); - assert_arrays_eq!(expected, actual); + assert_arrays_eq!(expected, actual, &mut assertion_ctx); Ok(()) } /// Extension over already-canonical storage executes unchanged. #[crate::test] async fn test_ext_canonical_storage() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -443,7 +457,7 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(ext.into_array(), actual); + assert_arrays_eq!(ext.into_array(), actual, &mut assertion_ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/arrays/constant.rs b/vortex-cuda/src/kernel/arrays/constant.rs index 0f769bacbf9..7854e5f12e5 100644 --- a/vortex-cuda/src/kernel/arrays/constant.rs +++ b/vortex-cuda/src/kernel/arrays/constant.rs @@ -199,6 +199,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::scalar::Scalar; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -223,6 +224,8 @@ mod tests { async fn test_cuda_constant_materialization( #[case] constant_array: ConstantArray, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -236,13 +239,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_cuda_constant_empty_array() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -257,13 +262,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_cuda_constant_small_array() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -279,7 +286,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index 2b21d470f09..312bc2b4ea2 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -308,6 +308,7 @@ mod tests { use vortex::dtype::DecimalDType; use vortex::dtype::i256; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -324,6 +325,8 @@ mod tests { #[crate::test] async fn test_cuda_dict_u32_values_u8_codes() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -350,12 +353,18 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_u64_values_u16_codes() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -385,12 +394,18 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_i32_values_u32_codes() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -416,12 +431,18 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_large_array() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -448,12 +469,18 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_values_with_validity() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -480,12 +507,18 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_codes_with_validity() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -517,12 +550,18 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_both_with_validity() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -561,12 +600,18 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_i64_values_with_validity() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -606,12 +651,18 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_all_valid_matches_baseline() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -639,7 +690,11 @@ mod tests { let cuda_result = cuda_primitive_to_host(cuda_result)?; // Compare CUDA result with baseline - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } @@ -655,6 +710,8 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i8_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -677,12 +734,18 @@ mod tests { .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i16_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -705,12 +768,18 @@ mod tests { .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i32_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -733,12 +802,18 @@ mod tests { .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i64_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -764,12 +839,18 @@ mod tests { .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i128_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -800,7 +881,11 @@ mod tests { .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } @@ -814,6 +899,8 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_u8_codes() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -833,12 +920,18 @@ mod tests { .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_string_values_u16_codes() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -858,12 +951,18 @@ mod tests { .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_string_max_inlined_12_bytes() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -885,12 +984,18 @@ mod tests { .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_string_outlined_views() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -915,12 +1020,18 @@ mod tests { .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_string_empty_strings() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -940,12 +1051,18 @@ mod tests { .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_string_values_with_validity() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -966,12 +1083,18 @@ mod tests { .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_string_outlined_with_validity() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -999,12 +1122,18 @@ mod tests { .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i256_values() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1035,7 +1164,11 @@ mod tests { .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!(cuda_result.into_array(), baseline.into_array()); + assert_arrays_eq!( + cuda_result.into_array(), + baseline.into_array(), + &mut assertion_ctx + ); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/alp.rs b/vortex-cuda/src/kernel/encodings/alp.rs index 37c83af5d93..bf2ded1007a 100644 --- a/vortex-cuda/src/kernel/encodings/alp.rs +++ b/vortex-cuda/src/kernel/encodings/alp.rs @@ -170,6 +170,8 @@ mod tests { /// Patches must carry `chunk_offsets` — the fused kernel requires them. #[crate::test] async fn test_cuda_alp_decompression_f32() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -206,7 +208,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } @@ -217,6 +219,8 @@ mod tests { /// preserved through the standalone ALP GPU executor. #[crate::test] async fn test_cuda_alp_nullable_with_patches() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -250,7 +254,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } @@ -258,6 +262,8 @@ mod tests { /// elements are actually null. #[crate::test] async fn test_cuda_alp_all_valid_nullable() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -281,7 +287,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } @@ -291,6 +297,8 @@ mod tests { /// (zero patches) via the offset math rather than the NULL sentinel. #[crate::test] async fn test_cuda_alp_multi_chunk_sparse_patches() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -327,7 +335,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } @@ -336,6 +344,8 @@ mod tests { /// so this guards the fast-path for the (i64, f64) kernel variant. #[crate::test] async fn test_cuda_alp_f64_multi_chunk_with_patches() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -370,7 +380,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } @@ -380,6 +390,8 @@ mod tests { /// (existing tests have ≤ 6 patches per chunk). #[crate::test] async fn test_cuda_alp_dense_patches_single_chunk() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -419,7 +431,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } @@ -428,6 +440,8 @@ mod tests { /// loop. Includes a patch in the tail. #[crate::test] async fn test_cuda_alp_partial_tail_chunk() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -455,7 +469,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 72161505b0a..e6cdf5fcba2 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -250,6 +250,8 @@ mod tests { #[case] iter: impl Iterator, #[case] bw: u8, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -275,13 +277,15 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] fn test_patches() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -310,7 +314,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } @@ -325,6 +329,8 @@ mod tests { #[case::bw_7(7)] #[crate::test] fn test_cuda_bitunpack_u8(#[case] bit_width: u8) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -355,7 +361,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } @@ -378,6 +384,8 @@ mod tests { #[case::bw_15(15)] #[crate::test] fn test_cuda_bitunpack_u16(#[case] bit_width: u8) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -408,7 +416,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } @@ -447,6 +455,8 @@ mod tests { #[case::bw_31(31)] #[crate::test] fn test_cuda_bitunpack_u32(#[case] bit_width: u8) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -477,7 +487,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } @@ -548,6 +558,8 @@ mod tests { #[case::bw_63(63)] #[crate::test] fn test_cuda_bitunpack_u64(#[case] bit_width: u8) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -577,13 +589,15 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] fn test_cuda_bitunpack_sliced() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let bit_width = 32; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -615,7 +629,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } @@ -667,6 +681,8 @@ mod tests { /// offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_sliced_patches_offset_within_chunk() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -701,7 +717,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } @@ -709,6 +725,8 @@ mod tests { /// Test slicing a bitpacked array multiple times, accumulating offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_double_sliced_patches() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -755,7 +773,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } @@ -763,6 +781,8 @@ mod tests { /// Test slicing to skip an entire chunk's worth of patches. #[crate::test] fn test_cuda_bitunpack_sliced_skip_first_chunk_patches() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -807,7 +827,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/date_time_parts.rs b/vortex-cuda/src/kernel/encodings/date_time_parts.rs index 4a1e9cc458f..582de845d94 100644 --- a/vortex-cuda/src/kernel/encodings/date_time_parts.rs +++ b/vortex-cuda/src/kernel/encodings/date_time_parts.rs @@ -214,6 +214,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::extension::datetime::TimeUnit; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -279,6 +280,8 @@ mod tests { #[case] subseconds: Vec, #[case] time_unit: TimeUnit, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -293,13 +296,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_cuda_datetimeparts_large_array() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -319,13 +324,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_cuda_datetimeparts_with_nulls() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -367,7 +374,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index 9e28737db79..a018b0d33f8 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -70,6 +70,7 @@ mod tests { use vortex::dtype::DecimalDType; use vortex::encodings::decimal_byte_parts::DecimalByteParts; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::session::CudaSession; @@ -85,6 +86,8 @@ mod tests { #[case] precision: u8, #[case] scale: i8, ) { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("create execution context"); @@ -103,6 +106,10 @@ mod tests { .await .vortex_expect("GPU decode"); - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!( + cpu_result.into_array(), + gpu_result.into_array(), + &mut assertion_ctx + ); } } diff --git a/vortex-cuda/src/kernel/encodings/for_.rs b/vortex-cuda/src/kernel/encodings/for_.rs index b0e231060b0..f71bd1738ac 100644 --- a/vortex-cuda/src/kernel/encodings/for_.rs +++ b/vortex-cuda/src/kernel/encodings/for_.rs @@ -160,6 +160,8 @@ mod tests { #[case::u64(make_for_array((0..2050).map(|i| (i % 2050) as u64).collect(), 1000000u64))] #[crate::test] async fn test_cuda_for_decompression(#[case] for_array: FoRArray) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -173,13 +175,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_signed_ffor() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -204,6 +208,6 @@ mod tests { .vortex_expect("copying to host failed") .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); } } diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 6b978ebba23..0b6a65b7f8a 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -209,6 +209,7 @@ mod tests { use vortex::encodings::fsst::fsst_compress; use vortex::encodings::fsst::fsst_train_compressor; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -235,6 +236,8 @@ mod tests { #[case] strings: Vec>, #[case] nullability: Nullability, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -252,13 +255,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } /// Exercises the multi-block grid-stride path on a larger dataset. #[crate::test] async fn test_cuda_fsst_decompression_roundtrip_large() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); use vortex_fsst::test_utils::make_fsst_clickbench_urls; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) @@ -275,7 +280,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index 5af16ca8e04..2059f3da432 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -174,17 +174,15 @@ mod tests { use vortex::encodings::runend::RunEndArray; use vortex::error::VortexExpect; use vortex::error::VortexResult; + use vortex_array::ExecutionCtx; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; use crate::executor::CudaArrayExt; use crate::session::CudaSession; - fn make_runend_array( - ends: Vec, - values: Vec, - ctx: &mut vortex::array::ExecutionCtx, - ) -> RunEndArray + fn make_runend_array(ends: Vec, values: Vec, ctx: &mut ExecutionCtx) -> RunEndArray where V: NativePType, E: NativePType, @@ -196,17 +194,19 @@ mod tests { RunEnd::new(ends_array, values_array, ctx) } - type RunEndBuilder = fn(&mut vortex::array::ExecutionCtx) -> RunEndArray; + type RunEndBuilder = fn(&mut ExecutionCtx) -> RunEndArray; #[rstest] - #[case::u32_ends_u8_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![3u32, 6, 10], vec![10u8, 20, 30], ctx))] - #[case::u32_ends_u32_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u32, 5, 10], vec![1u32, 2, 3], ctx))] - #[case::u32_ends_f64_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u32, 5, 8], vec![1.5f64, 2.5, 3.5], ctx))] - #[case::u8_ends_i32_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u8, 5, 10], vec![1i32, 2, 3], ctx))] - #[case::u32_ends_i32_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u32, 5, 10], vec![1i32, 2, 3], ctx))] - #[case::u64_ends_i32_values(|ctx: &mut vortex::array::ExecutionCtx| make_runend_array(vec![2u64, 5, 10], vec![1i32, 2, 3], ctx))] + #[case::u32_ends_u8_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![3u32, 6, 10], vec![10u8, 20, 30], ctx))] + #[case::u32_ends_u32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u32, 5, 10], vec![1u32, 2, 3], ctx))] + #[case::u32_ends_f64_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u32, 5, 8], vec![1.5f64, 2.5, 3.5], ctx))] + #[case::u8_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u8, 5, 10], vec![1i32, 2, 3], ctx))] + #[case::u32_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u32, 5, 10], vec![1i32, 2, 3], ctx))] + #[case::u64_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u64, 5, 10], vec![1i32, 2, 3], ctx))] #[crate::test] async fn test_cuda_runend_types(#[case] build: RunEndBuilder) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -221,13 +221,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_large_array() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -251,13 +253,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_single_run() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -273,13 +277,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_many_small_runs() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -300,13 +306,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_nullable_values_falls_back_to_cpu() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -332,7 +340,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/sequence.rs b/vortex-cuda/src/kernel/encodings/sequence.rs index b9363fdd9dc..9fff943c520 100644 --- a/vortex-cuda/src/kernel/encodings/sequence.rs +++ b/vortex-cuda/src/kernel/encodings/sequence.rs @@ -92,6 +92,7 @@ mod tests { use vortex::dtype::Nullability; use vortex::encodings::sequence::Sequence; use vortex::scalar::PValue; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::CudaSession; @@ -124,6 +125,8 @@ mod tests { len: usize, nullability: Nullability, ) { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).unwrap(); let array = Sequence::try_new_typed(base, multiplier, nullability, len).unwrap(); @@ -139,6 +142,6 @@ mod tests { .unwrap() .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); } } diff --git a/vortex-cuda/src/kernel/encodings/zigzag.rs b/vortex-cuda/src/kernel/encodings/zigzag.rs index 3c4cf8e08ee..7f8b5738e9a 100644 --- a/vortex-cuda/src/kernel/encodings/zigzag.rs +++ b/vortex-cuda/src/kernel/encodings/zigzag.rs @@ -105,6 +105,7 @@ mod tests { use vortex::buffer::Buffer; use vortex::encodings::zigzag::ZigZag; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -112,6 +113,8 @@ mod tests { #[crate::test] async fn test_cuda_zigzag_decompression_u32() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -133,7 +136,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result); + assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index 703a0ddf0e1..9e1bd991d40 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -356,6 +356,7 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::encodings::zstd::Zstd; use vortex::error::VortexResult; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -364,6 +365,8 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_utf8() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -384,12 +387,18 @@ mod tests { .execute(zstd_array.into_array(), &mut cuda_ctx) .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!( + cpu_result.into_array(), + gpu_result.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_zstd_decompression_multiple_frames() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -420,12 +429,18 @@ mod tests { .execute(zstd_array.into_array(), &mut cuda_ctx) .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!( + cpu_result.into_array(), + gpu_result.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_zstd_decompression_sliced() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -452,7 +467,11 @@ mod tests { .execute(sliced_zstd.clone(), &mut cuda_ctx) .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!( + cpu_result.into_array(), + gpu_result.into_array(), + &mut assertion_ctx + ); Ok(()) } @@ -461,6 +480,8 @@ mod tests { /// correct results instead of panicking. #[crate::test] async fn test_cuda_zstd_nullable_falls_back_to_cpu() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -486,7 +507,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs index 691f1ee23e7..7da492879e7 100644 --- a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs +++ b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs @@ -226,6 +226,7 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::error::VortexExpect; use vortex::error::VortexResult; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -233,6 +234,8 @@ mod tests { #[crate::test] async fn test_cuda_zstd_buffers_decompression_primitive() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -246,12 +249,18 @@ mod tests { .into_host() .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!( + cpu_result.into_array(), + gpu_result.into_array(), + &mut assertion_ctx + ); Ok(()) } #[crate::test] async fn test_cuda_zstd_buffers_decompression_varbinview() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -273,7 +282,11 @@ mod tests { .into_host() .await?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array()); + assert_arrays_eq!( + cpu_result.into_array(), + gpu_result.into_array(), + &mut assertion_ctx + ); Ok(()) } } diff --git a/vortex-cuda/src/kernel/filter/decimal.rs b/vortex-cuda/src/kernel/filter/decimal.rs index cd3be77f423..b5a308f6e19 100644 --- a/vortex-cuda/src/kernel/filter/decimal.rs +++ b/vortex-cuda/src/kernel/filter/decimal.rs @@ -48,6 +48,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -88,6 +89,8 @@ mod tests { #[case] input: DecimalArray, #[case] mask: Mask, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -103,13 +106,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_gpu_filter_decimal_large_array() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -133,7 +138,7 @@ mod tests { .into_array(); assert_eq!(cpu_result.len(), gpu_result.len()); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/filter/primitive.rs b/vortex-cuda/src/kernel/filter/primitive.rs index 2d943f590b6..4c701f81e9b 100644 --- a/vortex-cuda/src/kernel/filter/primitive.rs +++ b/vortex-cuda/src/kernel/filter/primitive.rs @@ -46,6 +46,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -82,6 +83,8 @@ mod tests { #[case] input: PrimitiveArray, #[case] mask: Mask, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -97,13 +100,15 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } #[crate::test] async fn test_gpu_filter_large_array() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -127,7 +132,7 @@ mod tests { .into_array(); assert_eq!(cpu_result.len(), gpu_result.len()); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/filter/varbinview.rs b/vortex-cuda/src/kernel/filter/varbinview.rs index e4ca3fa4df4..704d448c6f5 100644 --- a/vortex-cuda/src/kernel/filter/varbinview.rs +++ b/vortex-cuda/src/kernel/filter/varbinview.rs @@ -46,6 +46,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -68,6 +69,8 @@ mod tests { #[case] input: VarBinViewArray, #[case] mask: Mask, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -83,7 +86,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/patches/mod.rs b/vortex-cuda/src/kernel/patches/mod.rs index 993526af1fb..4408a3003f2 100644 --- a/vortex-cuda/src/kernel/patches/mod.rs +++ b/vortex-cuda/src/kernel/patches/mod.rs @@ -264,7 +264,7 @@ mod tests { .unwrap() .into_primitive(); - assert_arrays_eq!(cpu_result, gpu_result); + assert_arrays_eq!(cpu_result, gpu_result, &mut ctx); } fn force_cast(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> PrimitiveArray { diff --git a/vortex-cuda/src/kernel/patches/types.rs b/vortex-cuda/src/kernel/patches/types.rs index 6234250189a..63c3b731253 100644 --- a/vortex-cuda/src/kernel/patches/types.rs +++ b/vortex-cuda/src/kernel/patches/types.rs @@ -228,6 +228,7 @@ mod tests { use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::patches::Patches; use vortex_error::VortexResult; @@ -317,6 +318,8 @@ mod tests { #[case] chunk_offsets: ArrayRef, #[case] expected: ArrayRef, ) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let indices = PrimitiveArray::from_iter([100u32, 1100, 2100, 3100, 4100]); let values = PrimitiveArray::from_iter([10u32, 11, 12, 13, 14]); @@ -337,7 +340,7 @@ mod tests { NonNullable, ) .into_array(); - assert_arrays_eq!(expected, actual); + assert_arrays_eq!(expected, actual, &mut assertion_ctx); assert_eq!(device_patches.n_chunks, 3); assert_eq!(device_patches.offset, 1024); assert_eq!(device_patches.offset_within_chunk, 0); diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 03bd45c0074..4cfcdbca5f4 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -415,6 +415,8 @@ mod tests { #[test] fn test_integer_vector_conversion() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let values = vec![1i32, 2, 3, 4, 5]; let len = values.len(); @@ -431,7 +433,7 @@ mod tests { let result = flat_vector_to_vortex(&vector, len).unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4), Some(5)]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[test] @@ -642,7 +644,7 @@ mod tests { let result = flat_vector_to_vortex(&vector, len).unwrap(); let vortex_array = result.execute::(&mut ctx).unwrap(); let expected = BoolArray::new(BitBuffer::from(values), Validity::AllValid); - assert_arrays_eq!(vortex_array, expected); + assert_arrays_eq!(vortex_array, expected, &mut ctx); } #[test] @@ -712,7 +714,8 @@ mod tests { assert_eq!(vortex_array.len(), len); assert_arrays_eq!( vortex_array.list_elements_at(0).unwrap(), - PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]) + PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]), + &mut ctx ); } @@ -741,7 +744,8 @@ mod tests { assert_eq!(vortex_array.len(), len); assert_arrays_eq!( vortex_array.fixed_size_list_elements_at(0).unwrap(), - PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]) + PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]), + &mut ctx ); } @@ -797,11 +801,13 @@ mod tests { assert_eq!(vortex_array.struct_fields().nfields(), 2); assert_arrays_eq!( vortex_array.unmasked_field(0), - PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]) + PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]), + &mut ctx ); assert_arrays_eq!( vortex_array.unmasked_field(1), - PrimitiveArray::from_option_iter([Some(5i32), Some(6), Some(7), Some(8)]) + PrimitiveArray::from_option_iter([Some(5i32), Some(6), Some(7), Some(8)]), + &mut ctx ); } @@ -846,7 +852,8 @@ mod tests { assert_eq!(vortex_array.len(), len); assert_arrays_eq!( vortex_array.list_elements_at(0).unwrap(), - PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]) + PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4)]), + &mut ctx ); assert_eq!( vortex_array @@ -897,11 +904,13 @@ mod tests { assert_eq!(vortex_array.len(), len); assert_arrays_eq!( vortex_array.list_elements_at(0).unwrap(), - PrimitiveArray::from_option_iter([Some(3i32), Some(4)]) + PrimitiveArray::from_option_iter([Some(3i32), Some(4)]), + &mut ctx ); assert_arrays_eq!( vortex_array.list_elements_at(1).unwrap(), - PrimitiveArray::from_option_iter([Some(1i32), Some(2)]) + PrimitiveArray::from_option_iter([Some(1i32), Some(2)]), + &mut ctx ); } @@ -955,11 +964,13 @@ mod tests { // Valid entries should work correctly. assert_arrays_eq!( vortex_array.list_elements_at(0).unwrap(), - PrimitiveArray::from_option_iter([Some(1i32), Some(2)]) + PrimitiveArray::from_option_iter([Some(1i32), Some(2)]), + &mut ctx ); assert_arrays_eq!( vortex_array.list_elements_at(2).unwrap(), - PrimitiveArray::from_option_iter([Some(3i32), Some(4)]) + PrimitiveArray::from_option_iter([Some(3i32), Some(4)]), + &mut ctx ); // Verify the null entry has sanitized offset/size (offset=2, size=0) rather than garbage. @@ -1048,7 +1059,7 @@ mod tests { None, Some(wkb_b.as_slice()), ]); - assert_arrays_eq!(storage, expected); + assert_arrays_eq!(storage, expected, &mut ctx); Ok(()) } diff --git a/vortex-duckdb/src/exporter/dict.rs b/vortex-duckdb/src/exporter/dict.rs index 45aa724dd86..b167df7e192 100644 --- a/vortex-duckdb/src/exporter/dict.rs +++ b/vortex-duckdb/src/exporter/dict.rs @@ -158,7 +158,6 @@ impl> ColumnExporter for DictExporter { #[cfg(test)] mod tests { use vortex::VortexSessionDefault; - use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ConstantArray; @@ -220,7 +219,7 @@ mod tests { let mut chunk = DataChunk::new([LogicalType::new(cpp::duckdb_type::DUCKDB_TYPE_INTEGER)]); - let mut ctx = ExecutionCtx::new(VortexSession::default()); + let mut ctx = VortexSession::default().create_execution_ctx(); new_exporter_with_flatten(&arr, &ConversionCache::default(), &mut ctx, false)?.export( 0, 2, diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 4b8191bd4bb..32051f17d96 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -453,7 +453,7 @@ mod tests { use vortex::VortexSessionDefault; use vortex::array::arrays::StructArray; use vortex::session::VortexSession; - use vortex_array::ExecutionCtx; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; @@ -532,17 +532,21 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_no_options() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let (array, struct_array) = scan(ptr::null()); - assert_arrays_eq!(vx_array::as_ref(array), struct_array); + assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); unsafe { vx_array_free(array) }; } #[test] #[cfg_attr(miri, ignore)] fn test_project_all() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let opts = vx_scan_options::default(); let (array, struct_array) = scan(&raw const opts); - assert_arrays_eq!(vx_array::as_ref(array), struct_array); + assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); unsafe { vx_array_free(array) }; } @@ -550,6 +554,8 @@ mod tests { #[cfg_attr(miri, ignore)] fn test_project_single_field() { unsafe { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let root = vx_expression_root(); let mut opts = vx_scan_options::default(); @@ -560,7 +566,8 @@ mod tests { let (array, struct_array) = scan(&raw const opts); assert_arrays_eq!( vx_array::as_ref(array), - struct_array.unmasked_field_by_name(field).unwrap() + struct_array.unmasked_field_by_name(field).unwrap(), + &mut assertion_ctx ); vx_array_free(array); vx_expression_free(field_expr); @@ -573,7 +580,7 @@ mod tests { #[cfg_attr(miri, ignore)] fn test_project_sum() { let session = VortexSession::default(); - let mut ctx = ExecutionCtx::new(session); + let mut ctx = session.create_execution_ctx(); unsafe { let root = vx_expression_root(); let mut opts = vx_scan_options::default(); @@ -710,12 +717,14 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_ordered() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let opts = vx_scan_options { ordered: true, ..Default::default() }; let (array, struct_array) = scan(&raw const opts); - assert_arrays_eq!(vx_array::as_ref(array), struct_array); + assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); unsafe { vx_array_free(array) }; } diff --git a/vortex-ffi/src/struct_array.rs b/vortex-ffi/src/struct_array.rs index bfac6107319..99e10c29c34 100644 --- a/vortex-ffi/src/struct_array.rs +++ b/vortex-ffi/src/struct_array.rs @@ -242,7 +242,7 @@ mod tests { .clone() .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(array, struct_array); + assert_arrays_eq!(array, struct_array, &mut ctx); } vx_array_free(array); diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 17020e11c04..01afe39663d 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -75,7 +75,6 @@ use crate::VERSION; use crate::VortexFile; use crate::WriteOptionsSessionExt; use crate::footer::SegmentSpec; - static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session() .with::() @@ -312,7 +311,7 @@ async fn test_read_projection() { .unmasked_field(0) .clone(); let expected = VarBinArray::from(strings_expected.to_vec()).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); let array = file .scan() @@ -338,7 +337,7 @@ async fn test_read_projection() { .unmasked_field(0) .clone(); let expected = Buffer::copy_from(numbers_expected).into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } #[tokio::test] @@ -478,6 +477,8 @@ async fn test_empty_varbin_array_roundtrip() { #[tokio::test] #[cfg_attr(miri, ignore)] async fn issue_5385_filter_casted_column() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = StructArray::try_from_iter([("x", buffer![1u8, 2, 3, 4, 5])]) .unwrap() .into_array(); @@ -510,7 +511,8 @@ async fn issue_5385_filter_casted_column() { assert_arrays_eq!( result, - StructArray::try_from_iter([("x", buffer![1u8])]).unwrap() + StructArray::try_from_iter([("x", buffer![1u8])]).unwrap(), + &mut assertion_ctx ); } @@ -563,7 +565,7 @@ async fn filter_string() { let names_expected = VarBinArray::from_iter(vec![Some("Joseph")], DType::Utf8(Nullability::Nullable)) .into_array(); - assert_arrays_eq!(names_actual, names_expected); + assert_arrays_eq!(names_actual, names_expected, &mut ctx); let ages_actual = result[0] .clone() @@ -572,7 +574,7 @@ async fn filter_string() { .unmasked_field(1) .clone(); let ages_expected = PrimitiveArray::from_option_iter([Some(25i32)]).into_array(); - assert_arrays_eq!(ages_actual, ages_expected); + assert_arrays_eq!(ages_actual, ages_expected, &mut ctx); } #[tokio::test] @@ -631,7 +633,7 @@ async fn filter_or() { DType::Utf8(Nullability::Nullable), ) .into_array(); - assert_arrays_eq!(names_actual, names_expected); + assert_arrays_eq!(names_actual, names_expected, &mut ctx); let ages_actual = result[0] .clone() @@ -640,7 +642,7 @@ async fn filter_or() { .unmasked_field(1) .clone(); let ages_expected = PrimitiveArray::from_option_iter([Some(25i32), None]).into_array(); - assert_arrays_eq!(ages_actual, ages_expected); + assert_arrays_eq!(ages_actual, ages_expected, &mut ctx); } #[tokio::test] @@ -696,7 +698,7 @@ async fn filter_and() { DType::Utf8(Nullability::Nullable), ) .into_array(); - assert_arrays_eq!(names_actual, names_expected); + assert_arrays_eq!(names_actual, names_expected, &mut ctx); let ages_actual = result[0] .clone() @@ -705,7 +707,7 @@ async fn filter_and() { .unmasked_field(1) .clone(); let ages_expected = PrimitiveArray::from_option_iter([Some(25i32), Some(31i32)]).into_array(); - assert_arrays_eq!(ages_actual, ages_expected); + assert_arrays_eq!(ages_actual, ages_expected, &mut ctx); } #[tokio::test] @@ -775,7 +777,7 @@ async fn test_with_indices_simple() { .map(|&x| expected_numbers[x as usize]) .collect(); let expected_array = Buffer::copy_from(&expected_kept_numbers).into_array(); - assert_arrays_eq!(actual_kept_numbers_array, expected_array); + assert_arrays_eq!(actual_kept_numbers_array, expected_array, &mut ctx); // test all indices let actual_array = file @@ -791,7 +793,7 @@ async fn test_with_indices_simple() { .unwrap(); let actual_numbers_array = actual_array.unmasked_field(0).clone(); let expected_array = Buffer::copy_from(&expected_numbers).into_array(); - assert_arrays_eq!(actual_numbers_array, expected_array); + assert_arrays_eq!(actual_numbers_array, expected_array, &mut ctx); } #[tokio::test] @@ -841,7 +843,7 @@ async fn test_with_indices_on_two_columns() { .map(|&x| strings_expected[x as usize]) .collect(); let strings_expected_array = VarBinArray::from(strings_expected_vec).into_array(); - assert_arrays_eq!(strings_actual, strings_expected_array); + assert_arrays_eq!(strings_actual, strings_expected_array, &mut ctx); let numbers_actual = array.unmasked_field(1).clone(); let numbers_expected_vec: Vec = kept_indices @@ -849,7 +851,7 @@ async fn test_with_indices_on_two_columns() { .map(|&x| numbers_expected[x as usize]) .collect(); let numbers_expected_array = Buffer::copy_from(&numbers_expected_vec).into_array(); - assert_arrays_eq!(numbers_actual, numbers_expected_array); + assert_arrays_eq!(numbers_actual, numbers_expected_array, &mut ctx); } #[tokio::test] @@ -922,7 +924,7 @@ async fn test_with_indices_and_with_row_filter_simple() { .filter(|&x| x > 50) .collect(); let expected_array = expected_kept_numbers.into_array(); - assert_arrays_eq!(actual_kept_numbers_array, expected_array); + assert_arrays_eq!(actual_kept_numbers_array, expected_array, &mut ctx); // test all indices let actual_array = file @@ -945,7 +947,7 @@ async fn test_with_indices_and_with_row_filter_simple() { .cloned() .collect(); let expected_numbers_array = expected_filtered.into_array(); - assert_arrays_eq!(actual_numbers_array, expected_numbers_array); + assert_arrays_eq!(actual_numbers_array, expected_numbers_array, &mut ctx); } #[tokio::test] @@ -1004,11 +1006,11 @@ async fn filter_string_chunked() { let names_expected = VarBinArray::from_iter(vec![Some("Joseph")], DType::Utf8(Nullability::Nullable)) .into_array(); - assert_arrays_eq!(names_actual, names_expected); + assert_arrays_eq!(names_actual, names_expected, &mut ctx); let ages_actual = actual_array.unmasked_field(1).clone(); let ages_expected = PrimitiveArray::from_option_iter([Some(25i32)]).into_array(); - assert_arrays_eq!(ages_actual, ages_expected); + assert_arrays_eq!(ages_actual, ages_expected, &mut ctx); } #[tokio::test] @@ -1107,7 +1109,7 @@ async fn test_pruning_with_or() { Some("P".to_owned()), ]) .into_array(); - assert_arrays_eq!(letters_actual, letters_expected); + assert_arrays_eq!(letters_actual, letters_expected, &mut ctx); let numbers_actual = actual_array.unmasked_field(1).clone(); let numbers_expected = PrimitiveArray::from_option_iter([ @@ -1123,7 +1125,7 @@ async fn test_pruning_with_or() { Some(22), ]) .into_array(); - assert_arrays_eq!(numbers_actual, numbers_expected); + assert_arrays_eq!(numbers_actual, numbers_expected, &mut ctx); } #[tokio::test] @@ -1164,7 +1166,7 @@ async fn test_repeated_projection() { .execute::(&mut ctx) .unwrap(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut ctx); } async fn chunked_file() -> VortexResult { @@ -1186,11 +1188,13 @@ async fn chunked_file() -> VortexResult { #[tokio::test] async fn basic_file_roundtrip() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let vxf = chunked_file().await?; let result = vxf.scan()?.into_array_stream()?.read_all().await?; let expected = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } @@ -1229,6 +1233,8 @@ async fn file_excluding_dtype() -> VortexResult<()> { #[tokio::test] async fn file_take() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let vxf = chunked_file().await?; let result = vxf .scan()? @@ -1238,7 +1244,7 @@ async fn file_take() -> VortexResult<()> { .await?; let expected = buffer![0i32, 1, 8].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); Ok(()) } @@ -1460,7 +1466,7 @@ async fn test_writer_multiple_pushes() -> VortexResult<()> { .unmasked_field_by_name("numbers")? .clone(); let expected = buffer![1u32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); - assert_arrays_eq!(numbers, expected); + assert_arrays_eq!(numbers, expected, &mut ctx); Ok(()) } @@ -1495,7 +1501,7 @@ async fn test_writer_push_stream() -> VortexResult<()> { .unmasked_field_by_name("numbers")? .clone(); let expected = buffer![1u32, 2, 3, 4, 5, 6].into_array(); - assert_arrays_eq!(numbers, expected); + assert_arrays_eq!(numbers, expected, &mut ctx); Ok(()) } @@ -1560,7 +1566,7 @@ async fn test_writer_empty_chunks() -> VortexResult<()> { .unmasked_field_by_name("numbers")? .clone(); let expected = buffer![1u32, 2].into_array(); - assert_arrays_eq!(numbers, expected); + assert_arrays_eq!(numbers, expected, &mut ctx); Ok(()) } @@ -1599,7 +1605,7 @@ async fn test_writer_mixed_push_and_stream() -> VortexResult<()> { .unmasked_field_by_name("numbers")? .clone(); let expected = buffer![1u32, 2, 3, 4, 5, 6].into_array(); - assert_arrays_eq!(numbers, expected); + assert_arrays_eq!(numbers, expected, &mut ctx); Ok(()) } diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index 2550606d86a..d2a815e9758 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -29,7 +29,6 @@ use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::table::TableStrategy; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; - static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session() .with::() @@ -124,6 +123,8 @@ async fn test_file_roundtrip() { /// with "Array vortex.fill_null does not support serialization". #[tokio::test] async fn test_dict_listview_validity_roundtrip() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]).into_array(); let offsets = PrimitiveArray::from_iter(vec![0u32, 2, 4]).into_array(); let sizes = PrimitiveArray::from_iter(vec![2u32, 2, 1]).into_array(); @@ -164,6 +165,6 @@ async fn test_dict_listview_validity_roundtrip() { .await .unwrap() .expect("read back should succeed"); - vortex_array::assert_arrays_eq!(data, chunk); + vortex_array::assert_arrays_eq!(data, chunk, &mut assertion_ctx); assert!(stream.next().await.is_none(), "expected a single chunk"); } diff --git a/vortex-ipc/src/iterator.rs b/vortex-ipc/src/iterator.rs index 318d062127c..4714f055dc5 100644 --- a/vortex-ipc/src/iterator.rs +++ b/vortex-ipc/src/iterator.rs @@ -172,6 +172,7 @@ mod test { use std::io::Cursor; use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorExt; @@ -182,6 +183,8 @@ mod test { #[test] fn test_sync_stream() -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let ipc_buffer = array .to_array_iterator() @@ -192,7 +195,7 @@ mod test { assert_eq!(reader.dtype(), array.dtype()); let result = reader.read_all()?; - assert_arrays_eq!(result, array); + assert_arrays_eq!(result, array, &mut assertion_ctx); Ok(()) } diff --git a/vortex-ipc/src/stream.rs b/vortex-ipc/src/stream.rs index 8b6d9f69e3e..1ab4b87a64d 100644 --- a/vortex-ipc/src/stream.rs +++ b/vortex-ipc/src/stream.rs @@ -225,6 +225,7 @@ mod test { use futures::io::Cursor; use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamExt; @@ -235,6 +236,8 @@ mod test { #[tokio::test] async fn test_async_stream() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![1, 2, 3].into_array(); let ipc_buffer = array .to_array_stream() @@ -249,7 +252,7 @@ mod test { assert_eq!(reader.dtype(), array.dtype()); let result = reader.read_all().await.unwrap(); - assert_arrays_eq!(result, array); + assert_arrays_eq!(result, array, &mut assertion_ctx); } /// Wrapper that limits reads to small chunks to simulate network behavior @@ -271,6 +274,8 @@ mod test { #[tokio::test] async fn test_async_stream_chunked() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(); let ipc_buffer = array .to_array_stream() @@ -288,12 +293,14 @@ mod test { let result = reader.read_all().await.unwrap(); let expected = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } /// Test with 1-byte chunks to stress-test partial read handling. #[tokio::test] async fn test_async_stream_single_byte_chunks() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = buffer![42i64, -1, 0, i64::MAX, i64::MIN].into_array(); let ipc_buffer = array .to_array_stream() @@ -311,6 +318,6 @@ mod test { let result = reader.read_all().await.unwrap(); let expected = buffer![42i64, -1, 0, i64::MAX, i64::MIN].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } } diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index 4d095cfee37..94de11f0096 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -353,6 +353,7 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; @@ -471,6 +472,8 @@ mod test { #[from(chunked_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|_h| async { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let result = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() @@ -484,7 +487,7 @@ mod test { .unwrap(); let expected = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); }) } } diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 685444a61e5..46f1da82b05 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -422,6 +422,8 @@ mod tests { #[test] fn reading_nested_packs_works() { block_on(|handle| async move { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), @@ -500,7 +502,7 @@ mod tests { ) .unwrap() .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut assertion_ctx); }) } @@ -521,6 +523,8 @@ mod tests { #[case] expected: Vec, ) { block_on(|handle| async move { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), @@ -564,13 +568,19 @@ mod tests { .await .unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter(expected)); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter(expected), + &mut assertion_ctx + ); }) } #[test] fn reading_is_null_works() { block_on(|handle| async move { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let mut ctx_exec = LEGACY_SESSION.create_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( @@ -638,12 +648,14 @@ mod tests { .execute::(&mut ctx_exec) .vortex_expect("to_canonical failed") .into_array(); - assert_arrays_eq!(actual_canonical, expected); + assert_arrays_eq!(actual_canonical, expected, &mut assertion_ctx); }) } #[test] fn reading_byte_length_pushdown_works() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = VarBinArray::from_iter( [ Some("abc"), @@ -682,7 +694,7 @@ mod tests { .await .unwrap() .into_array(); - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut assertion_ctx); }) } @@ -701,6 +713,8 @@ mod tests { } fn test_apply(original: Expression, outer: Expression, inner: Expression) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = VarBinArray::from_iter( [Some("abc"), Some("def"), None], DType::Utf8(Nullability::Nullable), @@ -713,7 +727,7 @@ mod tests { ))?; let actual = pushed.apply(&outer)?; let expected = array.apply(&original)?; - assert_arrays_eq!(actual, expected); + assert_arrays_eq!(actual, expected, &mut assertion_ctx); Ok(()) } diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 6881d3c814a..7af89a19f4a 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -233,6 +233,7 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -255,6 +256,8 @@ mod test { #[test] fn flat_identity() -> VortexResult<()> { block_on(|handle| async { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -285,7 +288,7 @@ mod test { )? .await?; - assert_arrays_eq!(result, array); + assert_arrays_eq!(result, array, &mut assertion_ctx); Ok(()) }) @@ -294,6 +297,8 @@ mod test { #[test] fn flat_expr() { block_on(|handle| async { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); @@ -326,13 +331,15 @@ mod test { .unwrap(); let expected = BoolArray::from_iter([false, false, false, true, true].map(Some)); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); }) } #[test] fn flat_unaligned_row_mask() { block_on(|handle| async { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -359,7 +366,7 @@ mod test { .unwrap(); let expected = PrimitiveArray::new(buffer![3i32, 4], Validity::AllValid).into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); }) } } diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index c810078e55b..cba2046ebe8 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -277,6 +277,7 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::LEGACY_SESSION; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index fb826a1b401..8ba7ebf007d 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -327,6 +327,7 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::IntoArray as _; use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; use vortex_array::expr::eq; @@ -351,6 +352,8 @@ mod tests { #[test] fn flat_expr_no_row_id() { block_on(|handle| async { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -386,7 +389,8 @@ mod tests { assert_arrays_eq!( result, - BoolArray::from_iter([false, false, true, false, false]) + BoolArray::from_iter([false, false, true, false, false]), + &mut assertion_ctx ); }) } @@ -394,6 +398,8 @@ mod tests { #[test] fn flat_expr_row_id() { block_on(|handle| async { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -429,7 +435,8 @@ mod tests { assert_arrays_eq!( result, - BoolArray::from_iter([false, false, false, false, true]) + BoolArray::from_iter([false, false, false, false, true]), + &mut assertion_ctx ); }) } @@ -437,6 +444,8 @@ mod tests { #[test] fn flat_expr_or() { block_on(|handle| async { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -476,7 +485,8 @@ mod tests { assert_arrays_eq!( result, - BoolArray::from_iter([true, false, true, false, true]) + BoolArray::from_iter([true, false, true, false, true]), + &mut assertion_ctx ); }) } diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index d27bd3746e7..50e7b409237 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -634,6 +634,8 @@ mod tests { fn test_struct_layout( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -645,13 +647,15 @@ mod tests { }) .unwrap(); let expected = BoolArray::from_iter([true, false, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[rstest] fn test_struct_layout_row_mask( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -668,7 +672,7 @@ mod tests { .unwrap(); let expected = BoolArray::from_iter([true, false]); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); } #[rstest] @@ -701,14 +705,16 @@ mod tests { let result_struct_a = result.clone().execute::(&mut ctx).unwrap(); assert_arrays_eq!( result_struct_a.unmasked_field_by_name("a").unwrap(), - expected_a + expected_a, + &mut ctx ); let expected_b = PrimitiveArray::from_iter([4i32, 5]); let result_struct_b = result.execute::(&mut ctx).unwrap(); assert_arrays_eq!( result_struct_b.unmasked_field_by_name("b").unwrap(), - expected_b + expected_b, + &mut ctx ); } @@ -716,6 +722,8 @@ mod tests { fn test_struct_layout_nulls( #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), ) { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); // Read the layout source from the top. let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) @@ -739,8 +747,8 @@ mod tests { .unwrap(), Scalar::null(result.dtype().clone()), ); - assert_nth_scalar!(result, 1, 2); - assert_nth_scalar!(result, 2, 3); + assert_nth_scalar!(result, 1, 2, &mut assertion_ctx); + assert_nth_scalar!(result, 2, 3, &mut assertion_ctx); } #[rstest] diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index f6346d0760f..b6cbc9b52da 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -226,6 +226,7 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -309,6 +310,8 @@ mod test { #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|handle| async { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let session = session_with_handle(handle); let result = layout .new_reader("".into(), segments, &session, &Default::default()) @@ -323,7 +326,7 @@ mod test { .unwrap(); let expected = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); - assert_arrays_eq!(result, expected); + assert_arrays_eq!(result, expected, &mut assertion_ctx); }) } diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index dbb2b3c333c..e5f2af494de 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -359,6 +359,7 @@ mod tests { use std::sync::Arc; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; use vortex_array::aggregate_fn::NumericalAggregateOpts; @@ -446,6 +447,7 @@ mod tests { 10, ) .unwrap(); + let ctx = &mut SESSION.create_execution_ctx(); // A >= 6 // => A.max < 6 @@ -454,7 +456,8 @@ mod tests { let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([true, false, false]) + BoolArray::from_iter([true, false, false]), + ctx ); // A > 5 @@ -464,7 +467,8 @@ mod tests { let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([true, false, false]) + BoolArray::from_iter([true, false, false]), + ctx ); // A < 2 @@ -472,7 +476,11 @@ mod tests { let expr = lt(root(), lit(2i32)); let pruning_expr = falsify(&expr, PType::I32.into()); let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true, true])); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([false, true, true]), + ctx + ); } #[test] @@ -512,19 +520,25 @@ mod tests { 10, ) .unwrap(); + let ctx = &mut SESSION.create_execution_ctx(); let expr = gt(root(), lit(5i32)); let pruning_expr = falsify(&expr, PType::I32.into()); let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([true, false, false]) + BoolArray::from_iter([true, false, false]), + ctx ); let expr = lt(root(), lit(2i32)); let pruning_expr = falsify(&expr, PType::I32.into()); let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true, true])); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([false, true, true]), + ctx + ); } #[test] @@ -548,7 +562,8 @@ mod tests { let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([false, false, true]) + BoolArray::from_iter([false, false, true]), + &mut SESSION.create_execution_ctx() ); } @@ -595,7 +610,8 @@ mod tests { let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([true, false, false]) + BoolArray::from_iter([true, false, false]), + &mut SESSION.create_execution_ctx() ); } @@ -615,7 +631,11 @@ mod tests { .unwrap(); let mask = zone_map.prune(&all_null(root()), &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true, true])); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([false, true, true]), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -636,7 +656,8 @@ mod tests { let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([true, false, false]) + BoolArray::from_iter([true, false, false]), + &mut SESSION.create_execution_ctx() ); } @@ -655,14 +676,20 @@ mod tests { 10, ) .unwrap(); + let ctx = &mut SESSION.create_execution_ctx(); let mask = zone_map.prune(&all_null(root()), &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([true, false, true])); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([true, false, true]), + ctx + ); let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([false, true, false]) + BoolArray::from_iter([false, true, false]), + ctx ); } @@ -681,14 +708,20 @@ mod tests { 10, ) .unwrap(); + let ctx = &mut SESSION.create_execution_ctx(); let mask = zone_map.prune(&all_nan(root()), &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([true, false, true])); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([true, false, true]), + ctx + ); let mask = zone_map.prune(&all_non_nan(root()), &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([false, true, false]) + BoolArray::from_iter([false, true, false]), + ctx ); } @@ -702,12 +735,13 @@ mod tests { 8, ) .unwrap(); + let ctx = &mut SESSION.create_execution_ctx(); let mask = zone_map.prune(&all_nan(root()), &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, false])); + assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, false]), ctx); let mask = zone_map.prune(&all_non_nan(root()), &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([true, true])); + assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([true, true]), ctx); } #[test] @@ -720,11 +754,13 @@ mod tests { 10, ) .unwrap(); + let ctx = &mut SESSION.create_execution_ctx(); let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([false, false, false]) + BoolArray::from_iter([false, false, false]), + ctx ); let expr = gt(root(), lit(5u64)); @@ -732,7 +768,8 @@ mod tests { let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([false, false, false]) + BoolArray::from_iter([false, false, false]), + ctx ); } @@ -751,13 +788,15 @@ mod tests { 12, ) .unwrap(); + let ctx = &mut SESSION.create_execution_ctx(); let expr = gt(root(), lit(5.0f32)); let pruning_expr = falsify(&expr, PType::F32.into()); let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([false, false, false]) + BoolArray::from_iter([false, false, false]), + ctx ); let nan_count = NanCount.bind(EmptyOptions); @@ -783,7 +822,8 @@ mod tests { let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([true, false, false]) + BoolArray::from_iter([true, false, false]), + ctx ); } @@ -825,7 +865,11 @@ mod tests { let pruning_expr = falsify(&expr, PType::F32.into()); let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true])); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([false, true]), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -855,7 +899,11 @@ mod tests { // Missing StatFn lowers to a nullable null literal, so `is_null(...)` is true for every zone. let mask = zone_map.prune(&predicate, &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([true, true, true])); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([true, true, true]), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -905,7 +953,8 @@ mod tests { let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([false, true, false]) + BoolArray::from_iter([false, true, false]), + &mut SESSION.create_execution_ctx() ); } @@ -926,7 +975,11 @@ mod tests { .unwrap(); let mask = zone_map.prune(&all_null(root()), &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, true, true])); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([false, true, true]), + &mut SESSION.create_execution_ctx() + ); } #[test] @@ -948,7 +1001,8 @@ mod tests { let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap(); assert_arrays_eq!( mask.into_array(), - BoolArray::from_iter([true, false, false]) + BoolArray::from_iter([true, false, false]), + &mut SESSION.create_execution_ctx() ); } } diff --git a/vortex-layout/src/scan/arrow.rs b/vortex-layout/src/scan/arrow.rs index d3042c8b510..485c2685934 100644 --- a/vortex-layout/src/scan/arrow.rs +++ b/vortex-layout/src/scan/arrow.rs @@ -127,6 +127,7 @@ mod tests { use arrow_schema::Field; use arrow_schema::Schema; use vortex_array::ArrayRef; + use vortex_array::VortexSessionExecute; use vortex_array::arrow::FromArrowArray; use vortex_error::VortexResult; diff --git a/vortex-python/src/arrays/range_to_sequence.rs b/vortex-python/src/arrays/range_to_sequence.rs index 9bfc0d331ee..dcb04d13ee4 100644 --- a/vortex-python/src/arrays/range_to_sequence.rs +++ b/vortex-python/src/arrays/range_to_sequence.rs @@ -78,6 +78,7 @@ mod test { use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; + use vortex_array::VortexSessionExecute; use crate::arrays::range_to_sequence::range_len; use crate::arrays::range_to_sequence::sequence_array_from_range; @@ -98,21 +99,35 @@ mod test { #[test] fn test_sequence_array_from_len() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let dtype = DType::Primitive(PType::U16, Nullability::NonNullable); let arr = sequence_array_from_range::(0, 10, 1, dtype).unwrap(); - assert_arrays_eq!(arr, buffer![0u16, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array()); + assert_arrays_eq!( + arr, + buffer![0u16, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(), + &mut assertion_ctx + ); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let arr = sequence_array_from_range::(0, 10, 5, dtype).unwrap(); - assert_arrays_eq!(arr, buffer![0i32, 5].into_array()); + assert_arrays_eq!(arr, buffer![0i32, 5].into_array(), &mut assertion_ctx); let dtype = DType::Primitive(PType::I8, Nullability::NonNullable); let arr = sequence_array_from_range::(-5, 5, 3, dtype).unwrap(); - assert_arrays_eq!(arr, buffer![-5i8, -2, 1, 4].into_array()); + assert_arrays_eq!( + arr, + buffer![-5i8, -2, 1, 4].into_array(), + &mut assertion_ctx + ); let dtype = DType::Primitive(PType::I8, Nullability::NonNullable); let arr = sequence_array_from_range::(3, -3, -1, dtype).unwrap(); - assert_arrays_eq!(arr, buffer![3i8, 2, 1, 0, -1, -2].into_array()); + assert_arrays_eq!( + arr, + buffer![3i8, 2, 1, 0, -1, -2].into_array(), + &mut assertion_ctx + ); let dtype = DType::Primitive(PType::U32, Nullability::NonNullable); let result = sequence_array_from_range::(1_000_000, 10, -500_000, dtype); @@ -122,6 +137,10 @@ mod test { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let arr = sequence_array_from_range::(1_000_000, 10, -500_000, dtype).unwrap(); - assert_arrays_eq!(arr, buffer![1_000_000i32, 500_000].into_array()); + assert_arrays_eq!( + arr, + buffer![1_000_000i32, 500_000].into_array(), + &mut assertion_ctx + ); } } diff --git a/vortex-tensor/src/types/vector/arrow.rs b/vortex-tensor/src/types/vector/arrow.rs index 9a148839f3d..a011dc50c49 100644 --- a/vortex-tensor/src/types/vector/arrow.rs +++ b/vortex-tensor/src/types/vector/arrow.rs @@ -319,7 +319,7 @@ mod tests { let imported = session.from_arrow_array(arrow, &field)?; assert_eq!(imported.dtype(), original.dtype()); - vortex_array::assert_arrays_eq!(imported, original); + vortex_array::assert_arrays_eq!(imported, original, &mut ctx); Ok(()) } @@ -411,7 +411,7 @@ mod tests { let field = session.to_arrow_field("v", original.dtype())?; let imported = session.from_arrow_array(arrow, &field)?; assert_eq!(imported.dtype(), original.dtype()); - vortex_array::assert_arrays_eq!(imported, original); + vortex_array::assert_arrays_eq!(imported, original, &mut ctx); Ok(()) } } diff --git a/vortex-test/compat-gen/src/check.rs b/vortex-test/compat-gen/src/check.rs index 7533ee9dab7..6923a2be65d 100644 --- a/vortex-test/compat-gen/src/check.rs +++ b/vortex-test/compat-gen/src/check.rs @@ -5,6 +5,7 @@ use std::path::Path; use clap::ValueEnum; use serde::Serialize; +use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -13,7 +14,6 @@ use vortex_error::vortex_err; use crate::adapter; use crate::fixtures::all_fixtures; - /// How to handle mismatches between directory and known fixtures. #[derive(Clone, ValueEnum)] pub enum Mode { @@ -47,6 +47,8 @@ struct FailedFixture { /// Prints JSON result to stdout, human-readable progress to stderr. /// Returns error if any fixture failed or if mode constraints are violated. pub fn check(dir: &Path, mode: Mode, exclude: &[String]) -> VortexResult<()> { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let fixtures = all_fixtures(); let fixtures: Vec<_> = fixtures .into_iter() @@ -190,7 +192,7 @@ pub fn check(dir: &Path, mode: Mode, exclude: &[String]) -> VortexResult<()> { } }; - assert_arrays_eq!(stored_array, fresh_array); + assert_arrays_eq!(stored_array, fresh_array, &mut assertion_ctx); eprintln!(" pass {fresh_name}"); result.passed.push(fresh_name.clone()); } diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/patched.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/patched.rs index 861c77c293e..389e3b1687e 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/patched.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/patched.rs @@ -3,8 +3,8 @@ use vortex_array::ArrayId; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; use vortex_array::arrays::Patched; use vortex_array::arrays::PrimitiveArray; use vortex_array::patches::Patches; @@ -29,7 +29,7 @@ impl FlatLayoutFixture for PatchedFixture { } fn build(&self) -> VortexResult { - let mut ctx = ExecutionCtx::new(vortex_array::array_session()); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter((0u64..2048).map(Some)).into_array(); let patches = Patches::new( diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs index dd707779aea..6f036525c58 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs @@ -17,6 +17,7 @@ pub fn fixtures() -> Vec> { #[cfg(test)] mod tests { + use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use super::fixtures; @@ -26,11 +27,13 @@ mod tests { #[test] fn roundtrip_fixtures_to_bytes() { for fixture in fixtures() { + let assertion_session = vortex_array::array_session(); + let mut assertion_ctx = assertion_session.create_execution_ctx(); let array = fixture.build().unwrap(); check_expected_encodings(&array, fixture.as_ref()).unwrap(); let bytes = adapter::write_file_to_bytes(array.clone()).unwrap(); let roundtripped = adapter::read_file(bytes).unwrap(); - assert_arrays_eq!(array, roundtripped); + assert_arrays_eq!(array, roundtripped, &mut assertion_ctx); } } } diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index 26ba402a534..4025c43fc46 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -72,6 +72,7 @@ fn with_byte_counter<'a, 'b>(bencher: Bencher<'a, 'b>, bytes: u64) -> Bencher<'a mod setup { use rand::rngs::StdRng; + use vortex_array::VortexSessionExecute; use super::*; From 7777f437bd186d8a8d58565add935827ed04235d Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 10:33:18 +0100 Subject: [PATCH 02/12] more Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/constant/vtable/mod.rs | 3 +-- vortex-array/src/arrays/patched/vtable/mod.rs | 17 ++++++++--------- vortex-array/src/normalize.rs | 7 +++---- 3 files changed, 12 insertions(+), 15 deletions(-) diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index 6eb599153a7..b2996e4985c 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -270,7 +270,6 @@ fn append_value_or_nulls( mod tests { use rstest::rstest; use vortex_error::VortexResult; - use vortex_session::VortexSession; use crate::IntoArray; use crate::VortexSessionExecute; @@ -286,7 +285,7 @@ mod tests { /// Appends `array` into a fresh builder and asserts the result matches `constant_canonicalize`. fn assert_append_matches_canonical(array: ConstantArray) -> VortexResult<()> { - let mut ctx = VortexSession::empty().create_execution_ctx(); + let mut ctx = crate::array_session().create_execution_ctx(); let expected = constant_canonicalize(array.as_view(), &mut ctx)?.into_array(); let mut builder = builder_with_capacity(array.dtype(), array.len()); diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index aff9f426457..58452857476 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -348,7 +348,6 @@ mod tests { use vortex_buffer::buffer; use vortex_buffer::buffer_mut; use vortex_error::VortexResult; - use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use crate::ArrayContext; @@ -383,7 +382,7 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); + let session = crate::array_session(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -416,7 +415,7 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); + let session = crate::array_session(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -449,7 +448,7 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); + let session = crate::array_session(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -482,7 +481,7 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); + let session = crate::array_session(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -519,7 +518,7 @@ mod tests { ) .unwrap(); - let session = VortexSession::empty(); + let session = crate::array_session(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -563,7 +562,7 @@ mod tests { let patches = Patches::new(len, 0, indices, patch_vals, None)?; - let session = VortexSession::empty(); + let session = crate::array_session(); let mut ctx = session.create_execution_ctx(); Patched::from_array_and_patches(array, &patches, &mut ctx) @@ -640,7 +639,7 @@ mod tests { assert_eq!(array_ref.dtype(), new_array.dtype()); // Execute both and compare results - let mut ctx = VortexSession::empty().create_execution_ctx(); + let mut ctx = crate::array_session().create_execution_ctx(); let original_executed = array_ref.execute::(&mut ctx)?.into_primitive(); let new_executed = new_array.execute::(&mut ctx)?.into_primitive(); @@ -666,7 +665,7 @@ mod tests { let new_array = array_ref.with_slots(slots.into_slots())?; // Execute and verify the inner values changed (except at patch positions) - let mut ctx = VortexSession::empty().create_execution_ctx(); + let mut ctx = crate::array_session().create_execution_ctx(); let executed = new_array.execute::(&mut ctx)?.into_primitive(); // Expected: all 5s except indices 1, 2, 3 which are patched to 10, 20, 30 diff --git a/vortex-array/src/normalize.rs b/vortex-array/src/normalize.rs index 9de1603fd3f..7957adc5297 100644 --- a/vortex-array/src/normalize.rs +++ b/vortex-array/src/normalize.rs @@ -98,7 +98,6 @@ impl ArrayRef { #[cfg(test)] mod tests { use vortex_error::VortexResult; - use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; use super::NormalizeOptions; @@ -128,7 +127,7 @@ mod tests { )? .into_array(); let allowed = HashSet::from_iter([array.encoding_id(), field.encoding_id()]); - let mut ctx = VortexSession::empty().create_execution_ctx(); + let mut ctx = crate::array_session().create_execution_ctx(); let normalized = array.clone().normalize(&mut NormalizeOptions { allowed: &allowed, @@ -173,7 +172,7 @@ mod tests { )? .into_array(); let allowed = HashSet::from_iter([array.encoding_id(), unchanged.encoding_id()]); - let mut ctx = VortexSession::empty().create_execution_ctx(); + let mut ctx = crate::array_session().create_execution_ctx(); let normalized = array.clone().normalize(&mut NormalizeOptions { allowed: &allowed, @@ -212,7 +211,7 @@ mod tests { assert_eq!(sliced.encoding_id(), Slice.id()); let allowed = HashSet::from_iter([Dict.id(), Primitive.id()]); - let mut ctx = VortexSession::empty().create_execution_ctx(); + let mut ctx = crate::array_session().create_execution_ctx(); let normalized = sliced.normalize(&mut NormalizeOptions { allowed: &allowed, From 96b4f8ff73d34067993da0d8448db102a7792be3 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 11:06:22 +0100 Subject: [PATCH 03/12] fixmore Signed-off-by: Robert Kruszewski --- .../aggregate_fn/fns/all_non_distinct/mod.rs | 6 ++++-- .../fns/all_non_distinct/variant.rs | 19 +++++++++++++++++++ vortex-array/src/arrays/variant/vtable/mod.rs | 1 - 3 files changed, 23 insertions(+), 3 deletions(-) create mode 100644 vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index 445f4e17b6f..fcb9b18346d 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -12,6 +12,7 @@ mod struct_; #[cfg(test)] mod tests; mod varbin; +mod variant; use std::sync::LazyLock; @@ -38,6 +39,7 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::fns::all_non_distinct::variant::check_variant_identical; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; @@ -264,8 +266,8 @@ fn check_canonical_identical( (Canonical::Extension(lhs), Canonical::Extension(rhs)) => { check_extension_identical(lhs, rhs, ctx) } - (Canonical::Variant(_), _) | (_, Canonical::Variant(_)) => { - vortex_bail!("Variant arrays don't support AllNonDistinct") + (Canonical::Variant(lhs), Canonical::Variant(rhs)) => { + check_variant_identical(lhs, rhs, ctx) } _ => Err(vortex_err!( "Canonical type mismatch in AllNonDistinct: {:?} vs {:?}", diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs new file mode 100644 index 00000000000..83ddd00b2b1 --- /dev/null +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs @@ -0,0 +1,19 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct; +use crate::ExecutionCtx; +use crate::arrays::VariantArray; +use crate::arrays::variant::VariantArrayExt; + +pub(super) fn check_variant_identical( + lhs: &VariantArray, + rhs: &VariantArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let lhs_core_storage = lhs.core_storage(); + let rhs_core_storage = rhs.core_storage(); + + all_non_distinct(lhs_core_storage, rhs_core_storage, ctx) +} diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 68c12cb634e..19ade687875 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -143,7 +143,6 @@ impl VTable for Variant { dtype: &DType, len: usize, metadata: &[u8], - buffers: &[BufferHandle], children: &dyn ArrayChildren, session: &VortexSession, From c6f27d09124491a295cca6619cd5c8ade90a2a97 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 12:29:11 +0100 Subject: [PATCH 04/12] rebase Signed-off-by: Robert Kruszewski --- .../fns/all_non_distinct/variant.rs | 3 +- .../src/aggregate_fn/fns/count/grouped.rs | 56 ++++++++++--------- 2 files changed, 31 insertions(+), 28 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs index 83ddd00b2b1..b28e71fb82a 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs @@ -2,8 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_error::VortexResult; -use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct; + use crate::ExecutionCtx; +use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct; use crate::arrays::VariantArray; use crate::arrays::variant::VariantArrayExt; diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index e5e780e42f0..39a957530bf 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -91,13 +91,14 @@ mod tests { use vortex_error::VortexResult; use crate::ArrayRef; + use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::DynGroupedAccumulator; use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::count::Count; + use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -110,13 +111,17 @@ mod tests { use crate::validity::Validity; /// Run a grouped count through the accumulator. - fn grouped_count_actual(groups: &ArrayRef, elem_dtype: &DType) -> VortexResult { + fn grouped_count_actual( + groups: &ArrayRef, + elem_dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult { let mut acc = GroupedAccumulator::try_new( Count, NumericalAggregateOpts::default(), elem_dtype.clone(), )?; - acc.accumulate_list(groups, &mut LEGACY_SESSION.create_execution_ctx())?; + acc.accumulate_list(groups, ctx)?; acc.finish() } @@ -125,7 +130,7 @@ mod tests { elements: &ArrayRef, ranges: &[(usize, usize)], ) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let counts: Buffer = ranges .iter() .map(|&(offset, size)| { @@ -151,28 +156,26 @@ mod tests { #[test] fn listview_counts_all_valid() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I32, NonNullable); let ranges = [(0, 2), (2, 1), (3, 3), (6, 0)]; let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let expected = grouped_count_reference(&elements, &ranges)?; let direct = PrimitiveArray::new(buffer![2u64, 1, 3, 0], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct, &mut assertion_ctx); - assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct, &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_counts_with_nulls() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) .into_array(); @@ -180,20 +183,19 @@ mod tests { let ranges = [(0, 3), (3, 2), (5, 1)]; let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let expected = grouped_count_reference(&elements, &ranges)?; // Group 0: {1, null, 3} -> 2. Group 1: {null, null} -> 0. Group 2: {9} -> 1. let direct = PrimitiveArray::new(buffer![2u64, 0, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct, &mut assertion_ctx); - assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct, &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_counts_varbinview_with_nulls() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = VarBinViewArray::from_iter_nullable_str([ Some("a"), None, @@ -206,17 +208,18 @@ mod tests { let ranges = [(0, 2), (2, 2), (4, 1)]; let groups = listview(elements.clone(), &ranges)?; - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let expected = grouped_count_reference(&elements, &ranges)?; let direct = PrimitiveArray::new(buffer![1u64, 1, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct, &mut assertion_ctx); - assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct, &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn fixed_size_counts_float_nans() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(2.0)]) .into_array(); @@ -225,32 +228,31 @@ mod tests { FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into_array(); // NaNs are excluded by default and counted otherwise. - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let expected = PrimitiveArray::new(buffer![1u64, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &expected, &mut ctx); let mut acc = GroupedAccumulator::try_new(Count, NumericalAggregateOpts::include_nans(), elem_dtype)?; - acc.accumulate_list(&groups, &mut LEGACY_SESSION.create_execution_ctx())?; + acc.accumulate_list(&groups, &mut ctx)?; let actual = acc.finish()?; let expected = PrimitiveArray::new(buffer![2u64, 1], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &expected); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn fixed_size_counts_with_nulls() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(); let elem_dtype = DType::Primitive(PType::I32, Nullable); let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?.into_array(); - let actual = grouped_count_actual(&groups, &elem_dtype)?; + let actual = grouped_count_actual(&groups, &elem_dtype, &mut ctx)?; let direct = PrimitiveArray::new(buffer![1u64, 2], Validity::NonNullable).into_array(); - assert_arrays_eq!(&actual, &direct, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct, &mut ctx); Ok(()) } } From 8a01038686e8c8ba88248bc012f6c96089e2f116 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 15:37:57 +0100 Subject: [PATCH 05/12] fixes Signed-off-by: Robert Kruszewski --- .../src/compute/allnondistinct.rs | 65 +++++++++++++++++++ encodings/parquet-variant/src/compute/mod.rs | 6 ++ encodings/parquet-variant/src/kernel.rs | 14 ++++ encodings/parquet-variant/src/lib.rs | 1 + .../fns/all_non_distinct/variant.rs | 24 +++++-- 5 files changed, 104 insertions(+), 6 deletions(-) create mode 100644 encodings/parquet-variant/src/compute/allnondistinct.rs create mode 100644 encodings/parquet-variant/src/compute/mod.rs diff --git a/encodings/parquet-variant/src/compute/allnondistinct.rs b/encodings/parquet-variant/src/compute/allnondistinct.rs new file mode 100644 index 00000000000..766741bb75c --- /dev/null +++ b/encodings/parquet-variant/src/compute/allnondistinct.rs @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Debug; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct; +use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::arrays::Struct; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use crate::ParquetVariant; +use crate::ParquetVariantArrayExt; + +/// Lets `AllNonDistinct` compare two `ParquetVariant` arrays without canonicalizing them. +/// +/// `AllNonDistinct` accumulates over a `Struct{lhs, rhs}` batch, so this kernel is registered for +/// the struct encoding and inspects the two children. When both are `ParquetVariant`, we compare +/// the typed (`typed_value`) arrays if both sides are shredded, and fall back to the raw `value` +/// arrays otherwise. Comparing these child arrays directly avoids re-canonicalizing the variant +/// (which would recurse through the `Variant` canonical form). +#[derive(Debug)] +pub struct AllNonDistinctParquetVariant; + +impl DynAggregateKernel for AllNonDistinctParquetVariant { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !aggregate_fn.is::() { + return Ok(None); + } + + let Some(batch) = batch.as_opt::() else { + return Ok(None); + }; + let lhs = batch.unmasked_field(0); + let rhs = batch.unmasked_field(1); + let (Some(lhs), Some(rhs)) = ( + lhs.as_opt::(), + rhs.as_opt::(), + ) else { + return Ok(None); + }; + + let identical = match (lhs.typed_value_array(), rhs.typed_value_array()) { + (Some(lhs_typed), Some(rhs_typed)) => all_non_distinct(lhs_typed, rhs_typed, ctx)?, + _ => match (lhs.value_array(), rhs.value_array()) { + (Some(lhs_value), Some(rhs_value)) => all_non_distinct(lhs_value, rhs_value, ctx)?, + // Mixed shredding layouts: let the generic canonical path handle it. + _ => return Ok(None), + }, + }; + + Ok(Some(Scalar::bool(identical, Nullability::NonNullable))) + } +} diff --git a/encodings/parquet-variant/src/compute/mod.rs b/encodings/parquet-variant/src/compute/mod.rs new file mode 100644 index 00000000000..f69f592b89d --- /dev/null +++ b/encodings/parquet-variant/src/compute/mod.rs @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod allnondistinct; + +pub use allnondistinct::*; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 684ce0217f4..f4aa2b145ed 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -18,9 +18,13 @@ use vortex_array::ArrayVTable; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnVTable; +use vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrays::Dict; use vortex_array::arrays::Filter; use vortex_array::arrays::Slice; +use vortex_array::arrays::Struct; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; @@ -45,6 +49,7 @@ use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; +use crate::compute::AllNonDistinctParquetVariant; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -64,6 +69,15 @@ pub(crate) fn initialize(session: &VortexSession) { TakeExecuteAdaptor(ParquetVariant), ); kernels.register_execute_parent_kernel(VariantGet.id(), ParquetVariant, VariantGetKernel); + + // `AllNonDistinct` accumulates over a `Struct{lhs, rhs}` batch, so the kernel is keyed on the + // struct encoding and matches `ParquetVariant` children. See `AllNonDistinctParquetVariant`. + let aggregates = session.aggregate_fns(); + aggregates.register_aggregate_kernel( + Struct.id(), + Some(AllNonDistinct.id()), + &AllNonDistinctParquetVariant, + ); } #[derive(Default, Debug)] diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 0076af48d22..91045a65e48 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -26,6 +26,7 @@ mod array; mod arrow; +mod compute; mod kernel; mod operations; mod validity; diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs index b28e71fb82a..1576b8c3d6a 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs @@ -4,17 +4,29 @@ use vortex_error::VortexResult; use crate::ExecutionCtx; -use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct; use crate::arrays::VariantArray; -use crate::arrays::variant::VariantArrayExt; +/// Checks whether two canonical variant arrays are element-wise non-distinct. +/// +/// Variant values cannot be routed back through [`all_non_distinct`]: canonicalizing a variant +/// value array yields another canonical variant (with no shredded tree), which would recurse +/// forever. The generic fallback therefore compares logical variant scalars row-by-row. Encodings +/// that can compare their typed/value children more cheaply (e.g. `ParquetVariant`) register an +/// aggregate kernel that intercepts the comparison before it reaches this fallback. +/// +/// [`all_non_distinct`]: super::all_non_distinct pub(super) fn check_variant_identical( lhs: &VariantArray, rhs: &VariantArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - let lhs_core_storage = lhs.core_storage(); - let rhs_core_storage = rhs.core_storage(); - - all_non_distinct(lhs_core_storage, rhs_core_storage, ctx) + if lhs.len() != rhs.len() { + return Ok(false); + } + for idx in 0..lhs.len() { + if lhs.execute_scalar(idx, ctx)? != rhs.execute_scalar(idx, ctx)? { + return Ok(false); + } + } + Ok(true) } From 42d09bc2bb8ea47ee28fa68ca9b55e1a86a349e1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 16:47:55 +0100 Subject: [PATCH 06/12] inline Signed-off-by: Robert Kruszewski --- .../src/aggregate_fn/fns/sum/grouped.rs | 15 ++-- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 18 ++--- vortex-array/src/array/typed.rs | 7 +- vortex-array/src/arrays/bool/array.rs | 12 +-- .../src/arrays/bool/compute/filter.rs | 13 +--- vortex-array/src/arrays/bool/compute/take.rs | 16 ++-- vortex-array/src/arrays/bool/patch.rs | 7 +- vortex-array/src/arrays/bool/vtable/mod.rs | 4 +- .../src/arrays/bool/vtable/operations.rs | 4 +- vortex-array/src/arrays/chunked/array.rs | 12 +-- .../src/arrays/chunked/compute/cast.rs | 4 +- .../src/arrays/chunked/compute/take.rs | 22 ++---- vortex-array/src/arrays/chunked/tests.rs | 30 +++---- .../src/arrays/chunked/vtable/operations.rs | 13 +--- .../src/arrays/constant/compute/take.rs | 6 +- .../src/arrays/constant/vtable/canonical.rs | 3 +- vortex-array/src/arrays/datetime/test.rs | 7 +- .../src/arrays/decimal/compute/fill_null.rs | 12 +-- .../src/arrays/decimal/compute/take.rs | 7 +- vortex-array/src/arrays/decimal/vtable/mod.rs | 4 +- vortex-array/src/arrays/dict/array.rs | 3 +- vortex-array/src/arrays/dict/compute/cast.rs | 9 +-- .../src/arrays/dict/compute/fill_null.rs | 3 +- vortex-array/src/arrays/dict/compute/like.rs | 4 +- vortex-array/src/arrays/dict/compute/mod.rs | 12 +-- vortex-array/src/arrays/dict/compute/slice.rs | 7 +- vortex-array/src/arrays/dict/tests.rs | 4 +- .../src/arrays/extension/compute/cast.rs | 3 +- .../arrays/filter/execute/fixed_size_list.rs | 10 +-- .../src/arrays/filter/execute/listview.rs | 9 +-- .../src/arrays/filter/execute/struct_.rs | 10 +-- .../arrays/fixed_size_list/tests/filter.rs | 22 ++---- .../src/arrays/fixed_size_list/tests/take.rs | 12 +-- vortex-array/src/arrays/list/tests.rs | 24 ++---- .../src/arrays/listview/conversion.rs | 21 ++--- .../src/arrays/listview/tests/basic.rs | 9 +-- .../src/arrays/listview/tests/filter.rs | 6 +- .../src/arrays/listview/tests/operations.rs | 9 +-- .../src/arrays/listview/tests/take.rs | 6 +- vortex-array/src/arrays/masked/tests.rs | 3 +- .../src/arrays/patched/compute/take.rs | 15 ++-- .../src/arrays/primitive/array/patch.rs | 4 +- .../src/arrays/primitive/compute/cast.rs | 18 ++--- .../src/arrays/primitive/compute/fill_null.rs | 12 +-- .../src/arrays/primitive/compute/take/mod.rs | 4 +- .../src/arrays/primitive/vtable/mod.rs | 4 +- .../src/arrays/scalar_fn/vtable/operations.rs | 12 +-- vortex-array/src/arrays/slice/vtable.rs | 4 +- .../src/arrays/struct_/compute/mod.rs | 6 +- .../src/arrays/struct_/compute/rules.rs | 9 +-- vortex-array/src/arrays/struct_/tests.rs | 9 +-- .../src/arrays/varbin/compute/compare.rs | 7 +- .../src/arrays/varbin/compute/filter.rs | 15 ++-- .../src/arrays/varbin/compute/take.rs | 4 +- vortex-array/src/arrays/varbin/tests.rs | 7 +- .../src/arrays/varbin/vtable/canonical.rs | 3 +- vortex-array/src/arrays/varbinview/compact.rs | 27 +++---- vortex-array/src/arrays/varbinview/tests.rs | 10 +-- .../src/arrays/varbinview/vtable/mod.rs | 4 +- vortex-array/src/arrays/variant/mod.rs | 6 +- vortex-array/src/builders/bool.rs | 3 +- vortex-array/src/builders/decimal.rs | 3 +- vortex-array/src/builders/dict/primitive.rs | 6 +- vortex-array/src/builders/extension.rs | 4 +- vortex-array/src/builders/listview.rs | 3 +- vortex-array/src/builders/primitive.rs | 6 +- vortex-array/src/builders/struct_.rs | 3 +- vortex-array/src/builders/varbinview.rs | 9 +-- .../src/compute/conformance/filter.rs | 3 +- vortex-array/src/lib.rs | 8 ++ vortex-array/src/patches.rs | 63 +++++---------- vortex-array/src/scalar_fn/fns/between/mod.rs | 3 +- .../src/scalar_fn/fns/binary/boolean.rs | 6 +- .../src/scalar_fn/fns/binary/compare.rs | 18 ++--- vortex-array/src/scalar_fn/fns/binary/mod.rs | 15 ++-- .../src/scalar_fn/fns/binary/numeric.rs | 24 ++---- vortex-array/src/scalar_fn/fns/byte_length.rs | 6 +- vortex-array/src/scalar_fn/fns/case_when.rs | 69 ++++++---------- vortex-array/src/scalar_fn/fns/dynamic.rs | 13 +--- .../src/scalar_fn/fns/fill_null/mod.rs | 10 +-- .../src/scalar_fn/fns/list_contains/mod.rs | 24 ++---- vortex-array/src/scalar_fn/fns/merge.rs | 4 +- vortex-array/src/scalar_fn/fns/pack.rs | 7 +- vortex-array/src/scalar_fn/fns/zip/mod.rs | 9 +-- vortex-array/src/stats/expr.rs | 48 ++++-------- vortex-btrblocks/src/canonical_compressor.rs | 24 ++---- vortex-btrblocks/src/schemes/float/tests.rs | 3 +- .../schemes/integer/scheme_selection_tests.rs | 3 +- vortex-btrblocks/src/schemes/integer/tests.rs | 9 +-- vortex-compressor/src/sample.rs | 4 +- vortex-cuda/src/dynamic_dispatch/mod.rs | 78 +++++++------------ vortex-cuda/src/hybrid_dispatch/mod.rs | 21 ++--- vortex-cuda/src/kernel/arrays/constant.rs | 10 +-- vortex-cuda/src/kernel/arrays/dict.rs | 67 ++++++---------- vortex-cuda/src/kernel/encodings/alp.rs | 21 ++--- vortex-cuda/src/kernel/encodings/bitpacked.rs | 30 +++---- .../src/kernel/encodings/date_time_parts.rs | 10 +-- .../kernel/encodings/decimal_byte_parts.rs | 4 +- vortex-cuda/src/kernel/encodings/for_.rs | 6 +- vortex-cuda/src/kernel/encodings/fsst.rs | 7 +- vortex-cuda/src/kernel/encodings/runend.rs | 16 ++-- vortex-cuda/src/kernel/encodings/sequence.rs | 4 +- vortex-cuda/src/kernel/encodings/zigzag.rs | 4 +- vortex-cuda/src/kernel/encodings/zstd.rs | 13 +--- .../src/kernel/encodings/zstd_buffers.rs | 6 +- vortex-cuda/src/kernel/filter/decimal.rs | 7 +- vortex-cuda/src/kernel/filter/primitive.rs | 7 +- vortex-cuda/src/kernel/filter/varbinview.rs | 4 +- vortex-cuda/src/kernel/patches/types.rs | 4 +- vortex-duckdb/src/convert/vector.rs | 3 +- vortex-ffi/src/scan.rs | 12 +-- vortex-file/src/tests.rs | 9 +-- vortex-file/tests/test_write_table.rs | 3 +- vortex-ipc/src/iterator.rs | 4 +- vortex-ipc/src/stream.rs | 10 +-- vortex-layout/src/layouts/chunked/reader.rs | 4 +- vortex-layout/src/layouts/dict/reader.rs | 15 ++-- vortex-layout/src/layouts/flat/reader.rs | 10 +-- vortex-layout/src/layouts/row_idx/mod.rs | 10 +-- vortex-layout/src/layouts/struct_/reader.rs | 9 +-- vortex-layout/src/layouts/zoned/reader.rs | 4 +- vortex-python/src/arrays/range_to_sequence.rs | 4 +- vortex-test/compat-gen/src/check.rs | 4 +- .../src/fixtures/arrays/synthetic/mod.rs | 5 +- 124 files changed, 460 insertions(+), 957 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 51394873ad5..597615dbad9 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -236,8 +236,7 @@ mod tests { #[test] fn listview_matches_reference_unsigned() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::new(buffer![1u32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::U32, NonNullable); @@ -257,8 +256,7 @@ mod tests { #[test] fn listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Offsets are not in group order and a group is null: the group validity must be indexed by // group index, not by element offset. let elements = @@ -280,8 +278,7 @@ mod tests { #[test] fn listview_interior_and_full_nulls() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Group 1 has an interior null, group 2 is entirely null, group 3 is empty. let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) @@ -303,8 +300,7 @@ mod tests { #[test] fn listview_overflow_group_is_null() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I64, NonNullable); @@ -388,8 +384,7 @@ mod tests { #[test] fn fixed_size_overflow_and_nan() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // FixedSize path: first group overflows -> null sum, second sums normally. let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 5262cc72a34..26d3d4c19b5 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -586,8 +586,7 @@ mod tests { #[test] fn grouped_sum_fixed_size_list() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; @@ -602,8 +601,7 @@ mod tests { #[test] fn grouped_sum_with_null_elements() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) .into_array(); @@ -619,8 +617,7 @@ mod tests { #[test] fn grouped_sum_with_null_group() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) .into_array(); @@ -638,8 +635,7 @@ mod tests { #[test] fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; @@ -654,8 +650,7 @@ mod tests { #[test] fn grouped_sum_bool() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); let groups = FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; @@ -696,8 +691,7 @@ mod tests { #[test] fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 06a7ca538fb..830ff7451d9 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -489,7 +489,6 @@ mod tests { use vortex_buffer::buffer; use super::Array; - use crate::VortexSessionExecute; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -497,8 +496,7 @@ mod tests { #[test] fn typed_array_into_parts_roundtrips() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let expected = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); @@ -510,8 +508,7 @@ mod tests { #[test] fn typed_array_try_into_parts_requires_unique_owner() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let alias = array.clone(); diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 100487a607f..fa7239967ac 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -449,8 +449,7 @@ mod tests { #[test] fn patch_sliced_bools() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.slice(4..12).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut assertion_ctx); @@ -498,8 +497,7 @@ mod tests { #[test] fn slice_array_in_middle() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let sliced = arr.slice(4..12).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut assertion_ctx); @@ -507,8 +505,7 @@ mod tests { #[test] fn patch_bools_owned() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let buf_ptr = arr.to_bit_buffer().inner().as_ptr(); @@ -533,8 +530,7 @@ mod tests { #[test] fn patch_sliced_bools_offset() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.slice(4..15).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11]), &mut assertion_ctx); diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index 6c1b692006e..666c33a2067 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -303,15 +303,13 @@ mod tests { use super::*; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; #[test] fn filter_bool_test() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_iter([true, false, true]); @@ -325,8 +323,7 @@ mod tests { #[test] fn filter_bool_sparse_index_mask() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_indices(3, [0, 2]); @@ -340,8 +337,7 @@ mod tests { #[test] fn filter_bool_sparse_slice_mask() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_slices(3, vec![(0, 1), (2, 3)]); @@ -355,8 +351,7 @@ mod tests { #[test] fn filter_bool_sparse_buffer_mask() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_buffer(BitBuffer::from_iter([true, false, true])); diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 1981e3aad85..2894bde97f6 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -90,7 +90,6 @@ mod test { use crate::IntoArray as _; #[expect(deprecated)] use crate::ToCanonical as _; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; @@ -100,8 +99,7 @@ mod test { #[test] fn take_nullable() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let reference = BoolArray::from_iter(vec![ Some(false), Some(true), @@ -131,8 +129,7 @@ mod test { #[test] fn test_bool_array_take_with_null_out_of_bounds_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -150,8 +147,7 @@ mod test { #[test] fn test_non_null_bool_array_take_with_null_out_of_bounds_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -168,8 +164,7 @@ mod test { #[test] fn test_bool_array_take_all_null_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -185,8 +180,7 @@ mod test { #[test] fn test_non_null_bool_array_take_all_null_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], diff --git a/vortex-array/src/arrays/bool/patch.rs b/vortex-array/src/arrays/bool/patch.rs index 840e8388b07..188bcfd39d0 100644 --- a/vortex-array/src/arrays/bool/patch.rs +++ b/vortex-array/src/arrays/bool/patch.rs @@ -47,14 +47,12 @@ mod tests { use vortex_buffer::BitBuffer; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; #[test] fn patch_sliced_bools() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.into_array().slice(4..12).unwrap(); let expected = BoolArray::from_iter([true; 8]); @@ -63,8 +61,7 @@ mod tests { #[test] fn patch_sliced_bools_offset() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.into_array().slice(4..15).unwrap(); let expected = BoolArray::from_iter([true; 11]); diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 8d390c4b7ff..d2bfef2ab26 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -212,7 +212,6 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -220,8 +219,7 @@ mod tests { #[test] fn test_nullable_bool_serde_roundtrip() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = BoolArray::from_iter([Some(true), None, Some(false), None]); let dtype = array.dtype().clone(); let len = array.len(); diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index 5a7f0334ae6..f8927dfbaec 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -30,7 +30,6 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; @@ -47,8 +46,7 @@ mod tests { #[test] fn test_slice() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); #[expect(deprecated)] let sliced_arr = arr.into_array().slice(1..4).unwrap().to_bool(); diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 4b828c7262b..503771ac9d5 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -291,8 +291,7 @@ mod test { #[test] fn test_rechunk_one_chunk() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -306,8 +305,7 @@ mod test { #[test] fn test_rechunk_two_chunks() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array(), buffer![5u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -322,8 +320,7 @@ mod test { #[test] fn test_rechunk_tiny_target_chunks() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2, 3].into_array(), @@ -342,8 +339,7 @@ mod test { #[test] fn test_rechunk_with_too_big_chunk() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2].into_array(), diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index b74a04f8ba2..90c7f8c0a95 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -35,7 +35,6 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -47,8 +46,7 @@ mod test { #[test] fn test_cast_chunked() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr0 = buffer![0u32, 1].into_array(); let arr1 = buffer![2u32, 3].into_array(); diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 3ab6ed147db..7dcc9a57a07 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -128,7 +128,6 @@ mod test { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; @@ -142,8 +141,7 @@ mod test { #[test] fn test_take() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -161,8 +159,7 @@ mod test { #[test] fn test_take_nullable_values() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let a = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -180,8 +177,7 @@ mod test { #[test] fn test_take_nullable_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -202,8 +198,7 @@ mod test { #[test] fn test_take_nullable_struct() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let struct_array = StructArray::try_new(FieldNames::default(), vec![], 100, Validity::NonNullable) .unwrap(); @@ -229,8 +224,7 @@ mod test { #[test] fn test_empty_take() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -251,8 +245,7 @@ mod test { #[test] fn test_take_shuffled_indices() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let c0 = buffer![0i32, 1, 2].into_array(); let c1 = buffer![3i32, 4, 5].into_array(); let c2 = buffer![6i32, 7, 8].into_array(); @@ -321,8 +314,7 @@ mod test { #[test] fn test_take_null_indices() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let c0 = buffer![10i32, 20, 30].into_array(); let c1 = buffer![40i32, 50, 60].into_array(); let arr = ChunkedArray::try_new( diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index bd1a21cf74e..255349adc68 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -203,8 +203,7 @@ fn execute_path_nested_chunked_dict_of_dict_into_canonical() { #[test] fn with_slot_rewrites_chunk_and_offsets() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = chunked_array().into_array(); let replacement = buffer![10u64, 11, 12].into_array(); @@ -237,8 +236,7 @@ fn with_slot_rejects_len_mismatch() { #[test] fn slice_middle() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); assert_arrays_eq!( chunked_array().slice(2..5).unwrap(), PrimitiveArray::from_iter([3u64, 4, 5]), @@ -248,8 +246,7 @@ fn slice_middle() { #[test] fn slice_begin() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); assert_arrays_eq!( chunked_array().slice(1..3).unwrap(), PrimitiveArray::from_iter([2u64, 3]), @@ -259,8 +256,7 @@ fn slice_begin() { #[test] fn slice_aligned() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); assert_arrays_eq!( chunked_array().slice(3..6).unwrap(), PrimitiveArray::from_iter([4u64, 5, 6]), @@ -270,8 +266,7 @@ fn slice_aligned() { #[test] fn slice_many_aligned() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); assert_arrays_eq!( chunked_array().slice(0..6).unwrap(), PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]), @@ -281,8 +276,7 @@ fn slice_many_aligned() { #[test] fn slice_end() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); assert_arrays_eq!( chunked_array().slice(7..8).unwrap(), PrimitiveArray::from_iter([8u64]), @@ -292,8 +286,7 @@ fn slice_end() { #[test] fn slice_exactly_end() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); assert_arrays_eq!( chunked_array().slice(6..9).unwrap(), PrimitiveArray::from_iter([7u64, 8, 9]), @@ -311,8 +304,7 @@ fn slice_empty() { #[test] fn scalar_at_empty_children_both_sides() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -333,8 +325,7 @@ fn scalar_at_empty_children_both_sides() { #[test] fn scalar_at_empty_children_trailing() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -354,8 +345,7 @@ fn scalar_at_empty_children_trailing() { #[test] fn scalar_at_empty_children_leading() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 1df291aa5dc..8960ff47540 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -30,7 +30,6 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -58,8 +57,7 @@ mod tests { #[case::end(7..8, &[8u64])] #[case::exactly_end(6..9, &[7u64, 8, 9])] fn slice(#[case] range: Range, #[case] expected: &[u64]) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); assert_arrays_eq!( chunked_array().slice(range).unwrap(), PrimitiveArray::from_iter(expected.iter().copied()), @@ -77,8 +75,7 @@ mod tests { #[test] fn scalar_at_empty_children_both_sides() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -99,8 +96,7 @@ mod tests { #[test] fn scalar_at_empty_children_trailing() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -120,8 +116,7 @@ mod tests { #[test] fn scalar_at_empty_children_leading() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index eb2debc0594..f7f6b6c51f2 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -87,8 +87,7 @@ mod tests { #[test] fn take_nullable_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take( @@ -126,8 +125,7 @@ mod tests { #[test] fn take_all_valid_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take(PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array()) diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 8e9b78c9584..bcd5ca99d6b 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -384,8 +384,7 @@ mod tests { #[test] fn test_canonicalize_const_str() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let const_array = ConstantArray::new("four".to_string(), 4); let expected = VarBinArray::from(vec!["four", "four", "four", "four"]); diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index 0ca0a8abfbf..084b6e517bf 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -10,7 +10,6 @@ use crate::EqMode; use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; -use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::datetime::TemporalData; use crate::assert_arrays_eq; @@ -27,8 +26,7 @@ use crate::validity::Validity; macro_rules! test_temporal_roundtrip { ($prim:ty, $constructor:expr, $unit:expr) => {{ - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = buffer![100 as $prim].into_array(); let temporal: TemporalData = $constructor(array, $unit); @@ -149,8 +147,7 @@ test_fail_case!( // We test Timestamp explicitly to avoid the macro getting too complex. #[test] fn test_timestamp() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let ts = buffer![100i64].into_array(); let ts_array = ts.into_array(); diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index b28a97933c0..48d71545bcc 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -104,8 +104,7 @@ mod tests { #[test] fn fill_null_leading_none() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( [None, Some(800i128), None, Some(1000i128), None], @@ -142,8 +141,7 @@ mod tests { #[test] fn fill_null_all_none() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( @@ -171,8 +169,7 @@ mod tests { /// fill_null with a value that overflows the array's storage type should upcast the array. #[test] fn fill_null_overflow_upcasts() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let decimal_dtype = DecimalDType::new(3, 0); let arr = DecimalArray::from_option_iter([None, Some(10i8), None], decimal_dtype); // i8 max is 127, so 200 doesn't fit — the array should be widened to i16. @@ -195,8 +192,7 @@ mod tests { #[test] fn fill_null_non_nullable() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::new( diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 8ddd8c84b38..faffbf3433d 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -53,7 +53,6 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -63,8 +62,7 @@ mod tests { #[test] fn test_take() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![10i128, 11i128, 12i128, 13i128], @@ -81,8 +79,7 @@ mod tests { #[test] fn test_take_null_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![i128::MAX, 11i128, 12i128, 13i128], diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index 1e83d022f4c..bc0c3295a80 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -209,7 +209,6 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; - use crate::VortexSessionExecute; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; @@ -249,8 +248,7 @@ mod tests { #[test] fn test_nullable_decimal_serde_roundtrip() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = DecimalArray::new( buffer![1234567i32, 0i32, -9999999i32], DecimalDType::new(7, 3), diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index c97b137c514..bc8649c1962 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -462,8 +462,7 @@ mod test { #[test] fn test_dict_array_from_primitive_chunks() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let len = 2; let chunk_count = 2; let array = make_dict_primitive_chunks::(len, 2, chunk_count); diff --git a/vortex-array/src/arrays/dict/compute/cast.rs b/vortex-array/src/arrays/dict/compute/cast.rs index ed881d5db53..a86f197d629 100644 --- a/vortex-array/src/arrays/dict/compute/cast.rs +++ b/vortex-array/src/arrays/dict/compute/cast.rs @@ -71,8 +71,7 @@ mod tests { #[test] fn test_cast_dict_to_wider_type() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![1i32, 2, 3, 2, 1].into_array(); let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); @@ -112,8 +111,7 @@ mod tests { #[test] fn test_cast_dict_allvalid_to_nonnullable_and_back() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create an AllValid dict array (no nulls) let values = buffer![10i32, 20, 30, 40].into_array(); let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); @@ -195,8 +193,7 @@ mod tests { #[test] fn test_cast_dict_with_unreferenced_null_values_to_nonnullable() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use crate::arrays::DictArray; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/dict/compute/fill_null.rs b/vortex-array/src/arrays/dict/compute/fill_null.rs index 57f5ab8c34b..70a6cef1716 100644 --- a/vortex-array/src/arrays/dict/compute/fill_null.rs +++ b/vortex-array/src/arrays/dict/compute/fill_null.rs @@ -108,8 +108,7 @@ mod tests { #[test] fn nullable_codes_fill_in_values() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::new( buffer![0u32, 1, 2], diff --git a/vortex-array/src/arrays/dict/compute/like.rs b/vortex-array/src/arrays/dict/compute/like.rs index dfc4e5e8803..8319c28c2e3 100644 --- a/vortex-array/src/arrays/dict/compute/like.rs +++ b/vortex-array/src/arrays/dict/compute/like.rs @@ -56,7 +56,6 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::DictArray; use crate::arrays::VarBinArray; @@ -69,8 +68,7 @@ mod tests { #[test] fn like_reduce_dict() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let dict = DictArray::try_new( buffer![0u8, 1, 0, 2].into_array(), VarBinArray::from(vec!["hello", "world", "help"]).into_array(), diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index adc3470a3dc..421642a8a18 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -86,8 +86,7 @@ mod test { #[test] fn canonicalise_nullable_primitive() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values: Vec> = (0..65) .map(|i| match i % 3 { 0 => Some(42), @@ -112,8 +111,7 @@ mod test { #[test] fn canonicalise_non_nullable_primitive_32_unique_values() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let unique_values: Vec = (0..32).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 32])); @@ -130,8 +128,7 @@ mod test { #[test] fn canonicalise_non_nullable_primitive_100_unique_values() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let unique_values: Vec = (0..100).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 100])); @@ -186,8 +183,7 @@ mod test { #[test] fn compare_sliced_dict() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use crate::arrays::BoolArray; let sliced = sliced_dict_array(); let compared = sliced diff --git a/vortex-array/src/arrays/dict/compute/slice.rs b/vortex-array/src/arrays/dict/compute/slice.rs index ba92370f16c..485346bf5d8 100644 --- a/vortex-array/src/arrays/dict/compute/slice.rs +++ b/vortex-array/src/arrays/dict/compute/slice.rs @@ -91,7 +91,6 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::dict::compute::slice::ConstantArray; @@ -103,8 +102,7 @@ mod tests { #[test] fn slice_constant_valid_code() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let dict = DictArray::new( ConstantArray::new(1u8, 5).into_array(), buffer![10i32, 20, 30].into_array(), @@ -117,8 +115,7 @@ mod tests { #[test] fn slice_constant_null_code() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let dict = DictArray::new( ConstantArray::new(Scalar::null(DType::Primitive(PType::U8, Nullable)), 5).into_array(), buffer![10i32, 20, 30].into_array(), diff --git a/vortex-array/src/arrays/dict/tests.rs b/vortex-array/src/arrays/dict/tests.rs index ae0454db757..5cbb08d5627 100644 --- a/vortex-array/src/arrays/dict/tests.rs +++ b/vortex-array/src/arrays/dict/tests.rs @@ -5,7 +5,6 @@ use vortex_buffer::buffer; use super::DictArray; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; @@ -15,8 +14,7 @@ use crate::validity::Validity; #[test] fn test_scalar_at_null_code() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::from_option_iter(vec![None, Some(0u32), None]).into_array(), buffer![1i32].into_array(), diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index 5b50dbacc85..d7378f7f095 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -121,8 +121,7 @@ mod tests { #[test] fn cast_timestamp_to_i64() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let ext_dtype = Timestamp::new_with_tz( TimeUnit::Nanoseconds, Some("UTC".into()), diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index 0c30352b564..d2626460ba1 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -124,7 +124,6 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -150,8 +149,7 @@ mod test { #[test] fn filter_fixed_size_list_selects_correct_lists() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let array = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -168,8 +166,7 @@ mod test { #[test] fn filter_degenerate_list_size_zero() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::empty::(Nullability::NonNullable); let array = FixedSizeListArray::new(elements.into_array(), 0, Validity::NonNullable, 5); @@ -185,8 +182,7 @@ mod test { #[test] fn filter_nested_fixed_size_lists() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Inner lists of size 2, outer lists of size 2 (so 2 outer lists, each with 2 inner lists). let inner_elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); let inner_fsl = FixedSizeListArray::new(inner_elements, 2, Validity::NonNullable, 4); diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index ce0fb1cd1d8..6f0f54edd66 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -135,8 +135,7 @@ mod test { #[test] fn filter_listview_selects_correct_lists() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // 3 lists: [10,20], [30,40], [50,60] let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let offsets = buffer![0u32, 2, 4].into_array(); @@ -164,8 +163,7 @@ mod test { #[test] fn test_filter_preserves_unreferenced_elements() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -200,8 +198,7 @@ mod test { #[test] fn test_filter_with_gaps() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index 06ac484bafa..71f573a5881 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -49,7 +49,6 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -87,8 +86,7 @@ mod test { #[test] fn filter_struct_selects_correct_rows() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = StructArray::try_new( ["x", "y"].into(), vec![ @@ -119,8 +117,7 @@ mod test { #[test] fn filter_empty_struct() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let mask = Mask::from_iter([ @@ -135,8 +132,7 @@ mod test { #[test] fn filter_empty_struct_with_empty_filter() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); let filtered = struct_arr.filter(Mask::from_iter::<[bool; 0]>([])).unwrap(); diff --git a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs index 5615753e90c..72c4dbda15b 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs @@ -8,7 +8,6 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -52,8 +51,7 @@ fn test_filter_degenerate_list_size_zero( #[case] mask_values: Vec, #[case] expected_len: usize, ) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let new_validity = if matches!(validity, Validity::NonNullable) { Validity::NonNullable } else { @@ -82,8 +80,7 @@ fn test_filter_degenerate_list_size_zero( #[test] fn test_filter_with_nulls() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4), Some(5), Some(6)]); let validity = Validity::from_iter([true, false, true]); @@ -111,8 +108,7 @@ fn test_filter_with_nulls() { #[test] fn test_filter_all_null_array() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create an array where all elements are null. let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let validity = Validity::AllInvalid; @@ -129,8 +125,7 @@ fn test_filter_all_null_array() { #[test] fn test_filter_nested_fixed_size_lists() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create nested fixed-size lists: FSL>. // Inner lists are of size 2, outer lists are of size 3. // So we have 2 outer lists, each containing 3 inner lists, each containing 2 i32s. @@ -256,8 +251,7 @@ fn create_fsl_empty() -> ArrayRef { #[test] fn test_filter_all_null_various_list_sizes() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test filtering with all-null arrays of different list sizes. // The implementation returns ConstantArray only when validity_mask() is Mask::AllFalse. @@ -296,8 +290,7 @@ fn test_filter_all_null_various_list_sizes() { #[test] fn test_mask_expansion_threshold_boundary() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test with list_size == 8 (the FSL_SPARSE_MASK_LIST_SIZE_THRESHOLD). let list_size = 8u32; let num_lists = 100; @@ -374,8 +367,7 @@ fn test_mask_expansion_threshold_boundary() { // Test FSL-specific behavior with very large list sizes. #[test] fn test_filter_large_list_size() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test with list_size=100, which is significantly larger than typical use cases. let list_size = 100u32; let num_lists = 5; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 5875b380267..8f6713fe98c 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -40,8 +40,7 @@ fn test_take_fsl_conformance(#[case] fsl: FixedSizeListArray) { #[test] fn test_take_basic_smoke_test() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -109,8 +108,7 @@ fn test_take_degenerate_lists( #[test] fn test_take_large_list_size() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = buffer![0i32..300].into_array(); let fsl = FixedSizeListArray::new(elements, 100, Validity::NonNullable, 3); @@ -125,8 +123,7 @@ fn test_take_large_list_size() { #[test] fn test_take_fsl_with_null_indices_preserves_elements() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -173,8 +170,7 @@ fn test_element_index_overflow( #[case] indices: ArrayRef, #[case] expected: FixedSizeListArray, ) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let result = fsl.take(indices).unwrap(); assert_arrays_eq!(result, expected, &mut assertion_ctx); } diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 26e2a03dbbf..48a04f8522e 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -113,8 +113,7 @@ fn test_simple_list_filter() { #[test] fn test_list_filter_dense_mask() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test filtering with a dense mask (high density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -146,8 +145,7 @@ fn test_list_filter_dense_mask() { #[test] fn test_list_filter_sparse_mask() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test filtering with a sparse mask (low density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -181,8 +179,7 @@ fn test_list_filter_sparse_mask() { #[test] fn test_list_filter_empty_lists() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test filtering arrays that contain empty lists. let elements = buffer![0..10].into_array(); let offsets = buffer![0, 0, 3, 3, 7, 10, 10].into_array(); // Lists at indices 0, 2, 5 are empty. @@ -239,8 +236,7 @@ fn test_list_filter_with_nulls() { #[test] fn test_list_filter_all_true() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test filtering with an all-true mask. let elements = buffer![0..20].into_array(); let offsets = buffer![0, 5, 10, 15, 20].into_array(); @@ -283,8 +279,7 @@ fn test_list_filter_all_false() { #[test] fn test_list_filter_single_element() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test filtering to keep only one element. let elements = buffer![0..50].into_array(); let offsets = buffer![0, 10, 20, 30, 40, 50].into_array(); @@ -313,8 +308,7 @@ fn test_list_filter_single_element() { #[test] fn test_list_filter_alternating_pattern() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test filtering with an alternating pattern. let elements = buffer![0..60].into_array(); let offsets = buffer![0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60].into_array(); @@ -350,8 +344,7 @@ fn test_list_filter_alternating_pattern() { #[test] fn test_list_filter_variable_sizes() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test filtering lists with highly variable sizes. let elements = buffer![0..100].into_array(); let offsets = buffer![0, 1, 2, 5, 10, 20, 35, 60, 100].into_array(); @@ -588,8 +581,7 @@ fn create_list_of_lists_nullable(data: OptVec>>) -> ListArray #[test] #[expect(clippy::cognitive_complexity)] fn test_list_of_lists() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let data = vec![ Some(vec![Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]), Some(vec![Some(vec![Some(4), Some(5), Some(6)])]), diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index ae85225c440..b3f8807821e 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -348,8 +348,7 @@ mod tests { #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let list_view = create_basic_listview(); let list_array = list_from_list_view(list_view.clone(), &mut SESSION.create_execution_ctx())?; @@ -425,8 +424,7 @@ mod tests { #[test] fn test_non_zero_copy_listview_to_list() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create ListViewArray with overlapping lists (not zero-copyable). let list_view = create_overlapping_listview(); let list_array = @@ -560,8 +558,7 @@ mod tests { #[test] fn test_recursive_simple_listview() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let list_view = create_basic_listview(); let result = recursive_list_from_list_view( list_view.clone().into_array(), @@ -575,8 +572,7 @@ mod tests { #[test] fn test_recursive_nested_listview() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let inner_elements = buffer![1i32, 2, 3].into_array(); let inner_offsets = buffer![0u32, 2].into_array(); let inner_sizes = buffer![2u32, 1].into_array(); @@ -614,8 +610,7 @@ mod tests { #[test] fn test_recursive_struct_with_listview_fields() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let listview_field = create_basic_listview().into_array(); let primitive_field = buffer![10i32, 20, 30, 40].into_array(); @@ -638,8 +633,7 @@ mod tests { #[test] fn test_recursive_fixed_size_list_with_listview_elements() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let lv1_elements = buffer![1i32, 2].into_array(); let lv1_offsets = buffer![0u32].into_array(); let lv1_sizes = buffer![2u32].into_array(); @@ -685,8 +679,7 @@ mod tests { #[test] fn test_recursive_deep_nesting() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let innermost_elements = buffer![1i32, 2, 3].into_array(); let innermost_offsets = buffer![0u32, 2].into_array(); let innermost_sizes = buffer![2u32, 1].into_array(); diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index c596b108f07..dddcf169ace 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -26,8 +26,7 @@ use crate::validity::Validity; #[test] fn test_basic_listview_comprehensive() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Comprehensive test for basic ListView functionality including scalar_at. // Logical lists: [[1,2,3], [4,5], [6,7,8,9]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -84,8 +83,7 @@ fn test_basic_listview_comprehensive() { #[test] fn test_out_of_order_offsets() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // ListView-specific: Tests that offsets can be non-sequential and out-of-order. // Logical lists: [[7,8,9], [1,2,3], [4,5,6]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -169,8 +167,7 @@ fn test_from_list_array() -> VortexResult<()> { #[case::constant_offsets(false, true)] // Varying sizes, constant offsets #[case::both_constant(true, true)] // Both constant fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_offsets: bool) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Logical lists vary by case: // - constant_sizes: [[1,2,3], [4,5,6], [7,8,9]] (size 3 each, varying offsets) // - constant_offsets: [[1,2,3], [1,2], [1]] (all start at 0, varying sizes) diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index 83504a1a225..e1568d14d4b 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -39,8 +39,7 @@ fn test_filter_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_filter_preserves_unreferenced_elements() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -74,8 +73,7 @@ fn test_filter_preserves_unreferenced_elements() { #[test] fn test_filter_with_gaps() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index ba28419e3de..f488b285e6c 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -84,8 +84,7 @@ fn test_slice_comprehensive() { #[test] fn test_slice_out_of_order() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // ListView-specific: Test slicing with out-of-order offsets. // Logical lists: [[70,80], [10,20,30], [40,50,60], [90], [30]] let elements = buffer![10i32, 20, 30, 40, 50, 60, 70, 80, 90].into_array(); @@ -395,8 +394,7 @@ fn test_cast_large_dataset() { #[test] fn test_zip_widens_false_element_nullability() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // [[1, 2], [3], [4]] let if_true = ListViewArray::new( buffer![1i32, 2, 3, 4].into_array(), @@ -442,8 +440,7 @@ fn test_zip_widens_false_element_nullability() -> VortexResult<()> { #[test] fn test_zip_widens_true_element_nullability() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // [[1, null], [3], [4]] let if_true = ListViewArray::new( PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(), diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 7c86cb360fb..1c2fd05b3a4 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -40,8 +40,7 @@ fn test_take_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_take_preserves_unreferenced_elements() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // ListView-specific: Test that take preserves the entire elements array // even when taking only a subset of lists. let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -73,8 +72,7 @@ fn test_take_preserves_unreferenced_elements() { #[test] fn test_take_with_gaps() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // ListView-specific: Test with gaps in elements array. // Elements with gaps (999 values are "gaps" between used ranges). let elements = buffer![1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12].into_array(); diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index fb5f2e2727e..19d35bfa4c6 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -95,8 +95,7 @@ fn test_masked_child_with_validity() { #[test] fn test_masked_child_all_valid() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // When validity is AllValid, masked_child should invert to AllInvalid. let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); diff --git a/vortex-array/src/arrays/patched/compute/take.rs b/vortex-array/src/arrays/patched/compute/take.rs index 248ec8dc616..078bfc2b4f0 100644 --- a/vortex-array/src/arrays/patched/compute/take.rs +++ b/vortex-array/src/arrays/patched/compute/take.rs @@ -168,8 +168,7 @@ mod tests { #[test] fn test_take_basic() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -186,8 +185,7 @@ mod tests { #[test] fn test_take_sliced() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = make_patched_array(&[0; 10], &[1, 3], &[100, 200], 2..10)?; let indices = buffer![0u32, 1, 2, 3, 7].into_array(); @@ -202,8 +200,7 @@ mod tests { #[test] fn test_take_out_of_order() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -220,8 +217,7 @@ mod tests { #[test] fn test_take_duplicates() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at index [2] with value [99] let array = make_patched_array(&[0; 5], &[2], &[99], 0..5)?; @@ -242,8 +238,7 @@ mod tests { #[test] fn test_take_with_null_indices() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use crate::arrays::BoolArray; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index 3161b2562e5..32c5964f3bc 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -137,7 +137,6 @@ mod tests { use super::*; #[expect(deprecated)] use crate::ToCanonical as _; - use crate::VortexSessionExecute; use crate::assert_arrays_eq; use crate::validity::Validity; @@ -174,8 +173,7 @@ mod tests { #[test] fn patch_sliced() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid); let sliced = input.slice(2..8).unwrap(); #[expect(deprecated)] diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 64ec5542a7c..2657b68f73c 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -286,8 +286,7 @@ mod test { #[test] fn cast_u32_u8() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); // cast from u32 to u8 @@ -359,8 +358,7 @@ mod test { #[test] fn cast_u32_f32() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); #[expect(deprecated)] let u8arr = arr.cast(PType::F32.into()).unwrap().to_primitive(); @@ -402,8 +400,7 @@ mod test { #[test] fn cast_with_invalid_nulls() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = PrimitiveArray::new( buffer![-1i32, 0, 10], Validity::from_iter([false, true, true]), @@ -433,8 +430,7 @@ mod test { /// buffer without allocation (pointer identity). #[test] fn cast_same_width_int_reinterprets_buffer() -> vortex_error::VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let src = PrimitiveArray::from_iter([0u32, 10, 100]); let src_ptr = src.as_slice::().as_ptr(); @@ -484,8 +480,7 @@ mod test { /// not prevent the cast from succeeding. #[test] fn cast_same_width_int_nullable_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // The null position holds u32::MAX which doesn't fit in i32, but it's // masked as invalid so the cast should still succeed via reinterpret. let arr = PrimitiveArray::new( @@ -507,8 +502,7 @@ mod test { #[test] fn cast_u32_to_u8_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = PrimitiveArray::new( buffer![1000u32, 10u32, 42u32], Validity::from_iter([false, true, true]), diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index 93d815a9365..0a1d5c01ff0 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -64,8 +64,7 @@ mod test { #[test] fn fill_null_leading_none() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = PrimitiveArray::from_option_iter([None, Some(8u8), None, Some(10), None]); #[expect(deprecated)] let p = arr @@ -90,8 +89,7 @@ mod test { #[test] fn fill_null_all_none() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Option::::None, None, None, None, None]); #[expect(deprecated)] @@ -117,8 +115,7 @@ mod test { #[test] fn fill_null_nullable_non_null() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = PrimitiveArray::new( buffer![8u8, 10, 12, 14, 16], Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()), @@ -146,8 +143,7 @@ mod test { #[test] fn fill_null_non_nullable() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = buffer![8u8, 10, 12, 14, 16].into_array(); #[expect(deprecated)] let p = arr.fill_null(Scalar::from(255u8)).unwrap().to_primitive(); diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index ee0547244b2..7bc5ff453be 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -222,7 +222,6 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -230,8 +229,7 @@ mod tests { #[test] fn take_null_index_skips_out_of_bounds_value() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = PrimitiveArray::from_iter([10i32, 20, 30]); let indices = PrimitiveArray::new( buffer![1u64, 3], diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index d669f861fcf..e64ac0558bb 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -226,7 +226,6 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; - use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -235,8 +234,7 @@ mod tests { #[test] fn test_nullable_primitive_serde_roundtrip() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::new( buffer![1i32, 2, 3, 4], Validity::from_iter([true, false, true, false]), diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index eba022d1e5d..2e063ba85f5 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -76,8 +76,7 @@ mod tests { #[test] fn test_scalar_fn_add() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let lhs = buffer![1i32, 2, 3].into_array(); let rhs = buffer![10i32, 20, 30].into_array(); @@ -132,8 +131,7 @@ mod tests { #[test] fn test_scalar_fn_mul() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let lhs = buffer![2i32, 3, 4].into_array(); let rhs = buffer![5i32, 6, 7].into_array(); @@ -152,8 +150,7 @@ mod tests { #[test] fn test_scalar_fn_with_nullable() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let rhs = PrimitiveArray::new( buffer![10i32, 20, 30], @@ -180,8 +177,7 @@ mod tests { #[test] fn test_scalar_fn_comparison() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let lhs = buffer![1i32, 5, 3].into_array(); let rhs = buffer![2i32, 5, 1].into_array(); diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index efbfb029734..ef939d8ed1d 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -188,15 +188,13 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::assert_arrays_eq; #[test] fn test_slice_slice() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Slice(1..4, Slice(2..8, base)) combines to Slice(3..6, base) let arr = PrimitiveArray::from_iter(0i32..10).into_array(); let inner_slice = SliceArray::new(arr, 2..8).into_array(); diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index 9614d601616..ed4ebc782b0 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -40,8 +40,7 @@ mod tests { #[test] fn take_empty_struct() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let indices = PrimitiveArray::from_option_iter([Some(1), None]); @@ -95,8 +94,7 @@ mod tests { #[test] fn take_field_struct() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let struct_arr = StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).into_array())]) .unwrap(); diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 6a4b803096a..6f46515433e 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -165,8 +165,7 @@ mod tests { #[test] fn test_struct_cast_field_reorder() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Source: {a, b}, Target: {c, b, a} - reordered + new null field let source = StructArray::try_new( FieldNames::from(["a", "b"]), @@ -303,8 +302,7 @@ mod tests { #[test] fn cast_struct_drop_field() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Casting to a struct with a subset of fields should succeed. let source = StructArray::try_new( FieldNames::from(["a", "b", "c"]), @@ -350,8 +348,7 @@ mod tests { #[test] fn cast_struct_field_type_widening() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Casting struct fields to wider types (i32 -> i64). let source = StructArray::try_new( FieldNames::from(["val"]), diff --git a/vortex-array/src/arrays/struct_/tests.rs b/vortex-array/src/arrays/struct_/tests.rs index da198c7c1e2..947223bc7a7 100644 --- a/vortex-array/src/arrays/struct_/tests.rs +++ b/vortex-array/src/arrays/struct_/tests.rs @@ -24,8 +24,7 @@ use crate::validity::Validity; #[test] fn test_project() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = VarBinArray::from_vec( vec!["a", "b", "c", "d", "e"], @@ -68,8 +67,7 @@ fn test_project() { #[test] fn test_remove_column() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = PrimitiveArray::new(buffer![4u64, 5, 6, 7, 8], Validity::NonNullable); @@ -115,8 +113,7 @@ fn test_remove_column() { #[test] fn test_duplicate_field_names() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test that StructArray allows duplicate field names and returns the first match let field1 = buffer![1i32, 2, 3].into_array(); let field2 = buffer![10i32, 20, 30].into_array(); diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index b6055cfb144..22b61adaefb 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -254,7 +254,6 @@ mod tests { use vortex_buffer::ByteBuffer; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; @@ -290,8 +289,7 @@ mod tests { /// [`CompareKernel`]: super::CompareKernel #[test] fn varbin_i64_offsets_compare_constant() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -312,8 +310,7 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant_binary() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index ddddc3fbd37..8784540b996 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -240,8 +240,7 @@ mod test { #[test] fn filter_var_bin_test() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -268,8 +267,7 @@ mod test { #[test] fn filter_var_bin_slice_test() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -299,8 +297,7 @@ mod test { #[test] fn filter_var_bin_slice_null() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let bytes = [ b"one".as_slice(), b"two".as_slice(), @@ -343,8 +340,7 @@ mod test { #[test] fn filter_varbin_nulls() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let bytes = [b"".as_slice(), b"two".as_slice(), b"two".as_slice()] .into_iter() .flat_map(|x| x.iter().cloned()) @@ -372,8 +368,7 @@ mod test { #[test] fn filter_varbin_all_null() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let offsets = buffer![0, 0, 0, 0].into_array(); let validity = Validity::Array(BoolArray::from_iter([false, false, false]).into_array()); let arr = VarBinArray::try_new( diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 3345eb6d952..b0e5d6ec731 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -256,7 +256,6 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::compute::take::PrimitiveArray; @@ -305,8 +304,7 @@ mod tests { #[test] fn test_take_overflow() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let scream = std::iter::once("a").cycle().take(128).collect::(); let bytes = ByteBuffer::copy_from(scream.as_bytes()); let offsets = buffer![0u8, 128u8].into_array(); diff --git a/vortex-array/src/arrays/varbin/tests.rs b/vortex-array/src/arrays/varbin/tests.rs index eb5bb97e5de..5b1e9b0b529 100644 --- a/vortex-array/src/arrays/varbin/tests.rs +++ b/vortex-array/src/arrays/varbin/tests.rs @@ -8,7 +8,6 @@ use vortex_buffer::buffer; use crate::ArrayRef; use crate::IntoArray; -use crate::VortexSessionExecute; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; @@ -33,8 +32,7 @@ fn binary_array() -> ArrayRef { #[rstest] pub fn test_scalar_at(binary_array: ArrayRef) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); assert_arrays_eq!( binary_array, VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), @@ -44,8 +42,7 @@ pub fn test_scalar_at(binary_array: ArrayRef) { #[rstest] pub fn slice_array(binary_array: ArrayRef) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let binary_arr = binary_array.slice(1..2).unwrap(); assert_arrays_eq!( binary_arr, diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index f4739bf70a1..c0a69629945 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -96,8 +96,7 @@ mod tests { #[case(DType::Utf8(Nullability::NonNullable))] #[case(DType::Binary(Nullability::NonNullable))] fn test_canonical_varbin_unsliced(#[case] dtype: DType) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let varbin = VarBinArray::from_iter_nonnull(["foo", "bar", "baz"], dtype.clone()); #[expect(deprecated)] let canonical = varbin.as_array().to_varbinview(); diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index ebcceb55652..cdc1a480b57 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -202,8 +202,7 @@ mod tests { use crate::dtype::Nullability; #[test] fn test_optimize_compacts_buffers() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create a VarBinViewArray with some long strings that will create multiple buffers let original = VarBinViewArray::from_iter_nullable_str([ Some("short"), @@ -244,8 +243,7 @@ mod tests { #[test] fn test_optimize_with_long_strings() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create strings that are definitely longer than 12 bytes let long_string_1 = "this is definitely a very long string that exceeds the inline limit"; let long_string_2 = "another extremely long string that also needs external buffer storage"; @@ -282,8 +280,7 @@ mod tests { #[test] fn test_optimize_no_buffers() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create an array with only short strings (all inlined) let original = VarBinViewArray::from_iter_str(["a", "bb", "ccc", "dddd"]); @@ -300,8 +297,7 @@ mod tests { #[test] fn test_optimize_single_buffer() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create an array that naturally has only one buffer let str1 = "this is a long string that goes into a buffer"; let str2 = "another long string in the same buffer"; @@ -321,8 +317,7 @@ mod tests { #[test] fn test_selective_compaction_with_threshold_zero() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // threshold=0 should keep all buffers (no compaction) let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -350,8 +345,7 @@ mod tests { #[test] fn test_selective_compaction_with_high_threshold() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // threshold=1.0 should compact any buffer with waste let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -380,8 +374,7 @@ mod tests { #[test] fn test_selective_compaction_preserves_well_utilized_buffers() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create an array with multiple strings in one buffer (well-utilized) let str1 = "first long string that needs external buffer storage"; let str2 = "second long string also in buffer"; @@ -404,8 +397,7 @@ mod tests { #[test] fn test_selective_compaction_with_mixed_utilization() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create array with some long strings let strings: Vec = (0..10) .map(|i| { @@ -437,8 +429,7 @@ mod tests { #[test] fn test_slice_strategy_with_contiguous_range() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Create array with strings that will be in one buffer let strings: Vec = (0..20) .map(|i| format!("this is a long string number {} for slice test", i)) diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index 82bcb01b939..d8772540a87 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -3,15 +3,13 @@ #[expect(deprecated)] use crate::ToCanonical as _; -use crate::VortexSessionExecute; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; #[test] pub fn varbin_view() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]); assert_arrays_eq!( @@ -23,8 +21,7 @@ pub fn varbin_view() { #[test] pub fn slice_array() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) .slice(1..2) @@ -38,8 +35,7 @@ pub fn slice_array() { #[test] pub fn flatten_array() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["string1", "string2"]); #[expect(deprecated)] let var_bin = binary_arr.as_array().to_varbinview(); diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 9545aef2c9c..194b1145573 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -249,7 +249,6 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; - use crate::VortexSessionExecute; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -267,8 +266,7 @@ mod tests { let len = array.len(); let ctx = ArrayContext::empty(); - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let serialized = array .clone() .into_array() diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index ed89b75f4ca..aabdac06d95 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -379,8 +379,7 @@ mod tests { #[test] fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([( "a", @@ -408,8 +407,7 @@ mod tests { #[test] fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([ ( diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index 51fcfa2a288..80aad067648 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -221,8 +221,7 @@ mod tests { #[test] fn test_append_scalar() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10); // Test appending true value. diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index cd9f6bbc93a..0e3acd28774 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -345,8 +345,7 @@ mod tests { #[test] fn test_append_scalar() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use crate::scalar::Scalar; // Simply test that the builder accepts its own finish output via scalar. diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index a97f0327742..a6af40ebc57 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -173,8 +173,7 @@ mod test { #[test] fn encode_primitive() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = buffer![1, 1, 3, 3, 3].into_array(); let dict = dict_encode(&arr, &mut SESSION.create_execution_ctx()).unwrap(); @@ -187,8 +186,7 @@ mod test { #[test] fn encode_primitive_nulls() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let arr = PrimitiveArray::from_option_iter([ Some(1), Some(1), diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index ba844e3f178..dd3eba25a17 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -127,7 +127,6 @@ impl ArrayBuilder for ExtensionBuilder { #[cfg(test)] mod tests { use super::*; - use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -138,8 +137,7 @@ mod tests { #[test] fn test_append_scalar() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); let mut builder = ExtensionBuilder::new(ext_dtype.clone()); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 49b4338c466..846dfe3201c 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -529,8 +529,7 @@ mod tests { #[test] fn test_different_offset_size_types() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test u32 offsets with u8 sizes. let dtype: Arc = Arc::new(I32.into()); let mut builder = diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 79b96a7a762..568c41c010d 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -385,8 +385,7 @@ mod tests { /// buffer. #[test] fn test_multiple_uninit_ranges_correct_offsets() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); // First range. @@ -497,8 +496,7 @@ mod tests { /// This verifies the new simplified API without the redundant `len` parameter. #[test] fn test_copy_from_slice_with_offsets() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); let mut range = builder.uninit_range(6); diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index fb35bf2fdcf..a292b8c5e4f 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -267,8 +267,7 @@ mod tests { #[test] fn test_append_scalar() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use crate::scalar::Scalar; let dtype = DType::Struct( diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 6c653b48c3b..e3767fbccac 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -884,8 +884,7 @@ mod tests { #[test] fn test_utf8_builder() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); builder.append_value("Hello"); @@ -913,8 +912,7 @@ mod tests { #[test] fn test_utf8_builder_with_extend() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = { let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); @@ -994,8 +992,7 @@ mod tests { #[test] fn test_append_scalar() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use crate::scalar::Scalar; // Test with Utf8 builder. diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index bcbae669675..1eadc1671c1 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -62,8 +62,7 @@ pub fn create_runs_pattern(len: usize, run_length: usize) -> Vec { /// Tests that filtering with an all-true mask returns all elements unchanged fn test_all_filter(array: &ArrayRef) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let len = array.len(); let mask = Mask::new_true(len); let filtered = array diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 16d29c69efe..c94b3de161b 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -117,6 +117,14 @@ pub fn array_session() -> VortexSession { .build() } +/// Builds a fresh session via [`array_session`] and returns an [`ExecutionCtx`](crate::ExecutionCtx) +/// backed by it. +/// +/// Convenience for the common `array_session().create_execution_ctx()` pattern. +pub fn array_execution_ctx() -> ExecutionCtx { + array_session().create_execution_ctx() +} + // TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot invoke execute // from the new array encodings to support back-compat for legacy encodings. So we hold a session // here... diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 558b3329ed8..784f77d133f 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -1222,8 +1222,7 @@ mod test { #[test] fn test_filter() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 100, 0, @@ -1255,8 +1254,7 @@ mod test { #[test] fn take_with_nulls() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1305,8 +1303,7 @@ mod test { #[test] fn take_search_with_nulls_chunked() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1355,8 +1352,7 @@ mod test { #[test] fn take_search_chunked_multiple_chunks() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 2048, 0, @@ -1407,8 +1403,7 @@ mod test { #[test] fn take_search_chunked_interleaved() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 30, 0, @@ -1460,8 +1455,7 @@ mod test { #[test] fn test_slice() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1478,8 +1472,7 @@ mod test { #[test] fn doubly_sliced() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1521,8 +1514,7 @@ mod test { #[test] fn test_mask_all_false() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1573,8 +1565,7 @@ mod test { #[test] fn test_mask_partial() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1611,8 +1602,7 @@ mod test { #[test] fn test_mask_with_offset() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1647,8 +1637,7 @@ mod test { #[test] fn test_mask_nullable_values() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1701,8 +1690,7 @@ mod test { #[test] fn test_filter_keep_all() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1752,8 +1740,7 @@ mod test { #[test] fn test_filter_with_indices() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1784,8 +1771,7 @@ mod test { #[test] fn test_slice_full_range() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1811,8 +1797,7 @@ mod test { #[test] fn test_slice_partial() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1857,8 +1842,7 @@ mod test { #[test] fn test_slice_with_offset() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1966,8 +1950,7 @@ mod test { #[test] fn test_mask_boundary_patches() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test masking patches at array boundaries let patches = Patches::new( 10, @@ -2022,8 +2005,7 @@ mod test { #[test] fn test_mask_no_patches_removed() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test when no patches are masked let patches = Patches::new( 10, @@ -2057,8 +2039,7 @@ mod test { #[test] fn test_mask_single_patch() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test with a single patch let patches = Patches::new( 5, @@ -2091,8 +2072,7 @@ mod test { #[test] fn test_mask_contiguous_patches() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test with contiguous patches let patches = Patches::new( 10, @@ -2126,8 +2106,7 @@ mod test { #[test] fn test_mask_with_large_offset() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test with a large offset that shifts all indices let patches = Patches::new( 20, diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index bdb3f1a5276..af80381e105 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -469,8 +469,7 @@ mod tests { #[test] fn test_between_decimal() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![100i128, 200i128, 300i128, 400i128]; let decimal_type = DecimalDType::new(3, 2); let array = DecimalArray::new(values, decimal_type, Validity::NonNullable).into_array(); diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index a5ad61bb973..cb9c324293d 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -772,8 +772,7 @@ mod tests { #[test] fn test_kleene_truth_table() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let lhs = BoolArray::from_iter([ Some(true), Some(true), @@ -836,8 +835,7 @@ mod tests { #[test] fn test_null_constant_kleene() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let lhs = BoolArray::from_iter([Some(false), Some(true), None]).into_array(); let null = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len()) .into_array(); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare.rs b/vortex-array/src/scalar_fn/fns/binary/compare.rs index 02dc2d2ac19..e0985fd342f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare.rs @@ -369,8 +369,7 @@ mod tests { #[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())] #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())] fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let res = left.binary(right, Operator::Eq).unwrap(); let expected = BoolArray::from_iter([true, true]); assert_arrays_eq!(res, expected, &mut assertion_ctx); @@ -378,8 +377,7 @@ mod tests { #[test] fn test_list_array_comparison() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values1 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets1 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list1 = ListArray::try_new( @@ -424,8 +422,7 @@ mod tests { #[test] fn test_list_array_constant_comparison() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list = ListArray::try_new( @@ -452,8 +449,7 @@ mod tests { #[test] fn test_struct_array_comparison() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let bool_field1 = BoolArray::from_iter([Some(true), Some(false), Some(true)]); let int_field1 = PrimitiveArray::from_iter([1i32, 2, 3]); @@ -490,8 +486,7 @@ mod tests { #[test] fn test_empty_struct_compare() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let empty1 = StructArray::try_new( FieldNames::from(Vec::::new()), Vec::new(), @@ -520,8 +515,7 @@ mod tests { /// different Vortex encodings (VarBinArray vs VarBinViewArray) must not panic. #[test] fn struct_compare_mixed_binary_encodings() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // LHS: struct with a VarBinArray (offset-based) binary field let bin_field1 = VarBinArray::from(vec![ "apple".as_bytes(), diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 09938e13cf1..5987e2cb8ca 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -519,8 +519,7 @@ mod tests { #[test] fn test_or_kleene_validity() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::StructArray; @@ -548,8 +547,7 @@ mod tests { #[test] fn test_scalar_subtract_unsigned() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -568,8 +566,7 @@ mod tests { #[test] fn test_scalar_subtract_signed() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -588,8 +585,7 @@ mod tests { #[test] fn test_scalar_subtract_nullable() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -606,8 +602,7 @@ mod tests { #[test] fn test_scalar_subtract_float() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index 17af9dd8289..52dc78d987a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -951,8 +951,7 @@ mod test { #[test] fn test_scalar_subtract_unsigned() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![1u16, 2, 3].into_array(); let result = sub_scalar(&values, 1u16).unwrap(); assert_arrays_eq!( @@ -964,8 +963,7 @@ mod test { #[test] fn test_scalar_subtract_signed() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![1i64, 2, 3].into_array(); let result = sub_scalar(&values, -1i64).unwrap(); assert_arrays_eq!( @@ -977,8 +975,7 @@ mod test { #[test] fn test_scalar_subtract_nullable() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); let result = sub_scalar(&values.into_array(), Some(1u16)).unwrap(); assert_arrays_eq!( @@ -990,8 +987,7 @@ mod test { #[test] fn test_scalar_subtract_float() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![1.0f64, 2.0, 3.0].into_array(); let result = sub_scalar(&values, -1f64).unwrap(); assert_arrays_eq!( @@ -1010,8 +1006,7 @@ mod test { #[test] fn test_float_divide_by_zero_is_ok() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![1.0f64, -1.0].into_array(); let result = values .binary( @@ -1069,8 +1064,7 @@ mod test { #[test] fn test_integer_divide_errors_ignore_null_lanes() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let lhs = PrimitiveArray::new(buffer![10i32, 10], Validity::from_iter([false, true])) .into_array(); let rhs = buffer![0i32, 2].into_array(); @@ -1091,8 +1085,7 @@ mod test { #[test] fn test_integer_errors_ignore_null_lanes() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true])) .into_array(); let result = values @@ -1130,8 +1123,7 @@ mod test { #[test] fn test_present_nullable_constant_preserves_nullable_output() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let values = buffer![1u8, 2].into_array(); let result = values .binary( diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-array/src/scalar_fn/fns/byte_length.rs index 9e0f491bf21..bfe2b11af4a 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-array/src/scalar_fn/fns/byte_length.rs @@ -202,8 +202,7 @@ mod tests { #[case] array: ArrayRef, #[case] expected_lens: Vec, ) -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(expected_lens); assert_arrays_eq!(result, expected, &mut assertion_ctx); @@ -212,8 +211,7 @@ mod tests { #[test] fn test_varbinview_byte_length() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = VarBinViewArray::from_iter_str(["short", "a longer string here"]).into_array(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 20]); diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index 2e41daa1069..5c37b8227ab 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -783,8 +783,7 @@ mod tests { #[test] fn test_evaluate_simple_condition() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -806,8 +805,7 @@ mod tests { #[test] fn test_evaluate_nary_multiple_conditions() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Test n-ary via nested_case_when let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) @@ -832,8 +830,7 @@ mod tests { #[test] fn test_evaluate_nary_first_match_wins() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -858,8 +855,7 @@ mod tests { #[test] fn test_evaluate_no_else_returns_null() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -879,8 +875,7 @@ mod tests { #[test] fn test_evaluate_all_conditions_false() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -902,8 +897,7 @@ mod tests { #[test] fn test_evaluate_all_conditions_true() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -925,8 +919,7 @@ mod tests { #[test] fn test_evaluate_all_true_no_else_returns_correct_dtype() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // CASE WHEN value > 0 THEN 100 END — condition is always true, no ELSE. // Result must be Nullable because the implicit ELSE is NULL. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -950,8 +943,7 @@ mod tests { #[test] fn test_merge_case_branches_widens_nullability_of_later_branch() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // When a later THEN branch is Nullable and branches[0] and ELSE are NonNullable, // the result dtype must still be Nullable. // @@ -989,8 +981,7 @@ mod tests { #[test] fn test_evaluate_with_literal_condition() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = case_when(lit(true), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); @@ -1004,8 +995,7 @@ mod tests { #[test] fn test_evaluate_with_bool_column_result() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1027,8 +1017,7 @@ mod tests { #[test] fn test_evaluate_with_nullable_condition() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]).into_array(), @@ -1048,8 +1037,7 @@ mod tests { #[test] fn test_evaluate_with_nullable_result_values() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[ ("value", buffer![1i32, 2, 3, 4, 5].into_array()), ( @@ -1078,8 +1066,7 @@ mod tests { #[test] fn test_evaluate_with_all_null_condition() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([None, None, None]).into_array(), @@ -1097,8 +1084,7 @@ mod tests { #[test] fn test_evaluate_nary_no_else_returns_null() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1125,8 +1111,7 @@ mod tests { #[test] fn test_evaluate_nary_many_conditions() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1154,8 +1139,7 @@ mod tests { #[test] fn test_evaluate_nary_all_false_no_else() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) .unwrap() .into_array(); @@ -1180,8 +1164,7 @@ mod tests { #[test] fn test_evaluate_nary_overlapping_conditions_first_wins() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![10i32, 20, 30].into_array())]) .unwrap() @@ -1206,8 +1189,7 @@ mod tests { #[test] fn test_evaluate_nary_early_exit_when_remaining_empty() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // After branch 0 claims all rows, remaining becomes all_false. // The loop breaks before evaluating branch 1's condition. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1233,8 +1215,7 @@ mod tests { #[test] fn test_evaluate_nary_skips_branch_with_empty_effective_mask() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Branch 0 claims value=1. Branch 1 targets the same rows but they are already // matched → effective_mask is all_false → branch 1 is skipped (THEN not used). let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1300,8 +1281,7 @@ mod tests { #[test] fn test_evaluate_nary_with_nullable_conditions() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[ ( "cond1", @@ -1422,8 +1402,7 @@ mod tests { #[test] fn test_simplify_null_fill_semantic_equivalence() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // The collapse-to-input rewrite must preserve values (and `x`'s nullability). let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1484,8 +1463,7 @@ mod tests { #[test] fn test_simplify_semantic_equivalence() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // The optimized expression must produce the same values as the original CASE WHEN. let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1514,8 +1492,7 @@ mod tests { #[test] fn test_merge_case_branches_alternating_mask() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); // Exercises the scalar path: alternating rows produce one slice per row (no runs), // triggering the per-row cursor path in merge_case_branches. let n = 100usize; diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 9ea3a17de51..a9ebac9a2a2 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -277,7 +277,6 @@ mod tests { use super::*; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -304,8 +303,7 @@ mod tests { #[test] fn execute_with_value() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -325,8 +323,7 @@ mod tests { #[test] fn execute_without_value_default_true() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -346,8 +343,7 @@ mod tests { #[test] fn execute_without_value_default_false() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -367,8 +363,7 @@ mod tests { #[test] fn execute_value_flips() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let threshold = Arc::new(AtomicI32::new(5)); let threshold_clone = Arc::clone(&threshold); let expr = dynamic( diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index 9069d82cab4..da5946a6228 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -179,7 +179,6 @@ mod tests { use vortex_error::VortexExpect; use crate::IntoArray; - use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::assert_arrays_eq; @@ -209,8 +208,7 @@ mod tests { #[test] fn evaluate() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]) .into_array(); @@ -231,8 +229,7 @@ mod tests { #[test] fn evaluate_struct_field() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = StructArray::from_fields(&[( "a", PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), @@ -256,8 +253,7 @@ mod tests { #[test] fn evaluate_non_nullable_input() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = fill_null(root(), lit(0i32)); let result = test_array.apply(&expr).unwrap(); diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index c2afb16a458..bcd663aa002 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -752,8 +752,7 @@ mod tests { #[case] value: Option<&str>, #[case] expected: BoolArray, ) { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let element_nullability = list_array .dtype() .as_list_element_opt() @@ -771,8 +770,7 @@ mod tests { #[test] fn test_constant_list() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let list_array = ConstantArray::new( Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -791,8 +789,7 @@ mod tests { #[test] fn test_all_nulls() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let list_array = ConstantArray::new( Scalar::null(DType::List( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -814,8 +811,7 @@ mod tests { #[test] fn test_list_array_element() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let list_scalar = Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), vec![1.into(), 3.into(), 6.into()], @@ -832,8 +828,7 @@ mod tests { #[test] fn test_list_contains_empty_listview() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let empty_elements = PrimitiveArray::empty::(Nullability::NonNullable); let offsets = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); let sizes = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); @@ -857,8 +852,7 @@ mod tests { #[test] fn test_list_contains_all_null_elements() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = PrimitiveArray::from_option_iter::([None, None, None, None, None]); let offsets = Buffer::from_iter([0u32, 2, 4]).into_array(); let sizes = Buffer::from_iter([2u32, 2, 1]).into_array(); @@ -894,8 +888,7 @@ mod tests { #[test] fn test_list_contains_large_offsets() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = Buffer::from_iter([1i32, 2, 3, 4, 5]).into_array(); let offsets = Buffer::from_iter([0u32, 1, 4, 0]).into_array(); @@ -919,8 +912,7 @@ mod tests { #[test] fn test_list_contains_offset_size_boundary() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let elements = Buffer::from_iter(0..256).into_array(); let offsets = Buffer::from_iter([0u8, 100, 200, 254]).into_array(); let sizes = Buffer::from_iter([50u8, 50, 54, 2]).into_array(); diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index d152f55ad3e..0a13034b9f6 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -279,7 +279,6 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; - use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -319,8 +318,7 @@ mod tests { #[test] pub fn test_merge_right_most() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let expr = merge_opts( vec![ get_item("0", root()), diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 209621109fe..b3131429536 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -172,7 +172,6 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; - use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -231,8 +230,7 @@ mod tests { #[test] pub fn test_simple_pack() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -266,8 +264,7 @@ mod tests { #[test] pub fn test_nested_pack() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 4fcff8c728c..580f24bc036 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -317,8 +317,7 @@ mod tests { #[test] fn test_zip_basic() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mask = Mask::from_iter([true, false, false, true, false]); let if_true = buffer![10, 20, 30, 40, 50].into_array(); let if_false = buffer![1, 2, 3, 4, 5].into_array(); @@ -331,8 +330,7 @@ mod tests { #[test] fn test_zip_all_true() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mask = Mask::new_true(4); let if_true = buffer![10, 20, 30, 40].into_array(); let if_false = @@ -348,8 +346,7 @@ mod tests { #[test] fn test_zip_all_false_widens_nullability() { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let mask = Mask::new_false(4); let if_true = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array(); diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 036d72abe30..6e3a3f7a41a 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -109,8 +109,7 @@ mod tests { #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let sum_scalar = Scalar::primitive(6i64, Nullability::Nullable); array.statistics().set( @@ -132,8 +131,7 @@ mod tests { #[test] fn stat_expr_returns_null_when_sum_is_missing() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let result = array @@ -153,8 +151,7 @@ mod tests { #[test] fn stat_expr_reads_cached_sum_per_chunk() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let chunk0 = buffer![1i32, 2].into_array(); let sum_scalar = Scalar::primitive(3i64, Nullability::Nullable); chunk0.statistics().set( @@ -190,8 +187,7 @@ mod tests { #[test] fn stat_expr_reads_cached_null_count() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); let null_count_scalar = Scalar::primitive(2u64, Nullability::NonNullable); @@ -218,8 +214,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_null_from_null_count() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); array .statistics() @@ -239,8 +234,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_null_false_from_inexact_low_null_count() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, Some(2), None]).into_array(); array .statistics() @@ -260,8 +254,7 @@ mod tests { #[test] fn stat_expr_returns_null_for_inexact_full_null_count_as_all_null() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, Some(2), None]).into_array(); array .statistics() @@ -281,8 +274,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_non_null_from_null_count() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); array .statistics() @@ -302,8 +294,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_non_null_true_from_inexact_zero_null_count() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); array .statistics() @@ -323,8 +314,7 @@ mod tests { #[test] fn stat_expr_returns_null_for_inexact_nonzero_null_count_as_all_non_null() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); array @@ -358,8 +348,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_nan_from_nan_count() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(f32::NAN), Some(f32::NAN)]) .into_array(); @@ -381,8 +370,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_nan_false_from_inexact_low_nan_count() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32), Some(f32::NAN)]) .into_array(); @@ -404,8 +392,7 @@ mod tests { #[test] fn stat_expr_returns_null_for_inexact_full_nan_count_as_all_nan() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32), Some(f32::NAN)]) .into_array(); @@ -427,8 +414,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_non_nan_true_from_inexact_zero_nan_count() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = buffer![1.0f32, 2.0, 3.0].into_array(); array .statistics() @@ -448,8 +434,7 @@ mod tests { #[test] fn stat_expr_returns_null_for_inexact_nonzero_nan_count_as_all_non_nan() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1.0f32), Some(f32::NAN), Some(3.0)]) .into_array(); array @@ -470,8 +455,7 @@ mod tests { #[test] fn stat_expr_reads_cached_min_and_max() -> VortexResult<()> { - let assertion_session = crate::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = crate::array_execution_ctx(); let array = buffer![3i32, 1, 2].into_array(); array .statistics() diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index 37fab4f7c07..7c4e0c7992d 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -112,8 +112,7 @@ mod tests { #[case] input: ListViewArray, #[case] expect_list: bool, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array_ref = input.clone().into_array(); let result = BtrBlocksCompressor::default() .compress(&array_ref, &mut SESSION.create_execution_ctx())?; @@ -128,8 +127,7 @@ mod tests { #[test] fn test_constant_all_true() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -143,8 +141,7 @@ mod tests { #[test] fn test_constant_all_false() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -158,8 +155,7 @@ mod tests { #[test] fn test_nullable_all_valid_compressed() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true; 100]), Validity::from(BitBuffer::from(vec![true; 100])), @@ -176,8 +172,7 @@ mod tests { #[test] fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0))); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity); let btr = BtrBlocksCompressor::default(); @@ -192,8 +187,7 @@ mod tests { #[test] fn test_mixed_not_constant() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true, false, true]), Validity::NonNullable, @@ -210,8 +204,7 @@ mod tests { #[test] fn test_binary_constant_compressed() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let values = vec![Some(b"constant-bytes".as_slice()); 100]; let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); let btr = BtrBlocksCompressor::default(); @@ -226,8 +219,7 @@ mod tests { #[test] fn test_binary_dict_compressed() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"]; let values = (0..1000) .map(|idx| Some(distinct_values[idx % distinct_values.len()])) diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index 7752989ea08..a2745b2bce8 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -57,8 +57,7 @@ fn test_compress() -> VortexResult<()> { #[test] fn test_rle_compression() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(1.5f32, 100)); values.extend(iter::repeat_n(2.7f32, 200)); diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index 48e7bbd6bb8..ad8f29c5cdd 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -161,8 +161,7 @@ fn test_rle_compressed() -> VortexResult<()> { #[cfg(feature = "unstable_encodings")] #[test] fn test_delta_compressed() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); use vortex_array::assert_arrays_eq; use vortex_fastlanes::Delta; diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index 27c9afb4150..518a4d7b193 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -71,8 +71,7 @@ fn test_dict_encodable() -> VortexResult<()> { #[test] fn constant_mostly_nulls() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = PrimitiveArray::new( buffer![189u8, 189, 189, 189, 189, 189, 189, 189, 189, 0, 46], Validity::from_iter(vec![ @@ -96,8 +95,7 @@ fn constant_mostly_nulls() -> VortexResult<()> { #[test] fn nullable_sequence() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let values = (0i32..20).step_by(7).collect_vec(); let array = PrimitiveArray::from_option_iter(values.clone().into_iter().map(Some)); @@ -113,8 +111,7 @@ fn nullable_sequence() -> VortexResult<()> { #[test] fn test_rle_compression() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(42i32, 100)); values.extend(iter::repeat_n(123i32, 200)); diff --git a/vortex-compressor/src/sample.rs b/vortex-compressor/src/sample.rs index 4b117cb1ae7..d12a60a3625 100644 --- a/vortex-compressor/src/sample.rs +++ b/vortex-compressor/src/sample.rs @@ -133,7 +133,6 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> #[cfg(test)] mod tests { use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; @@ -144,8 +143,7 @@ mod tests { #[test] fn sample_is_deterministic() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); // Create a deterministic array with linear-with-noise pattern let values: Vec = (0i64..100_000).map(|i| i + (i * 7 + 3) % 11).collect(); diff --git a/vortex-cuda/src/dynamic_dispatch/mod.rs b/vortex-cuda/src/dynamic_dispatch/mod.rs index 0d8519dfd3d..dc1b0d0d01a 100644 --- a/vortex-cuda/src/dynamic_dispatch/mod.rs +++ b/vortex-cuda/src/dynamic_dispatch/mod.rs @@ -1145,8 +1145,7 @@ mod tests { #[crate::test] async fn test_dict_mixed_width_u8_codes_u32_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let len = 3000; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u8).collect(); @@ -1179,8 +1178,7 @@ mod tests { #[crate::test] async fn test_dict_mixed_width_u16_codes_u32_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let dict_values: Vec = vec![1000, 2000, 3000, 4000, 5000]; let len = 2048; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u16).collect(); @@ -1213,8 +1211,7 @@ mod tests { #[crate::test] async fn test_runend_mixed_width_u64_ends_u32_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let ends: Vec = vec![1000, 2000, 3000]; let values: Vec = vec![10, 20, 30]; let len = 3000; @@ -1637,8 +1634,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u8() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let bit_width: u8 = 4; let len = 3000; let reference = 100u8; @@ -1673,8 +1669,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u16() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let reference = 1000u16; @@ -1711,8 +1706,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u64() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let bit_width: u8 = 20; let len = 3000; let reference = 100_000u64; @@ -1758,8 +1752,7 @@ mod tests { #[crate::test] async fn test_single_element() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let values: Vec = vec![42]; let primitive = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); let bp = BitPacked::encode( @@ -1785,8 +1778,7 @@ mod tests { #[crate::test] async fn test_exactly_elements_per_block() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); // Exactly 2048 elements — one full block, no remainder let bit_width: u8 = 6; let len = 2048; @@ -1945,8 +1937,7 @@ mod tests { #[crate::test] async fn test_alp_f64_for_bitpacked() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut ctx = LEGACY_SESSION.create_execution_ctx(); // ALP(FoR(BitPacked)) with f64: same structure as the f32 test. let len = 3000; @@ -2047,8 +2038,7 @@ mod tests { #[crate::test] async fn alp_slice_device_patches() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); // Regression test for https://github.com/vortex-data/vortex/issues/7838#issuecomment-4452796116. let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 4096; @@ -2111,8 +2101,7 @@ mod tests { #[crate::test] async fn test_runend_u32_ends_u16_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); // RunEnd with u32 ends, u16 values. Output type = u16. // Ends (u32) differ from output (u16) → pending subtree. let ends: Vec = vec![500, 1000, 1500, 2000]; @@ -2157,8 +2146,7 @@ mod tests { #[crate::test] async fn test_dict_bitpacked_u8_codes_u32_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); // Dict with BitPacked u8 codes (narrower than u32 output) and u32 values. // The kernel's bitunpack_typed decodes at the source's native width and // widens to T, so this fuses into a single kernel launch. @@ -2208,8 +2196,7 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); @@ -2253,8 +2240,7 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); @@ -2308,8 +2294,7 @@ mod tests { #[case] ends: Vec, #[case] values: Vec, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let ends_u64: Vec = ends.iter().map(|e| (*e).into()).collect(); let len = *ends_u64.last().unwrap() as usize; let bit_width = 64 - ends_u64.iter().max().unwrap().leading_zeros() as u8; @@ -2354,8 +2339,7 @@ mod tests { /// at native width and widens to T, fusing everything. #[crate::test] async fn test_runend_mixed_width_for_bp_u16_ends_u32_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let ends: Vec = vec![500, 1000, 1500, 2000]; let values: Vec = vec![100, 200, 300, 400]; let len = 2000usize; @@ -2398,8 +2382,7 @@ mod tests { #[crate::test] async fn test_sliced_dict_mixed_width() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); // Sliced Dict with u8 codes and u32 values — combines PartiallyFused + slice handling. let dict_values: Vec = vec![100, 200, 300, 400]; let full_len = 4096; @@ -2510,8 +2493,7 @@ mod tests { /// Nullable Primitive array — LOAD source with validity propagated. #[crate::test] async fn test_nullable_primitive() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let array = PrimitiveArray::from_option_iter( @@ -2535,8 +2517,7 @@ mod tests { /// validity, so this produces a real nullable FoR(BitPacked) tree. #[crate::test] async fn test_nullable_for_bitpacked() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 2048; @@ -2609,8 +2590,7 @@ mod tests { /// AllValid nullable array — should fuse and produce AllValid output. #[crate::test] async fn test_all_valid_nullable() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let values: Vec = (0..2048).collect(); @@ -2647,8 +2627,7 @@ mod tests { /// Dict with non-nullable codes but nullable values should still fuse. #[crate::test] async fn test_dict_nullable_values_fuses() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); use vortex::buffer::buffer; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; @@ -2674,8 +2653,7 @@ mod tests { /// Validity must survive through fused dispatch and into the filter. #[crate::test] async fn test_nullable_fused_then_filter() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); use vortex::array::arrays::FilterArray; use vortex::mask::Mask; @@ -2956,8 +2934,7 @@ mod tests { /// u8 BitPacked with patches (bit_width=3, patch values > 7). #[crate::test] async fn test_bitpacked_with_patches_u8() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let bit_width: u8 = 3; let len = 3000usize; let max_val = (1u8 << bit_width) - 1; @@ -2991,8 +2968,7 @@ mod tests { /// u16 BitPacked with patches (bit_width=6, patch values > 63). #[crate::test] async fn test_bitpacked_with_patches_u16() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let bit_width: u8 = 6; let len = 3000usize; let max_val = (1u16 << bit_width) - 1; @@ -3026,8 +3002,7 @@ mod tests { /// u64 BitPacked with patches (bit_width=4, patch values > 15). #[crate::test] async fn test_bitpacked_with_patches_u64() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let bit_width: u8 = 4; let len = 3000usize; let max_val = (1u64 << bit_width) - 1; @@ -3161,8 +3136,7 @@ mod tests { /// dispatch alongside patch application. #[crate::test] async fn test_nullable_bitpacked_with_patches() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 3000usize; diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index 5cd4dede718..66851f85682 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -194,8 +194,7 @@ mod tests { /// FoR(BitPacked) u32 — entire tree compiles into a single fused plan. #[crate::test] async fn test_fused() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let values: Vec = (0..2048).map(|i| (i % 128) as u32).collect(); @@ -223,8 +222,7 @@ mod tests { /// Exercises the unsigned type reinterpretation in CudaDispatchPlan::execute. #[crate::test] async fn test_fused_f32() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); use vortex::encodings::alp::ALP; use vortex::encodings::alp::Exponents; @@ -260,8 +258,7 @@ mod tests { /// ALP with patches — plan builder rejects it, falls back to ALPExecutor. #[crate::test] async fn test_fallback() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); use vortex::array::patches::Patches; use vortex::array::validity::Validity::NonNullable as NN; use vortex::buffer::buffer; @@ -303,8 +300,7 @@ mod tests { #[cfg(feature = "unstable_encodings")] #[crate::test] async fn test_partial_fusion() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); use vortex::array::arrays::DictArray; use vortex::array::session::ArraySessionExt; use vortex::encodings::fastlanes; @@ -368,8 +364,7 @@ mod tests { /// Filter(FoR(BP), mask) — FoR+BP fuses via dyn dispatch, then CUB filters the result. #[crate::test] async fn test_filter_fused_child() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -413,8 +408,7 @@ mod tests { ))] #[crate::test] async fn test_ext_storage_gpu_decode(#[case] ext: ExtensionArray) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -438,8 +432,7 @@ mod tests { /// Extension over already-canonical storage executes unchanged. #[crate::test] async fn test_ext_canonical_storage() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); diff --git a/vortex-cuda/src/kernel/arrays/constant.rs b/vortex-cuda/src/kernel/arrays/constant.rs index 7854e5f12e5..f1591d8159a 100644 --- a/vortex-cuda/src/kernel/arrays/constant.rs +++ b/vortex-cuda/src/kernel/arrays/constant.rs @@ -199,7 +199,6 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::scalar::Scalar; - use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -224,8 +223,7 @@ mod tests { async fn test_cuda_constant_materialization( #[case] constant_array: ConstantArray, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -246,8 +244,7 @@ mod tests { #[crate::test] async fn test_cuda_constant_empty_array() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -269,8 +266,7 @@ mod tests { #[crate::test] async fn test_cuda_constant_small_array() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index 312bc2b4ea2..74070f684b7 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -308,7 +308,6 @@ mod tests { use vortex::dtype::DecimalDType; use vortex::dtype::i256; use vortex::error::VortexExpect; - use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -325,8 +324,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_u32_values_u8_codes() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -363,8 +361,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_u64_values_u16_codes() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -404,8 +401,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_i32_values_u32_codes() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -441,8 +437,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_large_array() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -479,8 +474,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_values_with_validity() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -517,8 +511,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_codes_with_validity() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -560,8 +553,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_both_with_validity() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -610,8 +602,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_i64_values_with_validity() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -661,8 +652,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_all_valid_matches_baseline() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -710,8 +700,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i8_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -744,8 +733,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i16_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -778,8 +766,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i32_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -812,8 +799,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i64_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -849,8 +835,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i128_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -899,8 +884,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_u8_codes() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -930,8 +914,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_u16_codes() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -961,8 +944,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_max_inlined_12_bytes() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -994,8 +976,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_outlined_views() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1030,8 +1011,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_empty_strings() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1061,8 +1041,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_with_validity() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1093,8 +1072,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_outlined_with_validity() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1132,8 +1110,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i256_values() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/alp.rs b/vortex-cuda/src/kernel/encodings/alp.rs index bf2ded1007a..610f8c0bbda 100644 --- a/vortex-cuda/src/kernel/encodings/alp.rs +++ b/vortex-cuda/src/kernel/encodings/alp.rs @@ -170,8 +170,7 @@ mod tests { /// Patches must carry `chunk_offsets` — the fused kernel requires them. #[crate::test] async fn test_cuda_alp_decompression_f32() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -219,8 +218,7 @@ mod tests { /// preserved through the standalone ALP GPU executor. #[crate::test] async fn test_cuda_alp_nullable_with_patches() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -262,8 +260,7 @@ mod tests { /// elements are actually null. #[crate::test] async fn test_cuda_alp_all_valid_nullable() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -297,8 +294,7 @@ mod tests { /// (zero patches) via the offset math rather than the NULL sentinel. #[crate::test] async fn test_cuda_alp_multi_chunk_sparse_patches() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -344,8 +340,7 @@ mod tests { /// so this guards the fast-path for the (i64, f64) kernel variant. #[crate::test] async fn test_cuda_alp_f64_multi_chunk_with_patches() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -390,8 +385,7 @@ mod tests { /// (existing tests have ≤ 6 patches per chunk). #[crate::test] async fn test_cuda_alp_dense_patches_single_chunk() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -440,8 +434,7 @@ mod tests { /// loop. Includes a patch in the tail. #[crate::test] async fn test_cuda_alp_partial_tail_chunk() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index e6cdf5fcba2..875fdec09a3 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -250,8 +250,7 @@ mod tests { #[case] iter: impl Iterator, #[case] bw: u8, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -284,8 +283,7 @@ mod tests { #[crate::test] fn test_patches() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -329,8 +327,7 @@ mod tests { #[case::bw_7(7)] #[crate::test] fn test_cuda_bitunpack_u8(#[case] bit_width: u8) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -384,8 +381,7 @@ mod tests { #[case::bw_15(15)] #[crate::test] fn test_cuda_bitunpack_u16(#[case] bit_width: u8) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -455,8 +451,7 @@ mod tests { #[case::bw_31(31)] #[crate::test] fn test_cuda_bitunpack_u32(#[case] bit_width: u8) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -558,8 +553,7 @@ mod tests { #[case::bw_63(63)] #[crate::test] fn test_cuda_bitunpack_u64(#[case] bit_width: u8) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -596,8 +590,7 @@ mod tests { #[crate::test] fn test_cuda_bitunpack_sliced() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let bit_width = 32; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -681,8 +674,7 @@ mod tests { /// offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_sliced_patches_offset_within_chunk() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -725,8 +717,7 @@ mod tests { /// Test slicing a bitpacked array multiple times, accumulating offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_double_sliced_patches() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -781,8 +772,7 @@ mod tests { /// Test slicing to skip an entire chunk's worth of patches. #[crate::test] fn test_cuda_bitunpack_sliced_skip_first_chunk_patches() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/date_time_parts.rs b/vortex-cuda/src/kernel/encodings/date_time_parts.rs index 582de845d94..e3cbc1fda59 100644 --- a/vortex-cuda/src/kernel/encodings/date_time_parts.rs +++ b/vortex-cuda/src/kernel/encodings/date_time_parts.rs @@ -214,7 +214,6 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::extension::datetime::TimeUnit; - use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -280,8 +279,7 @@ mod tests { #[case] subseconds: Vec, #[case] time_unit: TimeUnit, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -303,8 +301,7 @@ mod tests { #[crate::test] async fn test_cuda_datetimeparts_large_array() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -331,8 +328,7 @@ mod tests { #[crate::test] async fn test_cuda_datetimeparts_with_nulls() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index a018b0d33f8..1752127d423 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -70,7 +70,6 @@ mod tests { use vortex::dtype::DecimalDType; use vortex::encodings::decimal_byte_parts::DecimalByteParts; use vortex::error::VortexExpect; - use vortex_array::VortexSessionExecute; use super::*; use crate::session::CudaSession; @@ -86,8 +85,7 @@ mod tests { #[case] precision: u8, #[case] scale: i8, ) { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/for_.rs b/vortex-cuda/src/kernel/encodings/for_.rs index f71bd1738ac..4de3fb79ed9 100644 --- a/vortex-cuda/src/kernel/encodings/for_.rs +++ b/vortex-cuda/src/kernel/encodings/for_.rs @@ -160,8 +160,7 @@ mod tests { #[case::u64(make_for_array((0..2050).map(|i| (i % 2050) as u64).collect(), 1000000u64))] #[crate::test] async fn test_cuda_for_decompression(#[case] for_array: FoRArray) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -182,8 +181,7 @@ mod tests { #[crate::test] async fn test_signed_ffor() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 0b6a65b7f8a..6c2b080058e 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -209,7 +209,6 @@ mod tests { use vortex::encodings::fsst::fsst_compress; use vortex::encodings::fsst::fsst_train_compressor; use vortex::error::VortexExpect; - use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -236,8 +235,7 @@ mod tests { #[case] strings: Vec>, #[case] nullability: Nullability, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -262,8 +260,7 @@ mod tests { /// Exercises the multi-block grid-stride path on a larger dataset. #[crate::test] async fn test_cuda_fsst_decompression_roundtrip_large() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); use vortex_fsst::test_utils::make_fsst_clickbench_urls; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index 2059f3da432..52849f109dc 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -175,7 +175,6 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex_array::ExecutionCtx; - use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -205,8 +204,7 @@ mod tests { #[case::u64_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u64, 5, 10], vec![1i32, 2, 3], ctx))] #[crate::test] async fn test_cuda_runend_types(#[case] build: RunEndBuilder) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -228,8 +226,7 @@ mod tests { #[crate::test] async fn test_cuda_runend_large_array() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -260,8 +257,7 @@ mod tests { #[crate::test] async fn test_cuda_runend_single_run() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -284,8 +280,7 @@ mod tests { #[crate::test] async fn test_cuda_runend_many_small_runs() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -313,8 +308,7 @@ mod tests { #[crate::test] async fn test_cuda_runend_nullable_values_falls_back_to_cpu() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/sequence.rs b/vortex-cuda/src/kernel/encodings/sequence.rs index 9fff943c520..db36d068f57 100644 --- a/vortex-cuda/src/kernel/encodings/sequence.rs +++ b/vortex-cuda/src/kernel/encodings/sequence.rs @@ -92,7 +92,6 @@ mod tests { use vortex::dtype::Nullability; use vortex::encodings::sequence::Sequence; use vortex::scalar::PValue; - use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::CudaSession; @@ -125,8 +124,7 @@ mod tests { len: usize, nullability: Nullability, ) { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).unwrap(); let array = Sequence::try_new_typed(base, multiplier, nullability, len).unwrap(); diff --git a/vortex-cuda/src/kernel/encodings/zigzag.rs b/vortex-cuda/src/kernel/encodings/zigzag.rs index 7f8b5738e9a..b6321ed5178 100644 --- a/vortex-cuda/src/kernel/encodings/zigzag.rs +++ b/vortex-cuda/src/kernel/encodings/zigzag.rs @@ -105,7 +105,6 @@ mod tests { use vortex::buffer::Buffer; use vortex::encodings::zigzag::ZigZag; use vortex::error::VortexExpect; - use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -113,8 +112,7 @@ mod tests { #[crate::test] async fn test_cuda_zigzag_decompression_u32() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index 9e1bd991d40..71597f59e8c 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -356,7 +356,6 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::encodings::zstd::Zstd; use vortex::error::VortexResult; - use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -365,8 +364,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_utf8() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -397,8 +395,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_multiple_frames() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -439,8 +436,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_sliced() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -480,8 +476,7 @@ mod tests { /// correct results instead of panicking. #[crate::test] async fn test_cuda_zstd_nullable_falls_back_to_cpu() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs index 7da492879e7..5a67fb4ea12 100644 --- a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs +++ b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs @@ -234,8 +234,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_buffers_decompression_primitive() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -259,8 +258,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_buffers_decompression_varbinview() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/filter/decimal.rs b/vortex-cuda/src/kernel/filter/decimal.rs index b5a308f6e19..e40b1445bf4 100644 --- a/vortex-cuda/src/kernel/filter/decimal.rs +++ b/vortex-cuda/src/kernel/filter/decimal.rs @@ -48,7 +48,6 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; - use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -89,8 +88,7 @@ mod tests { #[case] input: DecimalArray, #[case] mask: Mask, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -113,8 +111,7 @@ mod tests { #[crate::test] async fn test_gpu_filter_decimal_large_array() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); diff --git a/vortex-cuda/src/kernel/filter/primitive.rs b/vortex-cuda/src/kernel/filter/primitive.rs index 4c701f81e9b..6c099d864c4 100644 --- a/vortex-cuda/src/kernel/filter/primitive.rs +++ b/vortex-cuda/src/kernel/filter/primitive.rs @@ -46,7 +46,6 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; - use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -83,8 +82,7 @@ mod tests { #[case] input: PrimitiveArray, #[case] mask: Mask, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -107,8 +105,7 @@ mod tests { #[crate::test] async fn test_gpu_filter_large_array() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); diff --git a/vortex-cuda/src/kernel/filter/varbinview.rs b/vortex-cuda/src/kernel/filter/varbinview.rs index 704d448c6f5..87267e368be 100644 --- a/vortex-cuda/src/kernel/filter/varbinview.rs +++ b/vortex-cuda/src/kernel/filter/varbinview.rs @@ -46,7 +46,6 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; - use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -69,8 +68,7 @@ mod tests { #[case] input: VarBinViewArray, #[case] mask: Mask, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); diff --git a/vortex-cuda/src/kernel/patches/types.rs b/vortex-cuda/src/kernel/patches/types.rs index 63c3b731253..2bb4c0056e0 100644 --- a/vortex-cuda/src/kernel/patches/types.rs +++ b/vortex-cuda/src/kernel/patches/types.rs @@ -228,7 +228,6 @@ mod tests { use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::patches::Patches; use vortex_error::VortexResult; @@ -318,8 +317,7 @@ mod tests { #[case] chunk_offsets: ArrayRef, #[case] expected: ArrayRef, ) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let indices = PrimitiveArray::from_iter([100u32, 1100, 2100, 3100, 4100]); let values = PrimitiveArray::from_iter([10u32, 11, 12, 13, 14]); diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 4cfcdbca5f4..a86e375d37b 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -415,8 +415,7 @@ mod tests { #[test] fn test_integer_vector_conversion() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let values = vec![1i32, 2, 3, 4, 5]; let len = values.len(); diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 32051f17d96..f656755ab65 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -532,8 +532,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_no_options() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let (array, struct_array) = scan(ptr::null()); assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); unsafe { vx_array_free(array) }; @@ -542,8 +541,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_project_all() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let opts = vx_scan_options::default(); let (array, struct_array) = scan(&raw const opts); assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); @@ -554,8 +552,7 @@ mod tests { #[cfg_attr(miri, ignore)] fn test_project_single_field() { unsafe { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let root = vx_expression_root(); let mut opts = vx_scan_options::default(); @@ -717,8 +714,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_ordered() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let opts = vx_scan_options { ordered: true, ..Default::default() diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 01afe39663d..cba3a4ec430 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -477,8 +477,7 @@ async fn test_empty_varbin_array_roundtrip() { #[tokio::test] #[cfg_attr(miri, ignore)] async fn issue_5385_filter_casted_column() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = StructArray::try_from_iter([("x", buffer![1u8, 2, 3, 4, 5])]) .unwrap() .into_array(); @@ -1188,8 +1187,7 @@ async fn chunked_file() -> VortexResult { #[tokio::test] async fn basic_file_roundtrip() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let vxf = chunked_file().await?; let result = vxf.scan()?.into_array_stream()?.read_all().await?; @@ -1233,8 +1231,7 @@ async fn file_excluding_dtype() -> VortexResult<()> { #[tokio::test] async fn file_take() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let vxf = chunked_file().await?; let result = vxf .scan()? diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index d2a815e9758..3aa66034068 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -123,8 +123,7 @@ async fn test_file_roundtrip() { /// with "Array vortex.fill_null does not support serialization". #[tokio::test] async fn test_dict_listview_validity_roundtrip() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]).into_array(); let offsets = PrimitiveArray::from_iter(vec![0u32, 2, 4]).into_array(); let sizes = PrimitiveArray::from_iter(vec![2u32, 2, 1]).into_array(); diff --git a/vortex-ipc/src/iterator.rs b/vortex-ipc/src/iterator.rs index 4714f055dc5..04bbbf21d52 100644 --- a/vortex-ipc/src/iterator.rs +++ b/vortex-ipc/src/iterator.rs @@ -172,7 +172,6 @@ mod test { use std::io::Cursor; use vortex_array::IntoArray as _; - use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorExt; @@ -183,8 +182,7 @@ mod test { #[test] fn test_sync_stream() -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let ipc_buffer = array .to_array_iterator() diff --git a/vortex-ipc/src/stream.rs b/vortex-ipc/src/stream.rs index 1ab4b87a64d..68f77e7b80c 100644 --- a/vortex-ipc/src/stream.rs +++ b/vortex-ipc/src/stream.rs @@ -225,7 +225,6 @@ mod test { use futures::io::Cursor; use vortex_array::IntoArray as _; - use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamExt; @@ -236,8 +235,7 @@ mod test { #[tokio::test] async fn test_async_stream() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = buffer![1, 2, 3].into_array(); let ipc_buffer = array .to_array_stream() @@ -274,8 +272,7 @@ mod test { #[tokio::test] async fn test_async_stream_chunked() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(); let ipc_buffer = array .to_array_stream() @@ -299,8 +296,7 @@ mod test { /// Test with 1-byte chunks to stress-test partial read handling. #[tokio::test] async fn test_async_stream_single_byte_chunks() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = buffer![42i64, -1, 0, i64::MAX, i64::MIN].into_array(); let ipc_buffer = array .to_array_stream() diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index 94de11f0096..39c8c0223b1 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -353,7 +353,6 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; - use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; @@ -472,8 +471,7 @@ mod test { #[from(chunked_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|_h| async { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let result = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 46f1da82b05..29e0b9ab3f9 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -422,8 +422,7 @@ mod tests { #[test] fn reading_nested_packs_works() { block_on(|handle| async move { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), @@ -523,8 +522,7 @@ mod tests { #[case] expected: Vec, ) { block_on(|handle| async move { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), @@ -579,8 +577,7 @@ mod tests { #[test] fn reading_is_null_works() { block_on(|handle| async move { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let mut ctx_exec = LEGACY_SESSION.create_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( @@ -654,8 +651,7 @@ mod tests { #[test] fn reading_byte_length_pushdown_works() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = VarBinArray::from_iter( [ Some("abc"), @@ -713,8 +709,7 @@ mod tests { } fn test_apply(original: Expression, outer: Expression, inner: Expression) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = VarBinArray::from_iter( [Some("abc"), Some("def"), None], DType::Utf8(Nullability::Nullable), diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 7af89a19f4a..0df63df4ddf 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -233,7 +233,6 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; - use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -256,8 +255,7 @@ mod test { #[test] fn flat_identity() -> VortexResult<()> { block_on(|handle| async { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -297,8 +295,7 @@ mod test { #[test] fn flat_expr() { block_on(|handle| async { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); @@ -338,8 +335,7 @@ mod test { #[test] fn flat_unaligned_row_mask() { block_on(|handle| async { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 8ba7ebf007d..6651b8aa3fd 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -327,7 +327,6 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::IntoArray as _; use vortex_array::MaskFuture; - use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; use vortex_array::expr::eq; @@ -352,8 +351,7 @@ mod tests { #[test] fn flat_expr_no_row_id() { block_on(|handle| async { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -398,8 +396,7 @@ mod tests { #[test] fn flat_expr_row_id() { block_on(|handle| async { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -444,8 +441,7 @@ mod tests { #[test] fn flat_expr_or() { block_on(|handle| async { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 50e7b409237..1d4d4958889 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -634,8 +634,7 @@ mod tests { fn test_struct_layout( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -654,8 +653,7 @@ mod tests { fn test_struct_layout_row_mask( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -722,8 +720,7 @@ mod tests { fn test_struct_layout_nulls( #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); // Read the layout source from the top. let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index b6cbc9b52da..eb6b0ebdbc4 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -226,7 +226,6 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; - use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -310,8 +309,7 @@ mod test { #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|handle| async { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let session = session_with_handle(handle); let result = layout .new_reader("".into(), segments, &session, &Default::default()) diff --git a/vortex-python/src/arrays/range_to_sequence.rs b/vortex-python/src/arrays/range_to_sequence.rs index dcb04d13ee4..a04a70e32a9 100644 --- a/vortex-python/src/arrays/range_to_sequence.rs +++ b/vortex-python/src/arrays/range_to_sequence.rs @@ -78,7 +78,6 @@ mod test { use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; - use vortex_array::VortexSessionExecute; use crate::arrays::range_to_sequence::range_len; use crate::arrays::range_to_sequence::sequence_array_from_range; @@ -99,8 +98,7 @@ mod test { #[test] fn test_sequence_array_from_len() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let dtype = DType::Primitive(PType::U16, Nullability::NonNullable); let arr = sequence_array_from_range::(0, 10, 1, dtype).unwrap(); assert_arrays_eq!( diff --git a/vortex-test/compat-gen/src/check.rs b/vortex-test/compat-gen/src/check.rs index 6923a2be65d..102c949b2f0 100644 --- a/vortex-test/compat-gen/src/check.rs +++ b/vortex-test/compat-gen/src/check.rs @@ -5,7 +5,6 @@ use std::path::Path; use clap::ValueEnum; use serde::Serialize; -use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -47,8 +46,7 @@ struct FailedFixture { /// Prints JSON result to stdout, human-readable progress to stderr. /// Returns error if any fixture failed or if mode constraints are violated. pub fn check(dir: &Path, mode: Mode, exclude: &[String]) -> VortexResult<()> { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let fixtures = all_fixtures(); let fixtures: Vec<_> = fixtures .into_iter() diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs index 6f036525c58..01a9c4cb898 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs @@ -17,7 +17,7 @@ pub fn fixtures() -> Vec> { #[cfg(test)] mod tests { - use vortex_array::VortexSessionExecute; + use vortex_array::assert_arrays_eq; use super::fixtures; @@ -27,8 +27,7 @@ mod tests { #[test] fn roundtrip_fixtures_to_bytes() { for fixture in fixtures() { - let assertion_session = vortex_array::array_session(); - let mut assertion_ctx = assertion_session.create_execution_ctx(); + let mut assertion_ctx = vortex_array::array_execution_ctx(); let array = fixture.build().unwrap(); check_expected_encodings(&array, fixture.as_ref()).unwrap(); let bytes = adapter::write_file_to_bytes(array.clone()).unwrap(); From de4c3a4a8ba4a48a9b5c36c7ab151b58564c5d92 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 16:58:36 +0100 Subject: [PATCH 07/12] Move AllNonDistinct-based equality + variant impl out of assert_arrays_eq assert_arrays_eq now uses find_mismatched_indices (scalar comparison) for detection again; the AllNonDistinct-based comparison and its variant-type implementation are extracted to a separate (stacked) PR. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Robert Kruszewski --- .../src/compute/allnondistinct.rs | 65 ------------------- encodings/parquet-variant/src/compute/mod.rs | 6 -- encodings/parquet-variant/src/kernel.rs | 14 ---- encodings/parquet-variant/src/lib.rs | 1 - .../aggregate_fn/fns/all_non_distinct/mod.rs | 6 +- .../fns/all_non_distinct/primitive.rs | 7 +- .../fns/all_non_distinct/variant.rs | 32 --------- vortex-array/src/arrays/assertions.rs | 11 +--- 8 files changed, 6 insertions(+), 136 deletions(-) delete mode 100644 encodings/parquet-variant/src/compute/allnondistinct.rs delete mode 100644 encodings/parquet-variant/src/compute/mod.rs delete mode 100644 vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs diff --git a/encodings/parquet-variant/src/compute/allnondistinct.rs b/encodings/parquet-variant/src/compute/allnondistinct.rs deleted file mode 100644 index 766741bb75c..00000000000 --- a/encodings/parquet-variant/src/compute/allnondistinct.rs +++ /dev/null @@ -1,65 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fmt::Debug; - -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct; -use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct; -use vortex_array::aggregate_fn::kernels::DynAggregateKernel; -use vortex_array::arrays::Struct; -use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::dtype::Nullability; -use vortex_array::scalar::Scalar; -use vortex_error::VortexResult; - -use crate::ParquetVariant; -use crate::ParquetVariantArrayExt; - -/// Lets `AllNonDistinct` compare two `ParquetVariant` arrays without canonicalizing them. -/// -/// `AllNonDistinct` accumulates over a `Struct{lhs, rhs}` batch, so this kernel is registered for -/// the struct encoding and inspects the two children. When both are `ParquetVariant`, we compare -/// the typed (`typed_value`) arrays if both sides are shredded, and fall back to the raw `value` -/// arrays otherwise. Comparing these child arrays directly avoids re-canonicalizing the variant -/// (which would recurse through the `Variant` canonical form). -#[derive(Debug)] -pub struct AllNonDistinctParquetVariant; - -impl DynAggregateKernel for AllNonDistinctParquetVariant { - fn aggregate( - &self, - aggregate_fn: &AggregateFnRef, - batch: &ArrayRef, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - if !aggregate_fn.is::() { - return Ok(None); - } - - let Some(batch) = batch.as_opt::() else { - return Ok(None); - }; - let lhs = batch.unmasked_field(0); - let rhs = batch.unmasked_field(1); - let (Some(lhs), Some(rhs)) = ( - lhs.as_opt::(), - rhs.as_opt::(), - ) else { - return Ok(None); - }; - - let identical = match (lhs.typed_value_array(), rhs.typed_value_array()) { - (Some(lhs_typed), Some(rhs_typed)) => all_non_distinct(lhs_typed, rhs_typed, ctx)?, - _ => match (lhs.value_array(), rhs.value_array()) { - (Some(lhs_value), Some(rhs_value)) => all_non_distinct(lhs_value, rhs_value, ctx)?, - // Mixed shredding layouts: let the generic canonical path handle it. - _ => return Ok(None), - }, - }; - - Ok(Some(Scalar::bool(identical, Nullability::NonNullable))) - } -} diff --git a/encodings/parquet-variant/src/compute/mod.rs b/encodings/parquet-variant/src/compute/mod.rs deleted file mode 100644 index f69f592b89d..00000000000 --- a/encodings/parquet-variant/src/compute/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -mod allnondistinct; - -pub use allnondistinct::*; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index f4aa2b145ed..684ce0217f4 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -18,13 +18,9 @@ use vortex_array::ArrayVTable; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::aggregate_fn::AggregateFnVTable; -use vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::arrays::Dict; use vortex_array::arrays::Filter; use vortex_array::arrays::Slice; -use vortex_array::arrays::Struct; use vortex_array::arrays::dict::TakeExecute; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; @@ -49,7 +45,6 @@ use vortex_session::VortexSession; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; -use crate::compute::AllNonDistinctParquetVariant; pub(crate) fn initialize(session: &VortexSession) { let kernels = session.kernels(); @@ -69,15 +64,6 @@ pub(crate) fn initialize(session: &VortexSession) { TakeExecuteAdaptor(ParquetVariant), ); kernels.register_execute_parent_kernel(VariantGet.id(), ParquetVariant, VariantGetKernel); - - // `AllNonDistinct` accumulates over a `Struct{lhs, rhs}` batch, so the kernel is keyed on the - // struct encoding and matches `ParquetVariant` children. See `AllNonDistinctParquetVariant`. - let aggregates = session.aggregate_fns(); - aggregates.register_aggregate_kernel( - Struct.id(), - Some(AllNonDistinct.id()), - &AllNonDistinctParquetVariant, - ); } #[derive(Default, Debug)] diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 91045a65e48..0076af48d22 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -26,7 +26,6 @@ mod array; mod arrow; -mod compute; mod kernel; mod operations; mod validity; diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index fcb9b18346d..445f4e17b6f 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs @@ -12,7 +12,6 @@ mod struct_; #[cfg(test)] mod tests; mod varbin; -mod variant; use std::sync::LazyLock; @@ -39,7 +38,6 @@ use crate::aggregate_fn::AggregateFnId; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::fns::all_non_distinct::variant::check_variant_identical; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; @@ -266,8 +264,8 @@ fn check_canonical_identical( (Canonical::Extension(lhs), Canonical::Extension(rhs)) => { check_extension_identical(lhs, rhs, ctx) } - (Canonical::Variant(lhs), Canonical::Variant(rhs)) => { - check_variant_identical(lhs, rhs, ctx) + (Canonical::Variant(_), _) | (_, Canonical::Variant(_)) => { + vortex_bail!("Variant arrays don't support AllNonDistinct") } _ => Err(vortex_err!( "Canonical type mismatch in AllNonDistinct: {:?} vs {:?}", diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs index 1ad0ca6d646..f11a9a2bc52 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::dtype::NativePType; use vortex_error::VortexResult; use crate::arrays::primitive::PrimitiveArrayExt; @@ -13,10 +12,6 @@ where R: PrimitiveArrayExt, { match_each_native_ptype!(lhs.ptype(), |P| { - Ok(lhs - .as_slice::

() - .iter() - .zip(rhs.as_slice::

()) - .all(|(l, r)| l.is_eq(*r))) + Ok(lhs.as_slice::

() == rhs.as_slice::

()) }) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs deleted file mode 100644 index 1576b8c3d6a..00000000000 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs +++ /dev/null @@ -1,32 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use vortex_error::VortexResult; - -use crate::ExecutionCtx; -use crate::arrays::VariantArray; - -/// Checks whether two canonical variant arrays are element-wise non-distinct. -/// -/// Variant values cannot be routed back through [`all_non_distinct`]: canonicalizing a variant -/// value array yields another canonical variant (with no shredded tree), which would recurse -/// forever. The generic fallback therefore compares logical variant scalars row-by-row. Encodings -/// that can compare their typed/value children more cheaply (e.g. `ParquetVariant`) register an -/// aggregate kernel that intercepts the comparison before it reaches this fallback. -/// -/// [`all_non_distinct`]: super::all_non_distinct -pub(super) fn check_variant_identical( - lhs: &VariantArray, - rhs: &VariantArray, - ctx: &mut ExecutionCtx, -) -> VortexResult { - if lhs.len() != rhs.len() { - return Ok(false); - } - for idx in 0..lhs.len() { - if lhs.execute_scalar(idx, ctx)? != rhs.execute_scalar(idx, ctx)? { - return Ok(false); - } - } - Ok(true) -} diff --git a/vortex-array/src/arrays/assertions.rs b/vortex-array/src/arrays/assertions.rs index 7d2948e5d79..a445f26ec15 100644 --- a/vortex-array/src/arrays/assertions.rs +++ b/vortex-array/src/arrays/assertions.rs @@ -10,7 +10,6 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::RecursiveCanonical; -use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct; fn format_indices>(indices: I) -> impl Display { indices.into_iter().format(",") @@ -118,14 +117,10 @@ macro_rules! assert_arrays_eq { pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef, ctx: &mut ExecutionCtx) { let executed = execute_to_canonical(left.clone(), ctx); - let left_right_the_same = - all_non_distinct(left, right, ctx).vortex_expect("failed to compare left and right"); - let executed_right_the_same = all_non_distinct(&executed, right, ctx) - .vortex_expect("failed to compare executed left and right"); + let left_right = find_mismatched_indices(left, right, ctx); + let executed_right = find_mismatched_indices(&executed, right, ctx); - if !left_right_the_same || !executed_right_the_same { - let left_right = find_mismatched_indices(left, right, ctx); - let executed_right = find_mismatched_indices(&executed, right, ctx); + if !left_right.is_empty() || !executed_right.is_empty() { let mut msg = String::new(); if !left_right.is_empty() { msg.push_str(&format!( From ee5e96ac8e0957116acfa81b266459fc933d7fc4 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 17:19:12 +0100 Subject: [PATCH 08/12] Pass ExecutionCtx to assert_arrays_eq in ext_storage tests The new ext_storage compute function's tests used the pre-ctx 2-arg assert_arrays_eq!; thread a ctx via array_execution_ctx() to match the ExecutionCtx-requiring macro on this branch. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Robert Kruszewski --- vortex-array/src/scalar_fn/fns/ext_storage.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/ext_storage.rs b/vortex-array/src/scalar_fn/fns/ext_storage.rs index b646c9e53f3..ec3b453a576 100644 --- a/vortex-array/src/scalar_fn/fns/ext_storage.rs +++ b/vortex-array/src/scalar_fn/fns/ext_storage.rs @@ -138,7 +138,7 @@ mod tests { result.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_arrays_eq!(result, storage); + assert_arrays_eq!(result, storage, &mut crate::array_execution_ctx()); Ok(()) } @@ -154,7 +154,7 @@ mod tests { result.dtype(), &DType::Primitive(PType::I64, Nullability::Nullable) ); - assert_arrays_eq!(result, storage); + assert_arrays_eq!(result, storage, &mut crate::array_execution_ctx()); Ok(()) } @@ -171,7 +171,11 @@ mod tests { result.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_arrays_eq!(result, ConstantArray::new(storage_scalar, 3)); + assert_arrays_eq!( + result, + ConstantArray::new(storage_scalar, 3), + &mut crate::array_execution_ctx() + ); Ok(()) } From 8c78821a0d12d6d222d8a6209805e34038ae4689 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 18:09:24 +0100 Subject: [PATCH 09/12] more Signed-off-by: Robert Kruszewski --- .../src/aggregate_fn/fns/sum/grouped.rs | 10 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 12 +- vortex-array/src/array/typed.rs | 5 +- vortex-array/src/arrays/bool/array.rs | 8 +- .../src/arrays/bool/compute/filter.rs | 9 +- vortex-array/src/arrays/bool/compute/take.rs | 11 +- vortex-array/src/arrays/bool/patch.rs | 5 +- vortex-array/src/arrays/bool/vtable/mod.rs | 3 +- .../src/arrays/bool/vtable/operations.rs | 3 +- vortex-array/src/arrays/chunked/array.rs | 8 +- .../src/arrays/chunked/compute/cast.rs | 3 +- .../src/arrays/chunked/compute/take.rs | 15 +- vortex-array/src/arrays/chunked/tests.rs | 20 +- .../src/arrays/chunked/vtable/operations.rs | 9 +- .../src/arrays/constant/compute/take.rs | 4 +- .../src/arrays/constant/vtable/canonical.rs | 2 +- vortex-array/src/arrays/datetime/test.rs | 5 +- .../src/arrays/decimal/compute/fill_null.rs | 8 +- .../src/arrays/decimal/compute/take.rs | 5 +- vortex-array/src/arrays/decimal/vtable/mod.rs | 3 +- vortex-array/src/arrays/dict/array.rs | 2 +- vortex-array/src/arrays/dict/compute/cast.rs | 6 +- .../src/arrays/dict/compute/fill_null.rs | 2 +- vortex-array/src/arrays/dict/compute/like.rs | 3 +- vortex-array/src/arrays/dict/compute/mod.rs | 8 +- vortex-array/src/arrays/dict/compute/slice.rs | 5 +- vortex-array/src/arrays/dict/tests.rs | 3 +- .../src/arrays/extension/compute/cast.rs | 2 +- .../arrays/filter/execute/fixed_size_list.rs | 7 +- .../src/arrays/filter/execute/listview.rs | 6 +- .../src/arrays/filter/execute/struct_.rs | 7 +- .../arrays/fixed_size_list/tests/filter.rs | 15 +- .../src/arrays/fixed_size_list/tests/take.rs | 8 +- vortex-array/src/arrays/list/tests.rs | 16 +- .../src/arrays/listview/conversion.rs | 14 +- .../src/arrays/listview/tests/basic.rs | 6 +- .../src/arrays/listview/tests/filter.rs | 4 +- .../src/arrays/listview/tests/operations.rs | 6 +- .../src/arrays/listview/tests/take.rs | 4 +- vortex-array/src/arrays/masked/tests.rs | 2 +- .../src/arrays/patched/compute/take.rs | 10 +- .../src/arrays/primitive/array/patch.rs | 3 +- .../src/arrays/primitive/compute/cast.rs | 12 +- .../src/arrays/primitive/compute/fill_null.rs | 8 +- .../src/arrays/primitive/compute/take/mod.rs | 3 +- .../src/arrays/primitive/vtable/mod.rs | 3 +- .../src/arrays/scalar_fn/vtable/operations.rs | 8 +- vortex-array/src/arrays/slice/vtable.rs | 3 +- .../src/arrays/struct_/compute/mod.rs | 4 +- .../src/arrays/struct_/compute/rules.rs | 6 +- vortex-array/src/arrays/struct_/tests.rs | 6 +- .../src/arrays/varbin/compute/compare.rs | 5 +- .../src/arrays/varbin/compute/filter.rs | 10 +- .../src/arrays/varbin/compute/take.rs | 3 +- vortex-array/src/arrays/varbin/tests.rs | 5 +- .../src/arrays/varbin/vtable/canonical.rs | 2 +- vortex-array/src/arrays/varbinview/compact.rs | 18 +- vortex-array/src/arrays/varbinview/tests.rs | 9 +- .../src/arrays/varbinview/vtable/mod.rs | 3 +- vortex-array/src/arrays/variant/mod.rs | 4 +- vortex-array/src/builders/bool.rs | 2 +- vortex-array/src/builders/decimal.rs | 2 +- vortex-array/src/builders/dict/primitive.rs | 4 +- vortex-array/src/builders/extension.rs | 3 +- vortex-array/src/builders/listview.rs | 2 +- vortex-array/src/builders/primitive.rs | 4 +- vortex-array/src/builders/struct_.rs | 2 +- vortex-array/src/builders/varbinview.rs | 6 +- .../src/compute/conformance/filter.rs | 2 +- vortex-array/src/lib.rs | 8 - vortex-array/src/patches.rs | 42 +-- vortex-array/src/scalar_fn/fns/between/mod.rs | 2 +- .../src/scalar_fn/fns/binary/boolean.rs | 4 +- .../src/scalar_fn/fns/binary/compare.rs | 12 +- vortex-array/src/scalar_fn/fns/binary/mod.rs | 10 +- .../src/scalar_fn/fns/binary/numeric.rs | 16 +- vortex-array/src/scalar_fn/fns/byte_length.rs | 4 +- vortex-array/src/scalar_fn/fns/case_when.rs | 46 +-- vortex-array/src/scalar_fn/fns/dynamic.rs | 9 +- vortex-array/src/scalar_fn/fns/ext_storage.rs | 15 +- .../src/scalar_fn/fns/fill_null/mod.rs | 7 +- .../src/scalar_fn/fns/list_contains/mod.rs | 16 +- vortex-array/src/scalar_fn/fns/merge.rs | 3 +- vortex-array/src/scalar_fn/fns/pack.rs | 5 +- vortex-array/src/scalar_fn/fns/zip/mod.rs | 6 +- vortex-array/src/stats/expr.rs | 32 +- vortex-btrblocks/src/canonical_compressor.rs | 16 +- vortex-btrblocks/src/schemes/float/tests.rs | 2 +- .../schemes/integer/scheme_selection_tests.rs | 2 +- vortex-btrblocks/src/schemes/integer/tests.rs | 6 +- vortex-compressor/src/sample.rs | 3 +- vortex-cuda/src/dynamic_dispatch/mod.rs | 118 ++++--- vortex-cuda/src/hybrid_dispatch/mod.rs | 39 ++- vortex-cuda/src/kernel/arrays/constant.rs | 25 +- vortex-cuda/src/kernel/arrays/dict.rs | 292 +++++------------- vortex-cuda/src/kernel/encodings/alp.rs | 54 ++-- vortex-cuda/src/kernel/encodings/bitpacked.rs | 157 +++------- .../src/kernel/encodings/date_time_parts.rs | 25 +- .../kernel/encodings/decimal_byte_parts.rs | 14 +- vortex-cuda/src/kernel/encodings/for_.rs | 16 +- vortex-cuda/src/kernel/encodings/fsst.rs | 15 +- vortex-cuda/src/kernel/encodings/runend.rs | 40 +-- vortex-cuda/src/kernel/encodings/sequence.rs | 13 +- vortex-cuda/src/kernel/encodings/zigzag.rs | 12 +- vortex-cuda/src/kernel/encodings/zstd.rs | 33 +- .../src/kernel/encodings/zstd_buffers.rs | 26 +- vortex-cuda/src/kernel/filter/decimal.rs | 19 +- vortex-cuda/src/kernel/filter/primitive.rs | 19 +- vortex-cuda/src/kernel/filter/varbinview.rs | 9 +- vortex-cuda/src/kernel/patches/types.rs | 5 +- vortex-cuda/src/lib.rs | 12 - vortex-duckdb/src/convert/vector.rs | 2 +- vortex-ffi/src/scan.rs | 8 +- vortex-file/src/tests.rs | 6 +- vortex-file/tests/test_write_table.rs | 2 +- vortex-ipc/src/iterator.rs | 3 +- vortex-ipc/src/stream.rs | 7 +- vortex-layout/src/layouts/chunked/reader.rs | 3 +- vortex-layout/src/layouts/dict/reader.rs | 10 +- vortex-layout/src/layouts/flat/reader.rs | 7 +- vortex-layout/src/layouts/row_idx/mod.rs | 7 +- vortex-layout/src/layouts/struct_/reader.rs | 6 +- vortex-layout/src/layouts/zoned/reader.rs | 3 +- vortex-python/src/arrays/range_to_sequence.rs | 3 +- vortex-test/compat-gen/src/check.rs | 3 +- .../src/fixtures/arrays/synthetic/mod.rs | 4 +- 126 files changed, 731 insertions(+), 978 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 597615dbad9..4f13c76777f 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -236,7 +236,7 @@ mod tests { #[test] fn listview_matches_reference_unsigned() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1u32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::U32, NonNullable); @@ -256,7 +256,7 @@ mod tests { #[test] fn listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Offsets are not in group order and a group is null: the group validity must be indexed by // group index, not by element offset. let elements = @@ -278,7 +278,7 @@ mod tests { #[test] fn listview_interior_and_full_nulls() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Group 1 has an interior null, group 2 is entirely null, group 3 is empty. let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) @@ -300,7 +300,7 @@ mod tests { #[test] fn listview_overflow_group_is_null() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I64, NonNullable); @@ -384,7 +384,7 @@ mod tests { #[test] fn fixed_size_overflow_and_nan() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // FixedSize path: first group overflows -> null sum, second sums normally. let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 26d3d4c19b5..242f49f5073 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -586,7 +586,7 @@ mod tests { #[test] fn grouped_sum_fixed_size_list() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; @@ -601,7 +601,7 @@ mod tests { #[test] fn grouped_sum_with_null_elements() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) .into_array(); @@ -617,7 +617,7 @@ mod tests { #[test] fn grouped_sum_with_null_group() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) .into_array(); @@ -635,7 +635,7 @@ mod tests { #[test] fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; @@ -650,7 +650,7 @@ mod tests { #[test] fn grouped_sum_bool() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); let groups = FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; @@ -691,7 +691,7 @@ mod tests { #[test] fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 830ff7451d9..8ed98a254b3 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -489,6 +489,7 @@ mod tests { use vortex_buffer::buffer; use super::Array; + use crate::VortexSessionExecute; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -496,7 +497,7 @@ mod tests { #[test] fn typed_array_into_parts_roundtrips() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let expected = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); @@ -508,7 +509,7 @@ mod tests { #[test] fn typed_array_try_into_parts_requires_unique_owner() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let alias = array.clone(); diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index fa7239967ac..ab8e9611d36 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -449,7 +449,7 @@ mod tests { #[test] fn patch_sliced_bools() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.slice(4..12).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut assertion_ctx); @@ -497,7 +497,7 @@ mod tests { #[test] fn slice_array_in_middle() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let sliced = arr.slice(4..12).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut assertion_ctx); @@ -505,7 +505,7 @@ mod tests { #[test] fn patch_bools_owned() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let buf_ptr = arr.to_bit_buffer().inner().as_ptr(); @@ -530,7 +530,7 @@ mod tests { #[test] fn patch_sliced_bools_offset() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.slice(4..15).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11]), &mut assertion_ctx); diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index 666c33a2067..e5a091c7e29 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -303,13 +303,14 @@ mod tests { use super::*; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; #[test] fn filter_bool_test() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_iter([true, false, true]); @@ -323,7 +324,7 @@ mod tests { #[test] fn filter_bool_sparse_index_mask() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_indices(3, [0, 2]); @@ -337,7 +338,7 @@ mod tests { #[test] fn filter_bool_sparse_slice_mask() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_slices(3, vec![(0, 1), (2, 3)]); @@ -351,7 +352,7 @@ mod tests { #[test] fn filter_bool_sparse_buffer_mask() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_buffer(BitBuffer::from_iter([true, false, true])); diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index 2894bde97f6..f14f13a29c3 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -90,6 +90,7 @@ mod test { use crate::IntoArray as _; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; @@ -99,7 +100,7 @@ mod test { #[test] fn take_nullable() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let reference = BoolArray::from_iter(vec![ Some(false), Some(true), @@ -129,7 +130,7 @@ mod test { #[test] fn test_bool_array_take_with_null_out_of_bounds_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -147,7 +148,7 @@ mod test { #[test] fn test_non_null_bool_array_take_with_null_out_of_bounds_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -164,7 +165,7 @@ mod test { #[test] fn test_bool_array_take_all_null_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -180,7 +181,7 @@ mod test { #[test] fn test_non_null_bool_array_take_all_null_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], diff --git a/vortex-array/src/arrays/bool/patch.rs b/vortex-array/src/arrays/bool/patch.rs index 188bcfd39d0..96652c4656d 100644 --- a/vortex-array/src/arrays/bool/patch.rs +++ b/vortex-array/src/arrays/bool/patch.rs @@ -47,12 +47,13 @@ mod tests { use vortex_buffer::BitBuffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; #[test] fn patch_sliced_bools() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.into_array().slice(4..12).unwrap(); let expected = BoolArray::from_iter([true; 8]); @@ -61,7 +62,7 @@ mod tests { #[test] fn patch_sliced_bools_offset() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.into_array().slice(4..15).unwrap(); let expected = BoolArray::from_iter([true; 11]); diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index d2bfef2ab26..b48f651ff5b 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -212,6 +212,7 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -219,7 +220,7 @@ mod tests { #[test] fn test_nullable_bool_serde_roundtrip() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = BoolArray::from_iter([Some(true), None, Some(false), None]); let dtype = array.dtype().clone(); let len = array.len(); diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index f8927dfbaec..a6ac9641ba7 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -30,6 +30,7 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; @@ -46,7 +47,7 @@ mod tests { #[test] fn test_slice() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); #[expect(deprecated)] let sliced_arr = arr.into_array().slice(1..4).unwrap().to_bool(); diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 503771ac9d5..d06a16556eb 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -291,7 +291,7 @@ mod test { #[test] fn test_rechunk_one_chunk() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -305,7 +305,7 @@ mod test { #[test] fn test_rechunk_two_chunks() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array(), buffer![5u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -320,7 +320,7 @@ mod test { #[test] fn test_rechunk_tiny_target_chunks() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2, 3].into_array(), @@ -339,7 +339,7 @@ mod test { #[test] fn test_rechunk_with_too_big_chunk() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2].into_array(), diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index 90c7f8c0a95..dbccad6ae02 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -35,6 +35,7 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -46,7 +47,7 @@ mod test { #[test] fn test_cast_chunked() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr0 = buffer![0u32, 1].into_array(); let arr1 = buffer![2u32, 3].into_array(); diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 7dcc9a57a07..e19811b7379 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -128,6 +128,7 @@ mod test { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; @@ -141,7 +142,7 @@ mod test { #[test] fn test_take() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -159,7 +160,7 @@ mod test { #[test] fn test_take_nullable_values() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let a = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -177,7 +178,7 @@ mod test { #[test] fn test_take_nullable_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -198,7 +199,7 @@ mod test { #[test] fn test_take_nullable_struct() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let struct_array = StructArray::try_new(FieldNames::default(), vec![], 100, Validity::NonNullable) .unwrap(); @@ -224,7 +225,7 @@ mod test { #[test] fn test_empty_take() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -245,7 +246,7 @@ mod test { #[test] fn test_take_shuffled_indices() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let c0 = buffer![0i32, 1, 2].into_array(); let c1 = buffer![3i32, 4, 5].into_array(); let c2 = buffer![6i32, 7, 8].into_array(); @@ -314,7 +315,7 @@ mod test { #[test] fn test_take_null_indices() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let c0 = buffer![10i32, 20, 30].into_array(); let c1 = buffer![40i32, 50, 60].into_array(); let arr = ChunkedArray::try_new( diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index 255349adc68..5f3ef6db648 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -203,7 +203,7 @@ fn execute_path_nested_chunked_dict_of_dict_into_canonical() { #[test] fn with_slot_rewrites_chunk_and_offsets() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = chunked_array().into_array(); let replacement = buffer![10u64, 11, 12].into_array(); @@ -236,7 +236,7 @@ fn with_slot_rejects_len_mismatch() { #[test] fn slice_middle() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(2..5).unwrap(), PrimitiveArray::from_iter([3u64, 4, 5]), @@ -246,7 +246,7 @@ fn slice_middle() { #[test] fn slice_begin() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(1..3).unwrap(), PrimitiveArray::from_iter([2u64, 3]), @@ -256,7 +256,7 @@ fn slice_begin() { #[test] fn slice_aligned() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(3..6).unwrap(), PrimitiveArray::from_iter([4u64, 5, 6]), @@ -266,7 +266,7 @@ fn slice_aligned() { #[test] fn slice_many_aligned() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(0..6).unwrap(), PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]), @@ -276,7 +276,7 @@ fn slice_many_aligned() { #[test] fn slice_end() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(7..8).unwrap(), PrimitiveArray::from_iter([8u64]), @@ -286,7 +286,7 @@ fn slice_end() { #[test] fn slice_exactly_end() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(6..9).unwrap(), PrimitiveArray::from_iter([7u64, 8, 9]), @@ -304,7 +304,7 @@ fn slice_empty() { #[test] fn scalar_at_empty_children_both_sides() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -325,7 +325,7 @@ fn scalar_at_empty_children_both_sides() { #[test] fn scalar_at_empty_children_trailing() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -345,7 +345,7 @@ fn scalar_at_empty_children_trailing() { #[test] fn scalar_at_empty_children_leading() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 8960ff47540..12855434f13 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -30,6 +30,7 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -57,7 +58,7 @@ mod tests { #[case::end(7..8, &[8u64])] #[case::exactly_end(6..9, &[7u64, 8, 9])] fn slice(#[case] range: Range, #[case] expected: &[u64]) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(range).unwrap(), PrimitiveArray::from_iter(expected.iter().copied()), @@ -75,7 +76,7 @@ mod tests { #[test] fn scalar_at_empty_children_both_sides() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -96,7 +97,7 @@ mod tests { #[test] fn scalar_at_empty_children_trailing() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -116,7 +117,7 @@ mod tests { #[test] fn scalar_at_empty_children_leading() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index f7f6b6c51f2..3b2298b0bf1 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -87,7 +87,7 @@ mod tests { #[test] fn take_nullable_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take( @@ -125,7 +125,7 @@ mod tests { #[test] fn take_all_valid_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take(PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array()) diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index bcd5ca99d6b..8605877f5ac 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -384,7 +384,7 @@ mod tests { #[test] fn test_canonicalize_const_str() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let const_array = ConstantArray::new("four".to_string(), 4); let expected = VarBinArray::from(vec!["four", "four", "four", "four"]); diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index 084b6e517bf..e091ff3848d 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -10,6 +10,7 @@ use crate::EqMode; use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; +use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::datetime::TemporalData; use crate::assert_arrays_eq; @@ -26,7 +27,7 @@ use crate::validity::Validity; macro_rules! test_temporal_roundtrip { ($prim:ty, $constructor:expr, $unit:expr) => {{ - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![100 as $prim].into_array(); let temporal: TemporalData = $constructor(array, $unit); @@ -147,7 +148,7 @@ test_fail_case!( // We test Timestamp explicitly to avoid the macro getting too complex. #[test] fn test_timestamp() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let ts = buffer![100i64].into_array(); let ts_array = ts.into_array(); diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index 48d71545bcc..f12a1428eb5 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -104,7 +104,7 @@ mod tests { #[test] fn fill_null_leading_none() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( [None, Some(800i128), None, Some(1000i128), None], @@ -141,7 +141,7 @@ mod tests { #[test] fn fill_null_all_none() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( @@ -169,7 +169,7 @@ mod tests { /// fill_null with a value that overflows the array's storage type should upcast the array. #[test] fn fill_null_overflow_upcasts() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(3, 0); let arr = DecimalArray::from_option_iter([None, Some(10i8), None], decimal_dtype); // i8 max is 127, so 200 doesn't fit — the array should be widened to i16. @@ -192,7 +192,7 @@ mod tests { #[test] fn fill_null_non_nullable() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::new( diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index faffbf3433d..8f63637b1e9 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -53,6 +53,7 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -62,7 +63,7 @@ mod tests { #[test] fn test_take() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![10i128, 11i128, 12i128, 13i128], @@ -79,7 +80,7 @@ mod tests { #[test] fn test_take_null_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![i128::MAX, 11i128, 12i128, 13i128], diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index bc0c3295a80..e7f5e8c5f19 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -209,6 +209,7 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; @@ -248,7 +249,7 @@ mod tests { #[test] fn test_nullable_decimal_serde_roundtrip() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = DecimalArray::new( buffer![1234567i32, 0i32, -9999999i32], DecimalDType::new(7, 3), diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index bc8649c1962..79f1b538a63 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -462,7 +462,7 @@ mod test { #[test] fn test_dict_array_from_primitive_chunks() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let len = 2; let chunk_count = 2; let array = make_dict_primitive_chunks::(len, 2, chunk_count); diff --git a/vortex-array/src/arrays/dict/compute/cast.rs b/vortex-array/src/arrays/dict/compute/cast.rs index a86f197d629..6ce6846c389 100644 --- a/vortex-array/src/arrays/dict/compute/cast.rs +++ b/vortex-array/src/arrays/dict/compute/cast.rs @@ -71,7 +71,7 @@ mod tests { #[test] fn test_cast_dict_to_wider_type() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![1i32, 2, 3, 2, 1].into_array(); let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); @@ -111,7 +111,7 @@ mod tests { #[test] fn test_cast_dict_allvalid_to_nonnullable_and_back() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create an AllValid dict array (no nulls) let values = buffer![10i32, 20, 30, 40].into_array(); let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); @@ -193,7 +193,7 @@ mod tests { #[test] fn test_cast_dict_with_unreferenced_null_values_to_nonnullable() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::arrays::DictArray; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/dict/compute/fill_null.rs b/vortex-array/src/arrays/dict/compute/fill_null.rs index 70a6cef1716..40176f03f28 100644 --- a/vortex-array/src/arrays/dict/compute/fill_null.rs +++ b/vortex-array/src/arrays/dict/compute/fill_null.rs @@ -108,7 +108,7 @@ mod tests { #[test] fn nullable_codes_fill_in_values() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::new( buffer![0u32, 1, 2], diff --git a/vortex-array/src/arrays/dict/compute/like.rs b/vortex-array/src/arrays/dict/compute/like.rs index 8319c28c2e3..c87696597c3 100644 --- a/vortex-array/src/arrays/dict/compute/like.rs +++ b/vortex-array/src/arrays/dict/compute/like.rs @@ -56,6 +56,7 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::DictArray; use crate::arrays::VarBinArray; @@ -68,7 +69,7 @@ mod tests { #[test] fn like_reduce_dict() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let dict = DictArray::try_new( buffer![0u8, 1, 0, 2].into_array(), VarBinArray::from(vec!["hello", "world", "help"]).into_array(), diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index 421642a8a18..ed0f3758009 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -86,7 +86,7 @@ mod test { #[test] fn canonicalise_nullable_primitive() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values: Vec> = (0..65) .map(|i| match i % 3 { 0 => Some(42), @@ -111,7 +111,7 @@ mod test { #[test] fn canonicalise_non_nullable_primitive_32_unique_values() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let unique_values: Vec = (0..32).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 32])); @@ -128,7 +128,7 @@ mod test { #[test] fn canonicalise_non_nullable_primitive_100_unique_values() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let unique_values: Vec = (0..100).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 100])); @@ -183,7 +183,7 @@ mod test { #[test] fn compare_sliced_dict() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::arrays::BoolArray; let sliced = sliced_dict_array(); let compared = sliced diff --git a/vortex-array/src/arrays/dict/compute/slice.rs b/vortex-array/src/arrays/dict/compute/slice.rs index 485346bf5d8..f557a36e0da 100644 --- a/vortex-array/src/arrays/dict/compute/slice.rs +++ b/vortex-array/src/arrays/dict/compute/slice.rs @@ -91,6 +91,7 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::dict::compute::slice::ConstantArray; @@ -102,7 +103,7 @@ mod tests { #[test] fn slice_constant_valid_code() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(1u8, 5).into_array(), buffer![10i32, 20, 30].into_array(), @@ -115,7 +116,7 @@ mod tests { #[test] fn slice_constant_null_code() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(Scalar::null(DType::Primitive(PType::U8, Nullable)), 5).into_array(), buffer![10i32, 20, 30].into_array(), diff --git a/vortex-array/src/arrays/dict/tests.rs b/vortex-array/src/arrays/dict/tests.rs index 5cbb08d5627..2e52ab20bf0 100644 --- a/vortex-array/src/arrays/dict/tests.rs +++ b/vortex-array/src/arrays/dict/tests.rs @@ -5,6 +5,7 @@ use vortex_buffer::buffer; use super::DictArray; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; @@ -14,7 +15,7 @@ use crate::validity::Validity; #[test] fn test_scalar_at_null_code() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::from_option_iter(vec![None, Some(0u32), None]).into_array(), buffer![1i32].into_array(), diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index d7378f7f095..4347594b797 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -121,7 +121,7 @@ mod tests { #[test] fn cast_timestamp_to_i64() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let ext_dtype = Timestamp::new_with_tz( TimeUnit::Nanoseconds, Some("UTC".into()), diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index d2626460ba1..7e293987d13 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -124,6 +124,7 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -149,7 +150,7 @@ mod test { #[test] fn filter_fixed_size_list_selects_correct_lists() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let array = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -166,7 +167,7 @@ mod test { #[test] fn filter_degenerate_list_size_zero() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::empty::(Nullability::NonNullable); let array = FixedSizeListArray::new(elements.into_array(), 0, Validity::NonNullable, 5); @@ -182,7 +183,7 @@ mod test { #[test] fn filter_nested_fixed_size_lists() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Inner lists of size 2, outer lists of size 2 (so 2 outer lists, each with 2 inner lists). let inner_elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); let inner_fsl = FixedSizeListArray::new(inner_elements, 2, Validity::NonNullable, 4); diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index 6f0f54edd66..0419a439183 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -135,7 +135,7 @@ mod test { #[test] fn filter_listview_selects_correct_lists() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // 3 lists: [10,20], [30,40], [50,60] let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let offsets = buffer![0u32, 2, 4].into_array(); @@ -163,7 +163,7 @@ mod test { #[test] fn test_filter_preserves_unreferenced_elements() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -198,7 +198,7 @@ mod test { #[test] fn test_filter_with_gaps() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index 71f573a5881..58600097343 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -49,6 +49,7 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -86,7 +87,7 @@ mod test { #[test] fn filter_struct_selects_correct_rows() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = StructArray::try_new( ["x", "y"].into(), vec![ @@ -117,7 +118,7 @@ mod test { #[test] fn filter_empty_struct() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let mask = Mask::from_iter([ @@ -132,7 +133,7 @@ mod test { #[test] fn filter_empty_struct_with_empty_filter() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); let filtered = struct_arr.filter(Mask::from_iter::<[bool; 0]>([])).unwrap(); diff --git a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs index 72c4dbda15b..5c0ac468254 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs @@ -8,6 +8,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -51,7 +52,7 @@ fn test_filter_degenerate_list_size_zero( #[case] mask_values: Vec, #[case] expected_len: usize, ) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let new_validity = if matches!(validity, Validity::NonNullable) { Validity::NonNullable } else { @@ -80,7 +81,7 @@ fn test_filter_degenerate_list_size_zero( #[test] fn test_filter_with_nulls() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4), Some(5), Some(6)]); let validity = Validity::from_iter([true, false, true]); @@ -108,7 +109,7 @@ fn test_filter_with_nulls() { #[test] fn test_filter_all_null_array() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create an array where all elements are null. let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let validity = Validity::AllInvalid; @@ -125,7 +126,7 @@ fn test_filter_all_null_array() { #[test] fn test_filter_nested_fixed_size_lists() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create nested fixed-size lists: FSL>. // Inner lists are of size 2, outer lists are of size 3. // So we have 2 outer lists, each containing 3 inner lists, each containing 2 i32s. @@ -251,7 +252,7 @@ fn create_fsl_empty() -> ArrayRef { #[test] fn test_filter_all_null_various_list_sizes() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test filtering with all-null arrays of different list sizes. // The implementation returns ConstantArray only when validity_mask() is Mask::AllFalse. @@ -290,7 +291,7 @@ fn test_filter_all_null_various_list_sizes() { #[test] fn test_mask_expansion_threshold_boundary() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test with list_size == 8 (the FSL_SPARSE_MASK_LIST_SIZE_THRESHOLD). let list_size = 8u32; let num_lists = 100; @@ -367,7 +368,7 @@ fn test_mask_expansion_threshold_boundary() { // Test FSL-specific behavior with very large list sizes. #[test] fn test_filter_large_list_size() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test with list_size=100, which is significantly larger than typical use cases. let list_size = 100u32; let num_lists = 5; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 8f6713fe98c..6748f99c1e8 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -40,7 +40,7 @@ fn test_take_fsl_conformance(#[case] fsl: FixedSizeListArray) { #[test] fn test_take_basic_smoke_test() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -108,7 +108,7 @@ fn test_take_degenerate_lists( #[test] fn test_take_large_list_size() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = buffer![0i32..300].into_array(); let fsl = FixedSizeListArray::new(elements, 100, Validity::NonNullable, 3); @@ -123,7 +123,7 @@ fn test_take_large_list_size() { #[test] fn test_take_fsl_with_null_indices_preserves_elements() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -170,7 +170,7 @@ fn test_element_index_overflow( #[case] indices: ArrayRef, #[case] expected: FixedSizeListArray, ) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let result = fsl.take(indices).unwrap(); assert_arrays_eq!(result, expected, &mut assertion_ctx); } diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 48a04f8522e..382e0d9ff14 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -113,7 +113,7 @@ fn test_simple_list_filter() { #[test] fn test_list_filter_dense_mask() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test filtering with a dense mask (high density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -145,7 +145,7 @@ fn test_list_filter_dense_mask() { #[test] fn test_list_filter_sparse_mask() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test filtering with a sparse mask (low density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -179,7 +179,7 @@ fn test_list_filter_sparse_mask() { #[test] fn test_list_filter_empty_lists() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test filtering arrays that contain empty lists. let elements = buffer![0..10].into_array(); let offsets = buffer![0, 0, 3, 3, 7, 10, 10].into_array(); // Lists at indices 0, 2, 5 are empty. @@ -236,7 +236,7 @@ fn test_list_filter_with_nulls() { #[test] fn test_list_filter_all_true() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test filtering with an all-true mask. let elements = buffer![0..20].into_array(); let offsets = buffer![0, 5, 10, 15, 20].into_array(); @@ -279,7 +279,7 @@ fn test_list_filter_all_false() { #[test] fn test_list_filter_single_element() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test filtering to keep only one element. let elements = buffer![0..50].into_array(); let offsets = buffer![0, 10, 20, 30, 40, 50].into_array(); @@ -308,7 +308,7 @@ fn test_list_filter_single_element() { #[test] fn test_list_filter_alternating_pattern() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test filtering with an alternating pattern. let elements = buffer![0..60].into_array(); let offsets = buffer![0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60].into_array(); @@ -344,7 +344,7 @@ fn test_list_filter_alternating_pattern() { #[test] fn test_list_filter_variable_sizes() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test filtering lists with highly variable sizes. let elements = buffer![0..100].into_array(); let offsets = buffer![0, 1, 2, 5, 10, 20, 35, 60, 100].into_array(); @@ -581,7 +581,7 @@ fn create_list_of_lists_nullable(data: OptVec>>) -> ListArray #[test] #[expect(clippy::cognitive_complexity)] fn test_list_of_lists() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let data = vec![ Some(vec![Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]), Some(vec![Some(vec![Some(4), Some(5), Some(6)])]), diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index b3f8807821e..e820812351c 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -348,7 +348,7 @@ mod tests { #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let list_view = create_basic_listview(); let list_array = list_from_list_view(list_view.clone(), &mut SESSION.create_execution_ctx())?; @@ -424,7 +424,7 @@ mod tests { #[test] fn test_non_zero_copy_listview_to_list() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create ListViewArray with overlapping lists (not zero-copyable). let list_view = create_overlapping_listview(); let list_array = @@ -558,7 +558,7 @@ mod tests { #[test] fn test_recursive_simple_listview() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let list_view = create_basic_listview(); let result = recursive_list_from_list_view( list_view.clone().into_array(), @@ -572,7 +572,7 @@ mod tests { #[test] fn test_recursive_nested_listview() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let inner_elements = buffer![1i32, 2, 3].into_array(); let inner_offsets = buffer![0u32, 2].into_array(); let inner_sizes = buffer![2u32, 1].into_array(); @@ -610,7 +610,7 @@ mod tests { #[test] fn test_recursive_struct_with_listview_fields() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let listview_field = create_basic_listview().into_array(); let primitive_field = buffer![10i32, 20, 30, 40].into_array(); @@ -633,7 +633,7 @@ mod tests { #[test] fn test_recursive_fixed_size_list_with_listview_elements() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let lv1_elements = buffer![1i32, 2].into_array(); let lv1_offsets = buffer![0u32].into_array(); let lv1_sizes = buffer![2u32].into_array(); @@ -679,7 +679,7 @@ mod tests { #[test] fn test_recursive_deep_nesting() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let innermost_elements = buffer![1i32, 2, 3].into_array(); let innermost_offsets = buffer![0u32, 2].into_array(); let innermost_sizes = buffer![2u32, 1].into_array(); diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index dddcf169ace..67a462ab910 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -26,7 +26,7 @@ use crate::validity::Validity; #[test] fn test_basic_listview_comprehensive() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Comprehensive test for basic ListView functionality including scalar_at. // Logical lists: [[1,2,3], [4,5], [6,7,8,9]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -83,7 +83,7 @@ fn test_basic_listview_comprehensive() { #[test] fn test_out_of_order_offsets() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // ListView-specific: Tests that offsets can be non-sequential and out-of-order. // Logical lists: [[7,8,9], [1,2,3], [4,5,6]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -167,7 +167,7 @@ fn test_from_list_array() -> VortexResult<()> { #[case::constant_offsets(false, true)] // Varying sizes, constant offsets #[case::both_constant(true, true)] // Both constant fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_offsets: bool) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Logical lists vary by case: // - constant_sizes: [[1,2,3], [4,5,6], [7,8,9]] (size 3 each, varying offsets) // - constant_offsets: [[1,2,3], [1,2], [1]] (all start at 0, varying sizes) diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index e1568d14d4b..6796297bf1c 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -39,7 +39,7 @@ fn test_filter_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_filter_preserves_unreferenced_elements() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -73,7 +73,7 @@ fn test_filter_preserves_unreferenced_elements() { #[test] fn test_filter_with_gaps() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index f488b285e6c..15716dcb883 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -84,7 +84,7 @@ fn test_slice_comprehensive() { #[test] fn test_slice_out_of_order() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // ListView-specific: Test slicing with out-of-order offsets. // Logical lists: [[70,80], [10,20,30], [40,50,60], [90], [30]] let elements = buffer![10i32, 20, 30, 40, 50, 60, 70, 80, 90].into_array(); @@ -394,7 +394,7 @@ fn test_cast_large_dataset() { #[test] fn test_zip_widens_false_element_nullability() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // [[1, 2], [3], [4]] let if_true = ListViewArray::new( buffer![1i32, 2, 3, 4].into_array(), @@ -440,7 +440,7 @@ fn test_zip_widens_false_element_nullability() -> VortexResult<()> { #[test] fn test_zip_widens_true_element_nullability() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // [[1, null], [3], [4]] let if_true = ListViewArray::new( PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(), diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 1c2fd05b3a4..53c1eb68f9d 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -40,7 +40,7 @@ fn test_take_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_take_preserves_unreferenced_elements() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // ListView-specific: Test that take preserves the entire elements array // even when taking only a subset of lists. let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -72,7 +72,7 @@ fn test_take_preserves_unreferenced_elements() { #[test] fn test_take_with_gaps() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // ListView-specific: Test with gaps in elements array. // Elements with gaps (999 values are "gaps" between used ranges). let elements = buffer![1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12].into_array(); diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index 19d35bfa4c6..316020a0761 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -95,7 +95,7 @@ fn test_masked_child_with_validity() { #[test] fn test_masked_child_all_valid() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // When validity is AllValid, masked_child should invert to AllInvalid. let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); diff --git a/vortex-array/src/arrays/patched/compute/take.rs b/vortex-array/src/arrays/patched/compute/take.rs index 078bfc2b4f0..b985493ee83 100644 --- a/vortex-array/src/arrays/patched/compute/take.rs +++ b/vortex-array/src/arrays/patched/compute/take.rs @@ -168,7 +168,7 @@ mod tests { #[test] fn test_take_basic() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -185,7 +185,7 @@ mod tests { #[test] fn test_take_sliced() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = make_patched_array(&[0; 10], &[1, 3], &[100, 200], 2..10)?; let indices = buffer![0u32, 1, 2, 3, 7].into_array(); @@ -200,7 +200,7 @@ mod tests { #[test] fn test_take_out_of_order() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -217,7 +217,7 @@ mod tests { #[test] fn test_take_duplicates() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at index [2] with value [99] let array = make_patched_array(&[0; 5], &[2], &[99], 0..5)?; @@ -238,7 +238,7 @@ mod tests { #[test] fn test_take_with_null_indices() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::arrays::BoolArray; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index 32c5964f3bc..c9d9eab81b4 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -137,6 +137,7 @@ mod tests { use super::*; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::assert_arrays_eq; use crate::validity::Validity; @@ -173,7 +174,7 @@ mod tests { #[test] fn patch_sliced() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid); let sliced = input.slice(2..8).unwrap(); #[expect(deprecated)] diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 2657b68f73c..02e15fb08bc 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -286,7 +286,7 @@ mod test { #[test] fn cast_u32_u8() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); // cast from u32 to u8 @@ -358,7 +358,7 @@ mod test { #[test] fn cast_u32_f32() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); #[expect(deprecated)] let u8arr = arr.cast(PType::F32.into()).unwrap().to_primitive(); @@ -400,7 +400,7 @@ mod test { #[test] fn cast_with_invalid_nulls() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![-1i32, 0, 10], Validity::from_iter([false, true, true]), @@ -430,7 +430,7 @@ mod test { /// buffer without allocation (pointer identity). #[test] fn cast_same_width_int_reinterprets_buffer() -> vortex_error::VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let src = PrimitiveArray::from_iter([0u32, 10, 100]); let src_ptr = src.as_slice::().as_ptr(); @@ -480,7 +480,7 @@ mod test { /// not prevent the cast from succeeding. #[test] fn cast_same_width_int_nullable_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // The null position holds u32::MAX which doesn't fit in i32, but it's // masked as invalid so the cast should still succeed via reinterpret. let arr = PrimitiveArray::new( @@ -502,7 +502,7 @@ mod test { #[test] fn cast_u32_to_u8_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1000u32, 10u32, 42u32], Validity::from_iter([false, true, true]), diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index 0a1d5c01ff0..3fc83be2e3f 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -64,7 +64,7 @@ mod test { #[test] fn fill_null_leading_none() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([None, Some(8u8), None, Some(10), None]); #[expect(deprecated)] let p = arr @@ -89,7 +89,7 @@ mod test { #[test] fn fill_null_all_none() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Option::::None, None, None, None, None]); #[expect(deprecated)] @@ -115,7 +115,7 @@ mod test { #[test] fn fill_null_nullable_non_null() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![8u8, 10, 12, 14, 16], Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()), @@ -143,7 +143,7 @@ mod test { #[test] fn fill_null_non_nullable() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = buffer![8u8, 10, 12, 14, 16].into_array(); #[expect(deprecated)] let p = arr.fill_null(Scalar::from(255u8)).unwrap().to_primitive(); diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 7bc5ff453be..5836d2b26f3 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -222,6 +222,7 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -229,7 +230,7 @@ mod tests { #[test] fn take_null_index_skips_out_of_bounds_value() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter([10i32, 20, 30]); let indices = PrimitiveArray::new( buffer![1u64, 3], diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index e64ac0558bb..a10bffc9fef 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -226,6 +226,7 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -234,7 +235,7 @@ mod tests { #[test] fn test_nullable_primitive_serde_roundtrip() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::new( buffer![1i32, 2, 3, 4], Validity::from_iter([true, false, true, false]), diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index 2e063ba85f5..f3dd679bcc3 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -76,7 +76,7 @@ mod tests { #[test] fn test_scalar_fn_add() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let lhs = buffer![1i32, 2, 3].into_array(); let rhs = buffer![10i32, 20, 30].into_array(); @@ -131,7 +131,7 @@ mod tests { #[test] fn test_scalar_fn_mul() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let lhs = buffer![2i32, 3, 4].into_array(); let rhs = buffer![5i32, 6, 7].into_array(); @@ -150,7 +150,7 @@ mod tests { #[test] fn test_scalar_fn_with_nullable() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let rhs = PrimitiveArray::new( buffer![10i32, 20, 30], @@ -177,7 +177,7 @@ mod tests { #[test] fn test_scalar_fn_comparison() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let lhs = buffer![1i32, 5, 3].into_array(); let rhs = buffer![2i32, 5, 1].into_array(); diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index ef939d8ed1d..a4193a9502f 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -188,13 +188,14 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::assert_arrays_eq; #[test] fn test_slice_slice() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Slice(1..4, Slice(2..8, base)) combines to Slice(3..6, base) let arr = PrimitiveArray::from_iter(0i32..10).into_array(); let inner_slice = SliceArray::new(arr, 2..8).into_array(); diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index ed4ebc782b0..ed533948490 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -40,7 +40,7 @@ mod tests { #[test] fn take_empty_struct() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let indices = PrimitiveArray::from_option_iter([Some(1), None]); @@ -94,7 +94,7 @@ mod tests { #[test] fn take_field_struct() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let struct_arr = StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).into_array())]) .unwrap(); diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 6f46515433e..dd2ee10b400 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -165,7 +165,7 @@ mod tests { #[test] fn test_struct_cast_field_reorder() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Source: {a, b}, Target: {c, b, a} - reordered + new null field let source = StructArray::try_new( FieldNames::from(["a", "b"]), @@ -302,7 +302,7 @@ mod tests { #[test] fn cast_struct_drop_field() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Casting to a struct with a subset of fields should succeed. let source = StructArray::try_new( FieldNames::from(["a", "b", "c"]), @@ -348,7 +348,7 @@ mod tests { #[test] fn cast_struct_field_type_widening() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Casting struct fields to wider types (i32 -> i64). let source = StructArray::try_new( FieldNames::from(["val"]), diff --git a/vortex-array/src/arrays/struct_/tests.rs b/vortex-array/src/arrays/struct_/tests.rs index 947223bc7a7..8272dd439c8 100644 --- a/vortex-array/src/arrays/struct_/tests.rs +++ b/vortex-array/src/arrays/struct_/tests.rs @@ -24,7 +24,7 @@ use crate::validity::Validity; #[test] fn test_project() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = VarBinArray::from_vec( vec!["a", "b", "c", "d", "e"], @@ -67,7 +67,7 @@ fn test_project() { #[test] fn test_remove_column() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = PrimitiveArray::new(buffer![4u64, 5, 6, 7, 8], Validity::NonNullable); @@ -113,7 +113,7 @@ fn test_remove_column() { #[test] fn test_duplicate_field_names() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test that StructArray allows duplicate field names and returns the first match let field1 = buffer![1i32, 2, 3].into_array(); let field2 = buffer![10i32, 20, 30].into_array(); diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index 22b61adaefb..ab75006b3bb 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -254,6 +254,7 @@ mod tests { use vortex_buffer::ByteBuffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; @@ -289,7 +290,7 @@ mod tests { /// [`CompareKernel`]: super::CompareKernel #[test] fn varbin_i64_offsets_compare_constant() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -310,7 +311,7 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant_binary() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index 8784540b996..cef5e6a41c2 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -240,7 +240,7 @@ mod test { #[test] fn filter_var_bin_test() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -267,7 +267,7 @@ mod test { #[test] fn filter_var_bin_slice_test() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -297,7 +297,7 @@ mod test { #[test] fn filter_var_bin_slice_null() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let bytes = [ b"one".as_slice(), b"two".as_slice(), @@ -340,7 +340,7 @@ mod test { #[test] fn filter_varbin_nulls() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let bytes = [b"".as_slice(), b"two".as_slice(), b"two".as_slice()] .into_iter() .flat_map(|x| x.iter().cloned()) @@ -368,7 +368,7 @@ mod test { #[test] fn filter_varbin_all_null() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let offsets = buffer![0, 0, 0, 0].into_array(); let validity = Validity::Array(BoolArray::from_iter([false, false, false]).into_array()); let arr = VarBinArray::try_new( diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index b0e5d6ec731..97fc205cfa8 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -256,6 +256,7 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::compute::take::PrimitiveArray; @@ -304,7 +305,7 @@ mod tests { #[test] fn test_take_overflow() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let scream = std::iter::once("a").cycle().take(128).collect::(); let bytes = ByteBuffer::copy_from(scream.as_bytes()); let offsets = buffer![0u8, 128u8].into_array(); diff --git a/vortex-array/src/arrays/varbin/tests.rs b/vortex-array/src/arrays/varbin/tests.rs index 5b1e9b0b529..59b4d8354d1 100644 --- a/vortex-array/src/arrays/varbin/tests.rs +++ b/vortex-array/src/arrays/varbin/tests.rs @@ -8,6 +8,7 @@ use vortex_buffer::buffer; use crate::ArrayRef; use crate::IntoArray; +use crate::VortexSessionExecute; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; @@ -32,7 +33,7 @@ fn binary_array() -> ArrayRef { #[rstest] pub fn test_scalar_at(binary_array: ArrayRef) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); assert_arrays_eq!( binary_array, VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), @@ -42,7 +43,7 @@ pub fn test_scalar_at(binary_array: ArrayRef) { #[rstest] pub fn slice_array(binary_array: ArrayRef) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let binary_arr = binary_array.slice(1..2).unwrap(); assert_arrays_eq!( binary_arr, diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index c0a69629945..12f25183e12 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -96,7 +96,7 @@ mod tests { #[case(DType::Utf8(Nullability::NonNullable))] #[case(DType::Binary(Nullability::NonNullable))] fn test_canonical_varbin_unsliced(#[case] dtype: DType) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let varbin = VarBinArray::from_iter_nonnull(["foo", "bar", "baz"], dtype.clone()); #[expect(deprecated)] let canonical = varbin.as_array().to_varbinview(); diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index cdc1a480b57..894f604f567 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -202,7 +202,7 @@ mod tests { use crate::dtype::Nullability; #[test] fn test_optimize_compacts_buffers() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create a VarBinViewArray with some long strings that will create multiple buffers let original = VarBinViewArray::from_iter_nullable_str([ Some("short"), @@ -243,7 +243,7 @@ mod tests { #[test] fn test_optimize_with_long_strings() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create strings that are definitely longer than 12 bytes let long_string_1 = "this is definitely a very long string that exceeds the inline limit"; let long_string_2 = "another extremely long string that also needs external buffer storage"; @@ -280,7 +280,7 @@ mod tests { #[test] fn test_optimize_no_buffers() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create an array with only short strings (all inlined) let original = VarBinViewArray::from_iter_str(["a", "bb", "ccc", "dddd"]); @@ -297,7 +297,7 @@ mod tests { #[test] fn test_optimize_single_buffer() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create an array that naturally has only one buffer let str1 = "this is a long string that goes into a buffer"; let str2 = "another long string in the same buffer"; @@ -317,7 +317,7 @@ mod tests { #[test] fn test_selective_compaction_with_threshold_zero() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // threshold=0 should keep all buffers (no compaction) let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -345,7 +345,7 @@ mod tests { #[test] fn test_selective_compaction_with_high_threshold() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // threshold=1.0 should compact any buffer with waste let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -374,7 +374,7 @@ mod tests { #[test] fn test_selective_compaction_preserves_well_utilized_buffers() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create an array with multiple strings in one buffer (well-utilized) let str1 = "first long string that needs external buffer storage"; let str2 = "second long string also in buffer"; @@ -397,7 +397,7 @@ mod tests { #[test] fn test_selective_compaction_with_mixed_utilization() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create array with some long strings let strings: Vec = (0..10) .map(|i| { @@ -429,7 +429,7 @@ mod tests { #[test] fn test_slice_strategy_with_contiguous_range() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create array with strings that will be in one buffer let strings: Vec = (0..20) .map(|i| format!("this is a long string number {} for slice test", i)) diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index d8772540a87..24574942d60 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -1,15 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#[expect(deprecated)] use crate::ToCanonical as _; +#[expect(deprecated)] +use crate::VortexSessionExecute; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; #[test] pub fn varbin_view() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]); assert_arrays_eq!( @@ -21,7 +22,7 @@ pub fn varbin_view() { #[test] pub fn slice_array() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) .slice(1..2) @@ -35,7 +36,7 @@ pub fn slice_array() { #[test] pub fn flatten_array() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["string1", "string2"]); #[expect(deprecated)] let var_bin = binary_arr.as_array().to_varbinview(); diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 194b1145573..b557922951c 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -249,6 +249,7 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::LEGACY_SESSION; + use crate::VortexSessionExecute; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -266,7 +267,7 @@ mod tests { let len = array.len(); let ctx = ArrayContext::empty(); - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let serialized = array .clone() .into_array() diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index aabdac06d95..e1f8831ab6b 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -379,7 +379,7 @@ mod tests { #[test] fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([( "a", @@ -407,7 +407,7 @@ mod tests { #[test] fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([ ( diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index 80aad067648..e0411e69166 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -221,7 +221,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10); // Test appending true value. diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index 0e3acd28774..f6986960f97 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -345,7 +345,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::scalar::Scalar; // Simply test that the builder accepts its own finish output via scalar. diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index a6af40ebc57..f6cb5f2621e 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -173,7 +173,7 @@ mod test { #[test] fn encode_primitive() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = buffer![1, 1, 3, 3, 3].into_array(); let dict = dict_encode(&arr, &mut SESSION.create_execution_ctx()).unwrap(); @@ -186,7 +186,7 @@ mod test { #[test] fn encode_primitive_nulls() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([ Some(1), Some(1), diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index dd3eba25a17..0888e899d23 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -127,6 +127,7 @@ impl ArrayBuilder for ExtensionBuilder { #[cfg(test)] mod tests { use super::*; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -137,7 +138,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); let mut builder = ExtensionBuilder::new(ext_dtype.clone()); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 846dfe3201c..a639594fc54 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -529,7 +529,7 @@ mod tests { #[test] fn test_different_offset_size_types() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test u32 offsets with u8 sizes. let dtype: Arc = Arc::new(I32.into()); let mut builder = diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 568c41c010d..12022925b75 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -385,7 +385,7 @@ mod tests { /// buffer. #[test] fn test_multiple_uninit_ranges_correct_offsets() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); // First range. @@ -496,7 +496,7 @@ mod tests { /// This verifies the new simplified API without the redundant `len` parameter. #[test] fn test_copy_from_slice_with_offsets() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); let mut range = builder.uninit_range(6); diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index a292b8c5e4f..1bddfbca369 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -267,7 +267,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::scalar::Scalar; let dtype = DType::Struct( diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index e3767fbccac..b5ff523fbc3 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -884,7 +884,7 @@ mod tests { #[test] fn test_utf8_builder() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); builder.append_value("Hello"); @@ -912,7 +912,7 @@ mod tests { #[test] fn test_utf8_builder_with_extend() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = { let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); @@ -992,7 +992,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::scalar::Scalar; // Test with Utf8 builder. diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index 1eadc1671c1..97d86481963 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -62,7 +62,7 @@ pub fn create_runs_pattern(len: usize, run_length: usize) -> Vec { /// Tests that filtering with an all-true mask returns all elements unchanged fn test_all_filter(array: &ArrayRef) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let len = array.len(); let mask = Mask::new_true(len); let filtered = array diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index c94b3de161b..16d29c69efe 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -117,14 +117,6 @@ pub fn array_session() -> VortexSession { .build() } -/// Builds a fresh session via [`array_session`] and returns an [`ExecutionCtx`](crate::ExecutionCtx) -/// backed by it. -/// -/// Convenience for the common `array_session().create_execution_ctx()` pattern. -pub fn array_execution_ctx() -> ExecutionCtx { - array_session().create_execution_ctx() -} - // TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot invoke execute // from the new array encodings to support back-compat for legacy encodings. So we hold a session // here... diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 784f77d133f..237da17f861 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -1222,7 +1222,7 @@ mod test { #[test] fn test_filter() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 100, 0, @@ -1254,7 +1254,7 @@ mod test { #[test] fn take_with_nulls() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1303,7 +1303,7 @@ mod test { #[test] fn take_search_with_nulls_chunked() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1352,7 +1352,7 @@ mod test { #[test] fn take_search_chunked_multiple_chunks() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 2048, 0, @@ -1403,7 +1403,7 @@ mod test { #[test] fn take_search_chunked_interleaved() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 30, 0, @@ -1455,7 +1455,7 @@ mod test { #[test] fn test_slice() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1472,7 +1472,7 @@ mod test { #[test] fn doubly_sliced() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1514,7 +1514,7 @@ mod test { #[test] fn test_mask_all_false() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1565,7 +1565,7 @@ mod test { #[test] fn test_mask_partial() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1602,7 +1602,7 @@ mod test { #[test] fn test_mask_with_offset() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1637,7 +1637,7 @@ mod test { #[test] fn test_mask_nullable_values() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1690,7 +1690,7 @@ mod test { #[test] fn test_filter_keep_all() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1740,7 +1740,7 @@ mod test { #[test] fn test_filter_with_indices() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1771,7 +1771,7 @@ mod test { #[test] fn test_slice_full_range() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1797,7 +1797,7 @@ mod test { #[test] fn test_slice_partial() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1842,7 +1842,7 @@ mod test { #[test] fn test_slice_with_offset() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1950,7 +1950,7 @@ mod test { #[test] fn test_mask_boundary_patches() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test masking patches at array boundaries let patches = Patches::new( 10, @@ -2005,7 +2005,7 @@ mod test { #[test] fn test_mask_no_patches_removed() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test when no patches are masked let patches = Patches::new( 10, @@ -2039,7 +2039,7 @@ mod test { #[test] fn test_mask_single_patch() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test with a single patch let patches = Patches::new( 5, @@ -2072,7 +2072,7 @@ mod test { #[test] fn test_mask_contiguous_patches() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test with contiguous patches let patches = Patches::new( 10, @@ -2106,7 +2106,7 @@ mod test { #[test] fn test_mask_with_large_offset() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test with a large offset that shifts all indices let patches = Patches::new( 20, diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index af80381e105..ad787456410 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -469,7 +469,7 @@ mod tests { #[test] fn test_between_decimal() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![100i128, 200i128, 300i128, 400i128]; let decimal_type = DecimalDType::new(3, 2); let array = DecimalArray::new(values, decimal_type, Validity::NonNullable).into_array(); diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index cb9c324293d..ebe9c59b442 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -772,7 +772,7 @@ mod tests { #[test] fn test_kleene_truth_table() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let lhs = BoolArray::from_iter([ Some(true), Some(true), @@ -835,7 +835,7 @@ mod tests { #[test] fn test_null_constant_kleene() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let lhs = BoolArray::from_iter([Some(false), Some(true), None]).into_array(); let null = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len()) .into_array(); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare.rs b/vortex-array/src/scalar_fn/fns/binary/compare.rs index e0985fd342f..20f8b70360c 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare.rs @@ -369,7 +369,7 @@ mod tests { #[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())] #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())] fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let res = left.binary(right, Operator::Eq).unwrap(); let expected = BoolArray::from_iter([true, true]); assert_arrays_eq!(res, expected, &mut assertion_ctx); @@ -377,7 +377,7 @@ mod tests { #[test] fn test_list_array_comparison() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values1 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets1 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list1 = ListArray::try_new( @@ -422,7 +422,7 @@ mod tests { #[test] fn test_list_array_constant_comparison() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list = ListArray::try_new( @@ -449,7 +449,7 @@ mod tests { #[test] fn test_struct_array_comparison() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let bool_field1 = BoolArray::from_iter([Some(true), Some(false), Some(true)]); let int_field1 = PrimitiveArray::from_iter([1i32, 2, 3]); @@ -486,7 +486,7 @@ mod tests { #[test] fn test_empty_struct_compare() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let empty1 = StructArray::try_new( FieldNames::from(Vec::::new()), Vec::new(), @@ -515,7 +515,7 @@ mod tests { /// different Vortex encodings (VarBinArray vs VarBinViewArray) must not panic. #[test] fn struct_compare_mixed_binary_encodings() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // LHS: struct with a VarBinArray (offset-based) binary field let bin_field1 = VarBinArray::from(vec![ "apple".as_bytes(), diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 5987e2cb8ca..8c7a425b169 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -519,7 +519,7 @@ mod tests { #[test] fn test_or_kleene_validity() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::StructArray; @@ -547,7 +547,7 @@ mod tests { #[test] fn test_scalar_subtract_unsigned() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -566,7 +566,7 @@ mod tests { #[test] fn test_scalar_subtract_signed() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -585,7 +585,7 @@ mod tests { #[test] fn test_scalar_subtract_nullable() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -602,7 +602,7 @@ mod tests { #[test] fn test_scalar_subtract_float() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index 52dc78d987a..2e96c49f830 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -951,7 +951,7 @@ mod test { #[test] fn test_scalar_subtract_unsigned() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![1u16, 2, 3].into_array(); let result = sub_scalar(&values, 1u16).unwrap(); assert_arrays_eq!( @@ -963,7 +963,7 @@ mod test { #[test] fn test_scalar_subtract_signed() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![1i64, 2, 3].into_array(); let result = sub_scalar(&values, -1i64).unwrap(); assert_arrays_eq!( @@ -975,7 +975,7 @@ mod test { #[test] fn test_scalar_subtract_nullable() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); let result = sub_scalar(&values.into_array(), Some(1u16)).unwrap(); assert_arrays_eq!( @@ -987,7 +987,7 @@ mod test { #[test] fn test_scalar_subtract_float() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![1.0f64, 2.0, 3.0].into_array(); let result = sub_scalar(&values, -1f64).unwrap(); assert_arrays_eq!( @@ -1006,7 +1006,7 @@ mod test { #[test] fn test_float_divide_by_zero_is_ok() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![1.0f64, -1.0].into_array(); let result = values .binary( @@ -1064,7 +1064,7 @@ mod test { #[test] fn test_integer_divide_errors_ignore_null_lanes() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![10i32, 10], Validity::from_iter([false, true])) .into_array(); let rhs = buffer![0i32, 2].into_array(); @@ -1085,7 +1085,7 @@ mod test { #[test] fn test_integer_errors_ignore_null_lanes() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true])) .into_array(); let result = values @@ -1123,7 +1123,7 @@ mod test { #[test] fn test_present_nullable_constant_preserves_nullable_output() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![1u8, 2].into_array(); let result = values .binary( diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-array/src/scalar_fn/fns/byte_length.rs index bfe2b11af4a..94447f356f5 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-array/src/scalar_fn/fns/byte_length.rs @@ -202,7 +202,7 @@ mod tests { #[case] array: ArrayRef, #[case] expected_lens: Vec, ) -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(expected_lens); assert_arrays_eq!(result, expected, &mut assertion_ctx); @@ -211,7 +211,7 @@ mod tests { #[test] fn test_varbinview_byte_length() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["short", "a longer string here"]).into_array(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 20]); diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index 5c37b8227ab..bdba538d410 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -783,7 +783,7 @@ mod tests { #[test] fn test_evaluate_simple_condition() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -805,7 +805,7 @@ mod tests { #[test] fn test_evaluate_nary_multiple_conditions() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Test n-ary via nested_case_when let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) @@ -830,7 +830,7 @@ mod tests { #[test] fn test_evaluate_nary_first_match_wins() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -855,7 +855,7 @@ mod tests { #[test] fn test_evaluate_no_else_returns_null() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -875,7 +875,7 @@ mod tests { #[test] fn test_evaluate_all_conditions_false() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -897,7 +897,7 @@ mod tests { #[test] fn test_evaluate_all_conditions_true() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -919,7 +919,7 @@ mod tests { #[test] fn test_evaluate_all_true_no_else_returns_correct_dtype() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // CASE WHEN value > 0 THEN 100 END — condition is always true, no ELSE. // Result must be Nullable because the implicit ELSE is NULL. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -943,7 +943,7 @@ mod tests { #[test] fn test_merge_case_branches_widens_nullability_of_later_branch() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // When a later THEN branch is Nullable and branches[0] and ELSE are NonNullable, // the result dtype must still be Nullable. // @@ -981,7 +981,7 @@ mod tests { #[test] fn test_evaluate_with_literal_condition() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = case_when(lit(true), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); @@ -995,7 +995,7 @@ mod tests { #[test] fn test_evaluate_with_bool_column_result() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1017,7 +1017,7 @@ mod tests { #[test] fn test_evaluate_with_nullable_condition() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]).into_array(), @@ -1037,7 +1037,7 @@ mod tests { #[test] fn test_evaluate_with_nullable_result_values() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[ ("value", buffer![1i32, 2, 3, 4, 5].into_array()), ( @@ -1066,7 +1066,7 @@ mod tests { #[test] fn test_evaluate_with_all_null_condition() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([None, None, None]).into_array(), @@ -1084,7 +1084,7 @@ mod tests { #[test] fn test_evaluate_nary_no_else_returns_null() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1111,7 +1111,7 @@ mod tests { #[test] fn test_evaluate_nary_many_conditions() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1139,7 +1139,7 @@ mod tests { #[test] fn test_evaluate_nary_all_false_no_else() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) .unwrap() .into_array(); @@ -1164,7 +1164,7 @@ mod tests { #[test] fn test_evaluate_nary_overlapping_conditions_first_wins() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![10i32, 20, 30].into_array())]) .unwrap() @@ -1189,7 +1189,7 @@ mod tests { #[test] fn test_evaluate_nary_early_exit_when_remaining_empty() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // After branch 0 claims all rows, remaining becomes all_false. // The loop breaks before evaluating branch 1's condition. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1215,7 +1215,7 @@ mod tests { #[test] fn test_evaluate_nary_skips_branch_with_empty_effective_mask() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Branch 0 claims value=1. Branch 1 targets the same rows but they are already // matched → effective_mask is all_false → branch 1 is skipped (THEN not used). let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1281,7 +1281,7 @@ mod tests { #[test] fn test_evaluate_nary_with_nullable_conditions() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[ ( "cond1", @@ -1402,7 +1402,7 @@ mod tests { #[test] fn test_simplify_null_fill_semantic_equivalence() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // The collapse-to-input rewrite must preserve values (and `x`'s nullability). let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1463,7 +1463,7 @@ mod tests { #[test] fn test_simplify_semantic_equivalence() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // The optimized expression must produce the same values as the original CASE WHEN. let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1492,7 +1492,7 @@ mod tests { #[test] fn test_merge_case_branches_alternating_mask() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Exercises the scalar path: alternating rows produce one slice per row (no runs), // triggering the per-row cursor path in merge_case_branches. let n = 100usize; diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index a9ebac9a2a2..700e6f7db21 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -277,6 +277,7 @@ mod tests { use super::*; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -303,7 +304,7 @@ mod tests { #[test] fn execute_with_value() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -323,7 +324,7 @@ mod tests { #[test] fn execute_without_value_default_true() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -343,7 +344,7 @@ mod tests { #[test] fn execute_without_value_default_false() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -363,7 +364,7 @@ mod tests { #[test] fn execute_value_flips() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let threshold = Arc::new(AtomicI32::new(5)); let threshold_clone = Arc::clone(&threshold); let expr = dynamic( diff --git a/vortex-array/src/scalar_fn/fns/ext_storage.rs b/vortex-array/src/scalar_fn/fns/ext_storage.rs index ec3b453a576..4a014a3325b 100644 --- a/vortex-array/src/scalar_fn/fns/ext_storage.rs +++ b/vortex-array/src/scalar_fn/fns/ext_storage.rs @@ -108,6 +108,7 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::PrimitiveArray; @@ -138,7 +139,11 @@ mod tests { result.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_arrays_eq!(result, storage, &mut crate::array_execution_ctx()); + assert_arrays_eq!( + result, + storage, + &mut crate::array_session().create_execution_ctx() + ); Ok(()) } @@ -154,7 +159,11 @@ mod tests { result.dtype(), &DType::Primitive(PType::I64, Nullability::Nullable) ); - assert_arrays_eq!(result, storage, &mut crate::array_execution_ctx()); + assert_arrays_eq!( + result, + storage, + &mut crate::array_session().create_execution_ctx() + ); Ok(()) } @@ -174,7 +183,7 @@ mod tests { assert_arrays_eq!( result, ConstantArray::new(storage_scalar, 3), - &mut crate::array_execution_ctx() + &mut crate::array_session().create_execution_ctx() ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index da5946a6228..2e5e934e550 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -179,6 +179,7 @@ mod tests { use vortex_error::VortexExpect; use crate::IntoArray; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::assert_arrays_eq; @@ -208,7 +209,7 @@ mod tests { #[test] fn evaluate() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]) .into_array(); @@ -229,7 +230,7 @@ mod tests { #[test] fn evaluate_struct_field() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[( "a", PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), @@ -253,7 +254,7 @@ mod tests { #[test] fn evaluate_non_nullable_input() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = fill_null(root(), lit(0i32)); let result = test_array.apply(&expr).unwrap(); diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index bcd663aa002..8edc4f13c97 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -752,7 +752,7 @@ mod tests { #[case] value: Option<&str>, #[case] expected: BoolArray, ) { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let element_nullability = list_array .dtype() .as_list_element_opt() @@ -770,7 +770,7 @@ mod tests { #[test] fn test_constant_list() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let list_array = ConstantArray::new( Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -789,7 +789,7 @@ mod tests { #[test] fn test_all_nulls() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let list_array = ConstantArray::new( Scalar::null(DType::List( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -811,7 +811,7 @@ mod tests { #[test] fn test_list_array_element() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let list_scalar = Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), vec![1.into(), 3.into(), 6.into()], @@ -828,7 +828,7 @@ mod tests { #[test] fn test_list_contains_empty_listview() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let empty_elements = PrimitiveArray::empty::(Nullability::NonNullable); let offsets = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); let sizes = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); @@ -852,7 +852,7 @@ mod tests { #[test] fn test_list_contains_all_null_elements() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter::([None, None, None, None, None]); let offsets = Buffer::from_iter([0u32, 2, 4]).into_array(); let sizes = Buffer::from_iter([2u32, 2, 1]).into_array(); @@ -888,7 +888,7 @@ mod tests { #[test] fn test_list_contains_large_offsets() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = Buffer::from_iter([1i32, 2, 3, 4, 5]).into_array(); let offsets = Buffer::from_iter([0u32, 1, 4, 0]).into_array(); @@ -912,7 +912,7 @@ mod tests { #[test] fn test_list_contains_offset_size_boundary() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let elements = Buffer::from_iter(0..256).into_array(); let offsets = Buffer::from_iter([0u8, 100, 200, 254]).into_array(); let sizes = Buffer::from_iter([50u8, 50, 54, 2]).into_array(); diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 0a13034b9f6..9dcfa7b5184 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -279,6 +279,7 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -318,7 +319,7 @@ mod tests { #[test] pub fn test_merge_right_most() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let expr = merge_opts( vec![ get_item("0", root()), diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index b3131429536..5ff70ffedc4 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -172,6 +172,7 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; + use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -230,7 +231,7 @@ mod tests { #[test] pub fn test_simple_pack() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -264,7 +265,7 @@ mod tests { #[test] pub fn test_nested_pack() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 580f24bc036..8d892c3add9 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -317,7 +317,7 @@ mod tests { #[test] fn test_zip_basic() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mask = Mask::from_iter([true, false, false, true, false]); let if_true = buffer![10, 20, 30, 40, 50].into_array(); let if_false = buffer![1, 2, 3, 4, 5].into_array(); @@ -330,7 +330,7 @@ mod tests { #[test] fn test_zip_all_true() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mask = Mask::new_true(4); let if_true = buffer![10, 20, 30, 40].into_array(); let if_false = @@ -346,7 +346,7 @@ mod tests { #[test] fn test_zip_all_false_widens_nullability() { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let mask = Mask::new_false(4); let if_true = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array(); diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 6e3a3f7a41a..0f84de369c7 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -109,7 +109,7 @@ mod tests { #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let sum_scalar = Scalar::primitive(6i64, Nullability::Nullable); array.statistics().set( @@ -131,7 +131,7 @@ mod tests { #[test] fn stat_expr_returns_null_when_sum_is_missing() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let result = array @@ -151,7 +151,7 @@ mod tests { #[test] fn stat_expr_reads_cached_sum_per_chunk() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let chunk0 = buffer![1i32, 2].into_array(); let sum_scalar = Scalar::primitive(3i64, Nullability::Nullable); chunk0.statistics().set( @@ -187,7 +187,7 @@ mod tests { #[test] fn stat_expr_reads_cached_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); let null_count_scalar = Scalar::primitive(2u64, Nullability::NonNullable); @@ -214,7 +214,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_null_from_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); array .statistics() @@ -234,7 +234,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_null_false_from_inexact_low_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, Some(2), None]).into_array(); array .statistics() @@ -254,7 +254,7 @@ mod tests { #[test] fn stat_expr_returns_null_for_inexact_full_null_count_as_all_null() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, Some(2), None]).into_array(); array .statistics() @@ -274,7 +274,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_non_null_from_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); array .statistics() @@ -294,7 +294,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_non_null_true_from_inexact_zero_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); array .statistics() @@ -314,7 +314,7 @@ mod tests { #[test] fn stat_expr_returns_null_for_inexact_nonzero_null_count_as_all_non_null() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); array @@ -348,7 +348,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_nan_from_nan_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(f32::NAN), Some(f32::NAN)]) .into_array(); @@ -370,7 +370,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_nan_false_from_inexact_low_nan_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32), Some(f32::NAN)]) .into_array(); @@ -392,7 +392,7 @@ mod tests { #[test] fn stat_expr_returns_null_for_inexact_full_nan_count_as_all_nan() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32), Some(f32::NAN)]) .into_array(); @@ -414,7 +414,7 @@ mod tests { #[test] fn stat_expr_reads_cached_all_non_nan_true_from_inexact_zero_nan_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1.0f32, 2.0, 3.0].into_array(); array .statistics() @@ -434,7 +434,7 @@ mod tests { #[test] fn stat_expr_returns_null_for_inexact_nonzero_nan_count_as_all_non_nan() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1.0f32), Some(f32::NAN), Some(3.0)]) .into_array(); array @@ -455,7 +455,7 @@ mod tests { #[test] fn stat_expr_reads_cached_min_and_max() -> VortexResult<()> { - let mut assertion_ctx = crate::array_execution_ctx(); + let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![3i32, 1, 2].into_array(); array .statistics() diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index 7c4e0c7992d..e182281ebb5 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -112,7 +112,7 @@ mod tests { #[case] input: ListViewArray, #[case] expect_list: bool, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array_ref = input.clone().into_array(); let result = BtrBlocksCompressor::default() .compress(&array_ref, &mut SESSION.create_execution_ctx())?; @@ -127,7 +127,7 @@ mod tests { #[test] fn test_constant_all_true() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -141,7 +141,7 @@ mod tests { #[test] fn test_constant_all_false() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -155,7 +155,7 @@ mod tests { #[test] fn test_nullable_all_valid_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true; 100]), Validity::from(BitBuffer::from(vec![true; 100])), @@ -172,7 +172,7 @@ mod tests { #[test] fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0))); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity); let btr = BtrBlocksCompressor::default(); @@ -187,7 +187,7 @@ mod tests { #[test] fn test_mixed_not_constant() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true, false, true]), Validity::NonNullable, @@ -204,7 +204,7 @@ mod tests { #[test] fn test_binary_constant_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let values = vec![Some(b"constant-bytes".as_slice()); 100]; let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); let btr = BtrBlocksCompressor::default(); @@ -219,7 +219,7 @@ mod tests { #[test] fn test_binary_dict_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"]; let values = (0..1000) .map(|idx| Some(distinct_values[idx % distinct_values.len()])) diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index a2745b2bce8..dfa210264c8 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -57,7 +57,7 @@ fn test_compress() -> VortexResult<()> { #[test] fn test_rle_compression() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(1.5f32, 100)); values.extend(iter::repeat_n(2.7f32, 200)); diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index ad8f29c5cdd..54f0565c5bf 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -161,7 +161,7 @@ fn test_rle_compressed() -> VortexResult<()> { #[cfg(feature = "unstable_encodings")] #[test] fn test_delta_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); use vortex_array::assert_arrays_eq; use vortex_fastlanes::Delta; diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index 518a4d7b193..ed03c017a6e 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -71,7 +71,7 @@ fn test_dict_encodable() -> VortexResult<()> { #[test] fn constant_mostly_nulls() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = PrimitiveArray::new( buffer![189u8, 189, 189, 189, 189, 189, 189, 189, 189, 0, 46], Validity::from_iter(vec![ @@ -95,7 +95,7 @@ fn constant_mostly_nulls() -> VortexResult<()> { #[test] fn nullable_sequence() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let values = (0i32..20).step_by(7).collect_vec(); let array = PrimitiveArray::from_option_iter(values.clone().into_iter().map(Some)); @@ -111,7 +111,7 @@ fn nullable_sequence() -> VortexResult<()> { #[test] fn test_rle_compression() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(42i32, 100)); values.extend(iter::repeat_n(123i32, 200)); diff --git a/vortex-compressor/src/sample.rs b/vortex-compressor/src/sample.rs index d12a60a3625..27d7383ade8 100644 --- a/vortex-compressor/src/sample.rs +++ b/vortex-compressor/src/sample.rs @@ -133,6 +133,7 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> #[cfg(test)] mod tests { use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; @@ -143,7 +144,7 @@ mod tests { #[test] fn sample_is_deterministic() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); // Create a deterministic array with linear-with-noise pattern let values: Vec = (0i64..100_000).map(|i| i + (i * 7 + 3) % 11).collect(); diff --git a/vortex-cuda/src/dynamic_dispatch/mod.rs b/vortex-cuda/src/dynamic_dispatch/mod.rs index dc1b0d0d01a..2d351d3ee58 100644 --- a/vortex-cuda/src/dynamic_dispatch/mod.rs +++ b/vortex-cuda/src/dynamic_dispatch/mod.rs @@ -1145,7 +1145,7 @@ mod tests { #[crate::test] async fn test_dict_mixed_width_u8_codes_u32_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let len = 3000; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u8).collect(); @@ -1171,14 +1171,14 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_dict_mixed_width_u16_codes_u32_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let dict_values: Vec = vec![1000, 2000, 3000, 4000, 5000]; let len = 2048; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u16).collect(); @@ -1204,14 +1204,14 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_runend_mixed_width_u64_ends_u32_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let ends: Vec = vec![1000, 2000, 3000]; let values: Vec = vec![10, 20, 30]; let len = 3000; @@ -1246,7 +1246,7 @@ mod tests { }) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -1634,7 +1634,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u8() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let bit_width: u8 = 4; let len = 3000; let reference = 100u8; @@ -1663,13 +1663,13 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_for_bitpacked_u16() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let reference = 1000u16; @@ -1700,13 +1700,13 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_for_bitpacked_u64() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let bit_width: u8 = 20; let len = 3000; let reference = 100_000u64; @@ -1735,7 +1735,7 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -1752,7 +1752,7 @@ mod tests { #[crate::test] async fn test_single_element() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let values: Vec = vec![42]; let primitive = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); let bp = BitPacked::encode( @@ -1772,13 +1772,13 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_exactly_elements_per_block() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); // Exactly 2048 elements — one full block, no remainder let bit_width: u8 = 6; let len = 2048; @@ -1807,7 +1807,7 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -1937,7 +1937,7 @@ mod tests { #[crate::test] async fn test_alp_f64_for_bitpacked() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut ctx = LEGACY_SESSION.create_execution_ctx(); // ALP(FoR(BitPacked)) with f64: same structure as the f32 test. let len = 3000; @@ -1974,7 +1974,7 @@ mod tests { let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let gpu = CanonicalCudaExt::into_host(canonical).await?.into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(cpu, gpu, &mut cpu_ctx); Ok(()) } @@ -2038,7 +2038,7 @@ mod tests { #[crate::test] async fn alp_slice_device_patches() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); // Regression test for https://github.com/vortex-data/vortex/issues/7838#issuecomment-4452796116. let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 4096; @@ -2094,14 +2094,14 @@ mod tests { expected[2048 - 100] = std::f64::consts::LN_2; let expected = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected, gpu, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected, gpu, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_runend_u32_ends_u16_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); // RunEnd with u32 ends, u16 values. Output type = u16. // Ends (u32) differ from output (u16) → pending subtree. let ends: Vec = vec![500, 1000, 1500, 2000]; @@ -2139,14 +2139,14 @@ mod tests { }) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_dict_bitpacked_u8_codes_u32_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); // Dict with BitPacked u8 codes (narrower than u32 output) and u32 values. // The kernel's bitunpack_typed decodes at the source's native width and // widens to T, so this fuses into a single kernel launch. @@ -2178,7 +2178,7 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2196,7 +2196,7 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); @@ -2223,7 +2223,7 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2240,7 +2240,7 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); @@ -2268,7 +2268,7 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2294,7 +2294,7 @@ mod tests { #[case] ends: Vec, #[case] values: Vec, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let ends_u64: Vec = ends.iter().map(|e| (*e).into()).collect(); let len = *ends_u64.last().unwrap() as usize; let bit_width = 64 - ends_u64.iter().max().unwrap().leading_zeros() as u8; @@ -2329,7 +2329,7 @@ mod tests { prev = *end; } let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2339,7 +2339,7 @@ mod tests { /// at native width and widens to T, fusing everything. #[crate::test] async fn test_runend_mixed_width_for_bp_u16_ends_u32_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let ends: Vec = vec![500, 1000, 1500, 2000]; let values: Vec = vec![100, 200, 300, 400]; let len = 2000usize; @@ -2375,14 +2375,14 @@ mod tests { prev = end; } let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_sliced_dict_mixed_width() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); // Sliced Dict with u8 codes and u32 values — combines PartiallyFused + slice handling. let dict_values: Vec = vec![100, 200, 300, 400]; let full_len = 4096; @@ -2406,7 +2406,7 @@ mod tests { .map(|&c| dict_values[c as usize]) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2493,21 +2493,20 @@ mod tests { /// Nullable Primitive array — LOAD source with validity propagated. #[crate::test] async fn test_nullable_primitive() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let array = PrimitiveArray::from_option_iter( (0..2048u32).map(|i| if i % 3 == 0 { None } else { Some(i) }), ); - let cpu = crate::canonicalize_cpu(array.clone())?.into_array(); - let gpu = try_gpu_dispatch(&array.into_array(), &mut cuda_ctx) + let gpu = try_gpu_dispatch(&array.clone().into_array(), &mut cuda_ctx) .await? .into_host() .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(array, gpu, &mut cpu_ctx); Ok(()) } @@ -2517,7 +2516,7 @@ mod tests { /// validity, so this produces a real nullable FoR(BitPacked) tree. #[crate::test] async fn test_nullable_for_bitpacked() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 2048; @@ -2534,7 +2533,6 @@ mod tests { }) .collect(); let prim = PrimitiveArray::from_option_iter(values.iter().copied()); - let cpu = crate::canonicalize_cpu(prim.clone())?.into_array(); // FoR encoding: subtract reference to get residuals [0..63]. // Null positions get 0 (from from_option_iter), which is fine — @@ -2565,7 +2563,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(prim, gpu, &mut cpu_ctx); Ok(()) } @@ -2590,20 +2588,19 @@ mod tests { /// AllValid nullable array — should fuse and produce AllValid output. #[crate::test] async fn test_all_valid_nullable() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let values: Vec = (0..2048).collect(); let array = PrimitiveArray::new(Buffer::from(values.clone()), Validity::AllValid); - let cpu = crate::canonicalize_cpu(array.clone())?.into_array(); - let gpu = try_gpu_dispatch(&array.into_array(), &mut cuda_ctx) + let gpu = try_gpu_dispatch(&array.clone().into_array(), &mut cuda_ctx) .await? .into_host() .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(array, gpu, &mut cpu_ctx); Ok(()) } @@ -2627,7 +2624,7 @@ mod tests { /// Dict with non-nullable codes but nullable values should still fuse. #[crate::test] async fn test_dict_nullable_values_fuses() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); use vortex::buffer::buffer; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; @@ -2636,8 +2633,8 @@ mod tests { let values = PrimitiveArray::from_option_iter([Some(10u32), None, Some(30)]); let dict = DictArray::try_new(codes.into_array(), values.into_array())?; - let cpu = crate::canonicalize_cpu(dict.clone())?.into_array(); let gpu = dict + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -2645,7 +2642,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(dict, gpu, &mut cpu_ctx); Ok(()) } @@ -2653,7 +2650,7 @@ mod tests { /// Validity must survive through fused dispatch and into the filter. #[crate::test] async fn test_nullable_fused_then_filter() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); use vortex::array::arrays::FilterArray; use vortex::mask::Mask; @@ -2675,8 +2672,8 @@ mod tests { let mask = Mask::from_iter((0..len).map(|i| i % 2 == 0)); let filter_array = FilterArray::try_new(prim.into_array(), mask)?; - let cpu = crate::canonicalize_cpu(filter_array.clone())?.into_array(); let gpu = filter_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -2684,7 +2681,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(filter_array, gpu, &mut cpu_ctx); Ok(()) } @@ -2934,7 +2931,7 @@ mod tests { /// u8 BitPacked with patches (bit_width=3, patch values > 7). #[crate::test] async fn test_bitpacked_with_patches_u8() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let bit_width: u8 = 3; let len = 3000usize; let max_val = (1u8 << bit_width) - 1; @@ -2961,14 +2958,14 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } /// u16 BitPacked with patches (bit_width=6, patch values > 63). #[crate::test] async fn test_bitpacked_with_patches_u16() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let bit_width: u8 = 6; let len = 3000usize; let max_val = (1u16 << bit_width) - 1; @@ -2995,14 +2992,14 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } /// u64 BitPacked with patches (bit_width=4, patch values > 15). #[crate::test] async fn test_bitpacked_with_patches_u64() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let bit_width: u8 = 4; let len = 3000usize; let max_val = (1u64 << bit_width) - 1; @@ -3029,7 +3026,7 @@ mod tests { let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -3136,7 +3133,7 @@ mod tests { /// dispatch alongside patch application. #[crate::test] async fn test_nullable_bitpacked_with_patches() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let len = 3000usize; @@ -3157,10 +3154,9 @@ mod tests { .collect(); let prim = PrimitiveArray::from_option_iter(values.iter().copied()); - let cpu = crate::canonicalize_cpu(prim.clone())?.into_array(); let bp = BitPacked::encode( - &prim.into_array(), + &prim.clone().into_array(), bit_width, &mut LEGACY_SESSION.create_execution_ctx(), )?; @@ -3172,7 +3168,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + vortex::array::assert_arrays_eq!(prim, gpu, &mut cpu_ctx); Ok(()) } diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index 66851f85682..c9ae05b4467 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -175,7 +175,6 @@ mod tests { use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; - use crate::canonicalize_cpu; use crate::executor::CudaArrayExt; use crate::session::CudaSession; @@ -194,7 +193,7 @@ mod tests { /// FoR(BitPacked) u32 — entire tree compiles into a single fused plan. #[crate::test] async fn test_fused() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let values: Vec = (0..2048).map(|i| (i % 128) as u32).collect(); @@ -206,15 +205,15 @@ mod tests { .vortex_expect("bp"); let arr = FoR::try_new(bp.into_array(), 1000u32.into()).vortex_expect("for"); - let cpu = canonicalize_cpu(arr.clone())?.into_array(); let gpu = arr + .clone() .into_array() .execute_cuda(&mut ctx) .await? .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + assert_arrays_eq!(arr, gpu, &mut cpu_ctx); Ok(()) } @@ -222,7 +221,7 @@ mod tests { /// Exercises the unsigned type reinterpretation in CudaDispatchPlan::execute. #[crate::test] async fn test_fused_f32() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); use vortex::encodings::alp::ALP; use vortex::encodings::alp::Exponents; @@ -243,22 +242,22 @@ mod tests { None, )?; - let cpu = canonicalize_cpu(alp.clone())?.into_array(); let gpu = alp + .clone() .into_array() .execute_cuda(&mut ctx) .await? .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + assert_arrays_eq!(alp, gpu, &mut cpu_ctx); Ok(()) } /// ALP with patches — plan builder rejects it, falls back to ALPExecutor. #[crate::test] async fn test_fallback() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); use vortex::array::patches::Patches; use vortex::array::validity::Validity::NonNullable as NN; use vortex::buffer::buffer; @@ -282,15 +281,15 @@ mod tests { .unwrap(); let arr = ALP::try_new(encoded, Exponents { e: 0, f: 2 }, Some(patches))?; - let cpu = canonicalize_cpu(arr.clone())?.into_array(); let gpu = arr + .clone() .into_array() .execute_cuda(&mut ctx) .await? .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + assert_arrays_eq!(arr, gpu, &mut cpu_ctx); Ok(()) } @@ -300,7 +299,7 @@ mod tests { #[cfg(feature = "unstable_encodings")] #[crate::test] async fn test_partial_fusion() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); use vortex::array::arrays::DictArray; use vortex::array::session::ArraySessionExt; use vortex::encodings::fastlanes; @@ -357,14 +356,14 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + assert_arrays_eq!(cpu, gpu, &mut cpu_ctx); Ok(()) } /// Filter(FoR(BP), mask) — FoR+BP fuses via dyn dispatch, then CUB filters the result. #[crate::test] async fn test_filter_fused_child() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -390,7 +389,7 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(cpu, gpu, &mut assertion_ctx); + assert_arrays_eq!(cpu, gpu, &mut cpu_ctx); Ok(()) } @@ -408,12 +407,12 @@ mod tests { ))] #[crate::test] async fn test_ext_storage_gpu_decode(#[case] ext: ExtensionArray) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); - let expected_storage = canonicalize_cpu(ext.storage_array().clone())?.into_array(); - let expected = ExtensionArray::new(ext.ext_dtype().clone(), expected_storage).into_array(); + let expected = + ExtensionArray::new(ext.ext_dtype().clone(), ext.storage_array().clone()).into_array(); let actual = ext.into_array().execute_cuda(&mut ctx).await?; let storage = actual.as_extension().storage_array(); @@ -425,14 +424,14 @@ mod tests { assert!(!storage.is_host(), "storage was not decoded on the device"); let actual = actual.into_host().await?.into_array(); - assert_arrays_eq!(expected, actual, &mut assertion_ctx); + assert_arrays_eq!(expected, actual, &mut cpu_ctx); Ok(()) } /// Extension over already-canonical storage executes unchanged. #[crate::test] async fn test_ext_canonical_storage() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -450,7 +449,7 @@ mod tests { .into_host() .await? .into_array(); - assert_arrays_eq!(ext.into_array(), actual, &mut assertion_ctx); + assert_arrays_eq!(ext.into_array(), actual, &mut cpu_ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/arrays/constant.rs b/vortex-cuda/src/kernel/arrays/constant.rs index f1591d8159a..d385333a024 100644 --- a/vortex-cuda/src/kernel/arrays/constant.rs +++ b/vortex-cuda/src/kernel/arrays/constant.rs @@ -199,6 +199,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::scalar::Scalar; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -223,66 +224,60 @@ mod tests { async fn test_cuda_constant_materialization( #[case] constant_array: ConstantArray, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); - let cpu_result = crate::canonicalize_cpu(constant_array.clone())?; - let gpu_result = ConstantNumericExecutor - .execute(constant_array.into_array(), &mut cuda_ctx) + .execute(constant_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU materialization failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(constant_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_constant_empty_array() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let constant_array = ConstantArray::new(42i32, 0); - let cpu_result = crate::canonicalize_cpu(constant_array.clone())?; - let gpu_result = ConstantNumericExecutor - .execute(constant_array.into_array(), &mut cuda_ctx) + .execute(constant_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU materialization failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(constant_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_constant_small_array() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); // Test with array smaller than one block (< 2048 elements) let constant_array = ConstantArray::new(99i32, 100); - let cpu_result = crate::canonicalize_cpu(constant_array.clone())?; - let gpu_result = ConstantNumericExecutor - .execute(constant_array.into_array(), &mut cuda_ctx) + .execute(constant_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU materialization failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(constant_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index 74070f684b7..394de699f49 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -308,6 +308,7 @@ mod tests { use vortex::dtype::DecimalDType; use vortex::dtype::i256; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -324,7 +325,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_u32_values_u8_codes() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -338,30 +339,23 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_u64_values_u16_codes() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -378,30 +372,23 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_i32_values_u32_codes() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -415,29 +402,22 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_large_array() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -451,30 +431,23 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values_array.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_values_with_validity() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -489,29 +462,22 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_codes_with_validity() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -531,29 +497,22 @@ mod tests { let dict_array = DictArray::try_new(codes.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_both_with_validity() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -580,29 +539,22 @@ mod tests { let dict_array = DictArray::try_new(codes.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_i64_values_with_validity() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -630,29 +582,22 @@ mod tests { let dict_array = DictArray::try_new(codes.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_all_valid_matches_baseline() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -668,23 +613,16 @@ mod tests { let dict_array = DictArray::try_new(codes.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - // Get baseline from CPU canonicalization - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - // Execute on CUDA let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_primitive(); let cuda_result = cuda_primitive_to_host(cuda_result)?; - // Compare CUDA result with baseline - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + // Compare CUDA result with the encoded dict array + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } @@ -700,7 +638,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i8_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -714,26 +652,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i16_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -747,26 +679,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i32_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -780,26 +706,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i64_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -816,26 +736,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i128_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -857,20 +771,14 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } @@ -884,7 +792,7 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_u8_codes() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -895,26 +803,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_values_u16_codes() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -925,26 +827,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_max_inlined_12_bytes() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -957,26 +853,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_outlined_views() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -992,26 +882,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_empty_strings() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1022,26 +906,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_values_with_validity() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1053,26 +931,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_string_outlined_with_validity() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1091,26 +963,20 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_varbinview(); let cuda_result = cuda_varbinview_to_host(cuda_result).await?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_dict_decimal_i256_values() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -1132,20 +998,14 @@ mod tests { let dict_array = DictArray::try_new(codes_array.into_array(), values.into_array()) .vortex_expect("failed to create Dict array"); - let baseline = crate::canonicalize_cpu(dict_array.clone())?; - let cuda_result = DictExecutor - .execute(dict_array.into_array(), &mut cuda_ctx) + .execute(dict_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_decimal(); let cuda_result = cuda_decimal_to_host(cuda_result)?; - assert_arrays_eq!( - cuda_result.into_array(), - baseline.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cuda_result.into_array(), dict_array, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/alp.rs b/vortex-cuda/src/kernel/encodings/alp.rs index 610f8c0bbda..5c6f70498a3 100644 --- a/vortex-cuda/src/kernel/encodings/alp.rs +++ b/vortex-cuda/src/kernel/encodings/alp.rs @@ -142,7 +142,6 @@ mod tests { use vortex::array::IntoArray; use vortex::array::LEGACY_SESSION; - use vortex::array::VortexSessionExecute; use vortex::array::arrays::PrimitiveArray; use vortex::array::assert_arrays_eq; use vortex::array::patches::Patches; @@ -153,10 +152,10 @@ mod tests { use vortex::encodings::alp::Exponents; use vortex::encodings::alp::alp_encode; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; - use crate::canonicalize_cpu; use crate::executor::CudaArrayExt; use crate::session::CudaSession; @@ -170,7 +169,7 @@ mod tests { /// Patches must carry `chunk_offsets` — the fused kernel requires them. #[crate::test] async fn test_cuda_alp_decompression_f32() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -197,17 +196,15 @@ mod tests { Some(patches), )?; - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = ALPExecutor - .execute(alp_array.into_array(), &mut cuda_ctx) + .execute(alp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -218,7 +215,7 @@ mod tests { /// preserved through the standalone ALP GPU executor. #[crate::test] async fn test_cuda_alp_nullable_with_patches() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -242,9 +239,8 @@ mod tests { &mut LEGACY_SESSION.create_execution_ctx(), )?; - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -252,7 +248,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -260,7 +256,7 @@ mod tests { /// elements are actually null. #[crate::test] async fn test_cuda_alp_all_valid_nullable() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -274,9 +270,8 @@ mod tests { &mut LEGACY_SESSION.create_execution_ctx(), )?; - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -284,7 +279,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -294,7 +289,7 @@ mod tests { /// (zero patches) via the offset math rather than the NULL sentinel. #[crate::test] async fn test_cuda_alp_multi_chunk_sparse_patches() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -321,9 +316,8 @@ mod tests { "expected patches from ALP exceptions" ); - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -331,7 +325,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -340,7 +334,7 @@ mod tests { /// so this guards the fast-path for the (i64, f64) kernel variant. #[crate::test] async fn test_cuda_alp_f64_multi_chunk_with_patches() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -365,9 +359,8 @@ mod tests { "expected patches from ALP exceptions" ); - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -375,7 +368,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -385,7 +378,7 @@ mod tests { /// (existing tests have ≤ 6 patches per chunk). #[crate::test] async fn test_cuda_alp_dense_patches_single_chunk() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -415,17 +408,15 @@ mod tests { Some(patches), )?; - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = ALPExecutor - .execute(alp_array.into_array(), &mut cuda_ctx) + .execute(alp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } @@ -434,7 +425,7 @@ mod tests { /// loop. Includes a patch in the tail. #[crate::test] async fn test_cuda_alp_partial_tail_chunk() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -452,9 +443,8 @@ mod tests { "expected patches from ALP exceptions" ); - let cpu_result = canonicalize_cpu(alp_array.clone())?.into_array(); - let gpu_result = alp_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -462,7 +452,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(alp_array, gpu_result, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 875fdec09a3..9f99be0eea2 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -250,25 +250,19 @@ mod tests { #[case] iter: impl Iterator, #[case] bw: u8, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let array = PrimitiveArray::new(iter.collect::>(), NonNullable); // Last two items should be patched - let bp_with_patches = BitPacked::encode( - &array.into_array(), - bw, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp_with_patches = BitPacked::encode(&array.into_array(), bw, &mut ctx)?; assert!(bp_with_patches.patches().is_some()); - let cpu_result = crate::canonicalize_cpu(bp_with_patches.clone())?.into_array(); - let gpu_result = block_on(async { BitPackedExecutor - .execute(bp_with_patches.into_array(), &mut cuda_ctx) + .execute(bp_with_patches.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -276,14 +270,14 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(bp_with_patches, gpu_result, &mut ctx); Ok(()) } #[crate::test] fn test_patches() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -293,18 +287,12 @@ mod tests { ); // Last two items should be patched - let bp_with_patches = BitPacked::encode( - &array.into_array(), - 9, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp_with_patches = BitPacked::encode(&array.into_array(), 9, &mut ctx)?; assert!(bp_with_patches.patches().is_some()); - let cpu_result = crate::canonicalize_cpu(bp_with_patches.clone())?.into_array(); - let gpu_result = block_on(async { BitPackedExecutor - .execute(bp_with_patches.into_array(), &mut cuda_ctx) + .execute(bp_with_patches.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -312,7 +300,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(bp_with_patches, gpu_result, &mut ctx); Ok(()) } @@ -327,7 +315,7 @@ mod tests { #[case::bw_7(7)] #[crate::test] fn test_cuda_bitunpack_u8(#[case] bit_width: u8) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -340,17 +328,11 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let cpu_result = crate::canonicalize_cpu(bitpacked_array.clone())?; - + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let gpu_result = block_on(async { BitPackedExecutor - .execute(bitpacked_array.into_array(), &mut cuda_ctx) + .execute(bitpacked_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -358,7 +340,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(bitpacked_array, gpu_result, &mut ctx); Ok(()) } @@ -381,7 +363,7 @@ mod tests { #[case::bw_15(15)] #[crate::test] fn test_cuda_bitunpack_u16(#[case] bit_width: u8) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -394,17 +376,11 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let cpu_result = crate::canonicalize_cpu(bitpacked_array.clone())?; - + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let gpu_result = block_on(async { BitPackedExecutor - .execute(bitpacked_array.into_array(), &mut cuda_ctx) + .execute(bitpacked_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -412,7 +388,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(bitpacked_array, gpu_result, &mut ctx); Ok(()) } @@ -451,7 +427,7 @@ mod tests { #[case::bw_31(31)] #[crate::test] fn test_cuda_bitunpack_u32(#[case] bit_width: u8) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -464,17 +440,11 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let cpu_result = crate::canonicalize_cpu(bitpacked_array.clone())?; - + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let gpu_result = block_on(async { BitPackedExecutor - .execute(bitpacked_array.into_array(), &mut cuda_ctx) + .execute(bitpacked_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -482,7 +452,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(bitpacked_array, gpu_result, &mut ctx); Ok(()) } @@ -553,7 +523,7 @@ mod tests { #[case::bw_63(63)] #[crate::test] fn test_cuda_bitunpack_u64(#[case] bit_width: u8) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -566,16 +536,11 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); - let cpu_result = crate::canonicalize_cpu(bitpacked_array.clone())?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let gpu_result = block_on(async { BitPackedExecutor - .execute(bitpacked_array.into_array(), &mut cuda_ctx) + .execute(bitpacked_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -583,14 +548,14 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(bitpacked_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] fn test_cuda_bitunpack_sliced() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let bit_width = 32; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -604,17 +569,12 @@ mod tests { NonNullable, ); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("operation should succeed in test"); + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), bit_width, &mut ctx) + .vortex_expect("operation should succeed in test"); let sliced_array = bitpacked_array.into_array().slice(67..3969)?; - let cpu_result = crate::canonicalize_cpu(sliced_array.clone())?; let gpu_result = block_on(async { BitPackedExecutor - .execute(sliced_array, &mut cuda_ctx) + .execute(sliced_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -622,7 +582,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(sliced_array, gpu_result, &mut ctx); Ok(()) } @@ -646,11 +606,7 @@ mod tests { .collect::>(), NonNullable, ); - let bitpacked = BitPacked::encode( - &values.into_array(), - 9, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bitpacked = BitPacked::encode(&values.into_array(), 9, &mut ctx)?; assert!(bitpacked.patches().is_some()); let array = if let Some(range) = range { bitpacked.into_array().slice(range)? @@ -674,7 +630,7 @@ mod tests { /// offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_sliced_patches_offset_within_chunk() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -684,11 +640,7 @@ mod tests { let primitive_array = PrimitiveArray::new(buffer![100u8, 101, 102, 3, 4, 5], NonNullable); // Encode with bit width 4. First 3 elements patched, remainder will pack. - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - 4, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), 4, &mut ctx)?; assert!( bitpacked_array.patches().is_some(), "Expected patches to be present" @@ -696,12 +648,9 @@ mod tests { let sliced_array = bitpacked_array.into_array().slice(2..6)?; - let cpu_result = sliced_array - .clone() - .execute::(cuda_ctx.execution_ctx())?; let gpu_result = block_on(async { BitPackedExecutor - .execute(sliced_array, &mut cuda_ctx) + .execute(sliced_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -709,7 +658,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(cpu_result, gpu_result, &mut ctx); Ok(()) } @@ -717,7 +666,7 @@ mod tests { /// Test slicing a bitpacked array multiple times, accumulating offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_double_sliced_patches() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -734,11 +683,7 @@ mod tests { let primitive_array = PrimitiveArray::new(Buffer::from_iter(values.iter().copied()), NonNullable); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - 9, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), 9, &mut ctx)?; assert!( bitpacked_array.patches().is_some(), "Expected patches to be present" @@ -751,12 +696,9 @@ mod tests { // covers every chunk in the packed buffer. let second_slice = first_slice.slice(50..2900)?; - let cpu_result = second_slice - .clone() - .execute::(cuda_ctx.execution_ctx())?; let gpu_result = block_on(async { BitPackedExecutor - .execute(second_slice, &mut cuda_ctx) + .execute(second_slice.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -764,7 +706,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(second_slice, gpu_result, &mut ctx); Ok(()) } @@ -772,7 +714,7 @@ mod tests { /// Test slicing to skip an entire chunk's worth of patches. #[crate::test] fn test_cuda_bitunpack_sliced_skip_first_chunk_patches() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -791,11 +733,7 @@ mod tests { let primitive_array = PrimitiveArray::new(Buffer::from_iter(values.iter().copied()), NonNullable); - let bitpacked_array = BitPacked::encode( - &primitive_array.into_array(), - 9, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bitpacked_array = BitPacked::encode(&primitive_array.into_array(), 9, &mut ctx)?; assert!( bitpacked_array.patches().is_some(), "Expected patches to be present" @@ -804,12 +742,9 @@ mod tests { // Slice to skip past all first chunk patches let sliced_array = bitpacked_array.into_array().slice(1024..3072)?; - let cpu_result = sliced_array - .clone() - .execute::(cuda_ctx.execution_ctx())?; let gpu_result = block_on(async { BitPackedExecutor - .execute(sliced_array, &mut cuda_ctx) + .execute(sliced_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -817,7 +752,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(sliced_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/date_time_parts.rs b/vortex-cuda/src/kernel/encodings/date_time_parts.rs index e3cbc1fda59..dcb1fd1f69e 100644 --- a/vortex-cuda/src/kernel/encodings/date_time_parts.rs +++ b/vortex-cuda/src/kernel/encodings/date_time_parts.rs @@ -214,6 +214,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::extension::datetime::TimeUnit; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -279,29 +280,27 @@ mod tests { #[case] subseconds: Vec, #[case] time_unit: TimeUnit, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let dtp_array = make_datetimeparts_array(days, seconds, subseconds, time_unit); - let cpu_result = crate::canonicalize_cpu(dtp_array.clone())?; - let gpu_result = DateTimePartsExecutor - .execute(dtp_array.into_array(), &mut cuda_ctx) + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_datetimeparts_large_array() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -311,24 +310,22 @@ mod tests { let subseconds: Vec = (0..len).map(|i| (i % 1000) as i64).collect(); let dtp_array = make_datetimeparts_array(days, seconds, subseconds, TimeUnit::Milliseconds); - let cpu_result = crate::canonicalize_cpu(dtp_array.clone())?; - let gpu_result = DateTimePartsExecutor - .execute(dtp_array.into_array(), &mut cuda_ctx) + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_datetimeparts_with_nulls() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -360,17 +357,15 @@ mod tests { ) .vortex_expect("Failed to create DateTimePartsArray"); - let cpu_result = crate::canonicalize_cpu(dtp_array.clone())?; - let gpu_result = DateTimePartsExecutor - .execute(dtp_array.into_array(), &mut cuda_ctx) + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index 1752127d423..5f7e639b8a1 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -70,6 +70,7 @@ mod tests { use vortex::dtype::DecimalDType; use vortex::encodings::decimal_byte_parts::DecimalByteParts; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::session::CudaSession; @@ -85,7 +86,7 @@ mod tests { #[case] precision: u8, #[case] scale: i8, ) { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("create execution context"); @@ -96,18 +97,11 @@ mod tests { ) .vortex_expect("create DecimalBytePartsArray"); - let cpu_result = - crate::canonicalize_cpu(dbp_array.clone()).vortex_expect("CPU canonicalize"); - let gpu_result = DecimalBytePartsExecutor - .execute(dbp_array.into_array(), &mut cuda_ctx) + .execute(dbp_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decode"); - assert_arrays_eq!( - cpu_result.into_array(), - gpu_result.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(dbp_array, gpu_result.into_array(), &mut ctx); } } diff --git a/vortex-cuda/src/kernel/encodings/for_.rs b/vortex-cuda/src/kernel/encodings/for_.rs index 4de3fb79ed9..424b358c937 100644 --- a/vortex-cuda/src/kernel/encodings/for_.rs +++ b/vortex-cuda/src/kernel/encodings/for_.rs @@ -160,28 +160,26 @@ mod tests { #[case::u64(make_for_array((0..2050).map(|i| (i % 2050) as u64).collect(), 1000000u64))] #[crate::test] async fn test_cuda_for_decompression(#[case] for_array: FoRArray) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); - let cpu_result = crate::canonicalize_cpu(for_array.clone())?; - let gpu_result = FoRExecutor - .execute(for_array.into_array(), &mut cuda_ctx) + .execute(for_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(for_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_signed_ffor() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -195,10 +193,8 @@ mod tests { .into_array(); let for_array = FoR::try_new(packed, (-8i8).into()).unwrap(); - let cpu_result = crate::canonicalize_cpu(for_array.clone()).unwrap(); - let gpu_result = FoRExecutor - .execute(for_array.into_array(), &mut cuda_ctx) + .execute(for_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() @@ -206,6 +202,6 @@ mod tests { .vortex_expect("copying to host failed") .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(for_array, gpu_result, &mut ctx); } } diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 6c2b080058e..eef0db049b8 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -209,6 +209,7 @@ mod tests { use vortex::encodings::fsst::fsst_compress; use vortex::encodings::fsst::fsst_train_compressor; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -235,7 +236,7 @@ mod tests { #[case] strings: Vec>, #[case] nullability: Nullability, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -244,23 +245,22 @@ mod tests { let fsst_array = fsst_compress(&varbin, &compressor, cuda_ctx.execution_ctx())?.into_array(); - let cpu_result = crate::canonicalize_cpu(fsst_array.clone())?; let gpu_result = FSSTExecutor - .execute(fsst_array, &mut cuda_ctx) + .execute(fsst_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(fsst_array, gpu_result, &mut ctx); Ok(()) } /// Exercises the multi-block grid-stride path on a larger dataset. #[crate::test] async fn test_cuda_fsst_decompression_roundtrip_large() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); use vortex_fsst::test_utils::make_fsst_clickbench_urls; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) @@ -268,16 +268,15 @@ mod tests { let fsst_array = make_fsst_clickbench_urls(100_000, cuda_ctx.execution_ctx()).into_array(); - let cpu_result = crate::canonicalize_cpu(fsst_array.clone())?; let gpu_result = FSSTExecutor - .execute(fsst_array, &mut cuda_ctx) + .execute(fsst_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(fsst_array, gpu_result, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index 52849f109dc..4ac042e2972 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -175,6 +175,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex_array::ExecutionCtx; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -204,29 +205,27 @@ mod tests { #[case::u64_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u64, 5, 10], vec![1i32, 2, 3], ctx))] #[crate::test] async fn test_cuda_runend_types(#[case] build: RunEndBuilder) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let runend_array = build(cuda_ctx.execution_ctx()); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?; - let gpu_result = RunEndExecutor - .execute(runend_array.into_array(), &mut cuda_ctx) + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_large_array() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -240,47 +239,43 @@ mod tests { let runend_array = make_runend_array(ends, values, cuda_ctx.execution_ctx()); assert_eq!(runend_array.len(), total_len); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?; - let gpu_result = RunEndExecutor - .execute(runend_array.into_array(), &mut cuda_ctx) + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_single_run() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let runend_array = make_runend_array(vec![100u32], vec![42i32], cuda_ctx.execution_ctx()); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?; - let gpu_result = RunEndExecutor - .execute(runend_array.into_array(), &mut cuda_ctx) + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_many_small_runs() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -291,24 +286,22 @@ mod tests { let runend_array = make_runend_array(ends, values, cuda_ctx.execution_ctx()); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?; - let gpu_result = RunEndExecutor - .execute(runend_array.into_array(), &mut cuda_ctx) + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_runend_nullable_values_falls_back_to_cpu() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -322,10 +315,9 @@ mod tests { PrimitiveArray::new(Buffer::from(vec![10i32, 0, 30]), validity).into_array(); let runend_array = RunEnd::new(ends_array, values_array, cuda_ctx.execution_ctx()); - let cpu_result = crate::canonicalize_cpu(runend_array.clone())?.into_array(); - // execute_cuda should fall back to CPU and still produce the correct result. let gpu_result = runend_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await @@ -334,7 +326,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/sequence.rs b/vortex-cuda/src/kernel/encodings/sequence.rs index db36d068f57..e52a706c8b7 100644 --- a/vortex-cuda/src/kernel/encodings/sequence.rs +++ b/vortex-cuda/src/kernel/encodings/sequence.rs @@ -92,6 +92,7 @@ mod tests { use vortex::dtype::Nullability; use vortex::encodings::sequence::Sequence; use vortex::scalar::PValue; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::CudaSession; @@ -124,15 +125,15 @@ mod tests { len: usize, nullability: Nullability, ) { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).unwrap(); - let array = Sequence::try_new_typed(base, multiplier, nullability, len).unwrap(); - - let cpu_result = crate::canonicalize_cpu(array.clone()).unwrap().into_array(); + let array = Sequence::try_new_typed(base, multiplier, nullability, len) + .unwrap() + .into_array(); let gpu_result = SequenceExecutor - .execute(array.into_array(), &mut cuda_ctx) + .execute(array.clone(), &mut cuda_ctx) .await .unwrap() .into_host() @@ -140,6 +141,6 @@ mod tests { .unwrap() .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(array, gpu_result, &mut ctx); } } diff --git a/vortex-cuda/src/kernel/encodings/zigzag.rs b/vortex-cuda/src/kernel/encodings/zigzag.rs index b6321ed5178..107ad2b36a1 100644 --- a/vortex-cuda/src/kernel/encodings/zigzag.rs +++ b/vortex-cuda/src/kernel/encodings/zigzag.rs @@ -105,6 +105,7 @@ mod tests { use vortex::buffer::Buffer; use vortex::encodings::zigzag::ZigZag; use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -112,7 +113,7 @@ mod tests { #[crate::test] async fn test_cuda_zigzag_decompression_u32() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -122,19 +123,18 @@ mod tests { let zigzag_array = ZigZag::try_new( PrimitiveArray::new(Buffer::from(encoded_data), NonNullable).into_array(), - )?; - - let cpu_result = crate::canonicalize_cpu(zigzag_array.clone())?; + )? + .into_array(); let gpu_result = ZigZagExecutor - .execute(zigzag_array.into_array(), &mut cuda_ctx) + .execute(zigzag_array.clone(), &mut cuda_ctx) .await .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result.into_array(), gpu_result, &mut assertion_ctx); + assert_arrays_eq!(zigzag_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index 71597f59e8c..02bfc0e4c8a 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -356,6 +356,7 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::encodings::zstd::Zstd; use vortex::error::VortexResult; + use vortex_array::VortexSessionExecute; use super::*; use crate::CanonicalCudaExt; @@ -364,7 +365,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_utf8() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -385,17 +386,13 @@ mod tests { .execute(zstd_array.into_array(), &mut cuda_ctx) .await?; - assert_arrays_eq!( - cpu_result.into_array(), - gpu_result.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array(), &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_zstd_decompression_multiple_frames() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -426,17 +423,13 @@ mod tests { .execute(zstd_array.into_array(), &mut cuda_ctx) .await?; - assert_arrays_eq!( - cpu_result.into_array(), - gpu_result.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(cpu_result.into_array(), gpu_result.into_array(), &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_zstd_decompression_sliced() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -458,16 +451,11 @@ mod tests { // Slice the array to get a subset (indices 2..7) let sliced_zstd = zstd_array.slice(2..7)?; - let cpu_result = crate::canonicalize_cpu(sliced_zstd.clone())?; let gpu_result = ZstdExecutor .execute(sliced_zstd.clone(), &mut cuda_ctx) .await?; - assert_arrays_eq!( - cpu_result.into_array(), - gpu_result.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(sliced_zstd, gpu_result.into_array(), &mut ctx); Ok(()) } @@ -476,7 +464,7 @@ mod tests { /// correct results instead of panicking. #[crate::test] async fn test_cuda_zstd_nullable_falls_back_to_cpu() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -491,10 +479,9 @@ mod tests { let zstd_array = Zstd::from_var_bin_view(&strings, 3, 0, cuda_ctx.execution_ctx())?; - let cpu_result = crate::canonicalize_cpu(zstd_array.clone())?.into_array(); - // execute_cuda should fall back to CPU and still produce the correct result. let gpu_result = zstd_array + .clone() .into_array() .execute_cuda(&mut cuda_ctx) .await? @@ -502,7 +489,7 @@ mod tests { .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(zstd_array, gpu_result, &mut ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs index 5a67fb4ea12..71f0af259c1 100644 --- a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs +++ b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs @@ -234,31 +234,26 @@ mod tests { #[crate::test] async fn test_cuda_zstd_buffers_decompression_primitive() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); let input = PrimitiveArray::from_iter(0i64..1024).into_array(); - let compressed = ZstdBuffers::compress(&input, 3, &crate::cuda_session())?; + let compressed = ZstdBuffers::compress(&input, 3, &crate::cuda_session())?.into_array(); - let cpu_result = crate::canonicalize_cpu(compressed.clone())?; let gpu_result = ZstdBuffersExecutor - .execute(compressed.into_array(), &mut cuda_ctx) + .execute(compressed.clone(), &mut cuda_ctx) .await? .into_host() .await?; - assert_arrays_eq!( - cpu_result.into_array(), - gpu_result.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(compressed, gpu_result.into_array(), &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_cuda_zstd_buffers_decompression_varbinview() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -271,20 +266,15 @@ mod tests { "baz", ]) .into_array(); - let compressed = ZstdBuffers::compress(&input, 3, &crate::cuda_session())?; + let compressed = ZstdBuffers::compress(&input, 3, &crate::cuda_session())?.into_array(); - let cpu_result = crate::canonicalize_cpu(compressed.clone())?; let gpu_result = ZstdBuffersExecutor - .execute(compressed.into_array(), &mut cuda_ctx) + .execute(compressed.clone(), &mut cuda_ctx) .await? .into_host() .await?; - assert_arrays_eq!( - cpu_result.into_array(), - gpu_result.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(compressed, gpu_result.into_array(), &mut cpu_ctx); Ok(()) } } diff --git a/vortex-cuda/src/kernel/filter/decimal.rs b/vortex-cuda/src/kernel/filter/decimal.rs index e40b1445bf4..dc10b908461 100644 --- a/vortex-cuda/src/kernel/filter/decimal.rs +++ b/vortex-cuda/src/kernel/filter/decimal.rs @@ -48,6 +48,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -88,30 +89,28 @@ mod tests { #[case] input: DecimalArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); let filter_array = FilterArray::try_new(input.clone().into_array(), mask.clone())?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_gpu_filter_decimal_large_array() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -124,18 +123,16 @@ mod tests { let filter_array = FilterArray::try_new(input.into_array(), mask)?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_eq!(cpu_result.len(), gpu_result.len()); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_eq!(filter_array.len(), gpu_result.len()); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/filter/primitive.rs b/vortex-cuda/src/kernel/filter/primitive.rs index 6c099d864c4..c2a21ccc7d7 100644 --- a/vortex-cuda/src/kernel/filter/primitive.rs +++ b/vortex-cuda/src/kernel/filter/primitive.rs @@ -46,6 +46,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -82,30 +83,28 @@ mod tests { #[case] input: PrimitiveArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); let filter_array = FilterArray::try_new(input.clone().into_array(), mask.clone())?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } #[crate::test] async fn test_gpu_filter_large_array() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -118,18 +117,16 @@ mod tests { let filter_array = FilterArray::try_new(input.into_array(), mask)?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_eq!(cpu_result.len(), gpu_result.len()); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_eq!(filter_array.len(), gpu_result.len()); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/filter/varbinview.rs b/vortex-cuda/src/kernel/filter/varbinview.rs index 87267e368be..7bbafff824c 100644 --- a/vortex-cuda/src/kernel/filter/varbinview.rs +++ b/vortex-cuda/src/kernel/filter/varbinview.rs @@ -46,6 +46,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::mask::Mask; + use vortex_array::VortexSessionExecute; use crate::CanonicalCudaExt; use crate::FilterExecutor; @@ -68,23 +69,21 @@ mod tests { #[case] input: VarBinViewArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); let filter_array = FilterArray::try_new(input.into_array(), mask.clone())?; - let cpu_result = crate::canonicalize_cpu(filter_array.clone())?.into_array(); - let gpu_result = FilterExecutor - .execute(filter_array.into_array(), &mut cuda_ctx) + .execute(filter_array.clone().into_array(), &mut cuda_ctx) .await .vortex_expect("GPU filter failed") .into_host() .await? .into_array(); - assert_arrays_eq!(cpu_result, gpu_result, &mut assertion_ctx); + assert_arrays_eq!(filter_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/patches/types.rs b/vortex-cuda/src/kernel/patches/types.rs index 2bb4c0056e0..3bfe2270b66 100644 --- a/vortex-cuda/src/kernel/patches/types.rs +++ b/vortex-cuda/src/kernel/patches/types.rs @@ -228,6 +228,7 @@ mod tests { use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::patches::Patches; use vortex_error::VortexResult; @@ -317,7 +318,7 @@ mod tests { #[case] chunk_offsets: ArrayRef, #[case] expected: ArrayRef, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let indices = PrimitiveArray::from_iter([100u32, 1100, 2100, 3100, 4100]); let values = PrimitiveArray::from_iter([10u32, 11, 12, 13, 14]); @@ -338,7 +339,7 @@ mod tests { NonNullable, ) .into_array(); - assert_arrays_eq!(expected, actual, &mut assertion_ctx); + assert_arrays_eq!(expected, actual, &mut ctx); assert_eq!(device_patches.n_chunks, 3); assert_eq!(device_patches.offset, 1024); assert_eq!(device_patches.offset_within_chunk, 0); diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index a09301299fa..021c811a90e 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -89,18 +89,6 @@ pub use vortex_nvcomp as nvcomp; use crate::kernel::SequenceExecutor; use crate::kernel::SliceExecutor; -#[cfg(test)] -pub(crate) fn canonicalize_cpu( - array: impl vortex::array::IntoArray, -) -> vortex::error::VortexResult { - use vortex::array::LEGACY_SESSION; - use vortex::array::VortexSessionExecute; - - array - .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) -} - /// Checks if CUDA is available on the system by looking for nvcc. pub fn cuda_available() -> bool { Command::new("nvcc") diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index a86e375d37b..19475e17065 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -415,7 +415,7 @@ mod tests { #[test] fn test_integer_vector_conversion() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let values = vec![1i32, 2, 3, 4, 5]; let len = values.len(); diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index f656755ab65..139faa7aaa7 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -532,7 +532,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_no_options() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let (array, struct_array) = scan(ptr::null()); assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); unsafe { vx_array_free(array) }; @@ -541,7 +541,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_project_all() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let opts = vx_scan_options::default(); let (array, struct_array) = scan(&raw const opts); assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); @@ -552,7 +552,7 @@ mod tests { #[cfg_attr(miri, ignore)] fn test_project_single_field() { unsafe { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let root = vx_expression_root(); let mut opts = vx_scan_options::default(); @@ -714,7 +714,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_ordered() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let opts = vx_scan_options { ordered: true, ..Default::default() diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index cba3a4ec430..8eeea97fff2 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -477,7 +477,7 @@ async fn test_empty_varbin_array_roundtrip() { #[tokio::test] #[cfg_attr(miri, ignore)] async fn issue_5385_filter_casted_column() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = StructArray::try_from_iter([("x", buffer![1u8, 2, 3, 4, 5])]) .unwrap() .into_array(); @@ -1187,7 +1187,7 @@ async fn chunked_file() -> VortexResult { #[tokio::test] async fn basic_file_roundtrip() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let vxf = chunked_file().await?; let result = vxf.scan()?.into_array_stream()?.read_all().await?; @@ -1231,7 +1231,7 @@ async fn file_excluding_dtype() -> VortexResult<()> { #[tokio::test] async fn file_take() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let vxf = chunked_file().await?; let result = vxf .scan()? diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index 3aa66034068..1dbe3a80d32 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -123,7 +123,7 @@ async fn test_file_roundtrip() { /// with "Array vortex.fill_null does not support serialization". #[tokio::test] async fn test_dict_listview_validity_roundtrip() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]).into_array(); let offsets = PrimitiveArray::from_iter(vec![0u32, 2, 4]).into_array(); let sizes = PrimitiveArray::from_iter(vec![2u32, 2, 1]).into_array(); diff --git a/vortex-ipc/src/iterator.rs b/vortex-ipc/src/iterator.rs index 04bbbf21d52..7f8adfef986 100644 --- a/vortex-ipc/src/iterator.rs +++ b/vortex-ipc/src/iterator.rs @@ -172,6 +172,7 @@ mod test { use std::io::Cursor; use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorExt; @@ -182,7 +183,7 @@ mod test { #[test] fn test_sync_stream() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let ipc_buffer = array .to_array_iterator() diff --git a/vortex-ipc/src/stream.rs b/vortex-ipc/src/stream.rs index 68f77e7b80c..43e170e969f 100644 --- a/vortex-ipc/src/stream.rs +++ b/vortex-ipc/src/stream.rs @@ -225,6 +225,7 @@ mod test { use futures::io::Cursor; use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamExt; @@ -235,7 +236,7 @@ mod test { #[tokio::test] async fn test_async_stream() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = buffer![1, 2, 3].into_array(); let ipc_buffer = array .to_array_stream() @@ -272,7 +273,7 @@ mod test { #[tokio::test] async fn test_async_stream_chunked() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(); let ipc_buffer = array .to_array_stream() @@ -296,7 +297,7 @@ mod test { /// Test with 1-byte chunks to stress-test partial read handling. #[tokio::test] async fn test_async_stream_single_byte_chunks() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = buffer![42i64, -1, 0, i64::MAX, i64::MIN].into_array(); let ipc_buffer = array .to_array_stream() diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index 39c8c0223b1..aee6d6575b0 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -353,6 +353,7 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; @@ -471,7 +472,7 @@ mod test { #[from(chunked_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|_h| async { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let result = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 29e0b9ab3f9..87d4197bfa4 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -422,7 +422,7 @@ mod tests { #[test] fn reading_nested_packs_works() { block_on(|handle| async move { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), @@ -522,7 +522,7 @@ mod tests { #[case] expected: Vec, ) { block_on(|handle| async move { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), @@ -577,7 +577,7 @@ mod tests { #[test] fn reading_is_null_works() { block_on(|handle| async move { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let mut ctx_exec = LEGACY_SESSION.create_execution_ctx(); let session = session_with_handle(handle); let strategy = DictStrategy::new( @@ -651,7 +651,7 @@ mod tests { #[test] fn reading_byte_length_pushdown_works() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = VarBinArray::from_iter( [ Some("abc"), @@ -709,7 +709,7 @@ mod tests { } fn test_apply(original: Expression, outer: Expression, inner: Expression) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = VarBinArray::from_iter( [Some("abc"), Some("def"), None], DType::Utf8(Nullability::Nullable), diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 0df63df4ddf..a0269786d51 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -233,6 +233,7 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -255,7 +256,7 @@ mod test { #[test] fn flat_identity() -> VortexResult<()> { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -295,7 +296,7 @@ mod test { #[test] fn flat_expr() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); @@ -335,7 +336,7 @@ mod test { #[test] fn flat_unaligned_row_mask() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 6651b8aa3fd..19e8ac81376 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -327,6 +327,7 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::IntoArray as _; use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; use vortex_array::expr::eq; @@ -351,7 +352,7 @@ mod tests { #[test] fn flat_expr_no_row_id() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -396,7 +397,7 @@ mod tests { #[test] fn flat_expr_row_id() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -441,7 +442,7 @@ mod tests { #[test] fn flat_expr_or() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 1d4d4958889..3310d9eb834 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -634,7 +634,7 @@ mod tests { fn test_struct_layout( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -653,7 +653,7 @@ mod tests { fn test_struct_layout_row_mask( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -720,7 +720,7 @@ mod tests { fn test_struct_layout_nulls( #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); // Read the layout source from the top. let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index eb6b0ebdbc4..0e4fc623715 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -226,6 +226,7 @@ mod test { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; + use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -309,7 +310,7 @@ mod test { #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = session_with_handle(handle); let result = layout .new_reader("".into(), segments, &session, &Default::default()) diff --git a/vortex-python/src/arrays/range_to_sequence.rs b/vortex-python/src/arrays/range_to_sequence.rs index a04a70e32a9..d83662b4387 100644 --- a/vortex-python/src/arrays/range_to_sequence.rs +++ b/vortex-python/src/arrays/range_to_sequence.rs @@ -78,6 +78,7 @@ mod test { use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex::dtype::PType; + use vortex_array::VortexSessionExecute; use crate::arrays::range_to_sequence::range_len; use crate::arrays::range_to_sequence::sequence_array_from_range; @@ -98,7 +99,7 @@ mod test { #[test] fn test_sequence_array_from_len() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::U16, Nullability::NonNullable); let arr = sequence_array_from_range::(0, 10, 1, dtype).unwrap(); assert_arrays_eq!( diff --git a/vortex-test/compat-gen/src/check.rs b/vortex-test/compat-gen/src/check.rs index 102c949b2f0..c788ed17cb4 100644 --- a/vortex-test/compat-gen/src/check.rs +++ b/vortex-test/compat-gen/src/check.rs @@ -5,6 +5,7 @@ use std::path::Path; use clap::ValueEnum; use serde::Serialize; +use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -46,7 +47,7 @@ struct FailedFixture { /// Prints JSON result to stdout, human-readable progress to stderr. /// Returns error if any fixture failed or if mode constraints are violated. pub fn check(dir: &Path, mode: Mode, exclude: &[String]) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let fixtures = all_fixtures(); let fixtures: Vec<_> = fixtures .into_iter() diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs index 01a9c4cb898..919ccfb39b7 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs @@ -17,7 +17,7 @@ pub fn fixtures() -> Vec> { #[cfg(test)] mod tests { - + use vortex_array::VortexSessionExecute; use vortex_array::assert_arrays_eq; use super::fixtures; @@ -27,7 +27,7 @@ mod tests { #[test] fn roundtrip_fixtures_to_bytes() { for fixture in fixtures() { - let mut assertion_ctx = vortex_array::array_execution_ctx(); + let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = fixture.build().unwrap(); check_expected_encodings(&array, fixture.as_ref()).unwrap(); let bytes = adapter::write_file_to_bytes(array.clone()).unwrap(); From d988366bac206804c53a8fda85e0877106a35170 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 23:10:55 +0100 Subject: [PATCH 10/12] fixes Signed-off-by: Robert Kruszewski --- vortex-cuda/src/kernel/encodings/bitpacked.rs | 7 ++++--- vortex-cuda/src/kernel/encodings/zstd.rs | 4 ++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 9f99be0eea2..51df660a501 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -600,13 +600,14 @@ mod tests { #[case] expected_offset: u16, #[case] expected_packed_len: usize, ) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); let values = PrimitiveArray::new( (0u16..4096) .map(|i| if i % 1000 == 0 { 600 } else { i % 512 }) .collect::>(), NonNullable, - ); - let bitpacked = BitPacked::encode(&values.into_array(), 9, &mut ctx)?; + ).into_array(); + let bitpacked = BitPacked::encode(&values, 9, &mut ctx)?; assert!(bitpacked.patches().is_some()); let array = if let Some(range) = range { bitpacked.into_array().slice(range)? @@ -658,7 +659,7 @@ mod tests { .map(|a| a.into_array()) })?; - assert_arrays_eq!(cpu_result, gpu_result, &mut ctx); + assert_arrays_eq!(sliced_array, gpu_result, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index 02bfc0e4c8a..6d14cc46e86 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -477,12 +477,12 @@ mod tests { Some("another string"), ]); - let zstd_array = Zstd::from_var_bin_view(&strings, 3, 0, cuda_ctx.execution_ctx())?; + let zstd_array = + Zstd::from_var_bin_view(&strings, 3, 0, cuda_ctx.execution_ctx())?.into_array(); // execute_cuda should fall back to CPU and still produce the correct result. let gpu_result = zstd_array .clone() - .into_array() .execute_cuda(&mut cuda_ctx) .await? .into_host() From 8091c1a10622278a8f87ef053443700d42889d92 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 22 Jun 2026 23:19:33 +0100 Subject: [PATCH 11/12] fixes Signed-off-by: Robert Kruszewski --- encodings/fsst/src/tests.rs | 3 ++- vortex-array/src/arrays/varbinview/tests.rs | 19 ++++++++----------- vortex-cuda/src/kernel/encodings/bitpacked.rs | 4 ++-- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 26e0975da22..43cdeb1a02c 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -4,7 +4,8 @@ use std::sync::LazyLock; use fsst::CompressorBuilder; -use vortex_array::{ArrayRef, ExecutionCtx}; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::VarBinViewArray; diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index 24574942d60..981d4f602a1 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -1,28 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use crate::ToCanonical as _; -#[expect(deprecated)] use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; #[test] pub fn varbin_view() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]); assert_arrays_eq!( binary_arr, VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), - &mut assertion_ctx + &mut ctx ); } #[test] pub fn slice_array() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) .slice(1..2) @@ -30,20 +29,18 @@ pub fn slice_array() { assert_arrays_eq!( binary_arr, VarBinViewArray::from_iter_str(["hello world this is a long string"]), - &mut assertion_ctx + &mut ctx ); } #[test] pub fn flatten_array() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["string1", "string2"]); - #[expect(deprecated)] - let var_bin = binary_arr.as_array().to_varbinview(); assert_arrays_eq!( - var_bin, + binary_arr, VarBinViewArray::from_iter_str(["string1", "string2"]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 51df660a501..6820050b72f 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -233,7 +233,6 @@ mod tests { use vortex::buffer::buffer; use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::error::VortexExpect; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use super::*; @@ -606,7 +605,8 @@ mod tests { .map(|i| if i % 1000 == 0 { 600 } else { i % 512 }) .collect::>(), NonNullable, - ).into_array(); + ) + .into_array(); let bitpacked = BitPacked::encode(&values, 9, &mut ctx)?; assert!(bitpacked.patches().is_some()); let array = if let Some(range) = range { From 18cfdab436007268dd7e0458c1815fa79c4baf61 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 23 Jun 2026 00:05:57 +0100 Subject: [PATCH 12/12] more Signed-off-by: Robert Kruszewski --- .../src/aggregate_fn/fns/sum/grouped.rs | 31 +- vortex-array/src/aggregate_fn/fns/sum/mod.rs | 25 +- vortex-array/src/array/typed.rs | 9 +- vortex-array/src/arrays/bool/array.rs | 93 +-- .../src/arrays/bool/compute/filter.rs | 33 +- vortex-array/src/arrays/bool/compute/take.rs | 33 +- vortex-array/src/arrays/bool/patch.rs | 9 +- vortex-array/src/arrays/bool/vtable/mod.rs | 18 +- .../src/arrays/bool/vtable/operations.rs | 5 +- vortex-array/src/arrays/chunked/array.rs | 17 +- .../src/arrays/chunked/compute/cast.rs | 9 +- .../src/arrays/chunked/compute/take.rs | 33 +- vortex-array/src/arrays/chunked/tests.rs | 54 +- .../src/arrays/chunked/vtable/operations.rs | 29 +- .../src/arrays/constant/compute/take.rs | 9 +- .../src/arrays/constant/vtable/canonical.rs | 4 +- vortex-array/src/arrays/datetime/test.rs | 9 +- .../src/arrays/decimal/compute/fill_null.rs | 17 +- .../src/arrays/decimal/compute/take.rs | 9 +- vortex-array/src/arrays/decimal/vtable/mod.rs | 25 +- vortex-array/src/arrays/dict/array.rs | 5 +- vortex-array/src/arrays/dict/compute/cast.rs | 13 +- .../src/arrays/dict/compute/fill_null.rs | 17 +- vortex-array/src/arrays/dict/compute/like.rs | 5 +- vortex-array/src/arrays/dict/compute/mod.rs | 16 +- vortex-array/src/arrays/dict/compute/slice.rs | 9 +- vortex-array/src/arrays/dict/tests.rs | 5 +- .../src/arrays/extension/compute/cast.rs | 4 +- .../arrays/filter/execute/fixed_size_list.rs | 13 +- .../src/arrays/filter/execute/listview.rs | 17 +- .../src/arrays/filter/execute/struct_.rs | 13 +- .../arrays/fixed_size_list/tests/filter.rs | 45 +- .../src/arrays/fixed_size_list/tests/take.rs | 17 +- vortex-array/src/arrays/list/tests.rs | 53 +- .../src/arrays/listview/conversion.rs | 32 +- .../src/arrays/listview/tests/basic.rs | 23 +- .../src/arrays/listview/tests/filter.rs | 10 +- .../src/arrays/listview/tests/operations.rs | 17 +- .../src/arrays/listview/tests/take.rs | 10 +- vortex-array/src/arrays/masked/tests.rs | 5 +- .../src/arrays/patched/compute/take.rs | 21 +- .../src/arrays/primitive/array/patch.rs | 5 +- .../src/arrays/primitive/compute/cast.rs | 52 +- .../src/arrays/primitive/compute/fill_null.rs | 17 +- .../src/arrays/primitive/compute/take/mod.rs | 5 +- .../src/arrays/primitive/vtable/mod.rs | 18 +- .../src/arrays/scalar_fn/vtable/operations.rs | 17 +- vortex-array/src/arrays/slice/vtable.rs | 9 +- .../src/arrays/struct_/compute/mod.rs | 9 +- .../src/arrays/struct_/compute/rules.rs | 18 +- vortex-array/src/arrays/struct_/tests.rs | 25 +- .../src/arrays/varbin/compute/compare.rs | 9 +- .../src/arrays/varbin/compute/filter.rs | 75 +-- .../src/arrays/varbin/compute/take.rs | 9 +- vortex-array/src/arrays/varbin/tests.rs | 9 +- .../src/arrays/varbin/vtable/canonical.rs | 5 +- vortex-array/src/arrays/varbinview/compact.rs | 37 +- .../src/arrays/varbinview/vtable/mod.rs | 18 +- vortex-array/src/arrays/variant/mod.rs | 11 +- vortex-array/src/builders/bool.rs | 5 +- vortex-array/src/builders/decimal.rs | 8 +- vortex-array/src/builders/dict/primitive.rs | 12 +- vortex-array/src/builders/extension.rs | 5 +- vortex-array/src/builders/listview.rs | 9 +- vortex-array/src/builders/primitive.rs | 9 +- vortex-array/src/builders/struct_.rs | 6 +- vortex-array/src/builders/varbinview.rs | 19 +- .../src/compute/conformance/filter.rs | 5 +- vortex-array/src/patches.rs | 287 ++++------ vortex-array/src/scalar_fn/fns/between/mod.rs | 6 +- .../src/scalar_fn/fns/binary/boolean.rs | 13 +- .../src/scalar_fn/fns/binary/compare.rs | 31 +- vortex-array/src/scalar_fn/fns/binary/mod.rs | 29 +- .../src/scalar_fn/fns/binary/numeric.rs | 41 +- vortex-array/src/scalar_fn/fns/byte_length.rs | 9 +- vortex-array/src/scalar_fn/fns/case_when.rs | 132 ++--- vortex-array/src/scalar_fn/fns/dynamic.rs | 35 +- .../src/scalar_fn/fns/fill_null/mod.rs | 21 +- .../src/scalar_fn/fns/list_contains/mod.rs | 39 +- vortex-array/src/scalar_fn/fns/merge.rs | 13 +- vortex-array/src/scalar_fn/fns/pack.rs | 19 +- vortex-array/src/scalar_fn/fns/zip/mod.rs | 13 +- vortex-array/src/stats/expr.rs | 61 +- vortex-btrblocks/src/canonical_compressor.rs | 32 +- vortex-btrblocks/src/schemes/float/tests.rs | 4 +- .../schemes/integer/scheme_selection_tests.rs | 4 +- vortex-btrblocks/src/schemes/integer/tests.rs | 12 +- vortex-compressor/src/sample.rs | 5 +- vortex-cuda/src/dynamic_dispatch/mod.rs | 539 +++++++----------- vortex-cuda/src/hybrid_dispatch/mod.rs | 29 +- .../src/kernel/encodings/zstd_buffers.rs | 9 +- vortex-duckdb/src/convert/vector.rs | 5 +- vortex-ffi/src/scan.rs | 17 +- vortex-file/src/tests.rs | 12 +- vortex-file/tests/test_write_table.rs | 3 +- vortex-ipc/src/iterator.rs | 3 +- vortex-ipc/src/stream.rs | 9 +- vortex-layout/src/layouts/chunked/reader.rs | 4 +- vortex-layout/src/layouts/dict/reader.rs | 45 +- vortex-layout/src/layouts/flat/reader.rs | 30 +- vortex-layout/src/layouts/row_idx/mod.rs | 36 +- vortex-layout/src/layouts/struct_/reader.rs | 14 +- vortex-layout/src/layouts/zoned/reader.rs | 9 +- vortex-python/src/arrays/range_to_sequence.rs | 25 +- vortex-test/compat-gen/src/check.rs | 5 +- .../src/fixtures/arrays/synthetic/mod.rs | 5 +- 106 files changed, 1198 insertions(+), 1654 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index 4f13c76777f..8f65737c48e 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -166,6 +166,7 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; + use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -236,7 +237,7 @@ mod tests { #[test] fn listview_matches_reference_unsigned() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1u32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::U32, NonNullable); @@ -249,14 +250,14 @@ mod tests { // Unsigned input sums to U64. let direct = PrimitiveArray::from_option_iter([Some(3u64), Some(3u64), Some(15u64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); - assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Offsets are not in group order and a group is null: the group validity must be indexed by // group index, not by element offset. let elements = @@ -271,14 +272,14 @@ mod tests { let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; let direct = PrimitiveArray::from_option_iter([Some(110i64), None, Some(70i64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); - assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_interior_and_full_nulls() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Group 1 has an interior null, group 2 is entirely null, group 3 is empty. let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) @@ -293,14 +294,14 @@ mod tests { let direct = PrimitiveArray::from_option_iter([Some(4i64), Some(0i64), Some(0i64), Some(9i64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); - assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } #[test] fn listview_overflow_group_is_null() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I64, NonNullable); @@ -313,8 +314,8 @@ mod tests { // First group overflows -> null sum; second group sums normally. let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); - assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } @@ -384,7 +385,7 @@ mod tests { #[test] fn fixed_size_overflow_and_nan() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // FixedSize path: first group overflows -> null sum, second sums normally. let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); @@ -396,8 +397,8 @@ mod tests { let expected = grouped_sum_reference(&elements, &[(0, 2), (2, 2)], &[true, true], &elem_dtype)?; let direct = PrimitiveArray::from_option_iter([None, Some(5i64)]); - assert_arrays_eq!(&actual, &direct.into_array(), &mut assertion_ctx); - assert_arrays_eq!(&actual, &expected, &mut assertion_ctx); + assert_arrays_eq!(&actual, &direct.into_array(), &mut ctx); + assert_arrays_eq!(&actual, &expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index 242f49f5073..e31fc7b8058 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -404,6 +404,7 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -586,7 +587,7 @@ mod tests { #[test] fn grouped_sum_fixed_size_list() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; @@ -595,13 +596,13 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(6i64), Some(15i64)]).into_array(); - assert_arrays_eq!(&result, &expected, &mut assertion_ctx); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_with_null_elements() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) .into_array(); @@ -611,13 +612,13 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(4i64), Some(11i64)]).into_array(); - assert_arrays_eq!(&result, &expected, &mut assertion_ctx); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_with_null_group() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) .into_array(); @@ -629,13 +630,13 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(6i64), None, Some(24i64)]).into_array(); - assert_arrays_eq!(&result, &expected, &mut assertion_ctx); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; @@ -644,13 +645,13 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(0i64), Some(7i64)]).into_array(); - assert_arrays_eq!(&result, &expected, &mut assertion_ctx); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } #[test] fn grouped_sum_bool() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); let groups = FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; @@ -659,7 +660,7 @@ mod tests { let result = run_grouped_sum(&groups.into_array(), &elem_dtype)?; let expected = PrimitiveArray::from_option_iter([Some(2u64), Some(3u64)]).into_array(); - assert_arrays_eq!(&result, &expected, &mut assertion_ctx); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } @@ -691,7 +692,7 @@ mod tests { #[test] fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); @@ -705,7 +706,7 @@ mod tests { // group 0 -> elements[2..3] = 300; group 1 -> null; group 2 -> elements[1..2] = 200. let expected = PrimitiveArray::from_option_iter([Some(300i64), None, Some(200i64)]).into_array(); - assert_arrays_eq!(&result, &expected, &mut assertion_ctx); + assert_arrays_eq!(&result, &expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 8ed98a254b3..4218ac1ae5d 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -490,6 +490,7 @@ mod tests { use super::Array; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -497,19 +498,19 @@ mod tests { #[test] fn typed_array_into_parts_roundtrips() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let expected = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let parts = array.try_into_parts().unwrap(); let rebuilt = Array::::try_from_parts(parts).unwrap(); - assert_arrays_eq!(rebuilt, expected, &mut assertion_ctx); + assert_arrays_eq!(rebuilt, expected, &mut ctx); } #[test] fn typed_array_try_into_parts_requires_unique_owner() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let alias = array.clone(); @@ -518,6 +519,6 @@ mod tests { Err(array) => array, }; - assert_arrays_eq!(array, alias, &mut assertion_ctx); + assert_arrays_eq!(array, alias, &mut ctx); } } diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index ab8e9611d36..8b83d0f8128 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -52,7 +52,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// ``` /// # fn main() -> vortex_error::VortexResult<()> { /// use vortex_array::arrays::BoolArray; -/// use vortex_array::{IntoArray, LEGACY_SESSION, VortexSessionExecute}; +/// use vortex_array::{IntoArray, array_session, VortexSessionExecute}; /// /// // Create from iterator using FromIterator impl /// let array: BoolArray = [true, false, true, false].into_iter().collect(); @@ -62,7 +62,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// assert_eq!(sliced.len(), 2); /// /// // Access individual values -/// let mut ctx = LEGACY_SESSION.create_execution_ctx(); +/// let mut ctx = array_session().create_execution_ctx(); /// let value = array.execute_scalar(0, &mut ctx).unwrap(); /// assert_eq!(value, true.into()); /// # Ok(()) @@ -371,8 +371,8 @@ mod tests { use vortex_buffer::buffer; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; @@ -382,77 +382,54 @@ mod tests { #[test] fn bool_array() { + let mut ctx = array_session().create_execution_ctx(); + let arr = BoolArray::from_iter([true, false, true]); - let scalar = bool::try_from( - &arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(0, &mut ctx).unwrap()).unwrap(); assert!(scalar); } #[test] fn test_all_some_iter() { + let mut ctx = array_session().create_execution_ctx(); + let arr = BoolArray::from_iter([Some(true), Some(false)]); assert!(matches!(arr.validity(), Ok(Validity::AllValid))); - let scalar = bool::try_from( - &arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(0, &mut ctx).unwrap()).unwrap(); assert!(scalar); - let scalar = bool::try_from( - &arr.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(1, &mut ctx).unwrap()).unwrap(); assert!(!scalar); } #[test] fn test_bool_from_iter() { + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); - let scalar = bool::try_from( - &arr.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(0, &mut ctx).unwrap()).unwrap(); assert!(scalar); - let scalar = bool::try_from( - &arr.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(1, &mut ctx).unwrap()).unwrap(); assert!(scalar); - let scalar = arr - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let scalar = arr.execute_scalar(2, &mut ctx).unwrap(); assert!(scalar.is_null()); - let scalar = bool::try_from( - &arr.execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - ) - .unwrap(); + let scalar = bool::try_from(&arr.execute_scalar(3, &mut ctx).unwrap()).unwrap(); assert!(!scalar); - let scalar = arr - .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let scalar = arr.execute_scalar(4, &mut ctx).unwrap(); assert!(scalar.is_null()); } #[test] fn patch_sliced_bools() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.slice(4..12).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut assertion_ctx); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut ctx); let arr = { let mut builder = BitBufferMut::new_unset(12); @@ -464,7 +441,7 @@ mod tests { assert_arrays_eq!( sliced, BoolArray::from_iter(expected_slice.clone()), - &mut assertion_ctx + &mut ctx ); // patch the underlying array at index 4 to false @@ -476,36 +453,26 @@ mod tests { None, ) .unwrap(); - let arr = arr - .patch(&patches, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let arr = arr.patch(&patches, &mut ctx).unwrap(); // After patching index 4 to false: indices 1-3 and 5-11 are true, index 0 and 4 are false let expected_patched: Vec = (0..12).map(|i| (1..12).contains(&i) && i != 4).collect(); - assert_arrays_eq!( - arr, - BoolArray::from_iter(expected_patched), - &mut assertion_ctx - ); + assert_arrays_eq!(arr, BoolArray::from_iter(expected_patched), &mut ctx); // the slice should be unchanged (still has original values before patch) - assert_arrays_eq!( - sliced, - BoolArray::from_iter(expected_slice), - &mut assertion_ctx - ); + assert_arrays_eq!(sliced, BoolArray::from_iter(expected_slice), &mut ctx); } #[test] fn slice_array_in_middle() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let sliced = arr.slice(4..12).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut assertion_ctx); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut ctx); } #[test] fn patch_bools_owned() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let buf_ptr = arr.to_bit_buffer().inner().as_ptr(); @@ -517,22 +484,20 @@ mod tests { None, ) .unwrap(); - let arr = arr - .patch(&patches, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let arr = arr.patch(&patches, &mut ctx).unwrap(); // Verify buffer was reused in place assert_eq!(arr.to_bit_buffer().inner().as_ptr(), buf_ptr); // After patching index 0 to false: [false, true, true, ..., true] (16 values) let expected: BoolArray = once(false).chain(repeat_n(true, 15)).collect(); - assert_arrays_eq!(arr, expected, &mut assertion_ctx); + assert_arrays_eq!(arr, expected, &mut ctx); } #[test] fn patch_sliced_bools_offset() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.slice(4..15).unwrap(); - assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11]), &mut assertion_ctx); + assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11]), &mut ctx); } } diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index e5a091c7e29..f463f27f437 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -304,64 +304,49 @@ mod tests { use super::*; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; #[test] fn filter_bool_test() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_iter([true, false, true]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!( - filtered, - BoolArray::from_iter([true, false]), - &mut assertion_ctx - ); + assert_arrays_eq!(filtered, BoolArray::from_iter([true, false]), &mut ctx); } #[test] fn filter_bool_sparse_index_mask() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_indices(3, [0, 2]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!( - filtered, - BoolArray::from_iter([true, false]), - &mut assertion_ctx - ); + assert_arrays_eq!(filtered, BoolArray::from_iter([true, false]), &mut ctx); } #[test] fn filter_bool_sparse_slice_mask() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_slices(3, vec![(0, 1), (2, 3)]); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!( - filtered, - BoolArray::from_iter([true, false]), - &mut assertion_ctx - ); + assert_arrays_eq!(filtered, BoolArray::from_iter([true, false]), &mut ctx); } #[test] fn filter_bool_sparse_buffer_mask() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_buffer(BitBuffer::from_iter([true, false, true])); let filtered = arr.filter(mask).unwrap(); - assert_arrays_eq!( - filtered, - BoolArray::from_iter([true, false]), - &mut assertion_ctx - ); + assert_arrays_eq!(filtered, BoolArray::from_iter([true, false]), &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index f14f13a29c3..abff3527bfe 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -91,6 +91,7 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; @@ -100,7 +101,7 @@ mod test { #[test] fn take_nullable() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let reference = BoolArray::from_iter(vec![ Some(false), Some(true), @@ -121,16 +122,12 @@ mod test { let all_invalid_indices = PrimitiveArray::from_option_iter([None::, None, None]); let b = reference.take(all_invalid_indices.into_array()).unwrap(); - assert_arrays_eq!( - b, - BoolArray::from_iter([None, None, None]), - &mut assertion_ctx - ); + assert_arrays_eq!(b, BoolArray::from_iter([None, None, None]), &mut ctx); } #[test] fn test_bool_array_take_with_null_out_of_bounds_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -142,13 +139,13 @@ mod test { assert_arrays_eq!( actual, BoolArray::from_iter([Some(false), None, None]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_non_null_bool_array_take_with_null_out_of_bounds_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -159,40 +156,32 @@ mod test { assert_arrays_eq!( actual, BoolArray::from_iter([Some(false), Some(true), None]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_bool_array_take_all_null_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([false, false, false]).into_array()), ); let actual = values.take(indices.into_array()).unwrap(); - assert_arrays_eq!( - actual, - BoolArray::from_iter([None, None, None]), - &mut assertion_ctx - ); + assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None]), &mut ctx); } #[test] fn test_non_null_bool_array_take_all_null_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], Validity::Array(BoolArray::from_iter([false, false, false]).into_array()), ); let actual = values.take(indices.into_array()).unwrap(); - assert_arrays_eq!( - actual, - BoolArray::from_iter([None, None, None]), - &mut assertion_ctx - ); + assert_arrays_eq!(actual, BoolArray::from_iter([None, None, None]), &mut ctx); } #[rstest] diff --git a/vortex-array/src/arrays/bool/patch.rs b/vortex-array/src/arrays/bool/patch.rs index 96652c4656d..babf5a0ae1d 100644 --- a/vortex-array/src/arrays/bool/patch.rs +++ b/vortex-array/src/arrays/bool/patch.rs @@ -48,24 +48,25 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; #[test] fn patch_sliced_bools() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.into_array().slice(4..12).unwrap(); let expected = BoolArray::from_iter([true; 8]); - assert_arrays_eq!(sliced, expected, &mut assertion_ctx); + assert_arrays_eq!(sliced, expected, &mut ctx); } #[test] fn patch_sliced_bools_offset() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.into_array().slice(4..15).unwrap(); let expected = BoolArray::from_iter([true; 11]); - assert_arrays_eq!(sliced, expected, &mut assertion_ctx); + assert_arrays_eq!(sliced, expected, &mut ctx); } } diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index b48f651ff5b..348a1f5f69a 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -211,8 +211,8 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -220,16 +220,17 @@ mod tests { #[test] fn test_nullable_bool_serde_roundtrip() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = BoolArray::from_iter([Some(true), None, Some(false), None]); let dtype = array.dtype().clone(); let len = array.len(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let serialized = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); let mut concat = ByteBufferMut::empty(); @@ -238,14 +239,9 @@ mod tests { } let parts = SerializedArray::try_from(concat.freeze()).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array, &mut assertion_ctx); + assert_arrays_eq!(decoded, array, &mut ctx); } } diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index a6ac9641ba7..9bae2b9e19c 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -31,6 +31,7 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; @@ -47,7 +48,7 @@ mod tests { #[test] fn test_slice() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); #[expect(deprecated)] let sliced_arr = arr.into_array().slice(1..4).unwrap().to_bool(); @@ -55,7 +56,7 @@ mod tests { assert_arrays_eq!( sliced_arr, BoolArray::from_iter([Some(true), None, Some(false)]), - &mut assertion_ctx + &mut ctx ); } } diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index d06a16556eb..bd6b3d2d06d 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -280,6 +280,7 @@ mod test { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; @@ -291,7 +292,7 @@ mod test { #[test] fn test_rechunk_one_chunk() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -300,12 +301,12 @@ mod test { let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap(); - assert_arrays_eq!(chunked, rechunked, &mut assertion_ctx); + assert_arrays_eq!(chunked, rechunked, &mut ctx); } #[test] fn test_rechunk_two_chunks() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array(), buffer![5u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -315,12 +316,12 @@ mod test { let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap(); assert_eq!(rechunked.nchunks(), 1); - assert_arrays_eq!(chunked, rechunked, &mut assertion_ctx); + assert_arrays_eq!(chunked, rechunked, &mut ctx); } #[test] fn test_rechunk_tiny_target_chunks() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2, 3].into_array(), @@ -334,12 +335,12 @@ mod test { assert_eq!(rechunked.nchunks(), 2); assert!(rechunked.iter_chunks().all(|c| c.len() < 5)); - assert_arrays_eq!(chunked, rechunked, &mut assertion_ctx); + assert_arrays_eq!(chunked, rechunked, &mut ctx); } #[test] fn test_rechunk_with_too_big_chunk() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2].into_array(), @@ -356,7 +357,7 @@ mod test { // greedy so should be: [0, 1, 2] [42, 42, 42, 42, 42, 42] [4, 5, 6, 7] [8, 9] assert_eq!(rechunked.nchunks(), 4); - assert_arrays_eq!(chunked, rechunked, &mut assertion_ctx); + assert_arrays_eq!(chunked, rechunked, &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index dbccad6ae02..80c6848b973 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -36,6 +36,7 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -47,7 +48,7 @@ mod test { #[test] fn test_cast_chunked() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr0 = buffer![0u32, 1].into_array(); let arr1 = buffer![2u32, 3].into_array(); @@ -69,11 +70,7 @@ mod test { let result = root .cast(DType::Primitive(PType::U64, Nullability::NonNullable)) .unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([0u64, 1, 2, 3]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, PrimitiveArray::from_iter([0u64, 1, 2, 3]), &mut ctx); } #[rstest] diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index e19811b7379..34f443f1ce1 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -129,6 +129,7 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; @@ -142,7 +143,7 @@ mod test { #[test] fn test_take() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -151,16 +152,12 @@ mod test { let indices = buffer![0u64, 0, 6, 4].into_array(); let result = arr.take(indices).unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([1i32, 1, 1, 2]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 1, 1, 2]), &mut ctx); } #[test] fn test_take_nullable_values() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let a = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -172,13 +169,13 @@ mod test { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([1i32, 1, 1, 2].map(Some)), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_take_nullable_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -193,13 +190,13 @@ mod test { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(1i32), None, None, Some(2)]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_take_nullable_struct() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let struct_array = StructArray::try_new(FieldNames::default(), vec![], 100, Validity::NonNullable) .unwrap(); @@ -220,12 +217,12 @@ mod test { Validity::Array(BoolArray::from_iter(vec![true, false, true]).into_array()), ) .unwrap(); - assert_arrays_eq!(result, expect, &mut assertion_ctx); + assert_arrays_eq!(result, expect, &mut ctx); } #[test] fn test_empty_take() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -240,13 +237,13 @@ mod test { assert_arrays_eq!( result, PrimitiveArray::empty::(Nullability::NonNullable), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_take_shuffled_indices() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let c0 = buffer![0i32, 1, 2].into_array(); let c1 = buffer![3i32, 4, 5].into_array(); let c2 = buffer![6i32, 7, 8].into_array(); @@ -264,7 +261,7 @@ mod test { assert_arrays_eq!( result, PrimitiveArray::from_iter([8i32, 0, 5, 3, 2, 7, 1, 6, 4]), - &mut assertion_ctx + &mut ctx ); Ok(()) } @@ -315,7 +312,7 @@ mod test { #[test] fn test_take_null_indices() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let c0 = buffer![10i32, 20, 30].into_array(); let c1 = buffer![40i32, 50, 60].into_array(); let arr = ChunkedArray::try_new( @@ -340,7 +337,7 @@ mod test { None, Some(30) ]), - &mut assertion_ctx + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index 5f3ef6db648..d3cae8be8f0 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -203,7 +203,7 @@ fn execute_path_nested_chunked_dict_of_dict_into_canonical() { #[test] fn with_slot_rewrites_chunk_and_offsets() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = chunked_array().into_array(); let replacement = buffer![10u64, 11, 12].into_array(); @@ -215,12 +215,12 @@ fn with_slot_rewrites_chunk_and_offsets() { assert_arrays_eq!( array.chunk(0).clone(), PrimitiveArray::from_iter([10u64, 11, 12]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( array.array().clone(), PrimitiveArray::from_iter([10u64, 11, 12, 4, 5, 6, 7, 8, 9]), - &mut assertion_ctx + &mut ctx ); } @@ -236,61 +236,61 @@ fn with_slot_rejects_len_mismatch() { #[test] fn slice_middle() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(2..5).unwrap(), PrimitiveArray::from_iter([3u64, 4, 5]), - &mut assertion_ctx + &mut ctx ); } #[test] fn slice_begin() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(1..3).unwrap(), PrimitiveArray::from_iter([2u64, 3]), - &mut assertion_ctx + &mut ctx ); } #[test] fn slice_aligned() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(3..6).unwrap(), PrimitiveArray::from_iter([4u64, 5, 6]), - &mut assertion_ctx + &mut ctx ); } #[test] fn slice_many_aligned() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(0..6).unwrap(), PrimitiveArray::from_iter([1u64, 2, 3, 4, 5, 6]), - &mut assertion_ctx + &mut ctx ); } #[test] fn slice_end() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(7..8).unwrap(), PrimitiveArray::from_iter([8u64]), - &mut assertion_ctx + &mut ctx ); } #[test] fn slice_exactly_end() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(6..9).unwrap(), PrimitiveArray::from_iter([7u64, 8, 9]), - &mut assertion_ctx + &mut ctx ); } @@ -304,7 +304,7 @@ fn slice_empty() { #[test] fn scalar_at_empty_children_both_sides() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -316,16 +316,12 @@ fn scalar_at_empty_children_both_sides() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!( - array, - PrimitiveArray::from_iter([1u64, 2]), - &mut assertion_ctx - ); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2]), &mut ctx); } #[test] fn scalar_at_empty_children_trailing() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -336,16 +332,12 @@ fn scalar_at_empty_children_trailing() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!( - array, - PrimitiveArray::from_iter([1u64, 2, 3, 4]), - &mut assertion_ctx - ); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx); } #[test] fn scalar_at_empty_children_leading() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -356,11 +348,7 @@ fn scalar_at_empty_children_leading() { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!( - array, - PrimitiveArray::from_iter([1u64, 2, 3, 4]), - &mut assertion_ctx - ); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 12855434f13..8f9e0867a88 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -31,6 +31,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -58,11 +59,11 @@ mod tests { #[case::end(7..8, &[8u64])] #[case::exactly_end(6..9, &[7u64, 8, 9])] fn slice(#[case] range: Range, #[case] expected: &[u64]) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(range).unwrap(), PrimitiveArray::from_iter(expected.iter().copied()), - &mut assertion_ctx + &mut ctx ); } @@ -76,7 +77,7 @@ mod tests { #[test] fn scalar_at_empty_children_both_sides() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -88,16 +89,12 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!( - array, - PrimitiveArray::from_iter([1u64, 2]), - &mut assertion_ctx - ); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2]), &mut ctx); } #[test] fn scalar_at_empty_children_trailing() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -108,16 +105,12 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!( - array, - PrimitiveArray::from_iter([1u64, 2, 3, 4]), - &mut assertion_ctx - ); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx); } #[test] fn scalar_at_empty_children_leading() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -128,10 +121,6 @@ mod tests { DType::Primitive(PType::U64, Nullability::NonNullable), ) .unwrap(); - assert_arrays_eq!( - array, - PrimitiveArray::from_iter([1u64, 2, 3, 4]), - &mut assertion_ctx - ); + assert_arrays_eq!(array, PrimitiveArray::from_iter([1u64, 2, 3, 4]), &mut ctx); } } diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index 3b2298b0bf1..2330cd28f88 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -77,6 +77,7 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -87,7 +88,7 @@ mod tests { #[test] fn take_nullable_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take( @@ -110,7 +111,7 @@ mod tests { buffer![42i32, 42, 42], Validity::from_iter([false, true, false]) ), - &mut assertion_ctx + &mut ctx ); assert_eq!( taken @@ -125,7 +126,7 @@ mod tests { #[test] fn take_all_valid_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take(PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array()) @@ -138,7 +139,7 @@ mod tests { #[expect(deprecated)] taken.to_primitive(), PrimitiveArray::new(buffer![42i32, 42, 42], Validity::AllValid), - &mut assertion_ctx + &mut ctx ); assert_eq!( taken diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 8605877f5ac..e160bcc16a9 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -384,11 +384,11 @@ mod tests { #[test] fn test_canonicalize_const_str() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let const_array = ConstantArray::new("four".to_string(), 4); let expected = VarBinArray::from(vec!["four", "four", "four", "four"]); - assert_arrays_eq!(const_array, expected, &mut assertion_ctx); + assert_arrays_eq!(const_array, expected, &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index e091ff3848d..1ae114874d9 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -11,6 +11,7 @@ use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::datetime::TemporalData; use crate::assert_arrays_eq; @@ -27,14 +28,14 @@ use crate::validity::Validity; macro_rules! test_temporal_roundtrip { ($prim:ty, $constructor:expr, $unit:expr) => {{ - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = buffer![100 as $prim].into_array(); let temporal: TemporalData = $constructor(array, $unit); assert_arrays_eq!( temporal.temporal_values(), PrimitiveArray::from_iter([100 as $prim]), - &mut assertion_ctx + &mut ctx ); assert_eq!(temporal.temporal_metadata().time_unit(), $unit); }}; @@ -148,7 +149,7 @@ test_fail_case!( // We test Timestamp explicitly to avoid the macro getting too complex. #[test] fn test_timestamp() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let ts = buffer![100i64].into_array(); let ts_array = ts.into_array(); @@ -164,7 +165,7 @@ fn test_timestamp() { assert_arrays_eq!( temporal_array.temporal_values(), PrimitiveArray::from_iter([100i64]), - &mut assertion_ctx + &mut ctx ); assert_eq!( temporal_array.temporal_metadata(), diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index f12a1428eb5..3ab8762e875 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -91,6 +91,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; @@ -104,7 +105,7 @@ mod tests { #[test] fn fill_null_leading_none() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( [None, Some(800i128), None, Some(1000i128), None], @@ -123,7 +124,7 @@ mod tests { assert_arrays_eq!( p, DecimalArray::from_iter([4200, 800, 4200, 1000, 4200], decimal_dtype), - &mut assertion_ctx + &mut ctx ); assert_eq!( p.buffer::().as_slice(), @@ -141,7 +142,7 @@ mod tests { #[test] fn fill_null_all_none() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( @@ -162,14 +163,14 @@ mod tests { assert_arrays_eq!( p, DecimalArray::from_iter([25500, 25500, 25500, 25500, 25500], decimal_dtype), - &mut assertion_ctx + &mut ctx ); } /// fill_null with a value that overflows the array's storage type should upcast the array. #[test] fn fill_null_overflow_upcasts() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(3, 0); let arr = DecimalArray::from_option_iter([None, Some(10i8), None], decimal_dtype); // i8 max is 127, so 200 doesn't fit — the array should be widened to i16. @@ -186,13 +187,13 @@ mod tests { assert_arrays_eq!( result, DecimalArray::from_iter([200i16, 10, 200], decimal_dtype), - &mut assertion_ctx + &mut ctx ); } #[test] fn fill_null_non_nullable() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::new( @@ -213,7 +214,7 @@ mod tests { assert_arrays_eq!( p, DecimalArray::from_iter([800i128, 1000, 1200, 1400, 1600], decimal_dtype), - &mut assertion_ctx + &mut ctx ); } } diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 8f63637b1e9..762ac3d7365 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -54,6 +54,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -63,7 +64,7 @@ mod tests { #[test] fn test_take() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![10i128, 11i128, 12i128, 13i128], @@ -75,12 +76,12 @@ mod tests { let taken = array.take(indices).unwrap(); let expected = DecimalArray::from_iter([10i128, 12, 13], ddtype); - assert_arrays_eq!(expected, taken, &mut assertion_ctx); + assert_arrays_eq!(expected, taken, &mut ctx); } #[test] fn test_take_null_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![i128::MAX, 11i128, 12i128, 13i128], @@ -92,7 +93,7 @@ mod tests { let taken = array.take(indices).unwrap(); let expected = DecimalArray::from_option_iter([None, Some(12i128), Some(13)], ddtype); - assert_arrays_eq!(expected, taken, &mut assertion_ctx); + assert_arrays_eq!(expected, taken, &mut ctx); } #[rstest] diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index e7f5e8c5f19..a52f2882238 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -208,8 +208,8 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; @@ -220,6 +220,7 @@ mod tests { #[test] fn test_array_serde() { + let session = array_session(); let array = DecimalArray::new( buffer![100i128, 200i128, 300i128, 400i128, 500i128], DecimalDType::new(10, 2), @@ -227,10 +228,10 @@ mod tests { ); let dtype = array.dtype().clone(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let out = array .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); // Concat into a single buffer let mut concat = ByteBufferMut::empty(); @@ -242,14 +243,15 @@ mod tests { let parts = SerializedArray::try_from(concat).unwrap(); let decoded = parts - .decode(&dtype, 5, &ReadContext::new(ctx.to_ids()), &LEGACY_SESSION) + .decode(&dtype, 5, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); assert!(decoded.is::()); } #[test] fn test_nullable_decimal_serde_roundtrip() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = DecimalArray::new( buffer![1234567i32, 0i32, -9999999i32], DecimalDType::new(7, 3), @@ -258,11 +260,11 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let out = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); let mut concat = ByteBufferMut::empty(); for buf in out { @@ -271,14 +273,9 @@ mod tests { let parts = SerializedArray::try_from(concat.freeze()).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array, &mut assertion_ctx); + assert_arrays_eq!(decoded, array, &mut ctx); } } diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index 79f1b538a63..ad777ddcad9 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -311,6 +311,7 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; @@ -462,7 +463,7 @@ mod test { #[test] fn test_dict_array_from_primitive_chunks() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let len = 2; let chunk_count = 2; let array = make_dict_primitive_chunks::(len, 2, chunk_count); @@ -477,7 +478,7 @@ mod test { let into_prim = array.to_primitive(); let prim_into = builder.finish_into_canonical().into_primitive(); - assert_arrays_eq!(into_prim, prim_into, &mut assertion_ctx); + assert_arrays_eq!(into_prim, prim_into, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/dict/compute/cast.rs b/vortex-array/src/arrays/dict/compute/cast.rs index 6ce6846c389..c4438891176 100644 --- a/vortex-array/src/arrays/dict/compute/cast.rs +++ b/vortex-array/src/arrays/dict/compute/cast.rs @@ -71,7 +71,6 @@ mod tests { #[test] fn test_cast_dict_to_wider_type() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let values = buffer![1i32, 2, 3, 2, 1].into_array(); let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); @@ -89,7 +88,7 @@ mod tests { assert_arrays_eq!( decoded, PrimitiveArray::from_iter([1i64, 2, 3, 2, 1]), - &mut assertion_ctx + &mut SESSION.create_execution_ctx() ); } @@ -111,7 +110,6 @@ mod tests { #[test] fn test_cast_dict_allvalid_to_nonnullable_and_back() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); // Create an AllValid dict array (no nulls) let values = buffer![10i32, 20, 30, 40].into_array(); let dict = dict_encode(&values, &mut SESSION.create_execution_ctx()).unwrap(); @@ -179,7 +177,11 @@ mod tests { let original_values = dict.as_array().to_primitive(); #[expect(deprecated)] let final_values = back_to_non_nullable.to_primitive(); - assert_arrays_eq!(original_values, final_values, &mut assertion_ctx); + assert_arrays_eq!( + original_values, + final_values, + &mut SESSION.create_execution_ctx() + ); } #[rstest] @@ -193,7 +195,6 @@ mod tests { #[test] fn test_cast_dict_with_unreferenced_null_values_to_nonnullable() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); use crate::arrays::DictArray; use crate::validity::Validity; @@ -232,7 +233,7 @@ mod tests { assert_arrays_eq!( casted_prim, PrimitiveArray::from_iter([1.0f64, 3.0, 1.0]), - &mut assertion_ctx + &mut SESSION.create_execution_ctx() ); } } diff --git a/vortex-array/src/arrays/dict/compute/fill_null.rs b/vortex-array/src/arrays/dict/compute/fill_null.rs index 40176f03f28..290fdd95bda 100644 --- a/vortex-array/src/arrays/dict/compute/fill_null.rs +++ b/vortex-array/src/arrays/dict/compute/fill_null.rs @@ -94,10 +94,8 @@ mod tests { use vortex_error::VortexExpect; use crate::IntoArray; - use crate::LEGACY_SESSION; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -108,7 +106,7 @@ mod tests { #[test] fn nullable_codes_fill_in_values() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::new( buffer![0u32, 1, 2], @@ -123,17 +121,12 @@ mod tests { .into_array() .fill_null(Scalar::primitive(20, Nullability::NonNullable)) .vortex_expect("operation should succeed in test"); - #[expect(deprecated)] - let filled_primitive = filled.to_primitive(); + let filled_primitive = filled.execute::(&mut ctx).unwrap(); assert_arrays_eq!( filled_primitive, PrimitiveArray::from_iter([10, 20, 20]), - &mut assertion_ctx - ); - assert!( - filled_primitive - .all_valid(&mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + &mut ctx ); + assert!(filled_primitive.all_valid(&mut ctx).unwrap()); } } diff --git a/vortex-array/src/arrays/dict/compute/like.rs b/vortex-array/src/arrays/dict/compute/like.rs index c87696597c3..a4ca7c1f76e 100644 --- a/vortex-array/src/arrays/dict/compute/like.rs +++ b/vortex-array/src/arrays/dict/compute/like.rs @@ -57,6 +57,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::DictArray; use crate::arrays::VarBinArray; @@ -69,7 +70,7 @@ mod tests { #[test] fn like_reduce_dict() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::try_new( buffer![0u8, 1, 0, 2].into_array(), VarBinArray::from(vec!["hello", "world", "help"]).into_array(), @@ -84,7 +85,7 @@ mod tests { assert_arrays_eq!( result, BoolArray::from_iter([true, false, true, false]), - &mut assertion_ctx + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index ed0f3758009..0e330c8d820 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -86,7 +86,7 @@ mod test { #[test] fn canonicalise_nullable_primitive() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values: Vec> = (0..65) .map(|i| match i % 3 { 0 => Some(42), @@ -106,12 +106,12 @@ mod test { let expected = PrimitiveArray::from_option_iter(values); - assert_arrays_eq!(actual, expected, &mut assertion_ctx); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn canonicalise_non_nullable_primitive_32_unique_values() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let unique_values: Vec = (0..32).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 32])); @@ -123,12 +123,12 @@ mod test { #[expect(deprecated)] let actual = dict.as_array().to_primitive(); - assert_arrays_eq!(actual, expected, &mut assertion_ctx); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn canonicalise_non_nullable_primitive_100_unique_values() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let unique_values: Vec = (0..100).collect(); let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 100])); @@ -140,7 +140,7 @@ mod test { #[expect(deprecated)] let actual = dict.as_array().to_primitive(); - assert_arrays_eq!(actual, expected, &mut assertion_ctx); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] @@ -183,7 +183,7 @@ mod test { #[test] fn compare_sliced_dict() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); use crate::arrays::BoolArray; let sliced = sliced_dict_array(); let compared = sliced @@ -191,7 +191,7 @@ mod test { .unwrap(); let expected = BoolArray::from_iter([Some(false), None, Some(true)]); - assert_arrays_eq!(compared, expected.into_array(), &mut assertion_ctx); + assert_arrays_eq!(compared, expected.into_array(), &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/dict/compute/slice.rs b/vortex-array/src/arrays/dict/compute/slice.rs index f557a36e0da..376621fba2f 100644 --- a/vortex-array/src/arrays/dict/compute/slice.rs +++ b/vortex-array/src/arrays/dict/compute/slice.rs @@ -92,6 +92,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::dict::compute::slice::ConstantArray; @@ -103,20 +104,20 @@ mod tests { #[test] fn slice_constant_valid_code() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(1u8, 5).into_array(), buffer![10i32, 20, 30].into_array(), ); let sliced = dict.slice(1..4)?; let expected = PrimitiveArray::from_iter([20i32, 20, 20]).into_array(); - assert_arrays_eq!(sliced, expected, &mut assertion_ctx); + assert_arrays_eq!(sliced, expected, &mut ctx); Ok(()) } #[test] fn slice_constant_null_code() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(Scalar::null(DType::Primitive(PType::U8, Nullable)), 5).into_array(), buffer![10i32, 20, 30].into_array(), @@ -124,7 +125,7 @@ mod tests { let sliced = dict.slice(1..4)?; let expected = PrimitiveArray::from_option_iter([Option::::None, None, None]).into_array(); - assert_arrays_eq!(sliced, expected, &mut assertion_ctx); + assert_arrays_eq!(sliced, expected, &mut ctx); Ok(()) } } diff --git a/vortex-array/src/arrays/dict/tests.rs b/vortex-array/src/arrays/dict/tests.rs index 2e52ab20bf0..7961e2b7037 100644 --- a/vortex-array/src/arrays/dict/tests.rs +++ b/vortex-array/src/arrays/dict/tests.rs @@ -6,6 +6,7 @@ use vortex_buffer::buffer; use super::DictArray; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; @@ -15,7 +16,7 @@ use crate::validity::Validity; #[test] fn test_scalar_at_null_code() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::from_option_iter(vec![None, Some(0u32), None]).into_array(), buffer![1i32].into_array(), @@ -23,7 +24,7 @@ fn test_scalar_at_null_code() { .unwrap(); let expected = PrimitiveArray::from_option_iter(vec![None, Some(1i32), None]).into_array(); - assert_arrays_eq!(dict, expected, &mut assertion_ctx); + assert_arrays_eq!(dict, expected, &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index 4347594b797..dfce610f43a 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -121,7 +121,7 @@ mod tests { #[test] fn cast_timestamp_to_i64() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let ext_dtype = Timestamp::new_with_tz( TimeUnit::Nanoseconds, Some("UTC".into()), @@ -136,7 +136,7 @@ mod tests { result.dtype(), &DType::Primitive(PType::I64, Nullability::NonNullable) ); - assert_arrays_eq!(result, buffer![1i64, 2, 3].into_array(), &mut assertion_ctx); + assert_arrays_eq!(result, buffer![1i64, 2, 3].into_array(), &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index 7e293987d13..392a6373f76 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -125,6 +125,7 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -150,7 +151,7 @@ mod test { #[test] fn filter_fixed_size_list_selects_correct_lists() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let array = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -162,12 +163,12 @@ mod test { let expected = FixedSizeListArray::new(expected_elements.into_array(), 2, Validity::NonNullable, 2); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn filter_degenerate_list_size_zero() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::empty::(Nullability::NonNullable); let array = FixedSizeListArray::new(elements.into_array(), 0, Validity::NonNullable, 5); @@ -178,12 +179,12 @@ mod test { let expected = FixedSizeListArray::new(expected_elements.into_array(), 0, Validity::NonNullable, 3); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn filter_nested_fixed_size_lists() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Inner lists of size 2, outer lists of size 2 (so 2 outer lists, each with 2 inner lists). let inner_elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); let inner_fsl = FixedSizeListArray::new(inner_elements, 2, Validity::NonNullable, 4); @@ -200,6 +201,6 @@ mod test { let expected_outer = FixedSizeListArray::new(expected_inner.into_array(), 2, Validity::NonNullable, 1); - assert_arrays_eq!(filtered, expected_outer, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected_outer, &mut ctx); } } diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index 0419a439183..c8118143e94 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -67,6 +67,7 @@ mod test { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::filter::execute::ConstantArray; @@ -75,7 +76,7 @@ mod test { use crate::compute::conformance::filter::test_filter_conformance; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = LazyLock::new(array_session); #[test] fn test_filter_listview_conformance() { @@ -135,7 +136,7 @@ mod test { #[test] fn filter_listview_selects_correct_lists() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // 3 lists: [10,20], [30,40], [50,60] let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let offsets = buffer![0u32, 2, 4].into_array(); @@ -158,12 +159,12 @@ mod test { Validity::NonNullable, ); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_filter_preserves_unreferenced_elements() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -188,7 +189,7 @@ mod test { assert_arrays_eq!( result_list.elements(), PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), - &mut assertion_ctx + &mut ctx ); // Verify offsets are unchanged. @@ -198,7 +199,7 @@ mod test { #[test] fn test_filter_with_gaps() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] @@ -223,7 +224,7 @@ mod test { assert_arrays_eq!( result_list.elements(), PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), - &mut assertion_ctx + &mut ctx ); // Verify offsets are unchanged. @@ -235,7 +236,7 @@ mod test { assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), PrimitiveArray::from_iter([7i32, 8, 9]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index 58600097343..8980c422ef3 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -50,6 +50,7 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -87,7 +88,7 @@ mod test { #[test] fn filter_struct_selects_correct_rows() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = StructArray::try_new( ["x", "y"].into(), vec![ @@ -113,12 +114,12 @@ mod test { ) .unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn filter_empty_struct() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let mask = Mask::from_iter([ @@ -128,19 +129,19 @@ mod test { let expected = StructArray::try_new(FieldNames::empty(), vec![], 5, Validity::NonNullable).unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn filter_empty_struct_with_empty_filter() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); let filtered = struct_arr.filter(Mask::from_iter::<[bool; 0]>([])).unwrap(); let expected = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs index 5c0ac468254..6f609120901 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs @@ -9,6 +9,7 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -52,7 +53,7 @@ fn test_filter_degenerate_list_size_zero( #[case] mask_values: Vec, #[case] expected_len: usize, ) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let new_validity = if matches!(validity, Validity::NonNullable) { Validity::NonNullable } else { @@ -76,12 +77,12 @@ fn test_filter_degenerate_list_size_zero( new_validity, expected_len, ); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_filter_with_nulls() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4), Some(5), Some(6)]); let validity = Validity::from_iter([true, false, true]); @@ -104,12 +105,12 @@ fn test_filter_with_nulls() { 2, ); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_filter_all_null_array() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create an array where all elements are null. let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let validity = Validity::AllInvalid; @@ -120,13 +121,13 @@ fn test_filter_all_null_array() { // Verify the result is an array of nulls. assert_eq!(filtered.len(), 2, "All-null FSL should produce 2 elements"); - assert_nth_scalar_is_null!(filtered, 0, &mut assertion_ctx); - assert_nth_scalar_is_null!(filtered, 1, &mut assertion_ctx); + assert_nth_scalar_is_null!(filtered, 0, &mut ctx); + assert_nth_scalar_is_null!(filtered, 1, &mut ctx); } #[test] fn test_filter_nested_fixed_size_lists() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create nested fixed-size lists: FSL>. // Inner lists are of size 2, outer lists are of size 3. // So we have 2 outer lists, each containing 3 inner lists, each containing 2 i32s. @@ -165,7 +166,7 @@ fn test_filter_nested_fixed_size_lists() { let expected_outer = FixedSizeListArray::new(expected_inner.into_array(), 3, Validity::NonNullable, 1); - assert_arrays_eq!(filtered, expected_outer, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected_outer, &mut ctx); } // Conformance tests using rstest for various array configurations. @@ -252,7 +253,7 @@ fn create_fsl_empty() -> ArrayRef { #[test] fn test_filter_all_null_various_list_sizes() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test filtering with all-null arrays of different list sizes. // The implementation returns ConstantArray only when validity_mask() is Mask::AllFalse. @@ -263,8 +264,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered0 = fsl0.filter(mask0).unwrap(); assert_eq!(filtered0.len(), 2); // Check that all elements are null (might be ConstantArray or FixedSizeListArray). - assert_nth_scalar_is_null!(filtered0, 0, &mut assertion_ctx); - assert_nth_scalar_is_null!(filtered0, 1, &mut assertion_ctx); + assert_nth_scalar_is_null!(filtered0, 0, &mut ctx); + assert_nth_scalar_is_null!(filtered0, 1, &mut ctx); // Case 2: list_size == 1. let elements1 = buffer![1i32, 2, 3].into_array(); @@ -273,8 +274,8 @@ fn test_filter_all_null_various_list_sizes() { let filtered1 = fsl1.filter(mask1).unwrap(); assert_eq!(filtered1.len(), 2); // Check that all elements are null. - assert_nth_scalar_is_null!(filtered1, 0, &mut assertion_ctx); - assert_nth_scalar_is_null!(filtered1, 1, &mut assertion_ctx); + assert_nth_scalar_is_null!(filtered1, 0, &mut ctx); + assert_nth_scalar_is_null!(filtered1, 1, &mut ctx); // Case 3: list_size == 10 (large). let elements10 = buffer![0..50i32].into_array(); @@ -283,15 +284,15 @@ fn test_filter_all_null_various_list_sizes() { let filtered10 = fsl10.filter(mask10).unwrap(); assert_eq!(filtered10.len(), 5); // Check that all elements are null. - assert_nth_scalar_is_null!(filtered10, 0, &mut assertion_ctx); - assert_nth_scalar_is_null!(filtered10, 4, &mut assertion_ctx); + assert_nth_scalar_is_null!(filtered10, 0, &mut ctx); + assert_nth_scalar_is_null!(filtered10, 4, &mut ctx); } // Note: test_filter_to_empty_degenerate has been consolidated into test_filter_degenerate_list_size_zero above. #[test] fn test_mask_expansion_threshold_boundary() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test with list_size == 8 (the FSL_SPARSE_MASK_LIST_SIZE_THRESHOLD). let list_size = 8u32; let num_lists = 100; @@ -330,7 +331,7 @@ fn test_mask_expansion_threshold_boundary() { 3, ); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); // Test with list_size == 7 (just below threshold). let list_size_7 = 7u32; @@ -362,13 +363,13 @@ fn test_mask_expansion_threshold_boundary() { 3, ); - assert_arrays_eq!(filtered7, expected7, &mut assertion_ctx); + assert_arrays_eq!(filtered7, expected7, &mut ctx); } // Test FSL-specific behavior with very large list sizes. #[test] fn test_filter_large_list_size() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test with list_size=100, which is significantly larger than typical use cases. let list_size = 100u32; let num_lists = 5; @@ -402,7 +403,7 @@ fn test_filter_large_list_size() { 3, ); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); // Test edge case: filter out all but one large list. let mask_single = Mask::from_iter([false, false, true, false, false]); @@ -422,5 +423,5 @@ fn test_filter_large_list_size() { 1, ); - assert_arrays_eq!(filtered_single, expected_single, &mut assertion_ctx); + assert_arrays_eq!(filtered_single, expected_single, &mut ctx); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index 6748f99c1e8..b0c41c30d27 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -13,6 +13,7 @@ use crate::ArrayRef; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -40,7 +41,7 @@ fn test_take_fsl_conformance(#[case] fsl: FixedSizeListArray) { #[test] fn test_take_basic_smoke_test() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -54,7 +55,7 @@ fn test_take_basic_smoke_test() { Validity::NonNullable, 3, ); - assert_arrays_eq!(expected, result, &mut assertion_ctx); + assert_arrays_eq!(expected, result, &mut ctx); } // Parameterized test for FSL-specific degenerate (list_size=0) cases. @@ -108,7 +109,7 @@ fn test_take_degenerate_lists( #[test] fn test_take_large_list_size() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = buffer![0i32..300].into_array(); let fsl = FixedSizeListArray::new(elements, 100, Validity::NonNullable, 3); @@ -118,12 +119,12 @@ fn test_take_large_list_size() { // Expected: [[200..300], [0..100]] let expected_elems = PrimitiveArray::from_iter((200i32..300).chain(0..100)).into_array(); let expected = FixedSizeListArray::new(expected_elems, 100, Validity::NonNullable, 2); - assert_arrays_eq!(expected, result, &mut assertion_ctx); + assert_arrays_eq!(expected, result, &mut ctx); } #[test] fn test_take_fsl_with_null_indices_preserves_elements() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -138,7 +139,7 @@ fn test_take_fsl_with_null_indices_preserves_elements() { Validity::from_iter([true, false, true]), 3, ); - assert_arrays_eq!(expected, result, &mut assertion_ctx); + assert_arrays_eq!(expected, result, &mut ctx); } // Element index overflow: with u8 indices and list_size=16, data_idx=16 produces element index @@ -170,9 +171,9 @@ fn test_element_index_overflow( #[case] indices: ArrayRef, #[case] expected: FixedSizeListArray, ) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = fsl.take(indices).unwrap(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } // Parameterized test for nullable array scenarios that are specific to FSL's implementation. diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 382e0d9ff14..906f4e03588 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -15,6 +15,7 @@ use super::*; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::FilterArray; use crate::arrays::List; use crate::arrays::PrimitiveArray; @@ -28,7 +29,7 @@ use crate::scalar::Scalar; use crate::validity::Validity; /// A shared session for `List` tests, used to create execution contexts. -static SESSION: LazyLock = LazyLock::new(crate::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); #[test] fn test_empty_list_array() { @@ -113,7 +114,7 @@ fn test_simple_list_filter() { #[test] fn test_list_filter_dense_mask() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test filtering with a dense mask (high density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -140,12 +141,12 @@ fn test_list_filter_dense_mask() { ) .unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_list_filter_sparse_mask() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test filtering with a sparse mask (low density of true values). let elements = buffer![0..100].into_array(); let offsets = buffer![0, 10, 25, 40, 60, 85, 100].into_array(); @@ -174,12 +175,12 @@ fn test_list_filter_sparse_mask() { ) .unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_list_filter_empty_lists() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test filtering arrays that contain empty lists. let elements = buffer![0..10].into_array(); let offsets = buffer![0, 0, 3, 3, 7, 10, 10].into_array(); // Lists at indices 0, 2, 5 are empty. @@ -203,7 +204,7 @@ fn test_list_filter_empty_lists() { ) .unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] @@ -236,7 +237,7 @@ fn test_list_filter_with_nulls() { #[test] fn test_list_filter_all_true() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test filtering with an all-true mask. let elements = buffer![0..20].into_array(); let offsets = buffer![0, 5, 10, 15, 20].into_array(); @@ -254,7 +255,7 @@ fn test_list_filter_all_true() { assert_eq!(filtered.len(), 4); let expected = ListArray::try_new(elements, offsets, validity).unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] @@ -279,7 +280,7 @@ fn test_list_filter_all_false() { #[test] fn test_list_filter_single_element() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test filtering to keep only one element. let elements = buffer![0..50].into_array(); let offsets = buffer![0, 10, 20, 30, 40, 50].into_array(); @@ -303,12 +304,12 @@ fn test_list_filter_single_element() { ) .unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_list_filter_alternating_pattern() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test filtering with an alternating pattern. let elements = buffer![0..60].into_array(); let offsets = buffer![0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55, 60].into_array(); @@ -339,12 +340,12 @@ fn test_list_filter_alternating_pattern() { ) .unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] fn test_list_filter_variable_sizes() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test filtering lists with highly variable sizes. let elements = buffer![0..100].into_array(); let offsets = buffer![0, 1, 2, 5, 10, 20, 35, 60, 100].into_array(); @@ -380,7 +381,7 @@ fn test_list_filter_variable_sizes() { ) .unwrap(); - assert_arrays_eq!(filtered, expected, &mut assertion_ctx); + assert_arrays_eq!(filtered, expected, &mut ctx); } #[test] @@ -581,7 +582,7 @@ fn create_list_of_lists_nullable(data: OptVec>>) -> ListArray #[test] #[expect(clippy::cognitive_complexity)] fn test_list_of_lists() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let data = vec![ Some(vec![Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]), Some(vec![Some(vec![Some(4), Some(5), Some(6)])]), @@ -612,19 +613,11 @@ fn test_list_of_lists() { // Check first inner list [1, 2]. let first_inner = first_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!( - first_inner, - PrimitiveArray::from_iter([1, 2]), - &mut assertion_ctx - ); + assert_arrays_eq!(first_inner, PrimitiveArray::from_iter([1, 2]), &mut ctx); // Check second inner list [3]. let second_inner = first_outer_list.list_elements_at(1).unwrap(); - assert_arrays_eq!( - second_inner, - PrimitiveArray::from_iter([3]), - &mut assertion_ctx - ); + assert_arrays_eq!(second_inner, PrimitiveArray::from_iter([3]), &mut ctx); // Check the second list of lists [[4, 5, 6]]. let second_outer = list_of_lists.list_elements_at(1).unwrap(); @@ -632,11 +625,7 @@ fn test_list_of_lists() { assert_eq!(second_outer_list.len(), 1); let inner = second_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!( - inner, - PrimitiveArray::from_iter([4, 5, 6]), - &mut assertion_ctx - ); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([4, 5, 6]), &mut ctx); // Check the third list of lists (empty). let third_outer = list_of_lists.list_elements_at(2).unwrap(); @@ -649,7 +638,7 @@ fn test_list_of_lists() { assert_eq!(fourth_outer_list.len(), 1); let inner = fourth_outer_list.list_elements_at(0).unwrap(); - assert_arrays_eq!(inner, PrimitiveArray::from_iter([7]), &mut assertion_ctx); + assert_arrays_eq!(inner, PrimitiveArray::from_iter([7]), &mut ctx); // Test scalar conversion. let scalar = list_of_lists diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index e820812351c..e85567a9ed9 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -348,7 +348,7 @@ mod tests { #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let list_view = create_basic_listview(); let list_array = list_from_list_view(list_view.clone(), &mut SESSION.create_execution_ctx())?; @@ -357,7 +357,7 @@ mod tests { assert_arrays_eq!( list_view.elements().clone(), list_array.elements().clone(), - &mut assertion_ctx + &mut ctx ); // ListArray offsets should have n+1 elements for n lists (add the final offset). @@ -366,11 +366,11 @@ mod tests { assert_arrays_eq!( list_view.offsets().clone(), list_array_offsets_without_last, - &mut assertion_ctx + &mut ctx ); // Verify data integrity. - assert_arrays_eq!(list_view, list_array, &mut assertion_ctx); + assert_arrays_eq!(list_view, list_array, &mut ctx); Ok(()) } @@ -424,7 +424,7 @@ mod tests { #[test] fn test_non_zero_copy_listview_to_list() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create ListViewArray with overlapping lists (not zero-copyable). let list_view = create_overlapping_listview(); let list_array = @@ -438,7 +438,7 @@ mod tests { } // The data should still be correct even though it required a rebuild. - assert_arrays_eq!(list_view, list_array, &mut assertion_ctx); + assert_arrays_eq!(list_view, list_array, &mut ctx); Ok(()) } @@ -558,7 +558,7 @@ mod tests { #[test] fn test_recursive_simple_listview() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let list_view = create_basic_listview(); let result = recursive_list_from_list_view( list_view.clone().into_array(), @@ -566,13 +566,13 @@ mod tests { )?; assert_eq!(result.len(), list_view.len()); - assert_arrays_eq!(list_view.into_array(), result, &mut assertion_ctx); + assert_arrays_eq!(list_view.into_array(), result, &mut ctx); Ok(()) } #[test] fn test_recursive_nested_listview() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let inner_elements = buffer![1i32, 2, 3].into_array(); let inner_offsets = buffer![0u32, 2].into_array(); let inner_sizes = buffer![2u32, 1].into_array(); @@ -604,13 +604,13 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(outer_listview.into_array(), result, &mut assertion_ctx); + assert_arrays_eq!(outer_listview.into_array(), result, &mut ctx); Ok(()) } #[test] fn test_recursive_struct_with_listview_fields() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let listview_field = create_basic_listview().into_array(); let primitive_field = buffer![10i32, 20, 30, 40].into_array(); @@ -627,13 +627,13 @@ mod tests { )?; assert_eq!(result.len(), 4); - assert_arrays_eq!(struct_array.into_array(), result, &mut assertion_ctx); + assert_arrays_eq!(struct_array.into_array(), result, &mut ctx); Ok(()) } #[test] fn test_recursive_fixed_size_list_with_listview_elements() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let lv1_elements = buffer![1i32, 2].into_array(); let lv1_offsets = buffer![0u32].into_array(); let lv1_sizes = buffer![2u32].into_array(); @@ -673,13 +673,13 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(fixed_list.into_array(), result, &mut assertion_ctx); + assert_arrays_eq!(fixed_list.into_array(), result, &mut ctx); Ok(()) } #[test] fn test_recursive_deep_nesting() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let innermost_elements = buffer![1i32, 2, 3].into_array(); let innermost_offsets = buffer![0u32, 2].into_array(); let innermost_sizes = buffer![2u32, 1].into_array(); @@ -718,7 +718,7 @@ mod tests { )?; assert_eq!(result.len(), 2); - assert_arrays_eq!(outer_listview.into_array(), result, &mut assertion_ctx); + assert_arrays_eq!(outer_listview.into_array(), result, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index 67a462ab910..7fb90bc5e04 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -10,6 +10,7 @@ use vortex_error::VortexResult; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::ListArray; @@ -26,7 +27,7 @@ use crate::validity::Validity; #[test] fn test_basic_listview_comprehensive() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Comprehensive test for basic ListView functionality including scalar_at. // Logical lists: [[1,2,3], [4,5], [6,7,8,9]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -52,7 +53,7 @@ fn test_basic_listview_comprehensive() { assert_arrays_eq!( listview.list_elements_at(0).unwrap(), PrimitiveArray::from_iter([1i32, 2, 3]), - &mut assertion_ctx + &mut ctx ); // Test scalar_at which returns entire lists as Scalar values. @@ -71,19 +72,19 @@ fn test_basic_listview_comprehensive() { assert_arrays_eq!( listview.list_elements_at(1).unwrap(), PrimitiveArray::from_iter([4i32, 5]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( listview.list_elements_at(2).unwrap(), PrimitiveArray::from_iter([6i32, 7, 8, 9]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_out_of_order_offsets() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // ListView-specific: Tests that offsets can be non-sequential and out-of-order. // Logical lists: [[7,8,9], [1,2,3], [4,5,6]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -98,14 +99,14 @@ fn test_out_of_order_offsets() { assert_arrays_eq!( listview.list_elements_at(0).unwrap(), PrimitiveArray::from_iter([7i32, 8, 9]), - &mut assertion_ctx + &mut ctx ); // Second list starts at offset 0: [1, 2, 3]. assert_arrays_eq!( listview.list_elements_at(1).unwrap(), PrimitiveArray::from_iter([1i32, 2, 3]), - &mut assertion_ctx + &mut ctx ); } @@ -167,7 +168,7 @@ fn test_from_list_array() -> VortexResult<()> { #[case::constant_offsets(false, true)] // Varying sizes, constant offsets #[case::both_constant(true, true)] // Both constant fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_offsets: bool) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Logical lists vary by case: // - constant_sizes: [[1,2,3], [4,5,6], [7,8,9]] (size 3 each, varying offsets) // - constant_offsets: [[1,2,3], [1,2], [1]] (all start at 0, varying sizes) @@ -200,11 +201,7 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o // All lists are identical [1, 2, 3] (overlapping). let expected = PrimitiveArray::from_iter([1i32, 2, 3]); for i in 0..3 { - assert_arrays_eq!( - listview.list_elements_at(i).unwrap(), - expected, - &mut assertion_ctx - ); + assert_arrays_eq!(listview.list_elements_at(i).unwrap(), expected, &mut ctx); } } else if const_sizes { // All lists have size 3, different offsets (no overlap). diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index 6796297bf1c..97b73d3591a 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -39,7 +39,7 @@ fn test_filter_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_filter_preserves_unreferenced_elements() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -63,7 +63,7 @@ fn test_filter_preserves_unreferenced_elements() { assert_arrays_eq!( result_list.elements(), PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), - &mut assertion_ctx + &mut ctx ); // Verify offsets are unchanged. @@ -73,7 +73,7 @@ fn test_filter_preserves_unreferenced_elements() { #[test] fn test_filter_with_gaps() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] @@ -97,7 +97,7 @@ fn test_filter_with_gaps() { assert_arrays_eq!( result_list.elements(), PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), - &mut assertion_ctx + &mut ctx ); // Verify offsets are unchanged. @@ -109,7 +109,7 @@ fn test_filter_with_gaps() { assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), PrimitiveArray::from_iter([7i32, 8, 9]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index 15716dcb883..8460fd956ac 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -18,6 +18,7 @@ use crate::LEGACY_SESSION; use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::ListView; @@ -84,7 +85,7 @@ fn test_slice_comprehensive() { #[test] fn test_slice_out_of_order() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // ListView-specific: Test slicing with out-of-order offsets. // Logical lists: [[70,80], [10,20,30], [40,50,60], [90], [30]] let elements = buffer![10i32, 20, 30, 40, 50, 60, 70, 80, 90].into_array(); @@ -125,17 +126,17 @@ fn test_slice_out_of_order() { assert_arrays_eq!( sliced_list.list_elements_at(0).unwrap(), PrimitiveArray::from_iter([10i32, 20, 30]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( sliced_list.list_elements_at(1).unwrap(), PrimitiveArray::from_iter([40i32, 50, 60]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( sliced_list.list_elements_at(2).unwrap(), PrimitiveArray::from_iter([90i32]), - &mut assertion_ctx + &mut ctx ); } @@ -394,7 +395,7 @@ fn test_cast_large_dataset() { #[test] fn test_zip_widens_false_element_nullability() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // [[1, 2], [3], [4]] let if_true = ListViewArray::new( buffer![1i32, 2, 3, 4].into_array(), @@ -434,13 +435,13 @@ fn test_zip_widens_false_element_nullability() -> VortexResult<()> { Validity::NonNullable, ) .into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } #[test] fn test_zip_widens_true_element_nullability() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // [[1, null], [3], [4]] let if_true = ListViewArray::new( PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(), @@ -480,7 +481,7 @@ fn test_zip_widens_true_element_nullability() -> VortexResult<()> { Validity::NonNullable, ) .into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 53c1eb68f9d..b33243b56ee 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -40,7 +40,7 @@ fn test_take_listview_conformance(#[case] listview: ListViewArray) { #[test] fn test_take_preserves_unreferenced_elements() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // ListView-specific: Test that take preserves the entire elements array // even when taking only a subset of lists. let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -62,7 +62,7 @@ fn test_take_preserves_unreferenced_elements() { assert_arrays_eq!( result_list.elements(), PrimitiveArray::from_iter([0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9]), - &mut assertion_ctx + &mut ctx ); // Verify offsets are preserved. @@ -72,7 +72,7 @@ fn test_take_preserves_unreferenced_elements() { #[test] fn test_take_with_gaps() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // ListView-specific: Test with gaps in elements array. // Elements with gaps (999 values are "gaps" between used ranges). let elements = buffer![1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12].into_array(); @@ -91,14 +91,14 @@ fn test_take_with_gaps() { assert_arrays_eq!( result_list.elements(), PrimitiveArray::from_iter([1i32, 2, 3, 999, 999, 999, 7, 8, 9, 999, 11, 12]), - &mut assertion_ctx + &mut ctx ); // Verify the lists still read correctly despite gaps. assert_arrays_eq!( result_list.list_elements_at(0).unwrap(), PrimitiveArray::from_iter([7i32, 8, 9]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index 316020a0761..7fc5f6f8dc4 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -12,6 +12,7 @@ use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -95,7 +96,7 @@ fn test_masked_child_with_validity() { #[test] fn test_masked_child_all_valid() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // When validity is AllValid, masked_child should invert to AllInvalid. let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); @@ -110,7 +111,7 @@ fn test_masked_child_all_valid() { assert_arrays_eq!( PrimitiveArray::from_option_iter([10i32, 20, 30].map(Some)), array, - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/arrays/patched/compute/take.rs b/vortex-array/src/arrays/patched/compute/take.rs index b985493ee83..993428808d4 100644 --- a/vortex-array/src/arrays/patched/compute/take.rs +++ b/vortex-array/src/arrays/patched/compute/take.rs @@ -138,6 +138,7 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -168,7 +169,7 @@ mod tests { #[test] fn test_take_basic() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -178,14 +179,14 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 10, 0, 30, 0]).into_array(); - assert_arrays_eq!(expected, result, &mut assertion_ctx); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } #[test] fn test_take_sliced() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = make_patched_array(&[0; 10], &[1, 3], &[100, 200], 2..10)?; let indices = buffer![0u32, 1, 2, 3, 7].into_array(); @@ -193,14 +194,14 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 200, 0, 0, 0]).into_array(); - assert_arrays_eq!(expected, result, &mut assertion_ctx); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } #[test] fn test_take_out_of_order() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -210,14 +211,14 @@ mod tests { let result = array.take(indices)?.to_canonical()?.into_array(); let expected = PrimitiveArray::from_iter([0u16, 30, 0, 10, 0]).into_array(); - assert_arrays_eq!(expected, result, &mut assertion_ctx); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } #[test] fn test_take_duplicates() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at index [2] with value [99] let array = make_patched_array(&[0; 5], &[2], &[99], 0..5)?; @@ -231,14 +232,14 @@ mod tests { let _canonical = result.to_canonical()?.into_primitive(); let expected = PrimitiveArray::from_iter([99u16, 99, 0, 99]).into_array(); - assert_arrays_eq!(expected, result, &mut assertion_ctx); + assert_arrays_eq!(expected, result, &mut ctx); Ok(()) } #[test] fn test_take_with_null_indices() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use crate::arrays::BoolArray; use crate::validity::Validity; @@ -276,7 +277,7 @@ mod tests { .into_array(), ), ); - assert_arrays_eq!(expected.into_array(), result, &mut assertion_ctx); + assert_arrays_eq!(expected.into_array(), result, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index c9d9eab81b4..d183552b66f 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -138,6 +138,7 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::validity::Validity; @@ -174,7 +175,7 @@ mod tests { #[test] fn patch_sliced() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid); let sliced = input.slice(2..8).unwrap(); #[expect(deprecated)] @@ -182,7 +183,7 @@ mod tests { assert_arrays_eq!( sliced_primitive, PrimitiveArray::new(buffer![2u32; 6], Validity::AllValid), - &mut assertion_ctx + &mut ctx ); } } diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 02e15fb08bc..a83cc293e4d 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -270,9 +270,11 @@ mod test { use vortex_error::VortexError; use vortex_mask::Mask; + use crate::ArrayRef; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; @@ -286,17 +288,13 @@ mod test { #[test] fn cast_u32_u8() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); // cast from u32 to u8 #[expect(deprecated)] let p = arr.cast(PType::U8.into()).unwrap().to_primitive(); - assert_arrays_eq!( - p, - PrimitiveArray::from_iter([0u8, 10, 200]), - &mut assertion_ctx - ); + assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); // to nullable @@ -309,7 +307,7 @@ mod test { assert_arrays_eq!( p, PrimitiveArray::new(buffer![0u8, 10, 200], Validity::AllValid), - &mut assertion_ctx + &mut ctx ); assert!(matches!(p.validity(), Ok(Validity::AllValid))); @@ -320,11 +318,7 @@ mod test { .cast(DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() .to_primitive(); - assert_arrays_eq!( - p, - PrimitiveArray::from_iter([0u8, 10, 200]), - &mut assertion_ctx - ); + assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); // to nullable u32 @@ -337,7 +331,7 @@ mod test { assert_arrays_eq!( p, PrimitiveArray::new(buffer![0u32, 10, 200], Validity::AllValid), - &mut assertion_ctx + &mut ctx ); assert!(matches!(p.validity(), Ok(Validity::AllValid))); @@ -348,24 +342,20 @@ mod test { .cast(DType::Primitive(PType::U8, Nullability::NonNullable)) .unwrap() .to_primitive(); - assert_arrays_eq!( - p, - PrimitiveArray::from_iter([0u8, 10, 200]), - &mut assertion_ctx - ); + assert_arrays_eq!(p, PrimitiveArray::from_iter([0u8, 10, 200]), &mut ctx); assert!(matches!(p.validity(), Ok(Validity::NonNullable))); } #[test] fn cast_u32_f32() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); #[expect(deprecated)] let u8arr = arr.cast(PType::F32.into()).unwrap().to_primitive(); assert_arrays_eq!( u8arr, PrimitiveArray::from_iter([0.0f32, 10., 200.]), - &mut assertion_ctx + &mut ctx ); } @@ -400,7 +390,7 @@ mod test { #[test] fn cast_with_invalid_nulls() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![-1i32, 0, 10], Validity::from_iter([false, true, true]), @@ -414,7 +404,7 @@ mod test { assert_arrays_eq!( p, PrimitiveArray::from_option_iter([None, Some(0u32), Some(10)]), - &mut assertion_ctx + &mut ctx ); assert_eq!( p.as_ref() @@ -430,7 +420,7 @@ mod test { /// buffer without allocation (pointer identity). #[test] fn cast_same_width_int_reinterprets_buffer() -> vortex_error::VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let src = PrimitiveArray::from_iter([0u32, 10, 100]); let src_ptr = src.as_slice::().as_ptr(); @@ -440,11 +430,7 @@ mod test { // Zero-copy: the data pointer should be identical. assert_eq!(src_ptr as usize, dst_ptr as usize); - assert_arrays_eq!( - dst, - PrimitiveArray::from_iter([0i32, 10, 100]), - &mut assertion_ctx - ); + assert_arrays_eq!(dst, PrimitiveArray::from_iter([0i32, 10, 100]), &mut ctx); Ok(()) } @@ -480,7 +466,7 @@ mod test { /// not prevent the cast from succeeding. #[test] fn cast_same_width_int_nullable_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // The null position holds u32::MAX which doesn't fit in i32, but it's // masked as invalid so the cast should still succeed via reinterpret. let arr = PrimitiveArray::new( @@ -495,14 +481,14 @@ mod test { assert_arrays_eq!( casted, PrimitiveArray::from_option_iter([None, Some(0i32), Some(42)]), - &mut assertion_ctx + &mut ctx ); Ok(()) } #[test] fn cast_u32_to_u8_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1000u32, 10u32, 42u32], Validity::from_iter([false, true, true]), @@ -515,7 +501,7 @@ mod test { assert_arrays_eq!( casted, PrimitiveArray::from_option_iter([None, Some(10u8), Some(42)]), - &mut assertion_ctx + &mut ctx ); Ok(()) } @@ -536,7 +522,7 @@ mod test { #[case(PrimitiveArray::from_option_iter([Some(1u8), None, Some(255), Some(0), None]).into_array())] #[case(PrimitiveArray::from_option_iter([Some(1i32), None, Some(-100), Some(0), None]).into_array())] #[case(buffer![42u32].into_array())] - fn test_cast_primitive_conformance(#[case] array: crate::ArrayRef) { + fn test_cast_primitive_conformance(#[case] array: ArrayRef) { test_cast_conformance(&array); } } diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index 3fc83be2e3f..c9376de2ee5 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -53,6 +53,7 @@ mod test { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::compute::fill_null::BoolArray; use crate::assert_arrays_eq; @@ -64,7 +65,7 @@ mod test { #[test] fn fill_null_leading_none() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([None, Some(8u8), None, Some(10), None]); #[expect(deprecated)] let p = arr @@ -75,7 +76,7 @@ mod test { assert_arrays_eq!( p, PrimitiveArray::from_iter([42u8, 8, 42, 10, 42]), - &mut assertion_ctx + &mut ctx ); assert!( p.as_ref() @@ -89,7 +90,7 @@ mod test { #[test] fn fill_null_all_none() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Option::::None, None, None, None, None]); #[expect(deprecated)] @@ -101,7 +102,7 @@ mod test { assert_arrays_eq!( p, PrimitiveArray::from_iter([255u8, 255, 255, 255, 255]), - &mut assertion_ctx + &mut ctx ); assert!( p.as_ref() @@ -115,7 +116,7 @@ mod test { #[test] fn fill_null_nullable_non_null() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![8u8, 10, 12, 14, 16], Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()), @@ -129,7 +130,7 @@ mod test { assert_arrays_eq!( p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16]), - &mut assertion_ctx + &mut ctx ); assert!( p.as_ref() @@ -143,14 +144,14 @@ mod test { #[test] fn fill_null_non_nullable() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = buffer![8u8, 10, 12, 14, 16].into_array(); #[expect(deprecated)] let p = arr.fill_null(Scalar::from(255u8)).unwrap().to_primitive(); assert_arrays_eq!( p, PrimitiveArray::from_iter([8u8, 10, 12, 14, 16]), - &mut assertion_ctx + &mut ctx ); assert!( p.as_ref() diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 5836d2b26f3..db37a820a37 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -223,6 +223,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -230,7 +231,7 @@ mod tests { #[test] fn take_null_index_skips_out_of_bounds_value() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter([10i32, 20, 30]); let indices = PrimitiveArray::new( buffer![1u64, 3], @@ -242,7 +243,7 @@ mod tests { assert_arrays_eq!( taken, PrimitiveArray::from_option_iter([Some(20i32), None]).into_array(), - &mut assertion_ctx + &mut ctx ); } } diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index a10bffc9fef..b5eba2cb1a0 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -225,8 +225,8 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; @@ -235,7 +235,8 @@ mod tests { #[test] fn test_nullable_primitive_serde_roundtrip() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); let array = PrimitiveArray::new( buffer![1i32, 2, 3, 4], Validity::from_iter([true, false, true, false]), @@ -243,11 +244,11 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let serialized = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); let mut concat = ByteBufferMut::empty(); @@ -256,14 +257,9 @@ mod tests { } let parts = SerializedArray::try_from(concat.freeze()).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array, &mut assertion_ctx); + assert_arrays_eq!(decoded, array, &mut ctx); } } diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index f3dd679bcc3..2e6b0240004 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -62,6 +62,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::ScalarFnArray; @@ -76,7 +77,7 @@ mod tests { #[test] fn test_scalar_fn_add() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = buffer![1i32, 2, 3].into_array(); let rhs = buffer![10i32, 20, 30].into_array(); @@ -90,7 +91,7 @@ mod tests { .execute::(&mut LEGACY_SESSION.create_execution_ctx())? .into_array(); let expected = buffer![11i32, 22, 33].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } @@ -131,7 +132,7 @@ mod tests { #[test] fn test_scalar_fn_mul() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = buffer![2i32, 3, 4].into_array(); let rhs = buffer![5i32, 6, 7].into_array(); @@ -143,14 +144,14 @@ mod tests { .execute::(&mut LEGACY_SESSION.create_execution_ctx())? .into_array(); let expected = buffer![10i32, 18, 28].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } #[test] fn test_scalar_fn_with_nullable() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let rhs = PrimitiveArray::new( buffer![10i32, 20, 30], @@ -170,14 +171,14 @@ mod tests { Validity::from_iter([true, false, true]), ) .into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } #[test] fn test_scalar_fn_comparison() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = buffer![1i32, 5, 3].into_array(); let rhs = buffer![2i32, 5, 1].into_array(); @@ -189,7 +190,7 @@ mod tests { .execute::(&mut LEGACY_SESSION.create_execution_ctx())? .into_array(); let expected = BoolArray::from_iter([false, true, false]).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index a4193a9502f..ba716bee024 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -189,23 +189,20 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::assert_arrays_eq; #[test] fn test_slice_slice() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Slice(1..4, Slice(2..8, base)) combines to Slice(3..6, base) let arr = PrimitiveArray::from_iter(0i32..10).into_array(); let inner_slice = SliceArray::new(arr, 2..8).into_array(); let slice = inner_slice.slice(1..4)?; - assert_arrays_eq!( - slice, - PrimitiveArray::from_iter([3i32, 4, 5]), - &mut assertion_ctx - ); + assert_arrays_eq!(slice, PrimitiveArray::from_iter([3i32, 4, 5]), &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index ed533948490..36a355bd3de 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -22,6 +22,7 @@ mod tests { use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -40,7 +41,7 @@ mod tests { #[test] fn take_empty_struct() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let indices = PrimitiveArray::from_option_iter([Some(1), None]); @@ -54,7 +55,7 @@ mod tests { 2, Validity::from_iter([true, false]) ), - &mut assertion_ctx + &mut ctx ); } @@ -94,7 +95,7 @@ mod tests { #[test] fn take_field_struct() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let struct_arr = StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).into_array())]) .unwrap(); @@ -107,7 +108,7 @@ mod tests { Validity::from_iter([true, false]) ) .unwrap(), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index dd2ee10b400..01981803239 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -165,7 +165,7 @@ mod tests { #[test] fn test_struct_cast_field_reorder() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Source: {a, b}, Target: {c, b, a} - reordered + new null field let source = StructArray::try_new( FieldNames::from(["a", "b"]), @@ -198,17 +198,17 @@ mod tests { assert_arrays_eq!( result.unmasked_field_by_name("a").unwrap(), VarBinViewArray::from_iter_nullable_str([Some("A")]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( result.unmasked_field_by_name("b").unwrap(), VarBinViewArray::from_iter_nullable_str([Some("B")]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( result.unmasked_field_by_name("c").unwrap(), ConstantArray::new(Scalar::null(utf8_null), 1), - &mut assertion_ctx + &mut ctx ); } @@ -302,7 +302,7 @@ mod tests { #[test] fn cast_struct_drop_field() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Casting to a struct with a subset of fields should succeed. let source = StructArray::try_new( FieldNames::from(["a", "b", "c"]), @@ -337,18 +337,18 @@ mod tests { assert_arrays_eq!( result.unmasked_field_by_name("a").unwrap(), buffer![1i32, 2, 3].into_array(), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( result.unmasked_field_by_name("c").unwrap(), buffer![100u8, 200, 255].into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn cast_struct_field_type_widening() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Casting struct fields to wider types (i32 -> i64). let source = StructArray::try_new( FieldNames::from(["val"]), @@ -379,7 +379,7 @@ mod tests { assert_arrays_eq!( result.unmasked_field_by_name("val").unwrap(), buffer![1i64, 2, 3].into_array(), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/arrays/struct_/tests.rs b/vortex-array/src/arrays/struct_/tests.rs index 8272dd439c8..62b0b3c8701 100644 --- a/vortex-array/src/arrays/struct_/tests.rs +++ b/vortex-array/src/arrays/struct_/tests.rs @@ -8,6 +8,7 @@ use crate::Canonical; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -24,7 +25,7 @@ use crate::validity::Validity; #[test] fn test_project() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = VarBinArray::from_vec( vec!["a", "b", "c", "d", "e"], @@ -54,20 +55,20 @@ fn test_project() { assert_arrays_eq!( bools, BoolArray::from_iter([true, true, true, false, false]), - &mut assertion_ctx + &mut ctx ); let prims = struct_b.unmasked_field(1); assert_arrays_eq!( prims, PrimitiveArray::from_iter([0i64, 1, 2, 3, 4]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_remove_column() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = PrimitiveArray::new(buffer![4u64, 5, 6, 7, 8], Validity::NonNullable); @@ -87,7 +88,7 @@ fn test_remove_column() { assert_arrays_eq!( removed, PrimitiveArray::from_iter([0i64, 1, 2, 3, 4]), - &mut assertion_ctx + &mut ctx ); assert_eq!(data.names(), &["ys"]); @@ -100,7 +101,7 @@ fn test_remove_column() { assert_arrays_eq!( data.unmasked_field(0), PrimitiveArray::from_iter([4u64, 5, 6, 7, 8]), - &mut assertion_ctx + &mut ctx ); let empty = data.remove_column("non_existent"); @@ -113,7 +114,7 @@ fn test_remove_column() { #[test] fn test_duplicate_field_names() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test that StructArray allows duplicate field names and returns the first match let field1 = buffer![1i32, 2, 3].into_array(); let field2 = buffer![10i32, 20, 30].into_array(); @@ -133,23 +134,19 @@ fn test_duplicate_field_names() { assert_arrays_eq!( first_value_field, PrimitiveArray::from_iter([1i32, 2, 3]), - &mut assertion_ctx + &mut ctx ); // Verify field_by_name_opt also returns the first match let opt_field = struct_array.unmasked_field_by_name_opt("value").unwrap(); - assert_arrays_eq!( - opt_field, - PrimitiveArray::from_iter([1i32, 2, 3]), - &mut assertion_ctx - ); + assert_arrays_eq!(opt_field, PrimitiveArray::from_iter([1i32, 2, 3]), &mut ctx); // Verify the third field (second "value") can be accessed by index let third_field = struct_array.unmasked_field(2); assert_arrays_eq!( third_field, PrimitiveArray::from_iter([100i32, 200, 300]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index ab75006b3bb..32b795009f1 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -255,6 +255,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; @@ -290,7 +291,7 @@ mod tests { /// [`CompareKernel`]: super::CompareKernel #[test] fn varbin_i64_offsets_compare_constant() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -306,12 +307,12 @@ mod tests { .unwrap(); let expected = BoolArray::from_iter([true, false, true]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn varbin_i64_offsets_compare_constant_binary() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -331,6 +332,6 @@ mod tests { .unwrap(); let expected = BoolArray::from_iter([true, false, true]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } } diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index cef5e6a41c2..ab969cf3def 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -225,8 +225,8 @@ mod test { use vortex_buffer::buffer; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::VarBinArray; use crate::arrays::varbin::compute::filter::filter_select_var_bin_by_index; @@ -240,7 +240,7 @@ mod test { #[test] fn filter_var_bin_test() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -250,24 +250,14 @@ mod test { DType::Utf8(NonNullable), ); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_index( - arr, - &[0, 2], - 2, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = filter_select_var_bin_by_index(arr, &[0, 2], 2, &mut ctx).unwrap(); - assert_arrays_eq!( - buf, - VarBinArray::from(vec!["hello", "filter"]), - &mut assertion_ctx - ); + assert_arrays_eq!(buf, VarBinArray::from(vec!["hello", "filter"]), &mut ctx); } #[test] fn filter_var_bin_slice_test() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -280,24 +270,19 @@ mod test { ); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_slice( - arr, - &[(0, 1), (2, 3), (4, 5)], - 3, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = + filter_select_var_bin_by_slice(arr, &[(0, 1), (2, 3), (4, 5)], 3, &mut ctx).unwrap(); assert_arrays_eq!( buf, VarBinArray::from(vec!["hello", "filter", "filter3"]), - &mut assertion_ctx + &mut ctx ); } #[test] fn filter_var_bin_slice_null() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let bytes = [ b"one".as_slice(), b"two".as_slice(), @@ -317,13 +302,7 @@ mod test { let arr = VarBinArray::try_new(offsets, bytes, DType::Utf8(Nullable), validity).unwrap(); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_slice( - arr, - &[(0, 3), (4, 6)], - 5, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = filter_select_var_bin_by_slice(arr, &[(0, 3), (4, 6)], 5, &mut ctx).unwrap(); assert_arrays_eq!( buf, @@ -334,13 +313,13 @@ mod test { Some("five"), Some("six") ]), - &mut assertion_ctx + &mut ctx ); } #[test] fn filter_varbin_nulls() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let bytes = [b"".as_slice(), b"two".as_slice(), b"two".as_slice()] .into_iter() .flat_map(|x| x.iter().cloned()) @@ -351,24 +330,14 @@ mod test { let arr = VarBinArray::try_new(offsets, bytes, DType::Utf8(Nullable), validity).unwrap(); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_slice( - arr, - &[(0, 1), (2, 3)], - 2, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = filter_select_var_bin_by_slice(arr, &[(0, 1), (2, 3)], 2, &mut ctx).unwrap(); - assert_arrays_eq!( - buf, - VarBinArray::from(vec![None, Some("two")]), - &mut assertion_ctx - ); + assert_arrays_eq!(buf, VarBinArray::from(vec![None, Some("two")]), &mut ctx); } #[test] fn filter_varbin_all_null() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let offsets = buffer![0, 0, 0, 0].into_array(); let validity = Validity::Array(BoolArray::from_iter([false, false, false]).into_array()); let arr = VarBinArray::try_new( @@ -380,19 +349,9 @@ mod test { .unwrap(); let arr = arr.as_view(); - let buf = filter_select_var_bin_by_slice( - arr, - &[(0, 1), (2, 3)], - 2, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let buf = filter_select_var_bin_by_slice(arr, &[(0, 1), (2, 3)], 2, &mut ctx).unwrap(); - assert_arrays_eq!( - buf, - VarBinArray::from(vec![None::<&str>, None]), - &mut assertion_ctx - ); + assert_arrays_eq!(buf, VarBinArray::from(vec![None::<&str>, None]), &mut ctx); } #[test] diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 97fc205cfa8..74e0bd6af9d 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -251,12 +251,15 @@ fn take_nullable(); + let mut ctx = array_session().create_execution_ctx(); + let scream = iter::once("a").cycle().take(128).collect::(); let bytes = ByteBuffer::copy_from(scream.as_bytes()); let offsets = buffer![0u8, 128u8].into_array(); @@ -324,6 +327,6 @@ mod tests { [Some(scream.clone()), Some(scream.clone()), Some(scream)], DType::Utf8(Nullability::NonNullable), ); - assert_arrays_eq!(expected, taken, &mut assertion_ctx); + assert_arrays_eq!(expected, taken, &mut ctx); } } diff --git a/vortex-array/src/arrays/varbin/tests.rs b/vortex-array/src/arrays/varbin/tests.rs index 59b4d8354d1..e2eb79fff54 100644 --- a/vortex-array/src/arrays/varbin/tests.rs +++ b/vortex-array/src/arrays/varbin/tests.rs @@ -9,6 +9,7 @@ use vortex_buffer::buffer; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; @@ -33,21 +34,21 @@ fn binary_array() -> ArrayRef { #[rstest] pub fn test_scalar_at(binary_array: ArrayRef) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); assert_arrays_eq!( binary_array, VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), - &mut assertion_ctx + &mut ctx ); } #[rstest] pub fn slice_array(binary_array: ArrayRef) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let binary_arr = binary_array.slice(1..2).unwrap(); assert_arrays_eq!( binary_arr, VarBinViewArray::from_iter_str(["hello world this is a long string"]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index 12f25183e12..9a2256fe05a 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -49,6 +49,7 @@ mod tests { use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::builder::VarBinBuilder; @@ -96,7 +97,7 @@ mod tests { #[case(DType::Utf8(Nullability::NonNullable))] #[case(DType::Binary(Nullability::NonNullable))] fn test_canonical_varbin_unsliced(#[case] dtype: DType) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let varbin = VarBinArray::from_iter_nonnull(["foo", "bar", "baz"], dtype.clone()); #[expect(deprecated)] let canonical = varbin.as_array().to_varbinview(); @@ -104,7 +105,7 @@ mod tests { DType::Utf8(_) => VarBinViewArray::from_iter_str(["foo", "bar", "baz"]), _ => VarBinViewArray::from_iter_bin(["foo", "bar", "baz"]), }; - assert_arrays_eq!(canonical, expected, &mut assertion_ctx); + assert_arrays_eq!(canonical, expected, &mut ctx); } // Empty array: offsets has exactly one element; no elements to canonicalize. diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index 894f604f567..2f6d459113f 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -195,6 +195,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; @@ -202,7 +203,7 @@ mod tests { use crate::dtype::Nullability; #[test] fn test_optimize_compacts_buffers() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create a VarBinViewArray with some long strings that will create multiple buffers let original = VarBinViewArray::from_iter_nullable_str([ Some("short"), @@ -237,13 +238,13 @@ mod tests { assert_arrays_eq!( optimized_array, >::from_iter([Some("short"), Some("tiny")]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_optimize_with_long_strings() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create strings that are definitely longer than 12 bytes let long_string_1 = "this is definitely a very long string that exceeds the inline limit"; let long_string_2 = "another extremely long string that also needs external buffer storage"; @@ -274,13 +275,13 @@ mod tests { assert_arrays_eq!( optimized_array, VarBinArray::from(vec![long_string_1, long_string_3]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_optimize_no_buffers() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create an array with only short strings (all inlined) let original = VarBinViewArray::from_iter_str(["a", "bb", "ccc", "dddd"]); @@ -292,12 +293,12 @@ mod tests { assert_eq!(optimized_array.data_buffers().len(), 0); - assert_arrays_eq!(optimized_array, original, &mut assertion_ctx); + assert_arrays_eq!(optimized_array, original, &mut ctx); } #[test] fn test_optimize_single_buffer() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create an array that naturally has only one buffer let str1 = "this is a long string that goes into a buffer"; let str2 = "another long string in the same buffer"; @@ -312,12 +313,12 @@ mod tests { assert_eq!(optimized_array.data_buffers().len(), 1); - assert_arrays_eq!(optimized_array, original, &mut assertion_ctx); + assert_arrays_eq!(optimized_array, original, &mut ctx); } #[test] fn test_selective_compaction_with_threshold_zero() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // threshold=0 should keep all buffers (no compaction) let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -340,12 +341,12 @@ mod tests { assert_eq!(compacted.data_buffers().len(), taken.data_buffers().len()); // Verify correctness - assert_arrays_eq!(compacted, taken, &mut assertion_ctx); + assert_arrays_eq!(compacted, taken, &mut ctx); } #[test] fn test_selective_compaction_with_high_threshold() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // threshold=1.0 should compact any buffer with waste let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -369,12 +370,12 @@ mod tests { assert!(compacted.data_buffers().len() <= original_buffers); // Verify correctness - assert_arrays_eq!(compacted, taken, &mut assertion_ctx); + assert_arrays_eq!(compacted, taken, &mut ctx); } #[test] fn test_selective_compaction_preserves_well_utilized_buffers() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create an array with multiple strings in one buffer (well-utilized) let str1 = "first long string that needs external buffer storage"; let str2 = "second long string also in buffer"; @@ -392,12 +393,12 @@ mod tests { assert_eq!(compacted.data_buffers().len(), 1); // Verify all data is correct - assert_arrays_eq!(compacted, original, &mut assertion_ctx); + assert_arrays_eq!(compacted, original, &mut ctx); } #[test] fn test_selective_compaction_with_mixed_utilization() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create array with some long strings let strings: Vec = (0..10) .map(|i| { @@ -424,12 +425,12 @@ mod tests { [0, 2, 4, 6, 8].map(|i| Some(strings[i].as_str())), DType::Utf8(Nullability::NonNullable), ); - assert_arrays_eq!(expected, compacted, &mut assertion_ctx); + assert_arrays_eq!(expected, compacted, &mut ctx); } #[test] fn test_slice_strategy_with_contiguous_range() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create array with strings that will be in one buffer let strings: Vec = (0..20) .map(|i| format!("this is a long string number {} for slice test", i)) @@ -458,7 +459,7 @@ mod tests { ); // Verify correctness - assert_arrays_eq!(&compacted, taken, &mut assertion_ctx); + assert_arrays_eq!(&compacted, taken, &mut ctx); // Verify that if there was only one buffer, the compacted version also has one // (it was sliced, not rewritten into multiple buffers) diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index b557922951c..386625adcf1 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -248,8 +248,8 @@ mod tests { use super::*; use crate::ArrayContext; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -266,12 +266,13 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - let ctx = ArrayContext::empty(); - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let session = array_session(); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); let serialized = array .clone() .into_array() - .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default()) + .serialize(&array_ctx, &session, &SerializeOptions::default()) .unwrap(); let mut concat = ByteBufferMut::empty(); @@ -280,14 +281,9 @@ mod tests { } let parts = SerializedArray::try_from(concat.freeze()).unwrap(); let decoded = parts - .decode( - &dtype, - len, - &ReadContext::new(ctx.to_ids()), - &LEGACY_SESSION, - ) + .decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session) .unwrap(); - assert_arrays_eq!(decoded, array, &mut assertion_ctx); + assert_arrays_eq!(decoded, array, &mut ctx); } } diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index e1f8831ab6b..d159570f146 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -90,6 +90,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -379,7 +380,7 @@ mod tests { #[test] fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([( "a", @@ -400,14 +401,14 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), - &mut assertion_ctx + &mut ctx ); Ok(()) } #[test] fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([ ( @@ -434,7 +435,7 @@ mod tests { assert_arrays_eq!( value_result, PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), - &mut assertion_ctx + &mut ctx ); let typed_value_expr = variant_get( @@ -449,7 +450,7 @@ mod tests { assert_arrays_eq!( typed_value_result, PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]), - &mut assertion_ctx + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index e0411e69166..e7a64089940 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -163,6 +163,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; @@ -221,7 +222,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10); // Test appending true value. @@ -238,7 +239,7 @@ mod tests { let array = builder.finish_into_bool(); let expected = BoolArray::from_iter([Some(true), Some(false), None]); - assert_arrays_eq!(&array, &expected, &mut assertion_ctx); + assert_arrays_eq!(&array, &expected, &mut ctx); // Test wrong dtype error. let mut builder = BoolBuilder::with_capacity(Nullability::NonNullable, 10); diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index f6986960f97..2a6de9661e8 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -309,9 +309,9 @@ impl Default for DecimalBuffer { #[cfg(test)] mod tests { - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; @@ -345,7 +345,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use crate::scalar::Scalar; // Simply test that the builder accepts its own finish output via scalar. @@ -359,7 +359,7 @@ mod tests { [Some(1234i64), Some(5678), None], DecimalDType::new(10, 2), ); - assert_arrays_eq!(&array, &expected, &mut assertion_ctx); + assert_arrays_eq!(&array, &expected, &mut ctx); // Test by taking a scalar from the array and appending it to a new builder. let mut builder2 = DecimalBuilder::new::(DecimalDType::new(10, 2), true.into()); @@ -371,7 +371,7 @@ mod tests { } let array2 = builder2.finish(); - assert_arrays_eq!(&array2, &array, &mut assertion_ctx); + assert_arrays_eq!(&array2, &array, &mut ctx); // Test wrong dtype error. let mut builder = DecimalBuilder::new::(DecimalDType::new(10, 2), false.into()); diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index f6cb5f2621e..5e6d0be67ba 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -173,20 +173,20 @@ mod test { #[test] fn encode_primitive() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = buffer![1, 1, 3, 3, 3].into_array(); let dict = dict_encode(&arr, &mut SESSION.create_execution_ctx()).unwrap(); let expected_codes = buffer![0u8, 0, 1, 1, 1].into_array(); - assert_arrays_eq!(dict.codes(), expected_codes, &mut assertion_ctx); + assert_arrays_eq!(dict.codes(), expected_codes, &mut ctx); let expected_values = buffer![1i32, 3].into_array(); - assert_arrays_eq!(dict.values(), expected_values, &mut assertion_ctx); + assert_arrays_eq!(dict.values(), expected_values, &mut ctx); } #[test] fn encode_primitive_nulls() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([ Some(1), Some(1), @@ -200,10 +200,10 @@ mod test { let dict = dict_encode(&arr.into_array(), &mut SESSION.create_execution_ctx()).unwrap(); let expected_codes = buffer![0u8, 0, 1, 2, 2, 1, 2, 1].into_array(); - assert_arrays_eq!(dict.codes(), expected_codes, &mut assertion_ctx); + assert_arrays_eq!(dict.codes(), expected_codes, &mut ctx); let expected_values = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(); - assert_arrays_eq!(dict.values(), expected_values, &mut assertion_ctx); + assert_arrays_eq!(dict.values(), expected_values, &mut ctx); } } diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index 0888e899d23..61d5fe8d7b0 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -128,6 +128,7 @@ impl ArrayBuilder for ExtensionBuilder { mod tests { use super::*; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; @@ -138,7 +139,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); let mut builder = ExtensionBuilder::new(ext_dtype.clone()); @@ -167,7 +168,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(42i32), Some(84), None]).into_array(), ); - assert_arrays_eq!(&array, &expected, &mut assertion_ctx); + assert_arrays_eq!(&array, &expected, &mut ctx); assert_eq!(array.len(), 3); // Test wrong dtype error. diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index a639594fc54..5f54c6f0256 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -441,6 +441,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; @@ -529,7 +530,7 @@ mod tests { #[test] fn test_different_offset_size_types() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test u32 offsets with u8 sizes. let dtype: Arc = Arc::new(I32.into()); let mut builder = @@ -564,14 +565,14 @@ mod tests { assert_arrays_eq!( listview.list_elements_at(0).unwrap(), PrimitiveArray::from_iter([1i32, 2]), - &mut assertion_ctx + &mut ctx ); // Verify second list: [3, 4, 5]. assert_arrays_eq!( listview.list_elements_at(1).unwrap(), PrimitiveArray::from_iter([3i32, 4, 5]), - &mut assertion_ctx + &mut ctx ); // Test u64 offsets with u16 sizes. @@ -595,7 +596,7 @@ mod tests { assert_arrays_eq!( listview2.list_elements_at(i as usize).unwrap(), PrimitiveArray::from_iter([i * 10]), - &mut assertion_ctx + &mut ctx ); } } diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 12022925b75..77dfb4342bc 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -377,6 +377,7 @@ mod tests { use super::*; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; /// REGRESSION TEST: This test verifies that multiple sequential ranges have correct offsets. @@ -385,7 +386,7 @@ mod tests { /// buffer. #[test] fn test_multiple_uninit_ranges_correct_offsets() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); // First range. @@ -418,7 +419,7 @@ mod tests { assert_arrays_eq!( array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]), - &mut assertion_ctx + &mut ctx ); } @@ -496,7 +497,7 @@ mod tests { /// This verifies the new simplified API without the redundant `len` parameter. #[test] fn test_copy_from_slice_with_offsets() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); let mut range = builder.uninit_range(6); @@ -514,7 +515,7 @@ mod tests { assert_arrays_eq!( array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 1bddfbca369..53c521b4b07 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -210,10 +210,10 @@ impl ArrayBuilder for StructBuilder { #[cfg(test)] mod tests { - use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; @@ -267,7 +267,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use crate::scalar::Scalar; let dtype = DType::Struct( @@ -330,7 +330,7 @@ mod tests { Validity::from_iter([true, true, false]), ) .unwrap(); - assert_arrays_eq!(&array, &expected, &mut assertion_ctx); + assert_arrays_eq!(&array, &expected, &mut ctx); // Test wrong dtype error. let struct_fields = match &dtype { diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index b5ff523fbc3..f3ad1ffd355 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -875,6 +875,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; @@ -884,7 +885,7 @@ mod tests { #[test] fn test_utf8_builder() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); builder.append_value("Hello"); @@ -907,12 +908,12 @@ mod tests { Some(""), Some("test"), ]); - assert_arrays_eq!(actual, expected, &mut assertion_ctx); + assert_arrays_eq!(actual, expected, &mut ctx); } #[test] fn test_utf8_builder_with_extend() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = { let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); @@ -936,11 +937,7 @@ mod tests { None, Some("Hello3"), ]); - assert_arrays_eq!( - actual.into_array(), - expected.into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(actual.into_array(), expected.into_array(), &mut ctx); } #[test] @@ -992,7 +989,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use crate::scalar::Scalar; // Test with Utf8 builder. @@ -1014,7 +1011,7 @@ mod tests { let array = utf8_builder.finish(); let expected = >::from_iter([Some("hello"), Some("world"), None]); - assert_arrays_eq!(&array, &expected, &mut assertion_ctx); + assert_arrays_eq!(&array, &expected, &mut ctx); // Test with Binary builder. let mut binary_builder = @@ -1029,7 +1026,7 @@ mod tests { let binary_array = binary_builder.finish(); let expected = >::from_iter([Some(vec![1u8, 2, 3]), None]); - assert_arrays_eq!(&binary_array, &expected, &mut assertion_ctx); + assert_arrays_eq!(&binary_array, &expected, &mut ctx); // Test wrong dtype error. let mut builder = diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index 97d86481963..dc87726e095 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -8,6 +8,7 @@ use crate::ArrayRef; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; +use crate::array_session; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -62,13 +63,13 @@ pub fn create_runs_pattern(len: usize, run_length: usize) -> Vec { /// Tests that filtering with an all-true mask returns all elements unchanged fn test_all_filter(array: &ArrayRef) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let len = array.len(); let mask = Mask::new_true(len); let filtered = array .filter(mask) .vortex_expect("filter should succeed in conformance test"); - assert_arrays_eq!(filtered, array, &mut assertion_ctx); + assert_arrays_eq!(filtered, array, &mut ctx); } /// Tests that filtering with an all-false mask returns an empty array with the same dtype diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 237da17f861..a9952a7e99b 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -1210,10 +1210,10 @@ mod test { use vortex_mask::Mask; use crate::IntoArray; - use crate::LEGACY_SESSION; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::patches::Patches; use crate::patches::PrimitiveArray; @@ -1222,7 +1222,7 @@ mod test { #[test] fn test_filter() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 100, 0, @@ -1233,28 +1233,25 @@ mod test { .unwrap(); let filtered = patches - .filter( - &Mask::from_indices(100, vec![10, 20, 30]), - &mut LEGACY_SESSION.create_execution_ctx(), - ) + .filter(&Mask::from_indices(100, vec![10, 20, 30]), &mut ctx) .unwrap() .unwrap(); assert_arrays_eq!( filtered.indices(), PrimitiveArray::from_iter([0u64, 1]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( filtered.values(), PrimitiveArray::from_iter([100i32, 200]), - &mut assertion_ctx + &mut ctx ); } #[test] fn take_with_nulls() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1268,7 +1265,7 @@ mod test { .take( &PrimitiveArray::new(buffer![9, 0], Validity::from_iter(vec![true, false])) .into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1280,22 +1277,19 @@ mod test { assert_arrays_eq!( primitive_values, PrimitiveArray::from_option_iter([Some(44i32)]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_indices, PrimitiveArray::from_iter([0u64]), - &mut assertion_ctx + &mut ctx ); assert_eq!( primitive_values .as_ref() .validity() .unwrap() - .execute_mask( - primitive_values.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx() - ) + .execute_mask(primitive_values.as_ref().len(), &mut ctx) .unwrap(), Mask::from_iter(vec![true]) ); @@ -1303,7 +1297,7 @@ mod test { #[test] fn take_search_with_nulls_chunked() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1317,7 +1311,7 @@ mod test { .take_search( PrimitiveArray::new(buffer![9, 0], Validity::from_iter([true, false])), true, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1328,12 +1322,12 @@ mod test { assert_arrays_eq!( primitive_values, PrimitiveArray::from_option_iter([Some(44i32), None]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( taken.indices(), PrimitiveArray::from_iter([0u64, 1]), - &mut assertion_ctx + &mut ctx ); assert_eq!( @@ -1341,10 +1335,7 @@ mod test { .as_ref() .validity() .unwrap() - .execute_mask( - primitive_values.as_ref().len(), - &mut LEGACY_SESSION.create_execution_ctx() - ) + .execute_mask(primitive_values.as_ref().len(), &mut ctx) .unwrap(), Mask::from_iter([true, false]) ); @@ -1352,7 +1343,7 @@ mod test { #[test] fn take_search_chunked_multiple_chunks() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 2048, 0, @@ -1366,7 +1357,7 @@ mod test { .take_search( PrimitiveArray::new(buffer![500, 1200, 999], Validity::AllValid), true, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1375,12 +1366,13 @@ mod test { assert_arrays_eq!( taken.values(), PrimitiveArray::from_option_iter([Some(20i32), Some(30)]), - &mut assertion_ctx + &mut ctx ); } #[test] fn take_search_chunked_indices_with_no_patches() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1394,7 +1386,7 @@ mod test { .take_search( PrimitiveArray::new(buffer![3, 4, 5], Validity::AllValid), true, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap(); @@ -1403,7 +1395,7 @@ mod test { #[test] fn take_search_chunked_interleaved() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 30, 0, @@ -1417,7 +1409,7 @@ mod test { .take_search( PrimitiveArray::new(buffer![10, 15, 20, 99], Validity::AllValid), true, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1426,12 +1418,13 @@ mod test { assert_arrays_eq!( taken.values(), PrimitiveArray::from_option_iter([Some(200i32), Some(300)]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_take_search_multiple_chunk_offsets() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 1500, 0, @@ -1445,7 +1438,7 @@ mod test { .take_search( PrimitiveArray::new(BufferMut::from_iter(0..1500u64), Validity::AllValid), false, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .unwrap() .unwrap(); @@ -1455,7 +1448,7 @@ mod test { #[test] fn test_slice() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1466,13 +1459,13 @@ mod test { assert_arrays_eq!( sliced.values(), PrimitiveArray::from_iter([13531u32]), - &mut assertion_ctx + &mut ctx ); } #[test] fn doubly_sliced() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1483,19 +1476,20 @@ mod test { assert_arrays_eq!( sliced.values(), PrimitiveArray::from_iter([13531u32]), - &mut assertion_ctx + &mut ctx ); let doubly_sliced = sliced.slice(35..36).unwrap().unwrap(); assert_arrays_eq!( doubly_sliced.values(), PrimitiveArray::from_iter([13531u32]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_mask_all_true() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1506,15 +1500,13 @@ mod test { .unwrap(); let mask = Mask::new_true(10); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap(); assert!(masked.is_none()); } #[test] fn test_mask_all_false() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1525,47 +1517,29 @@ mod test { .unwrap(); let mask = Mask::new_false(10); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); // No patch values should be masked assert_arrays_eq!( masked.values(), PrimitiveArray::from_iter([100i32, 200, 300]), - &mut assertion_ctx - ); - assert!( - masked - .values() - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ); - assert!( - masked - .values() - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ); - assert!( - masked - .values() - .is_valid(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + &mut ctx ); + assert!(masked.values().is_valid(0, &mut ctx).unwrap()); + assert!(masked.values().is_valid(1, &mut ctx).unwrap()); + assert!(masked.values().is_valid(2, &mut ctx).unwrap()); // Indices should remain unchanged assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([2u64, 5, 8]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_mask_partial() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1579,30 +1553,27 @@ mod test { let mask = Mask::from_iter([ false, false, true, false, false, false, false, false, true, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); // Only the patch at index 5 should remain assert_eq!(masked.values().len(), 1); assert_arrays_eq!( masked.values(), PrimitiveArray::from_iter([200i32]), - &mut assertion_ctx + &mut ctx ); // Only index 5 should remain assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([5u64]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_mask_with_offset() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1617,27 +1588,24 @@ mod test { false, false, true, false, false, false, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); assert_eq!(masked.array_len(), 10); assert_eq!(masked.offset(), 5); assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([10u64, 13]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( masked.values(), PrimitiveArray::from_iter([200i32, 300]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_mask_nullable_values() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1651,46 +1619,30 @@ mod test { let mask = Mask::from_iter([ false, false, true, false, false, false, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); // Patches at indices 5 and 8 should remain assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([5u64, 8]), - &mut assertion_ctx + &mut ctx ); // Values should be the null and 300 #[expect(deprecated)] let masked_values = masked.values().to_primitive(); assert_eq!(masked_values.len(), 2); - assert!( - !masked_values - .is_valid(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ); // the null value at index 5 - assert!( - masked_values - .is_valid(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ); // the 300 value at index 8 + assert!(!masked_values.is_valid(0, &mut ctx).unwrap()); // the null value at index 5 + assert!(masked_values.is_valid(1, &mut ctx).unwrap()); // the 300 value at index 8 assert_eq!( - i32::try_from( - &masked_values - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ) - .unwrap(), + i32::try_from(&masked_values.execute_scalar(1, &mut ctx).unwrap()).unwrap(), 300i32 ); } #[test] fn test_filter_keep_all() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1702,25 +1654,23 @@ mod test { // Keep all indices (mask with indices 0-9) let mask = Mask::from_indices(10, 0..10); - let filtered = patches - .filter(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let filtered = patches.filter(&mask, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( filtered.indices(), PrimitiveArray::from_iter([2u64, 5, 8]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( filtered.values(), PrimitiveArray::from_iter([100i32, 200, 300]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_filter_none() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1732,15 +1682,13 @@ mod test { // Filter out all (empty mask means keep nothing) let mask = Mask::from_indices(10, vec![]); - let filtered = patches - .filter(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let filtered = patches.filter(&mask, &mut ctx).unwrap(); assert!(filtered.is_none()); } #[test] fn test_filter_with_indices() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1752,26 +1700,23 @@ mod test { // Keep indices 2, 5, 9 (so patches at 2 and 5 remain) let mask = Mask::from_indices(10, vec![2, 5, 9]); - let filtered = patches - .filter(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let filtered = patches.filter(&mask, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( filtered.indices(), PrimitiveArray::from_iter([0u64, 1]), - &mut assertion_ctx + &mut ctx ); // Adjusted indices assert_arrays_eq!( filtered.values(), PrimitiveArray::from_iter([100i32, 200]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_slice_full_range() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1786,18 +1731,18 @@ mod test { assert_arrays_eq!( sliced.indices(), PrimitiveArray::from_iter([2u64, 5, 8]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( sliced.values(), PrimitiveArray::from_iter([100i32, 200, 300]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_slice_partial() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1813,12 +1758,12 @@ mod test { assert_arrays_eq!( sliced.indices(), PrimitiveArray::from_iter([5u64]), - &mut assertion_ctx + &mut ctx ); // Index stays the same assert_arrays_eq!( sliced.values(), PrimitiveArray::from_iter([200i32]), - &mut assertion_ctx + &mut ctx ); assert_eq!(sliced.array_len(), 5); // 8 - 3 = 5 assert_eq!(sliced.offset(), 3); // New offset @@ -1842,7 +1787,7 @@ mod test { #[test] fn test_slice_with_offset() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1858,18 +1803,19 @@ mod test { assert_arrays_eq!( sliced.indices(), PrimitiveArray::from_iter([10u64]), - &mut assertion_ctx + &mut ctx ); // Index stays the same (offset + 5 = 10) assert_arrays_eq!( sliced.values(), PrimitiveArray::from_iter([200i32]), - &mut assertion_ctx + &mut ctx ); assert_eq!(sliced.offset(), 8); // New offset = 5 + 3 } #[test] fn test_patch_values() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1882,30 +1828,15 @@ mod test { #[expect(deprecated)] let values = patches.values().to_primitive(); assert_eq!( - i32::try_from( - &values - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ) - .unwrap(), + i32::try_from(&values.execute_scalar(0, &mut ctx).unwrap()).unwrap(), 100i32 ); assert_eq!( - i32::try_from( - &values - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ) - .unwrap(), + i32::try_from(&values.execute_scalar(1, &mut ctx).unwrap()).unwrap(), 200i32 ); assert_eq!( - i32::try_from( - &values - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - ) - .unwrap(), + i32::try_from(&values.execute_scalar(2, &mut ctx).unwrap()).unwrap(), 300i32 ); } @@ -1950,7 +1881,7 @@ mod test { #[test] fn test_mask_boundary_patches() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test masking patches at array boundaries let patches = Patches::new( 10, @@ -1964,25 +1895,24 @@ mod test { let mask = Mask::from_iter([ true, false, false, false, false, false, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap(); assert!(masked.is_some()); let masked = masked.unwrap(); assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([9u64]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( masked.values(), PrimitiveArray::from_iter([200i32]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_mask_all_patches_removed() { + let mut ctx = array_session().create_execution_ctx(); // Test when all patches are masked out let patches = Patches::new( 10, @@ -1997,15 +1927,13 @@ mod test { let mask = Mask::from_iter([ false, false, true, false, false, true, false, false, true, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap(); assert!(masked.is_none()); } #[test] fn test_mask_no_patches_removed() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test when no patches are masked let patches = Patches::new( 10, @@ -2020,26 +1948,23 @@ mod test { let mask = Mask::from_iter([ true, false, false, true, false, false, true, false, false, true, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([2u64, 5, 8]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( masked.values(), PrimitiveArray::from_iter([100i32, 200, 300]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_mask_single_patch() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test with a single patch let patches = Patches::new( 5, @@ -2052,27 +1977,22 @@ mod test { // Mask that removes the single patch let mask = Mask::from_iter([false, false, true, false, false]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap(); assert!(masked.is_none()); // Mask that keeps the single patch let mask = Mask::from_iter([true, false, false, true, false]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([2u64]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_mask_contiguous_patches() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test with contiguous patches let patches = Patches::new( 10, @@ -2087,26 +2007,23 @@ mod test { let mask = Mask::from_iter([ false, false, false, false, true, true, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([3u64, 6]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( masked.values(), PrimitiveArray::from_iter([100i32, 400]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_mask_with_large_offset() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Test with a large offset that shifts all indices let patches = Patches::new( 20, @@ -2122,26 +2039,24 @@ mod test { false, false, true, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, ]); - let masked = patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .unwrap(); + let masked = patches.mask(&mask, &mut ctx).unwrap().unwrap(); assert_arrays_eq!( masked.indices(), PrimitiveArray::from_iter([16u64, 19]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( masked.values(), PrimitiveArray::from_iter([100i32, 300]), - &mut assertion_ctx + &mut ctx ); } #[test] #[should_panic(expected = "Filter mask length 5 does not match array length 10")] fn test_mask_wrong_length() { + let mut ctx = array_session().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -2153,9 +2068,7 @@ mod test { // Mask with wrong length let mask = Mask::from_iter([false, false, true, false, false]); - patches - .mask(&mask, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + patches.mask(&mask, &mut ctx).unwrap(); } #[test] diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index ad787456410..7d15b250a50 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -469,7 +469,7 @@ mod tests { #[test] fn test_between_decimal() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = buffer![100i128, 200i128, 300i128, 400i128]; let decimal_type = DecimalDType::new(3, 2); let array = DecimalArray::new(values, decimal_type, Validity::NonNullable).into_array(); @@ -508,7 +508,7 @@ mod tests { assert_arrays_eq!( between_strict, BoolArray::from_iter([false, true, true, true]), - &mut assertion_ctx + &mut ctx ); // Non-strict lower bound, strict upper bound @@ -526,7 +526,7 @@ mod tests { assert_arrays_eq!( between_strict, BoolArray::from_iter([true, true, true, false]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index ebe9c59b442..d109ba3d01b 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -759,6 +759,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::assert_arrays_eq; @@ -772,7 +773,7 @@ mod tests { #[test] fn test_kleene_truth_table() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = BoolArray::from_iter([ Some(true), Some(true), @@ -811,7 +812,7 @@ mod tests { Some(false), None, ]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( @@ -827,7 +828,7 @@ mod tests { None, None, ]), - &mut assertion_ctx + &mut ctx ); Ok(()) @@ -835,7 +836,7 @@ mod tests { #[test] fn test_null_constant_kleene() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = BoolArray::from_iter([Some(false), Some(true), None]).into_array(); let null = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len()) .into_array(); @@ -843,12 +844,12 @@ mod tests { assert_arrays_eq!( lhs.binary(null.clone(), Operator::And)?, BoolArray::from_iter([Some(false), None, None]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( lhs.binary(null, Operator::Or)?, BoolArray::from_iter([None, Some(true), None]), - &mut assertion_ctx + &mut ctx ); Ok(()) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare.rs b/vortex-array/src/scalar_fn/fns/binary/compare.rs index 20f8b70360c..67f34bbddd6 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare.rs @@ -253,6 +253,7 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; @@ -369,15 +370,15 @@ mod tests { #[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())] #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())] fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let res = left.binary(right, Operator::Eq).unwrap(); let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(res, expected, &mut assertion_ctx); + assert_arrays_eq!(res, expected, &mut ctx); } #[test] fn test_list_array_comparison() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values1 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets1 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list1 = ListArray::try_new( @@ -402,7 +403,7 @@ mod tests { .binary(list2.clone().into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); let result = list1 .clone() @@ -410,19 +411,19 @@ mod tests { .binary(list2.clone().into_array(), Operator::NotEq) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); let result = list1 .into_array() .binary(list2.into_array(), Operator::Lt) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_list_array_constant_comparison() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list = ListArray::try_new( @@ -444,12 +445,12 @@ mod tests { .binary(constant.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([false, true, false]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_struct_array_comparison() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let bool_field1 = BoolArray::from_iter([Some(true), Some(false), Some(true)]); let int_field1 = PrimitiveArray::from_iter([1i32, 2, 3]); @@ -474,19 +475,19 @@ mod tests { .binary(struct2.clone().into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); let result = struct1 .into_array() .binary(struct2.into_array(), Operator::Gt) .unwrap(); let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_empty_struct_compare() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let empty1 = StructArray::try_new( FieldNames::from(Vec::::new()), Vec::new(), @@ -508,14 +509,14 @@ mod tests { .binary(empty2.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, true, true, true]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } /// Regression test: comparing struct arrays where the same logical field is backed by /// different Vortex encodings (VarBinArray vs VarBinViewArray) must not panic. #[test] fn struct_compare_mixed_binary_encodings() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // LHS: struct with a VarBinArray (offset-based) binary field let bin_field1 = VarBinArray::from(vec![ "apple".as_bytes(), @@ -537,7 +538,7 @@ mod tests { .binary(struct2.into_array(), Operator::Eq) .unwrap(); let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } /// Regression test: `scalar_cmp` must error when comparing scalars with incompatible diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 8c7a425b169..9dc142eb892 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -283,6 +283,7 @@ mod tests { use super::*; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; @@ -519,7 +520,7 @@ mod tests { #[test] fn test_or_kleene_validity() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::StructArray; @@ -541,13 +542,13 @@ mod tests { assert_arrays_eq!( result, BoolArray::from_iter([Some(true)]).into_array(), - &mut assertion_ctx + &mut ctx ) } #[test] fn test_scalar_subtract_unsigned() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -557,16 +558,12 @@ mod tests { let values = buffer![1u16, 2, 3].into_array(); let rhs = ConstantArray::new(Scalar::from(1u16), 3).into_array(); let result = values.binary(rhs, Operator::Sub).unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([0u16, 1, 2]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx); } #[test] fn test_scalar_subtract_signed() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -576,16 +573,12 @@ mod tests { let values = buffer![1i64, 2, 3].into_array(); let rhs = ConstantArray::new(Scalar::from(-1i64), 3).into_array(); let result = values.binary(rhs, Operator::Sub).unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([2i64, 3, 4]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx); } #[test] fn test_scalar_subtract_nullable() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -596,13 +589,13 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_scalar_subtract_float() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -615,7 +608,7 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index 2e96c49f830..0c824f59002 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -929,6 +929,7 @@ mod test { use crate::LEGACY_SESSION; use crate::RecursiveCanonical; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -951,49 +952,41 @@ mod test { #[test] fn test_scalar_subtract_unsigned() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = buffer![1u16, 2, 3].into_array(); let result = sub_scalar(&values, 1u16).unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([0u16, 1, 2]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx); } #[test] fn test_scalar_subtract_signed() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = buffer![1i64, 2, 3].into_array(); let result = sub_scalar(&values, -1i64).unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([2i64, 3, 4]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx); } #[test] fn test_scalar_subtract_nullable() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); let result = sub_scalar(&values.into_array(), Some(1u16)).unwrap(); assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_scalar_subtract_float() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = buffer![1.0f64, 2.0, 3.0].into_array(); let result = sub_scalar(&values, -1f64).unwrap(); assert_arrays_eq!( result, PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]), - &mut assertion_ctx + &mut ctx ); } @@ -1006,7 +999,7 @@ mod test { #[test] fn test_float_divide_by_zero_is_ok() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = buffer![1.0f64, -1.0].into_array(); let result = values .binary( @@ -1019,7 +1012,7 @@ mod test { assert_arrays_eq!( result, PrimitiveArray::from_iter([f64::INFINITY, f64::NEG_INFINITY]), - &mut assertion_ctx + &mut ctx ); } @@ -1064,7 +1057,7 @@ mod test { #[test] fn test_integer_divide_errors_ignore_null_lanes() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![10i32, 10], Validity::from_iter([false, true])) .into_array(); let rhs = buffer![0i32, 2].into_array(); @@ -1079,13 +1072,13 @@ mod test { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([None, Some(5i32)]), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_integer_errors_ignore_null_lanes() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true])) .into_array(); let result = values @@ -1102,7 +1095,7 @@ mod test { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([None, Some(2u8)]), - &mut assertion_ctx + &mut ctx ); } @@ -1123,7 +1116,7 @@ mod test { #[test] fn test_present_nullable_constant_preserves_nullable_output() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = buffer![1u8, 2].into_array(); let result = values .binary( @@ -1136,7 +1129,7 @@ mod test { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(2u8), Some(3)]), - &mut assertion_ctx + &mut ctx ); } } diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-array/src/scalar_fn/fns/byte_length.rs index 94447f356f5..064e3a95869 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-array/src/scalar_fn/fns/byte_length.rs @@ -182,6 +182,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; @@ -202,20 +203,20 @@ mod tests { #[case] array: ArrayRef, #[case] expected_lens: Vec, ) -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(expected_lens); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } #[test] fn test_varbinview_byte_length() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["short", "a longer string here"]).into_array(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 20]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index bdba538d410..7ff2190ccee 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -783,7 +783,7 @@ mod tests { #[test] fn test_evaluate_simple_condition() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -799,13 +799,13 @@ mod tests { assert_arrays_eq!( result, buffer![0i32, 0, 100, 100, 100].into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_nary_multiple_conditions() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test n-ary via nested_case_when let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) @@ -821,16 +821,12 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!( - result, - buffer![10i32, 0, 30, 0, 0].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(result, buffer![10i32, 0, 30, 0, 0].into_array(), &mut ctx); } #[test] fn test_evaluate_nary_first_match_wins() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -849,13 +845,13 @@ mod tests { assert_arrays_eq!( result, buffer![0i32, 0, 100, 100, 100].into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_no_else_returns_null() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -869,13 +865,13 @@ mod tests { result, PrimitiveArray::from_option_iter([None::, None, None, Some(100), Some(100)]) .into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_all_conditions_false() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -888,16 +884,12 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!( - result, - buffer![0i32, 0, 0, 0, 0].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(result, buffer![0i32, 0, 0, 0, 0].into_array(), &mut ctx); } #[test] fn test_evaluate_all_conditions_true() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -913,13 +905,13 @@ mod tests { assert_arrays_eq!( result, buffer![100i32, 100, 100, 100, 100].into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_all_true_no_else_returns_correct_dtype() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // CASE WHEN value > 0 THEN 100 END — condition is always true, no ELSE. // Result must be Nullable because the implicit ELSE is NULL. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -937,13 +929,13 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(100i32), Some(100), Some(100)]).into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_merge_case_branches_widens_nullability_of_later_branch() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // When a later THEN branch is Nullable and branches[0] and ELSE are NonNullable, // the result dtype must still be Nullable. // @@ -974,28 +966,24 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(10), Some(20), Some(0)]).into_array(), - &mut assertion_ctx + &mut ctx ); Ok(()) } #[test] fn test_evaluate_with_literal_condition() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = case_when(lit(true), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!( - result, - buffer![100i32, 100, 100].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(result, buffer![100i32, 100, 100].into_array(), &mut ctx); } #[test] fn test_evaluate_with_bool_column_result() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1011,13 +999,13 @@ mod tests { assert_arrays_eq!( result, BoolArray::from_iter([false, false, true, true, true]).into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_with_nullable_condition() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]).into_array(), @@ -1028,16 +1016,12 @@ mod tests { let expr = case_when(get_item("cond", root()), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!( - result, - buffer![100i32, 0, 0, 0, 100].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(result, buffer![100i32, 0, 0, 0, 100].into_array(), &mut ctx); } #[test] fn test_evaluate_with_nullable_result_values() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[ ("value", buffer![1i32, 2, 3, 4, 5].into_array()), ( @@ -1060,13 +1044,13 @@ mod tests { result, PrimitiveArray::from_option_iter([Some(0i32), Some(0), Some(30), Some(40), Some(50)]) .into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_with_all_null_condition() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[( "cond", BoolArray::from_iter([None, None, None]).into_array(), @@ -1077,14 +1061,14 @@ mod tests { let expr = case_when(get_item("cond", root()), lit(100i32), lit(0i32)); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!(result, buffer![0i32, 0, 0].into_array(), &mut assertion_ctx); + assert_arrays_eq!(result, buffer![0i32, 0, 0].into_array(), &mut ctx); } // ==================== N-ary Evaluate Tests ==================== #[test] fn test_evaluate_nary_no_else_returns_null() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1105,13 +1089,13 @@ mod tests { result, PrimitiveArray::from_option_iter([Some(10i32), None, Some(30), None, None]) .into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_nary_many_conditions() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3, 4, 5].into_array())]) .unwrap() @@ -1133,13 +1117,13 @@ mod tests { assert_arrays_eq!( result, buffer![10i32, 20, 30, 40, 50].into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_nary_all_false_no_else() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) .unwrap() .into_array(); @@ -1158,13 +1142,13 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter([None::, None, None]).into_array(), - &mut assertion_ctx + &mut ctx ); } #[test] fn test_evaluate_nary_overlapping_conditions_first_wins() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[("value", buffer![10i32, 20, 30].into_array())]) .unwrap() @@ -1184,12 +1168,12 @@ mod tests { let result = evaluate_expr(&expr, &test_array); // First matching condition always wins - assert_arrays_eq!(result, buffer![1i32, 1, 1].into_array(), &mut assertion_ctx); + assert_arrays_eq!(result, buffer![1i32, 1, 1].into_array(), &mut ctx); } #[test] fn test_evaluate_nary_early_exit_when_remaining_empty() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // After branch 0 claims all rows, remaining becomes all_false. // The loop breaks before evaluating branch 1's condition. let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1206,16 +1190,12 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!( - result, - buffer![100i32, 100, 100].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(result, buffer![100i32, 100, 100].into_array(), &mut ctx); } #[test] fn test_evaluate_nary_skips_branch_with_empty_effective_mask() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Branch 0 claims value=1. Branch 1 targets the same rows but they are already // matched → effective_mask is all_false → branch 1 is skipped (THEN not used). let test_array = StructArray::from_fields(&[("value", buffer![1i32, 2, 3].into_array())]) @@ -1234,11 +1214,7 @@ mod tests { ); let result = evaluate_expr(&expr, &test_array); - assert_arrays_eq!( - result, - buffer![10i32, 20, 0].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(result, buffer![10i32, 20, 0].into_array(), &mut ctx); } #[test] @@ -1281,7 +1257,7 @@ mod tests { #[test] fn test_evaluate_nary_with_nullable_conditions() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let test_array = StructArray::from_fields(&[ ( "cond1", @@ -1307,11 +1283,7 @@ mod tests { // row 0: cond1=true → 10 // row 1: cond1=NULL(→false), cond2=true → 20 // row 2: cond1=false, cond2=NULL(→false) → else=0 - assert_arrays_eq!( - result, - buffer![10i32, 20, 0].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(result, buffer![10i32, 20, 0].into_array(), &mut ctx); } // ==================== Simplify: COALESCE -> fill_null ==================== @@ -1402,7 +1374,7 @@ mod tests { #[test] fn test_simplify_null_fill_semantic_equivalence() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // The collapse-to-input rewrite must preserve values (and `x`'s nullability). let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1420,16 +1392,8 @@ mod tests { ); let expected = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); - assert_arrays_eq!( - evaluate_expr(&original, &array), - expected, - &mut assertion_ctx - ); - assert_arrays_eq!( - evaluate_expr(&optimized, &array), - expected, - &mut assertion_ctx - ); + assert_arrays_eq!(evaluate_expr(&original, &array), expected, &mut ctx); + assert_arrays_eq!(evaluate_expr(&optimized, &array), expected, &mut ctx); Ok(()) } @@ -1463,7 +1427,7 @@ mod tests { #[test] fn test_simplify_semantic_equivalence() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // The optimized expression must produce the same values as the original CASE WHEN. let array = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); let scope = DType::Primitive(PType::I64, Nullability::Nullable); @@ -1480,19 +1444,19 @@ mod tests { assert_arrays_eq!( evaluate_expr(&original, &array), PrimitiveArray::from_option_iter([Some(1i64), Some(0), Some(3)]).into_array(), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( evaluate_expr(&optimized, &array), buffer![1i64, 0, 3].into_array(), - &mut assertion_ctx + &mut ctx ); Ok(()) } #[test] fn test_merge_case_branches_alternating_mask() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Exercises the scalar path: alternating rows produce one slice per row (no runs), // triggering the per-row cursor path in merge_case_branches. let n = 100usize; @@ -1523,7 +1487,7 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_option_iter(expected).into_array(), - &mut assertion_ctx + &mut ctx ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 700e6f7db21..74058aa19f2 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -278,6 +278,7 @@ mod tests { use super::*; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -304,7 +305,7 @@ mod tests { #[test] fn execute_with_value() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -314,17 +315,13 @@ mod tests { root(), ); let result = input.apply(&expr)?; - assert_arrays_eq!( - result, - BoolArray::from_iter([true, false, false]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); Ok(()) } #[test] fn execute_without_value_default_true() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -334,17 +331,13 @@ mod tests { root(), ); let result = input.apply(&expr)?; - assert_arrays_eq!( - result, - BoolArray::from_iter([true, true, true]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, true]), &mut ctx); Ok(()) } #[test] fn execute_without_value_default_false() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -357,14 +350,14 @@ mod tests { assert_arrays_eq!( result, BoolArray::from_iter([false, false, false]), - &mut assertion_ctx + &mut ctx ); Ok(()) } #[test] fn execute_value_flips() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let threshold = Arc::new(AtomicI32::new(5)); let threshold_clone = Arc::clone(&threshold); let expr = dynamic( @@ -377,19 +370,11 @@ mod tests { let input = buffer![1i32, 5, 10].into_array(); let result = input.clone().apply(&expr)?; - assert_arrays_eq!( - result, - BoolArray::from_iter([true, false, false]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); threshold.store(10, Ordering::SeqCst); let result = input.apply(&expr)?; - assert_arrays_eq!( - result, - BoolArray::from_iter([true, true, false]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index 2e5e934e550..c9eae72ad08 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -180,6 +180,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::assert_arrays_eq; @@ -209,7 +210,7 @@ mod tests { #[test] fn evaluate() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let test_array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]) .into_array(); @@ -224,13 +225,13 @@ mod tests { assert_arrays_eq!( result, PrimitiveArray::from_iter([1i32, 42, 3, 42, 5]), - &mut assertion_ctx + &mut ctx ); } #[test] fn evaluate_struct_field() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let test_array = StructArray::from_fields(&[( "a", PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), @@ -245,24 +246,16 @@ mod tests { result.dtype(), &DType::Primitive(PType::I32, Nullability::NonNullable) ); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([1i32, 0, 3]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 0, 3]), &mut ctx); } #[test] fn evaluate_non_nullable_input() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = fill_null(root(), lit(0i32)); let result = test_array.apply(&expr).unwrap(); - assert_arrays_eq!( - result, - PrimitiveArray::from_iter([1i32, 2, 3]), - &mut assertion_ctx - ); + assert_arrays_eq!(result, PrimitiveArray::from_iter([1i32, 2, 3]), &mut ctx); } #[test] diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 8edc4f13c97..cf0704b68f3 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -407,6 +407,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ListArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; @@ -752,7 +753,7 @@ mod tests { #[case] value: Option<&str>, #[case] expected: BoolArray, ) { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let element_nullability = list_array .dtype() .as_list_element_opt() @@ -765,12 +766,12 @@ mod tests { let elem = ConstantArray::new(scalar, list_array.len()); let expr = list_contains(root(), lit(elem.scalar().clone())); let result = list_array.apply(&expr).unwrap(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_constant_list() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let list_array = ConstantArray::new( Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -784,12 +785,12 @@ mod tests { let expr = list_contains(root(), lit(2i32)); let contains = list_array.apply(&expr).unwrap(); let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(contains, expected, &mut assertion_ctx); + assert_arrays_eq!(contains, expected, &mut ctx); } #[test] fn test_all_nulls() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let list_array = ConstantArray::new( Scalar::null(DType::List( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -806,12 +807,12 @@ mod tests { [false, false, false, false, false].into_iter().collect(), Validity::AllInvalid, ); - assert_arrays_eq!(contains, expected, &mut assertion_ctx); + assert_arrays_eq!(contains, expected, &mut ctx); } #[test] fn test_list_array_element() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let list_scalar = Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), vec![1.into(), 3.into(), 6.into()], @@ -823,12 +824,12 @@ mod tests { let contains = arr.apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, true, false, true, false, false, true]); - assert_arrays_eq!(contains, expected, &mut assertion_ctx); + assert_arrays_eq!(contains, expected, &mut ctx); } #[test] fn test_list_contains_empty_listview() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let empty_elements = PrimitiveArray::empty::(Nullability::NonNullable); let offsets = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); let sizes = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); @@ -847,12 +848,12 @@ mod tests { let result = list_array.into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, false, false, false]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_list_contains_all_null_elements() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter::([None, None, None, None, None]); let offsets = Buffer::from_iter([0u32, 2, 4]).into_array(); let sizes = Buffer::from_iter([2u32, 2, 1]).into_array(); @@ -876,19 +877,19 @@ mod tests { [false, false, false].into_iter().collect(), Validity::AllInvalid, ); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); // Searching for non-null let expr2 = list_contains(root(), lit(42i32)); let result2 = list_array.into_array().apply(&expr2).unwrap(); let expected2 = BoolArray::from_iter([false, false, false]); - assert_arrays_eq!(result2, expected2, &mut assertion_ctx); + assert_arrays_eq!(result2, expected2, &mut ctx); } #[test] fn test_list_contains_large_offsets() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = Buffer::from_iter([1i32, 2, 3, 4, 5]).into_array(); let offsets = Buffer::from_iter([0u32, 1, 4, 0]).into_array(); @@ -901,18 +902,18 @@ mod tests { let result = list_array.clone().into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, true, false, false]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); let expr5 = list_contains(root(), lit(5i32)); let result5 = list_array.into_array().apply(&expr5).unwrap(); let expected5 = BoolArray::from_iter([false, false, true, false]); - assert_arrays_eq!(result5, expected5, &mut assertion_ctx); + assert_arrays_eq!(result5, expected5, &mut ctx); } #[test] fn test_list_contains_offset_size_boundary() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let elements = Buffer::from_iter(0..256).into_array(); let offsets = Buffer::from_iter([0u8, 100, 200, 254]).into_array(); let sizes = Buffer::from_iter([50u8, 50, 54, 2]).into_array(); @@ -924,12 +925,12 @@ mod tests { let result = list_array.clone().into_array().apply(&expr).unwrap(); let expected = BoolArray::from_iter([false, false, false, true]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); let expr_zero = list_contains(root(), lit(0i32)); let result_zero = list_array.into_array().apply(&expr_zero).unwrap(); let expected_zero = BoolArray::from_iter([true, false, false, false]); - assert_arrays_eq!(result_zero, expected_zero, &mut assertion_ctx); + assert_arrays_eq!(result_zero, expected_zero, &mut ctx); } } diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 9dcfa7b5184..c9f5783de27 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -280,6 +280,7 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -319,7 +320,7 @@ mod tests { #[test] pub fn test_merge_right_most() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let expr = merge_opts( vec![ get_item("0", root()), @@ -370,27 +371,27 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array, &["a"]).unwrap(), PrimitiveArray::from_iter([0i32, 0, 0]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["b"]).unwrap(), PrimitiveArray::from_iter([2i32, 2, 2]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["c"]).unwrap(), PrimitiveArray::from_iter([3i32, 3, 3]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["d"]).unwrap(), PrimitiveArray::from_iter([4i32, 4, 4]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array, &["e"]).unwrap(), PrimitiveArray::from_iter([5i32, 5, 5]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 5ff70ffedc4..81019e22d6f 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -173,6 +173,7 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; @@ -231,7 +232,7 @@ mod tests { #[test] pub fn test_simple_pack() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -249,23 +250,23 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(), PrimitiveArray::from_iter([0i32, 1, 2]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two"]).unwrap(), PrimitiveArray::from_iter([4i32, 5, 6]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.into_array(), &["three"]).unwrap(), PrimitiveArray::from_iter([0i32, 1, 2]), - &mut assertion_ctx + &mut ctx ); } #[test] pub fn test_nested_pack() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -292,22 +293,22 @@ mod tests { assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(), PrimitiveArray::from_iter([0i32, 1, 2]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two", "two_one"]).unwrap(), PrimitiveArray::from_iter([4i32, 5, 6]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.clone().into_array(), &["two", "two_two"]).unwrap(), PrimitiveArray::from_iter([4i32, 5, 6]), - &mut assertion_ctx + &mut ctx ); assert_arrays_eq!( primitive_field(&actual_array.into_array(), &["three"]).unwrap(), PrimitiveArray::from_iter([0i32, 1, 2]), - &mut assertion_ctx + &mut ctx ); } diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 8d892c3add9..249d8bfcf49 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -278,6 +278,7 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::Struct; @@ -317,7 +318,7 @@ mod tests { #[test] fn test_zip_basic() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mask = Mask::from_iter([true, false, false, true, false]); let if_true = buffer![10, 20, 30, 40, 50].into_array(); let if_false = buffer![1, 2, 3, 4, 5].into_array(); @@ -325,12 +326,12 @@ mod tests { let result = mask.into_array().zip(if_true, if_false).unwrap(); let expected = buffer![10, 2, 3, 40, 5].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] fn test_zip_all_true() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mask = Mask::new_true(4); let if_true = buffer![10, 20, 30, 40].into_array(); let if_false = @@ -340,13 +341,13 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), Some(40)]).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); assert_eq!(result.dtype(), if_false.dtype()) } #[test] fn test_zip_all_false_widens_nullability() { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mask = Mask::new_false(4); let if_true = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array(); @@ -356,7 +357,7 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), Some(4)]).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); assert_eq!(result.dtype(), if_true.dtype()); } diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 0f84de369c7..2038df0b613 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -89,6 +89,7 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; + use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -105,11 +106,10 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = LazyLock::new(array_session); #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let sum_scalar = Scalar::primitive(6i64, Nullability::Nullable); array.statistics().set( @@ -124,14 +124,13 @@ mod tests { let expected = ConstantArray::new(Scalar::primitive(6i64, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_returns_null_when_sum_is_missing() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let result = array @@ -144,14 +143,13 @@ mod tests { 3, ) .into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_sum_per_chunk() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let chunk0 = buffer![1i32, 2].into_array(); let sum_scalar = Scalar::primitive(3i64, Nullability::Nullable); chunk0.statistics().set( @@ -180,14 +178,13 @@ mod tests { Validity::from_iter([true, true, false, false, false]), ) .into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); let null_count_scalar = Scalar::primitive(2u64, Nullability::NonNullable); @@ -207,14 +204,13 @@ mod tests { let expected = ConstantArray::new(Scalar::primitive(2u64, Nullability::Nullable), 4).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_all_null_from_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); array .statistics() @@ -227,14 +223,13 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_all_null_false_from_inexact_low_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, Some(2), None]).into_array(); array .statistics() @@ -247,14 +242,13 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_returns_null_for_inexact_full_null_count_as_all_null() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter::([None, Some(2), None]).into_array(); array .statistics() @@ -267,14 +261,13 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_all_non_null_from_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); array .statistics() @@ -287,14 +280,13 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_all_non_null_true_from_inexact_zero_null_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); array .statistics() @@ -307,14 +299,13 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_returns_null_for_inexact_nonzero_null_count_as_all_non_null() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); array @@ -328,7 +319,7 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 4).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -348,7 +339,6 @@ mod tests { #[test] fn stat_expr_reads_cached_all_nan_from_nan_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(f32::NAN), Some(f32::NAN)]) .into_array(); @@ -363,14 +353,13 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_all_nan_false_from_inexact_low_nan_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32), Some(f32::NAN)]) .into_array(); @@ -385,14 +374,13 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(false, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_returns_null_for_inexact_full_nan_count_as_all_nan() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(f32::NAN), Some(1.0f32), Some(f32::NAN)]) .into_array(); @@ -407,14 +395,13 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_all_non_nan_true_from_inexact_zero_nan_count() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![1.0f32, 2.0, 3.0].into_array(); array .statistics() @@ -427,14 +414,13 @@ mod tests { let expected = ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_returns_null_for_inexact_nonzero_nan_count_as_all_non_nan() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1.0f32), Some(f32::NAN), Some(3.0)]) .into_array(); array @@ -448,14 +434,13 @@ mod tests { let expected = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), 3).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } #[test] fn stat_expr_reads_cached_min_and_max() -> VortexResult<()> { - let mut assertion_ctx = crate::array_session().create_execution_ctx(); let array = buffer![3i32, 1, 2].into_array(); array .statistics() @@ -476,7 +461,11 @@ mod tests { .into_array(); let expected_min = ConstantArray::new(Scalar::primitive(1i32, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(min_result, expected_min, &mut assertion_ctx); + assert_arrays_eq!( + min_result, + expected_min, + &mut SESSION.create_execution_ctx() + ); let max_result = array .apply(&stat( @@ -489,7 +478,11 @@ mod tests { .into_array(); let expected_max = ConstantArray::new(Scalar::primitive(3i32, Nullability::Nullable), 3).into_array(); - assert_arrays_eq!(max_result, expected_max, &mut assertion_ctx); + assert_arrays_eq!( + max_result, + expected_max, + &mut SESSION.create_execution_ctx() + ); Ok(()) } diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index e182281ebb5..bddb0be113a 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -112,7 +112,7 @@ mod tests { #[case] input: ListViewArray, #[case] expect_list: bool, ) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array_ref = input.clone().into_array(); let result = BtrBlocksCompressor::default() .compress(&array_ref, &mut SESSION.create_execution_ctx())?; @@ -121,13 +121,13 @@ mod tests { } else { assert!(result.as_opt::().is_some()); } - assert_arrays_eq!(result, input, &mut assertion_ctx); + assert_arrays_eq!(result, input, &mut ctx); Ok(()) } #[test] fn test_constant_all_true() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -135,13 +135,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut assertion_ctx); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_constant_all_false() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable); let btr = BtrBlocksCompressor::default(); let compressed = btr.compress( @@ -149,13 +149,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut assertion_ctx); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_nullable_all_valid_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true; 100]), Validity::from(BitBuffer::from(vec![true; 100])), @@ -166,13 +166,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut assertion_ctx); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0))); let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity); let btr = BtrBlocksCompressor::default(); @@ -181,13 +181,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(!compressed.is::()); - assert_arrays_eq!(compressed, array, &mut assertion_ctx); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_mixed_not_constant() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true, false, true]), Validity::NonNullable, @@ -198,13 +198,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(!compressed.is::()); - assert_arrays_eq!(compressed, array, &mut assertion_ctx); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_binary_constant_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = vec![Some(b"constant-bytes".as_slice()); 100]; let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); let btr = BtrBlocksCompressor::default(); @@ -213,13 +213,13 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut assertion_ctx); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } #[test] fn test_binary_dict_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"]; let values = (0..1000) .map(|idx| Some(distinct_values[idx % distinct_values.len()])) @@ -231,7 +231,7 @@ mod tests { &mut SESSION.create_execution_ctx(), )?; assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut assertion_ctx); + assert_arrays_eq!(compressed, array, &mut ctx); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index dfa210264c8..2d0a04542cc 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -57,7 +57,7 @@ fn test_compress() -> VortexResult<()> { #[test] fn test_rle_compression() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(1.5f32, 100)); values.extend(iter::repeat_n(2.7f32, 200)); @@ -71,7 +71,7 @@ fn test_rle_compression() -> VortexResult<()> { assert!(compressed.is::()); let expected = Buffer::copy_from(&values).into_array(); - assert_arrays_eq!(compressed, expected, &mut assertion_ctx); + assert_arrays_eq!(compressed, expected, &mut ctx); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index 54f0565c5bf..e4227a472ec 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -161,7 +161,7 @@ fn test_rle_compressed() -> VortexResult<()> { #[cfg(feature = "unstable_encodings")] #[test] fn test_delta_compressed() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); use vortex_array::assert_arrays_eq; use vortex_fastlanes::Delta; @@ -191,7 +191,7 @@ fn test_delta_compressed() -> VortexResult<()> { "Delta was applied more than once in the tree:\n{}", compressed.display_tree() ); - assert_arrays_eq!(compressed, array.into_array(), &mut assertion_ctx); + assert_arrays_eq!(compressed, array.into_array(), &mut ctx); Ok(()) } diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index ed03c017a6e..a9ef24dc0e0 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -71,7 +71,7 @@ fn test_dict_encodable() -> VortexResult<()> { #[test] fn constant_mostly_nulls() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::new( buffer![189u8, 189, 189, 189, 189, 189, 189, 189, 189, 0, 46], Validity::from_iter(vec![ @@ -89,13 +89,13 @@ fn constant_mostly_nulls() -> VortexResult<()> { let decoded = compressed; let expected = PrimitiveArray::new(buffer![0u8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 46], validity).into_array(); - assert_arrays_eq!(decoded, expected, &mut assertion_ctx); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn nullable_sequence() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let values = (0i32..20).step_by(7).collect_vec(); let array = PrimitiveArray::from_option_iter(values.clone().into_iter().map(Some)); @@ -105,13 +105,13 @@ fn nullable_sequence() -> VortexResult<()> { let decoded = compressed; let expected = PrimitiveArray::from_option_iter(values.into_iter().map(Some)).into_array(); - assert_arrays_eq!(decoded, expected, &mut assertion_ctx); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } #[test] fn test_rle_compression() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut values = Vec::new(); values.extend(iter::repeat_n(42i32, 100)); values.extend(iter::repeat_n(123i32, 200)); @@ -124,7 +124,7 @@ fn test_rle_compression() -> VortexResult<()> { assert!(compressed.is::()); let expected = Buffer::copy_from(&values).into_array(); - assert_arrays_eq!(compressed, expected, &mut assertion_ctx); + assert_arrays_eq!(compressed, expected, &mut ctx); Ok(()) } diff --git a/vortex-compressor/src/sample.rs b/vortex-compressor/src/sample.rs index 27d7383ade8..deb5bbe6f37 100644 --- a/vortex-compressor/src/sample.rs +++ b/vortex-compressor/src/sample.rs @@ -134,6 +134,7 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::validity::Validity; @@ -144,7 +145,7 @@ mod tests { #[test] fn sample_is_deterministic() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Create a deterministic array with linear-with-noise pattern let values: Vec = (0i64..100_000).map(|i| i + (i * 7 + 3) % 11).collect(); @@ -155,7 +156,7 @@ mod tests { for _ in 0..10 { let again = sample(&array, SAMPLE_SIZE, SAMPLE_COUNT); assert_eq!(first.nbytes(), again.nbytes()); - assert_arrays_eq!(&first, &again, &mut assertion_ctx); + assert_arrays_eq!(&first, &again, &mut ctx); } Ok(()) } diff --git a/vortex-cuda/src/dynamic_dispatch/mod.rs b/vortex-cuda/src/dynamic_dispatch/mod.rs index 2d351d3ee58..f4d374f7b00 100644 --- a/vortex-cuda/src/dynamic_dispatch/mod.rs +++ b/vortex-cuda/src/dynamic_dispatch/mod.rs @@ -545,6 +545,7 @@ mod tests { use vortex::array::validity::Validity; use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; + use vortex::buffer::buffer; use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::dtype::PType; @@ -560,11 +561,16 @@ mod tests { use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArrayExt; use vortex::encodings::runend::RunEnd; + use vortex::encodings::sequence::Sequence; use vortex::encodings::zigzag::ZigZag; use vortex::error::VortexExpect; use vortex::error::VortexResult; - use vortex_array::LEGACY_SESSION; + use vortex::mask::Mask; + use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::FilterArray; + use vortex_array::assert_arrays_eq; use vortex_array::patches::Patches; use super::*; @@ -572,23 +578,20 @@ mod tests { use crate::CudaBufferExt; use crate::CudaDeviceBuffer; use crate::CudaExecutionCtx; + use crate::cuda_session; use crate::executor::CudaArrayExt; use crate::executor::CudaDispatchMode; use crate::hybrid_dispatch::try_gpu_dispatch; use crate::session::CudaSession; - fn bitpacked_array_u32(bit_width: u8, len: usize) -> BitPackedArray { + fn bitpacked_array_u32(bit_width: u8, len: usize, ctx: &mut ExecutionCtx) -> BitPackedArray { let max_val = (1u64 << bit_width).saturating_sub(1); let values: Vec = (0..len) .map(|i| ((i as u64) % (max_val + 1)) as u32) .collect(); let primitive = PrimitiveArray::new(Buffer::from(values), NonNullable); - BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("failed to create BitPacked array") + BitPacked::encode(&primitive.into_array(), bit_width, ctx) + .vortex_expect("failed to create BitPacked array") } async fn dispatch_plan( @@ -624,6 +627,7 @@ mod tests { #[crate::test] fn test_max_scalar_ops() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 6; let len = 2050; let references: [u32; 4] = [1, 2, 4, 8]; @@ -634,8 +638,8 @@ mod tests { .map(|i| ((i as u64) % (max_val + 1)) as u32 + total_reference) .collect(); - let bitpacked = bitpacked_array_u32(bit_width, len); - let cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let bitpacked = bitpacked_array_u32(bit_width, len, &mut ctx); + let cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let packed = bitpacked.packed().clone(); let device_input = futures::executor::block_on(cuda_ctx.ensure_on_device(packed))?; let input_ptr = device_input.cuda_device_ptr()?; @@ -751,7 +755,7 @@ mod tests { }) .collect(); - let cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let (input_ptr, _di) = copy_raw_to_device(&cuda_ctx, &data)?; let plan = CudaDispatchPlan::new( @@ -844,6 +848,7 @@ mod tests { #[crate::test] async fn test_bitpacked() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let max_val = (1u64 << bit_width).saturating_sub(1); @@ -851,8 +856,8 @@ mod tests { .map(|i| ((i as u64) % (max_val + 1)) as u32) .collect(); - let bp = bitpacked_array_u32(bit_width, len); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let bp = bitpacked_array_u32(bit_width, len, &mut ctx); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&bp.into_array(), &mut cuda_ctx).await?; let actual = @@ -864,6 +869,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 6; let len = 3000; let reference = 42u32; @@ -874,10 +880,10 @@ mod tests { .collect(); let expected: Vec = raw.iter().map(|&v| v + reference).collect(); - let bp = bitpacked_array_u32(bit_width, len); + let bp = bitpacked_array_u32(bit_width, len, &mut ctx); let for_arr = FoR::try_new(bp.into_array(), Scalar::from(reference))?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&for_arr.into_array(), &mut cuda_ctx).await?; let actual = @@ -899,7 +905,7 @@ mod tests { expected.push(values[run]); } - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, cuda_ctx.execution_ctx()); @@ -915,6 +921,7 @@ mod tests { #[crate::test] async fn test_dict_for_bp_values_bp_codes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Dict where both codes and values are BitPacked+FoR. let dict_reference = 1_000_000u32; let dict_residuals: Vec = (0..64).collect(); @@ -927,24 +934,16 @@ mod tests { // BitPack+FoR the dict values let dict_prim = PrimitiveArray::new(Buffer::from(dict_residuals), NonNullable); - let dict_bp = BitPacked::encode( - &dict_prim.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let dict_bp = BitPacked::encode(&dict_prim.into_array(), 6, &mut ctx)?; let dict_for = FoR::try_new(dict_bp.into_array(), Scalar::from(dict_reference))?; // BitPack the codes let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), 6, &mut ctx)?; let dict = DictArray::try_new(codes_bp.into_array(), dict_for.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&dict.into_array(), &mut cuda_ctx).await?; let actual = @@ -956,7 +955,7 @@ mod tests { #[crate::test] async fn test_alp_for_bitpacked() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // ALP(FoR(BitPacked)): encode each layer, then reassemble the tree // bottom-up because encode() methods produce flat outputs. let len = 3000; @@ -969,11 +968,7 @@ mod tests { let alp = alp_encode(float_prim.as_view(), Some(exponents), &mut ctx)?; assert!(alp.patches().is_none()); let for_arr = FoR::encode(alp.encoded().clone().execute::(&mut ctx)?)?; - let bp = BitPacked::encode( - for_arr.encoded(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(for_arr.encoded(), 6, &mut ctx)?; let tree = ALP::new( FoR::try_new(bp.into_array(), for_arr.reference_scalar().clone())?.into_array(), @@ -981,7 +976,7 @@ mod tests { None, ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&tree.into_array(), &mut cuda_ctx).await?; let actual = @@ -993,6 +988,7 @@ mod tests { #[crate::test] async fn test_zigzag_bitpacked() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // ZigZag(BitPacked): unpack then zigzag-decode. let bit_width: u8 = 4; let len = 3000; @@ -1007,14 +1003,10 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(raw), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; let zz = ZigZag::try_new(bp.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&zz.into_array(), &mut cuda_ctx).await?; let actual = @@ -1038,7 +1030,7 @@ mod tests { expected.push(values[run] + reference); } - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, cuda_ctx.execution_ctx()); @@ -1072,7 +1064,7 @@ mod tests { let dict = DictArray::try_new(codes_prim.into_array(), values_prim.into_array())?; let for_arr = FoR::try_new(dict.into_array(), Scalar::from(reference))?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&for_arr.into_array(), &mut cuda_ctx).await?; let actual = @@ -1084,6 +1076,7 @@ mod tests { #[crate::test] async fn test_dict_for_bp_codes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // Dict(codes=FoR(BitPacked), values=primitive) let dict_values: Vec = (0..8).map(|i| i * 1000 + 7).collect(); let dict_size = dict_values.len(); @@ -1094,17 +1087,13 @@ mod tests { // BitPack codes, then wrap in FoR (reference=0 so values unchanged) let bit_width: u8 = 3; let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut ctx)?; let codes_for = FoR::try_new(codes_bp.into_array(), Scalar::from(0u32))?; let values_prim = PrimitiveArray::new(Buffer::from(dict_values), NonNullable); let dict = DictArray::try_new(codes_for.into_array(), values_prim.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&dict.into_array(), &mut cuda_ctx).await?; let actual = @@ -1116,6 +1105,7 @@ mod tests { #[crate::test] async fn test_dict_primitive_values_bp_codes() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let dict_size = dict_values.len(); let len = 3000; @@ -1124,16 +1114,12 @@ mod tests { let bit_width: u8 = 2; let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut ctx)?; let values_prim = PrimitiveArray::new(Buffer::from(dict_values), NonNullable); let dict = DictArray::try_new(codes_bp.into_array(), values_prim.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&dict.into_array(), &mut cuda_ctx).await?; let actual = @@ -1145,7 +1131,7 @@ mod tests { #[crate::test] async fn test_dict_mixed_width_u8_codes_u32_values() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let len = 3000; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u8).collect(); @@ -1165,20 +1151,20 @@ mod tests { ); // Execute through the hybrid dispatch path (handles widening). - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } #[crate::test] async fn test_dict_mixed_width_u16_codes_u32_values() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let dict_values: Vec = vec![1000, 2000, 3000, 4000, 5000]; let len = 2048; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u16).collect(); @@ -1198,25 +1184,25 @@ mod tests { ); // Execute through the hybrid dispatch path (handles widening). - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_runend_mixed_width_u64_ends_u32_values() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let ends: Vec = vec![1000, 2000, 3000]; let values: Vec = vec![10, 20, 30]; let len = 3000; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, cuda_ctx.execution_ctx()); @@ -1246,7 +1232,7 @@ mod tests { }) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -1280,7 +1266,7 @@ mod tests { let expected: Vec = data[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1314,6 +1300,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; let len = 5000; @@ -1325,17 +1312,13 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(raw), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; let zz = ZigZag::try_new(bp.into_array())?; let sliced = zz.into_array().slice(slice_start..slice_end)?; let expected: Vec = all_decoded[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1385,7 +1368,7 @@ mod tests { .map(|&c| dict_values[c as usize]) .collect(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1419,22 +1402,19 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; let len = 5000; let data: Vec = (0..len).map(|i| (i as u32) % max_val).collect(); let prim = PrimitiveArray::new(Buffer::from(data.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; let sliced = bp.into_array().slice(slice_start..slice_end)?; let expected: Vec = data[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1468,6 +1448,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let reference = 100u32; let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; @@ -1475,11 +1456,7 @@ mod tests { let encoded_data: Vec = (0..len).map(|i| (i as u32) % max_val).collect(); let prim = PrimitiveArray::new(Buffer::from(encoded_data.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; let for_arr = FoR::try_new(bp.into_array(), Scalar::from(reference))?; let all_decoded: Vec = encoded_data.iter().map(|&v| v + reference).collect(); @@ -1487,7 +1464,7 @@ mod tests { let sliced = for_arr.into_array().slice(slice_start..slice_end)?; let expected: Vec = all_decoded[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1521,6 +1498,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict_reference = 1_000_000u32; let dict_residuals: Vec = (0..64).collect(); let dict_expected: Vec = dict_residuals.iter().map(|&r| r + dict_reference).collect(); @@ -1532,27 +1510,19 @@ mod tests { // BitPack+FoR the dict values let dict_prim = PrimitiveArray::new(Buffer::from(dict_residuals), NonNullable); - let dict_bp = BitPacked::encode( - &dict_prim.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let dict_bp = BitPacked::encode(&dict_prim.into_array(), 6, &mut ctx)?; let dict_for = FoR::try_new(dict_bp.into_array(), Scalar::from(dict_reference))?; // BitPack the codes let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), 6, &mut ctx)?; let dict = DictArray::try_new(codes_bp.into_array(), dict_for.into_array())?; let sliced = dict.into_array().slice(slice_start..slice_end)?; let expected: Vec = all_decoded[slice_start..slice_end].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1584,7 +1554,7 @@ mod tests { let seq = Sequence::try_new_typed(base, multiplier, Nullability::NonNullable, len)?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&seq.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( @@ -1617,7 +1587,7 @@ mod tests { let seq = Sequence::try_new_typed(base, multiplier, Nullability::NonNullable, len)?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&seq.into_array(), &mut cuda_ctx).await?; let actual_u32 = run_dynamic_dispatch_plan( @@ -1634,7 +1604,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u8() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 4; let len = 3000; let reference = 100u8; @@ -1646,30 +1616,26 @@ mod tests { .collect(); let primitive = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack u8"); + let bp = BitPacked::encode(&primitive.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack u8"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(reference, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_for_bitpacked_u16() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let reference = 1000u16; @@ -1683,30 +1649,26 @@ mod tests { .collect(); let primitive = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack u16"); + let bp = BitPacked::encode(&primitive.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack u16"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(reference, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_for_bitpacked_u64() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 20; let len = 3000; let reference = 100_000u64; @@ -1718,24 +1680,20 @@ mod tests { .collect(); let primitive = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack u64"); + let bp = BitPacked::encode(&primitive.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack u64"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(reference, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -1743,7 +1701,7 @@ mod tests { async fn test_empty_array() -> VortexResult<()> { let values: Vec = vec![]; let primitive = PrimitiveArray::new(Buffer::from(values), NonNullable); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&primitive.into_array(), &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); assert_eq!(result.len(), 0); @@ -1752,33 +1710,29 @@ mod tests { #[crate::test] async fn test_single_element() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let values: Vec = vec![42]; let primitive = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack"); + let bp = + BitPacked::encode(&primitive.into_array(), 6, &mut cpu_ctx).vortex_expect("bitpack"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(0u32, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected, result, &mut cpu_ctx); + assert_arrays_eq!(expected, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_exactly_elements_per_block() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Exactly 2048 elements — one full block, no remainder let bit_width: u8 = 6; let len = 2048; @@ -1790,24 +1744,20 @@ mod tests { let expected: Vec = residuals.iter().map(|&r| r + reference).collect(); let primitive = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &primitive.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack"); + let bp = BitPacked::encode(&primitive.into_array(), bit_width, &mut ctx) + .vortex_expect("bitpack"); let for_arr = FoR::try_new( bp.into_array(), Scalar::primitive(reference, Nullability::NonNullable), )?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } @@ -1817,8 +1767,9 @@ mod tests { #[crate::test] fn test_has_standalone_kernel_true_cases() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); // BitPacked — leaf encoding, no children. - let bp = bitpacked_array_u32(6, 2048); + let bp = bitpacked_array_u32(6, 2048, &mut ctx); let bp_arr = bp.into_array(); assert!(plan_builder::has_standalone_kernel(&bp_arr)); @@ -1868,14 +1819,12 @@ mod tests { /// a kernel registered in the CUDA session. #[crate::test] fn test_has_standalone_kernel_implies_registered_kernel() -> VortexResult<()> { - use vortex::dtype::Nullability; - use vortex::encodings::sequence::Sequence; - - let session = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut ctx = array_session().create_execution_ctx(); + let session = CudaSession::create_execution_ctx(&cuda_session())?; let cuda_session = session.cuda_session(); // Leaf encodings. - let bp = bitpacked_array_u32(6, 2048); + let bp = bitpacked_array_u32(6, 2048, &mut ctx); let bp_arr = bp.into_array(); let seq = Sequence::try_new_typed(0u32, 1u32, Nullability::NonNullable, 2048)?; let seq_arr = seq.into_array(); @@ -1892,7 +1841,7 @@ mod tests { let patched_bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(), 4, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, )?; assert!(patched_bp.patches().is_some(), "expected patches"); let patched_bp_arr = patched_bp.into_array(); @@ -1937,8 +1886,7 @@ mod tests { #[crate::test] async fn test_alp_f64_for_bitpacked() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // ALP(FoR(BitPacked)) with f64: same structure as the f32 test. let len = 3000; let exponents = Exponents { e: 2, f: 0 }; @@ -1950,11 +1898,7 @@ mod tests { let alp = alp_encode(float_prim.as_view(), Some(exponents), &mut ctx)?; assert!(alp.patches().is_none()); let for_arr = FoR::encode(alp.encoded().clone().execute::(&mut ctx)?)?; - let bp = BitPacked::encode( - for_arr.encoded(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(for_arr.encoded(), 6, &mut ctx)?; let tree = ALP::new( FoR::try_new(bp.into_array(), for_arr.reference_scalar().clone())?.into_array(), @@ -1966,15 +1910,15 @@ mod tests { // CPU decode as ground truth. let cpu = array .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())? + .execute::(&mut ctx)? .into_array(); // GPU decode. - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let gpu = CanonicalCudaExt::into_host(canonical).await?.into_array(); - vortex::array::assert_arrays_eq!(cpu, gpu, &mut cpu_ctx); + assert_arrays_eq!(cpu, gpu, &mut ctx); Ok(()) } @@ -1987,6 +1931,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let mut values: Vec = (0..len).map(|i| (i as f64) * 1.1).collect(); // Insert exception values that ALP can't encode. values[0] = 99.9; @@ -2000,12 +1945,7 @@ mod tests { } let float_prim = PrimitiveArray::new(Buffer::from(values), NonNullable); - let encoded = alp_encode( - float_prim.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - )? - .into_array(); + let encoded = alp_encode(float_prim.as_view(), None, &mut ctx)?.into_array(); let (array, base_offset) = if let Some(range) = &slice_range { (encoded.slice(range.clone())?, range.start) @@ -2014,12 +1954,10 @@ mod tests { }; // Decode on CPU as ground truth (accounts for ALP precision loss + patches). - let cpu_decoded = array - .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + let cpu_decoded = array.clone().execute::(&mut ctx)?; let expected: Vec = cpu_decoded.as_slice::().to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?; let result_prim = result.as_primitive(); @@ -2038,9 +1976,9 @@ mod tests { #[crate::test] async fn alp_slice_device_patches() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); // Regression test for https://github.com/vortex-data/vortex/issues/7838#issuecomment-4452796116. - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 4096; let exponents = Exponents { e: 0, f: 0 }; @@ -2094,21 +2032,21 @@ mod tests { expected[2048 - 100] = std::f64::consts::LN_2; let expected = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected, gpu, &mut cpu_ctx); + assert_arrays_eq!(expected, gpu, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_runend_u32_ends_u16_values() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); // RunEnd with u32 ends, u16 values. Output type = u16. // Ends (u32) differ from output (u16) → pending subtree. let ends: Vec = vec![500, 1000, 1500, 2000]; let values: Vec = vec![100, 200, 300, 400]; let len = 2000; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, cuda_ctx.execution_ctx()); @@ -2139,14 +2077,14 @@ mod tests { }) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } #[crate::test] async fn test_dict_bitpacked_u8_codes_u32_values() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); // Dict with BitPacked u8 codes (narrower than u32 output) and u32 values. // The kernel's bitunpack_typed decodes at the source's native width and // widens to T, so this fuses into a single kernel launch. @@ -2156,12 +2094,8 @@ mod tests { let codes_prim = PrimitiveArray::new(Buffer::from(codes.clone()), NonNullable); // BitPack the u8 codes at 2 bits (4 values need 2 bits) - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - 2, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack codes"); + let codes_bp = BitPacked::encode(&codes_prim.into_array(), 2, &mut cpu_ctx) + .vortex_expect("bitpack codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); let dict = DictArray::try_new(codes_bp.into_array(), values_prim.into_array())?; let array = dict.into_array(); @@ -2172,13 +2106,13 @@ mod tests { "expected Fused for mixed-width Dict with BitPacked codes" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2196,17 +2130,13 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes.clone()), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack codes"); + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); let dict = DictArray::try_new(codes_bp.into_array(), values_prim.into_array())?; let array = dict.into_array(); @@ -2217,13 +2147,13 @@ mod tests { "expected Fused for mixed-width Dict with BitPacked codes" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2240,17 +2170,13 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes.clone()), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack codes"); + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut cpu_ctx) + .vortex_expect("bitpack codes"); let codes_for = FoR::try_new(codes_bp.into_array(), Scalar::from(0u8))?; let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); let dict = DictArray::try_new(codes_for.into_array(), values_prim.into_array())?; @@ -2262,13 +2188,13 @@ mod tests { "expected Fused for mixed-width Dict with FoR(BitPacked) codes" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } @@ -2294,20 +2220,15 @@ mod tests { #[case] ends: Vec, #[case] values: Vec, ) -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let ends_u64: Vec = ends.iter().map(|e| (*e).into()).collect(); let len = *ends_u64.last().unwrap() as usize; let bit_width = 64 - ends_u64.iter().max().unwrap().leading_zeros() as u8; let ends_prim = PrimitiveArray::new(Buffer::from(ends.clone()), NonNullable); - let ends_bp = BitPacked::encode( - &ends_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack ends"); + let ends_bp = BitPacked::encode(&ends_prim.into_array(), bit_width, &mut ctx) + .vortex_expect("bitpack ends"); let values_prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); let re = RunEnd::new(ends_bp.into_array(), values_prim.into_array(), &mut ctx); let array = re.into_array(); @@ -2317,7 +2238,7 @@ mod tests { "expected Fused for mixed-width RunEnd with BitPacked ends" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); @@ -2329,7 +2250,7 @@ mod tests { prev = *end; } let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } @@ -2339,22 +2260,17 @@ mod tests { /// at native width and widens to T, fusing everything. #[crate::test] async fn test_runend_mixed_width_for_bp_u16_ends_u32_values() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let ends: Vec = vec![500, 1000, 1500, 2000]; let values: Vec = vec![100, 200, 300, 400]; let len = 2000usize; let ends_prim = PrimitiveArray::new(Buffer::from(ends.clone()), NonNullable); - let ends_bp = BitPacked::encode( - &ends_prim.into_array(), - 11, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .vortex_expect("bitpack ends"); + let ends_bp = + BitPacked::encode(&ends_prim.into_array(), 11, &mut ctx).vortex_expect("bitpack ends"); let ends_for = FoR::try_new(ends_bp.into_array(), Scalar::from(0u16))?; let values_prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); let re = RunEnd::new(ends_for.into_array(), values_prim.into_array(), &mut ctx); let array = re.into_array(); @@ -2364,7 +2280,7 @@ mod tests { "expected Fused for mixed-width RunEnd with FoR(BitPacked) ends" ); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&array, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); @@ -2375,14 +2291,14 @@ mod tests { prev = end; } let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } #[crate::test] async fn test_sliced_dict_mixed_width() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); // Sliced Dict with u8 codes and u32 values — combines PartiallyFused + slice handling. let dict_values: Vec = vec![100, 200, 300, 400]; let full_len = 4096; @@ -2397,7 +2313,7 @@ mod tests { // Slice from 1000..3000 let sliced = dict.into_array().slice(1000..3000)?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&sliced, &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); @@ -2406,7 +2322,7 @@ mod tests { .map(|&c| dict_values[c as usize]) .collect(); let expected_arr = PrimitiveArray::new(Buffer::from(expected), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } @@ -2416,7 +2332,7 @@ mod tests { /// (the bit-pattern for i32(-1)), not u32(0x000000FF) = 255. #[crate::test] fn test_load_element_sign_extends_i8_to_u32() -> VortexResult<()> { - let cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let i8_values: Vec = vec![-1, -2, -3, 127, -128, 0, 1, 42]; let len = i8_values.len(); @@ -2454,7 +2370,7 @@ mod tests { /// Same as above but for i16 → u32 widening. #[crate::test] fn test_load_element_sign_extends_i16_to_u32() -> VortexResult<()> { - let cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let i16_values: Vec = vec![-1, -256, -32768, 32767, 0, 1, -100, 12345]; let len = i16_values.len(); @@ -2493,8 +2409,8 @@ mod tests { /// Nullable Primitive array — LOAD source with validity propagated. #[crate::test] async fn test_nullable_primitive() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let array = PrimitiveArray::from_option_iter( (0..2048u32).map(|i| if i % 3 == 0 { None } else { Some(i) }), @@ -2506,7 +2422,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(array, gpu, &mut cpu_ctx); + assert_arrays_eq!(array, gpu, &mut cpu_ctx); Ok(()) } @@ -2516,8 +2432,8 @@ mod tests { /// validity, so this produces a real nullable FoR(BitPacked) tree. #[crate::test] async fn test_nullable_for_bitpacked() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 2048; let reference = 1000u32; @@ -2541,11 +2457,7 @@ mod tests { PrimitiveArray::from_option_iter(values.iter().map(|v| v.map(|x| x - reference))); // BitPacked::encode preserves nullable validity from the input. - let bp = BitPacked::encode( - &residuals.into_array(), - 6, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&residuals.into_array(), 6, &mut cpu_ctx)?; let for_arr = FoR::try_new(bp.into_array(), reference.into())?; // Verify the plan actually fuses (not just a LOAD). @@ -2563,14 +2475,14 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(prim, gpu, &mut cpu_ctx); + assert_arrays_eq!(prim, gpu, &mut cpu_ctx); Ok(()) } /// AllInvalid array — kernel should be skipped entirely. #[crate::test] async fn test_all_invalid_skips_kernel() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let array = PrimitiveArray::new(Buffer::from(vec![0u32; 2048]), Validity::AllInvalid); @@ -2588,8 +2500,8 @@ mod tests { /// AllValid nullable array — should fuse and produce AllValid output. #[crate::test] async fn test_all_valid_nullable() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let values: Vec = (0..2048).collect(); let array = PrimitiveArray::new(Buffer::from(values.clone()), Validity::AllValid); @@ -2600,7 +2512,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(array, gpu, &mut cpu_ctx); + assert_arrays_eq!(array, gpu, &mut cpu_ctx); Ok(()) } @@ -2624,10 +2536,8 @@ mod tests { /// Dict with non-nullable codes but nullable values should still fuse. #[crate::test] async fn test_dict_nullable_values_fuses() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); - use vortex::buffer::buffer; - - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let codes = PrimitiveArray::new(buffer![0u32, 1, 2, 2, 1, 0], NonNullable); let values = PrimitiveArray::from_option_iter([Some(10u32), None, Some(30)]); @@ -2642,7 +2552,7 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(dict, gpu, &mut cpu_ctx); + assert_arrays_eq!(dict, gpu, &mut cpu_ctx); Ok(()) } @@ -2650,11 +2560,8 @@ mod tests { /// Validity must survive through fused dispatch and into the filter. #[crate::test] async fn test_nullable_fused_then_filter() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); - use vortex::array::arrays::FilterArray; - use vortex::mask::Mask; - - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 2048usize; let values: Vec> = (0..len) @@ -2681,14 +2588,14 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(filter_array, gpu, &mut cpu_ctx); + assert_arrays_eq!(filter_array, gpu, &mut cpu_ctx); Ok(()) } /// Empty nullable array should preserve nullability. #[crate::test] async fn test_empty_nullable_array() -> VortexResult<()> { - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let array = PrimitiveArray::new(Buffer::::empty(), Validity::AllValid); let result = try_gpu_dispatch(&array.into_array(), &mut cuda_ctx).await?; @@ -2704,6 +2611,7 @@ mod tests { #[crate::test] async fn test_bitpacked_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 3000; let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; @@ -2718,16 +2626,12 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let array = bp.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&array, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( &cuda_ctx, @@ -2748,6 +2652,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; let values: Vec = (0..len) @@ -2761,11 +2666,7 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let (array, expected) = if let Some(range) = slice_range { @@ -2775,7 +2676,7 @@ mod tests { (bp.into_array(), values) }; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&array, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( &cuda_ctx, @@ -2789,6 +2690,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 3000; let bit_width: u8 = 6; let reference = 42u32; @@ -2805,17 +2707,13 @@ mod tests { let all_values: Vec = residuals.iter().map(|&v| v + reference).collect(); let prim = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let for_arr = FoR::try_new(bp.into_array(), Scalar::from(reference))?; let array = for_arr.into_array(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&array, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( &cuda_ctx, @@ -2829,6 +2727,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_with_patches_sliced() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 5000; let bit_width: u8 = 6; let reference = 42u32; @@ -2845,11 +2744,7 @@ mod tests { let all_values: Vec = residuals.iter().map(|&v| v + reference).collect(); let prim = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let for_arr = FoR::try_new(bp.into_array(), Scalar::from(reference))?; @@ -2857,7 +2752,7 @@ mod tests { let sliced = for_arr.into_array().slice(range.clone())?; let expected = all_values[range].to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&sliced, &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan( &cuda_ctx, @@ -2878,7 +2773,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut values: Vec = (0..len).map(|i| (i as f32) * 1.1).collect(); // Insert exception values that ALP can't encode. values[0] = 99.9; @@ -2904,7 +2799,7 @@ mod tests { let cpu_decoded = array.clone().execute::(&mut ctx)?; let expected: Vec = cpu_decoded.as_slice::().to_vec(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&array, &mut cuda_ctx).await?; let actual = run_dispatch_plan_f32( &cuda_ctx, @@ -2931,7 +2826,7 @@ mod tests { /// u8 BitPacked with patches (bit_width=3, patch values > 7). #[crate::test] async fn test_bitpacked_with_patches_u8() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 3; let len = 3000usize; let max_val = (1u8 << bit_width) - 1; @@ -2946,26 +2841,22 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&bp.into_array(), &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut ctx); Ok(()) } /// u16 BitPacked with patches (bit_width=6, patch values > 63). #[crate::test] async fn test_bitpacked_with_patches_u16() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 6; let len = 3000usize; let max_val = (1u16 << bit_width) - 1; @@ -2980,26 +2871,22 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut cpu_ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&bp.into_array(), &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } /// u64 BitPacked with patches (bit_width=4, patch values > 15). #[crate::test] async fn test_bitpacked_with_patches_u64() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let bit_width: u8 = 4; let len = 3000usize; let max_val = (1u64 << bit_width) - 1; @@ -3014,25 +2901,22 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut cpu_ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let canonical = try_gpu_dispatch(&bp.into_array(), &mut cuda_ctx).await?; let result = CanonicalCudaExt::into_host(canonical).await?.into_array(); let expected_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - vortex::array::assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); + assert_arrays_eq!(expected_arr, result, &mut cpu_ctx); Ok(()) } /// Dict where codes are BitPacked u32 with patches exceeding the bit width. #[crate::test] async fn test_dict_bitpacked_codes_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let dict_values: Vec = (0..256).map(|i| i * 1000 + 42).collect(); let len = 3000; let bit_width: u8 = 4; @@ -3050,17 +2934,13 @@ mod tests { let expected: Vec = codes.iter().map(|&c| dict_values[c as usize]).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let codes_bp = BitPacked::encode( - &codes_prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let codes_bp = BitPacked::encode(&codes_prim.into_array(), bit_width, &mut ctx)?; assert!(codes_bp.patches().is_some(), "expected patches on codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values), NonNullable); let dict = DictArray::try_new(codes_bp.into_array(), values_prim.into_array())?; - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&dict.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan(&cuda_ctx, len, &plan.dispatch_plan, plan.shared_mem_bytes)?; @@ -3071,6 +2951,7 @@ mod tests { /// Patches placed exactly at FastLanes chunk boundaries (1024-element chunks). #[crate::test] async fn test_bitpacked_patches_at_chunk_boundaries() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 4096usize; let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; @@ -3082,14 +2963,10 @@ mod tests { values[2048] = 4000; // start of chunk 2 let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&bp.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan(&cuda_ctx, len, &plan.dispatch_plan, plan.shared_mem_bytes)?; @@ -3100,6 +2977,7 @@ mod tests { /// Large array (100k elements) spanning many blocks with sparse patches. #[crate::test] async fn test_bitpacked_large_array_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let len = 100_000usize; let bit_width: u8 = 6; let max_val = (1u32 << bit_width) - 1; @@ -3114,14 +2992,10 @@ mod tests { .collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&bp.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan(&cuda_ctx, len, &plan.dispatch_plan, plan.shared_mem_bytes)?; @@ -3133,8 +3007,8 @@ mod tests { /// dispatch alongside patch application. #[crate::test] async fn test_nullable_bitpacked_with_patches() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 3000usize; let bit_width: u8 = 4; @@ -3155,11 +3029,7 @@ mod tests { let prim = PrimitiveArray::from_option_iter(values.iter().copied()); - let bp = BitPacked::encode( - &prim.clone().into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.clone().into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); let gpu = try_gpu_dispatch(&bp.into_array(), &mut cuda_ctx) @@ -3168,13 +3038,14 @@ mod tests { .await? .into_array(); - vortex::array::assert_arrays_eq!(prim, gpu, &mut cpu_ctx); + assert_arrays_eq!(prim, gpu, &mut ctx); Ok(()) } /// Extreme case: ALL values are patches (bit_width=1, every value > 1). #[crate::test] async fn test_bitpacked_all_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let bit_width: u8 = 1; let len = 2000usize; // All values >= 2, so every single element exceeds max storable (1) and @@ -3182,14 +3053,10 @@ mod tests { let values: Vec = (0..len).map(|i| (i as u32) + 2).collect(); let prim = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); - let bp = BitPacked::encode( - &prim.into_array(), - bit_width, - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let bp = BitPacked::encode(&prim.into_array(), bit_width, &mut ctx)?; assert!(bp.patches().is_some(), "expected patches"); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let plan = dispatch_plan(&bp.into_array(), &mut cuda_ctx).await?; let actual = run_dynamic_dispatch_plan(&cuda_ctx, len, &plan.dispatch_plan, plan.shared_mem_bytes)?; diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index c9ae05b4467..b3a0151a727 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -171,18 +171,19 @@ mod tests { use vortex::error::VortexResult; use vortex::mask::Mask; use vortex::scalar::Scalar; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use crate::CanonicalCudaExt; use crate::executor::CudaArrayExt; use crate::session::CudaSession; fn for_bp>(values: Vec, reference: T) -> ArrayRef { + let mut ctx = array_session().create_execution_ctx(); let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(), 7, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut ctx, ) .vortex_expect("bp"); FoR::try_new(bp.into_array(), reference.into()) @@ -193,14 +194,14 @@ mod tests { /// FoR(BitPacked) u32 — entire tree compiles into a single fused plan. #[crate::test] async fn test_fused() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let values: Vec = (0..2048).map(|i| (i % 128) as u32).collect(); let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(), 7, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut cpu_ctx, ) .vortex_expect("bp"); let arr = FoR::try_new(bp.into_array(), 1000u32.into()).vortex_expect("for"); @@ -221,7 +222,7 @@ mod tests { /// Exercises the unsigned type reinterpretation in CudaDispatchPlan::execute. #[crate::test] async fn test_fused_f32() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); use vortex::encodings::alp::ALP; use vortex::encodings::alp::Exponents; @@ -231,7 +232,7 @@ mod tests { let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(encoded), NonNullable).into_array(), 9, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut cpu_ctx, ) .vortex_expect("bp"); let alp = ALP::try_new( @@ -257,7 +258,7 @@ mod tests { /// ALP with patches — plan builder rejects it, falls back to ALPExecutor. #[crate::test] async fn test_fallback() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); use vortex::array::patches::Patches; use vortex::array::validity::Validity::NonNullable as NN; use vortex::buffer::buffer; @@ -299,7 +300,7 @@ mod tests { #[cfg(feature = "unstable_encodings")] #[crate::test] async fn test_partial_fusion() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); use vortex::array::arrays::DictArray; use vortex::array::session::ArraySessionExt; use vortex::encodings::fastlanes; @@ -320,7 +321,7 @@ mod tests { ) .into_array(); let vals = FoR::try_new( - BitPacked::encode(&vals, 6, &mut LEGACY_SESSION.create_execution_ctx()) + BitPacked::encode(&vals, 6, &mut cpu_ctx) .vortex_expect("bp") .into_array(), 0u32.into(), @@ -335,7 +336,7 @@ mod tests { ) .into_array(); let codes = FoR::try_new( - BitPacked::encode(&codes, 6, &mut LEGACY_SESSION.create_execution_ctx()) + BitPacked::encode(&codes, 6, &mut cpu_ctx) .vortex_expect("bp") .into_array(), 0u32.into(), @@ -363,7 +364,7 @@ mod tests { /// Filter(FoR(BP), mask) — FoR+BP fuses via dyn dispatch, then CUB filters the result. #[crate::test] async fn test_filter_fused_child() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -372,7 +373,7 @@ mod tests { let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(data.clone()), NonNullable).into_array(), 7, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut cpu_ctx, ) .vortex_expect("bp"); let for_arr = FoR::try_new(bp.into_array(), 100u32.into()).vortex_expect("for"); @@ -407,7 +408,7 @@ mod tests { ))] #[crate::test] async fn test_ext_storage_gpu_decode(#[case] ext: ExtensionArray) -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -431,7 +432,7 @@ mod tests { /// Extension over already-canonical storage executes unchanged. #[crate::test] async fn test_ext_canonical_storage() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut cpu_ctx = array_session().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); diff --git a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs index 71f0af259c1..7de2158ac48 100644 --- a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs +++ b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs @@ -227,6 +227,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use super::*; use crate::CanonicalCudaExt; @@ -234,7 +235,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_buffers_decompression_primitive() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -247,13 +248,13 @@ mod tests { .into_host() .await?; - assert_arrays_eq!(compressed, gpu_result.into_array(), &mut cpu_ctx); + assert_arrays_eq!(compressed, gpu_result.into_array(), &mut ctx); Ok(()) } #[crate::test] async fn test_cuda_zstd_buffers_decompression_varbinview() -> VortexResult<()> { - let mut cpu_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -274,7 +275,7 @@ mod tests { .into_host() .await?; - assert_arrays_eq!(compressed, gpu_result.into_array(), &mut cpu_ctx); + assert_arrays_eq!(compressed, gpu_result.into_array(), &mut ctx); Ok(()) } } diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 19475e17065..5ee6b1e2541 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -403,6 +403,7 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::error::VortexExpect; use vortex::mask::Mask; + use vortex_array::array_session; use vortex_geo::extension::WellKnownBinaryData; use wkb::writer::WriteOptions; use wkb::writer::write_point; @@ -415,7 +416,7 @@ mod tests { #[test] fn test_integer_vector_conversion() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let values = vec![1i32, 2, 3, 4, 5]; let len = values.len(); @@ -432,7 +433,7 @@ mod tests { let result = flat_vector_to_vortex(&vector, len).unwrap(); let expected = PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4), Some(5)]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[test] diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 139faa7aaa7..619f82e3f3e 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -454,6 +454,7 @@ mod tests { use vortex::array::arrays::StructArray; use vortex::session::VortexSession; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; @@ -532,19 +533,19 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_no_options() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let (array, struct_array) = scan(ptr::null()); - assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); + assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut ctx); unsafe { vx_array_free(array) }; } #[test] #[cfg_attr(miri, ignore)] fn test_project_all() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let opts = vx_scan_options::default(); let (array, struct_array) = scan(&raw const opts); - assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); + assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut ctx); unsafe { vx_array_free(array) }; } @@ -552,7 +553,7 @@ mod tests { #[cfg_attr(miri, ignore)] fn test_project_single_field() { unsafe { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let root = vx_expression_root(); let mut opts = vx_scan_options::default(); @@ -564,7 +565,7 @@ mod tests { assert_arrays_eq!( vx_array::as_ref(array), struct_array.unmasked_field_by_name(field).unwrap(), - &mut assertion_ctx + &mut ctx ); vx_array_free(array); vx_expression_free(field_expr); @@ -714,13 +715,13 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_ordered() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let opts = vx_scan_options { ordered: true, ..Default::default() }; let (array, struct_array) = scan(&raw const opts); - assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut assertion_ctx); + assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut ctx); unsafe { vx_array_free(array) }; } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 8eeea97fff2..295fef36039 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -14,6 +14,7 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::accessor::ArrayAccessor; +use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; @@ -76,7 +77,7 @@ use crate::VortexFile; use crate::WriteOptionsSessionExt; use crate::footer::SegmentSpec; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session() + let session = array_session() .with::() .with::(); @@ -477,7 +478,6 @@ async fn test_empty_varbin_array_roundtrip() { #[tokio::test] #[cfg_attr(miri, ignore)] async fn issue_5385_filter_casted_column() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = StructArray::try_from_iter([("x", buffer![1u8, 2, 3, 4, 5])]) .unwrap() .into_array(); @@ -511,7 +511,7 @@ async fn issue_5385_filter_casted_column() { assert_arrays_eq!( result, StructArray::try_from_iter([("x", buffer![1u8])]).unwrap(), - &mut assertion_ctx + &mut SESSION.create_execution_ctx() ); } @@ -1187,12 +1187,11 @@ async fn chunked_file() -> VortexResult { #[tokio::test] async fn basic_file_roundtrip() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let vxf = chunked_file().await?; let result = vxf.scan()?.into_array_stream()?.read_all().await?; let expected = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } @@ -1231,7 +1230,6 @@ async fn file_excluding_dtype() -> VortexResult<()> { #[tokio::test] async fn file_take() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let vxf = chunked_file().await?; let result = vxf .scan()? @@ -1241,7 +1239,7 @@ async fn file_take() -> VortexResult<()> { .await?; let expected = buffer![0i32, 1, 8].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); Ok(()) } diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index 1dbe3a80d32..94f9519fd59 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -123,7 +123,6 @@ async fn test_file_roundtrip() { /// with "Array vortex.fill_null does not support serialization". #[tokio::test] async fn test_dict_listview_validity_roundtrip() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]).into_array(); let offsets = PrimitiveArray::from_iter(vec![0u32, 2, 4]).into_array(); let sizes = PrimitiveArray::from_iter(vec![2u32, 2, 1]).into_array(); @@ -164,6 +163,6 @@ async fn test_dict_listview_validity_roundtrip() { .await .unwrap() .expect("read back should succeed"); - vortex_array::assert_arrays_eq!(data, chunk, &mut assertion_ctx); + vortex_array::assert_arrays_eq!(data, chunk, &mut SESSION.create_execution_ctx()); assert!(stream.next().await.is_none(), "expected a single chunk"); } diff --git a/vortex-ipc/src/iterator.rs b/vortex-ipc/src/iterator.rs index 7f8adfef986..64f52b4bbf8 100644 --- a/vortex-ipc/src/iterator.rs +++ b/vortex-ipc/src/iterator.rs @@ -183,7 +183,6 @@ mod test { #[test] fn test_sync_stream() -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3].into_array(); let ipc_buffer = array .to_array_iterator() @@ -194,7 +193,7 @@ mod test { assert_eq!(reader.dtype(), array.dtype()); let result = reader.read_all()?; - assert_arrays_eq!(result, array, &mut assertion_ctx); + assert_arrays_eq!(result, array, &mut SESSION.create_execution_ctx()); Ok(()) } diff --git a/vortex-ipc/src/stream.rs b/vortex-ipc/src/stream.rs index 43e170e969f..d8da4935115 100644 --- a/vortex-ipc/src/stream.rs +++ b/vortex-ipc/src/stream.rs @@ -236,7 +236,6 @@ mod test { #[tokio::test] async fn test_async_stream() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = buffer![1, 2, 3].into_array(); let ipc_buffer = array .to_array_stream() @@ -251,7 +250,7 @@ mod test { assert_eq!(reader.dtype(), array.dtype()); let result = reader.read_all().await.unwrap(); - assert_arrays_eq!(result, array, &mut assertion_ctx); + assert_arrays_eq!(result, array, &mut SESSION.create_execution_ctx()); } /// Wrapper that limits reads to small chunks to simulate network behavior @@ -273,7 +272,6 @@ mod test { #[tokio::test] async fn test_async_stream_chunked() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(); let ipc_buffer = array .to_array_stream() @@ -291,13 +289,12 @@ mod test { let result = reader.read_all().await.unwrap(); let expected = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } /// Test with 1-byte chunks to stress-test partial read handling. #[tokio::test] async fn test_async_stream_single_byte_chunks() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = buffer![42i64, -1, 0, i64::MAX, i64::MIN].into_array(); let ipc_buffer = array .to_array_stream() @@ -315,6 +312,6 @@ mod test { let result = reader.read_all().await.unwrap(); let expected = buffer![42i64, -1, 0, i64::MAX, i64::MIN].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx()); } } diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index aee6d6575b0..a9870229792 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -472,7 +472,7 @@ mod test { #[from(chunked_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|_h| async { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let result = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() @@ -486,7 +486,7 @@ mod test { .unwrap(); let expected = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); }) } } diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 87d4197bfa4..a4b3e20cbff 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -337,9 +337,9 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray as _; - use vortex_array::LEGACY_SESSION; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinArray; @@ -386,7 +386,7 @@ mod tests { // FIXME(ngates): Deprecate the global `runtime::single::block_on` helper and require tests // to call `block_on` on an explicit runtime instance. fn session_with_handle(handle: Handle) -> VortexSession { - vortex_array::array_session() + array_session() .with::() .with::() .with_handle(handle) @@ -422,8 +422,8 @@ mod tests { #[test] fn reading_nested_packs_works() { block_on(|handle| async move { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = session_with_handle(handle); + let mut ctx = session.create_execution_ctx(); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), @@ -447,12 +447,12 @@ mod tests { ) .into_array(); let array_to_write = array.clone(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let layout: LayoutRef = strategy .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), SequentialStreamAdapter::new( DType::Utf8(Nullability::Nullable), @@ -501,7 +501,7 @@ mod tests { ) .unwrap() .into_array(); - assert_arrays_eq!(actual, expected, &mut assertion_ctx); + assert_arrays_eq!(actual, expected, &mut ctx); }) } @@ -522,8 +522,8 @@ mod tests { #[case] expected: Vec, ) { block_on(|handle| async move { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = session_with_handle(handle); + let mut ctx = session.create_execution_ctx(); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), @@ -533,12 +533,12 @@ mod tests { let array = VarBinArray::from_iter(data, DType::Utf8(Nullability::Nullable)).into_array(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let layout: LayoutRef = strategy .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), SequentialStreamAdapter::new( DType::Utf8(Nullability::Nullable), @@ -566,20 +566,15 @@ mod tests { .await .unwrap(); - assert_arrays_eq!( - mask.into_array(), - BoolArray::from_iter(expected), - &mut assertion_ctx - ); + assert_arrays_eq!(mask.into_array(), BoolArray::from_iter(expected), &mut ctx); }) } #[test] fn reading_is_null_works() { block_on(|handle| async move { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); - let mut ctx_exec = LEGACY_SESSION.create_execution_ctx(); let session = session_with_handle(handle); + let mut ctx = session.create_execution_ctx(); let strategy = DictStrategy::new( FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), @@ -603,13 +598,13 @@ mod tests { ) .into_array(); let array_to_write = array.clone(); - let ctx = ArrayContext::empty(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let layout: LayoutRef = strategy .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), SequentialStreamAdapter::new( DType::Utf8(Nullability::Nullable), @@ -638,20 +633,19 @@ mod tests { let expected = array .validity() .unwrap() - .execute_mask(array.len(), &mut ctx_exec) + .execute_mask(array.len(), &mut ctx) .unwrap() .into_array(); let actual_canonical = actual - .execute::(&mut ctx_exec) + .execute::(&mut ctx) .vortex_expect("to_canonical failed") .into_array(); - assert_arrays_eq!(actual_canonical, expected, &mut assertion_ctx); + assert_arrays_eq!(actual_canonical, expected, &mut ctx); }) } #[test] fn reading_byte_length_pushdown_works() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let array = VarBinArray::from_iter( [ Some("abc"), @@ -676,6 +670,7 @@ mod tests { block_on(|handle| async move { let session = session_with_handle(handle); + let mut ctx = session.create_execution_ctx(); let (layout, segments) = write_dict_layout(array, &session).await; assert_eq!(layout.encoding_id(), LayoutId::new("vortex.dict")); let actual = layout @@ -690,7 +685,7 @@ mod tests { .await .unwrap() .into_array(); - assert_arrays_eq!(actual, expected, &mut assertion_ctx); + assert_arrays_eq!(actual, expected, &mut ctx); }) } @@ -709,7 +704,7 @@ mod tests { } fn test_apply(original: Expression, outer: Expression, inner: Expression) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = VarBinArray::from_iter( [Some("abc"), Some("def"), None], DType::Utf8(Nullability::Nullable), @@ -722,7 +717,7 @@ mod tests { ))?; let actual = pushed.apply(&outer)?; let expected = array.apply(&original)?; - assert_arrays_eq!(actual, expected, &mut assertion_ctx); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index a0269786d51..50ec634a61a 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -256,16 +256,16 @@ mod test { #[test] fn flat_identity() -> VortexResult<()> { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); - let ctx = ArrayContext::empty(); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let array = PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid).into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -279,7 +279,7 @@ mod test { ); let result = layout - .new_reader("".into(), segments, &SESSION, &Default::default())? + .new_reader("".into(), segments, &session, &Default::default())? .projection_evaluation( &(0..layout.row_count()), &root(), @@ -287,7 +287,7 @@ mod test { )? .await?; - assert_arrays_eq!(result, array, &mut assertion_ctx); + assert_arrays_eq!(result, array, &mut ctx); Ok(()) }) @@ -296,9 +296,9 @@ mod test { #[test] fn flat_expr() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); - let ctx = ArrayContext::empty(); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -306,7 +306,7 @@ mod test { PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid).into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -317,7 +317,7 @@ mod test { let expr = gt(root(), lit(3i32)); let result = layout - .new_reader("".into(), segments, &SESSION, &Default::default()) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), @@ -329,23 +329,23 @@ mod test { .unwrap(); let expected = BoolArray::from_iter([false, false, false, true, true].map(Some)); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); }) } #[test] fn flat_unaligned_row_mask() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); - let ctx = ArrayContext::empty(); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let array = PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid).into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -355,7 +355,7 @@ mod test { .unwrap(); let result = layout - .new_reader("".into(), segments, &SESSION, &Default::default()) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap() .projection_evaluation(&(2..4), &root(), MaskFuture::new_true(2)) .unwrap() @@ -363,7 +363,7 @@ mod test { .unwrap(); let expected = PrimitiveArray::new(buffer![3i32, 4], Validity::AllValid).into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); }) } } diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 19e8ac81376..4e5e8e494fa 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -352,15 +352,15 @@ mod tests { #[test] fn flat_expr_no_row_id() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); - let ctx = ArrayContext::empty(); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let array = buffer![1..=5].into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -373,9 +373,9 @@ mod tests { let result = RowIdxLayoutReader::new( 0, layout - .new_reader("".into(), segments, &SESSION, &Default::default()) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap(), - SESSION.clone(), + session.clone(), ) .projection_evaluation( &(0..layout.row_count()), @@ -389,7 +389,7 @@ mod tests { assert_arrays_eq!( result, BoolArray::from_iter([false, false, true, false, false]), - &mut assertion_ctx + &mut ctx ); }) } @@ -397,15 +397,15 @@ mod tests { #[test] fn flat_expr_row_id() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); - let ctx = ArrayContext::empty(); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let array = buffer![1..=5].into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -418,9 +418,9 @@ mod tests { let result = RowIdxLayoutReader::new( 0, layout - .new_reader("".into(), segments, &SESSION, &Default::default()) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap(), - SESSION.clone(), + session.clone(), ) .projection_evaluation( &(0..layout.row_count()), @@ -434,7 +434,7 @@ mod tests { assert_arrays_eq!( result, BoolArray::from_iter([false, false, false, false, true]), - &mut assertion_ctx + &mut ctx ); }) } @@ -442,15 +442,15 @@ mod tests { #[test] fn flat_expr_or() { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); let session = SESSION.clone().with_handle(handle); - let ctx = ArrayContext::empty(); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let array = buffer![1..=5].into_array(); let layout = FlatLayoutStrategy::default() .write_stream( - ctx, + array_ctx, Arc::::clone(&segments), array.to_array_stream().sequenced(ptr), eof, @@ -467,9 +467,9 @@ mod tests { let result = RowIdxLayoutReader::new( 0, layout - .new_reader("".into(), segments, &SESSION, &Default::default()) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap(), - SESSION.clone(), + session.clone(), ) .projection_evaluation( &(0..layout.row_count()), @@ -483,7 +483,7 @@ mod tests { assert_arrays_eq!( result, BoolArray::from_iter([true, false, true, false, true]), - &mut assertion_ctx + &mut ctx ); }) } diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 3310d9eb834..f44c9178413 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -634,7 +634,7 @@ mod tests { fn test_struct_layout( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -646,14 +646,14 @@ mod tests { }) .unwrap(); let expected = BoolArray::from_iter([true, false, false]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[rstest] fn test_struct_layout_row_mask( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -670,7 +670,7 @@ mod tests { .unwrap(); let expected = BoolArray::from_iter([true, false]); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); } #[rstest] @@ -720,7 +720,7 @@ mod tests { fn test_struct_layout_nulls( #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Read the layout source from the top. let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) @@ -744,8 +744,8 @@ mod tests { .unwrap(), Scalar::null(result.dtype().clone()), ); - assert_nth_scalar!(result, 1, 2, &mut assertion_ctx); - assert_nth_scalar!(result, 2, 3, &mut assertion_ctx); + assert_nth_scalar!(result, 1, 2, &mut ctx); + assert_nth_scalar!(result, 2, 3, &mut ctx); } #[rstest] diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 0e4fc623715..8feaf46fa03 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -227,6 +227,7 @@ mod test { use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; @@ -266,7 +267,7 @@ mod test { use crate::session::LayoutSession; fn session_with_handle(handle: Handle) -> VortexSession { - vortex_array::array_session() + array_session() .with::() .with::() .with_handle(handle) @@ -310,7 +311,7 @@ mod test { #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|handle| async { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let session = session_with_handle(handle); let result = layout .new_reader("".into(), segments, &session, &Default::default()) @@ -325,7 +326,7 @@ mod test { .unwrap(); let expected = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); - assert_arrays_eq!(result, expected, &mut assertion_ctx); + assert_arrays_eq!(result, expected, &mut ctx); }) } @@ -431,7 +432,7 @@ mod test { let zoned_layout = layout.as_::(); let children = OwnedLayoutChildren::layout_children(vec![layout.child(0)?, layout.child(1)?]); - let session = vortex_array::array_session(); + let session = array_session(); let read_ctx = ReadContext::new([]); let build_ctx = LayoutBuildContext { session: &session, diff --git a/vortex-python/src/arrays/range_to_sequence.rs b/vortex-python/src/arrays/range_to_sequence.rs index d83662b4387..eed03c64d13 100644 --- a/vortex-python/src/arrays/range_to_sequence.rs +++ b/vortex-python/src/arrays/range_to_sequence.rs @@ -79,6 +79,7 @@ mod test { use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use crate::arrays::range_to_sequence::range_len; use crate::arrays::range_to_sequence::sequence_array_from_range; @@ -99,34 +100,26 @@ mod test { #[test] fn test_sequence_array_from_len() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::U16, Nullability::NonNullable); let arr = sequence_array_from_range::(0, 10, 1, dtype).unwrap(); assert_arrays_eq!( arr, buffer![0u16, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(), - &mut assertion_ctx + &mut ctx ); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let arr = sequence_array_from_range::(0, 10, 5, dtype).unwrap(); - assert_arrays_eq!(arr, buffer![0i32, 5].into_array(), &mut assertion_ctx); + assert_arrays_eq!(arr, buffer![0i32, 5].into_array(), &mut ctx); let dtype = DType::Primitive(PType::I8, Nullability::NonNullable); let arr = sequence_array_from_range::(-5, 5, 3, dtype).unwrap(); - assert_arrays_eq!( - arr, - buffer![-5i8, -2, 1, 4].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(arr, buffer![-5i8, -2, 1, 4].into_array(), &mut ctx); let dtype = DType::Primitive(PType::I8, Nullability::NonNullable); let arr = sequence_array_from_range::(3, -3, -1, dtype).unwrap(); - assert_arrays_eq!( - arr, - buffer![3i8, 2, 1, 0, -1, -2].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(arr, buffer![3i8, 2, 1, 0, -1, -2].into_array(), &mut ctx); let dtype = DType::Primitive(PType::U32, Nullability::NonNullable); let result = sequence_array_from_range::(1_000_000, 10, -500_000, dtype); @@ -136,10 +129,6 @@ mod test { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let arr = sequence_array_from_range::(1_000_000, 10, -500_000, dtype).unwrap(); - assert_arrays_eq!( - arr, - buffer![1_000_000i32, 500_000].into_array(), - &mut assertion_ctx - ); + assert_arrays_eq!(arr, buffer![1_000_000i32, 500_000].into_array(), &mut ctx); } } diff --git a/vortex-test/compat-gen/src/check.rs b/vortex-test/compat-gen/src/check.rs index c788ed17cb4..b458ab7fe9b 100644 --- a/vortex-test/compat-gen/src/check.rs +++ b/vortex-test/compat-gen/src/check.rs @@ -6,6 +6,7 @@ use std::path::Path; use clap::ValueEnum; use serde::Serialize; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::assert_arrays_eq; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -47,7 +48,7 @@ struct FailedFixture { /// Prints JSON result to stdout, human-readable progress to stderr. /// Returns error if any fixture failed or if mode constraints are violated. pub fn check(dir: &Path, mode: Mode, exclude: &[String]) -> VortexResult<()> { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let fixtures = all_fixtures(); let fixtures: Vec<_> = fixtures .into_iter() @@ -191,7 +192,7 @@ pub fn check(dir: &Path, mode: Mode, exclude: &[String]) -> VortexResult<()> { } }; - assert_arrays_eq!(stored_array, fresh_array, &mut assertion_ctx); + assert_arrays_eq!(stored_array, fresh_array, &mut ctx); eprintln!(" pass {fresh_name}"); result.passed.push(fresh_name.clone()); } diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs index 919ccfb39b7..374e753aa23 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs @@ -18,6 +18,7 @@ pub fn fixtures() -> Vec> { #[cfg(test)] mod tests { use vortex_array::VortexSessionExecute; + use vortex_array::array_session; use vortex_array::assert_arrays_eq; use super::fixtures; @@ -27,12 +28,12 @@ mod tests { #[test] fn roundtrip_fixtures_to_bytes() { for fixture in fixtures() { - let mut assertion_ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = array_session().create_execution_ctx(); let array = fixture.build().unwrap(); check_expected_encodings(&array, fixture.as_ref()).unwrap(); let bytes = adapter::write_file_to_bytes(array.clone()).unwrap(); let roundtripped = adapter::read_file(bytes).unwrap(); - assert_arrays_eq!(array, roundtripped, &mut assertion_ctx); + assert_arrays_eq!(array, roundtripped, &mut ctx); } } }