From 03c6d2fc70d1b0c3f597e033c87187fd63416a78 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 5 Jun 2026 11:49:29 -0400 Subject: [PATCH 1/2] pass execution context through list view rebuild Signed-off-by: Connor Tsui --- encodings/zstd/benches/listview_rebuild.rs | 20 +- vortex-array/benches/listview_rebuild.rs | 72 +++-- .../aggregate_fn/fns/all_non_distinct/list.rs | 4 +- .../uncompressed_size_in_bytes/list_view.rs | 2 +- .../src/arrays/chunked/vtable/canonical.rs | 98 +++---- .../src/arrays/constant/vtable/canonical.rs | 262 ++++++++++-------- vortex-array/src/arrays/list/array.rs | 15 +- vortex-array/src/arrays/list/tests.rs | 109 +++----- .../src/arrays/listview/conversion.rs | 116 +++++--- vortex-array/src/arrays/listview/rebuild.rs | 115 ++++---- .../src/arrays/listview/tests/common.rs | 9 + vortex-array/src/arrow/executor/list.rs | 20 +- vortex-array/src/arrow/executor/list_view.rs | 17 +- vortex-array/src/builders/listview.rs | 21 +- vortex-array/src/canonical.rs | 26 +- .../src/datasets/struct_list_of_ints.rs | 14 +- vortex-bench/src/vector_dataset/convert.rs | 5 +- vortex-compressor/src/compressor.rs | 8 +- vortex-cuda/src/arrow/canonical.rs | 3 +- vortex-duckdb/src/exporter/list_view.rs | 2 +- 20 files changed, 522 insertions(+), 416 deletions(-) diff --git a/encodings/zstd/benches/listview_rebuild.rs b/encodings/zstd/benches/listview_rebuild.rs index 4b8ca95eafc..d59af116188 100644 --- a/encodings/zstd/benches/listview_rebuild.rs +++ b/encodings/zstd/benches/listview_rebuild.rs @@ -3,18 +3,25 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ListViewArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewRebuildMode; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_session::VortexSession; use vortex_zstd::Zstd; use vortex_zstd::ZstdData; +/// A shared session for the `ListView` rebuild benchmark, used to create execution contexts. +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[divan::bench(sample_size = 1000)] fn rebuild_naive(bencher: Bencher) { let dudes = VarBinViewArray::from_iter_str(["Washington", "Adams", "Jefferson", "Madison"]) @@ -23,7 +30,7 @@ fn rebuild_naive(bencher: Bencher) { let validity = dudes.validity().unwrap(); let dudes = Zstd::try_new( dtype, - ZstdData::from_array(dudes, 9, 1024, &mut LEGACY_SESSION.create_execution_ctx()).unwrap(), + ZstdData::from_array(dudes, 9, 1024, &mut SESSION.create_execution_ctx()).unwrap(), validity, ) .unwrap() @@ -41,9 +48,12 @@ fn rebuild_naive(bencher: Bencher) { let list_view = ListViewArray::new(dudes, offsets, sizes, Validity::NonNullable); - bencher - .with_inputs(|| &list_view) - .bench_refs(|list_view| list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList)) + bencher.with_inputs(|| &list_view).bench_refs(|list_view| { + list_view.rebuild( + ListViewRebuildMode::MakeZeroCopyToList, + &mut SESSION.create_execution_ctx(), + ) + }) } fn main() { diff --git a/vortex-array/benches/listview_rebuild.rs b/vortex-array/benches/listview_rebuild.rs index 00781b1a534..f605a9d372c 100644 --- a/vortex-array/benches/listview_rebuild.rs +++ b/vortex-array/benches/listview_rebuild.rs @@ -6,10 +6,11 @@ #![expect(clippy::unwrap_used)] #![expect(clippy::cast_possible_truncation)] +use std::sync::LazyLock; + use divan::Bencher; use vortex_array::Canonical; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; @@ -20,13 +21,19 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArrayExt; use vortex_array::arrays::listview::ListViewRebuildMode; use vortex_array::dtype::FieldNames; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_session::VortexSession; fn main() { divan::main(); } +/// A shared session for the `ListView` rebuild benchmarks, used to create execution contexts. +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + fn make_primitive_lv(num_lists: usize, list_size: usize, step: usize) -> ListViewArray { let element_count = step * num_lists + list_size; let elements = PrimitiveArray::from_iter(0..element_count as i32).into_array(); @@ -111,11 +118,14 @@ fn make_nested_list_lv( fn i32_small(bencher: Bencher) { let lv = make_primitive_lv(50, 32, 32); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } @@ -124,11 +134,14 @@ fn i32_small(bencher: Bencher) { fn i32_small_overlapping(bencher: Bencher) { let lv = make_primitive_lv(50, 8, 1); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } @@ -137,11 +150,14 @@ fn i32_small_overlapping(bencher: Bencher) { fn varbinview_small(bencher: Bencher) { let lv = make_varbinview_lv(50, 32, 32); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } @@ -150,11 +166,14 @@ fn varbinview_small(bencher: Bencher) { fn struct_small(bencher: Bencher) { let lv = make_struct_lv(50, 32, 32); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } @@ -163,11 +182,14 @@ fn struct_small(bencher: Bencher) { fn i32_large(bencher: Bencher) { let lv = make_primitive_lv(50, 1_024, 1_024); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } @@ -176,11 +198,14 @@ fn i32_large(bencher: Bencher) { fn varbinview_large(bencher: Bencher) { let lv = make_varbinview_lv(5, 1_024, 1_024); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } @@ -189,11 +214,14 @@ fn varbinview_large(bencher: Bencher) { fn struct_large(bencher: Bencher) { let lv = make_struct_lv(25, 1_024, 1_024); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } @@ -215,11 +243,14 @@ fn fsl_large(bencher: Bencher) { Validity::NonNullable, ); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } @@ -228,11 +259,14 @@ fn fsl_large(bencher: Bencher) { fn list_i32_large(bencher: Bencher) { let lv = make_nested_list_lv(2, 512, 2); bencher.with_inputs(|| &lv).bench_refs(|lv| { - let rebuilt = lv.rebuild(ListViewRebuildMode::MakeZeroCopyToList).unwrap(); + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) + .unwrap(); rebuilt .elements() .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) + .execute::(&mut ctx) .unwrap() }); } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs index 9aa5e93d08c..2f4e3e52d91 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/list.rs @@ -28,8 +28,8 @@ pub(super) fn check_list_identical( return check_zero_copy_list_identical(lhs, rhs, ctx); } - let lhs = list_from_list_view(lhs.clone())?; - let rhs = list_from_list_view(rhs.clone())?; + let lhs = list_from_list_view(lhs.clone(), ctx)?; + let rhs = list_from_list_view(rhs.clone(), ctx)?; if !check_list_offsets_identical(&lhs, &rhs)? { return Ok(false); diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs index 919430195d9..f75fcb92025 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/list_view.rs @@ -18,7 +18,7 @@ pub(super) fn list_view_uncompressed_size_in_bytes( let mut size = if array.is_empty() { 0 } else { - let rebuilt = array.rebuild(ListViewRebuildMode::MakeExact)?; + let rebuilt = array.rebuild(ListViewRebuildMode::MakeExact, ctx)?; uncompressed_size_in_bytes_u64(rebuilt.elements(), ctx)? }; diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 6da34e3cb07..ea7e531992d 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -188,7 +188,7 @@ fn swizzle_list_chunks( let chunk_array = chunk.clone().execute::(ctx)?; // By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make // the final output `ListView` also zero-copyable to `List`. - let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact)?; + let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact, ctx)?; // Add the `elements` of the current array as a new chunk. list_elements_chunks.push(chunk_array.elements().clone()); @@ -285,6 +285,7 @@ fn swizzle_fixed_size_list_chunks( #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -298,15 +299,13 @@ mod tests { use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; - use crate::LEGACY_SESSION; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::accessor::ArrayAccessor; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; + use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; @@ -324,8 +323,13 @@ mod tests { use crate::memory::MemorySessionExt; use crate::memory::WritableHostBuffer; use crate::scalar::Scalar; + use crate::session::ArraySession; use crate::validity::Validity; + /// A shared session for these chunked-array tests, used to create execution contexts. + static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[derive(Debug)] struct CountingAllocator { allocations: Arc, @@ -375,7 +379,7 @@ mod tests { fn assert_variant_values(array: &VariantArray, expected: &[i32]) -> 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().copied().enumerate() { let scalar = array.execute_scalar(idx, &mut ctx)?; @@ -391,6 +395,7 @@ mod tests { #[test] fn pack_variant_chunks_without_shredded() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ variant_chunk([1, 2])?.into_array(), @@ -400,9 +405,7 @@ mod tests { )? .into_array(); - let variant = into_variant( - chunked.execute::(&mut LEGACY_SESSION.create_execution_ctx())?, - )?; + let variant = into_variant(chunked.execute::(&mut ctx)?)?; assert_eq!(variant.len(), 3); assert!(variant.shredded().is_none()); @@ -411,6 +414,7 @@ mod tests { #[test] fn pack_variant_chunks_all_shredded_same_dtype() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ variant_chunk_with_shredded( @@ -425,9 +429,7 @@ mod tests { )? .into_array(); - let variant = into_variant( - chunked.execute::(&mut LEGACY_SESSION.create_execution_ctx())?, - )?; + let variant = into_variant(chunked.execute::(&mut ctx)?)?; let shredded = variant .shredded() .ok_or_else(|| vortex_err!("expected shredded child"))?; @@ -436,15 +438,14 @@ mod tests { assert_eq!(shredded.len(), 3); assert_variant_values(&variant, &[10, 20, 30])?; - let shredded = shredded - .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + let shredded = shredded.clone().execute::(&mut ctx)?; assert_arrays_eq!(shredded, PrimitiveArray::from_iter([10i32, 20, 30])); Ok(()) } #[test] fn pack_variant_chunks_mixed_shredded_presence_errors() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ variant_chunk_with_shredded([1], PrimitiveArray::from_iter([10i32]).into_array())? @@ -455,9 +456,7 @@ mod tests { )? .into_array(); - let err = chunked - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) - .unwrap_err(); + let err = chunked.execute::(&mut ctx).unwrap_err(); assert!( err.to_string() .contains("chunks disagree on shredded presence") @@ -467,6 +466,7 @@ mod tests { #[test] fn pack_variant_chunks_mismatched_shredded_dtype_errors() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ variant_chunk_with_shredded([1], PrimitiveArray::from_iter([10i32]).into_array())? @@ -478,20 +478,17 @@ mod tests { )? .into_array(); - let err = chunked - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) - .unwrap_err(); + let err = chunked.execute::(&mut ctx).unwrap_err(); assert!(err.to_string().contains("shredded dtype mismatch")); Ok(()) } #[test] fn pack_variant_chunks_empty() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let chunked = ChunkedArray::try_new(vec![], VariantDType(NonNullable))?.into_array(); - let variant = into_variant( - chunked.execute::(&mut LEGACY_SESSION.create_execution_ctx())?, - )?; + let variant = into_variant(chunked.execute::(&mut ctx)?)?; assert_eq!(variant.len(), 0); assert!(variant.shredded().is_none()); @@ -500,6 +497,7 @@ mod tests { #[test] fn pack_variant_chunks_single_chunk() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ variant_chunk_with_shredded( @@ -512,9 +510,7 @@ mod tests { )? .into_array(); - let variant = into_variant( - chunked.execute::(&mut LEGACY_SESSION.create_execution_ctx())?, - )?; + let variant = into_variant(chunked.execute::(&mut ctx)?)?; assert_eq!(variant.len(), 2); assert!(variant.shredded().is_some()); @@ -523,6 +519,7 @@ mod tests { #[test] pub fn pack_nested_structs() { + let mut ctx = SESSION.create_execution_ctx(); let struct_array = StructArray::try_new( ["a"].into(), vec![VarBinViewArray::from_iter_str(["foo", "bar", "baz", "quak"]).into_array()], @@ -541,12 +538,17 @@ mod tests { ) .unwrap() .into_array(); - #[expect(deprecated)] - let canonical_struct = chunked.to_struct(); - #[expect(deprecated)] - let canonical_varbin = canonical_struct.unmasked_field(0).to_varbinview(); - #[expect(deprecated)] - let original_varbin = struct_array.unmasked_field(0).to_varbinview(); + let canonical_struct = chunked.execute::(&mut ctx).unwrap(); + let canonical_varbin = canonical_struct + .unmasked_field(0) + .clone() + .execute::(&mut ctx) + .unwrap(); + let original_varbin = struct_array + .unmasked_field(0) + .clone() + .execute::(&mut ctx) + .unwrap(); let orig_values = original_varbin .with_iterator(|it| it.map(|a| a.map(|v| v.to_vec())).collect::>()); let canon_values = canonical_varbin @@ -556,6 +558,7 @@ mod tests { #[test] pub fn pack_nested_lists() { + let mut ctx = SESSION.create_execution_ctx(); let l1 = ListArray::try_new( buffer![1, 2, 3, 4].into_array(), buffer![0, 3].into_array(), @@ -575,27 +578,26 @@ mod tests { List(Arc::new(Primitive(I32, NonNullable)), NonNullable), ); - #[expect(deprecated)] - let canon_values = chunked_list.unwrap().as_array().to_listview(); + let canon_values = chunked_list + .unwrap() + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!( - l1.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - canon_values - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + l1.execute_scalar(0, &mut ctx).unwrap(), + canon_values.execute_scalar(0, &mut ctx).unwrap() ); assert_eq!( - l2.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - canon_values - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + l2.execute_scalar(0, &mut ctx).unwrap(), + canon_values.execute_scalar(1, &mut ctx).unwrap() ); } #[test] fn pack_fixed_size_lists() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let f1 = FixedSizeListArray::try_new( buffer![1, 2, 3, 4, 5, 6].into_array(), 2, @@ -613,9 +615,7 @@ mod tests { let chunked = ChunkedArray::try_new(vec![f1.into_array(), f2.into_array()], dtype)?.into_array(); - let canonical = chunked - .clone() - .execute::(&mut LEGACY_SESSION.create_execution_ctx())?; + let canonical = chunked.clone().execute::(&mut ctx)?; let fsl = match canonical { Canonical::FixedSizeList(fsl) => fsl, other => vortex_bail!("expected FixedSizeList canonical array, got {other:?}"), @@ -630,8 +630,8 @@ mod tests { )?; for idx in 0..5 { assert_eq!( - chunked.execute_scalar(idx, &mut LEGACY_SESSION.create_execution_ctx())?, - expected.execute_scalar(idx, &mut LEGACY_SESSION.create_execution_ctx())?, + chunked.execute_scalar(idx, &mut ctx)?, + expected.execute_scalar(idx, &mut ctx)?, ); } Ok(()) diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 53dcd3d262d..3d8d0898049 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -330,25 +330,30 @@ fn constant_canonical_fixed_size_list_array( #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use enum_iterator::all; use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; + use vortex_session::VortexSession; + use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::ConstantArray; + use crate::arrays::FixedSizeListArray; + use crate::arrays::ListViewArray; + use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; + use crate::arrays::StructArray; use crate::arrays::VarBinArray; + use crate::arrays::VarBinViewArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::ListViewRebuildMode; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; - #[expect(deprecated)] - use crate::canonical::ToCanonical as _; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -356,18 +361,25 @@ mod tests { use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; use crate::scalar::Scalar; + use crate::session::ArraySession; use crate::validity::Validity; + /// A shared session for these constant-array tests, used to create execution contexts. + static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[test] fn test_canonicalize_null() { + let mut ctx = SESSION.create_execution_ctx(); let const_null = ConstantArray::new(Scalar::null(DType::Null), 42); - #[expect(deprecated)] - let actual = const_null.as_array().to_null(); + let actual = const_null + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(actual.len(), 42); assert_eq!( - actual - .execute_scalar(33, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + actual.execute_scalar(33, &mut ctx).unwrap(), Scalar::null(DType::Null) ); } @@ -382,14 +394,13 @@ mod tests { #[test] fn test_canonicalize_propagates_stats() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let scalar = Scalar::bool(true, Nullability::NonNullable); let const_array = ConstantArray::new(scalar, 4).into_array(); - let stats = const_array.statistics().compute_all( - &all::().collect_vec(), - &mut LEGACY_SESSION.create_execution_ctx(), - )?; - #[expect(deprecated)] - let canonical = const_array.to_canonical()?.into_array(); + let stats = const_array + .statistics() + .compute_all(&all::().collect_vec(), &mut ctx)?; + let canonical = const_array.execute::(&mut ctx)?.into_array(); let canonical_stats = canonical.statistics(); let stats_ref = stats.as_typed_ref(canonical.dtype()); @@ -409,47 +420,52 @@ mod tests { #[test] fn test_canonicalize_scalar_values() { + let mut ctx = SESSION.create_execution_ctx(); let f16_value = f16::from_f32(5.722046e-6); let f16_scalar = Scalar::primitive(f16_value, Nullability::NonNullable); // Create a ConstantArray with the f16 scalar let const_array = ConstantArray::new(f16_scalar.clone(), 1).into_array(); - #[expect(deprecated)] - let canonical_const = const_array.to_primitive(); + let canonical_const = const_array.execute::(&mut ctx).unwrap(); // Verify the scalar value is preserved through canonicalization assert_eq!( - canonical_const - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + canonical_const.execute_scalar(0, &mut ctx).unwrap(), f16_scalar ); } #[test] fn test_canonicalize_lists() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let list_scalar = Scalar::list( Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), vec![1u64.into(), 2u64.into()], Nullability::NonNullable, ); let const_array = ConstantArray::new(list_scalar, 2).into_array(); - #[expect(deprecated)] - let canonical_const = const_array.to_listview(); - let list_array = canonical_const.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?; + let canonical_const = const_array.execute::(&mut ctx)?; + let list_array = + canonical_const.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; assert_arrays_eq!( - #[expect(deprecated)] - list_array.elements().to_primitive(), + list_array + .elements() + .clone() + .execute::(&mut ctx)?, PrimitiveArray::from_iter([1u64, 2, 1, 2]) ); assert_arrays_eq!( - #[expect(deprecated)] - list_array.offsets().to_primitive(), + list_array + .offsets() + .clone() + .execute::(&mut ctx)?, PrimitiveArray::from_iter([0u64, 2]) ); assert_arrays_eq!( - #[expect(deprecated)] - list_array.sizes().to_primitive(), + list_array + .sizes() + .clone() + .execute::(&mut ctx)?, PrimitiveArray::from_iter([2u64, 2]) ); Ok(()) @@ -457,55 +473,74 @@ mod tests { #[test] fn test_canonicalize_empty_list() { + let mut ctx = SESSION.create_execution_ctx(); let list_scalar = Scalar::list( Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), vec![], Nullability::NonNullable, ); let const_array = ConstantArray::new(list_scalar, 2).into_array(); - #[expect(deprecated)] - let canonical_const = const_array.to_listview(); - #[expect(deprecated)] - let elements_prim = canonical_const.elements().to_primitive(); + let canonical_const = const_array.execute::(&mut ctx).unwrap(); + let elements_prim = canonical_const + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert!(elements_prim.is_empty()); assert_arrays_eq!( - #[expect(deprecated)] - canonical_const.offsets().to_primitive(), + canonical_const + .offsets() + .clone() + .execute::(&mut ctx) + .unwrap(), PrimitiveArray::from_iter([0u64, 0]) ); assert_arrays_eq!( - #[expect(deprecated)] - canonical_const.sizes().to_primitive(), + canonical_const + .sizes() + .clone() + .execute::(&mut ctx) + .unwrap(), PrimitiveArray::from_iter([0u64, 0]) ); } #[test] fn test_canonicalize_null_list() { + let mut ctx = SESSION.create_execution_ctx(); let list_scalar = Scalar::null(DType::List( Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), Nullability::Nullable, )); let const_array = ConstantArray::new(list_scalar, 2).into_array(); - #[expect(deprecated)] - let canonical_const = const_array.to_listview(); - #[expect(deprecated)] - let elements_prim = canonical_const.elements().to_primitive(); + let canonical_const = const_array.execute::(&mut ctx).unwrap(); + let elements_prim = canonical_const + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert!(elements_prim.is_empty()); assert_arrays_eq!( - #[expect(deprecated)] - canonical_const.offsets().to_primitive(), + canonical_const + .offsets() + .clone() + .execute::(&mut ctx) + .unwrap(), PrimitiveArray::from_iter([0u64, 0]) ); assert_arrays_eq!( - #[expect(deprecated)] - canonical_const.sizes().to_primitive(), + canonical_const + .sizes() + .clone() + .execute::(&mut ctx) + .unwrap(), PrimitiveArray::from_iter([0u64, 0]) ); } #[test] fn test_canonicalize_nullable_struct() { + let mut ctx = SESSION.create_execution_ctx(); let array = ConstantArray::new( Scalar::null(DType::struct_( [( @@ -517,15 +552,13 @@ mod tests { 3, ); - #[expect(deprecated)] - let struct_array = array.as_array().to_struct(); + let struct_array = array + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(struct_array.len(), 3); - assert_eq!( - struct_array - .valid_count(&mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - 0 - ); + assert_eq!(struct_array.valid_count(&mut ctx).unwrap(), 0); let field = struct_array .unmasked_field_by_name("non_null_field") @@ -539,6 +572,7 @@ mod tests { #[test] fn test_canonicalize_fixed_size_list_non_null() { + let mut ctx = SESSION.create_execution_ctx(); // Test with a non-null fixed-size list constant. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), @@ -551,8 +585,7 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 4).into_array(); - #[expect(deprecated)] - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.execute::(&mut ctx).unwrap(); assert_eq!(canonical.len(), 4); assert_eq!(canonical.list_size(), 3); @@ -561,14 +594,14 @@ mod tests { // Check that each list is [10, 20, 30]. for i in 0..4 { let list = canonical.fixed_size_list_elements_at(i).unwrap(); - #[expect(deprecated)] - let list_primitive = list.to_primitive(); + let list_primitive = list.execute::(&mut ctx).unwrap(); assert_arrays_eq!(list_primitive, PrimitiveArray::from_iter([10i32, 20, 30])); } } #[test] fn test_canonicalize_fixed_size_list_nullable() { + let mut ctx = SESSION.create_execution_ctx(); // Test with a nullable but non-null fixed-size list constant. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::F64, Nullability::NonNullable)), @@ -580,16 +613,18 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 3).into_array(); - #[expect(deprecated)] - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.execute::(&mut ctx).unwrap(); assert_eq!(canonical.len(), 3); assert_eq!(canonical.list_size(), 2); assert!(matches!(canonical.validity(), Ok(Validity::AllValid))); // Check elements. - #[expect(deprecated)] - let elements = canonical.elements().to_primitive(); + let elements = canonical + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!( elements, PrimitiveArray::from_iter([1.5f64, 2.5, 1.5, 2.5, 1.5, 2.5]) @@ -598,6 +633,7 @@ mod tests { #[test] fn test_canonicalize_fixed_size_list_null() { + let mut ctx = SESSION.create_execution_ctx(); // Test with a null fixed-size list constant. let fsl_scalar = Scalar::null(DType::FixedSizeList( Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), @@ -606,22 +642,25 @@ mod tests { )); let const_array = ConstantArray::new(fsl_scalar, 5).into_array(); - #[expect(deprecated)] - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.execute::(&mut ctx).unwrap(); assert_eq!(canonical.len(), 5); assert_eq!(canonical.list_size(), 4); assert!(matches!(canonical.validity(), Ok(Validity::AllInvalid))); // Elements should be defaults (zeros). - #[expect(deprecated)] - let elements = canonical.elements().to_primitive(); + let elements = canonical + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(elements.len(), 20); // 5 lists * 4 elements each assert!(elements.as_slice::().iter().all(|&x| x == 0)); } #[test] fn test_canonicalize_fixed_size_list_empty() { + let mut ctx = SESSION.create_execution_ctx(); // Test with size-0 lists (edge case). let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::I8, Nullability::NonNullable)), @@ -630,8 +669,7 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 10).into_array(); - #[expect(deprecated)] - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.execute::(&mut ctx).unwrap(); assert_eq!(canonical.len(), 10); assert_eq!(canonical.list_size(), 0); @@ -643,6 +681,7 @@ mod tests { #[test] fn test_canonicalize_fixed_size_list_nested() { + let mut ctx = SESSION.create_execution_ctx(); // Test with nested data types (list of strings). let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Utf8(Nullability::NonNullable)), @@ -651,43 +690,38 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 2).into_array(); - #[expect(deprecated)] - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.execute::(&mut ctx).unwrap(); assert_eq!(canonical.len(), 2); assert_eq!(canonical.list_size(), 2); // Check elements are repeated correctly. - #[expect(deprecated)] - let elements = canonical.elements().to_varbinview(); + let elements = canonical + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!( - elements - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(0, &mut ctx).unwrap(), "hello".into() ); assert_eq!( - elements - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(1, &mut ctx).unwrap(), "world".into() ); assert_eq!( - elements - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(2, &mut ctx).unwrap(), "hello".into() ); assert_eq!( - elements - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(3, &mut ctx).unwrap(), "world".into() ); } #[test] fn test_canonicalize_fixed_size_list_single_element() { + let mut ctx = SESSION.create_execution_ctx(); // Test with a single-element list. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::I16, Nullability::NonNullable)), @@ -696,19 +730,22 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 1).into_array(); - #[expect(deprecated)] - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.execute::(&mut ctx).unwrap(); assert_eq!(canonical.len(), 1); assert_eq!(canonical.list_size(), 1); - #[expect(deprecated)] - let elements = canonical.elements().to_primitive(); + let elements = canonical + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_arrays_eq!(elements, PrimitiveArray::from_iter([42i16])); } #[test] fn test_canonicalize_fixed_size_list_with_null_elements() { + let mut ctx = SESSION.create_execution_ctx(); // Test FSL with nullable element type where some elements are null. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)), @@ -721,32 +758,28 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 3).into_array(); - #[expect(deprecated)] - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.execute::(&mut ctx).unwrap(); assert_eq!(canonical.len(), 3); assert_eq!(canonical.list_size(), 3); assert!(matches!(canonical.validity(), Ok(Validity::NonNullable))); // Check elements including nulls. - #[expect(deprecated)] - let elements = canonical.elements().to_primitive(); + let elements = canonical + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!( - elements - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(0, &mut ctx).unwrap(), Scalar::from(100i32) ); assert_eq!( - elements - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(1, &mut ctx).unwrap(), Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)) ); assert_eq!( - elements - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(2, &mut ctx).unwrap(), Scalar::from(200i32) ); @@ -766,6 +799,7 @@ mod tests { #[test] fn test_canonicalize_fixed_size_list_large() { + let mut ctx = SESSION.create_execution_ctx(); // Test with a large constant array. let fsl_scalar = Scalar::fixed_size_list( Arc::new(DType::Primitive(PType::U8, Nullability::NonNullable)), @@ -780,47 +814,39 @@ mod tests { ); let const_array = ConstantArray::new(fsl_scalar, 1000).into_array(); - #[expect(deprecated)] - let canonical = const_array.to_fixed_size_list(); + let canonical = const_array.execute::(&mut ctx).unwrap(); assert_eq!(canonical.len(), 1000); assert_eq!(canonical.list_size(), 5); - #[expect(deprecated)] - let elements = canonical.elements().to_primitive(); + let elements = canonical + .elements() + .clone() + .execute::(&mut ctx) + .unwrap(); assert_eq!(elements.len(), 5000); // Check pattern repeats correctly. for i in 0..1000 { let base = i * 5; assert_eq!( - elements - .execute_scalar(base, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(base, &mut ctx).unwrap(), Scalar::from(1u8) ); assert_eq!( - elements - .execute_scalar(base + 1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(base + 1, &mut ctx).unwrap(), Scalar::from(2u8) ); assert_eq!( - elements - .execute_scalar(base + 2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(base + 2, &mut ctx).unwrap(), Scalar::from(3u8) ); assert_eq!( - elements - .execute_scalar(base + 3, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(base + 3, &mut ctx).unwrap(), Scalar::from(4u8) ); assert_eq!( - elements - .execute_scalar(base + 4, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), + elements.execute_scalar(base + 4, &mut ctx).unwrap(), Scalar::from(5u8) ); } diff --git a/vortex-array/src/arrays/list/array.rs b/vortex-array/src/arrays/list/array.rs index 3e1c43ea3e9..7788511aa25 100644 --- a/vortex-array/src/arrays/list/array.rs +++ b/vortex-array/src/arrays/list/array.rs @@ -15,6 +15,8 @@ use vortex_error::vortex_panic; use crate::ArrayRef; use crate::ArraySlots; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; @@ -329,21 +331,24 @@ pub trait ListArrayExt: TypedArrayRef { self.elements().dtype() } - fn reset_offsets(&self, recurse: bool) -> VortexResult> { + fn reset_offsets(&self, recurse: bool, ctx: &mut ExecutionCtx) -> VortexResult> { let mut elements = self.sliced_elements()?; if recurse && elements.is_canonical() { - #[expect(deprecated)] - let compacted = elements.to_canonical()?.compact()?.into_array(); + let compacted = elements + .clone() + .execute::(ctx)? + .compact(ctx)? + .into_array(); elements = compacted; } else if recurse && let Some(child_list_array) = elements.as_opt::() { elements = child_list_array .into_owned() - .reset_offsets(recurse)? + .reset_offsets(recurse, ctx)? .into_array(); } let offsets = self.offsets(); - let first_offset = offsets.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())?; + let first_offset = offsets.execute_scalar(0, ctx)?; let adjusted_offsets = offsets.clone().binary( ConstantArray::new(first_offset, offsets.len()).into_array(), Operator::Sub, diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index f27e0a9f305..2ed260a5f55 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -2,17 +2,18 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; +use std::sync::LazyLock; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; +use vortex_session::VortexSession; use super::*; use crate::Canonical; use crate::IntoArray; -use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::FilterArray; use crate::arrays::List; @@ -24,8 +25,13 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType::I32; use crate::scalar::Scalar; +use crate::session::ArraySession; use crate::validity::Validity; +/// A shared session for `List` tests, used to create execution contexts. +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[test] fn test_empty_list_array() { let elements = PrimitiveArray::empty::(Nullability::NonNullable); @@ -39,6 +45,7 @@ fn test_empty_list_array() { #[test] fn test_simple_list_array() { + let mut ctx = SESSION.create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5].into_array(); let offsets = buffer![0, 2, 4, 5].into_array(); let validity = Validity::AllValid; @@ -51,8 +58,7 @@ fn test_simple_list_array() { vec![1.into(), 2.into()], Nullability::Nullable ), - list.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + list.execute_scalar(0, &mut ctx).unwrap() ); assert_eq!( Scalar::list( @@ -60,18 +66,17 @@ fn test_simple_list_array() { vec![3.into(), 4.into()], Nullability::Nullable ), - list.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + list.execute_scalar(1, &mut ctx).unwrap() ); assert_eq!( Scalar::list(Arc::new(I32.into()), vec![5.into()], Nullability::Nullable), - list.execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + list.execute_scalar(2, &mut ctx).unwrap() ); } #[test] fn test_simple_list_array_from_iter() { + let mut ctx = SESSION.create_execution_ctx(); let elements = buffer![1i32, 2, 3].into_array(); let offsets = buffer![0, 2, 3].into_array(); let validity = Validity::NonNullable; @@ -84,18 +89,12 @@ fn test_simple_list_array_from_iter() { assert_eq!(list.len(), list_from_iter.len()); assert_eq!( - list.execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - list_from_iter - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + list.execute_scalar(0, &mut ctx).unwrap(), + list_from_iter.execute_scalar(0, &mut ctx).unwrap() ); assert_eq!( - list.execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - list_from_iter - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + list.execute_scalar(1, &mut ctx).unwrap(), + list_from_iter.execute_scalar(1, &mut ctx).unwrap() ); } @@ -208,6 +207,7 @@ fn test_list_filter_empty_lists() { #[test] fn test_list_filter_with_nulls() { + let mut ctx = SESSION.create_execution_ctx(); // Test filtering lists with null validity. let elements = buffer![0..15].into_array(); let offsets = buffer![0, 3, 7, 10, 12, 15].into_array(); @@ -227,30 +227,10 @@ fn test_list_filter_with_nulls() { assert_eq!(filtered.len(), 4); // Check validity of filtered array using scalar_at (works on any array). - assert!( - filtered - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_valid() - ); - assert!( - !filtered - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_valid() - ); // Was null. - assert!( - !filtered - .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_valid() - ); // Was null. - assert!( - filtered - .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - .is_valid() - ); + assert!(filtered.execute_scalar(0, &mut ctx).unwrap().is_valid()); + assert!(!filtered.execute_scalar(1, &mut ctx).unwrap().is_valid()); // Was null. + assert!(!filtered.execute_scalar(2, &mut ctx).unwrap().is_valid()); // Was null. + assert!(filtered.execute_scalar(3, &mut ctx).unwrap().is_valid()); } #[test] @@ -655,7 +635,7 @@ fn test_list_of_lists() { // Test scalar conversion. let scalar = list_of_lists - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(0, &mut SESSION.create_execution_ctx()) .unwrap(); assert!(matches!(scalar.dtype(), DType::List(_, _))); let list_scalar = scalar.as_list(); @@ -679,6 +659,7 @@ fn test_list_of_lists() { #[test] fn test_list_of_lists_nullable_outer() { + let mut ctx = SESSION.create_execution_ctx(); // Create list of lists with nullable outer, non-nullable inner. // Structure: [[[1, 2], [3]], null, [[4, 5, 6]], [[7]]] let data = vec![ @@ -701,15 +682,11 @@ fn test_list_of_lists_nullable_outer() { )); // First element should be [[1, 2], [3]]. - let first = list_of_lists - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let first = list_of_lists.execute_scalar(0, &mut ctx).unwrap(); assert!(!first.is_null()); // Second element should be null. - let second = list_of_lists - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let second = list_of_lists.execute_scalar(1, &mut ctx).unwrap(); assert!(second.is_null()); // Third element should be [[4, 5, 6]]. @@ -764,13 +741,14 @@ fn test_list_of_lists_nullable_inner() { // Check that second inner list is null. let second_inner = first_list .array() - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) + .execute_scalar(1, &mut SESSION.create_execution_ctx()) .unwrap(); assert!(second_inner.is_null()); } #[test] fn test_list_of_lists_both_nullable() { + let mut ctx = SESSION.create_execution_ctx(); // Create list of lists with both nullable. // Structure: [[[1, 2], null], null, [[3]], [null]] let data = vec![ @@ -793,9 +771,7 @@ fn test_list_of_lists_both_nullable() { )); // First outer list should have 2 elements, second is null inner list. - let first_outer = list_of_lists - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let first_outer = list_of_lists.execute_scalar(0, &mut ctx).unwrap(); assert!(!first_outer.is_null()); let first_outer_array = list_of_lists.list_elements_at(0).unwrap(); let first_list = first_outer_array.as_::(); @@ -806,16 +782,11 @@ fn test_list_of_lists_both_nullable() { assert_eq!(first_inner.len(), 2); // Second inner list should be null. - let second_inner = first_list - .array() - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let second_inner = first_list.array().execute_scalar(1, &mut ctx).unwrap(); assert!(second_inner.is_null()); // Second outer list should be null. - let second_outer = list_of_lists - .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let second_outer = list_of_lists.execute_scalar(1, &mut ctx).unwrap(); assert!(second_outer.is_null()); // Third outer list should have [3]. @@ -829,10 +800,7 @@ fn test_list_of_lists_both_nullable() { let fourth_outer = list_of_lists.list_elements_at(3).unwrap(); let fourth_list = fourth_outer.as_::(); assert_eq!(fourth_list.len(), 1); - let inner = fourth_list - .array() - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + let inner = fourth_list.array().execute_scalar(0, &mut ctx).unwrap(); assert!(inner.is_null()); } @@ -925,6 +893,7 @@ fn test_offsets_constant() { #[test] fn test_recursive_compact_list_of_lists() { + let mut ctx = SESSION.create_execution_ctx(); // Create a nested list structure: [[[1,2,3], [4,5]], [[6,7,8,9]], [[10], [11,12]]] let nested_data = vec![ Some(vec![ @@ -941,9 +910,9 @@ fn test_recursive_compact_list_of_lists() { let sliced_list = sliced.as_::(); // Test non-recursive compaction: only resets outer list offsets - let non_recursive = sliced_list.reset_offsets(false).unwrap(); + let non_recursive = sliced_list.reset_offsets(false, &mut ctx).unwrap(); // Test recursive compaction: resets offsets and compacts inner canonical arrays - let recursive = sliced_list.reset_offsets(true).unwrap(); + let recursive = sliced_list.reset_offsets(true, &mut ctx).unwrap(); assert_eq!(non_recursive.len(), 2); assert_eq!(recursive.len(), 2); @@ -964,12 +933,8 @@ fn test_recursive_compact_list_of_lists() { let non_recursive_array = non_recursive.into_array(); let recursive_array = recursive.into_array(); assert_eq!( - non_recursive_array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(), - recursive_array - .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() + non_recursive_array.execute_scalar(0, &mut ctx).unwrap(), + recursive_array.execute_scalar(0, &mut ctx).unwrap() ); } @@ -985,7 +950,7 @@ fn test_filter_sliced_list_array() -> VortexResult<()> { let mask = Mask::from(BitBuffer::from(vec![true, false, true])); let filter_array = FilterArray::new(list, mask).into_array(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let result = filter_array.execute::(&mut ctx)?; assert_eq!(result.len(), 2); diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index a80ee90ca9d..a673b10143a 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -8,8 +8,6 @@ use crate::ArrayRef; use crate::Canonical; use crate::ExecutionCtx; use crate::IntoArray; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; use crate::arrays::ListArray; @@ -41,7 +39,9 @@ pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexRes // We reset the offsets here because mostly for convenience, and also because callers of this // function might not expect the output `ListViewArray` to have a bunch of leading and trailing // garbage data when they turn it back into a `ListArray`. - let list = list.reset_offsets(false).vortex_expect("This can't fail"); + let list = list + .reset_offsets(false, ctx) + .vortex_expect("This can't fail"); let list_offsets = list.offsets().clone(); @@ -108,15 +108,18 @@ fn build_sizes_from_offsets( /// /// Otherwise, this function fall back to the (very) expensive path and will rebuild the /// `ListArray` from scratch. -pub fn list_from_list_view(list_view: ListViewArray) -> VortexResult { +pub fn list_from_list_view( + list_view: ListViewArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { // Rebuild as zero-copyable to list array and also trim all leading and trailing elements. - let zctl_array = list_view.rebuild(ListViewRebuildMode::MakeExact)?; + let zctl_array = list_view.rebuild(ListViewRebuildMode::MakeExact, ctx)?; debug_assert!(zctl_array.is_zero_copy_to_list()); let list_offsets = match_each_integer_ptype!(zctl_array.offsets().dtype().as_ptype(), |O| { // SAFETY: We just made the array zero-copyable to `ListArray`, so the safety contract is // upheld. - unsafe { build_list_offsets_from_list_view::(&zctl_array) } + unsafe { build_list_offsets_from_list_view::(&zctl_array, ctx) } }); // SAFETY: Because the shape of the `ListViewArray` is zero-copyable to a `ListArray`, we @@ -142,6 +145,7 @@ pub fn list_from_list_view(list_view: ListViewArray) -> VortexResult /// between `offset[i]` and `offset[i + 1]`. unsafe fn build_list_offsets_from_list_view( list_view: &ListViewArray, + ctx: &mut ExecutionCtx, ) -> ArrayRef { let len = list_view.len(); let mut offsets_builder = @@ -150,8 +154,11 @@ unsafe fn build_list_offsets_from_list_view( // Create uninit range for direct memory access. let mut offsets_range = offsets_builder.uninit_range(len + 1); - #[expect(deprecated)] - let offsets = list_view.offsets().to_primitive(); + let offsets = list_view + .offsets() + .clone() + .execute::(ctx) + .vortex_expect("list view offsets must be primitive after rebuild"); let offsets_slice = offsets.as_slice::(); debug_assert!(offsets_slice.is_sorted()); @@ -184,17 +191,20 @@ unsafe fn build_list_offsets_from_list_view( /// Recursively converts all `ListViewArray`s to `ListArray`s in a nested array structure. /// /// The conversion happens bottom-up, processing children before parents. -pub fn recursive_list_from_list_view(array: ArrayRef) -> VortexResult { +pub fn recursive_list_from_list_view( + array: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { if !array.dtype().is_nested() { return Ok(array); } - #[expect(deprecated)] - let canonical = array.to_canonical()?; + let canonical = array.execute::(ctx)?; Ok(match canonical { Canonical::List(listview) => { - let converted_elements = recursive_list_from_list_view(listview.elements().clone())?; + let converted_elements = + recursive_list_from_list_view(listview.elements().clone(), ctx)?; debug_assert_eq!(converted_elements.len(), listview.elements().len()); // Avoid cloning if elements didn't change. @@ -216,12 +226,12 @@ pub fn recursive_list_from_list_view(array: ArrayRef) -> VortexResult }; // Make the conversion to `ListArray`. - let list_array = list_from_list_view(listview_with_converted_elements)?; + let list_array = list_from_list_view(listview_with_converted_elements, ctx)?; list_array.into_array() } Canonical::FixedSizeList(fixed_size_list) => { let converted_elements = - recursive_list_from_list_view(fixed_size_list.elements().clone())?; + recursive_list_from_list_view(fixed_size_list.elements().clone(), ctx)?; // Avoid cloning if elements didn't change. if !ArrayRef::ptr_eq(&converted_elements, fixed_size_list.elements()) { @@ -245,7 +255,7 @@ pub fn recursive_list_from_list_view(array: ArrayRef) -> VortexResult let mut any_changed = false; for field in fields.iter() { - let converted_field = recursive_list_from_list_view(field.clone())?; + let converted_field = recursive_list_from_list_view(field.clone(), ctx)?; // Avoid cloning if elements didn't change. any_changed |= !ArrayRef::ptr_eq(&converted_field, field); converted_fields.push(converted_field); @@ -266,7 +276,7 @@ pub fn recursive_list_from_list_view(array: ArrayRef) -> VortexResult } Canonical::Extension(ext_array) => { let converted_storage = - recursive_list_from_list_view(ext_array.storage_array().clone())?; + recursive_list_from_list_view(ext_array.storage_array().clone(), ctx)?; // Avoid cloning if elements didn't change. if !ArrayRef::ptr_eq(&converted_storage, ext_array.storage_array()) { @@ -286,6 +296,7 @@ mod tests { use vortex_error::VortexExpect; use vortex_error::VortexResult; + use super::super::tests::common::SESSION; use super::super::tests::common::create_basic_listview; use super::super::tests::common::create_empty_lists_listview; use super::super::tests::common::create_nullable_listview; @@ -295,7 +306,6 @@ mod tests { use crate::ArrayRef; use crate::EqMode; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::FixedSizeListArray; @@ -319,7 +329,7 @@ mod tests { let offsets = buffer![0u32, 3, 5, 7, 10].into_array(); let list_array = ListArray::try_new(elements.clone(), offsets, Validity::NonNullable)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let list_view = list_view_from_list(list_array.clone(), &mut ctx)?; // Verify structure. @@ -342,7 +352,8 @@ mod tests { #[test] fn test_listview_to_list_zero_copy() -> VortexResult<()> { let list_view = create_basic_listview(); - let list_array = list_from_list_view(list_view.clone())?; + 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()); @@ -366,12 +377,12 @@ mod tests { // This conversion will create an empty ListViewArray. // Note: list_view_from_list handles the empty case specially. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let empty_list_view = list_view_from_list(empty_list.clone(), &mut ctx)?; assert_eq!(empty_list_view.len(), 0); // Convert back. - let converted_back = list_from_list_view(empty_list_view)?; + let converted_back = list_from_list_view(empty_list_view, &mut ctx)?; assert_eq!(converted_back.len(), 0); // For empty arrays, we can't use assert_arrays_eq directly since the offsets might differ. // Just check that it's empty. @@ -387,7 +398,7 @@ mod tests { let validity = Validity::Array(BoolArray::from_iter(vec![true, false, true]).into_array()); let nullable_list = ListArray::try_new(elements, offsets, validity.clone())?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let nullable_list_view = list_view_from_list(nullable_list.clone(), &mut ctx)?; // Verify validity is preserved. @@ -400,7 +411,7 @@ mod tests { assert_eq!(nullable_list_view.len(), 3); // Round-trip conversion. - let converted_back = list_from_list_view(nullable_list_view)?; + let converted_back = list_from_list_view(nullable_list_view, &mut ctx)?; assert_arrays_eq!(nullable_list, converted_back); Ok(()) } @@ -409,7 +420,8 @@ mod tests { fn test_non_zero_copy_listview_to_list() -> VortexResult<()> { // Create ListViewArray with overlapping lists (not zero-copyable). let list_view = create_overlapping_listview(); - let list_array = list_from_list_view(list_view.clone())?; + let list_array = + list_from_list_view(list_view.clone(), &mut SESSION.create_execution_ctx())?; // The resulting ListArray should have monotonic offsets. for i in 0..list_array.len() { @@ -425,10 +437,11 @@ mod tests { #[test] fn test_empty_sublists() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let empty_lists_view = create_empty_lists_listview(); // Convert to ListArray. - let list_array = list_from_list_view(empty_lists_view.clone())?; + let list_array = list_from_list_view(empty_lists_view.clone(), &mut ctx)?; assert_eq!(list_array.len(), 4); // All sublists should be empty. @@ -437,7 +450,6 @@ mod tests { } // Round-trip. - let mut ctx = LEGACY_SESSION.create_execution_ctx(); let converted_back = list_view_from_list(list_array, &mut ctx)?; assert_arrays_eq!(empty_lists_view, converted_back); Ok(()) @@ -451,7 +463,7 @@ mod tests { let list_i32 = ListArray::try_new(elements.clone(), i32_offsets.clone(), Validity::NonNullable)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let list_view_i32 = list_view_from_list(list_i32.clone(), &mut ctx)?; assert_eq!(list_view_i32.offsets().dtype(), i32_offsets.dtype()); assert_eq!(list_view_i32.sizes().dtype(), i32_offsets.dtype()); @@ -472,24 +484,24 @@ mod tests { #[test] fn test_round_trip_conversions() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Test 1: Basic round-trip. let original = create_basic_listview(); - let to_list = list_from_list_view(original.clone())?; + 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); // Test 2: Nullable round-trip. let nullable = create_nullable_listview(); - let nullable_to_list = list_from_list_view(nullable.clone())?; + 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); // Test 3: Non-zero-copyable round-trip. let overlapping = create_overlapping_listview(); - let overlapping_to_list = list_from_list_view(overlapping.clone())?; + 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); Ok(()) @@ -502,7 +514,7 @@ mod tests { let offsets = buffer![0u32, 1, 2, 3].into_array(); let single_elem_list = ListArray::try_new(elements, offsets, Validity::NonNullable)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let list_view = list_view_from_list(single_elem_list.clone(), &mut ctx)?; assert_eq!(list_view.len(), 3); @@ -511,7 +523,7 @@ mod tests { assert_arrays_eq!(expected_sizes, list_view.sizes().clone()); // Round-trip. - let converted_back = list_from_list_view(list_view)?; + let converted_back = list_from_list_view(list_view, &mut ctx)?; assert_arrays_eq!(single_elem_list, converted_back); Ok(()) } @@ -523,7 +535,7 @@ mod tests { let offsets = buffer![0u32, 2, 2, 3, 3, 6].into_array(); let mixed_list = ListArray::try_new(elements, offsets, Validity::NonNullable)?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let list_view = list_view_from_list(mixed_list.clone(), &mut ctx)?; assert_eq!(list_view.len(), 5); @@ -532,7 +544,7 @@ mod tests { assert_arrays_eq!(expected_sizes, list_view.sizes().clone()); // Round-trip. - let converted_back = list_from_list_view(list_view)?; + let converted_back = list_from_list_view(list_view, &mut ctx)?; assert_arrays_eq!(mixed_list, converted_back); Ok(()) } @@ -540,7 +552,10 @@ mod tests { #[test] fn test_recursive_simple_listview() -> VortexResult<()> { let list_view = create_basic_listview(); - let result = recursive_list_from_list_view(list_view.clone().into_array())?; + let result = recursive_list_from_list_view( + list_view.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; assert_eq!(result.len(), list_view.len()); assert_arrays_eq!(list_view.into_array(), result); @@ -574,7 +589,10 @@ mod tests { .with_zero_copy_to_list(true) }; - let result = recursive_list_from_list_view(outer_listview.clone().into_array())?; + let result = recursive_list_from_list_view( + outer_listview.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; assert_eq!(result.len(), 2); assert_arrays_eq!(outer_listview.into_array(), result); @@ -593,7 +611,10 @@ mod tests { Validity::NonNullable, )?; - let result = recursive_list_from_list_view(struct_array.clone().into_array())?; + let result = recursive_list_from_list_view( + struct_array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; assert_eq!(result.len(), 4); assert_arrays_eq!(struct_array.into_array(), result); @@ -635,7 +656,10 @@ mod tests { let fixed_list = FixedSizeListArray::new(chunked_listviews.into_array(), 1, Validity::NonNullable, 2); - let result = recursive_list_from_list_view(fixed_list.clone().into_array())?; + let result = recursive_list_from_list_view( + fixed_list.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; assert_eq!(result.len(), 2); assert_arrays_eq!(fixed_list.into_array(), result); @@ -676,7 +700,10 @@ mod tests { .with_zero_copy_to_list(true) }; - let result = recursive_list_from_list_view(outer_listview.clone().into_array())?; + let result = recursive_list_from_list_view( + outer_listview.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; assert_eq!(result.len(), 2); assert_arrays_eq!(outer_listview.into_array(), result); @@ -687,7 +714,7 @@ mod tests { fn test_recursive_primitive_unchanged() -> VortexResult<()> { let prim = buffer![1i32, 2, 3].into_array(); let prim_clone = prim.clone(); - let result = recursive_list_from_list_view(prim)?; + let result = recursive_list_from_list_view(prim, &mut SESSION.create_execution_ctx())?; assert!(ArrayRef::ptr_eq(&result, &prim_clone)); Ok(()) @@ -695,8 +722,9 @@ mod tests { #[test] fn test_recursive_mixed_listview_and_list() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); let listview = create_basic_listview(); - let list = list_from_list_view(listview.clone())?; + let list = list_from_list_view(listview.clone(), &mut ctx)?; let struct_array = StructArray::try_new( FieldNames::from(["listview_field", "list_field"]), @@ -705,7 +733,7 @@ mod tests { Validity::NonNullable, )?; - let result = recursive_list_from_list_view(struct_array.clone().into_array())?; + 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); @@ -727,7 +755,7 @@ mod tests { // `try_new` sets `is_zero_copy_to_list: false`. assert!(!list_view.is_zero_copy_to_list()); - let list_array = list_from_list_view(list_view)?; + let list_array = list_from_list_view(list_view, &mut SESSION.create_execution_ctx())?; assert_eq!(list_array.len(), 0); Ok(()) } diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs index 2bcb4ef0c1f..c42f1bf85a1 100644 --- a/vortex-array/src/arrays/listview/rebuild.rs +++ b/vortex-array/src/arrays/listview/rebuild.rs @@ -6,11 +6,9 @@ use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use crate::Canonical; +use crate::ExecutionCtx; use crate::IntoArray; -use crate::LEGACY_SESSION; -#[expect(deprecated)] -use crate::ToCanonical as _; -use crate::VortexSessionExecute; use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -99,16 +97,20 @@ pub enum ListViewRebuildMode { impl ListViewArray { /// Rebuilds the [`ListViewArray`] according to the specified mode. - pub fn rebuild(&self, mode: ListViewRebuildMode) -> VortexResult { + pub fn rebuild( + &self, + mode: ListViewRebuildMode, + ctx: &mut ExecutionCtx, + ) -> VortexResult { if self.is_empty() { // SAFETY: An empty array is trivially zero-copyable to a `ListArray`. return Ok(unsafe { self.clone().with_zero_copy_to_list(true) }); } match mode { - ListViewRebuildMode::MakeZeroCopyToList => self.rebuild_zero_copy_to_list(), - ListViewRebuildMode::TrimElements => self.rebuild_trim_elements(), - ListViewRebuildMode::MakeExact => self.rebuild_make_exact(), + ListViewRebuildMode::MakeZeroCopyToList => self.rebuild_zero_copy_to_list(ctx), + ListViewRebuildMode::TrimElements => self.rebuild_trim_elements(ctx), + ListViewRebuildMode::MakeExact => self.rebuild_make_exact(ctx), ListViewRebuildMode::OverlapCompression => unimplemented!("Does P=NP?"), } } @@ -120,7 +122,7 @@ impl ListViewArray { /// a [`ListArray`]. /// /// [`ListArray`]: crate::arrays::ListArray - fn rebuild_zero_copy_to_list(&self) -> VortexResult { + fn rebuild_zero_copy_to_list(&self, ctx: &mut ExecutionCtx) -> VortexResult { if self.is_zero_copy_to_list() { // Note that since everything in `ListViewArray` is `Arc`ed, this is quite cheap. return Ok(self.clone()); @@ -138,10 +140,10 @@ impl ListViewArray { // for sizes as well. match_each_unsigned_integer_ptype!(sizes_ptype.to_unsigned(), |S| { match offsets_ptype.to_unsigned() { - PType::U8 => self.naive_rebuild::(), - PType::U16 => self.naive_rebuild::(), - PType::U32 => self.naive_rebuild::(), - PType::U64 => self.naive_rebuild::(), + PType::U8 => self.naive_rebuild::(ctx), + PType::U16 => self.naive_rebuild::(ctx), + PType::U32 => self.naive_rebuild::(ctx), + PType::U64 => self.naive_rebuild::(ctx), _ => unreachable!("invalid offsets PType"), } }) @@ -152,9 +154,9 @@ impl ListViewArray { /// list size. fn naive_rebuild( &self, + ctx: &mut ExecutionCtx, ) -> VortexResult { - #[expect(deprecated)] - let sizes_canonical = self.sizes().to_primitive(); + let sizes_canonical = self.sizes().clone().execute::(ctx)?; let sizes_canonical = sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); let total: u64 = sizes_canonical @@ -163,9 +165,9 @@ impl ListViewArray { .map(|s| (*s).as_() as u64) .sum(); if Self::should_use_take(total, self.len()) { - self.rebuild_with_take::() + self.rebuild_with_take::(ctx) } else { - self.rebuild_list_by_list::() + self.rebuild_list_by_list::(ctx) } } @@ -188,17 +190,16 @@ impl ListViewArray { /// `BufferMut`, perform a single `take`. fn rebuild_with_take( &self, + ctx: &mut ExecutionCtx, ) -> VortexResult { let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype()); let size_ptype = self.sizes().dtype().as_ptype(); - #[expect(deprecated)] - let offsets_canonical = self.offsets().to_primitive(); + let offsets_canonical = self.offsets().clone().execute::(ctx)?; let offsets_canonical = offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned()); let offsets_slice = offsets_canonical.as_slice::(); - #[expect(deprecated)] - let sizes_canonical = self.sizes().to_primitive(); + let sizes_canonical = self.sizes().clone().execute::(ctx)?; let sizes_canonical = sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); let sizes_slice = sizes_canonical.as_slice::(); @@ -250,6 +251,7 @@ impl ListViewArray { /// the relevant range and `extend_from_array` into a typed builder. fn rebuild_list_by_list( &self, + ctx: &mut ExecutionCtx, ) -> VortexResult { let element_dtype = self .dtype() @@ -259,13 +261,11 @@ impl ListViewArray { let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype()); let size_ptype = self.sizes().dtype().as_ptype(); - #[expect(deprecated)] - let offsets_canonical = self.offsets().to_primitive(); + let offsets_canonical = self.offsets().clone().execute::(ctx)?; let offsets_canonical = offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned()); let offsets_slice = offsets_canonical.as_slice::(); - #[expect(deprecated)] - let sizes_canonical = self.sizes().to_primitive(); + let sizes_canonical = self.sizes().clone().execute::(ctx)?; let sizes_canonical = sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned()); let sizes_slice = sizes_canonical.as_slice::(); @@ -280,11 +280,10 @@ impl ListViewArray { let mut new_sizes = BufferMut::::with_capacity(len); // Canonicalize the elements up front as we will be slicing the elements quite a lot. - #[expect(deprecated)] let elements_canonical = self .elements() - .to_canonical() - .vortex_expect("canonicalize elements for rebuild") + .clone() + .execute::(ctx)? .into_array(); // Note that we do not know what the exact capacity should be of the new elements since @@ -347,9 +346,8 @@ impl ListViewArray { /// Rebuilds a [`ListViewArray`] by trimming any unused / unreferenced leading and trailing /// elements, which is defined as a contiguous run of values in the `elements` array that are /// not referenced by any views in the corresponding [`ListViewArray`]. - fn rebuild_trim_elements(&self) -> VortexResult { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let (start, end) = self.referenced_element_bounds(&mut ctx)?; + fn rebuild_trim_elements(&self, ctx: &mut ExecutionCtx) -> VortexResult { + let (start, end) = self.referenced_element_bounds(ctx)?; // SAFETY: we calculated valid start and end bounds unsafe { self.trim_elements(start, end) } @@ -395,13 +393,13 @@ impl ListViewArray { }) } - fn rebuild_make_exact(&self) -> VortexResult { + fn rebuild_make_exact(&self, ctx: &mut ExecutionCtx) -> VortexResult { if self.is_zero_copy_to_list() { - self.rebuild_trim_elements() + self.rebuild_trim_elements(ctx) } else { // When we completely rebuild the `ListViewArray`, we get the benefit that we also trim // any leading and trailing garbage data. - self.rebuild_zero_copy_to_list() + self.rebuild_zero_copy_to_list(ctx) } } } @@ -411,11 +409,9 @@ mod tests { use vortex_buffer::BitBuffer; use vortex_error::VortexResult; + use super::super::tests::common::SESSION; use super::ListViewRebuildMode; use crate::IntoArray; - use crate::LEGACY_SESSION; - #[expect(deprecated)] - use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; @@ -435,7 +431,8 @@ mod tests { let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?; + let mut ctx = SESSION.create_execution_ctx(); + let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; // After flatten: elements should be [A, B, C, B, C] = [1, 2, 3, 2, 3] // Lists should be sequential with no overlaps @@ -478,7 +475,8 @@ mod tests { let listview = ListViewArray::new(elements, offsets, sizes, validity); - let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?; + let mut ctx = SESSION.create_execution_ctx(); + let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; // Verify nullability is preserved assert_eq!(flattened.dtype().nullability(), Nullability::Nullable); @@ -515,7 +513,8 @@ mod tests { let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements)?; + let mut ctx = SESSION.create_execution_ctx(); + let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; // After trimming: elements should be [A, B, _, C, D] = [1, 2, 97, 3, 4]. assert_eq!(trimmed.elements().len(), 5); @@ -538,12 +537,11 @@ mod tests { ); // Note that element at index 2 (97) is preserved as a gap. - #[expect(deprecated)] - let all_elements = trimmed.elements().to_primitive(); - assert_eq!( - all_elements.execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())?, - 97i32.into() - ); + let all_elements = trimmed + .elements() + .clone() + .execute::(&mut ctx)?; + assert_eq!(all_elements.execute_scalar(2, &mut ctx)?, 97i32.into()); Ok(()) } @@ -562,7 +560,8 @@ mod tests { let listview = ListViewArray::new(elements, offsets, sizes, validity); // First rebuild to make it zero-copy-to-list - let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?; + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; assert!(rebuilt.is_zero_copy_to_list()); // Verify NULL items have correct offsets (should not reuse previous offsets) @@ -580,13 +579,13 @@ mod tests { // Now rebuild with MakeExact (which calls naive_rebuild then trim_elements) // This should not panic (issue #5412) - let exact = rebuilt.rebuild(ListViewRebuildMode::MakeExact)?; + let exact = rebuilt.rebuild(ListViewRebuildMode::MakeExact, &mut ctx)?; // Verify the result is still valid - assert!(exact.is_valid(0, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(exact.is_valid(1, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(!exact.is_valid(2, &mut LEGACY_SESSION.create_execution_ctx())?); - assert!(!exact.is_valid(3, &mut LEGACY_SESSION.create_execution_ctx())?); + assert!(exact.is_valid(0, &mut ctx)?); + assert!(exact.is_valid(1, &mut ctx)?); + assert!(!exact.is_valid(2, &mut ctx)?); + assert!(!exact.is_valid(3, &mut ctx)?); // Verify data is preserved assert_arrays_eq!( @@ -615,7 +614,8 @@ mod tests { let sizes = PrimitiveArray::from_iter(vec![2u16, 2]).into_array(); let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements)?; + let mut ctx = SESSION.create_execution_ctx(); + let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; assert_arrays_eq!( trimmed.list_elements_at(1)?, PrimitiveArray::from_iter([30i32, 40]) @@ -635,7 +635,8 @@ mod tests { let sizes = PrimitiveArray::from_iter(vec![70_000u32, 2]).into_array(); let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements)?; + let mut ctx = SESSION.create_execution_ctx(); + let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; assert_arrays_eq!( trimmed.list_elements_at(1)?, PrimitiveArray::from_iter([30i32, 40]) @@ -655,7 +656,8 @@ mod tests { let sizes = PrimitiveArray::from_iter(vec![3i16, 2]).into_array(); let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?; + let mut ctx = SESSION.create_execution_ctx(); + let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?; // Values: [1,2,3] and [2,3]. assert_arrays_eq!( @@ -697,7 +699,8 @@ mod tests { let sizes = PrimitiveArray::from_iter(vec![46u8, 10]).into_array(); let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable); - let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements)?; + let mut ctx = SESSION.create_execution_ctx(); + let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?; // min(offsets) = 0, so nothing to trim; output should equal input. assert_arrays_eq!(trimmed, listview); diff --git a/vortex-array/src/arrays/listview/tests/common.rs b/vortex-array/src/arrays/listview/tests/common.rs index dc8bc94fb18..be89b586209 100644 --- a/vortex-array/src/arrays/listview/tests/common.rs +++ b/vortex-array/src/arrays/listview/tests/common.rs @@ -3,14 +3,23 @@ //! Common test utilities for ListView tests. +use std::sync::LazyLock; + use vortex_buffer::buffer; +use vortex_session::VortexSession; use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; +use crate::session::ArraySession; use crate::validity::Validity; +/// A shared session for `ListView` tests, used to create execution contexts via +/// [`create_execution_ctx`](crate::VortexSessionExecute::create_execution_ctx). +pub static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + /// Creates a basic ListView for testing: [[0,1,2], [3,4], [5,6], [7,8,9]] pub fn create_basic_listview() -> ListViewArray { let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); diff --git a/vortex-array/src/arrow/executor/list.rs b/vortex-array/src/arrow/executor/list.rs index d59088ade34..ed8387fe423 100644 --- a/vortex-array/src/arrow/executor/list.rs +++ b/vortex-array/src/arrow/executor/list.rs @@ -61,7 +61,7 @@ pub(super) fn to_arrow_list( let zctl = if array.is_zero_copy_to_list() { array } else { - array.rebuild(ListViewRebuildMode::MakeZeroCopyToList)? + array.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? }; return list_view_zctl::(zctl, elements_field, ctx); } @@ -76,7 +76,7 @@ pub(super) fn to_arrow_list( let zctl = if list_view.is_zero_copy_to_list() { list_view } else { - list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList)? + list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? }; list_view_zctl::(zctl, elements_field, ctx) } @@ -201,6 +201,7 @@ fn list_view_zctl( #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::Array; use arrow_array::GenericListArray; @@ -209,21 +210,26 @@ mod tests { use arrow_schema::Field; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::Canonical; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::PrimitiveArray; use crate::arrow::ArrowArrayExecutor; use crate::arrow::executor::list::ListViewArray; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; + use crate::session::ArraySession; use crate::validity::Validity; + /// A shared session for these list-executor tests, used to create execution contexts. + static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[test] fn test_to_arrow_list_i32() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a ListViewArray with i32 elements: [[1, 2, 3], [4, 5]] let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::NonNullable); let offsets = PrimitiveArray::new(buffer![0i32, 3], Validity::NonNullable); @@ -278,7 +284,7 @@ mod tests { #[test] fn test_to_arrow_list_i64() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a ListViewArray with i64 offsets: [[10, 20], [30]] let elements = PrimitiveArray::new(buffer![10i64, 20, 30], Validity::NonNullable); let offsets = PrimitiveArray::new(buffer![0i64, 2], Validity::NonNullable); @@ -318,7 +324,7 @@ mod tests { #[test] fn test_to_arrow_list_non_zctl() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Overlapping lists are NOT zero-copy-to-list, so this exercises the rebuild path. // Elements: [1, 2, 3, 4], List 0: [1,2,3], List 1: [2,3,4] (overlap at indices 1-2) let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable); @@ -360,7 +366,7 @@ mod tests { #[test] fn test_to_arrow_list_empty_zctl() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let dtype = DType::List( Arc::new(DType::Primitive(crate::dtype::PType::I32, NonNullable)), NonNullable, diff --git a/vortex-array/src/arrow/executor/list_view.rs b/vortex-array/src/arrow/executor/list_view.rs index b77d079f509..5a2cd517b37 100644 --- a/vortex-array/src/arrow/executor/list_view.rs +++ b/vortex-array/src/arrow/executor/list_view.rs @@ -54,7 +54,7 @@ pub(super) fn to_arrow_list_view( } } else if array.upper_bound_density(ctx)? < DEFAULT_REBUILD_DENSITY_THRESHOLD { // Overlaps, gaps, or garbage may be present, so a full rebuild is needed to reclaim waste. - array.rebuild(ListViewRebuildMode::MakeZeroCopyToList)? + array.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? } else { array }; @@ -109,24 +109,31 @@ fn list_view_to_list_view( #[cfg(test)] mod tests { + use std::sync::LazyLock; + use arrow_array::Array; use arrow_array::GenericListViewArray; use arrow_schema::DataType; use arrow_schema::Field; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::VortexSession; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrow::ArrowArrayExecutor; use crate::arrow::executor::list_view::ListViewArray; use crate::arrow::executor::list_view::PrimitiveArray; + use crate::session::ArraySession; use crate::validity::Validity; + /// A shared session for these list-view-executor tests, used to create execution contexts. + static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[test] fn trims_zero_copy_with_significant_trailing_waste() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Zero-copy-to-list array with 10 elements but only [0, 4) referenced -> 60% waste. // The conversion should trim the elements buffer down to the referenced range. let elements = PrimitiveArray::new( @@ -161,7 +168,7 @@ mod tests { #[test] fn test_to_arrow_listview_i32() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a ListViewArray with overlapping views: [[1, 2], [2, 3], [3, 4]] let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable); let offsets = PrimitiveArray::new(buffer![0i32, 1, 2], Validity::NonNullable); @@ -216,7 +223,7 @@ mod tests { #[test] fn test_to_arrow_listview_i64() -> VortexResult<()> { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a ListViewArray with nullable elements: [[100], null, [200, 300]] let elements = PrimitiveArray::new(buffer![100i64, 200, 300], Validity::NonNullable); let offsets = PrimitiveArray::new(buffer![0i64, 1, 1], Validity::NonNullable); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index d4f8f908dcb..837bd50b7db 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -21,8 +21,6 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::Canonical; use crate::LEGACY_SESSION; -#[expect(deprecated)] -use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::array::IntoArray; use crate::arrays::ListViewArray; @@ -294,8 +292,12 @@ impl ArrayBuilder for ListViewBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { - #[expect(deprecated)] - let listview = array.to_listview(); + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + + let listview = array + .clone() + .execute::(&mut ctx) + .vortex_expect("failed to execute array into ListViewArray in extend_from_array"); if listview.is_empty() { return; } @@ -303,7 +305,7 @@ impl ArrayBuilder for ListViewBuilder { // Normalize to an exact zero-copy-to-list layout and then bulk append. This avoids the // very expensive scalar_at-per-list path for overlapping / out-of-order list views. let listview = listview - .rebuild(ListViewRebuildMode::MakeExact) + .rebuild(ListViewRebuildMode::MakeExact, &mut ctx) .vortex_expect("ListViewArray::rebuild(MakeExact) failed in extend_from_array"); debug_assert!(listview.is_zero_copy_to_list()); @@ -311,7 +313,7 @@ impl ArrayBuilder for ListViewBuilder { &array .validity() .vortex_expect("validity_mask in extend_from_array_unchecked") - .execute_mask(array.len(), &mut LEGACY_SESSION.create_execution_ctx()) + .execute_mask(array.len(), &mut ctx) .vortex_expect("Failed to compute validity mask"), ); @@ -343,8 +345,11 @@ impl ArrayBuilder for ListViewBuilder { let uninit_range = self.offsets_builder.uninit_range(extend_length); // This should be cheap because we didn't compress after rebuilding. - #[expect(deprecated)] - let new_offsets = listview.offsets().to_primitive(); + let new_offsets = listview + .offsets() + .clone() + .execute::(&mut ctx) + .vortex_expect("failed to execute list view offsets into a PrimitiveArray"); match_each_integer_ptype!(new_offsets.ptype(), |A| { adjust_and_extend_offsets::( diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 28589ec8c20..3c07650b2ef 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -258,11 +258,11 @@ impl Canonical { /// /// This operation is very expensive and can result in things like allocations, full-scans /// and copy operations. - pub fn compact(&self) -> VortexResult { + pub fn compact(&self, ctx: &mut ExecutionCtx) -> VortexResult { match self { Canonical::VarBinView(array) => Ok(Canonical::VarBinView(array.compact_buffers()?)), Canonical::List(array) => Ok(Canonical::List( - array.rebuild(ListViewRebuildMode::TrimElements)?, + array.rebuild(ListViewRebuildMode::TrimElements, ctx)?, )), _ => Ok(self.clone()), } @@ -1115,6 +1115,7 @@ impl Matcher for AnyCanonical { #[cfg(test)] mod test { use std::sync::Arc; + use std::sync::LazyLock; use arrow_array::Array as ArrowArray; use arrow_array::ArrayRef as ArrowArrayRef; @@ -1134,12 +1135,12 @@ mod test { use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; + use vortex_session::VortexSession; use crate::ArrayRef; use crate::Canonical; use crate::CanonicalValidity; use crate::IntoArray; - use crate::LEGACY_SESSION; use crate::VortexSessionExecute; use crate::arrays::Constant; use crate::arrays::ConstantArray; @@ -1154,6 +1155,11 @@ mod test { use crate::canonical::StructArray; use crate::dtype::Nullability; use crate::scalar::Scalar; + use crate::session::ArraySession; + + /// A shared session for these canonical tests, used to create execution contexts. + static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); fn variant_core_storage(len: usize) -> ArrayRef { ConstantArray::new( @@ -1175,7 +1181,7 @@ mod test { let outer_variant = VariantArray::try_new(variant_core_storage(len), Some(inner_variant.into_array()))?; - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let Canonical::Variant(canonical) = outer_variant .into_array() .execute::(&mut ctx)? @@ -1202,7 +1208,7 @@ mod test { #[test] fn test_canonicalize_nested_struct() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); // Create a struct array with multiple internal components. let nested_struct_array = StructArray::from_fields(&[ ("a", buffer![1u64].into_array()), @@ -1222,7 +1228,7 @@ mod test { ]) .unwrap(); - let arrow_struct = LEGACY_SESSION + let arrow_struct = SESSION .arrow() .execute_arrow(nested_struct_array.into_array(), None, &mut ctx) .unwrap() @@ -1259,7 +1265,7 @@ mod test { #[test] fn roundtrip_struct() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut nulls = NullBufferBuilder::new(6); nulls.append_n_non_nulls(4); nulls.append_null(); @@ -1292,7 +1298,7 @@ mod test { ); let vortex_struct = ArrayRef::from_arrow(&arrow_struct, true).unwrap(); - let vortex_struct = LEGACY_SESSION + let vortex_struct = SESSION .arrow() .execute_arrow(vortex_struct, None, &mut ctx) .unwrap(); @@ -1301,7 +1307,7 @@ mod test { #[test] fn roundtrip_list() { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let names = Arc::new(StringArray::from_iter(vec![ Some("Joseph"), Some("Angela"), @@ -1319,7 +1325,7 @@ mod test { let vortex_list = ArrayRef::from_arrow(&arrow_list, true).unwrap(); - let rt_arrow_list = LEGACY_SESSION + let rt_arrow_list = SESSION .arrow() .execute_arrow(vortex_list, Some(&list_field), &mut ctx) .unwrap(); diff --git a/vortex-bench/src/datasets/struct_list_of_ints.rs b/vortex-bench/src/datasets/struct_list_of_ints.rs index b15cbce6ca0..fe9aa991681 100644 --- a/vortex-bench/src/datasets/struct_list_of_ints.rs +++ b/vortex-bench/src/datasets/struct_list_of_ints.rs @@ -16,7 +16,6 @@ use rand::rngs::StdRng; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; -use vortex::array::LEGACY_SESSION; use vortex::array::VortexSessionExecute; use vortex::array::arrays::Chunked; use vortex::array::arrays::ChunkedArray; @@ -30,6 +29,7 @@ use vortex::array::validity::Validity; use vortex::dtype::FieldNames; use crate::IdempotentPath; +use crate::SESSION; use crate::datasets::Dataset; use crate::idempotent_async; @@ -135,7 +135,7 @@ impl Dataset for StructListOfInts { idempotent_async(&parquet_path, |temp_path| async move { // Generate the data - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let array = self.to_vortex_array(&mut ctx).await?; // Convert to Arrow RecordBatches and write to parquet @@ -145,14 +145,12 @@ impl Dataset for StructListOfInts { let mut writer: Option> = None; for chunk in chunked.iter_chunks() { - let converted = recursive_list_from_list_view(chunk.clone())?; + let converted = recursive_list_from_list_view(chunk.clone(), &mut ctx)?; let schema = converted.dtype().to_arrow_schema()?; let schema = Field::new_struct("", schema.fields, false); - let batch = LEGACY_SESSION.arrow().execute_arrow( - converted, - Some(&schema), - &mut LEGACY_SESSION.create_execution_ctx(), - )?; + let batch = SESSION + .arrow() + .execute_arrow(converted, Some(&schema), &mut ctx)?; let batch = RecordBatch::from(batch.as_struct()); if writer.is_none() { diff --git a/vortex-bench/src/vector_dataset/convert.rs b/vortex-bench/src/vector_dataset/convert.rs index cfef49dfbef..a0125bd923b 100644 --- a/vortex-bench/src/vector_dataset/convert.rs +++ b/vortex-bench/src/vector_dataset/convert.rs @@ -6,6 +6,7 @@ use vortex::array::ArrayRef; use vortex::array::EmptyMetadata; use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; use vortex::array::arrays::Chunked; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::ExtensionArray; @@ -26,6 +27,8 @@ use vortex::error::vortex_bail; use vortex::error::vortex_err; use vortex_tensor::vector::Vector; +use crate::SESSION; + /// Rewrap a list-of-float column as a [`vortex_tensor::vector::Vector`] extension array. /// /// Parquet has no fixed-size list logical type, so an embedding column ingested via @@ -75,7 +78,7 @@ pub fn list_to_vector_ext(input: ArrayRef) -> VortexResult { // `parquet_to_vortex_chunks` produces `ListView` arrays for list columns by default; // materialize them into a flat `List` representation before we validate offsets. if input.as_opt::().is_some() { - let flat = recursive_list_from_list_view(input)?; + let flat = recursive_list_from_list_view(input, &mut SESSION.create_execution_ctx())?; return list_to_vector_ext(flat); } diff --git a/vortex-compressor/src/compressor.rs b/vortex-compressor/src/compressor.rs index a557736573a..33502104b13 100644 --- a/vortex-compressor/src/compressor.rs +++ b/vortex-compressor/src/compressor.rs @@ -120,7 +120,7 @@ impl CascadingCompressor { let _enter = span.enter(); let canonical = array.clone().execute::(exec_ctx)?.0; - let compact = canonical.compact()?; + let compact = canonical.compact(exec_ctx)?; let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); @@ -151,7 +151,7 @@ impl CascadingCompressor { } let canonical = child.clone().execute::(exec_ctx)?.0; - let compact = canonical.compact()?; + let compact = canonical.compact(exec_ctx)?; let child_ctx = parent_ctx .clone() @@ -197,7 +197,7 @@ impl CascadingCompressor { } Canonical::List(list_view_array) => { if list_view_array.is_zero_copy_to_list() || list_view_array.elements().is_empty() { - let list_array = list_from_list_view(list_view_array)?; + let list_array = list_from_list_view(list_view_array, exec_ctx)?; self.compress_list_array(list_array, compress_ctx, exec_ctx) } else { self.compress_list_view_array(list_view_array, compress_ctx, exec_ctx) @@ -495,7 +495,7 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let list_array = list_array.reset_offsets(true)?; + let list_array = list_array.reset_offsets(true, exec_ctx)?; let compressed_elems = self.compress(list_array.elements(), exec_ctx)?; diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index af180400675..de20df48ecc 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -170,7 +170,8 @@ fn export_canonical( return Err(gpu_err); } - export_list(list_from_list_view(listview)?, ctx).await + let list = list_from_list_view(listview, ctx.execution_ctx())?; + export_list(list, ctx).await } Canonical::FixedSizeList(fixed_size_list) => { export_fixed_size_list(fixed_size_list, ctx).await diff --git a/vortex-duckdb/src/exporter/list_view.rs b/vortex-duckdb/src/exporter/list_view.rs index a0ee1614001..a4cb61895f3 100644 --- a/vortex-duckdb/src/exporter/list_view.rs +++ b/vortex-duckdb/src/exporter/list_view.rs @@ -75,7 +75,7 @@ pub(crate) fn new_exporter( } } else if array.upper_bound_density(ctx)? < DEFAULT_REBUILD_DENSITY_THRESHOLD { // Overlaps, gaps, or garbage may be present, so a full rebuild is needed to reclaim waste. - array.rebuild(ListViewRebuildMode::MakeZeroCopyToList)? + array.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)? } else { array }; From 01701be172528a30f074a9ac4e9156c0bee27489 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 8 Jun 2026 12:11:19 -0400 Subject: [PATCH 2/2] clean up execution context usage in benchmarks Hoist `create_execution_ctx()` out of the divan-timed closures into the untimed `with_inputs` setup (using `bench_values` / `bench_local_values` where a by-value context is needed) so the benchmarks measure the operation under test rather than per-sample context creation. Replace every non-CUDA benchmark use of the deprecated global `LEGACY_SESSION` with a local `static SESSION` built from `VortexSession::empty().with::()`, which is exactly what `LEGACY_SESSION` is defined as, so the change is behavior-preserving while reducing reliance on the legacy global. Document why `ListViewBuilder::extend_from_array_unchecked` still mints a `LEGACY_SESSION` context: the `ArrayBuilder` trait does not thread an `ExecutionCtx` through its extend methods. Signed-off-by: Connor Tsui --- encodings/alp/benches/alp_compress.rs | 32 ++- .../experimental/onpair/benches/decode.rs | 20 +- .../fastlanes/benches/bitpack_compare.rs | 10 +- .../benches/bitpack_compare_sweep.rs | 10 +- .../fastlanes/benches/bitpacking_take.rs | 78 +++---- .../fastlanes/benches/compute_between.rs | 24 +- encodings/fsst/benches/fsst_compress.rs | 29 +-- encodings/runend/benches/run_end_compress.rs | 28 +-- encodings/runend/benches/run_end_decode.rs | 23 +- .../runend/benches/run_end_null_count.rs | 18 +- encodings/sparse/benches/sparse_canonical.rs | 11 +- encodings/zstd/benches/listview_rebuild.rs | 11 +- vortex-array/benches/aggregate_max.rs | 60 ++--- vortex-array/benches/aggregate_sum.rs | 60 ++--- vortex-array/benches/cast_primitive.rs | 27 ++- .../benches/chunked_fsl_canonicalize.rs | 10 +- vortex-array/benches/dict_compare.rs | 16 +- vortex-array/benches/dict_compress.rs | 47 ++-- vortex-array/benches/filter_bool.rs | 40 ++-- vortex-array/benches/listview_rebuild.rs | 207 +++++++++--------- vortex-array/benches/scalar_at_struct.rs | 22 +- vortex-array/benches/scalar_subtract.rs | 10 +- vortex-array/benches/take_fsl.rs | 12 +- vortex-array/benches/take_patches.rs | 14 +- vortex-array/benches/take_primitive.rs | 26 +-- vortex-array/benches/take_struct.rs | 14 +- vortex-array/benches/to_arrow.rs | 32 +-- vortex-array/benches/varbinview_zip.rs | 12 +- .../src/arrays/listview/conversion.rs | 9 +- vortex-array/src/builders/listview.rs | 4 + vortex-compressor/benches/dict_encode.rs | 16 +- vortex-compressor/benches/stats_calc.rs | 12 +- vortex-row/benches/row_encode.rs | 31 +-- .../common_encoding_tree_throughput.rs | 25 ++- vortex/benches/single_encoding_throughput.rs | 69 ++---- 35 files changed, 571 insertions(+), 498 deletions(-) diff --git a/encodings/alp/benches/alp_compress.rs b/encodings/alp/benches/alp_compress.rs index 7fbe60d5997..70a566ab0f8 100644 --- a/encodings/alp/benches/alp_compress.rs +++ b/encodings/alp/benches/alp_compress.rs @@ -3,6 +3,8 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng as _; @@ -14,13 +16,14 @@ use vortex_alp::alp_encode; use vortex_alp::decompress_into_array; 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::dtype::NativePType; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; +use vortex_session::VortexSession; fn main() { divan::main(); @@ -48,6 +51,9 @@ const BENCH_ARGS: &[(usize, f64, f64)] = &[ (10_000, 0.1, 1.0), ]; +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[divan::bench(types = [f32, f64], args = BENCH_ARGS)] fn compress_alp(bencher: Bencher, args: (usize, f64, f64)) { let (n, fraction_patch, fraction_valid) = args; @@ -68,14 +74,9 @@ fn compress_alp(bencher: Bencher, args: (usize, f64, let values = values.freeze(); let array = PrimitiveArray::new(values, validity); - bencher.with_inputs(|| &array).bench_values(|array| { - alp_encode( - array.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap() - }) + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| alp_encode(array.as_view(), None, &mut ctx).unwrap()) } #[divan::bench(types = [f32, f64], args = BENCH_ARGS)] @@ -102,10 +103,10 @@ fn decompress_alp(bencher: Bencher, args: (usize, f64 alp_encode( PrimitiveArray::new(Buffer::copy_from(&values), validity.clone()).as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_values(|(v, mut ctx)| decompress_into_array(v, &mut ctx)); @@ -143,7 +144,7 @@ fn compress_rd(bencher: Bencher, args: (usize, f64) let encoder = RDEncoder::new(primitive.as_slice::()); bencher - .with_inputs(|| (&primitive, &encoder, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&primitive, &encoder, SESSION.create_execution_ctx())) .bench_refs(|(primitive, encoder, ctx)| encoder.encode(primitive.as_view(), ctx)) } @@ -152,12 +153,9 @@ fn decompress_rd(bencher: Bencher, args: (usize, f6 let (n, fraction_patch) = args; let primitive = make_rd_array::(n, fraction_patch); let encoder = RDEncoder::new(primitive.as_slice::()); - let encoded = encoder.encode( - primitive.as_view(), - &mut LEGACY_SESSION.create_execution_ctx(), - ); + let encoded = encoder.encode(primitive.as_view(), &mut SESSION.create_execution_ctx()); bencher - .with_inputs(|| (&encoded, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) .bench_refs(|(encoded, ctx)| (**encoded).clone().into_array().execute::(ctx)); } diff --git a/encodings/experimental/onpair/benches/decode.rs b/encodings/experimental/onpair/benches/decode.rs index 2f3fce9db23..ff8db8c22c2 100644 --- a/encodings/experimental/onpair/benches/decode.rs +++ b/encodings/experimental/onpair/benches/decode.rs @@ -213,9 +213,8 @@ fn canonicalize_to_varbinview(bencher: Bencher, case: (Shape, usize)) { let (shape, n) = case; let arr = compress(n, shape); bencher - .with_inputs(|| arr.clone().into_array()) - .bench_local_values(|arr| { - let mut ctx = SESSION.create_execution_ctx(); + .with_inputs(|| (arr.clone().into_array(), SESSION.create_execution_ctx())) + .bench_local_values(|(arr, mut ctx)| { divan::black_box( arr.execute::(&mut ctx) .unwrap_or_else(|e| panic!("canonicalize failed: {e}")), @@ -234,13 +233,14 @@ fn filter_share_dict(bencher: Bencher, case: (Shape, usize)) { let (shape, n) = case; let arr = compress(n, shape); let mask = Mask::from_iter((0..n).map(|i| i % 7 == 0)); - bencher.bench_local(|| { - let mut ctx = SESSION.create_execution_ctx(); - let result = ::filter(arr.as_view(), &mask, &mut ctx) - .unwrap() - .unwrap(); - divan::black_box(result); - }); + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| { + let result = ::filter(arr.as_view(), &mask, &mut ctx) + .unwrap() + .unwrap(); + divan::black_box(result); + }); } fn main() { diff --git a/encodings/fastlanes/benches/bitpack_compare.rs b/encodings/fastlanes/benches/bitpack_compare.rs index 939c401ee38..c704f2b35ed 100644 --- a/encodings/fastlanes/benches/bitpack_compare.rs +++ b/encodings/fastlanes/benches/bitpack_compare.rs @@ -10,33 +10,39 @@ #![expect(clippy::unwrap_used)] #![expect(clippy::cast_possible_truncation)] +use std::sync::LazyLock; + use divan::Bencher; use divan::counter::ItemsCount; use vortex_array::ArrayRef; 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; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_fastlanes::BitPackedData; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + const LENS: &[usize] = &[1024, 64 * 1024]; const BIT_WIDTHS: &[u8] = &[4, 16]; /// Build a packed array of varied in-range values, plus an out-of-range constant RHS for /// the fast-path benches. fn build_inputs(len: usize) -> (ArrayRef, ArrayRef, ExecutionCtx) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let buf: BufferMut = (0..len).map(|i| (i as u32) % (1 << BW)).collect(); let array = BitPackedData::encode( &PrimitiveArray::new(buf.freeze(), Validity::NonNullable).into_array(), diff --git a/encodings/fastlanes/benches/bitpack_compare_sweep.rs b/encodings/fastlanes/benches/bitpack_compare_sweep.rs index 3eb0ba3b9a2..0be525566e7 100644 --- a/encodings/fastlanes/benches/bitpack_compare_sweep.rs +++ b/encodings/fastlanes/benches/bitpack_compare_sweep.rs @@ -12,12 +12,13 @@ #![expect(clippy::unwrap_used)] #![expect(clippy::cast_possible_truncation)] +use std::sync::LazyLock; + use divan::Bencher; use divan::counter::ItemsCount; use vortex_array::ArrayRef; 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; @@ -26,14 +27,19 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::NativePType; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_fastlanes::BitPackedData; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + /// Number of elements per benchmarked array (64 full FastLanes blocks). const LEN: usize = 64 * 1024; @@ -62,7 +68,7 @@ impl_bench_int!(u8, u16, u32, u64, i8, i16, i32, i64); /// Encode `LEN` in-range values of type `T` at the given bit width, returning the packed array, a /// mid-range constant to compare against, and an execution context. fn setup(width: usize) -> (ArrayRef, ArrayRef, ExecutionCtx) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let cap = 1u64 << width; let buf: BufferMut = (0..LEN) .map(|i| T::from_counter((i as u64) % cap)) diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index a5a7803bcaf..2e01f1253e6 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -4,37 +4,42 @@ #![expect(clippy::unwrap_used)] #![expect(clippy::cast_possible_truncation)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; use vortex_array::IntoArray as _; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_fastlanes::BitPackedArrayExt; use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[divan::bench] fn take_10_stratified(bencher: Bencher) { let values = fixture(65_536, 8); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let indices = PrimitiveArray::from_iter((0..10).map(|i| i * 6_553)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -49,12 +54,11 @@ fn take_10_contiguous(bencher: Bencher) { let values = fixture(65_536, 8); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let indices = buffer![0..10].into_array(); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone()) @@ -70,14 +74,13 @@ fn take_10k_random(bencher: Bencher) { let range = Uniform::new(0, values.len()).unwrap(); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let rng = StdRng::seed_from_u64(0); let indices = PrimitiveArray::from_iter(rng.sample_iter(range).take(10_000).map(|i| i as u32)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -92,12 +95,11 @@ fn take_10k_contiguous(bencher: Bencher) { let values = fixture(65_536, 8); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let indices = PrimitiveArray::from_iter(0..10_000); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -112,12 +114,11 @@ fn take_10k_dispersed(bencher: Bencher) { let values = fixture(65_536, 8); let uncompressed = PrimitiveArray::new(values.clone(), Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let indices = PrimitiveArray::from_iter((0..10_000).map(|i| (i * 42) % values.len() as u64)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -132,12 +133,11 @@ fn take_10k_first_chunk_only(bencher: Bencher) { let values = fixture(65_536, 8); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let indices = PrimitiveArray::from_iter((0..10_000).map(|i| ((i * 42) % 1024) as u64)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -172,8 +172,7 @@ fn patched_take_10_stratified(bencher: Bencher) { let values = (0u32..BIG_BASE2 + NUM_EXCEPTIONS).collect::>(); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); assert!(packed.patches().is_some()); assert_eq!( @@ -184,7 +183,7 @@ fn patched_take_10_stratified(bencher: Bencher) { let indices = PrimitiveArray::from_iter((0..10).map(|i| i * 6_653)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -199,8 +198,7 @@ fn patched_take_10_contiguous(bencher: Bencher) { let values = (0u32..BIG_BASE2 + NUM_EXCEPTIONS).collect::>(); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); assert!(packed.patches().is_some()); assert_eq!( @@ -211,7 +209,7 @@ fn patched_take_10_contiguous(bencher: Bencher) { let indices = buffer![0..10].into_array(); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone()) @@ -226,15 +224,14 @@ fn patched_take_10k_random(bencher: Bencher) { let values = (0u32..BIG_BASE2 + NUM_EXCEPTIONS).collect::>(); let uncompressed = PrimitiveArray::new(values.clone(), Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let rng = StdRng::seed_from_u64(0); let range = Uniform::new(0, values.len()).unwrap(); let indices = PrimitiveArray::from_iter(rng.sample_iter(range).take(10_000).map(|i| i as u32)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -249,12 +246,11 @@ fn patched_take_10k_contiguous_not_patches(bencher: Bencher) { let values = (0u32..BIG_BASE2 + NUM_EXCEPTIONS).collect::>(); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let indices = PrimitiveArray::from_iter((0u32..NUM_EXCEPTIONS).cycle().take(10000)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -269,8 +265,7 @@ fn patched_take_10k_contiguous_patches(bencher: Bencher) { let values = (0u32..BIG_BASE2 + NUM_EXCEPTIONS).collect::>(); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); assert!(packed.patches().is_some()); assert_eq!( @@ -282,7 +277,7 @@ fn patched_take_10k_contiguous_patches(bencher: Bencher) { PrimitiveArray::from_iter((BIG_BASE2..BIG_BASE2 + NUM_EXCEPTIONS).cycle().take(10000)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -297,12 +292,11 @@ fn patched_take_10k_dispersed(bencher: Bencher) { let values = (0u32..BIG_BASE2 + NUM_EXCEPTIONS).collect::>(); let uncompressed = PrimitiveArray::new(values.clone(), Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let indices = PrimitiveArray::from_iter((0..10_000).map(|i| (i * 42) % values.len() as u64)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -317,12 +311,11 @@ fn patched_take_10k_first_chunk_only(bencher: Bencher) { let values = (0u32..BIG_BASE2 + NUM_EXCEPTIONS).collect::>(); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let indices = PrimitiveArray::from_iter((0..10_000).map(|i| ((i * 42) % 1024) as u64)); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) @@ -337,8 +330,7 @@ fn patched_take_10k_adversarial(bencher: Bencher) { let values = (0u32..BIG_BASE2 + NUM_EXCEPTIONS).collect::>(); let uncompressed = PrimitiveArray::new(values, Validity::NonNullable); let packed = - bitpack_to_best_bit_width(&uncompressed, &mut LEGACY_SESSION.create_execution_ctx()) - .unwrap(); + bitpack_to_best_bit_width(&uncompressed, &mut SESSION.create_execution_ctx()).unwrap(); let per_chunk_count = 100; let indices = PrimitiveArray::from_iter( (0..(NUM_EXCEPTIONS + 1024) / 1024) @@ -349,7 +341,7 @@ fn patched_take_10k_adversarial(bencher: Bencher) { ); bencher - .with_inputs(|| (&packed, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&packed, &indices, SESSION.create_execution_ctx())) .bench_refs(|(packed, indices, execution_ctx)| { packed .take(indices.clone().into_array()) diff --git a/encodings/fastlanes/benches/compute_between.rs b/encodings/fastlanes/benches/compute_between.rs index 1630764ff8e..5a1bdcee655 100644 --- a/encodings/fastlanes/benches/compute_between.rs +++ b/encodings/fastlanes/benches/compute_between.rs @@ -2,6 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use num_traits::NumCast; use rand::RngExt; use rand::rngs::StdRng; @@ -11,17 +13,21 @@ use vortex_alp::ALPArraySlotsExt; use vortex_alp::alp_encode; 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::dtype::NativePType; +use vortex_array::session::ArraySession; use vortex_error::VortexExpect; use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + fn generate_primitive_array( rng: &mut StdRng, len: usize, @@ -39,7 +45,7 @@ fn generate_bit_pack_primitive_array( .map(|_| T::from_usize(rng.random_range(0..10_000)).vortex_expect("")) .collect::(); - bitpack_to_best_bit_width(&a, &mut LEGACY_SESSION.create_execution_ctx()) + bitpack_to_best_bit_width(&a, &mut SESSION.create_execution_ctx()) .vortex_expect("") .into_array() } @@ -52,7 +58,7 @@ fn generate_alp_bit_pack_primitive_array( .map(|_| T::from_usize(rng.random_range(0..10_000)).vortex_expect("")) .collect::(); - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let alp = alp_encode(a.as_view(), None, &mut ctx).vortex_expect(""); let encoded = alp @@ -75,7 +81,6 @@ mod primitive { use rand::SeedableRng; use rand::prelude::StdRng; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; @@ -86,6 +91,7 @@ mod primitive { use vortex_error::VortexExpect; use crate::BENCH_ARGS; + use crate::SESSION; use crate::generate_primitive_array; #[divan::bench( @@ -103,7 +109,7 @@ mod primitive { let arr = generate_primitive_array::(&mut rng, len); bencher - .with_inputs(|| (&arr, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&arr, SESSION.create_execution_ctx())) .bench_refs(|(arr, ctx)| { arr.clone() .into_array() @@ -128,7 +134,6 @@ mod bitpack { use rand::SeedableRng; use rand::prelude::StdRng; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; @@ -139,6 +144,7 @@ mod bitpack { use vortex_error::VortexExpect; use crate::BENCH_ARGS; + use crate::SESSION; use crate::generate_bit_pack_primitive_array; #[divan::bench( @@ -156,7 +162,7 @@ mod bitpack { let arr = generate_bit_pack_primitive_array::(&mut rng, len); bencher - .with_inputs(|| (&arr, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&arr, SESSION.create_execution_ctx())) .bench_refs(|(arr, ctx)| { arr.clone() .between( @@ -180,7 +186,6 @@ mod alp { use rand::SeedableRng; use rand::prelude::StdRng; use vortex_array::IntoArray; - use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; @@ -191,6 +196,7 @@ mod alp { use vortex_error::VortexExpect; use crate::BENCH_ARGS; + use crate::SESSION; use crate::generate_alp_bit_pack_primitive_array; #[divan::bench( @@ -208,7 +214,7 @@ mod alp { let arr = generate_alp_bit_pack_primitive_array::(&mut rng, len); bencher - .with_inputs(|| (&arr, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&arr, SESSION.create_execution_ctx())) .bench_refs(|(arr, ctx)| { arr.clone() .between( diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index 45aedfa464f..88e4e26a903 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -11,7 +11,6 @@ use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::Canonical; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; @@ -57,7 +56,7 @@ fn compress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usize let array = generate_test_data(string_count, avg_len, unique_chars); let compressor = fsst_train_compressor(&array); bencher - .with_inputs(|| (&array, &compressor, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&array, &compressor, SESSION.create_execution_ctx())) .bench_refs(|(array, compressor, ctx)| { fsst_compress(*array, array.len(), array.dtype(), compressor, ctx) }) @@ -74,11 +73,11 @@ fn decompress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usi len, &dtype, &compressor, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ); bencher - .with_inputs(|| (&encoded, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&encoded, SESSION.create_execution_ctx())) .bench_refs(|(encoded, ctx)| (**encoded).clone().into_array().execute::(ctx)) } @@ -99,18 +98,12 @@ fn pushdown_compare(bencher: Bencher, (string_count, avg_len, unique_chars): (us array.len(), array.dtype(), &compressor, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ); let constant = ConstantArray::new(Scalar::from(&b"const"[..]), array.len()); bencher - .with_inputs(|| { - ( - &fsst_array, - &constant, - LEGACY_SESSION.create_execution_ctx(), - ) - }) + .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { fsst_array .clone() @@ -134,18 +127,12 @@ fn canonicalize_compare( array.len(), array.dtype(), &compressor, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ); let constant = ConstantArray::new(Scalar::from(&b"const"[..]), array.len()); bencher - .with_inputs(|| { - ( - &fsst_array, - &constant, - LEGACY_SESSION.create_execution_ctx(), - ) - }) + .with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx())) .bench_refs(|(fsst_array, constant, ctx)| { (*fsst_array) .clone() @@ -235,7 +222,7 @@ fn generate_chunked_test_data( avg_len: usize, unique_chars: u8, ) -> ChunkedArray { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); (0..chunk_size) .map(|_| { let array = generate_test_data(string_count, avg_len, unique_chars); diff --git a/encodings/runend/benches/run_end_compress.rs b/encodings/runend/benches/run_end_compress.rs index 1c183f3e5fd..a49a5cfcb50 100644 --- a/encodings/runend/benches/run_end_compress.rs +++ b/encodings/runend/benches/run_end_compress.rs @@ -3,24 +3,30 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use itertools::repeat_n; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::IntegerPType; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_runend::RunEnd; use vortex_runend::compress::runend_encode; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + const BENCH_ARGS: &[(usize, usize)] = &[ (1000, 4), (1000, 16), @@ -47,7 +53,7 @@ fn compress(bencher: Bencher, (length, run_step): (usize, usize)) { ); bencher - .with_inputs(|| (&values, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&values, SESSION.create_execution_ctx())) .bench_refs(|(values, ctx)| runend_encode(values.as_view(), ctx)); } @@ -64,11 +70,11 @@ fn decompress(bencher: Bencher, (length, run_step): (usize, usi .collect::>() .into_array(); - let run_end_array = RunEnd::new(ends, values, &mut LEGACY_SESSION.create_execution_ctx()); + let run_end_array = RunEnd::new(ends, values, &mut SESSION.create_execution_ctx()); let array = run_end_array.into_array(); bencher - .with_inputs(|| (array.clone(), LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) .bench_values(|(array, mut execution_ctx)| { array .execute::(&mut execution_ctx) @@ -88,20 +94,14 @@ fn take_indices(bencher: Bencher, (length, run_step): (usize, usize)) { ); let source_array = PrimitiveArray::from_iter(0..(length as i32)).into_array(); - let mut encode_ctx = LEGACY_SESSION.create_execution_ctx(); + let mut encode_ctx = SESSION.create_execution_ctx(); let (ends, values) = runend_encode(values.as_view(), &mut encode_ctx); let runend_array = RunEnd::try_new(ends.into_array(), values, &mut encode_ctx) .unwrap() .into_array(); bencher - .with_inputs(|| { - ( - &source_array, - &runend_array, - LEGACY_SESSION.create_execution_ctx(), - ) - }) + .with_inputs(|| (&source_array, &runend_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array .take(indices.clone()) @@ -122,11 +122,11 @@ fn decompress_utf8(bencher: Bencher, (length, run_step): (usize, usize)) { let values = VarBinViewArray::from_iter_str((0..num_runs).map(|i| format!("run_value_{i}"))) .into_array(); - let run_end_array = RunEnd::new(ends, values, &mut LEGACY_SESSION.create_execution_ctx()); + let run_end_array = RunEnd::new(ends, values, &mut SESSION.create_execution_ctx()); let array = run_end_array.into_array(); bencher - .with_inputs(|| (array.clone(), LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) .bench_values(|(array, mut execution_ctx)| { array .execute::(&mut execution_ctx) diff --git a/encodings/runend/benches/run_end_decode.rs b/encodings/runend/benches/run_end_decode.rs index 818c8eb4ed6..f509e55aed3 100644 --- a/encodings/runend/benches/run_end_decode.rs +++ b/encodings/runend/benches/run_end_decode.rs @@ -4,21 +4,26 @@ #![expect(clippy::cast_possible_truncation)] use std::fmt; +use std::sync::LazyLock; use divan::Bencher; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_runend::decompress_bool::runend_decode_bools; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + /// Distribution types for bool benchmarks #[derive(Clone, Copy)] enum BoolDistribution { @@ -210,13 +215,7 @@ fn decode_bool(bencher: Bencher, args: BoolBenchArgs) { } = args; let (ends, values) = create_bool_test_data(total_length, avg_run_length, distribution); bencher - .with_inputs(|| { - ( - ends.clone(), - values.clone(), - LEGACY_SESSION.create_execution_ctx(), - ) - }) + .with_inputs(|| (ends.clone(), values.clone(), SESSION.create_execution_ctx())) .bench_refs(|(ends, values, ctx)| { runend_decode_bools(ends.clone(), values.clone(), 0, total_length, ctx) }); @@ -379,13 +378,7 @@ fn decode_bool_nullable(bencher: Bencher, args: NullableBoolBenchArgs) { let (ends, values) = create_nullable_bool_test_data(total_length, avg_run_length, distribution, validity); bencher - .with_inputs(|| { - ( - ends.clone(), - values.clone(), - LEGACY_SESSION.create_execution_ctx(), - ) - }) + .with_inputs(|| (ends.clone(), values.clone(), SESSION.create_execution_ctx())) .bench_refs(|(ends, values, ctx)| { runend_decode_bools(ends.clone(), values.clone(), 0, total_length, ctx) }); diff --git a/encodings/runend/benches/run_end_null_count.rs b/encodings/runend/benches/run_end_null_count.rs index c60c72c5e0e..5392e04189b 100644 --- a/encodings/runend/benches/run_end_null_count.rs +++ b/encodings/runend/benches/run_end_null_count.rs @@ -3,17 +3,20 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; +use vortex_array::session::ArraySession; use vortex_buffer::Buffer; use vortex_runend::RunEnd; use vortex_runend::RunEndArray; +use vortex_session::VortexSession; fn main() { divan::main(); @@ -47,15 +50,16 @@ const BENCH_ARGS: &[(usize, usize, f64)] = &[ (100_000, 1024, 0.5), ]; +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[divan::bench(args = BENCH_ARGS)] fn null_count_run_end(bencher: Bencher, (n, run_step, valid_density): (usize, usize, f64)) { let array = fixture(n, run_step, valid_density).into_array(); - bencher.with_inputs(|| &array).bench_refs(|array| { - array - .invalid_count(&mut LEGACY_SESSION.create_execution_ctx()) - .unwrap() - }); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| array.invalid_count(ctx).unwrap()); } fn fixture(n: usize, run_step: usize, valid_density: f64) -> RunEndArray { @@ -72,5 +76,5 @@ fn fixture(n: usize, run_step: usize, valid_density: f64) -> RunEndArray { ) .into_array(); - RunEnd::new(ends, values, &mut LEGACY_SESSION.create_execution_ctx()) + RunEnd::new(ends, values, &mut SESSION.create_execution_ctx()) } diff --git a/encodings/sparse/benches/sparse_canonical.rs b/encodings/sparse/benches/sparse_canonical.rs index ba1489a0e5a..3861e26da06 100644 --- a/encodings/sparse/benches/sparse_canonical.rs +++ b/encodings/sparse/benches/sparse_canonical.rs @@ -4,12 +4,12 @@ #![expect(clippy::cast_possible_truncation)] use std::sync::Arc; +use std::sync::LazyLock; use divan::Bencher; 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::FixedSizeListArray; use vortex_array::arrays::ListViewArray; @@ -17,15 +17,20 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType::I32; use vortex_array::scalar::Scalar; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; +use vortex_session::VortexSession; use vortex_sparse::Sparse; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + const LIST_ARGS: &[(usize, usize, usize)] = &[ // len, patch_stride, list_size (512, 7, 4), @@ -93,7 +98,7 @@ fn canonicalize_sparse_list( let sparse = make_sparse_list(len, patch_stride, list_size); bencher - .with_inputs(|| (sparse.clone(), LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (sparse.clone(), SESSION.create_execution_ctx())) .bench_values(|(array, mut ctx)| { divan::black_box( array @@ -111,7 +116,7 @@ fn canonicalize_sparse_fixed_size_list( let sparse = make_sparse_fixed_size_list(len, patch_stride, list_size); bencher - .with_inputs(|| (sparse.clone(), LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (sparse.clone(), SESSION.create_execution_ctx())) .bench_values(|(array, mut ctx)| { divan::black_box( array diff --git a/encodings/zstd/benches/listview_rebuild.rs b/encodings/zstd/benches/listview_rebuild.rs index d59af116188..fe4e9f063d1 100644 --- a/encodings/zstd/benches/listview_rebuild.rs +++ b/encodings/zstd/benches/listview_rebuild.rs @@ -48,12 +48,11 @@ fn rebuild_naive(bencher: Bencher) { let list_view = ListViewArray::new(dudes, offsets, sizes, Validity::NonNullable); - bencher.with_inputs(|| &list_view).bench_refs(|list_view| { - list_view.rebuild( - ListViewRebuildMode::MakeZeroCopyToList, - &mut SESSION.create_execution_ctx(), - ) - }) + bencher + .with_inputs(|| (&list_view, SESSION.create_execution_ctx())) + .bench_refs(|(list_view, ctx)| { + list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + }) } fn main() { diff --git a/vortex-array/benches/aggregate_max.rs b/vortex-array/benches/aggregate_max.rs index 8962ade2974..14bf13792dd 100644 --- a/vortex-array/benches/aggregate_max.rs +++ b/vortex-array/benches/aggregate_max.rs @@ -25,11 +25,13 @@ fn max_i32(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(1); let data: Vec = (0..N).map(|_| rng.random::()).collect(); bencher - .with_inputs(|| PrimitiveArray::from_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_max::(&mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_max::(ctx)); } #[divan::bench] @@ -37,11 +39,13 @@ fn max_i64(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(2); let data: Vec = (0..N).map(|_| rng.random::()).collect(); bencher - .with_inputs(|| PrimitiveArray::from_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_max::(&mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_max::(ctx)); } #[divan::bench] @@ -49,11 +53,13 @@ fn max_f64(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(3); let data: Vec = (0..N).map(|_| rng.random::()).collect(); bencher - .with_inputs(|| PrimitiveArray::from_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_max::(&mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_max::(ctx)); } // Clustered nulls: long valid runs broken up by null blocks (run-based path's best case). @@ -70,11 +76,13 @@ fn max_i32_nulls_clustered(bencher: Bencher) { }) .collect(); bencher - .with_inputs(|| PrimitiveArray::from_option_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_max::(&mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_max::(ctx)); } // Scattered nulls: ~50% random nulls producing many short runs (run-based path's worst case). @@ -85,9 +93,11 @@ fn max_i32_nulls_scattered(bencher: Bencher) { .map(|_| rng.random_bool(0.5).then(|| rng.random::())) .collect(); bencher - .with_inputs(|| PrimitiveArray::from_option_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_max::(&mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_max::(ctx)); } diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index db4bf5b284a..7d0e249c1ff 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -26,11 +26,13 @@ fn sum_i32(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(1); let data: Vec = (0..N).map(|_| rng.random_range(-1000..1000)).collect(); bencher - .with_inputs(|| PrimitiveArray::from_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_as::(Stat::Sum, &mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } #[divan::bench] @@ -38,11 +40,13 @@ fn sum_u32(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(2); let data: Vec = (0..N).map(|_| rng.random_range(0..2000)).collect(); bencher - .with_inputs(|| PrimitiveArray::from_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_as::(Stat::Sum, &mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } #[divan::bench] @@ -50,11 +54,13 @@ fn sum_i64(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(3); let data: Vec = (0..N).map(|_| rng.random_range(-1000..1000)).collect(); bencher - .with_inputs(|| PrimitiveArray::from_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_as::(Stat::Sum, &mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } // Clustered nulls: long runs of valid values broken up by occasional null blocks. This is the @@ -72,11 +78,13 @@ fn sum_i32_nulls_clustered(bencher: Bencher) { }) .collect(); bencher - .with_inputs(|| PrimitiveArray::from_option_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_as::(Stat::Sum, &mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } // Scattered nulls: ~50% nulls placed at random, producing many short runs. This is the worst case @@ -88,9 +96,11 @@ fn sum_i32_nulls_scattered(bencher: Bencher) { .map(|_| rng.random_bool(0.5).then(|| rng.random_range(-1000..1000))) .collect(); bencher - .with_inputs(|| PrimitiveArray::from_option_iter(data.iter().copied()).into_array()) - .bench_refs(|a| { - a.statistics() - .compute_as::(Stat::Sum, &mut SESSION.create_execution_ctx()) - }); + .with_inputs(|| { + ( + PrimitiveArray::from_option_iter(data.iter().copied()).into_array(), + SESSION.create_execution_ctx(), + ) + }) + .bench_refs(|(a, ctx)| a.statistics().compute_as::(Stat::Sum, ctx)); } diff --git a/vortex-array/benches/cast_primitive.rs b/vortex-array/benches/cast_primitive.rs index 86895fb2ce7..060a8290a2a 100644 --- a/vortex-array/benches/cast_primitive.rs +++ b/vortex-array/benches/cast_primitive.rs @@ -1,11 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::LazyLock; + use divan::Bencher; use rand::prelude::*; 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::builtins::ArrayBuiltins; @@ -13,6 +14,8 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::stats::Stat; +use vortex_array::session::ArraySession; +use vortex_session::VortexSession; fn main() { divan::main(); @@ -20,6 +23,9 @@ fn main() { const N: usize = 100_000; +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[divan::bench] fn cast_u16_to_u32(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(42); @@ -34,15 +40,14 @@ fn cast_u16_to_u32(bencher: Bencher) { .into_array(); // Pre-compute min/max so values_fit_in is a cache hit during the benchmark. arr.statistics() - .compute_all( - &[Stat::Min, Stat::Max], - &mut LEGACY_SESSION.create_execution_ctx(), - ) + .compute_all(&[Stat::Min, Stat::Max], &mut SESSION.create_execution_ctx()) .ok(); - bencher.with_inputs(|| arr.clone()).bench_refs(|a| { - #[expect(clippy::unwrap_used)] - a.cast(DType::Primitive(PType::U32, Nullability::Nullable)) - .unwrap() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) - }); + bencher + .with_inputs(|| (arr.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| { + #[expect(clippy::unwrap_used)] + a.cast(DType::Primitive(PType::U32, Nullability::Nullable)) + .unwrap() + .execute::(ctx) + }); } diff --git a/vortex-array/benches/chunked_fsl_canonicalize.rs b/vortex-array/benches/chunked_fsl_canonicalize.rs index d267415cc4e..ac9962046de 100644 --- a/vortex-array/benches/chunked_fsl_canonicalize.rs +++ b/vortex-array/benches/chunked_fsl_canonicalize.rs @@ -10,20 +10,26 @@ #![expect(clippy::cast_possible_truncation)] #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use vortex_array::Canonical; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + /// Number of lists in each chunk. const LISTS_PER_CHUNK: usize = 1_000; @@ -60,7 +66,7 @@ fn canonicalize(bencher: Bencher, num_chunks: usize) { let chunked = create_chunked_fsl(LIST_SIZE, num_chunks).into_array(); bencher - .with_inputs(|| (&chunked, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&chunked, SESSION.create_execution_ctx())) .bench_refs(|(array, execution_ctx)| { array.clone().execute::(execution_ctx).unwrap() }); diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index 54a3faf8224..401393495dd 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -4,10 +4,10 @@ #![expect(clippy::unwrap_used)] use std::str::from_utf8; +use std::sync::LazyLock; use vortex_array::Canonical; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::accessor::ArrayAccessor; @@ -22,12 +22,16 @@ use vortex_array::expr::eq; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + const LENGTH_AND_UNIQUE_VALUES: &[(usize, usize)] = &[ // length, unique_values (10_000, 2), @@ -51,7 +55,7 @@ fn bench_compare_primitive(bencher: divan::Bencher, (len, uniqueness): (usize, u let primitive_arr = gen_primitive_for_dict::(len, uniqueness); let dict = dict_encode( &primitive_arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let value = primitive_arr.as_slice::()[0]; @@ -74,7 +78,7 @@ fn bench_compare_varbin(bencher: divan::Bencher, (len, uniqueness): (usize, usiz let varbin_arr = VarBinArray::from(gen_varbin_words(len, uniqueness)); let dict = dict_encode( &varbin_arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let bytes = varbin_arr.with_iterator(|i| i.next().unwrap().unwrap().to_vec()); @@ -98,7 +102,7 @@ fn bench_compare_varbinview(bencher: divan::Bencher, (len, uniqueness): (usize, let varbinview_arr = VarBinViewArray::from_iter_str(gen_varbin_words(len, uniqueness)); let dict = dict_encode( &varbinview_arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let bytes = varbinview_arr.with_iterator(|i| i.next().unwrap().unwrap().to_vec()); @@ -137,7 +141,7 @@ fn bench_compare_sliced_dict_primitive( let primitive_arr = gen_primitive_for_dict::(codes_len.max(values_len), values_len); let dict = dict_encode( &primitive_arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let dict = dict.into_array().slice(0..codes_len).unwrap(); @@ -163,7 +167,7 @@ fn bench_compare_sliced_dict_varbinview( let varbin_arr = VarBinArray::from(gen_varbin_words(codes_len.max(values_len), values_len)); let dict = dict_encode( &varbin_arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let dict = dict.into_array().slice(0..codes_len).unwrap(); diff --git a/vortex-array/benches/dict_compress.rs b/vortex-array/benches/dict_compress.rs index cf7f647a7e0..21d427bd7f8 100644 --- a/vortex-array/benches/dict_compress.rs +++ b/vortex-array/benches/dict_compress.rs @@ -3,12 +3,13 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::distr::Distribution; use rand::distr::StandardUniform; use vortex_array::Canonical; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; @@ -16,6 +17,8 @@ use vortex_array::arrays::dict_test::gen_primitive_for_dict; use vortex_array::arrays::dict_test::gen_varbin_words; use vortex_array::builders::dict::dict_encode; use vortex_array::dtype::NativePType; +use vortex_array::session::ArraySession; +use vortex_session::VortexSession; fn main() { divan::main(); @@ -35,6 +38,9 @@ const BENCH_ARGS: &[(usize, usize)] = &[ (10_000, 512), ]; +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[divan::bench(types = [u8, f32, i64], args = BENCH_ARGS)] fn encode_primitives(bencher: Bencher, (len, unique_values): (usize, usize)) where @@ -43,36 +49,27 @@ where { let primitive_arr = gen_primitive_for_dict::(len, unique_values); - bencher.with_inputs(|| &primitive_arr).bench_refs(|arr| { - dict_encode( - &arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - }); + bencher + .with_inputs(|| (&primitive_arr, SESSION.create_execution_ctx())) + .bench_refs(|(arr, ctx)| dict_encode(&arr.clone().into_array(), ctx)); } #[divan::bench(args = BENCH_ARGS)] fn encode_varbin(bencher: Bencher, (len, unique_values): (usize, usize)) { let varbin_arr = VarBinArray::from(gen_varbin_words(len, unique_values)); - bencher.with_inputs(|| &varbin_arr).bench_refs(|arr| { - dict_encode( - &arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - }); + bencher + .with_inputs(|| (&varbin_arr, SESSION.create_execution_ctx())) + .bench_refs(|(arr, ctx)| dict_encode(&arr.clone().into_array(), ctx)); } #[divan::bench(args = BENCH_ARGS)] fn encode_varbinview(bencher: Bencher, (len, unique_values): (usize, usize)) { let varbinview_arr = VarBinViewArray::from_iter_str(gen_varbin_words(len, unique_values)); - bencher.with_inputs(|| &varbinview_arr).bench_refs(|arr| { - dict_encode( - &arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - }); + bencher + .with_inputs(|| (&varbinview_arr, SESSION.create_execution_ctx())) + .bench_refs(|(arr, ctx)| dict_encode(&arr.clone().into_array(), ctx)); } #[divan::bench(types = [u8, f32, i64], args = BENCH_ARGS)] @@ -84,13 +81,13 @@ where let primitive_arr = gen_primitive_for_dict::(len, unique_values); let dict = dict_encode( &primitive_arr.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() .into_array(); bencher - .with_inputs(|| (&dict, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| (**dict).clone().execute::(ctx)); } @@ -99,13 +96,13 @@ fn decode_varbin(bencher: Bencher, (len, unique_values): (usize, usize)) { let varbin_arr = VarBinArray::from(gen_varbin_words(len, unique_values)); let dict = dict_encode( &varbin_arr.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() .into_array(); bencher - .with_inputs(|| (&dict, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| (**dict).clone().execute::(ctx)); } @@ -114,12 +111,12 @@ fn decode_varbinview(bencher: Bencher, (len, unique_values): (usize, usize)) { let varbinview_arr = VarBinViewArray::from_iter_str(gen_varbin_words(len, unique_values)); let dict = dict_encode( &varbinview_arr.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap() .into_array(); bencher - .with_inputs(|| (&dict, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) .bench_refs(|(dict, ctx)| (**dict).clone().execute::(ctx)); } diff --git a/vortex-array/benches/filter_bool.rs b/vortex-array/benches/filter_bool.rs index 5699fdb19cb..6c240db6b60 100644 --- a/vortex-array/benches/filter_bool.rs +++ b/vortex-array/benches/filter_bool.rs @@ -11,20 +11,26 @@ #![expect(clippy::cast_sign_loss)] #![expect(clippy::cast_precision_loss)] +use std::sync::LazyLock; + use divan::Bencher; use rand::prelude::*; use rand_distr::Zipf; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; +use vortex_array::session::ArraySession; use vortex_buffer::BitBuffer; use vortex_mask::Mask; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + const SIZES: &[usize] = &[1_000, 10_000, 100_000, 250_000]; const DENSITY_SWEEP_SIZE: usize = 100_000; const ARRAY_SEED: u64 = 42; @@ -142,7 +148,7 @@ fn filter_random_by_mostly_true(bencher: Bencher, n: usize) { ( array.clone(), make_mostly_true(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -162,7 +168,7 @@ fn filter_random_by_mostly_false(bencher: Bencher, n: usize) { ( array.clone(), make_mostly_false(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -182,7 +188,7 @@ fn filter_random_by_random(bencher: Bencher, n: usize) { ( array.clone(), make_random(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -202,7 +208,7 @@ fn filter_random_by_correlated_runs(bencher: Bencher, n: usize) { ( array.clone(), make_correlated_runs(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -222,7 +228,7 @@ fn filter_random_by_power_law(bencher: Bencher, n: usize) { ( array.clone(), make_power_law(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -244,7 +250,7 @@ fn filter_powerlaw_by_mostly_true(bencher: Bencher, n: usize) { ( array.clone(), make_mostly_true(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -264,7 +270,7 @@ fn filter_powerlaw_by_mostly_false(bencher: Bencher, n: usize) { ( array.clone(), make_mostly_false(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -284,7 +290,7 @@ fn filter_powerlaw_by_random(bencher: Bencher, n: usize) { ( array.clone(), make_random(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -304,7 +310,7 @@ fn filter_powerlaw_by_correlated_runs(bencher: Bencher, n: usize) { ( array.clone(), make_correlated_runs(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -324,7 +330,7 @@ fn filter_powerlaw_by_power_law(bencher: Bencher, n: usize) { ( array.clone(), make_power_law(n, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -350,7 +356,7 @@ fn density_sweep_random(bencher: Bencher, density: f64) { ( array.clone(), make_density_mask(DENSITY_SWEEP_SIZE, density, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -371,7 +377,7 @@ fn density_sweep_dense_runs(bencher: Bencher, density: f64) { ( array.clone(), make_dense_runs(DENSITY_SWEEP_SIZE, false_rate, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -391,7 +397,7 @@ fn density_sweep_single_slice(bencher: Bencher, density: f64) { ( array.clone(), make_single_slice(DENSITY_SWEEP_SIZE, density), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -415,7 +421,7 @@ fn filter_all_true(bencher: Bencher, n: usize) { ( array.clone(), Mask::new_true(n), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -437,7 +443,7 @@ fn filter_one_false(bencher: Bencher, n: usize) { ( array.clone(), Mask::from_buffer(BitBuffer::from_iter(bits)), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { @@ -457,7 +463,7 @@ fn filter_ultra_sparse(bencher: Bencher, n: usize) { ( array.clone(), make_density_mask(n, 0.0001, &mut mask_rng()), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, m, ctx)| { diff --git a/vortex-array/benches/listview_rebuild.rs b/vortex-array/benches/listview_rebuild.rs index f605a9d372c..31fdaf5e890 100644 --- a/vortex-array/benches/listview_rebuild.rs +++ b/vortex-array/benches/listview_rebuild.rs @@ -117,113 +117,120 @@ fn make_nested_list_lv( #[divan::bench] fn i32_small(bencher: Bencher) { let lv = make_primitive_lv(50, 32, 32); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } #[divan::bench] fn i32_small_overlapping(bencher: Bencher) { let lv = make_primitive_lv(50, 8, 1); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } #[divan::bench] fn varbinview_small(bencher: Bencher) { let lv = make_varbinview_lv(50, 32, 32); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } #[divan::bench] fn struct_small(bencher: Bencher) { let lv = make_struct_lv(50, 32, 32); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } #[divan::bench] fn i32_large(bencher: Bencher) { let lv = make_primitive_lv(50, 1_024, 1_024); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } #[divan::bench] fn varbinview_large(bencher: Bencher) { let lv = make_varbinview_lv(5, 1_024, 1_024); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } #[divan::bench] fn struct_large(bencher: Bencher) { let lv = make_struct_lv(25, 1_024, 1_024); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } #[divan::bench] @@ -242,31 +249,33 @@ fn fsl_large(bencher: Bencher) { sizes.into_array(), Validity::NonNullable, ); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } #[divan::bench] fn list_i32_large(bencher: Bencher) { let lv = make_nested_list_lv(2, 512, 2); - bencher.with_inputs(|| &lv).bench_refs(|lv| { - let mut ctx = SESSION.create_execution_ctx(); - let rebuilt = lv - .rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx) - .unwrap(); - rebuilt - .elements() - .clone() - .execute::(&mut ctx) - .unwrap() - }); + bencher + .with_inputs(|| (&lv, SESSION.create_execution_ctx())) + .bench_refs(|(lv, ctx)| { + let rebuilt = lv + .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx) + .unwrap(); + rebuilt + .elements() + .clone() + .execute::(ctx) + .unwrap() + }); } diff --git a/vortex-array/benches/scalar_at_struct.rs b/vortex-array/benches/scalar_at_struct.rs index 757d6252bb9..5e435b1000a 100644 --- a/vortex-array/benches/scalar_at_struct.rs +++ b/vortex-array/benches/scalar_at_struct.rs @@ -3,6 +3,8 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; @@ -10,12 +12,13 @@ use rand::distr::Uniform; use rand::rngs::StdRng; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldNames; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_session::VortexSession; fn main() { divan::main(); @@ -24,6 +27,9 @@ fn main() { const ARRAY_SIZE: usize = 100_000; const NUM_ACCESSES: usize = 1000; +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[divan::bench] fn execute_scalar_struct_simple(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -49,11 +55,10 @@ fn execute_scalar_struct_simple(bencher: Bencher) { .collect(); bencher - .with_inputs(|| (&struct_array, &indices)) - .bench_refs(|(array, indices)| { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + .with_inputs(|| (&struct_array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { for &idx in indices.iter() { - divan::black_box(array.execute_scalar(idx, &mut ctx).unwrap()); + divan::black_box(array.execute_scalar(idx, ctx).unwrap()); } }); } @@ -87,11 +92,10 @@ fn execute_scalar_struct_wide(bencher: Bencher) { .collect(); bencher - .with_inputs(|| (&struct_array, &indices)) - .bench_refs(|(array, indices)| { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + .with_inputs(|| (&struct_array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, ctx)| { for &idx in indices.iter() { - divan::black_box(array.execute_scalar(idx, &mut ctx).unwrap()); + divan::black_box(array.execute_scalar(idx, ctx).unwrap()); } }); } diff --git a/vortex-array/benches/scalar_subtract.rs b/vortex-array/benches/scalar_subtract.rs index bdbcdd36f3c..2eee8f3d9ff 100644 --- a/vortex-array/benches/scalar_subtract.rs +++ b/vortex-array/benches/scalar_subtract.rs @@ -3,25 +3,31 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::rngs::StdRng; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::session::ArraySession; use vortex_buffer::Buffer; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + #[divan::bench] fn scalar_subtract(bencher: Bencher) { let mut rng = StdRng::seed_from_u64(0); @@ -41,7 +47,7 @@ fn scalar_subtract(bencher: Bencher) { let chunked = ChunkedArray::from_iter([data1, data2]).into_array(); bencher - .with_inputs(|| (&chunked, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&chunked, SESSION.create_execution_ctx())) .bench_refs(|(chunked, ctx)| { chunked .clone() diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index baf2d4ede37..4764f5538be 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -10,22 +10,28 @@ #![expect(clippy::cast_possible_truncation)] #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::FixedSizeListArray; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + /// Number of lists in the source array. const NUM_LISTS: usize = 500; @@ -62,7 +68,7 @@ fn take_fsl_random(bencher: Bencher, num_indices: usize) let indices_array = indices.into_array(); bencher - .with_inputs(|| (&fsl, &indices_array, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array .clone() @@ -88,7 +94,7 @@ fn take_fsl_nullable_random(bencher: Bencher, num_indice let indices_array = indices.into_array(); bencher - .with_inputs(|| (&fsl, &indices_array, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&fsl, &indices_array, SESSION.create_execution_ctx())) .bench_refs(|(array, indices, execution_ctx)| { array .clone() diff --git a/vortex-array/benches/take_patches.rs b/vortex-array/benches/take_patches.rs index f929cbd3b69..4d87f8b3c63 100644 --- a/vortex-array/benches/take_patches.rs +++ b/vortex-array/benches/take_patches.rs @@ -4,23 +4,29 @@ #![expect(clippy::unwrap_used)] #![expect(clippy::cast_possible_truncation)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; #[expect(deprecated)] use vortex_array::ToCanonical as _; use vortex_array::VortexSessionExecute; use vortex_array::patches::Patches; +use vortex_array::session::ArraySession; use vortex_buffer::Buffer; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + const BENCH_ARGS: &[(f64, f64)] = &[ // patches_sparsity, index_multiple (0.1, 1.0), @@ -48,7 +54,7 @@ fn take_search(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) ); bencher - .with_inputs(|| (&patches, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&patches, &indices, SESSION.create_execution_ctx())) .bench_refs(|(patches, indices, ctx)| { #[expect(deprecated)] let prim = indices.to_primitive(); @@ -67,7 +73,7 @@ fn take_search_chunked(bencher: Bencher, (patches_sparsity, index_multiple): (f6 ); bencher - .with_inputs(|| (&patches, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&patches, &indices, SESSION.create_execution_ctx())) .bench_refs(|(patches, indices, ctx)| { #[expect(deprecated)] let prim = indices.to_primitive(); @@ -86,7 +92,7 @@ fn take_map(bencher: Bencher, (patches_sparsity, index_multiple): (f64, f64)) { ); bencher - .with_inputs(|| (&patches, &indices, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&patches, &indices, SESSION.create_execution_ctx())) .bench_refs(|(patches, indices, ctx)| { #[expect(deprecated)] let prim = indices.to_primitive(); diff --git a/vortex-array/benches/take_primitive.rs b/vortex-array/benches/take_primitive.rs index 4411cf1ac79..e211e54d5f7 100644 --- a/vortex-array/benches/take_primitive.rs +++ b/vortex-array/benches/take_primitive.rs @@ -8,16 +8,19 @@ #![expect(clippy::cast_possible_truncation)] #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::distr::Uniform; use rand::prelude::*; use rand_distr::Zipf; use vortex_array::Canonical; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::session::ArraySession; +use vortex_session::VortexSession; fn main() { divan::main(); @@ -29,6 +32,9 @@ const NUM_INDICES: &[usize] = &[1_000, 10_000, 100_000]; /// Size of the source vector / dictionary values. const VECTOR_SIZE: &[usize] = &[16, 256, 2048, 8192]; +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + // --- DictArray canonicalization benchmarks --- #[divan::bench(args = NUM_INDICES, consts = VECTOR_SIZE, sample_count = 100_000)] @@ -41,12 +47,9 @@ fn dict_canonicalize_uniform(bencher: Bencher, num_indi let dict = DictArray::try_new(codes.into_array(), values.into_array()).unwrap(); - bencher.with_inputs(|| &dict).bench_refs(|dict| { - (*dict) - .clone() - .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) - }); + bencher + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) + .bench_refs(|(dict, ctx)| (*dict).clone().into_array().execute::(ctx)); } #[divan::bench(args = NUM_INDICES, consts = VECTOR_SIZE, sample_count = 100_000)] @@ -63,10 +66,7 @@ fn dict_canonicalize_zipfian(bencher: Bencher, num_indi let dict = DictArray::try_new(codes.into_array(), values.into_array()).unwrap(); - bencher.with_inputs(|| &dict).bench_refs(|dict| { - (*dict) - .clone() - .into_array() - .execute::(&mut LEGACY_SESSION.create_execution_ctx()) - }); + bencher + .with_inputs(|| (&dict, SESSION.create_execution_ctx())) + .bench_refs(|(dict, ctx)| (*dict).clone().into_array().execute::(ctx)); } diff --git a/vortex-array/benches/take_struct.rs b/vortex-array/benches/take_struct.rs index 6346bb705ac..5f15894383a 100644 --- a/vortex-array/benches/take_struct.rs +++ b/vortex-array/benches/take_struct.rs @@ -3,24 +3,30 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::rngs::StdRng; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::StructArray; use vortex_array::dtype::FieldNames; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + const ARRAY_SIZE: usize = 100_000; const TAKE_SIZE: usize = 1000; @@ -54,7 +60,7 @@ fn take_struct_simple(bencher: Bencher) { ( &struct_array, &indices_array, - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, indices, ctx)| { @@ -97,7 +103,7 @@ fn take_struct_wide(bencher: Bencher, width: usize) { ( &struct_array, &indices_array, - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, indices, ctx)| { @@ -137,7 +143,7 @@ fn take_struct_sequential_indices(bencher: Bencher) { ( &struct_array, &indices_array, - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(array, indices, ctx)| { diff --git a/vortex-array/benches/to_arrow.rs b/vortex-array/benches/to_arrow.rs index 332d0b642c0..0f84feb7baa 100644 --- a/vortex-array/benches/to_arrow.rs +++ b/vortex-array/benches/to_arrow.rs @@ -4,11 +4,11 @@ #![expect(clippy::unwrap_used)] use std::sync::Arc; +use std::sync::LazyLock; use divan::Bencher; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::ListArray; @@ -25,11 +25,16 @@ use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; +use vortex_array::session::ArraySession; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + fn schema() -> DType { let fields = StructFields::from_iter([ ( @@ -87,22 +92,17 @@ fn to_arrow_dtype(bencher: Bencher) { #[divan::bench] fn ArrowExportVTable_to_arrow_field(bencher: Bencher) { // Warm the ArrowSession - drop( - LEGACY_SESSION - .arrow() - .to_arrow_field("", &schema()) - .unwrap(), - ); + drop(SESSION.arrow().to_arrow_field("", &schema()).unwrap()); bencher .with_inputs(schema) - .bench_values(|dtype| LEGACY_SESSION.arrow().to_arrow_field("", &dtype).unwrap()) + .bench_values(|dtype| SESSION.arrow().to_arrow_field("", &dtype).unwrap()) } #[divan::bench] fn to_arrow_array(bencher: Bencher) { bencher - .with_inputs(|| (array(), LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (array(), SESSION.create_execution_ctx())) .bench_values(|(array, mut ctx)| { #[expect(deprecated, reason = "benchmarking deprecated code path")] array.execute_arrow(None, &mut ctx).unwrap() @@ -113,16 +113,16 @@ fn to_arrow_array(bencher: Bencher) { #[divan::bench] fn ArrowExportVTable_execute_arrow(bencher: Bencher) { // Warm the ArrowSession - drop(LEGACY_SESSION.arrow().execute_arrow( - array(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - )); + drop( + SESSION + .arrow() + .execute_arrow(array(), None, &mut SESSION.create_execution_ctx()), + ); bencher - .with_inputs(|| (array(), LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (array(), SESSION.create_execution_ctx())) .bench_values(|(array, mut ctx)| { - LEGACY_SESSION + SESSION .arrow() .execute_arrow(array, None, &mut ctx) .unwrap() diff --git a/vortex-array/benches/varbinview_zip.rs b/vortex-array/benches/varbinview_zip.rs index b7283bd7c52..d18012d9914 100644 --- a/vortex-array/benches/varbinview_zip.rs +++ b/vortex-array/benches/varbinview_zip.rs @@ -3,21 +3,27 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::session::ArraySession; use vortex_mask::Mask; +use vortex_session::VortexSession; fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + /// Benchmarks zip on VarBinView arrays with a highly fragmented mask (worst case for per-slice lookup paths). #[divan::bench] fn varbinview_zip_fragmented_mask(bencher: Bencher) { @@ -32,7 +38,7 @@ fn varbinview_zip_fragmented_mask(bencher: Bencher) { if_true.clone(), if_false.clone(), mask.clone().into_array(), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(t, f, m, ctx)| { @@ -57,7 +63,7 @@ fn varbinview_zip_block_mask(bencher: Bencher) { if_true.clone(), if_false.clone(), mask.clone().into_array(), - LEGACY_SESSION.create_execution_ctx(), + SESSION.create_execution_ctx(), ) }) .bench_refs(|(t, f, m, ctx)| { diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index a673b10143a..7b553041c68 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -27,8 +27,8 @@ use crate::match_each_integer_ptype; /// Creates a `ListViewArray` from a `ListArray` by computing `sizes` from `offsets`. /// -/// The output `ListViewArray` will be zero-copyable back to a `ListArray`, and additionally it -/// will not have any leading or trailing garbage data. +/// The output `ListViewArray` will be zero-copyable back to a `ListArray`, and additionally it will +/// not have any leading or trailing garbage data. pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexResult { // If the list is empty, create an empty `ListViewArray` with the same offset `DType` as the // input. @@ -39,10 +39,7 @@ pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexRes // We reset the offsets here because mostly for convenience, and also because callers of this // function might not expect the output `ListViewArray` to have a bunch of leading and trailing // garbage data when they turn it back into a `ListArray`. - let list = list - .reset_offsets(false, ctx) - .vortex_expect("This can't fail"); - + let list = list.reset_offsets(false, ctx)?; let list_offsets = list.offsets().clone(); // Create `sizes` array by computing differences between consecutive offsets. diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index 837bd50b7db..eeae7886026 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -292,6 +292,10 @@ impl ArrayBuilder for ListViewBuilder { } unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) { + // TODO: The `ArrayBuilder` trait does not thread an `ExecutionCtx` through its extend + // methods, so we are forced to mint a fresh `LEGACY_SESSION` context here on every call + // (which for chunked input means once per chunk). Once the trait carries a `&mut + // ExecutionCtx`, the caller's session should be reused instead. let mut ctx = LEGACY_SESSION.create_execution_ctx(); let listview = array diff --git a/vortex-compressor/benches/dict_encode.rs b/vortex-compressor/benches/dict_encode.rs index b3314008c03..68b7cd98013 100644 --- a/vortex-compressor/benches/dict_encode.rs +++ b/vortex-compressor/benches/dict_encode.rs @@ -3,17 +3,23 @@ #![expect(clippy::unwrap_used)] +use std::sync::LazyLock; + use divan::Bencher; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builders::dict::dict_encode; +use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_compressor::builtins::integer_dictionary_encode; use vortex_compressor::stats::IntegerStats; +use vortex_session::VortexSession; + +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); fn make_array() -> PrimitiveArray { let values: BufferMut = (0..50).cycle().take(64_000).collect(); @@ -33,16 +39,16 @@ fn make_array() -> PrimitiveArray { #[divan::bench] fn encode_generic(bencher: Bencher) { let array = make_array().into_array(); - bencher.with_inputs(|| &array).bench_refs(|array| { - dict_encode(array, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() - }); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_refs(|(array, ctx)| dict_encode(array, ctx).unwrap()); } #[cfg(not(codspeed))] #[divan::bench] fn encode_specialized(bencher: Bencher) { let array = make_array(); - let stats = IntegerStats::generate(&array, &mut LEGACY_SESSION.create_execution_ctx()); + let stats = IntegerStats::generate(&array, &mut SESSION.create_execution_ctx()); bencher .with_inputs(|| &stats) .bench_refs(|stats| integer_dictionary_encode(array.as_view(), stats)); diff --git a/vortex-compressor/benches/stats_calc.rs b/vortex-compressor/benches/stats_calc.rs index c93a8373801..69a8f970608 100644 --- a/vortex-compressor/benches/stats_calc.rs +++ b/vortex-compressor/benches/stats_calc.rs @@ -4,15 +4,21 @@ #[cfg(not(codspeed))] #[divan::bench_group(items_count = 64_000u32, bytes_count = 256_000u32)] mod benchmarks { + use std::sync::LazyLock; + use divan::Bencher; - use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; + use vortex_array::session::ArraySession; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_compressor::stats::GenerateStatsOptions; use vortex_compressor::stats::IntegerStats; + use vortex_session::VortexSession; + + static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); fn generate_dataset(max_run: u32, distinct: u32) -> Buffer { let mut output = BufferMut::with_capacity(64_000); @@ -56,7 +62,7 @@ mod benchmarks { }; bencher - .with_inputs(|| (&values, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&values, SESSION.create_execution_ctx())) .bench_refs(|(values, ctx)| { IntegerStats::generate_opts(values, GenerateStatsOptions::default(), ctx); }); @@ -71,7 +77,7 @@ mod benchmarks { }; bencher - .with_inputs(|| (&values, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&values, SESSION.create_execution_ctx())) .bench_refs(|(values, ctx)| { IntegerStats::generate_opts( values, diff --git a/vortex-row/benches/row_encode.rs b/vortex-row/benches/row_encode.rs index 07493d6ad48..729b992d8e1 100644 --- a/vortex-row/benches/row_encode.rs +++ b/vortex-row/benches/row_encode.rs @@ -13,6 +13,7 @@ //! mixed-field struct. use std::sync::Arc; +use std::sync::LazyLock; use arrow_array::Int64Array; use arrow_array::StringArray; @@ -28,18 +29,22 @@ use rand::SeedableRng; use rand::distr::Alphanumeric; use rand::rngs::StdRng; use vortex_array::IntoArray; -use vortex_array::LEGACY_SESSION; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::session::ArraySession; use vortex_row::RowEncoder; +use vortex_session::VortexSession; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; const N: usize = 100_000; +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + fn main() { divan::main(); } @@ -83,10 +88,10 @@ fn primitive_i64_vortex(bencher: divan::Bencher) { let col = PrimitiveArray::from_iter(v.clone()).into_array(); let bytes = (N * (1 + 8)) as u64; let encoder = RowEncoder::default(); - bencher.counter(BytesCount::new(bytes)).bench_local(|| { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - encoder.encode(&[col.clone()], &mut ctx).unwrap() - }) + bencher + .counter(BytesCount::new(bytes)) + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| encoder.encode(&[col.clone()], &mut ctx).unwrap()) } // ---------- utf8 ---------- @@ -114,10 +119,10 @@ fn utf8_vortex(bencher: divan::Bencher) { .sum(); let col = VarBinViewArray::from_iter_str(words.iter().map(String::as_str)).into_array(); let encoder = RowEncoder::default(); - bencher.counter(BytesCount::new(total)).bench_local(|| { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - encoder.encode(&[col.clone()], &mut ctx).unwrap() - }) + bencher + .counter(BytesCount::new(total)) + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| encoder.encode(&[col.clone()], &mut ctx).unwrap()) } // ---------- struct_mixed ---------- @@ -169,8 +174,8 @@ fn struct_mixed_vortex(bencher: divan::Bencher) { .unwrap() .into_array(); let encoder = RowEncoder::default(); - bencher.counter(BytesCount::new(total)).bench_local(|| { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); - encoder.encode(&[struct_arr.clone()], &mut ctx).unwrap() - }) + bencher + .counter(BytesCount::new(total)) + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_local_values(|mut ctx| encoder.encode(&[struct_arr.clone()], &mut ctx).unwrap()) } diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index 552f66c321f..1507e0fb813 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -5,6 +5,7 @@ use std::fmt; use std::ops::Deref; +use std::sync::LazyLock; use divan::Bencher; #[cfg(not(codspeed))] @@ -15,7 +16,6 @@ use rand::SeedableRng; 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::DictArray; use vortex::array::arrays::PrimitiveArray; @@ -24,6 +24,7 @@ use vortex::array::arrays::VarBinArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::varbin::VarBinArrayExt; use vortex::array::builtins::ArrayBuiltins; +use vortex::array::session::ArraySession; use vortex::dtype::DType; use vortex::dtype::PType; use vortex::encodings::alp::ALP; @@ -43,6 +44,7 @@ use vortex::encodings::runend::RunEnd; use vortex::encodings::runend::RunEndArrayExt; use vortex::error::VortexExpect; use vortex::extension::datetime::TimeUnit; +use vortex_session::VortexSession; #[global_allocator] static GLOBAL: MiMalloc = MiMalloc; @@ -51,6 +53,9 @@ fn main() { divan::main(); } +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + const NUM_VALUES: u64 = 100_000; // Helper function to conditionally add counter based on codspeed cfg @@ -72,7 +77,7 @@ mod setup { use super::*; fn setup_primitive_arrays() -> (PrimitiveArray, PrimitiveArray, PrimitiveArray) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut rng = StdRng::seed_from_u64(0); let uint_array = PrimitiveArray::from_iter((0..NUM_VALUES).map(|_| rng.random_range(42u32..256))); @@ -95,7 +100,7 @@ mod setup { /// Create FoR <- BitPacked encoding tree for u64 pub fn for_bp_u64() -> ArrayRef { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let (uint_array, ..) = setup_primitive_arrays(); let compressed = FoR::encode(uint_array).unwrap(); let inner = compressed.encoded(); @@ -107,7 +112,7 @@ mod setup { /// Create ALP <- FoR <- BitPacked encoding tree for f64 pub fn alp_for_bp_f64() -> ArrayRef { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let (_, _, float_array) = setup_primitive_arrays(); let alp_compressed = alp_encode(float_array.as_view(), None, &mut ctx).unwrap(); @@ -154,7 +159,7 @@ mod setup { let codes_prim = PrimitiveArray::from_iter(codes); // Compress codes with BitPacked (6 bits should be enough for ~50 unique values) - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let codes_bp = BitPacked::encode(&codes_prim.into_array(), 6, &mut ctx) .unwrap() .into_array(); @@ -185,7 +190,7 @@ mod setup { run_length -= 1; } - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let prim_array = PrimitiveArray::from_iter(values); let runend = RunEnd::encode(prim_array.into_array(), &mut ctx).unwrap(); @@ -234,7 +239,7 @@ mod setup { .collect(); // Train and compress unique values with FSST - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let unique_varbinview = VarBinViewArray::from_iter_str(unique_strings); let fsst_compressor = fsst_train_compressor(&unique_varbinview); let fsst_values = fsst_compress( @@ -273,7 +278,7 @@ mod setup { .collect(); // Train and compress unique values with FSST - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let unique_varbinview = VarBinViewArray::from_iter_str(unique_strings); let fsst_compressor = fsst_train_compressor(&unique_varbinview); let fsst = fsst_compress( @@ -341,7 +346,7 @@ mod setup { let temporal_array = TemporalArray::new_timestamp(ts_array, TimeUnit::Microseconds, None); // Split into days, seconds, subseconds - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let parts = split_temporal(temporal_array.clone(), &mut ctx).unwrap(); // Compress days with FoR <- BitPacked @@ -454,6 +459,6 @@ fn decompress(bencher: Bencher, setup_fn: SetupFn) { let nbytes = compressed.nbytes(); with_byte_counter(bencher, nbytes) - .with_inputs(|| (&compressed, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) .bench_refs(|(a, ctx)| (**a).clone().execute::(ctx)); } diff --git a/vortex/benches/single_encoding_throughput.rs b/vortex/benches/single_encoding_throughput.rs index d078e9fed4e..ad94adf8341 100644 --- a/vortex/benches/single_encoding_throughput.rs +++ b/vortex/benches/single_encoding_throughput.rs @@ -17,7 +17,6 @@ use rand::rngs::StdRng; use vortex::array::Canonical; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; -use vortex::array::LEGACY_SESSION; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::builders::dict::dict_encode; @@ -73,7 +72,7 @@ fn canonicalize(array: impl IntoArray, ctx: &mut ExecutionCtx) -> VortexResult (PrimitiveArray, PrimitiveArray, PrimitiveArray) { - let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let mut rng = StdRng::seed_from_u64(0); let uint_array = PrimitiveArray::from_iter((0..NUM_VALUES).map(|_| rng.random_range(42u32..256))); @@ -148,18 +147,15 @@ fn bench_runend_compress_u32(bencher: Bencher) { let (uint_array, ..) = setup_primitive_arrays(); with_byte_counter(bencher, NUM_VALUES * 4) - .with_inputs(|| (uint_array.clone(), LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (uint_array.clone(), SESSION.create_execution_ctx())) .bench_values(|(a, mut ctx)| RunEnd::encode(a.into_array(), &mut ctx).unwrap()); } #[divan::bench(name = "runend_decompress_u32")] fn bench_runend_decompress_u32(bencher: Bencher) { let (uint_array, ..) = setup_primitive_arrays(); - let compressed = RunEnd::encode( - uint_array.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap(); + let compressed = + RunEnd::encode(uint_array.into_array(), &mut SESSION.create_execution_ctx()).unwrap(); with_byte_counter(bencher, NUM_VALUES * 4) .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) @@ -215,14 +211,8 @@ fn bench_dict_compress_u32(bencher: Bencher) { let (uint_array, ..) = setup_primitive_arrays(); with_byte_counter(bencher, NUM_VALUES * 4) - .with_inputs(|| &uint_array) - .bench_refs(|a| { - dict_encode( - &a.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap() - }); + .with_inputs(|| (&uint_array, SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| dict_encode(&a.clone().into_array(), ctx).unwrap()); } #[divan::bench(name = "dict_decompress_u32")] @@ -230,7 +220,7 @@ fn bench_dict_decompress_u32(bencher: Bencher) { let (uint_array, ..) = setup_primitive_arrays(); let compressed = dict_encode( &uint_array.into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); @@ -285,15 +275,8 @@ fn bench_alp_compress_f64(bencher: Bencher) { let (_, _, float_array) = setup_primitive_arrays(); with_byte_counter(bencher, NUM_VALUES * 8) - .with_inputs(|| &float_array) - .bench_refs(|a| { - alp_encode( - a.as_view(), - None, - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap() - }); + .with_inputs(|| (&float_array, SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| alp_encode(a.as_view(), None, ctx).unwrap()); } #[divan::bench(name = "alp_decompress_f64")] @@ -302,7 +285,7 @@ fn bench_alp_decompress_f64(bencher: Bencher) { let compressed = alp_encode( float_array.as_view(), None, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); @@ -366,10 +349,8 @@ fn bench_zstd_compress_u32(bencher: Bencher) { let array = uint_array.into_array(); with_byte_counter(bencher, NUM_VALUES * 4) - .with_inputs(|| array.clone()) - .bench_values(|a| { - ZstdData::from_array(a, 3, 8192, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() - }); + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(a, mut ctx)| ZstdData::from_array(a, 3, 8192, &mut ctx).unwrap()); } #[cfg(feature = "zstd")] @@ -384,7 +365,7 @@ fn bench_zstd_decompress_u32(bencher: Bencher) { uint_array.into_array(), 3, 8192, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(), validity, @@ -405,14 +386,8 @@ fn bench_dict_compress_string(bencher: Bencher) { let nbytes = varbinview_arr.nbytes() as u64; with_byte_counter(bencher, nbytes) - .with_inputs(|| &varbinview_arr) - .bench_refs(|a| { - dict_encode( - &a.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), - ) - .unwrap() - }); + .with_inputs(|| (&varbinview_arr, SESSION.create_execution_ctx())) + .bench_refs(|(a, ctx)| dict_encode(&a.clone().into_array(), ctx).unwrap()); } #[divan::bench(name = "dict_decompress_string")] @@ -421,7 +396,7 @@ fn bench_dict_decompress_string(bencher: Bencher) { VarBinViewArray::from_iter_str(gen_varbin_words(NUM_VALUES as usize, 0.00005)); let dict = dict_encode( &varbinview_arr.clone().into_array(), - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(); let nbytes = varbinview_arr.into_array().nbytes() as u64; @@ -439,7 +414,7 @@ fn bench_fsst_compress_string(bencher: Bencher) { let nbytes = varbinview_arr.nbytes() as u64; with_byte_counter(bencher, nbytes) - .with_inputs(|| (&varbinview_arr, LEGACY_SESSION.create_execution_ctx())) + .with_inputs(|| (&varbinview_arr, SESSION.create_execution_ctx())) .bench_refs(|(a, ctx)| fsst_compress(*a, a.len(), a.dtype(), &fsst_compressor, ctx)); } @@ -453,7 +428,7 @@ fn bench_fsst_decompress_string(bencher: Bencher) { varbinview_arr.len(), varbinview_arr.dtype(), &fsst_compressor, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ); let nbytes = varbinview_arr.into_array().nbytes() as u64; @@ -471,10 +446,8 @@ fn bench_zstd_compress_string(bencher: Bencher) { let array = varbinview_arr.into_array(); with_byte_counter(bencher, nbytes) - .with_inputs(|| array.clone()) - .bench_values(|a| { - ZstdData::from_array(a, 3, 8192, &mut LEGACY_SESSION.create_execution_ctx()).unwrap() - }); + .with_inputs(|| (array.clone(), SESSION.create_execution_ctx())) + .bench_values(|(a, mut ctx)| ZstdData::from_array(a, 3, 8192, &mut ctx).unwrap()); } #[cfg(feature = "zstd")] @@ -490,7 +463,7 @@ fn bench_zstd_decompress_string(bencher: Bencher) { varbinview_arr.clone().into_array(), 3, 8192, - &mut LEGACY_SESSION.create_execution_ctx(), + &mut SESSION.create_execution_ctx(), ) .unwrap(), validity,