diff --git a/docs/developer-guide/internals/execution.md b/docs/developer-guide/internals/execution.md index 15ba16a8135..dc3138517a2 100644 --- a/docs/developer-guide/internals/execution.md +++ b/docs/developer-guide/internals/execution.md @@ -132,10 +132,8 @@ pub trait ExecuteParentKernel { ) -> VortexResult>; } -pub fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(parent_id, Child, Kernel); +pub fn initialize(session: &mut VortexSessionBuilder) { + builder_kernels(session).register_execute_parent_kernel(parent_id, Child, Kernel); } ``` diff --git a/docs/developer-guide/internals/session.md b/docs/developer-guide/internals/session.md index 823183ae7cc..63d1a37a00b 100644 --- a/docs/developer-guide/internals/session.md +++ b/docs/developer-guide/internals/session.md @@ -96,7 +96,7 @@ let session = VortexSession::default(); ``` For tests or specialized use-cases, sessions can be assembled from individual components using -the `.with::()` builder: +`VortexSession::builder()`: ```rust let session = VortexSession::builder() diff --git a/encodings/alp/benches/alp_compress.rs b/encodings/alp/benches/alp_compress.rs index 02cc30e44a6..03d4d2094cd 100644 --- a/encodings/alp/benches/alp_compress.rs +++ b/encodings/alp/benches/alp_compress.rs @@ -51,9 +51,9 @@ const BENCH_ARGS: &[(usize, f64, f64)] = &[ ]; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_alp::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_alp::initialize(&mut builder); + builder.build() }); #[divan::bench(types = [f32, f64], args = BENCH_ARGS)] diff --git a/encodings/alp/src/alp/array.rs b/encodings/alp/src/alp/array.rs index 8e581fe51fc..feac1cffdd1 100644 --- a/encodings/alp/src/alp/array.rs +++ b/encodings/alp/src/alp/array.rs @@ -482,9 +482,9 @@ mod tests { use crate::decompress_into_array; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index 819fc4b3dc4..39e35f847b0 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -151,9 +151,9 @@ mod tests { use crate::decompress_into_array; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/alp/src/alp/compute/between.rs b/encodings/alp/src/alp/compute/between.rs index dbb6fec9a22..e52ba2e70c1 100644 --- a/encodings/alp/src/alp/compute/between.rs +++ b/encodings/alp/src/alp/compute/between.rs @@ -154,9 +154,9 @@ mod tests { use crate::alp_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn assert_between( diff --git a/encodings/alp/src/alp/compute/cast.rs b/encodings/alp/src/alp/compute/cast.rs index d77a7257544..74e4d9902fd 100644 --- a/encodings/alp/src/alp/compute/cast.rs +++ b/encodings/alp/src/alp/compute/cast.rs @@ -66,9 +66,9 @@ mod tests { use crate::alp_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/alp/src/alp/compute/compare.rs b/encodings/alp/src/alp/compute/compare.rs index 83e38e19bda..9595507c743 100644 --- a/encodings/alp/src/alp/compute/compare.rs +++ b/encodings/alp/src/alp/compute/compare.rs @@ -173,9 +173,9 @@ mod tests { use crate::alp_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn test_alp_compare>( diff --git a/encodings/alp/src/alp/compute/filter.rs b/encodings/alp/src/alp/compute/filter.rs index f38a87f19b0..c64d5099251 100644 --- a/encodings/alp/src/alp/compute/filter.rs +++ b/encodings/alp/src/alp/compute/filter.rs @@ -45,9 +45,9 @@ mod test { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::filter::test_filter_conformance; + use vortex_array::default_session_builder; use vortex_buffer::buffer; use crate::alp_encode; @@ -62,7 +62,7 @@ mod test { 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0 ].into_array())] fn test_filter_alp_conformance(#[case] array: ArrayRef) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array_primitive = array.execute::(&mut ctx).unwrap(); let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).unwrap(); test_filter_conformance(&alp.into_array()); diff --git a/encodings/alp/src/alp/compute/mask.rs b/encodings/alp/src/alp/compute/mask.rs index b117ac216d0..0b808ed9b73 100644 --- a/encodings/alp/src/alp/compute/mask.rs +++ b/encodings/alp/src/alp/compute/mask.rs @@ -57,10 +57,10 @@ mod test { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::mask::test_mask_conformance; + use vortex_array::default_session_builder; use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::fns::mask::MaskKernel; use vortex_buffer::buffer; @@ -78,7 +78,7 @@ mod test { 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0 ].into_array())] fn test_mask_alp_conformance(#[case] array: vortex_array::ArrayRef) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array_primitive = array.execute::(&mut ctx).unwrap(); let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).unwrap(); test_mask_conformance(&alp.into_array()); @@ -87,7 +87,7 @@ mod test { #[test] fn test_mask_alp_with_patches() { use std::f64::consts::PI; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // PI doesn't encode cleanly with ALP, so it creates patches. let values: Vec = (0..100) .map(|i| if i % 4 == 3 { PI } else { 1.0 }) @@ -100,7 +100,7 @@ mod test { #[test] fn test_mask_alp_with_patches_casts_surviving_patch_values_to_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = PrimitiveArray::from_iter([1.234f32, f32::NAN, 2.345, f32::INFINITY, 3.456]); let alp = alp_encode(values.as_view(), None, &mut ctx).unwrap(); assert!(alp.patches().is_some(), "expected patches"); diff --git a/encodings/alp/src/alp/compute/mod.rs b/encodings/alp/src/alp/compute/mod.rs index 83a6c667893..d68ba3e2024 100644 --- a/encodings/alp/src/alp/compute/mod.rs +++ b/encodings/alp/src/alp/compute/mod.rs @@ -17,44 +17,44 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use vortex_session::VortexSession; use crate::ALPArray; use crate::alp_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - crate::initialize(&session); - session + let mut session = default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] // Basic float arrays - #[case::f32_array(alp_encode(PrimitiveArray::from_iter([1.23f32, 4.56, 7.89, 10.11, 12.13]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] - #[case::f64_array(alp_encode(PrimitiveArray::from_iter([100.1f64, 200.2, 300.3, 400.4, 500.5]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::f32_array(alp_encode(PrimitiveArray::from_iter([1.23f32, 4.56, 7.89, 10.11, 12.13]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] + #[case::f64_array(alp_encode(PrimitiveArray::from_iter([100.1f64, 200.2, 300.3, 400.4, 500.5]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] // Nullable arrays - #[case::nullable_f32(alp_encode(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] - #[case::nullable_f64(alp_encode(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::nullable_f32(alp_encode(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] + #[case::nullable_f64(alp_encode(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] // Edge cases - #[case::single_element(alp_encode(PrimitiveArray::from_iter([42.42f64]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::single_element(alp_encode(PrimitiveArray::from_iter([42.42f64]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] // Large arrays - #[case::large_f32(alp_encode(PrimitiveArray::from_iter((0..1000).map(|i| i as f32 * 0.1)).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::large_f32(alp_encode(PrimitiveArray::from_iter((0..1000).map(|i| i as f32 * 0.1)).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] // Arrays with patterns - #[case::repeating_pattern(alp_encode(PrimitiveArray::from_iter([1.1f32, 2.2, 3.3, 1.1, 2.2, 3.3, 1.1, 2.2, 3.3]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] - #[case::close_values(alp_encode(PrimitiveArray::from_iter([100.001f64, 100.002, 100.003, 100.004, 100.005]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::repeating_pattern(alp_encode(PrimitiveArray::from_iter([1.1f32, 2.2, 3.3, 1.1, 2.2, 3.3, 1.1, 2.2, 3.3]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] + #[case::close_values(alp_encode(PrimitiveArray::from_iter([100.001f64, 100.002, 100.003, 100.004, 100.005]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] fn test_alp_consistency(#[case] array: ALPArray) { test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx()); } #[rstest] - #[case::f32_basic(alp_encode(PrimitiveArray::from_iter([1.23f32, 4.56, 7.89, 10.11, 12.13]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] - #[case::f64_basic(alp_encode(PrimitiveArray::from_iter([100.1f64, 200.2, 300.3, 400.4, 500.5]).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] - #[case::f32_large(alp_encode(PrimitiveArray::from_iter((0..100).map(|i| i as f32 * 1.5)).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] - #[case::f64_large(alp_encode(PrimitiveArray::from_iter((0..100).map(|i| i as f64 * 2.5)).as_view(), None, &mut array_session().create_execution_ctx()).unwrap())] + #[case::f32_basic(alp_encode(PrimitiveArray::from_iter([1.23f32, 4.56, 7.89, 10.11, 12.13]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] + #[case::f64_basic(alp_encode(PrimitiveArray::from_iter([100.1f64, 200.2, 300.3, 400.4, 500.5]).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] + #[case::f32_large(alp_encode(PrimitiveArray::from_iter((0..100).map(|i| i as f32 * 1.5)).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] + #[case::f64_large(alp_encode(PrimitiveArray::from_iter((0..100).map(|i| i as f64 * 2.5)).as_view(), None, &mut default_session_builder().build().create_execution_ctx()).unwrap())] fn test_alp_binary_numeric(#[case] array: ALPArray) { test_binary_numeric_array(&array.into_array(), &mut SESSION.create_execution_ctx()); } diff --git a/encodings/alp/src/alp/compute/take.rs b/encodings/alp/src/alp/compute/take.rs index 9b193a3cfbe..7f3d2373299 100644 --- a/encodings/alp/src/alp/compute/take.rs +++ b/encodings/alp/src/alp/compute/take.rs @@ -43,9 +43,9 @@ mod test { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::default_session_builder; use vortex_buffer::buffer; use crate::alp_encode; @@ -56,7 +56,7 @@ mod test { #[case(PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]).into_array())] #[case(buffer![42.42f64].into_array())] fn test_take_alp_conformance(#[case] array: vortex_array::ArrayRef) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array_primitive = array.execute::(&mut ctx).unwrap(); let alp = alp_encode(array_primitive.as_view(), None, &mut ctx).unwrap(); test_take_conformance(&alp.into_array()); diff --git a/encodings/alp/src/alp/mod.rs b/encodings/alp/src/alp/mod.rs index 9992a9ddbe8..20871e6f07f 100644 --- a/encodings/alp/src/alp/mod.rs +++ b/encodings/alp/src/alp/mod.rs @@ -102,11 +102,11 @@ use vortex_array::dtype::NativePType; use vortex_array::scalar::PValue; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; const SAMPLE_SIZE: usize = 32; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { rules::initialize(session); } diff --git a/encodings/alp/src/alp/plugin.rs b/encodings/alp/src/alp/plugin.rs index c14133109d1..584167c5c7a 100644 --- a/encodings/alp/src/alp/plugin.rs +++ b/encodings/alp/src/alp/plugin.rs @@ -110,7 +110,7 @@ mod tests { use crate::alp_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(ALPPatchedPlugin); session }); diff --git a/encodings/alp/src/alp/rules.rs b/encodings/alp/src/alp/rules.rs index 3b888ae8269..a51aad26288 100644 --- a/encodings/alp/src/alp/rules.rs +++ b/encodings/alp/src/alp/rules.rs @@ -8,7 +8,7 @@ use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::between::BetweenReduceAdaptor; @@ -18,12 +18,12 @@ use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::Mask; use vortex_array::scalar_fn::fns::mask::MaskExecuteAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ALP; -pub(super) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(super) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Binary.id(), ALP, CompareExecuteAdaptor(ALP)); kernels.register_execute_parent_kernel(Filter.id(), ALP, FilterExecuteAdaptor(ALP)); kernels.register_execute_parent_kernel(Mask.id(), ALP, MaskExecuteAdaptor(ALP)); diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index 160723e83b2..22f260f69d0 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -638,9 +638,9 @@ mod test { use crate::alp_rd; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/alp/src/alp_rd/compute/cast.rs b/encodings/alp/src/alp_rd/compute/cast.rs index 717b97b760e..0ef6badbf76 100644 --- a/encodings/alp/src/alp_rd/compute/cast.rs +++ b/encodings/alp/src/alp_rd/compute/cast.rs @@ -51,10 +51,10 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::cast::test_cast_conformance; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -63,7 +63,7 @@ mod tests { #[test] fn test_cast_alprd_f32_to_f64() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1.0f32, 1.1, 1.2, 1.3, 1.4]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); @@ -87,7 +87,7 @@ mod tests { #[test] fn test_cast_alprd_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Some(10.0f64), None, Some(10.1), Some(10.2), None]); let values = vec![10.0f64, 10.1, 10.2]; @@ -122,31 +122,31 @@ mod tests { let values = vec![1.23f32, 4.56, 7.89, 10.11, 12.13]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut array_session().create_execution_ctx()) + encoder.encode(arr.as_view(), &mut default_session_builder().build().create_execution_ctx()) })] #[case::f64({ let values = vec![100.1f64, 200.2, 300.3, 400.4, 500.5]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut array_session().create_execution_ctx()) + encoder.encode(arr.as_view(), &mut default_session_builder().build().create_execution_ctx()) })] #[case::single({ let values = vec![42.42f64]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut array_session().create_execution_ctx()) + encoder.encode(arr.as_view(), &mut default_session_builder().build().create_execution_ctx()) })] #[case::negative({ let values = vec![0.0f32, -1.5, 2.5, -3.5, 4.5]; let arr = PrimitiveArray::from_iter(values.clone()); let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut array_session().create_execution_ctx()) + encoder.encode(arr.as_view(), &mut default_session_builder().build().create_execution_ctx()) })] #[case::nullable({ let arr = PrimitiveArray::from_option_iter([Some(1.1f32), None, Some(2.2), Some(3.3), None]); let values = vec![1.1f32, 2.2, 3.3]; let encoder = RDEncoder::new(&values); - encoder.encode(arr.as_view(), &mut array_session().create_execution_ctx()) + encoder.encode(arr.as_view(), &mut default_session_builder().build().create_execution_ctx()) })] fn test_cast_alprd_conformance(#[case] alprd: crate::alp_rd::ALPRDArray) { test_cast_conformance(&alprd.into_array()); diff --git a/encodings/alp/src/alp_rd/compute/filter.rs b/encodings/alp/src/alp_rd/compute/filter.rs index 38d3fa96156..faf87bc0607 100644 --- a/encodings/alp/src/alp_rd/compute/filter.rs +++ b/encodings/alp/src/alp_rd/compute/filter.rs @@ -46,10 +46,10 @@ mod test { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::filter::test_filter_conformance; + use vortex_array::default_session_builder; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_mask::Mask; @@ -60,9 +60,9 @@ mod test { use crate::RDEncoder; static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - crate::initialize(&session); - session + let mut session = default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/alp/src/alp_rd/compute/mask.rs b/encodings/alp/src/alp_rd/compute/mask.rs index cac98683adb..64c80dc6d13 100644 --- a/encodings/alp/src/alp_rd/compute/mask.rs +++ b/encodings/alp/src/alp_rd/compute/mask.rs @@ -44,9 +44,9 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::mask::test_mask_conformance; + use vortex_array::default_session_builder; use crate::ALPRDFloat; use crate::RDEncoder; @@ -55,7 +55,7 @@ mod tests { #[case(0.1f32, 0.2f32, 3e25f32)] #[case(0.1f64, 0.2f64, 3e100f64)] fn test_mask_simple(#[case] a: T, #[case] b: T, #[case] outlier: T) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); test_mask_conformance( &RDEncoder::new(&[a, b]) .encode( @@ -70,7 +70,7 @@ mod tests { #[case(0.1f32, 3e25f32)] #[case(0.5f64, 1e100f64)] fn test_mask_with_nulls(#[case] a: T, #[case] outlier: T) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); test_mask_conformance( &RDEncoder::new(&[a]) .encode( diff --git a/encodings/alp/src/alp_rd/compute/mod.rs b/encodings/alp/src/alp_rd/compute/mod.rs index f74b23fc2c2..8febe666ed5 100644 --- a/encodings/alp/src/alp_rd/compute/mod.rs +++ b/encodings/alp/src/alp_rd/compute/mod.rs @@ -13,19 +13,19 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use vortex_session::VortexSession; use crate::ALPRDArray; use crate::RDEncoder; static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - crate::initialize(&session); - session + let mut session = default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/alp/src/alp_rd/compute/take.rs b/encodings/alp/src/alp_rd/compute/take.rs index dee2454ea05..25ecaa7ed07 100644 --- a/encodings/alp/src/alp_rd/compute/take.rs +++ b/encodings/alp/src/alp_rd/compute/take.rs @@ -62,10 +62,10 @@ mod test { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::default_session_builder; use vortex_session::VortexSession; use crate::ALPRDArrayExt; @@ -73,9 +73,9 @@ mod test { use crate::RDEncoder; static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - crate::initialize(&session); - session + let mut session = default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/alp/src/alp_rd/kernel.rs b/encodings/alp/src/alp_rd/kernel.rs index 73670163954..efb9740b560 100644 --- a/encodings/alp/src/alp_rd/kernel.rs +++ b/encodings/alp/src/alp_rd/kernel.rs @@ -8,13 +8,13 @@ use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; -use vortex_session::VortexSession; +use vortex_array::optimizer::kernels::builder_kernels; +use vortex_session::VortexSessionBuilder; use crate::alp_rd::ALPRD; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Slice.id(), ALPRD, SliceExecuteAdaptor(ALPRD)); kernels.register_execute_parent_kernel(Filter.id(), ALPRD, FilterExecuteAdaptor(ALPRD)); kernels.register_execute_parent_kernel(Dict.id(), ALPRD, TakeExecuteAdaptor(ALPRD)); diff --git a/encodings/alp/src/alp_rd/mod.rs b/encodings/alp/src/alp_rd/mod.rs index a9ac7ca82fe..8626f00c712 100644 --- a/encodings/alp/src/alp_rd/mod.rs +++ b/encodings/alp/src/alp_rd/mod.rs @@ -36,7 +36,7 @@ use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_utils::aliases::hash_map::HashMap; use crate::match_each_alp_float_ptype; @@ -56,7 +56,7 @@ const CUT_LIMIT: usize = 16; const MAX_DICT_SIZE: u8 = 8; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } diff --git a/encodings/alp/src/alp_rd/ops.rs b/encodings/alp/src/alp_rd/ops.rs index be8d8f88948..cab187992f0 100644 --- a/encodings/alp/src/alp_rd/ops.rs +++ b/encodings/alp/src/alp_rd/ops.rs @@ -79,9 +79,9 @@ mod test { use crate::RDEncoder; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/alp/src/lib.rs b/encodings/alp/src/lib.rs index 1286f3e6e2b..8eca362181a 100644 --- a/encodings/alp/src/lib.rs +++ b/encodings/alp/src/lib.rs @@ -21,31 +21,38 @@ pub use alp_rd::*; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::nan_count::NanCount; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::patched::use_experimental_patches; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; mod alp; mod alp_rd; /// Initialize ALP encoding in the given session. -pub fn initialize(session: &VortexSession) { - // If we're using the experimental Patched encoding, register a shim - // for ALP with interior patches to decode as Patched array. - if use_experimental_patches() { - session.arrays().register(ALPPatchedPlugin); - } else { - session.arrays().register(ALP); +pub fn initialize(session: &mut VortexSessionBuilder) { + { + let arrays = session.get_mut::(); + + // If we're using the experimental Patched encoding, register a shim + // for ALP with interior patches to decode as Patched array. + if use_experimental_patches() { + arrays.register(ALPPatchedPlugin); + } else { + arrays.register(ALP); + } + arrays.register(ALPRD); } - session.arrays().register(ALPRD); + alp::initialize(session); alp_rd::initialize(session); // Register the ALP-specific NaN count aggregate kernel. - session.aggregate_fns().register_aggregate_kernel( - ALP.id(), - Some(NanCount.id()), - &compute::nan_count::ALPNanCountKernel, - ); + session + .get_mut::() + .register_aggregate_kernel( + ALP.id(), + Some(NanCount.id()), + &compute::nan_count::ALPNanCountKernel, + ); } diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs index dd5db041c15..5cb35878eef 100644 --- a/encodings/bytebool/src/array.rs +++ b/encodings/bytebool/src/array.rs @@ -319,9 +319,9 @@ mod tests { use super::*; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] @@ -361,7 +361,7 @@ mod tests { let array = ByteBool::from_option_vec(vec![Some(true), None, Some(false), None]); let dtype = array.dtype().clone(); let len = array.len(); - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(ByteBool); let ctx = ArrayContext::empty(); diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index 112abc921a4..5ec375fb9d9 100644 --- a/encodings/bytebool/src/compute.rs +++ b/encodings/bytebool/src/compute.rs @@ -167,7 +167,6 @@ mod tests { use rstest::rstest; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; @@ -176,6 +175,7 @@ mod tests { use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::compute::conformance::mask::test_mask_conformance; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar_fn::fns::operators::Operator; @@ -186,9 +186,9 @@ mod tests { use crate::ByteBoolArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - crate::initialize(&session); - session + let mut session = default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn bb(v: Vec) -> ByteBoolArray { @@ -351,7 +351,7 @@ mod tests { #[case::single_null(bb_opt(vec![None]))] #[case::mixed_with_nulls(bb_opt(vec![Some(true), None, Some(false), None, Some(true)]))] fn test_bytebool_consistency(#[case] array: ByteBoolArray) { - let ctx = &mut array_session().create_execution_ctx(); + let ctx = &mut default_session_builder().build().create_execution_ctx(); test_array_consistency(&array.into_array(), ctx); } } diff --git a/encodings/bytebool/src/kernel.rs b/encodings/bytebool/src/kernel.rs index 678da8642a7..3b494b072b6 100644 --- a/encodings/bytebool/src/kernel.rs +++ b/encodings/bytebool/src/kernel.rs @@ -4,18 +4,18 @@ use vortex_array::ArrayVTable; use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::BooleanExecuteAdaptor; use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ByteBool; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Binary.id(), ByteBool, BooleanExecuteAdaptor(ByteBool)); kernels.register_execute_parent_kernel(Cast.id(), ByteBool, CastExecuteAdaptor(ByteBool)); kernels.register_execute_parent_kernel(Dict.id(), ByteBool, TakeExecuteAdaptor(ByteBool)); diff --git a/encodings/bytebool/src/lib.rs b/encodings/bytebool/src/lib.rs index a803dacd776..cfa50c06de3 100644 --- a/encodings/bytebool/src/lib.rs +++ b/encodings/bytebool/src/lib.rs @@ -16,7 +16,7 @@ //! to a canonical [`BoolArray`][vortex_array::arrays::BoolArray]: //! //! ``` -//! # use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +//! # use vortex_array::{IntoArray, VortexSessionExecute, default_session_builder}; //! # use vortex_array::arrays::BoolArray; //! # use vortex_array::arrays::bool::BoolArrayExt; //! # use vortex_array::buffer::BufferHandle; @@ -25,7 +25,7 @@ //! # use vortex_bytebool::ByteBool; //! # use vortex_error::VortexResult; //! # fn main() -> VortexResult<()> { -//! # let mut ctx = array_session().create_execution_ctx(); +//! # let mut ctx = default_session_builder().build().create_execution_ctx(); //! let handle = BufferHandle::new_host(ByteBuffer::from(vec![0u8, 1, 42, 0])); //! let array = ByteBool::new(handle, Validity::NonNullable); //! @@ -41,8 +41,8 @@ //! [spec]: https://arrow.apache.org/docs/format/CanonicalExtensions.html#bit-boolean pub use array::*; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; mod array; mod compute; @@ -51,7 +51,7 @@ mod rules; mod slice; /// Initialize bytebool encoding in the given session. -pub fn initialize(session: &VortexSession) { - session.arrays().register(ByteBool); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(ByteBool); kernel::initialize(session); } diff --git a/encodings/datetime-parts/src/canonical.rs b/encodings/datetime-parts/src/canonical.rs index b08b76b148c..7637d98acac 100644 --- a/encodings/datetime-parts/src/canonical.rs +++ b/encodings/datetime-parts/src/canonical.rs @@ -119,9 +119,9 @@ mod test { use crate::canonical::decode_to_temporal; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index a3ab3e6f23b..0c5243f26b9 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/encodings/datetime-parts/src/compress.rs @@ -63,9 +63,9 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; + use vortex_array::default_session_builder; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; use vortex_buffer::buffer; @@ -80,7 +80,7 @@ mod tests { #[case(Validity::AllInvalid)] #[case(Validity::from_iter([true, false, true]))] fn test_split_temporal(#[case] validity: Validity) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let milliseconds = PrimitiveArray::new( buffer![ 86_400i64, // element with only day component diff --git a/encodings/datetime-parts/src/compute/cast.rs b/encodings/datetime-parts/src/compute/cast.rs index 90aba343ae3..32420607cd0 100644 --- a/encodings/datetime-parts/src/compute/cast.rs +++ b/encodings/datetime-parts/src/compute/cast.rs @@ -38,10 +38,10 @@ mod tests { use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::builtins::ArrayBuiltins; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; @@ -66,7 +66,7 @@ mod tests { TimeUnit::Milliseconds, Some("UTC".into()), ), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() .into_array() @@ -95,7 +95,7 @@ mod tests { #[case(Validity::AllInvalid)] #[case(Validity::from_iter([true, false, true]))] fn test_bad_cast_fails(#[case] validity: Validity) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = date_time_array(validity); // Cast to incompatible type - force evaluation via execute:: let result = array @@ -121,7 +121,7 @@ mod tests { ].into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut array_session().create_execution_ctx()).unwrap())] + ), &mut default_session_builder().build().create_execution_ctx()).unwrap())] #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp( PrimitiveArray::from_option_iter([ Some(0i64), @@ -132,12 +132,12 @@ mod tests { ]).into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut array_session().create_execution_ctx()).unwrap())] + ), &mut default_session_builder().build().create_execution_ctx()).unwrap())] #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp( buffer![86_400_000_000_000i64].into_array(), // 1 day in ns TimeUnit::Nanoseconds, Some("UTC".into()) - ), &mut array_session().create_execution_ctx()).unwrap())] + ), &mut default_session_builder().build().create_execution_ctx()).unwrap())] fn test_cast_datetime_parts_conformance(#[case] array: DateTimePartsArray) { use vortex_array::compute::conformance::cast::test_cast_conformance; test_cast_conformance(&array.into_array()); diff --git a/encodings/datetime-parts/src/compute/compare.rs b/encodings/datetime-parts/src/compute/compare.rs index 0791861df79..94705d22115 100644 --- a/encodings/datetime-parts/src/compute/compare.rs +++ b/encodings/datetime-parts/src/compute/compare.rs @@ -206,9 +206,9 @@ mod test { use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; + use vortex_array::default_session_builder; use vortex_array::dtype::IntegerPType; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; @@ -228,7 +228,7 @@ mod test { TimeUnit::Seconds, Some("UTC".into()), ), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .expect("Failed to construct DateTimePartsArray from TemporalArray") } @@ -248,7 +248,7 @@ mod test { #[case(Validity::AllValid, Validity::NonNullable)] #[case(Validity::AllValid, Validity::AllValid)] fn compare_date_time_parts_eq(#[case] lhs_validity: Validity, #[case] rhs_validity: Validity) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = dtp_array_from_timestamp(86400i64, lhs_validity); // January 2, 1970, 00:00:00 UTC let rhs = dtp_array_from_timestamp(86400i64, rhs_validity.clone()); // January 2, 1970, 00:00:00 UTC let comparison = lhs @@ -272,7 +272,7 @@ mod test { #[case(Validity::AllValid, Validity::NonNullable)] #[case(Validity::AllValid, Validity::AllValid)] fn compare_date_time_parts_ne(#[case] lhs_validity: Validity, #[case] rhs_validity: Validity) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = dtp_array_from_timestamp(86400i64, lhs_validity); // January 2, 1970, 00:00:00 UTC let rhs = dtp_array_from_timestamp(86401i64, rhs_validity.clone()); // January 2, 1970, 00:00:01 UTC let comparison = lhs @@ -296,7 +296,7 @@ mod test { #[case(Validity::AllValid, Validity::NonNullable)] #[case(Validity::AllValid, Validity::AllValid)] fn compare_date_time_parts_lt(#[case] lhs_validity: Validity, #[case] rhs_validity: Validity) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = dtp_array_from_timestamp(0i64, lhs_validity); // January 1, 1970, 01:00:00 UTC let rhs = dtp_array_from_timestamp(86400i64, rhs_validity); // January 2, 1970, 00:00:00 UTC @@ -313,7 +313,7 @@ mod test { #[case(Validity::AllValid, Validity::NonNullable)] #[case(Validity::AllValid, Validity::AllValid)] fn compare_date_time_parts_gt(#[case] lhs_validity: Validity, #[case] rhs_validity: Validity) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = dtp_array_from_timestamp(86400i64, lhs_validity); // January 2, 1970, 02:00:00 UTC let rhs = dtp_array_from_timestamp(0i64, rhs_validity); // January 1, 1970, 01:00:00 UTC @@ -333,7 +333,7 @@ mod test { #[case] lhs_validity: Validity, #[case] rhs_validity: Validity, ) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let temporal_array = TemporalArray::new_timestamp( PrimitiveArray::new(buffer![0i64], lhs_validity.clone()).into_array(), TimeUnit::Seconds, diff --git a/encodings/datetime-parts/src/compute/filter.rs b/encodings/datetime-parts/src/compute/filter.rs index 696660510d5..54d2f231968 100644 --- a/encodings/datetime-parts/src/compute/filter.rs +++ b/encodings/datetime-parts/src/compute/filter.rs @@ -28,10 +28,10 @@ impl FilterReduce for DateTimeParts { mod test { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::compute::conformance::filter::test_filter_conformance; + use vortex_array::default_session_builder; use vortex_array::extension::datetime::TimeUnit; use vortex_buffer::buffer; @@ -39,7 +39,7 @@ mod test { #[test] fn test_filter_datetime_parts() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create temporal arrays and convert to DateTimePartsArray let timestamps = buffer![ 0i64, diff --git a/encodings/datetime-parts/src/compute/kernel.rs b/encodings/datetime-parts/src/compute/kernel.rs index 4bd0cf2e1b3..221dd258ba9 100644 --- a/encodings/datetime-parts/src/compute/kernel.rs +++ b/encodings/datetime-parts/src/compute/kernel.rs @@ -4,16 +4,16 @@ use vortex_array::ArrayVTable; use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::DateTimeParts; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel( Binary.id(), DateTimeParts, diff --git a/encodings/datetime-parts/src/compute/mod.rs b/encodings/datetime-parts/src/compute/mod.rs index 4196b6b9832..1da7b6caf15 100644 --- a/encodings/datetime-parts/src/compute/mod.rs +++ b/encodings/datetime-parts/src/compute/mod.rs @@ -16,10 +16,10 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use vortex_array::extension::datetime::TimeUnit; use vortex_buffer::buffer; @@ -27,8 +27,11 @@ mod tests { use crate::DateTimePartsArray; fn dtp_from_temporal(temporal: TemporalArray) -> DateTimePartsArray { - DateTimeParts::try_from_temporal(temporal, &mut array_session().create_execution_ctx()) - .unwrap() + DateTimeParts::try_from_temporal( + temporal, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } #[rstest] @@ -79,7 +82,7 @@ mod tests { )))] fn test_datetime_parts_consistency(#[case] array: DateTimePartsArray) { - let ctx = &mut array_session().create_execution_ctx(); + let ctx = &mut default_session_builder().build().create_execution_ctx(); test_array_consistency(&array.into_array(), ctx); } } diff --git a/encodings/datetime-parts/src/compute/rules.rs b/encodings/datetime-parts/src/compute/rules.rs index 84f58310ac1..47c060efc0a 100644 --- a/encodings/datetime-parts/src/compute/rules.rs +++ b/encodings/datetime-parts/src/compute/rules.rs @@ -182,10 +182,10 @@ mod tests { use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; + use vortex_array::default_session_builder; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::TimestampOptions; use vortex_array::optimizer::ArrayOptimizer; @@ -230,8 +230,11 @@ mod tests { time_unit, None, ); - DateTimeParts::try_from_temporal(temporal, &mut array_session().create_execution_ctx()) - .vortex_expect("TemporalArray must produce valid DateTimeParts") + DateTimeParts::try_from_temporal( + temporal, + &mut default_session_builder().build().create_execution_ctx(), + ) + .vortex_expect("TemporalArray must produce valid DateTimeParts") } /// Create a constant timestamp scalar at midnight for the given day. @@ -296,7 +299,7 @@ mod tests { ); // Verify correctness: days [0, 1, 2] <= 1 should give [true, true, false] - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(true_count(&optimized, &mut ctx), 2); } @@ -325,7 +328,7 @@ mod tests { let optimized = between.optimize().unwrap(); // Verify correctness: days [0, 1, 2, 3, 4] between 1 and 3 should give [false, true, true, true, false] - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(true_count(&optimized, &mut ctx), 3); } @@ -348,7 +351,7 @@ mod tests { // (optimization doesn't apply, so we keep the original structure) // Just verify it still computes correctly // days [0, 1, 2] at midnight <= day 1 at noon: [true, true, false] - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(true_count(&optimized, &mut ctx), 2); } @@ -366,7 +369,7 @@ mod tests { TimeUnit::Seconds, None, ); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtp = DateTimeParts::try_from_temporal(temporal, &mut ctx).unwrap(); let len = dtp.len(); diff --git a/encodings/datetime-parts/src/compute/take.rs b/encodings/datetime-parts/src/compute/take.rs index 9f5a50d3a9b..d4a9b224fcf 100644 --- a/encodings/datetime-parts/src/compute/take.rs +++ b/encodings/datetime-parts/src/compute/take.rs @@ -97,10 +97,10 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::default_session_builder; use vortex_array::extension::datetime::TimeUnit; use vortex_buffer::buffer; @@ -118,7 +118,7 @@ mod tests { ].into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut array_session().create_execution_ctx()).unwrap())] + ), &mut default_session_builder().build().create_execution_ctx()).unwrap())] #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp( PrimitiveArray::from_option_iter([ Some(0i64), @@ -129,12 +129,12 @@ mod tests { ]).into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut array_session().create_execution_ctx()).unwrap())] + ), &mut default_session_builder().build().create_execution_ctx()).unwrap())] #[case(DateTimeParts::try_from_temporal(TemporalArray::new_timestamp( buffer![86_400_000i64].into_array(), TimeUnit::Milliseconds, Some("UTC".into()) - ), &mut array_session().create_execution_ctx()).unwrap())] + ), &mut default_session_builder().build().create_execution_ctx()).unwrap())] fn test_take_datetime_parts_conformance(#[case] array: DateTimePartsArray) { test_take_conformance(&array.into_array()); } diff --git a/encodings/datetime-parts/src/lib.rs b/encodings/datetime-parts/src/lib.rs index cdc65a44a71..abeaed11d4c 100644 --- a/encodings/datetime-parts/src/lib.rs +++ b/encodings/datetime-parts/src/lib.rs @@ -14,20 +14,22 @@ mod timestamp; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; /// Initialize datetime-parts encoding in the given session. -pub fn initialize(session: &VortexSession) { - session.arrays().register(DateTimeParts); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(DateTimeParts); compute::kernel::initialize(session); - session.aggregate_fns().register_aggregate_kernel( - DateTimeParts.id(), - Some(IsConstant.id()), - &compute::is_constant::DateTimePartsIsConstantKernel, - ); + session + .get_mut::() + .register_aggregate_kernel( + DateTimeParts.id(), + Some(IsConstant.id()), + &compute::is_constant::DateTimePartsIsConstantKernel, + ); } #[cfg(test)] diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index 376ad12f564..4a6521e9237 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs @@ -41,11 +41,11 @@ mod tests { use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::cast::test_cast_conformance; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; @@ -56,7 +56,7 @@ mod tests { #[test] fn test_cast_decimal_byte_parts_nullability() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let decimal_dtype = DecimalDType::new(10, 2); let array = DecimalByteParts::try_new(buffer![100i32, 200, 300, 400].into_array(), decimal_dtype) @@ -79,7 +79,7 @@ mod tests { #[test] fn test_cast_decimal_byte_parts_nullable_to_non_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let decimal_dtype = DecimalDType::new(10, 2); let array = DecimalByteParts::try_new( PrimitiveArray::from_option_iter([Some(100i32), None, Some(300)]).into_array(), diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index 04f321cf88f..4c86c01b785 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -167,9 +167,9 @@ mod tests { use crate::DecimalByteParts; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs index 5e8d28e3526..9881e2262b2 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs @@ -4,16 +4,16 @@ use vortex_array::ArrayVTable; use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::DecimalByteParts; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel( Binary.id(), DecimalByteParts, diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs index 0a5477e93cf..90c832ea82d 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -14,9 +14,9 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use vortex_array::dtype::DecimalDType; use vortex_buffer::buffer; @@ -71,7 +71,7 @@ mod tests { ).unwrap())] fn test_decimal_byte_parts_consistency(#[case] array: DecimalBytePartsArray) { - let ctx = &mut array_session().create_execution_ctx(); + let ctx = &mut default_session_builder().build().create_execution_ctx(); test_array_consistency(&array.into_array(), ctx); } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 8f4e9fed409..85a6ed8b997 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -321,9 +321,9 @@ impl ValidityChild for DecimalByteParts { mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; @@ -353,7 +353,10 @@ mod tests { assert_eq!( Scalar::null(dtype.clone()), array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( @@ -363,13 +366,19 @@ mod tests { ) .unwrap(), array - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), array - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 36a53c3a614..9993f56e83e 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -16,18 +16,20 @@ pub use decimal_byte_parts::*; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; /// Initialize decimal-byte-parts encoding in the given session. -pub fn initialize(session: &VortexSession) { - session.arrays().register(DecimalByteParts); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(DecimalByteParts); compute::kernel::initialize(session); - session.aggregate_fns().register_aggregate_kernel( - DecimalByteParts.id(), - Some(IsConstant.id()), - &DecimalBytePartsIsConstantKernel, - ); + session + .get_mut::() + .register_aggregate_kernel( + DecimalByteParts.id(), + Some(IsConstant.id()), + &DecimalBytePartsIsConstantKernel, + ); } diff --git a/encodings/experimental/onpair/benches/decode.rs b/encodings/experimental/onpair/benches/decode.rs index 1ee96652891..6a50e27f620 100644 --- a/encodings/experimental/onpair/benches/decode.rs +++ b/encodings/experimental/onpair/benches/decode.rs @@ -80,9 +80,9 @@ use vortex_onpair::onpair_compress; use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_onpair::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_onpair::initialize(&mut builder); + builder.build() }); #[derive(Copy, Clone, Debug)] diff --git a/encodings/experimental/onpair/src/compute/compare.rs b/encodings/experimental/onpair/src/compute/compare.rs index 8c01a36977c..dc52b2cbc17 100644 --- a/encodings/experimental/onpair/src/compute/compare.rs +++ b/encodings/experimental/onpair/src/compute/compare.rs @@ -84,7 +84,8 @@ mod tests { use crate::compress::DEFAULT_DICT12_CONFIG; use crate::compress::onpair_compress; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[cfg_attr(miri, ignore)] #[rstest] diff --git a/encodings/experimental/onpair/src/kernel.rs b/encodings/experimental/onpair/src/kernel.rs index 8863d750a72..b7207ea3a4a 100644 --- a/encodings/experimental/onpair/src/kernel.rs +++ b/encodings/experimental/onpair/src/kernel.rs @@ -4,19 +4,19 @@ use vortex_array::ArrayVTable; use vortex_array::arrays::Filter; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::byte_length::ByteLength; use vortex_array::scalar_fn::fns::byte_length::ByteLengthExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::OnPair; // TODO: implement ListExecute & TakeExecute for OnPair -pub(super) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(super) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Filter.id(), OnPair, FilterExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel(Binary.id(), OnPair, CompareExecuteAdaptor(OnPair)); kernels.register_execute_parent_kernel( diff --git a/encodings/experimental/onpair/src/lib.rs b/encodings/experimental/onpair/src/lib.rs index 11b22f63bc1..806e997f7aa 100644 --- a/encodings/experimental/onpair/src/lib.rs +++ b/encodings/experimental/onpair/src/lib.rs @@ -27,11 +27,11 @@ pub use onpair::Bits; pub use onpair::Config; pub use onpair::Error as OnPairError; pub use onpair::Threshold; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; /// Initialize OnPair encoding in the given session. -pub fn initialize(session: &VortexSession) { - session.arrays().register(OnPair); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(OnPair); kernel::initialize(session); } diff --git a/encodings/experimental/onpair/src/tests.rs b/encodings/experimental/onpair/src/tests.rs index 571240f3206..6ffaca978c4 100644 --- a/encodings/experimental/onpair/src/tests.rs +++ b/encodings/experimental/onpair/src/tests.rs @@ -27,7 +27,8 @@ use crate::OnPairMetadata; use crate::compress::DEFAULT_DICT12_CONFIG; use crate::compress::onpair_compress; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn sample_input() -> VarBinArray { VarBinArray::from_iter( diff --git a/encodings/experimental/onpair/tests/big_data.rs b/encodings/experimental/onpair/tests/big_data.rs index 9a38788378c..af5c4e2bd1a 100644 --- a/encodings/experimental/onpair/tests/big_data.rs +++ b/encodings/experimental/onpair/tests/big_data.rs @@ -30,7 +30,8 @@ use vortex_onpair::DEFAULT_DICT12_CONFIG; use vortex_onpair::onpair_compress; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn corpus(n: usize) -> Vec { let templates: &[&str] = &[ diff --git a/encodings/fastlanes/benches/bitpack_compare.rs b/encodings/fastlanes/benches/bitpack_compare.rs index 739fcd72dc6..72442b42069 100644 --- a/encodings/fastlanes/benches/bitpack_compare.rs +++ b/encodings/fastlanes/benches/bitpack_compare.rs @@ -36,9 +36,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fastlanes::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fastlanes::initialize(&mut builder); + builder.build() }); const LENS: &[usize] = &[1024, 64 * 1024]; diff --git a/encodings/fastlanes/benches/bitpack_compare_sweep.rs b/encodings/fastlanes/benches/bitpack_compare_sweep.rs index ec7cf9b6892..b58e3b06160 100644 --- a/encodings/fastlanes/benches/bitpack_compare_sweep.rs +++ b/encodings/fastlanes/benches/bitpack_compare_sweep.rs @@ -40,9 +40,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fastlanes::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fastlanes::initialize(&mut builder); + builder.build() }); /// Number of elements per benchmarked array (64 full FastLanes blocks). diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index eb072017ae3..d9f37655ddd 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -27,9 +27,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fastlanes::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fastlanes::initialize(&mut builder); + builder.build() }); #[divan::bench] diff --git a/encodings/fastlanes/benches/canonicalize_bench.rs b/encodings/fastlanes/benches/canonicalize_bench.rs index 30ee0756960..bae54dc1063 100644 --- a/encodings/fastlanes/benches/canonicalize_bench.rs +++ b/encodings/fastlanes/benches/canonicalize_bench.rs @@ -35,9 +35,9 @@ const BENCH_ARGS: &[(usize, usize, f64)] = &[ ]; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fastlanes::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fastlanes::initialize(&mut builder); + builder.build() }); #[cfg(not(codspeed))] diff --git a/encodings/fastlanes/benches/cast_bitpacked.rs b/encodings/fastlanes/benches/cast_bitpacked.rs index baeff20f5e1..ead008857cd 100644 --- a/encodings/fastlanes/benches/cast_bitpacked.rs +++ b/encodings/fastlanes/benches/cast_bitpacked.rs @@ -39,9 +39,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fastlanes::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fastlanes::initialize(&mut builder); + builder.build() }); const U32: DType = DType::Primitive(PType::U32, Nullability::NonNullable); diff --git a/encodings/fastlanes/benches/compute_between.rs b/encodings/fastlanes/benches/compute_between.rs index 2b295157e68..5b1f343f1e9 100644 --- a/encodings/fastlanes/benches/compute_between.rs +++ b/encodings/fastlanes/benches/compute_between.rs @@ -25,9 +25,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fastlanes::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fastlanes::initialize(&mut builder); + builder.build() }); fn generate_primitive_array( diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index f7a3485c113..f0edbaa72a8 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -448,9 +448,9 @@ mod test { use crate::bitpacking::array::BitPackedArrayExt; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 78aac1aa86e..d010e66cb15 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -226,9 +226,9 @@ mod tests { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn unpack(bitpacked: &BitPackedArray) -> VortexResult { diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 772de2ac641..96c6df3b66c 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -334,9 +334,9 @@ mod test { use crate::bitpacking::array::BitPackedArrayExt; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs index 3c77e146ad0..303d75fde0d 100644 --- a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs +++ b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs @@ -62,7 +62,7 @@ impl> UnpackStrategy for BitPackingStr /// use vortex_fastlanes::BitPackedArrayExt; /// use vortex_fastlanes::unpack_iter::BitUnpackedChunks; /// -/// let mut ctx = vortex_array::array_session().create_execution_ctx(); +/// let mut ctx = vortex_array::default_session_builder().build().create_execution_ctx(); /// let array = BitPackedData::encode(&buffer![2, 3, 4, 5].into_array(), 2, &mut ctx).unwrap(); /// let mut unpacked_chunks: BitUnpackedChunks = array.unpacked_chunks().unwrap(); /// diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index 1dd3b61dbf4..dab3f630c19 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -163,9 +163,9 @@ mod tests { use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn opts(lower: StrictComparison, upper: StrictComparison) -> BetweenOptions { diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index 3b917aa36af..271c088a623 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -135,9 +135,9 @@ mod tests { use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray { diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index 51aca62a9aa..a2c8a35f82f 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -126,9 +126,9 @@ mod tests { use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); /// All six operators on a small in-range input. diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 21184d785a5..ee9d2024763 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -194,9 +194,9 @@ mod test { use crate::bitpacking::array::BitPackedArrayExt; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 3b4e2b9770a..1ca1952dc73 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -187,7 +187,7 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::is_constant::is_constant; - use vortex_array::array_session; + use vortex_array::default_session_builder; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -195,7 +195,7 @@ mod tests { #[test] fn is_constant_with_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = BitPackedData::encode(&buffer![4; 1025].into_array(), 2, &mut ctx)?; assert!(is_constant(&array.into_array(), &mut ctx)?); Ok(()) diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index 38f86f781bb..473bd08f239 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -47,10 +47,10 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use crate::BitPackedArray; use crate::bitpack_compress::bitpack_encode; @@ -61,7 +61,7 @@ mod tests { array, bit_width, None, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -98,7 +98,7 @@ mod tests { #[case::alternating_bits(bp(&PrimitiveArray::from_iter([0u16, 255, 0, 255, 0, 255]), 8))] fn test_bitpacked_consistency(#[case] array: BitPackedArray) { - let ctx = &mut array_session().create_execution_ctx(); + let ctx = &mut default_session_builder().build().create_execution_ctx(); test_array_consistency(&array.into_array(), ctx); } @@ -112,7 +112,7 @@ mod tests { fn test_bitpacked_binary_numeric(#[case] array: BitPackedArray) { test_binary_numeric_array( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/encodings/fastlanes/src/bitpacking/compute/slice.rs b/encodings/fastlanes/src/bitpacking/compute/slice.rs index 996565a2672..711110e1127 100644 --- a/encodings/fastlanes/src/bitpacking/compute/slice.rs +++ b/encodings/fastlanes/src/bitpacking/compute/slice.rs @@ -73,9 +73,9 @@ fn slice_bitpacked( mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SliceArray; + use vortex_array::default_session_builder; use vortex_error::VortexResult; use crate::BitPacked; @@ -83,7 +83,7 @@ mod tests { #[test] fn test_reduce_parent_returns_bitpacked_slice() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = PrimitiveArray::from_iter(0u32..2048); let bitpacked = bitpack_encode(&values, 11, None, &mut ctx)?; diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 7ffe051ca40..e1f067ed877 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -181,9 +181,9 @@ mod test { use crate::bitpacking::compute::take::take_primitive; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index efa0677a91e..af419dcdad9 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -20,6 +20,6 @@ pub(crate) use plugin::BitPackedPatchedPlugin; pub use vtable::BitPacked; pub use vtable::BitPackedArray; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/encodings/fastlanes/src/bitpacking/plugin.rs b/encodings/fastlanes/src/bitpacking/plugin.rs index a621d085514..bd9976329fc 100644 --- a/encodings/fastlanes/src/bitpacking/plugin.rs +++ b/encodings/fastlanes/src/bitpacking/plugin.rs @@ -113,7 +113,7 @@ mod tests { use crate::BitPackedData; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(BitPackedPatchedPlugin); session }); diff --git a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs index eb0dd9b7a23..4ad0044edc2 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/kernels.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/kernels.rs @@ -8,7 +8,7 @@ use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::between::Between; use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; @@ -16,12 +16,12 @@ use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::BitPacked; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel( Between.id(), BitPacked, diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 9f8d41de014..eb415af8d6f 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -57,7 +57,7 @@ mod validity; /// A [`BitPacked`]-encoded Vortex array. pub type BitPackedArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { kernels::initialize(session); } diff --git a/encodings/fastlanes/src/delta/array/delta_compress.rs b/encodings/fastlanes/src/delta/array/delta_compress.rs index 7f00eb172dc..378c86fb4eb 100644 --- a/encodings/fastlanes/src/delta/array/delta_compress.rs +++ b/encodings/fastlanes/src/delta/array/delta_compress.rs @@ -116,9 +116,9 @@ mod tests { use crate::delta_compress; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/fastlanes/src/delta/array/mod.rs b/encodings/fastlanes/src/delta/array/mod.rs index 7754bb37a59..d411ad3abb7 100644 --- a/encodings/fastlanes/src/delta/array/mod.rs +++ b/encodings/fastlanes/src/delta/array/mod.rs @@ -38,7 +38,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["bases", "deltas"]; /// use vortex_session::VortexSession; /// use vortex_fastlanes::Delta; /// -/// let session = vortex_array::array_session(); +/// let session = vortex_array::default_session_builder().build(); /// let primitive = PrimitiveArray::from_iter([1_u32, 2, 3, 5, 10, 11]); /// let array = Delta::try_from_primitive_array(&primitive, &mut session.create_execution_ctx()).unwrap(); /// ``` @@ -53,7 +53,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["bases", "deltas"]; /// use vortex_session::VortexSession; /// use vortex_fastlanes::Delta; /// -/// let session = vortex_array::array_session(); +/// let session = vortex_array::default_session_builder().build(); /// let primitive = PrimitiveArray::from_iter([-3_i32, -2, -1, 0, 1, 2]); /// let array = Delta::try_from_primitive_array(&primitive, &mut session.create_execution_ctx()).unwrap(); /// ``` diff --git a/encodings/fastlanes/src/delta/compute/cast.rs b/encodings/fastlanes/src/delta/compute/cast.rs index 2ee10e7619a..fce5b0481b4 100644 --- a/encodings/fastlanes/src/delta/compute/cast.rs +++ b/encodings/fastlanes/src/delta/compute/cast.rs @@ -57,9 +57,9 @@ mod tests { use crate::Delta; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fastlanes/src/delta/vtable/operations.rs b/encodings/fastlanes/src/delta/vtable/operations.rs index 7ed57a0886d..8d84e0afe74 100644 --- a/encodings/fastlanes/src/delta/vtable/operations.rs +++ b/encodings/fastlanes/src/delta/vtable/operations.rs @@ -31,11 +31,11 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -45,9 +45,9 @@ mod tests { use crate::DeltaArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - crate::initialize(&session); - session + let mut session = default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn da(array: &PrimitiveArray) -> DeltaArray { diff --git a/encodings/fastlanes/src/for/array/for_compress.rs b/encodings/fastlanes/src/for/array/for_compress.rs index e39b4924e78..1daf1567133 100644 --- a/encodings/fastlanes/src/for/array/for_compress.rs +++ b/encodings/fastlanes/src/for/array/for_compress.rs @@ -69,9 +69,9 @@ mod test { use crate::r#for::array::for_decompress::fused_decompress; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fastlanes/src/for/compute/cast.rs b/encodings/fastlanes/src/for/compute/cast.rs index b865af43d12..fe184b789f4 100644 --- a/encodings/fastlanes/src/for/compute/cast.rs +++ b/encodings/fastlanes/src/for/compute/cast.rs @@ -52,9 +52,9 @@ mod tests { use crate::FoRArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn for_arr(encoded: ArrayRef, reference: Scalar) -> FoRArray { diff --git a/encodings/fastlanes/src/for/compute/compare.rs b/encodings/fastlanes/src/for/compute/compare.rs index 384bdee8037..878efd9dc15 100644 --- a/encodings/fastlanes/src/for/compute/compare.rs +++ b/encodings/fastlanes/src/for/compute/compare.rs @@ -107,9 +107,9 @@ mod tests { use crate::FoRArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn for_arr(encoded: ArrayRef, reference: Scalar) -> FoRArray { diff --git a/encodings/fastlanes/src/for/compute/is_sorted.rs b/encodings/fastlanes/src/for/compute/is_sorted.rs index a850f25583e..a42b907fbcb 100644 --- a/encodings/fastlanes/src/for/compute/is_sorted.rs +++ b/encodings/fastlanes/src/for/compute/is_sorted.rs @@ -62,8 +62,8 @@ mod test { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::is_sorted::is_sorted; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::default_session_builder; use vortex_array::validity::Validity; use vortex_buffer::buffer; @@ -72,7 +72,7 @@ mod test { #[test] fn test_sorted() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let a = PrimitiveArray::new(buffer![-1, 0, i8::MAX], Validity::NonNullable); let b = FoRData::encode(a).unwrap(); diff --git a/encodings/fastlanes/src/for/compute/mod.rs b/encodings/fastlanes/src/for/compute/mod.rs index 84e8812b598..aea5bec0ff5 100644 --- a/encodings/fastlanes/src/for/compute/mod.rs +++ b/encodings/fastlanes/src/for/compute/mod.rs @@ -103,10 +103,10 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -149,7 +149,7 @@ mod tests { fn test_for_consistency(#[case] array: FoRArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } @@ -165,7 +165,7 @@ mod tests { fn test_for_binary_numeric(#[case] array: FoRArray) { test_binary_numeric_array( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/encodings/fastlanes/src/for/mod.rs b/encodings/fastlanes/src/for/mod.rs index a5d495e3578..286f749adaa 100644 --- a/encodings/fastlanes/src/for/mod.rs +++ b/encodings/fastlanes/src/for/mod.rs @@ -11,6 +11,6 @@ mod vtable; pub use vtable::FoR; pub use vtable::FoRArray; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/encodings/fastlanes/src/for/vtable/kernels.rs b/encodings/fastlanes/src/for/vtable/kernels.rs index 6c63292bb9f..cfe8139fe31 100644 --- a/encodings/fastlanes/src/for/vtable/kernels.rs +++ b/encodings/fastlanes/src/for/vtable/kernels.rs @@ -4,16 +4,16 @@ use vortex_array::ArrayVTable; use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::FoR; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Binary.id(), FoR, CompareExecuteAdaptor(FoR)); kernels.register_execute_parent_kernel(Dict.id(), FoR, TakeExecuteAdaptor(FoR)); } diff --git a/encodings/fastlanes/src/for/vtable/mod.rs b/encodings/fastlanes/src/for/vtable/mod.rs index d3a0bb84811..3943903e97c 100644 --- a/encodings/fastlanes/src/for/vtable/mod.rs +++ b/encodings/fastlanes/src/for/vtable/mod.rs @@ -48,7 +48,7 @@ mod validity; /// A [`FoR`]-encoded Vortex array. pub type FoRArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { kernels::initialize(session); } diff --git a/encodings/fastlanes/src/for/vtable/operations.rs b/encodings/fastlanes/src/for/vtable/operations.rs index b6deefbe038..ae0f197d0fa 100644 --- a/encodings/fastlanes/src/for/vtable/operations.rs +++ b/encodings/fastlanes/src/for/vtable/operations.rs @@ -50,9 +50,9 @@ mod test { use crate::FoRData; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index e6192a9cdfb..85a5dbe9d7a 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -17,8 +17,9 @@ //! session before deserializing or executing arrays that may contain these encodings. //! //! ```rust -//! let session = vortex_array::array_session(); -//! vortex_fastlanes::initialize(&session); +//! let mut session = vortex_array::default_session_builder(); +//! vortex_fastlanes::initialize(&mut session); +//! let _session = session.build(); //! ``` //! //! ## Paper @@ -53,43 +54,41 @@ use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; +use vortex_array::aggregate_fn::session::AggregateFnSession; use vortex_array::arrays::patched::use_experimental_patches; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; /// Initialize fastlanes encodings in the given session. -pub fn initialize(session: &VortexSession) { - // If we're using the experimental Patched encoding, register a shim - // for BitPacked with interior patches decode as Patched array. - if use_experimental_patches() { - session.arrays().register(BitPackedPatchedPlugin); - } else { - session.arrays().register(BitPacked); +pub fn initialize(session: &mut VortexSessionBuilder) { + { + let arrays = session.get_mut::(); + + // If we're using the experimental Patched encoding, register a shim + // for BitPacked with interior patches decode as Patched array. + if use_experimental_patches() { + arrays.register(BitPackedPatchedPlugin); + } else { + arrays.register(BitPacked); + } + arrays.register(Delta); + arrays.register(FoR); + arrays.register(RLE); } - session.arrays().register(Delta); - session.arrays().register(FoR); - session.arrays().register(RLE); + bitpacking::initialize(session); r#for::initialize(session); rle::initialize(session); // Register the encoding-specific aggregate kernels. - session.aggregate_fns().register_aggregate_kernel( + let aggregate_fns = session.get_mut::(); + aggregate_fns.register_aggregate_kernel( BitPacked.id(), Some(IsConstant.id()), &BitPackedIsConstantKernel, ); - session.aggregate_fns().register_aggregate_kernel( - FoR.id(), - Some(IsConstant.id()), - &FoRIsConstantKernel, - ); - session.aggregate_fns().register_aggregate_kernel( - FoR.id(), - Some(IsSorted.id()), - &FoRIsSortedKernel, - ); + aggregate_fns.register_aggregate_kernel(FoR.id(), Some(IsConstant.id()), &FoRIsConstantKernel); + aggregate_fns.register_aggregate_kernel(FoR.id(), Some(IsSorted.id()), &FoRIsSortedKernel); } /// Fill-forward null values in a buffer, replacing each null with the last valid value seen. @@ -169,9 +168,9 @@ mod test { use super::*; pub static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); #[test] diff --git a/encodings/fastlanes/src/rle/array/rle_compress.rs b/encodings/fastlanes/src/rle/array/rle_compress.rs index b7e04a13699..89e6af50325 100644 --- a/encodings/fastlanes/src/rle/array/rle_compress.rs +++ b/encodings/fastlanes/src/rle/array/rle_compress.rs @@ -173,9 +173,9 @@ mod tests { use crate::rle::array::RLEArrayExt; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fastlanes/src/rle/compute/cast.rs b/encodings/fastlanes/src/rle/compute/cast.rs index 92947d927bc..e4d4ef5c232 100644 --- a/encodings/fastlanes/src/rle/compute/cast.rs +++ b/encodings/fastlanes/src/rle/compute/cast.rs @@ -65,9 +65,9 @@ mod tests { use crate::rle::RLEArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn rle(primitive: &PrimitiveArray, ctx: &mut ExecutionCtx) -> RLEArray { diff --git a/encodings/fastlanes/src/rle/kernel.rs b/encodings/fastlanes/src/rle/kernel.rs index 6e89cd6045d..a9af9cff59d 100644 --- a/encodings/fastlanes/src/rle/kernel.rs +++ b/encodings/fastlanes/src/rle/kernel.rs @@ -11,18 +11,20 @@ use vortex_array::IntoArray; use vortex_array::arrays::Slice; use vortex_array::arrays::slice::SliceExecuteAdaptor; use vortex_array::arrays::slice::SliceKernel; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_error::VortexResult; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::FL_CHUNK_SIZE; use crate::RLE; use crate::rle::RLEArrayExt; -pub(crate) fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(Slice.id(), RLE, SliceExecuteAdaptor(RLE)); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + builder_kernels(session).register_execute_parent_kernel( + Slice.id(), + RLE, + SliceExecuteAdaptor(RLE), + ); } impl SliceKernel for RLE { diff --git a/encodings/fastlanes/src/rle/mod.rs b/encodings/fastlanes/src/rle/mod.rs index 649d11f307c..b1c7417a6b0 100644 --- a/encodings/fastlanes/src/rle/mod.rs +++ b/encodings/fastlanes/src/rle/mod.rs @@ -12,6 +12,6 @@ mod vtable; pub use vtable::RLE; pub use vtable::RLEArray; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { kernel::initialize(session); } diff --git a/encodings/fastlanes/src/rle/vtable/operations.rs b/encodings/fastlanes/src/rle/vtable/operations.rs index b8f79db95a6..4ac1dd65e11 100644 --- a/encodings/fastlanes/src/rle/vtable/operations.rs +++ b/encodings/fastlanes/src/rle/vtable/operations.rs @@ -59,9 +59,9 @@ mod tests { use crate::RLEData; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); mod fixture { diff --git a/encodings/fsst/benches/chunked_dict_fsst_builder.rs b/encodings/fsst/benches/chunked_dict_fsst_builder.rs index 8c2c6c2e388..9a766be02d3 100644 --- a/encodings/fsst/benches/chunked_dict_fsst_builder.rs +++ b/encodings/fsst/benches/chunked_dict_fsst_builder.rs @@ -29,9 +29,9 @@ const BENCH_ARGS: &[(usize, usize, usize)] = &[ ]; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fsst::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fsst::initialize(&mut builder); + builder.build() }); fn make_dict_fsst_chunks( diff --git a/encodings/fsst/benches/fsst_compress.rs b/encodings/fsst/benches/fsst_compress.rs index 5e03d1d003d..9e96c7af030 100644 --- a/encodings/fsst/benches/fsst_compress.rs +++ b/encodings/fsst/benches/fsst_compress.rs @@ -33,9 +33,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fsst::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fsst::initialize(&mut builder); + builder.build() }); // [(string_count, avg_len, unique_chars)] diff --git a/encodings/fsst/benches/fsst_like.rs b/encodings/fsst/benches/fsst_like.rs index 414c1b7afaf..a46ffa85606 100644 --- a/encodings/fsst/benches/fsst_like.rs +++ b/encodings/fsst/benches/fsst_like.rs @@ -30,9 +30,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fsst::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fsst::initialize(&mut builder); + builder.build() }); const N: usize = NUM_STRINGS; diff --git a/encodings/fsst/benches/fsst_url_compare.rs b/encodings/fsst/benches/fsst_url_compare.rs index a62dea8aa9e..d84d3760e68 100644 --- a/encodings/fsst/benches/fsst_url_compare.rs +++ b/encodings/fsst/benches/fsst_url_compare.rs @@ -32,9 +32,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_fsst::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_fsst::initialize(&mut builder); + builder.build() }); const NUM_URLS: usize = NUM_STRINGS; diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs index 2526f0a697f..f4ae80a6256 100644 --- a/encodings/fsst/src/array.rs +++ b/encodings/fsst/src/array.rs @@ -841,9 +841,9 @@ mod test { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::accessor::ArrayAccessor; - use vortex_array::array_session; use vortex_array::arrays::VarBinViewArray; use vortex_array::buffer::BufferHandle; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -867,7 +867,7 @@ mod test { let symbol_lengths = Buffer::::copy_from([3, 8]); let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]); let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?; @@ -911,7 +911,7 @@ mod test { let symbol_lengths = Buffer::::copy_from([3, 8]); let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]); let fsst_array = fsst_compress(&input.into_array(), &compressor, &mut ctx).unwrap(); @@ -949,12 +949,14 @@ mod test { .encode_to_vec(), &buffers, &children.as_slice(), - &array_session(), + &default_session_builder().build(), ) .unwrap(); let decompressed = fsst - .execute::(&mut array_session().create_execution_ctx()) + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); decompressed .with_iterator(|it| { diff --git a/encodings/fsst/src/canonical.rs b/encodings/fsst/src/canonical.rs index 86caf98825e..70b435d62f3 100644 --- a/encodings/fsst/src/canonical.rs +++ b/encodings/fsst/src/canonical.rs @@ -110,7 +110,8 @@ mod tests { use crate::fsst_compress; use crate::fsst_train_compressor; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn make_data() -> (VarBinArray, Vec>>) { const STRING_COUNT: usize = 1000; diff --git a/encodings/fsst/src/compress.rs b/encodings/fsst/src/compress.rs index b9f815b474f..36e34fce59c 100644 --- a/encodings/fsst/src/compress.rs +++ b/encodings/fsst/src/compress.rs @@ -326,9 +326,9 @@ impl<'c, O: IntegerPType + 'static> FsstSink<'c, O> { #[cfg(test)] mod tests { use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbin::VarBinArrayExt; + use vortex_array::default_session_builder; use vortex_array::dtype::PType; use vortex_error::VortexResult; @@ -353,7 +353,7 @@ mod tests { #[test] fn codes_offsets_dtype_small_input_is_i32() -> VortexResult<()> { let array = VarBinViewArray::from_iter_str(["hello", "world", "fsst encoded"]); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let compressor = fsst_train_compressor(array.as_array(), &mut ctx)?; let fsst = fsst_compress(array.as_array(), &compressor, &mut ctx)?; assert_eq!(fsst.codes().offsets().dtype().as_ptype(), PType::I32); diff --git a/encodings/fsst/src/compute/cast.rs b/encodings/fsst/src/compute/cast.rs index bdca68e841b..ef82ce46e87 100644 --- a/encodings/fsst/src/compute/cast.rs +++ b/encodings/fsst/src/compute/cast.rs @@ -106,9 +106,9 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs index ec63af3735f..2598eeed88f 100644 --- a/encodings/fsst/src/compute/compare.rs +++ b/encodings/fsst/src/compute/compare.rs @@ -142,9 +142,9 @@ mod tests { use crate::fsst_train_compressor; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/fsst/src/compute/like.rs b/encodings/fsst/src/compute/like.rs index c922fba088e..9c5dcb83f14 100644 --- a/encodings/fsst/src/compute/like.rs +++ b/encodings/fsst/src/compute/like.rs @@ -105,9 +105,9 @@ mod tests { use crate::fsst_train_compressor; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn make_fsst(strings: &[Option<&str>], nullability: Nullability) -> FSSTArray { diff --git a/encodings/fsst/src/compute/mod.rs b/encodings/fsst/src/compute/mod.rs index 0065e0f499a..ff7aa046ff1 100644 --- a/encodings/fsst/src/compute/mod.rs +++ b/encodings/fsst/src/compute/mod.rs @@ -62,11 +62,11 @@ mod tests { use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_error::VortexResult; @@ -77,7 +77,7 @@ mod tests { #[test] fn test_take_null() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable)).into_array(); let compr = fsst_train_compressor(&arr, &mut ctx)?; @@ -113,7 +113,7 @@ mod tests { DType::Utf8(Nullability::NonNullable), ))] fn test_take_fsst_conformance(#[case] varbin: VarBinArray) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let varbin = varbin.into_array(); let compressor = fsst_train_compressor(&varbin, &mut ctx)?; let array = fsst_compress(&varbin, &compressor, &mut ctx)?; @@ -190,11 +190,11 @@ mod tests { })] fn test_fsst_consistency(#[case] build: FsstBuilder) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = build(&mut ctx); test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/encodings/fsst/src/dfa/tests.rs b/encodings/fsst/src/dfa/tests.rs index 3ea9a507c20..b53ec492349 100644 --- a/encodings/fsst/src/dfa/tests.rs +++ b/encodings/fsst/src/dfa/tests.rs @@ -32,9 +32,9 @@ use crate::fsst_compress; use crate::fsst_train_compressor; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); /// Helper: make a Symbol from a byte string (up to 8 bytes, zero-padded). diff --git a/encodings/fsst/src/kernel.rs b/encodings/fsst/src/kernel.rs index 30f25006195..dee9b563f33 100644 --- a/encodings/fsst/src/kernel.rs +++ b/encodings/fsst/src/kernel.rs @@ -6,7 +6,7 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::Filter; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; @@ -16,12 +16,12 @@ use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; use vortex_array::scalar_fn::fns::like::Like; use vortex_array::scalar_fn::fns::like::LikeExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::FSST; -pub(super) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(super) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Cast.id(), FSST, CastExecuteAdaptor(FSST)); kernels.register_execute_parent_kernel(Binary.id(), FSST, CompareExecuteAdaptor(FSST)); kernels.register_execute_parent_kernel(Filter.id(), FSST, FilterExecuteAdaptor(FSST)); @@ -55,9 +55,9 @@ mod tests { use crate::fsst_train_compressor; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn build_test_fsst_array() -> ArrayRef { diff --git a/encodings/fsst/src/lib.rs b/encodings/fsst/src/lib.rs index 70dcc705249..c22fdf43c92 100644 --- a/encodings/fsst/src/lib.rs +++ b/encodings/fsst/src/lib.rs @@ -27,11 +27,11 @@ mod tests; pub use array::*; pub use compress::*; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; /// Initialize FSST encoding in the given session. -pub fn initialize(session: &VortexSession) { - session.arrays().register(FSST); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(FSST); kernel::initialize(session); } diff --git a/encodings/fsst/src/tests.rs b/encodings/fsst/src/tests.rs index 43cdeb1a02c..d11860038c4 100644 --- a/encodings/fsst/src/tests.rs +++ b/encodings/fsst/src/tests.rs @@ -23,9 +23,9 @@ use crate::fsst_compress; use crate::fsst_train_compressor; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); /// this function is VERY slow on miri, so we only want to run it once diff --git a/encodings/parquet-variant/src/array.rs b/encodings/parquet-variant/src/array.rs index 7e55357c78f..b4b3dc143f7 100644 --- a/encodings/parquet-variant/src/array.rs +++ b/encodings/parquet-variant/src/array.rs @@ -531,9 +531,9 @@ mod tests { use crate::array::parquet_typed_value_from_logical_shredded; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn assert_arrow_variant_storage_roundtrip(struct_array: StructArray) -> VortexResult<()> { diff --git a/encodings/parquet-variant/src/arrow.rs b/encodings/parquet-variant/src/arrow.rs index da4e648a138..a99ac9954f6 100644 --- a/encodings/parquet-variant/src/arrow.rs +++ b/encodings/parquet-variant/src/arrow.rs @@ -314,9 +314,9 @@ mod tests { #[fixture] fn session() -> VortexSession { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + crate::initialize(&mut builder); + builder.build() } fn arrow_variant_storage() -> StructArray { diff --git a/encodings/parquet-variant/src/json_to_variant_tests.rs b/encodings/parquet-variant/src/json_to_variant_tests.rs index cd3b2f035ee..1932a9020b7 100644 --- a/encodings/parquet-variant/src/json_to_variant_tests.rs +++ b/encodings/parquet-variant/src/json_to_variant_tests.rs @@ -41,9 +41,9 @@ use crate::ShreddingSpec; use crate::json_to_variant; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn i64_dtype() -> DType { diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index bd7d0bfaa7c..804cab554d8 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/encodings/parquet-variant/src/kernel.rs @@ -41,7 +41,7 @@ use vortex_array::arrow::ArrowSessionExt; use vortex_array::arrow::FromArrowArray; use vortex_array::dtype::DType; use vortex_array::kernel::ExecuteParentKernel; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::variant_get::VariantGet; use vortex_array::scalar_fn::fns::variant_get::VariantPath; @@ -52,14 +52,14 @@ use vortex_error::vortex_err; use vortex_json::Json; use vortex_json::JsonToVariant; use vortex_mask::Mask; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ParquetVariant; use crate::ParquetVariantArrayExt; use crate::compute::AllNonDistinctParquetVariant; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel( Filter.id(), ParquetVariant, @@ -340,9 +340,9 @@ mod tests { use crate::ParquetVariantArrayExt; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn make_unshredded_array() -> VortexResult { diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index c71ea81e734..9c16598999e 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/encodings/parquet-variant/src/lib.rs @@ -37,13 +37,13 @@ mod vtable; use std::sync::Arc; pub use array::ParquetVariantArrayExt; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::session::ArraySessionExt; +use vortex_array::arrow::ArrowSession; +use vortex_array::session::ArraySession; pub use vortex_json::JsonToVariant; pub use vortex_json::JsonToVariantOptions; pub use vortex_json::ShreddingSpec; pub use vortex_json::json_to_variant; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; pub use vtable::ParquetVariant; pub use vtable::ParquetVariantArray; @@ -52,10 +52,12 @@ pub use vtable::ParquetVariantArray; /// /// This also initializes [`vortex_json`], registering the `Json` extension dtype and the /// `json_to_variant` scalar function whose execution this crate provides. -pub fn initialize(session: &VortexSession) { +pub fn initialize(session: &mut VortexSessionBuilder) { vortex_json::initialize(session); - session.arrays().register(ParquetVariant); + session.get_mut::().register(ParquetVariant); kernel::initialize(session); - session.arrow().register_exporter(Arc::new(ParquetVariant)); - session.arrow().register_importer(Arc::new(ParquetVariant)); + + let arrow = session.get_mut::(); + arrow.register_exporter(Arc::new(ParquetVariant)); + arrow.register_importer(Arc::new(ParquetVariant)); } diff --git a/encodings/parquet-variant/src/operations.rs b/encodings/parquet-variant/src/operations.rs index 3cc3d2e5b4a..281bfee4c0a 100644 --- a/encodings/parquet-variant/src/operations.rs +++ b/encodings/parquet-variant/src/operations.rs @@ -388,7 +388,7 @@ mod tests { use parquet_variant_compute::VariantArray as ArrowVariantArray; use parquet_variant_compute::VariantArrayBuilder; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; @@ -420,7 +420,10 @@ mod tests { Some(ScalarValue::Variant(Box::new(expected_inner))), )?; assert_eq!( - vortex_arr.execute_scalar(index, &mut array_session().create_execution_ctx())?, + vortex_arr.execute_scalar( + index, + &mut default_session_builder().build().create_execution_ctx() + )?, expected ); } @@ -452,17 +455,26 @@ mod tests { assert!( vortex_arr - .execute_scalar(1, &mut array_session().create_execution_ctx())? + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + )? .is_null() ); assert!( !vortex_arr - .execute_scalar(0, &mut array_session().create_execution_ctx())? + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + )? .is_null() ); assert!( !vortex_arr - .execute_scalar(2, &mut array_session().create_execution_ctx())? + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + )? .is_null() ); @@ -486,21 +498,27 @@ mod tests { let arrow_variant = ArrowVariantArray::try_new(&null_struct)?; let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; - let present_variant_null = - vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; + let present_variant_null = vortex_arr.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert!(!present_variant_null.is_null()); assert_eq!( present_variant_null.as_variant().is_variant_null(), Some(true) ); - let outer_null = - vortex_arr.execute_scalar(1, &mut array_session().create_execution_ctx())?; + let outer_null = vortex_arr.execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert!(outer_null.is_null()); assert_eq!(outer_null.as_variant().is_variant_null(), None); - let present_value = - vortex_arr.execute_scalar(2, &mut array_session().create_execution_ctx())?; + let present_value = vortex_arr.execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert!(!present_value.is_null()); assert_eq!(present_value.as_variant().is_variant_null(), Some(false)); @@ -526,17 +544,23 @@ mod tests { assert_eq!(vortex_arr.dtype(), &DType::Variant(Nullability::Nullable)); assert!( vortex_arr - .execute_scalar(0, &mut array_session().create_execution_ctx())? + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + )? .is_null() ); assert!( vortex_arr - .execute_scalar(1, &mut array_session().create_execution_ctx())? + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + )? .is_null() ); let inner_pv = vortex_arr.as_opt::().unwrap(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let roundtripped = inner_pv.to_arrow(&mut ctx)?; assert_eq!(roundtripped.inner().null_count(), 2); @@ -558,12 +582,18 @@ mod tests { ); assert!( !vortex_arr - .execute_scalar(0, &mut array_session().create_execution_ctx())? + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + )? .is_null() ); assert!( !vortex_arr - .execute_scalar(1, &mut array_session().create_execution_ctx())? + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + )? .is_null() ); assert_scalar_at_matches_arrow_try_value(&arrow_variant, [0, 1])?; @@ -673,7 +703,10 @@ mod tests { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; - let row0 = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; + let row0 = vortex_arr.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + )?; let row0 = row0.as_variant().value().unwrap().as_list(); assert_eq!(row0.len(), 2); assert_eq!( @@ -697,7 +730,10 @@ mod tests { Some(20) ); - let row1 = vortex_arr.execute_scalar(1, &mut array_session().create_execution_ctx())?; + let row1 = vortex_arr.execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + )?; let row1 = row1.as_variant().value().unwrap().as_list(); assert_eq!(row1.len(), 1); assert_eq!( @@ -764,7 +800,10 @@ mod tests { let arrow_variant = ArrowVariantArray::try_new(&struct_array)?; let vortex_arr = ParquetVariant::from_arrow_variant(&arrow_variant)?; - let object = vortex_arr.execute_scalar(0, &mut array_session().create_execution_ctx())?; + let object = vortex_arr.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + )?; let object = object.as_variant().value().unwrap().as_struct(); assert_eq!( diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index 5aae70c7d24..ec320c96776 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -334,16 +334,16 @@ mod tests { use crate::array::ParquetVariantArrayExt; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn roundtrip(array: ArrayRef) -> VortexResult { let dtype = array.dtype().clone(); let len = array.len(); - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(ParquetVariant); let ctx = ArrayContext::empty(); @@ -382,10 +382,11 @@ mod tests { #[fixture] fn parquet_variant_file_session() -> VortexSession { - let session = vortex_array::array_session() + let mut builder = vortex_array::default_session_builder() .with::() .with::(); - vortex_file::register_default_encodings(&session); + vortex_file::register_default_encodings(&mut builder); + let session = builder.build(); session.arrays().register(ParquetVariant); session } diff --git a/encodings/pco/src/array.rs b/encodings/pco/src/array.rs index d0804cd80c1..341aeea932b 100644 --- a/encodings/pco/src/array.rs +++ b/encodings/pco/src/array.rs @@ -681,9 +681,9 @@ impl OperationsVTable for Pco { mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use vortex_array::validity::Validity; use vortex_buffer::buffer; @@ -691,7 +691,7 @@ mod tests { #[test] fn test_slice_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create a nullable array with some nulls let values = PrimitiveArray::new( buffer![10u32, 20, 30, 40, 50, 60], diff --git a/encodings/pco/src/compute/cast.rs b/encodings/pco/src/compute/cast.rs index 5f562fafca6..823c7941ef3 100644 --- a/encodings/pco/src/compute/cast.rs +++ b/encodings/pco/src/compute/cast.rs @@ -71,7 +71,8 @@ mod tests { use crate::Pco; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_cast_pco_f32_to_f64() { diff --git a/encodings/pco/src/compute/mod.rs b/encodings/pco/src/compute/mod.rs index d86c3d4b0a3..66940e2386f 100644 --- a/encodings/pco/src/compute/mod.rs +++ b/encodings/pco/src/compute/mod.rs @@ -8,9 +8,9 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use crate::Pco; use crate::PcoArray; @@ -21,7 +21,7 @@ mod tests { values.as_view(), 0, 128, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -32,7 +32,7 @@ mod tests { values.as_view(), 0, 128, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -43,7 +43,7 @@ mod tests { values.as_view(), 0, 128, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -54,7 +54,7 @@ mod tests { values.as_view(), 0, 128, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -65,7 +65,7 @@ mod tests { values.as_view(), 0, 128, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -76,7 +76,7 @@ mod tests { values.as_view(), 0, 128, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -87,7 +87,7 @@ mod tests { values.as_view(), 0, 128, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -98,7 +98,7 @@ mod tests { values.as_view(), 3, 128, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() } @@ -115,7 +115,7 @@ mod tests { fn test_pco_consistency(#[case] array: PcoArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index 601bd0826bd..2660e4a9e23 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -12,10 +12,11 @@ //! To deserialize arrays manually, register the encoding in the array session: //! //! ```rust -//! use vortex_array::session::ArraySessionExt; +//! use vortex_array::session::ArraySession; //! -//! let session = vortex_array::array_session(); -//! session.arrays().register(vortex_pco::Pco); +//! let mut session = vortex_array::default_session_builder(); +//! session.get_mut::().register(vortex_pco::Pco); +//! let _session = session.build(); //! ``` mod array; diff --git a/encodings/pco/src/tests.rs b/encodings/pco/src/tests.rs index 573bd6a0fb4..ff4a1d5f13c 100644 --- a/encodings/pco/src/tests.rs +++ b/encodings/pco/src/tests.rs @@ -31,7 +31,7 @@ use vortex_session::registry::ReadContext; use crate::PcoData; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(Pco); session }); diff --git a/encodings/runend/benches/run_end_compress.rs b/encodings/runend/benches/run_end_compress.rs index cd241ad42c6..ae811b7569e 100644 --- a/encodings/runend/benches/run_end_compress.rs +++ b/encodings/runend/benches/run_end_compress.rs @@ -24,9 +24,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_runend::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_runend::initialize(&mut builder); + builder.build() }); const BENCH_ARGS: &[(usize, usize)] = &[ diff --git a/encodings/runend/benches/run_end_decode.rs b/encodings/runend/benches/run_end_decode.rs index 9256a4b3a53..d01ece52b1f 100644 --- a/encodings/runend/benches/run_end_decode.rs +++ b/encodings/runend/benches/run_end_decode.rs @@ -21,9 +21,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_runend::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_runend::initialize(&mut builder); + builder.build() }); /// Distribution types for bool benchmarks diff --git a/encodings/runend/benches/run_end_null_count.rs b/encodings/runend/benches/run_end_null_count.rs index f214889daae..de7b1c9726c 100644 --- a/encodings/runend/benches/run_end_null_count.rs +++ b/encodings/runend/benches/run_end_null_count.rs @@ -50,9 +50,9 @@ const BENCH_ARGS: &[(usize, usize, f64)] = &[ ]; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_runend::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_runend::initialize(&mut builder); + builder.build() }); #[divan::bench(args = BENCH_ARGS)] diff --git a/encodings/runend/benches/run_end_take.rs b/encodings/runend/benches/run_end_take.rs index f0dccccfc48..ff55902a1a6 100644 --- a/encodings/runend/benches/run_end_take.rs +++ b/encodings/runend/benches/run_end_take.rs @@ -25,9 +25,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_runend::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_runend::initialize(&mut builder); + builder.build() }); #[derive(Clone, Copy)] diff --git a/encodings/runend/src/array.rs b/encodings/runend/src/array.rs index fde553a2fd4..764f2d65311 100644 --- a/encodings/runend/src/array.rs +++ b/encodings/runend/src/array.rs @@ -406,8 +406,9 @@ impl RunEndData { /// # use vortex_error::VortexResult; /// # use vortex_runend::RunEnd; /// # fn main() -> VortexResult<()> { - /// let session = vortex_array::array_session(); - /// vortex_runend::initialize(&session); + /// let mut session = vortex_array::default_session_builder(); + /// vortex_runend::initialize(&mut session); + /// let session = session.build(); /// let mut ctx = session.create_execution_ctx(); /// let ends = buffer![2u8, 3u8].into_array(); /// let values = BoolArray::from_iter([false, true]).into_array(); @@ -519,9 +520,9 @@ mod tests { use crate::RunEnd; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/runend/src/arrow.rs b/encodings/runend/src/arrow.rs index 193937122ef..34e0db3d36e 100644 --- a/encodings/runend/src/arrow.rs +++ b/encodings/runend/src/arrow.rs @@ -104,9 +104,9 @@ mod tests { use crate::ops::find_slice_end_index; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn decode_run_array( diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index 544b677dfa3..26cb0813c7e 100644 --- a/encodings/runend/src/compress.rs +++ b/encodings/runend/src/compress.rs @@ -335,9 +335,9 @@ mod tests { use crate::compress::runend_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/runend/src/compute/cast.rs b/encodings/runend/src/compute/cast.rs index fe740739ca9..c3ddaea9cb8 100644 --- a/encodings/runend/src/compute/cast.rs +++ b/encodings/runend/src/compute/cast.rs @@ -53,9 +53,9 @@ mod tests { use crate::RunEndArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/runend/src/compute/compare.rs b/encodings/runend/src/compute/compare.rs index ecdd03e3686..27d80254c79 100644 --- a/encodings/runend/src/compute/compare.rs +++ b/encodings/runend/src/compute/compare.rs @@ -65,9 +65,9 @@ mod test { use crate::RunEndArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn ree_array(ctx: &mut ExecutionCtx) -> RunEndArray { diff --git a/encodings/runend/src/compute/mod.rs b/encodings/runend/src/compute/mod.rs index fc7fc8804ec..d4cefa8e62c 100644 --- a/encodings/runend/src/compute/mod.rs +++ b/encodings/runend/src/compute/mod.rs @@ -27,9 +27,9 @@ mod tests { use crate::RunEndArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index c61854635e3..98b4284587f 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -508,9 +508,9 @@ mod tests { use crate::RunEndArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn ree_array() -> RunEndArray { diff --git a/encodings/runend/src/compute/take_from.rs b/encodings/runend/src/compute/take_from.rs index 0987cd3fc55..e3640926515 100644 --- a/encodings/runend/src/compute/take_from.rs +++ b/encodings/runend/src/compute/take_from.rs @@ -72,9 +72,9 @@ mod tests { use crate::compute::take_from::RunEndTakeFrom; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); /// Build a DictArray whose codes are run-end encoded. diff --git a/encodings/runend/src/decompress_bool.rs b/encodings/runend/src/decompress_bool.rs index d2b778fdbbf..535064ad8ef 100644 --- a/encodings/runend/src/decompress_bool.rs +++ b/encodings/runend/src/decompress_bool.rs @@ -261,9 +261,9 @@ mod tests { use super::runend_decode_bools; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/runend/src/kernel.rs b/encodings/runend/src/kernel.rs index 484a4233789..de5e0a7bbff 100644 --- a/encodings/runend/src/kernel.rs +++ b/encodings/runend/src/kernel.rs @@ -15,19 +15,19 @@ use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::kernel::ExecuteParentKernel; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_error::VortexResult; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::RunEnd; use crate::array::RunEndArrayExt; use crate::compute::take_from::RunEndTakeFrom; -pub(super) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(super) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Binary.id(), RunEnd, CompareExecuteAdaptor(RunEnd)); kernels.register_execute_parent_kernel(Slice.id(), RunEnd, RunEndSliceKernel); kernels.register_execute_parent_kernel(Filter.id(), RunEnd, FilterExecuteAdaptor(RunEnd)); diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 8169f5d8dfc..1257524a2ca 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -31,27 +31,28 @@ use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; use vortex_array::aggregate_fn::fns::min_max::MinMax; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; /// Initialize run-end encoding in the given session. -pub fn initialize(session: &VortexSession) { - session.arrays().register(RunEnd); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(RunEnd); kernel::initialize(session); // Register the RunEnd-specific aggregate kernels. - session.aggregate_fns().register_aggregate_kernel( + let aggregate_fns = session.get_mut::(); + aggregate_fns.register_aggregate_kernel( RunEnd.id(), Some(MinMax.id()), &compute::min_max::RunEndMinMaxKernel, ); - session.aggregate_fns().register_aggregate_kernel( + aggregate_fns.register_aggregate_kernel( RunEnd.id(), Some(IsConstant.id()), &compute::is_constant::RunEndIsConstantKernel, ); - session.aggregate_fns().register_aggregate_kernel( + aggregate_fns.register_aggregate_kernel( RunEnd.id(), Some(IsSorted.id()), &compute::is_sorted::RunEndIsSortedKernel, @@ -70,9 +71,9 @@ mod tests { use crate::RunEndMetadata; pub static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + crate::initialize(&mut builder); + builder.build() }); #[cfg_attr(miri, ignore)] diff --git a/encodings/runend/src/ops.rs b/encodings/runend/src/ops.rs index 56de024f988..f52716fe7b2 100644 --- a/encodings/runend/src/ops.rs +++ b/encodings/runend/src/ops.rs @@ -65,9 +65,9 @@ mod tests { use crate::RunEnd; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/sequence/src/array.rs b/encodings/sequence/src/array.rs index 56e62764ce4..da8b690c1f8 100644 --- a/encodings/sequence/src/array.rs +++ b/encodings/sequence/src/array.rs @@ -459,9 +459,9 @@ mod tests { use crate::Sequence; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/sequence/src/compress.rs b/encodings/sequence/src/compress.rs index b20980cf132..ab3bfea70e3 100644 --- a/encodings/sequence/src/compress.rs +++ b/encodings/sequence/src/compress.rs @@ -158,9 +158,9 @@ mod tests { use crate::sequence_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/sequence/src/compute/cast.rs b/encodings/sequence/src/compute/cast.rs index e0a63d0c922..83239464dc4 100644 --- a/encodings/sequence/src/compute/cast.rs +++ b/encodings/sequence/src/compute/cast.rs @@ -105,9 +105,9 @@ mod tests { use crate::SequenceArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/sequence/src/compute/compare.rs b/encodings/sequence/src/compute/compare.rs index 5895330ea81..7974b5316c6 100644 --- a/encodings/sequence/src/compute/compare.rs +++ b/encodings/sequence/src/compute/compare.rs @@ -151,9 +151,9 @@ mod tests { use crate::Sequence; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/sequence/src/compute/list_contains.rs b/encodings/sequence/src/compute/list_contains.rs index 93bd823a5f2..2e85315b8bb 100644 --- a/encodings/sequence/src/compute/list_contains.rs +++ b/encodings/sequence/src/compute/list_contains.rs @@ -69,9 +69,9 @@ mod tests { use crate::Sequence; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/sequence/src/compute/mod.rs b/encodings/sequence/src/compute/mod.rs index 22b4e3ac05b..972eafe8deb 100644 --- a/encodings/sequence/src/compute/mod.rs +++ b/encodings/sequence/src/compute/mod.rs @@ -15,8 +15,8 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use vortex_array::dtype::Nullability; use crate::Sequence; @@ -88,7 +88,7 @@ mod tests { fn test_sequence_consistency(#[case] array: SequenceArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/encodings/sequence/src/compute/take.rs b/encodings/sequence/src/compute/take.rs index 4b056d0ae7f..ad61bbd185e 100644 --- a/encodings/sequence/src/compute/take.rs +++ b/encodings/sequence/src/compute/take.rs @@ -105,8 +105,8 @@ mod test { use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::default_session_builder; use vortex_array::dtype::Nullability; use crate::Sequence; @@ -174,7 +174,7 @@ mod test { let _array = array .take(indices.into_array()) .unwrap() - .execute::(&mut array_session().create_execution_ctx()) + .execute::(&mut default_session_builder().build().create_execution_ctx()) .unwrap(); } } diff --git a/encodings/sequence/src/kernel.rs b/encodings/sequence/src/kernel.rs index 37521a261e7..7475783ff97 100644 --- a/encodings/sequence/src/kernel.rs +++ b/encodings/sequence/src/kernel.rs @@ -6,16 +6,16 @@ use vortex_array::arrays::Dict; use vortex_array::arrays::Filter; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::Sequence; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Binary.id(), Sequence, CompareExecuteAdaptor(Sequence)); kernels.register_execute_parent_kernel(Filter.id(), Sequence, FilterExecuteAdaptor(Sequence)); kernels.register_execute_parent_kernel(Dict.id(), Sequence, TakeExecuteAdaptor(Sequence)); diff --git a/encodings/sequence/src/lib.rs b/encodings/sequence/src/lib.rs index b898963c346..6b100eaab87 100644 --- a/encodings/sequence/src/lib.rs +++ b/encodings/sequence/src/lib.rs @@ -20,22 +20,23 @@ use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::is_sorted::IsSorted; use vortex_array::aggregate_fn::fns::min_max::MinMax; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; /// Initialize sequence encoding in the given session. -pub fn initialize(session: &VortexSession) { - session.arrays().register(Sequence); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(Sequence); kernel::initialize(session); // Register the Sequence-specific aggregate kernels. - session.aggregate_fns().register_aggregate_kernel( + let aggregate_fns = session.get_mut::(); + aggregate_fns.register_aggregate_kernel( Sequence.id(), Some(MinMax.id()), &compute::min_max::SequenceMinMaxKernel, ); - session.aggregate_fns().register_aggregate_kernel( + aggregate_fns.register_aggregate_kernel( Sequence.id(), Some(IsSorted.id()), &compute::is_sorted::SequenceIsSortedKernel, diff --git a/encodings/sparse/benches/sparse_canonical.rs b/encodings/sparse/benches/sparse_canonical.rs index 27f7340fb91..66a35a063b5 100644 --- a/encodings/sparse/benches/sparse_canonical.rs +++ b/encodings/sparse/benches/sparse_canonical.rs @@ -28,9 +28,9 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_sparse::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_sparse::initialize(&mut builder); + builder.build() }); const LIST_ARGS: &[(usize, usize, usize)] = &[ diff --git a/encodings/sparse/benches/sparse_pushdown.rs b/encodings/sparse/benches/sparse_pushdown.rs index a066e196d76..beae6fdf135 100644 --- a/encodings/sparse/benches/sparse_pushdown.rs +++ b/encodings/sparse/benches/sparse_pushdown.rs @@ -45,9 +45,9 @@ const LEN: usize = 1_000_000; /// Session with Sparse and its pushdown kernels registered. static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_sparse::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + vortex_sparse::initialize(&mut builder); + builder.build() }); /// Build a sparse `i32` array of `LEN` with `num_patches` uniformly-spaced patches and diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 1829d4fffd1..ca92877e248 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -636,9 +636,9 @@ mod test { use crate::Sparse; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/sparse/src/compute/between.rs b/encodings/sparse/src/compute/between.rs index 46fcd7f1e59..5bde8d7c12e 100644 --- a/encodings/sparse/src/compute/between.rs +++ b/encodings/sparse/src/compute/between.rs @@ -81,9 +81,9 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); #[rstest] diff --git a/encodings/sparse/src/compute/cast.rs b/encodings/sparse/src/compute/cast.rs index 39c21805426..418462b34d8 100644 --- a/encodings/sparse/src/compute/cast.rs +++ b/encodings/sparse/src/compute/cast.rs @@ -55,9 +55,9 @@ mod tests { use crate::SparseArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/sparse/src/compute/compare.rs b/encodings/sparse/src/compute/compare.rs index 75bb78622df..4f8dbb4e997 100644 --- a/encodings/sparse/src/compute/compare.rs +++ b/encodings/sparse/src/compute/compare.rs @@ -73,9 +73,9 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); #[rstest] diff --git a/encodings/sparse/src/compute/fill_null.rs b/encodings/sparse/src/compute/fill_null.rs index f36032265bc..5f6d32f1775 100644 --- a/encodings/sparse/src/compute/fill_null.rs +++ b/encodings/sparse/src/compute/fill_null.rs @@ -62,9 +62,9 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); fn nullable_i32() -> DType { diff --git a/encodings/sparse/src/compute/filter.rs b/encodings/sparse/src/compute/filter.rs index 91e67297fd2..f5605d76ced 100644 --- a/encodings/sparse/src/compute/filter.rs +++ b/encodings/sparse/src/compute/filter.rs @@ -57,9 +57,9 @@ mod tests { use crate::Sparse; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[fixture] diff --git a/encodings/sparse/src/compute/is_constant.rs b/encodings/sparse/src/compute/is_constant.rs index eee8683df0a..34e174f02b8 100644 --- a/encodings/sparse/src/compute/is_constant.rs +++ b/encodings/sparse/src/compute/is_constant.rs @@ -84,14 +84,14 @@ mod tests { /// Session with Sparse + its pushdown kernels. static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); /// Baseline session: Sparse registered but no pushdown kernels. static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/min_max.rs b/encodings/sparse/src/compute/min_max.rs index 553b8eaaee0..d46bcc50d9f 100644 --- a/encodings/sparse/src/compute/min_max.rs +++ b/encodings/sparse/src/compute/min_max.rs @@ -78,13 +78,13 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/mod.rs b/encodings/sparse/src/compute/mod.rs index bdae87e560d..db89319f672 100644 --- a/encodings/sparse/src/compute/mod.rs +++ b/encodings/sparse/src/compute/mod.rs @@ -23,13 +23,13 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::compute::conformance::mask::test_mask_conformance; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -43,9 +43,9 @@ mod tests { use crate::SparseArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session(); - crate::initialize(&session); - session + let mut session = default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[fixture] diff --git a/encodings/sparse/src/compute/nan_count.rs b/encodings/sparse/src/compute/nan_count.rs index 9a8a3a871c1..2f0a4556e91 100644 --- a/encodings/sparse/src/compute/nan_count.rs +++ b/encodings/sparse/src/compute/nan_count.rs @@ -86,13 +86,13 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/null_count.rs b/encodings/sparse/src/compute/null_count.rs index f63dad80205..8e7fec9923d 100644 --- a/encodings/sparse/src/compute/null_count.rs +++ b/encodings/sparse/src/compute/null_count.rs @@ -76,13 +76,13 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/sum.rs b/encodings/sparse/src/compute/sum.rs index 8f85124b3b5..f5c30f86996 100644 --- a/encodings/sparse/src/compute/sum.rs +++ b/encodings/sparse/src/compute/sum.rs @@ -81,13 +81,13 @@ mod tests { use crate::initialize; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + builder.build() }); static CANONICAL_SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); session.arrays().register(Sparse); session }); diff --git a/encodings/sparse/src/compute/take.rs b/encodings/sparse/src/compute/take.rs index e260eb26e9e..38bfab89661 100644 --- a/encodings/sparse/src/compute/take.rs +++ b/encodings/sparse/src/compute/take.rs @@ -74,9 +74,9 @@ mod test { use crate::SparseArray; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); fn test_array_fill_value() -> Scalar { diff --git a/encodings/sparse/src/kernel.rs b/encodings/sparse/src/kernel.rs index 30e1f9a6d6a..a05cea513a8 100644 --- a/encodings/sparse/src/kernel.rs +++ b/encodings/sparse/src/kernel.rs @@ -8,7 +8,7 @@ use vortex_array::arrays::Slice; use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::arrays::filter::FilterExecuteAdaptor; use vortex_array::arrays::slice::SliceExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::optimizer::kernels::builder_kernels; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::between::Between; use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; @@ -16,12 +16,12 @@ use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; use vortex_array::scalar_fn::fns::fill_null::FillNull; use vortex_array::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::Sparse; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Between.id(), Sparse, BetweenExecuteAdaptor(Sparse)); kernels.register_execute_parent_kernel(Binary.id(), Sparse, CompareExecuteAdaptor(Sparse)); kernels.register_execute_parent_kernel(FillNull.id(), Sparse, FillNullExecuteAdaptor(Sparse)); diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 8ae38b56052..ce0502bd237 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -73,18 +73,19 @@ use vortex_array::aggregate_fn::fns::min_max::MinMax; use vortex_array::aggregate_fn::fns::nan_count::NanCount; use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::aggregate_fn::fns::sum::Sum; -use vortex_array::aggregate_fn::session::AggregateFnSessionExt; -use vortex_array::session::ArraySessionExt; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; /// Initialize Sparse encoding in the given session. /// /// Registers the Sparse array vtable, parent execution kernels, and aggregate kernels /// (`IsConstant`, `Sum`, `MinMax`, `NullCount`, `NanCount`). -pub fn initialize(session: &VortexSession) { - session.arrays().register(Sparse); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(Sparse); kernel::initialize(session); - let aggregate_fns = session.aggregate_fns(); + let aggregate_fns = session.get_mut::(); aggregate_fns.register_aggregate_kernel( Sparse.id(), Some(IsConstant.id()), @@ -709,9 +710,9 @@ mod test { use crate::Sparse; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + initialize(&mut session); + session.build() }); fn nullable_fill() -> Scalar { diff --git a/encodings/sparse/src/ops.rs b/encodings/sparse/src/ops.rs index 568d8d377d1..f8a4da18cdb 100644 --- a/encodings/sparse/src/ops.rs +++ b/encodings/sparse/src/ops.rs @@ -37,9 +37,9 @@ mod tests { use crate::Sparse; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/zigzag/src/array.rs b/encodings/zigzag/src/array.rs index 8c4be0d5544..18051154549 100644 --- a/encodings/zigzag/src/array.rs +++ b/encodings/zigzag/src/array.rs @@ -262,8 +262,8 @@ impl ValidityChild for ZigZag { mod test { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::default_session_builder; use vortex_array::scalar::Scalar; use vortex_buffer::buffer; @@ -272,7 +272,7 @@ mod test { #[test] fn test_compute_statistics() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = buffer![1i32, -5i32, 2, 3, 4, 5, 6, 7, 8, 9, 10] .into_array() .execute::(&mut ctx)?; diff --git a/encodings/zigzag/src/compress.rs b/encodings/zigzag/src/compress.rs index e70650a56e1..66dfa8bb69a 100644 --- a/encodings/zigzag/src/compress.rs +++ b/encodings/zigzag/src/compress.rs @@ -89,9 +89,9 @@ mod test { use crate::ZigZag; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[rstest] diff --git a/encodings/zigzag/src/compute/cast.rs b/encodings/zigzag/src/compute/cast.rs index c3fae84e346..012d23f58e3 100644 --- a/encodings/zigzag/src/compute/cast.rs +++ b/encodings/zigzag/src/compute/cast.rs @@ -44,9 +44,9 @@ mod tests { use crate::zigzag_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/zigzag/src/compute/mod.rs b/encodings/zigzag/src/compute/mod.rs index 1da341989c5..9efc5e76a32 100644 --- a/encodings/zigzag/src/compute/mod.rs +++ b/encodings/zigzag/src/compute/mod.rs @@ -92,9 +92,9 @@ mod tests { use crate::zigzag_encode; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut session = vortex_array::default_session_builder(); + crate::initialize(&mut session); + session.build() }); #[test] diff --git a/encodings/zigzag/src/kernel.rs b/encodings/zigzag/src/kernel.rs index f1d3e96bb01..c2e14a0366e 100644 --- a/encodings/zigzag/src/kernel.rs +++ b/encodings/zigzag/src/kernel.rs @@ -4,13 +4,15 @@ use vortex_array::ArrayVTable; use vortex_array::arrays::Dict; use vortex_array::arrays::dict::TakeExecuteAdaptor; -use vortex_array::optimizer::kernels::ArrayKernelsExt; -use vortex_session::VortexSession; +use vortex_array::optimizer::kernels::builder_kernels; +use vortex_session::VortexSessionBuilder; use crate::ZigZag; -pub(crate) fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(Dict.id(), ZigZag, TakeExecuteAdaptor(ZigZag)); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + builder_kernels(session).register_execute_parent_kernel( + Dict.id(), + ZigZag, + TakeExecuteAdaptor(ZigZag), + ); } diff --git a/encodings/zigzag/src/lib.rs b/encodings/zigzag/src/lib.rs index 01c3f51541b..843d1f1b5fe 100644 --- a/encodings/zigzag/src/lib.rs +++ b/encodings/zigzag/src/lib.rs @@ -3,8 +3,8 @@ pub use array::*; pub use compress::*; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; mod array; mod compress; @@ -14,7 +14,7 @@ mod rules; mod slice; /// Initialize zigzag encoding in the given session. -pub fn initialize(session: &VortexSession) { - session.arrays().register(ZigZag); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(ZigZag); kernel::initialize(session); } diff --git a/encodings/zstd/benches/listview_rebuild.rs b/encodings/zstd/benches/listview_rebuild.rs index e3c0e6d5dd1..b43e74fd773 100644 --- a/encodings/zstd/benches/listview_rebuild.rs +++ b/encodings/zstd/benches/listview_rebuild.rs @@ -18,7 +18,8 @@ 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(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[divan::bench(sample_size = 1000)] fn rebuild_naive(bencher: Bencher) { diff --git a/encodings/zstd/src/compute/cast.rs b/encodings/zstd/src/compute/cast.rs index c5f15eba284..e48b03c15c3 100644 --- a/encodings/zstd/src/compute/cast.rs +++ b/encodings/zstd/src/compute/cast.rs @@ -89,7 +89,8 @@ mod tests { use crate::Zstd; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_cast_zstd_i32_to_i64() { diff --git a/encodings/zstd/src/compute/mod.rs b/encodings/zstd/src/compute/mod.rs index 75ef5273ca5..d31c940c001 100644 --- a/encodings/zstd/src/compute/mod.rs +++ b/encodings/zstd/src/compute/mod.rs @@ -8,9 +8,9 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::consistency::test_array_consistency; + use vortex_array::default_session_builder; use vortex_buffer::buffer; use crate::Zstd; @@ -18,23 +18,47 @@ mod tests { fn zstd_i32() -> ZstdArray { let values = PrimitiveArray::from_iter([100i32, 200, 300, 400, 500]); - Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() + Zstd::from_primitive( + &values, + 0, + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } fn zstd_f64() -> ZstdArray { let values = PrimitiveArray::from_iter([1.1f64, 2.2, 3.3, 4.4, 5.5]); - Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() + Zstd::from_primitive( + &values, + 0, + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } fn zstd_u32() -> ZstdArray { let values = PrimitiveArray::from_iter([10u32, 20, 30, 40, 50]); - Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() + Zstd::from_primitive( + &values, + 0, + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } fn zstd_nullable_i64() -> ZstdArray { let values = PrimitiveArray::from_option_iter([Some(1000i64), None, Some(3000), Some(4000), None]); - Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() + Zstd::from_primitive( + &values, + 0, + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } fn zstd_single() -> ZstdArray { @@ -42,7 +66,13 @@ mod tests { buffer![42i64], vortex_array::validity::Validity::NonNullable, ); - Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() + Zstd::from_primitive( + &values, + 0, + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } fn zstd_large() -> ZstdArray { @@ -50,7 +80,13 @@ mod tests { buffer![0u32..1000], vortex_array::validity::Validity::NonNullable, ); - Zstd::from_primitive(&values, 3, 0, &mut array_session().create_execution_ctx()).unwrap() + Zstd::from_primitive( + &values, + 3, + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } fn zstd_all_same() -> ZstdArray { @@ -58,12 +94,24 @@ mod tests { buffer![42i32; 100], vortex_array::validity::Validity::NonNullable, ); - Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() + Zstd::from_primitive( + &values, + 0, + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } fn zstd_negative() -> ZstdArray { let values = PrimitiveArray::from_iter([-100i32, -50, 0, 50, 100]); - Zstd::from_primitive(&values, 0, 0, &mut array_session().create_execution_ctx()).unwrap() + Zstd::from_primitive( + &values, + 0, + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() } #[rstest] @@ -78,7 +126,7 @@ mod tests { fn test_zstd_consistency(#[case] array: ZstdArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/encodings/zstd/src/lib.rs b/encodings/zstd/src/lib.rs index a178fdec142..83ad01a3078 100644 --- a/encodings/zstd/src/lib.rs +++ b/encodings/zstd/src/lib.rs @@ -15,10 +15,11 @@ //! array session: //! //! ```rust -//! use vortex_array::session::ArraySessionExt; +//! use vortex_array::session::ArraySession; //! -//! let session = vortex_array::array_session(); -//! session.arrays().register(vortex_zstd::Zstd); +//! let mut session = vortex_array::default_session_builder(); +//! session.get_mut::().register(vortex_zstd::Zstd); +//! let _session = session.build(); //! ``` pub use array::*; diff --git a/encodings/zstd/src/test.rs b/encodings/zstd/src/test.rs index cf45eb97c04..923afd9775b 100644 --- a/encodings/zstd/src/test.rs +++ b/encodings/zstd/src/test.rs @@ -4,12 +4,12 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; +use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; @@ -21,7 +21,7 @@ use crate::Zstd; #[test] fn test_zstd_compress_decompress() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data: Vec = (0..200).collect(); let array = PrimitiveArray::from_iter(data.clone()); @@ -55,7 +55,7 @@ fn test_zstd_compress_decompress() { #[test] fn test_zstd_empty() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data: Vec = vec![]; let array = PrimitiveArray::new( data.iter().cloned().collect::>(), @@ -69,7 +69,7 @@ fn test_zstd_empty() { #[test] fn test_zstd_with_validity_and_multi_frame() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data: Vec = (0..200).collect(); let mut validity: Vec = vec![false; 200]; validity[3] = true; @@ -123,7 +123,7 @@ fn test_zstd_with_validity_and_multi_frame() { #[test] fn test_zstd_with_dict() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data: Vec = (0..200).collect(); let array = PrimitiveArray::new( data.iter().cloned().collect::>(), @@ -153,7 +153,7 @@ fn test_zstd_with_dict() { #[test] fn test_validity_vtable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mask_bools = vec![false, true, true, false, true]; let array = PrimitiveArray::new( (0..5).collect::>(), @@ -164,7 +164,10 @@ fn test_validity_vtable() { assert_eq!( arr.validity() .unwrap() - .execute_mask(arr.len(), &mut array_session().create_execution_ctx()) + .execute_mask( + arr.len(), + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Mask::from_iter(mask_bools) ); @@ -173,7 +176,10 @@ fn test_validity_vtable() { sliced .validity() .unwrap() - .execute_mask(sliced.len(), &mut array_session().create_execution_ctx()) + .execute_mask( + sliced.len(), + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Mask::from_iter(vec![true, true, false]) ); @@ -181,7 +187,7 @@ fn test_validity_vtable() { #[test] fn test_zstd_var_bin_view() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data: [Option<&'static [u8]>; 5] = [ Some(b"foo"), Some(b"bar"), @@ -207,7 +213,7 @@ fn test_zstd_var_bin_view() { #[test] fn test_zstd_decompress_var_bin_view() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data: [Option<&'static [u8]>; 5] = [ Some(b"foo"), Some(b"bar"), @@ -238,7 +244,7 @@ fn test_zstd_decompress_var_bin_view() { #[test] fn test_sliced_array_children() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data: Vec> = (0..10).map(|v| (v != 5).then_some(v)).collect(); let compressed = Zstd::from_primitive(&PrimitiveArray::from_option_iter(data), 0, 100, &mut ctx).unwrap(); @@ -250,7 +256,7 @@ fn test_sliced_array_children() { /// the buffer alignment when compressing primitive arrays. #[test] fn test_zstd_frame_start_buffer_alignment() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data = vec![0u8; 2]; let aligned_buffer = Buffer::copy_from_aligned(&data, Alignment::new(8)); // u8 array now has a 8-byte alignment. diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index f6cb9af586c..072c5394cd2 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -529,10 +529,10 @@ mod tests { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProvider; @@ -585,12 +585,12 @@ mod tests { #[case::empty_primitive(make_empty_primitive_array())] #[case::inlined_varbinview(make_inlined_varbinview_array())] fn test_roundtrip(#[case] input: ArrayRef) -> VortexResult<()> { - let compressed = ZstdBuffers::compress(&input, 3, &array_session())?; + let compressed = ZstdBuffers::compress(&input, 3, &default_session_builder().build())?; assert_eq!(compressed.len(), input.len()); assert_eq!(compressed.dtype(), input.dtype()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let decompressed = compressed.into_array().execute::(&mut ctx)?; assert_arrays_eq!(input, decompressed, &mut ctx); @@ -634,7 +634,7 @@ mod tests { let input = make_primitive_array(); input.statistics().set(Stat::Min, Precision::exact(0i32)); - let compressed = ZstdBuffers::compress(&input, 3, &array_session())?; + let compressed = ZstdBuffers::compress(&input, 3, &default_session_builder().build())?; assert!(!compressed.statistics().get(Stat::Min).is_absent()); Ok(()) @@ -643,9 +643,10 @@ mod tests { #[test] fn test_validity_delegates_for_nullable_input() -> VortexResult<()> { let input = make_nullable_primitive_array(); - let compressed = ZstdBuffers::compress(&input, 3, &array_session())?.into_array(); + let compressed = + ZstdBuffers::compress(&input, 3, &default_session_builder().build())?.into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(compressed.all_valid(&mut ctx)?, input.all_valid(&mut ctx)?); assert_eq!( compressed.all_invalid(&mut ctx)?, diff --git a/fuzz/src/array/fill_null.rs b/fuzz/src/array/fill_null.rs index f3a49b25637..c619e3c85af 100644 --- a/fuzz/src/array/fill_null.rs +++ b/fuzz/src/array/fill_null.rs @@ -301,13 +301,13 @@ mod tests { use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; @@ -322,13 +322,13 @@ mod tests { fn canonical(array: impl IntoArray) -> Canonical { array .into_array() - .execute::(&mut array_session().create_execution_ctx()) + .execute::(&mut default_session_builder().build().create_execution_ctx()) .unwrap() } #[test] fn test_fill_null_primitive() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]); let fill_value = Scalar::from(42i32); @@ -340,7 +340,7 @@ mod tests { #[test] fn test_fill_null_bool() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let data_buffer = BitBuffer::from(vec![true, false, false, false]); let validity_buffer = BitBuffer::from(vec![true, false, true, false]); let array = BoolArray::new(data_buffer, Validity::from(validity_buffer)); @@ -354,7 +354,7 @@ mod tests { #[test] fn test_fill_null_string() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = VarBinViewArray::from_iter( [Some("hello"), None, Some("world")].iter().copied(), DType::Utf8(Nullability::Nullable), @@ -369,7 +369,7 @@ mod tests { #[test] fn test_fill_null_all_invalid() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([None::, None, None]); let fill_value = Scalar::from(100i32); @@ -381,7 +381,7 @@ mod tests { #[test] fn test_fill_null_no_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_iter([1i32, 2, 3]); let fill_value = Scalar::from(42i32); @@ -394,7 +394,7 @@ mod tests { #[test] #[should_panic] fn test_fill_null_with_null_value_errors() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]); let fill_value = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); @@ -411,7 +411,7 @@ mod tests { #[test] fn test_fill_null_decimal_i32() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = DecimalArray::from_option_iter( [Some(100i32), None, Some(300i32), None, Some(500i32)], DecimalDType::new(10, 2), @@ -433,7 +433,7 @@ mod tests { #[test] fn test_fill_null_decimal_i64() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = DecimalArray::from_option_iter( [Some(1000i64), None, Some(3000i64)], DecimalDType::new(15, 3), @@ -453,7 +453,7 @@ mod tests { #[test] fn test_fill_null_decimal_i128() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = DecimalArray::from_option_iter( [Some(10000i128), None, Some(30000i128), None], DecimalDType::new(20, 4), @@ -475,7 +475,7 @@ mod tests { #[test] fn test_fill_null_decimal_all_invalid() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = DecimalArray::from_option_iter([None::, None, None], DecimalDType::new(10, 2)); let fill_value = Scalar::decimal( @@ -498,7 +498,7 @@ mod tests { #[test] fn test_fill_null_decimal_no_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = DecimalArray::from_option_iter( [Some(100i32), Some(200i32), Some(300i32)], DecimalDType::new(10, 2), diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index f3cecec2b10..2f5183e67ea 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -159,7 +159,6 @@ mod tests { use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::FixedSizeListArray; @@ -169,6 +168,7 @@ mod tests { use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; @@ -182,7 +182,7 @@ mod tests { #[test] fn test_mask_null_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = NullArray::new(5); let mask = Mask::from_iter([true, false, true, false, true]); @@ -197,7 +197,7 @@ mod tests { #[test] fn test_mask_bool_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = BoolArray::from_iter([true, false, true, false, true]); let mask = Mask::from_iter([false, true, true, false, true]); @@ -209,7 +209,7 @@ mod tests { #[test] fn test_mask_primitive_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::from_iter([true, false, true, false, true]); @@ -221,7 +221,7 @@ mod tests { #[test] fn test_mask_primitive_array_with_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]); let mask = Mask::from_iter([false, true, true, false, true]); @@ -233,7 +233,7 @@ mod tests { #[test] fn test_mask_decimal_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DecimalDType::new(10, 2); let array = DecimalArray::from_option_iter( [Some(1i128), Some(2), Some(3), Some(4), Some(5)], @@ -250,7 +250,7 @@ mod tests { #[test] fn test_mask_varbinview_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["one", "two", "three", "four", "five"]); let mask = Mask::from_iter([false, true, false, true, false]); @@ -263,7 +263,7 @@ mod tests { #[test] fn test_mask_list_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]).into_array(); let offsets = PrimitiveArray::from_iter([0i32, 2, 4]).into_array(); let sizes = PrimitiveArray::from_iter([2i32, 2, 2]).into_array(); @@ -284,7 +284,7 @@ mod tests { #[test] fn test_mask_fixed_size_list_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]).into_array(); let array = FixedSizeListArray::try_new(elements, 2, Nullability::NonNullable.into(), 3).unwrap(); @@ -301,7 +301,7 @@ mod tests { #[test] fn test_mask_struct_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let field1 = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); let field2 = PrimitiveArray::from_iter([4i32, 5, 6]).into_array(); let fields = vec![field1, field2]; @@ -326,7 +326,7 @@ mod tests { #[test] fn test_mask_all_false() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::AllFalse(5); @@ -338,7 +338,7 @@ mod tests { #[test] fn test_mask_all_true() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]); let mask = Mask::AllTrue(5); @@ -351,7 +351,7 @@ mod tests { #[test] fn test_mask_empty_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_iter(Vec::::new()); for mask in [Mask::AllFalse(0), Mask::AllTrue(0)] { let result = diff --git a/fuzz/src/fsst_like.rs b/fuzz/src/fsst_like.rs index 9a8c9ae21ac..244c09ad06f 100644 --- a/fuzz/src/fsst_like.rs +++ b/fuzz/src/fsst_like.rs @@ -31,7 +31,8 @@ use crate::error::Backtrace; use crate::error::VortexFuzzError; use crate::error::VortexFuzzResult; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); /// A random string from a small alphabet (`a..=h`) with bounded length. #[derive(Debug)] diff --git a/fuzz/src/lib.rs b/fuzz/src/lib.rs index 8869a343f94..8933bb9dc9c 100644 --- a/fuzz/src/lib.rs +++ b/fuzz/src/lib.rs @@ -43,19 +43,23 @@ mod native_runtime { use vortex::VortexSessionDefault; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::current::CurrentThreadRuntime; - use vortex_io::session::RuntimeSessionExt; + use vortex_io::session::RuntimeSessionBuilderExt; use vortex_session::VortexSession; pub static RUNTIME: LazyLock = LazyLock::new(CurrentThreadRuntime::new); pub static SESSION: LazyLock = LazyLock::new(|| { #[allow(unused_mut)] - let mut session = VortexSession::default().with_handle(RUNTIME.handle()); + let mut builder = VortexSession::default_builder().with_handle(RUNTIME.handle()); #[cfg(all(feature = "cuda", target_os = "linux"))] // Even if the CUDA feature is enabled we need to check at // runtime whether CUDA is available in the current environment. + if vortex_cuda::cuda_available() { + builder = builder.with::(); + } + let session = builder.build(); + #[cfg(all(feature = "cuda", target_os = "linux"))] if vortex_cuda::cuda_available() { use vortex_cuda::CudaSessionExt; - session = session.with::(); vortex_cuda::initialize_cuda(session.cuda_session()); } session @@ -75,11 +79,14 @@ mod wasm_runtime { use vortex::VortexSessionDefault; use vortex_io::runtime::wasm::WasmRuntime; - use vortex_io::session::RuntimeSessionExt; + use vortex_io::session::RuntimeSessionBuilderExt; use vortex_session::VortexSession; - pub static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_handle(WasmRuntime::handle())); + pub static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::default_builder() + .with_handle(WasmRuntime::handle()) + .build() + }); } #[cfg(all(target_arch = "wasm32", target_os = "unknown"))] diff --git a/vortex-array/benches/aggregate_grouped.rs b/vortex-array/benches/aggregate_grouped.rs index fdc43dbcabd..b5b8ba76b8e 100644 --- a/vortex-array/benches/aggregate_grouped.rs +++ b/vortex-array/benches/aggregate_grouped.rs @@ -31,7 +31,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const GROUP_COUNT: usize = 128; const GROUP_SIZE_SEED: u64 = 42; diff --git a/vortex-array/benches/aggregate_max.rs b/vortex-array/benches/aggregate_max.rs index 4da7d1fc653..a7009b24333 100644 --- a/vortex-array/benches/aggregate_max.rs +++ b/vortex-array/benches/aggregate_max.rs @@ -16,7 +16,8 @@ fn main() { const N: usize = 100_000; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[divan::bench] fn max_i32(bencher: Bencher) { diff --git a/vortex-array/benches/aggregate_sum.rs b/vortex-array/benches/aggregate_sum.rs index 81402b24750..498449bb174 100644 --- a/vortex-array/benches/aggregate_sum.rs +++ b/vortex-array/benches/aggregate_sum.rs @@ -17,7 +17,8 @@ fn main() { const N: usize = 100_000; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[divan::bench] fn sum_i32(bencher: Bencher) { diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 98c69413193..85f397cc763 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -26,7 +26,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const LEN: usize = 65_536; diff --git a/vortex-array/benches/bool_zip.rs b/vortex-array/benches/bool_zip.rs index 40496fd1f0b..70435b43e63 100644 --- a/vortex-array/benches/bool_zip.rs +++ b/vortex-array/benches/bool_zip.rs @@ -19,7 +19,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const LEN: usize = 65_536; diff --git a/vortex-array/benches/cast_primitive.rs b/vortex-array/benches/cast_primitive.rs index 3fd89270517..ac6fdde37c7 100644 --- a/vortex-array/benches/cast_primitive.rs +++ b/vortex-array/benches/cast_primitive.rs @@ -26,7 +26,8 @@ fn main() { // the kernel cost shows up clearly rather than being hidden by DRAM bandwidth. const SIZES: &[usize] = &[65_536]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[divan::bench(args = SIZES)] fn cast_u16_to_u32(bencher: Bencher, n: usize) { diff --git a/vortex-array/benches/chunk_array_builder.rs b/vortex-array/benches/chunk_array_builder.rs index 6582071a6e1..6e8ab4929a2 100644 --- a/vortex-array/benches/chunk_array_builder.rs +++ b/vortex-array/benches/chunk_array_builder.rs @@ -32,7 +32,8 @@ const BENCH_ARGS: &[(usize, usize)] = &[ (1000, 10), ]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[divan::bench(args = BENCH_ARGS)] fn chunked_bool_canonical_into(bencher: Bencher, (len, chunk_count): (usize, usize)) { diff --git a/vortex-array/benches/chunked_dict_builder.rs b/vortex-array/benches/chunked_dict_builder.rs index 74722235e54..f637ee1ec53 100644 --- a/vortex-array/benches/chunked_dict_builder.rs +++ b/vortex-array/benches/chunked_dict_builder.rs @@ -18,7 +18,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const BENCH_ARGS: &[(usize, usize, usize)] = &[ (1000, 10, 10), diff --git a/vortex-array/benches/chunked_fsl_canonicalize.rs b/vortex-array/benches/chunked_fsl_canonicalize.rs index 4277b33382e..069526534f9 100644 --- a/vortex-array/benches/chunked_fsl_canonicalize.rs +++ b/vortex-array/benches/chunked_fsl_canonicalize.rs @@ -26,7 +26,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); /// Number of lists in each chunk. const LISTS_PER_CHUNK: usize = 1_000; diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 95702c9f8f5..8c6b49aa44b 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -29,7 +29,7 @@ fn compare_bool(bencher: Bencher) { let arr1 = BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.sample(range) == 0)).into_array(); let arr2 = BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.sample(range) == 0)).into_array(); - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); bencher .with_inputs(|| (&arr1, &arr2, session.create_execution_ctx())) @@ -57,7 +57,7 @@ fn compare_int(bencher: Bencher) { .map(|_| rng.sample(range)) .collect::>() .into_array(); - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); bencher .with_inputs(|| (&arr1, &arr2, session.create_execution_ctx())) diff --git a/vortex-array/benches/dict_compare.rs b/vortex-array/benches/dict_compare.rs index 29e1516e2b6..e4fbd1cd7e8 100644 --- a/vortex-array/benches/dict_compare.rs +++ b/vortex-array/benches/dict_compare.rs @@ -28,7 +28,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const LENGTH_AND_UNIQUE_VALUES: &[(usize, usize)] = &[ // length, unique_values diff --git a/vortex-array/benches/dict_compress.rs b/vortex-array/benches/dict_compress.rs index 2901d82c6b5..28cbfaff640 100644 --- a/vortex-array/benches/dict_compress.rs +++ b/vortex-array/benches/dict_compress.rs @@ -37,7 +37,8 @@ const BENCH_ARGS: &[(usize, usize)] = &[ (10_000, 512), ]; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[divan::bench(types = [u8, f32, i64], args = BENCH_ARGS)] fn encode_primitives(bencher: Bencher, (len, unique_values): (usize, usize)) diff --git a/vortex-array/benches/dict_mask.rs b/vortex-array/benches/dict_mask.rs index 04cf18cfd03..9eda341ca59 100644 --- a/vortex-array/benches/dict_mask.rs +++ b/vortex-array/benches/dict_mask.rs @@ -59,7 +59,7 @@ fn bench_dict_mask(bencher: Bencher, (fraction_valid, fraction_masked): (f64, f6 let values = PrimitiveArray::from_option_iter([None, Some(42i32)]).into_array(); let array = DictArray::try_new(codes, values).unwrap().into_array(); let filter_mask = filter_mask(len, fraction_masked, &mut rng); - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); bencher .with_inputs(|| (&array, &filter_mask, session.create_execution_ctx())) .bench_refs(|(array, filter_mask, ctx)| { diff --git a/vortex-array/benches/expr/case_when_bench.rs b/vortex-array/benches/expr/case_when_bench.rs index 1bada8fa9e1..4979936c8c7 100644 --- a/vortex-array/benches/expr/case_when_bench.rs +++ b/vortex-array/benches/expr/case_when_bench.rs @@ -25,7 +25,8 @@ use vortex_array::expr::root; use vortex_buffer::Buffer; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn main() { divan::main(); diff --git a/vortex-array/benches/filter_bool.rs b/vortex-array/benches/filter_bool.rs index cd31805bb77..ad4c36f1ee1 100644 --- a/vortex-array/benches/filter_bool.rs +++ b/vortex-array/benches/filter_bool.rs @@ -27,7 +27,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const SIZES: &[usize] = &[1_000, 10_000, 100_000, 250_000]; const DENSITY_SWEEP_SIZE: usize = 100_000; diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index d92c033c2a4..3e48dc89508 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -24,9 +24,9 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::InterleaveArray; +use vortex_array::default_session_builder; use vortex_buffer::Buffer; fn main() { @@ -133,7 +133,7 @@ fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { #[divan::bench(args = combos())] fn vortex(bencher: Bencher, combo: Combo) { let (values, array_indices, row_indices) = vortex_inputs(combo); - let session = array_session(); + let session = default_session_builder().build(); bencher .with_inputs(|| { ( diff --git a/vortex-array/benches/kleene_bool.rs b/vortex-array/benches/kleene_bool.rs index b23a0580a5a..9482ad711d1 100644 --- a/vortex-array/benches/kleene_bool.rs +++ b/vortex-array/benches/kleene_bool.rs @@ -28,10 +28,11 @@ fn main() { } static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::empty() + VortexSession::builder() .with::() .with::() .with::() + .build() }); const LEN: usize = 65_536; diff --git a/vortex-array/benches/listview_rebuild.rs b/vortex-array/benches/listview_rebuild.rs index 81157d52b5a..a1425a01ce4 100644 --- a/vortex-array/benches/listview_rebuild.rs +++ b/vortex-array/benches/listview_rebuild.rs @@ -30,7 +30,8 @@ fn main() { } /// A shared session for the `ListView` rebuild benchmarks, used to create execution contexts. -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn make_primitive_lv(num_lists: usize, list_size: usize, step: usize) -> ListViewArray { let element_count = step * num_lists + list_size; diff --git a/vortex-array/benches/primitive_zip.rs b/vortex-array/benches/primitive_zip.rs index adcabf51d7d..b7ddf39a56e 100644 --- a/vortex-array/benches/primitive_zip.rs +++ b/vortex-array/benches/primitive_zip.rs @@ -23,7 +23,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); // Sized so the bench stays well under a few hundred microseconds under CodSpeed's instruction-count // simulation, which runs ~10x the local walltime; the branchless value blend is still exercised. diff --git a/vortex-array/benches/scalar_at_struct.rs b/vortex-array/benches/scalar_at_struct.rs index 3921f0a2366..11eb3e83e14 100644 --- a/vortex-array/benches/scalar_at_struct.rs +++ b/vortex-array/benches/scalar_at_struct.rs @@ -26,7 +26,8 @@ fn main() { const ARRAY_SIZE: usize = 100_000; const NUM_ACCESSES: usize = 1000; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[divan::bench] fn execute_scalar_struct_simple(bencher: Bencher) { diff --git a/vortex-array/benches/scalar_subtract.rs b/vortex-array/benches/scalar_subtract.rs index ee9aba1055d..af3874aef9a 100644 --- a/vortex-array/benches/scalar_subtract.rs +++ b/vortex-array/benches/scalar_subtract.rs @@ -24,7 +24,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[divan::bench] fn scalar_subtract(bencher: Bencher) { diff --git a/vortex-array/benches/take_filter.rs b/vortex-array/benches/take_filter.rs index b6ec1b4040c..6db41baec53 100644 --- a/vortex-array/benches/take_filter.rs +++ b/vortex-array/benches/take_filter.rs @@ -57,7 +57,8 @@ const INDEX_SEED: u64 = 43; const LIST_SIZE: usize = 4; const NULL_INDEX_INTERVAL: usize = 8; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn primitive_array() -> ArrayRef { PrimitiveArray::from_iter(0..ARRAY_LEN as u32).into_array() diff --git a/vortex-array/benches/take_fsl.rs b/vortex-array/benches/take_fsl.rs index 438faae4983..a2b45cf5f8a 100644 --- a/vortex-array/benches/take_fsl.rs +++ b/vortex-array/benches/take_fsl.rs @@ -28,7 +28,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); /// Number of lists in the source array. const NUM_LISTS: usize = 500; diff --git a/vortex-array/benches/take_patches.rs b/vortex-array/benches/take_patches.rs index 5dc0c04c325..9d6eab98af0 100644 --- a/vortex-array/benches/take_patches.rs +++ b/vortex-array/benches/take_patches.rs @@ -23,7 +23,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const BENCH_ARGS: &[(f64, f64)] = &[ // patches_sparsity, index_multiple diff --git a/vortex-array/benches/take_primitive.rs b/vortex-array/benches/take_primitive.rs index 383b525dcc6..806e80ade7f 100644 --- a/vortex-array/benches/take_primitive.rs +++ b/vortex-array/benches/take_primitive.rs @@ -31,7 +31,8 @@ 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(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); // --- DictArray canonicalization benchmarks --- diff --git a/vortex-array/benches/take_struct.rs b/vortex-array/benches/take_struct.rs index d49ce9e88a6..7df47016fee 100644 --- a/vortex-array/benches/take_struct.rs +++ b/vortex-array/benches/take_struct.rs @@ -23,7 +23,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const ARRAY_SIZE: usize = 100_000; const TAKE_SIZE: usize = 1000; diff --git a/vortex-array/benches/to_arrow.rs b/vortex-array/benches/to_arrow.rs index 48e5e917476..1ba5670128c 100644 --- a/vortex-array/benches/to_arrow.rs +++ b/vortex-array/benches/to_arrow.rs @@ -31,7 +31,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn schema() -> DType { let fields = StructFields::from_iter([ diff --git a/vortex-array/benches/varbinview_zip.rs b/vortex-array/benches/varbinview_zip.rs index e010291e9dc..c44158506e8 100644 --- a/vortex-array/benches/varbinview_zip.rs +++ b/vortex-array/benches/varbinview_zip.rs @@ -20,7 +20,8 @@ fn main() { divan::main(); } -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); /// Benchmarks zip on VarBinView arrays with a highly fragmented mask (worst case for per-slice lookup paths). #[divan::bench] diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 8273779fe2b..5cff6b70de8 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -335,7 +335,7 @@ mod tests { } fn fresh_session() -> VortexSession { - crate::array_session() + crate::default_session_builder().build() } fn dict_of_seven() -> ArrayRef { diff --git a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs index 44bed7406f3..b7bf7711820 100644 --- a/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_nan/mod.rs @@ -138,15 +138,15 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::all_nan::AllNan; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; #[test] fn all_nan_aggregate_fn() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?; @@ -159,7 +159,7 @@ mod tests { #[test] fn all_nan_false_with_non_nan() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?; @@ -179,7 +179,7 @@ mod tests { #[test] fn all_nan_false_with_null() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNan, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs index ab513670c18..e2869ab9488 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_distinct/tests.rs @@ -10,7 +10,6 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::DecimalArray; @@ -20,6 +19,7 @@ use crate::arrays::NullArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; @@ -43,7 +43,7 @@ fn scalar_baseline(a: &ArrayRef, b: &ArrayRef, ctx: &mut ExecutionCtx) -> Vortex /// Assert that `all_non_distinct` agrees with the scalar baseline. fn assert_matches_baseline(a: &ArrayRef, b: &ArrayRef) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let expected = scalar_baseline(a, b, &mut ctx)?; let actual = all_non_distinct(a, b, &mut ctx)?; assert_eq!( @@ -88,7 +88,7 @@ fn bool_non_nullable( ) -> VortexResult<()> { let a = BoolArray::from_iter(a.iter().copied()).into_array(); let b = BoolArray::from_iter(b.iter().copied()).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -121,7 +121,7 @@ fn bool_nullable( ) -> VortexResult<()> { let a = BoolArray::from_iter(a.iter().copied()).into_array(); let b = BoolArray::from_iter(b.iter().copied()).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -140,7 +140,7 @@ fn primitive_i32( ) -> VortexResult<()> { let a = PrimitiveArray::from_iter(a).into_array(); let b = PrimitiveArray::from_iter(b).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -206,7 +206,7 @@ fn primitive_nullable( ) -> VortexResult<()> { let a = PrimitiveArray::from_option_iter(a.iter().copied()).into_array(); let b = PrimitiveArray::from_option_iter(b.iter().copied()).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -235,7 +235,7 @@ fn strings_non_nullable( ) -> VortexResult<()> { let a = VarBinViewArray::from_iter_str(a.iter().copied()).into_array(); let b = VarBinViewArray::from_iter_str(b.iter().copied()).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } @@ -258,7 +258,7 @@ fn strings_nullable( ) -> VortexResult<()> { let a = VarBinViewArray::from_iter_nullable_str(a.iter().copied()).into_array(); let b = VarBinViewArray::from_iter_nullable_str(b.iter().copied()).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(all_non_distinct(&a, &b, &mut ctx)?, expected); assert_matches_baseline(&a, &b) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs index 4ce3aed9da9..f78b3bf88a1 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_nan/mod.rs @@ -128,15 +128,15 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; #[test] fn all_non_nan_aggregate_fn() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?; @@ -149,7 +149,7 @@ mod tests { #[test] fn all_non_nan_false_with_nan() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?; @@ -178,7 +178,7 @@ mod tests { #[test] fn all_non_nan_true_with_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNan, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs index 9ce37669f02..8c0aef702c3 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_non_null/mod.rs @@ -118,15 +118,15 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::all_non_null::AllNonNull; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; #[test] fn all_non_null_aggregate_fn() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?; @@ -139,7 +139,7 @@ mod tests { #[test] fn all_non_null_false_with_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?; @@ -152,7 +152,7 @@ mod tests { #[test] fn all_non_null_true_for_empty_input() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNonNull, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs index 54bf37d502d..53cd5ed1001 100644 --- a/vortex-array/src/aggregate_fn/fns/all_null/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/all_null/mod.rs @@ -121,15 +121,15 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::all_null::AllNull; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; #[test] fn all_null_aggregate_fn() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?; @@ -142,7 +142,7 @@ mod tests { #[test] fn all_null_false_with_non_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?; @@ -155,7 +155,7 @@ mod tests { #[test] fn all_null_true_for_empty_input() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(AllNull, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index 8da704d7c9f..b6cea9885f4 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -333,9 +333,9 @@ mod tests { use crate::aggregate_fn::fns::bounded_max::make_bounded_max_partial_dtype; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::scalar::Scalar; use crate::validity::Validity; @@ -345,12 +345,12 @@ mod tests { } fn fresh_session() -> VortexSession { - array_session() + default_session_builder().build() } #[test] fn bounded_max_truncates_utf8_to_upper_bound() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["aardvark", "char🪩"]).into_array(); let mut acc = Accumulator::try_new( BoundedMax, @@ -368,7 +368,7 @@ mod tests { #[test] fn bounded_max_unknown_upper_bound_returns_null() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = VarBinViewArray::from_iter_bin([&[255u8, 255, 255][..]]).into_array(); let mut acc = Accumulator::try_new( BoundedMax, @@ -386,7 +386,7 @@ mod tests { #[test] fn bounded_max_empty_does_not_poison_later_values() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let empty = VarBinViewArray::from_iter_bin(Vec::<&[u8]>::new()).into_array(); let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array(); let mut acc = Accumulator::try_new( @@ -409,7 +409,7 @@ mod tests { #[test] fn bounded_max_unknown_poisons_later_values() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let unknown = VarBinViewArray::from_iter_bin([&[255u8, 255, 255][..]]).into_array(); let values = VarBinViewArray::from_iter_bin([&[1u8][..]]).into_array(); let mut acc = Accumulator::try_new( @@ -479,7 +479,7 @@ mod tests { #[test] fn bounded_max_keeps_fixed_width_values_exact() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array(); let mut acc = Accumulator::try_new( BoundedMax, diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index dbe1e0594bd..0c3c2bf102e 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -249,9 +249,9 @@ mod tests { use crate::aggregate_fn::fns::bounded_min::BoundedMinOptions; use crate::aggregate_fn::fns::max::Max; use crate::aggregate_fn::fns::min::Min; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::scalar::Scalar; use crate::validity::Validity; @@ -261,12 +261,12 @@ mod tests { } fn fresh_session() -> VortexSession { - array_session() + default_session_builder().build() } #[test] fn bounded_min_truncates_utf8_to_lower_bound() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["snowman⛄️snowman", "untruncated"]).into_array(); let mut acc = Accumulator::try_new( @@ -288,7 +288,7 @@ mod tests { #[test] fn bounded_min_keeps_fixed_width_values_exact() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new(buffer![10i32, 20, 5], Validity::NonNullable).into_array(); let mut acc = Accumulator::try_new( BoundedMin, diff --git a/vortex-array/src/aggregate_fn/fns/count/grouped.rs b/vortex-array/src/aggregate_fn/fns/count/grouped.rs index 39a957530bf..156c2cb4e16 100644 --- a/vortex-array/src/aggregate_fn/fns/count/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/count/grouped.rs @@ -98,12 +98,12 @@ mod tests { use crate::aggregate_fn::GroupedAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::count::Count; - use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::Nullability::Nullable; @@ -130,7 +130,7 @@ mod tests { elements: &ArrayRef, ranges: &[(usize, usize)], ) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let counts: Buffer = ranges .iter() .map(|&(offset, size)| { @@ -156,7 +156,7 @@ mod tests { #[test] fn listview_counts_all_valid() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I32, NonNullable); @@ -175,7 +175,7 @@ mod tests { #[test] fn listview_counts_with_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) .into_array(); @@ -195,7 +195,7 @@ mod tests { #[test] fn listview_counts_varbinview_with_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = VarBinViewArray::from_iter_nullable_str([ Some("a"), None, @@ -219,7 +219,7 @@ mod tests { #[test] fn fixed_size_counts_float_nans() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1.0f64), Some(f64::NAN), None, Some(2.0)]) .into_array(); @@ -243,7 +243,7 @@ mod tests { #[test] fn fixed_size_counts_with_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(); let elem_dtype = DType::Primitive(PType::I32, Nullable); diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs index 21b382adec0..7a84ab7763a 100644 --- a/vortex-array/src/aggregate_fn/fns/count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/count/mod.rs @@ -152,7 +152,8 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); pub fn count(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { let mut acc = Accumulator::try_new( diff --git a/vortex-array/src/aggregate_fn/fns/first/mod.rs b/vortex-array/src/aggregate_fn/fns/first/mod.rs index 8d04b62cdc6..b8cd058493d 100644 --- a/vortex-array/src/aggregate_fn/fns/first/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/first/mod.rs @@ -137,11 +137,11 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::first::First; use crate::aggregate_fn::fns::first::first; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::Nullability::Nullable; @@ -152,7 +152,7 @@ mod tests { #[test] fn first_non_null() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(10i32, Nullable)); Ok(()) } @@ -161,7 +161,7 @@ mod tests { fn first_skips_leading_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([None, None, Some(7i32), Some(8)]).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(7i32, Nullable)); Ok(()) } @@ -169,7 +169,7 @@ mod tests { #[test] fn first_all_null() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullable); assert_eq!(first(&array, &mut ctx)?, Scalar::null(dtype)); Ok(()) @@ -187,7 +187,7 @@ mod tests { #[test] fn first_constant() -> VortexResult<()> { let array = ConstantArray::new(42i32, 10).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(42i32, Nullable)); Ok(()) } @@ -196,7 +196,7 @@ mod tests { fn first_constant_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullable); let array = ConstantArray::new(Scalar::null(dtype.clone()), 10).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::null(dtype)); Ok(()) } @@ -208,14 +208,14 @@ mod tests { DType::Utf8(Nullable), ) .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(first(&array, &mut ctx)?, Scalar::utf8("hello", Nullable)); Ok(()) } #[test] fn first_multi_batch_picks_earliest_non_null() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullable); let mut acc = Accumulator::try_new(First, EmptyOptions, dtype)?; @@ -240,7 +240,7 @@ mod tests { #[test] fn first_finish_resets_state() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(First, EmptyOptions, dtype)?; @@ -278,7 +278,7 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter([None, Some(42i32), Some(100)]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!( first(&chunked.into_array(), &mut ctx)?, Scalar::primitive(42i32, Nullable) diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs index ae62f71ca45..6e68bdba7cc 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/mod.rs @@ -441,13 +441,13 @@ mod tests { use crate::IntoArray as _; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::DecimalArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; @@ -459,7 +459,7 @@ mod tests { // Tests migrated from compute/is_constant.rs #[test] fn is_constant_min_max_no_nan() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = buffer![0, 1].into_array(); arr.statistics() @@ -478,7 +478,7 @@ mod tests { #[test] fn is_constant_min_max_with_nan() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::from_iter([0.0, 0.0, f32::NAN]).into_array(); arr.statistics() @@ -505,7 +505,7 @@ mod tests { }, false)] fn test_bool_is_constant(#[case] input: Vec, #[case] expected: bool) -> VortexResult<()> { let array = BoolArray::from_iter(input); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(is_constant(&array.into_array(), &mut ctx)?, expected); Ok(()) } @@ -525,7 +525,7 @@ mod tests { )? .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!(is_constant(&chunked, &mut ctx)?); Ok(()) } @@ -533,7 +533,7 @@ mod tests { // Tests migrated from arrays/decimal/compute/is_constant.rs #[test] fn test_decimal_is_constant() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = DecimalArray::new( buffer![0i128, 1i128, 2i128], @@ -554,7 +554,7 @@ mod tests { // Tests migrated from arrays/list/compute/is_constant.rs #[test] fn test_is_constant_nested_list() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let xs = ListArray::try_new( buffer![0i32, 1, 0, 1].into_array(), @@ -612,7 +612,7 @@ mod tests { Validity::NonNullable, )?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(is_constant(&list_array.into_array(), &mut ctx)?, expected); Ok(()) } @@ -630,7 +630,7 @@ mod tests { Validity::NonNullable, )?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Both outer lists contain [[1], [2]], so should be constant assert!(is_constant(&outer_list.into_array(), &mut ctx)?); Ok(()) @@ -675,7 +675,7 @@ mod tests { Validity::NonNullable, )?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(is_constant(&list_array.into_array(), &mut ctx)?, expected); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs index 57ee2d342a0..800aaf8e6f5 100644 --- a/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/is_sorted/mod.rs @@ -561,15 +561,15 @@ mod tests { use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_sorted::is_sorted; use crate::aggregate_fn::fns::is_sorted::is_strict_sorted; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::validity::Validity; // Tests migrated from compute/is_sorted.rs #[test] fn test_is_sorted() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::new(buffer!(0, 1, 2, 3), Validity::AllValid).into_array(); assert!(is_sorted(&arr, &mut ctx)?); @@ -603,7 +603,7 @@ mod tests { #[test] fn test_is_strict_sorted() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::new(buffer!(0, 1, 2, 3), Validity::AllValid).into_array(); assert!(is_strict_sorted(&arr, &mut ctx)?); @@ -640,7 +640,7 @@ mod tests { #[case(PrimitiveArray::from_option_iter([None, None, Some(1i32), Some(1)]), true)] #[case(PrimitiveArray::from_option_iter([None, Some(5_u8), None]), false)] fn test_primitive_is_sorted(#[case] array: PrimitiveArray, #[case] expected: bool) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!( is_sorted(&array.into_array(), &mut ctx) .vortex_expect("operation should succeed in test"), @@ -655,7 +655,7 @@ mod tests { #[case(PrimitiveArray::from_option_iter([None, None, Some(1i32), Some(1), None]), false)] #[case(PrimitiveArray::from_option_iter([None, Some(5_u8), None]), false)] fn test_primitive_is_strict_sorted(#[case] array: PrimitiveArray, #[case] expected: bool) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!( is_strict_sorted(&array.into_array(), &mut ctx) .vortex_expect("operation should succeed in test"), @@ -672,7 +672,7 @@ mod tests { use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DecimalDType::new(19, 2); let i100 = parse_decimal::("100.00", dtype.precision(), dtype.scale())?; let i200 = parse_decimal::("200.00", dtype.precision(), dtype.scale())?; @@ -697,7 +697,7 @@ mod tests { use crate::arrays::DecimalArray; use crate::dtype::DecimalDType; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DecimalDType::new(19, 2); let i100 = parse_decimal::("100.00", dtype.precision(), dtype.scale())?; let i200 = parse_decimal::("200.00", dtype.precision(), dtype.scale())?; diff --git a/vortex-array/src/aggregate_fn/fns/last/mod.rs b/vortex-array/src/aggregate_fn/fns/last/mod.rs index 61c1277ad13..43fbc560027 100644 --- a/vortex-array/src/aggregate_fn/fns/last/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/last/mod.rs @@ -135,11 +135,11 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::last::Last; use crate::aggregate_fn::fns::last::last; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::Nullability::Nullable; @@ -150,7 +150,7 @@ mod tests { #[test] fn last_non_null() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(30i32, Nullable)); Ok(()) } @@ -159,7 +159,7 @@ mod tests { fn last_skips_trailing_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([Some(7i32), Some(8), None, None]).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(8i32, Nullable)); Ok(()) } @@ -167,7 +167,7 @@ mod tests { #[test] fn last_all_null() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullable); assert_eq!(last(&array, &mut ctx)?, Scalar::null(dtype)); Ok(()) @@ -185,7 +185,7 @@ mod tests { #[test] fn last_constant() -> VortexResult<()> { let array = ConstantArray::new(42i32, 10).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(42i32, Nullable)); Ok(()) } @@ -194,7 +194,7 @@ mod tests { fn last_constant_null() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullable); let array = ConstantArray::new(Scalar::null(dtype.clone()), 10).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::null(dtype)); Ok(()) } @@ -206,14 +206,14 @@ mod tests { DType::Utf8(Nullable), ) .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(last(&array, &mut ctx)?, Scalar::utf8("world", Nullable)); Ok(()) } #[test] fn last_multi_batch_picks_latest_non_null() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullable); let mut acc = Accumulator::try_new(Last, EmptyOptions, dtype)?; @@ -237,7 +237,7 @@ mod tests { #[test] fn last_finish_resets_state() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Last, EmptyOptions, dtype)?; @@ -275,7 +275,7 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter::([None, None]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!( last(&chunked.into_array(), &mut ctx)?, Scalar::primitive(100i32, Nullable) diff --git a/vortex-array/src/aggregate_fn/fns/max/mod.rs b/vortex-array/src/aggregate_fn/fns/max/mod.rs index 79194f56941..ecb2363bb79 100644 --- a/vortex-array/src/aggregate_fn/fns/max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/max/mod.rs @@ -223,8 +223,8 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::max::Max; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -236,7 +236,7 @@ mod tests { #[test] fn max_aggregate_fn() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Max, NumericalAggregateOpts::default(), dtype)?; @@ -267,7 +267,7 @@ mod tests { #[test] fn max_with_nan_not_skipping() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(Max, NumericalAggregateOpts::include_nans(), dtype)?; @@ -288,7 +288,7 @@ mod tests { #[test] fn max_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // The array has no NaNs; a planted exact NaNCount stat proves the poisoning came from // the stat rather than a scan. let batch = PrimitiveArray::new(buffer![1.0f64, 2.0], Validity::NonNullable).into_array(); @@ -315,7 +315,7 @@ mod tests { fn max_nan_including_nullable_cached_stat() -> VortexResult<()> { // A nullable float array's cached Max stat is reconstructed as a nullable scalar. The // NaN-including shortcircuit merges it as-is; `to_scalar` re-casts to the result dtype. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1.0f64), Some(2.0), Some(3.0)]).into_array(); array @@ -339,7 +339,7 @@ mod tests { #[test] fn max_casts_nonnullable_legacy_stat_to_nullable_partial() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let batch = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); batch .statistics() diff --git a/vortex-array/src/aggregate_fn/fns/mean/mod.rs b/vortex-array/src/aggregate_fn/fns/mean/mod.rs index 061ed92c4b1..4635b5e4c0b 100644 --- a/vortex-array/src/aggregate_fn/fns/mean/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/mean/mod.rs @@ -169,18 +169,18 @@ mod tests { use super::*; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::validity::Validity; #[test] fn mean_all_valid() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0, 4.0, 5.0], Validity::NonNullable) .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mean(&array, &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(3.0)); Ok(()) @@ -189,7 +189,7 @@ mod tests { #[test] fn mean_with_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([Some(2.0f64), None, Some(4.0)]).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mean(&array, &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(3.0)); Ok(()) @@ -198,7 +198,7 @@ mod tests { #[test] fn mean_integers() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mean(&array, &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(20.0)); Ok(()) @@ -207,7 +207,7 @@ mod tests { #[test] fn mean_bool() -> VortexResult<()> { let array: BoolArray = [true, false, true, true].into_iter().collect(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mean(&array.into_array(), &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(0.75)); Ok(()) @@ -216,7 +216,7 @@ mod tests { #[test] fn mean_constant_non_null() -> VortexResult<()> { let array = ConstantArray::new(5.0f64, 4); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mean(&array.into_array(), &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(5.0)); Ok(()) @@ -228,7 +228,7 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter([Some(5.0f64), None]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mean(&chunked.into_array(), &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(3.0)); Ok(()) @@ -239,7 +239,7 @@ mod tests { // NaNs are excluded from both the sum and the count. let array = PrimitiveArray::new(buffer![1.0f64, f64::NAN, 3.0], Validity::NonNullable).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mean(&array, &mut ctx)?; assert_eq!(result.as_primitive().as_::(), Some(2.0)); Ok(()) @@ -249,7 +249,7 @@ mod tests { fn mean_with_nan_not_skipping() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![1.0f64, f64::NAN, 3.0], Validity::NonNullable).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let keep_nans = NumericalAggregateOpts::include_nans(); let mut acc = Accumulator::try_new( Mean::combined(), @@ -265,7 +265,7 @@ mod tests { #[test] fn mean_all_null_returns_nan() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mean(&array, &mut ctx)?; assert!(result.as_primitive().as_::().is_some_and(f64::is_nan)); Ok(()) @@ -273,7 +273,7 @@ mod tests { #[test] fn mean_multi_batch() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new( Mean::combined(), diff --git a/vortex-array/src/aggregate_fn/fns/min/mod.rs b/vortex-array/src/aggregate_fn/fns/min/mod.rs index 8a0fbd889ed..5f12d602cf6 100644 --- a/vortex-array/src/aggregate_fn/fns/min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min/mod.rs @@ -223,8 +223,8 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min::Min; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -236,7 +236,7 @@ mod tests { #[test] fn min_aggregate_fn() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Min, NumericalAggregateOpts::default(), dtype)?; @@ -267,7 +267,7 @@ mod tests { #[test] fn min_with_nan_not_skipping() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(Min, NumericalAggregateOpts::include_nans(), dtype)?; @@ -288,7 +288,7 @@ mod tests { #[test] fn min_not_skipping_shortcircuits_on_exact_nan_count_stat() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // The array has no NaNs; a planted exact NaNCount stat proves the poisoning came from // the stat rather than a scan. let batch = PrimitiveArray::new(buffer![1.0f64, 2.0], Validity::NonNullable).into_array(); @@ -315,7 +315,7 @@ mod tests { fn min_nan_including_nullable_cached_stat() -> VortexResult<()> { // A nullable float array's cached Min stat is reconstructed as a nullable scalar. The // NaN-including shortcircuit merges it as-is; `to_scalar` re-casts to the result dtype. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::from_option_iter([Some(1.0f64), Some(2.0), Some(3.0)]).into_array(); array @@ -339,7 +339,7 @@ mod tests { #[test] fn min_casts_nonnullable_legacy_stat_to_nullable_partial() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let batch = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array(); batch .statistics() diff --git a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs index 097d3137f42..2bb9c68b0b8 100644 --- a/vortex-array/src/aggregate_fn/fns/min_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/min_max/mod.rs @@ -471,7 +471,8 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_prim_min_max() -> VortexResult<()> { diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs index e66b3b0fb33..12c000322af 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/mod.rs @@ -192,10 +192,10 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::nan_count::nan_count; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -204,7 +204,7 @@ mod tests { #[test] fn nan_count_multi_batch() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?; @@ -224,7 +224,7 @@ mod tests { #[test] fn nan_count_finish_resets_state() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::F64, Nullability::NonNullable); let mut acc = Accumulator::try_new(NanCount, EmptyOptions, dtype)?; @@ -262,7 +262,7 @@ mod tests { #[test] fn nan_count_constant_nan() -> VortexResult<()> { let array = ConstantArray::new(f64::NAN, 10); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(nan_count(&array.into_array(), &mut ctx)?, 10); Ok(()) } @@ -270,7 +270,7 @@ mod tests { #[test] fn nan_count_constant_non_nan() -> VortexResult<()> { let array = ConstantArray::new(1.0f64, 10); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(nan_count(&array.into_array(), &mut ctx)?, 0); Ok(()) } @@ -290,7 +290,7 @@ mod tests { let chunk2 = PrimitiveArray::from_option_iter([Some(f64::NAN), Some(f64::NAN), None]); let dtype = chunk1.dtype().clone(); let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(nan_count(&chunked.into_array(), &mut ctx)?, 3); Ok(()) } @@ -298,7 +298,7 @@ mod tests { #[test] fn nan_count_all_null() -> VortexResult<()> { let p = PrimitiveArray::from_option_iter::([None, None, None]); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 0); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs b/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs index c9da090c7ab..c1bf46cee95 100644 --- a/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/nan_count/primitive.rs @@ -44,8 +44,8 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::nan_count::nan_count; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::validity::Validity; #[test] @@ -62,7 +62,7 @@ mod tests { ], Validity::NonNullable, ); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 2); Ok(()) } @@ -70,7 +70,7 @@ mod tests { #[test] fn primitive_nan_count_with_nulls() -> VortexResult<()> { let p = PrimitiveArray::from_option_iter([Some(f64::NAN), None, Some(f64::NAN), Some(1.0)]); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 2); Ok(()) } @@ -78,7 +78,7 @@ mod tests { #[test] fn primitive_nan_count_all_valid_no_nans() -> VortexResult<()> { let p = PrimitiveArray::new(buffer![1.0f64, 2.0, 3.0], Validity::NonNullable); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(nan_count(&p.into_array(), &mut ctx)?, 0); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs index e06f570169d..9b30b695388 100644 --- a/vortex-array/src/aggregate_fn/fns/null_count/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/null_count/mod.rs @@ -161,8 +161,8 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::null_count::NullCount; use crate::aggregate_fn::fns::null_count::null_count; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -174,7 +174,7 @@ mod tests { fn null_count_with_nulls() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None]).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(null_count(&array, &mut ctx)?, 2); assert_eq!( @@ -186,7 +186,7 @@ mod tests { #[test] fn null_count_multi_batch() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::Nullable); let mut acc = Accumulator::try_new(NullCount, EmptyOptions, dtype)?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/bool.rs b/vortex-array/src/aggregate_fn/fns/sum/bool.rs index b3993840586..28fe677dd31 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/bool.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/bool.rs @@ -43,8 +43,8 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; - use crate::array_session; use crate::arrays::BoolArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -55,7 +55,7 @@ mod tests { let arr: BoolArray = [true, true, true].into_iter().collect(); let result = sum( &arr.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(3)); Ok(()) @@ -66,7 +66,7 @@ mod tests { let arr: BoolArray = [true, false, true, false, true].into_iter().collect(); let result = sum( &arr.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(3)); Ok(()) @@ -77,7 +77,7 @@ mod tests { let arr: BoolArray = [false, false, false].into_iter().collect(); let result = sum( &arr.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -88,7 +88,7 @@ mod tests { let arr = BoolArray::from_iter([Some(true), None, Some(true), Some(false)]); let result = sum( &arr.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(2)); Ok(()) @@ -99,7 +99,7 @@ mod tests { let arr = BoolArray::from_iter([None::, None, None]); let result = sum( &arr.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) @@ -116,7 +116,7 @@ mod tests { #[test] fn sum_bool_finish_resets_state() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Bool(Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; @@ -147,7 +147,10 @@ mod tests { #[test] fn sum_boolean_from_iter() -> VortexResult<()> { let arr = BoolArray::from_iter([true, false, false, true]).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().as_::(), Some(2)); Ok(()) } diff --git a/vortex-array/src/aggregate_fn/fns/sum/constant.rs b/vortex-array/src/aggregate_fn/fns/sum/constant.rs index 0f366620e5c..ac6f707b5c7 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/constant.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/constant.rs @@ -94,8 +94,8 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::sum::sum; - use crate::array_session; use crate::arrays::ConstantArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; @@ -109,7 +109,10 @@ mod tests { #[test] fn sum_constant_unsigned() -> VortexResult<()> { let array = ConstantArray::new(5u64, 10).into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result, 50u64.into()); Ok(()) } @@ -117,7 +120,10 @@ mod tests { #[test] fn sum_constant_signed() -> VortexResult<()> { let array = ConstantArray::new(-5i64, 10).into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result, (-50i64).into()); Ok(()) } @@ -126,7 +132,10 @@ mod tests { fn sum_constant_nullable_value() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Primitive(PType::U32, Nullable)), 10) .into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -134,7 +143,10 @@ mod tests { #[test] fn sum_constant_bool_false() -> VortexResult<()> { let array = ConstantArray::new(false, 10).into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result, 0u64.into()); Ok(()) } @@ -142,7 +154,10 @@ mod tests { #[test] fn sum_constant_bool_true() -> VortexResult<()> { let array = ConstantArray::new(true, 10).into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result, 10u64.into()); Ok(()) } @@ -150,7 +165,10 @@ mod tests { #[test] fn sum_constant_bool_null() -> VortexResult<()> { let array = ConstantArray::new(Scalar::null(DType::Bool(Nullable)), 10).into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result, Scalar::primitive(0u64, Nullable)); Ok(()) } @@ -168,7 +186,10 @@ mod tests { ) .into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!( result.as_decimal().decimal_value(), @@ -184,7 +205,10 @@ mod tests { let array = ConstantArray::new(Scalar::null(DType::Decimal(decimal_dtype, Nullable)), 10) .into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!( result, Scalar::decimal( @@ -209,7 +233,10 @@ mod tests { ) .into_array(); - let result = sum(&array, &mut array_session().create_execution_ctx())?; + let result = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!( result.as_decimal().decimal_value(), Some(DecimalValue::I256(i256::from_i128(99_999_999_900))) diff --git a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs index 872e5769a01..2dcde853806 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/decimal.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/decimal.rs @@ -117,8 +117,8 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; - use crate::array_session; use crate::arrays::DecimalArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; @@ -139,7 +139,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -161,7 +161,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -183,7 +183,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -206,7 +206,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected_sum = near_max as i64 + 500 + 400; @@ -230,7 +230,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected_sum = (large_val as i128) * 4 + 1; @@ -253,7 +253,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -272,7 +272,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -294,7 +294,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected = Scalar::try_new( @@ -317,7 +317,7 @@ mod tests { let result = sum( &decimal.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let expected_sum = @@ -343,7 +343,7 @@ mod tests { assert_eq!( sum( &decimal.into_array(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .vortex_expect("operation should succeed in test"), Scalar::null(DType::Decimal(decimal_dtype, Nullable)) @@ -460,7 +460,7 @@ mod tests { // Drive accumulate through the vtable directly. let columnar = crate::Columnar::Canonical(crate::Canonical::Decimal(decimal)); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); Sum.accumulate(&mut state, &columnar, &mut ctx)?; let result = Sum.to_scalar(&state)?; diff --git a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs index efe0825d4d6..441c2590861 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/grouped.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/grouped.rs @@ -165,12 +165,12 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; - use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::builder_with_capacity; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::Nullability::Nullable; @@ -184,7 +184,10 @@ mod tests { NumericalAggregateOpts::default(), elem_dtype.clone(), )?; - acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; + acc.accumulate_list( + groups, + &mut default_session_builder().build().create_execution_ctx(), + )?; acc.finish() } @@ -198,7 +201,7 @@ mod tests { ) -> VortexResult { use crate::aggregate_fn::AggregateFnVTable; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let sum_dtype = Sum .partial_dtype(&NumericalAggregateOpts::default(), elem_dtype) .expect("sum partial dtype"); @@ -236,7 +239,7 @@ mod tests { #[test] fn listview_matches_reference_unsigned() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1u32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::U32, NonNullable); @@ -256,7 +259,7 @@ mod tests { #[test] fn listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Offsets are not in group order and a group is null: the group validity must be indexed by // group index, not by element offset. let elements = @@ -278,7 +281,7 @@ mod tests { #[test] fn listview_interior_and_full_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Group 1 has an interior null, group 2 is entirely null, group 3 is empty. let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, None, Some(9)]) @@ -300,7 +303,7 @@ mod tests { #[test] fn listview_overflow_group_is_null() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); let elem_dtype = DType::Primitive(PType::I64, NonNullable); @@ -334,7 +337,7 @@ mod tests { // Group 0: NaN skipped -> 3.0. Group 1: INF + -INF = NaN. (Avoid array equality here since // NaN != NaN; compare element scalars against the reference path instead.) - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let expected = grouped_sum_reference(&elements, &ranges, &valid, &elem_dtype)?; let g0 = actual.execute_scalar(0, &mut ctx)?; assert_eq!(g0.as_primitive().typed_value::(), Some(3.0)); @@ -370,10 +373,13 @@ mod tests { let mut acc = GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::include_nans(), elem_dtype)?; - acc.accumulate_list(&groups, &mut array_session().create_execution_ctx())?; + acc.accumulate_list( + &groups, + &mut default_session_builder().build().create_execution_ctx(), + )?; let actual = acc.finish()?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Group 0 contains a NaN -> NaN sum; group 1 sums normally. let g0 = actual.execute_scalar(0, &mut ctx)?; assert!(g0.as_primitive().typed_value::().unwrap().is_nan()); @@ -384,7 +390,7 @@ mod tests { #[test] fn fixed_size_overflow_and_nan() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // FixedSize path: first group overflows -> null sum, second sums normally. let elements = PrimitiveArray::new(buffer![i64::MAX, 1, 2, 3], Validity::NonNullable).into_array(); diff --git a/vortex-array/src/aggregate_fn/fns/sum/mod.rs b/vortex-array/src/aggregate_fn/fns/sum/mod.rs index f1922aa95fa..eb4a9cf59ee 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/mod.rs @@ -403,7 +403,6 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -412,6 +411,7 @@ mod tests { use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; @@ -428,7 +428,7 @@ mod tests { /// Sum an array with an initial value (test-only helper). fn sum_with_accumulator(array: &ArrayRef, accumulator: &Scalar) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); if accumulator.is_null() { return Ok(accumulator.clone()); } @@ -484,7 +484,7 @@ mod tests { #[test] fn sum_multi_batch() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; @@ -501,7 +501,7 @@ mod tests { #[test] fn sum_finish_resets_state() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = Accumulator::try_new(Sum, NumericalAggregateOpts::default(), dtype)?; @@ -552,7 +552,10 @@ mod tests { // compute sum with accumulator to populate stats sum_with_accumulator(&array, &Scalar::primitive(2i64, Nullable))?; - let sum_without_acc = sum(&array, &mut array_session().create_execution_ctx())?; + let sum_without_acc = sum( + &array, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(sum_without_acc, Scalar::primitive(9i64, Nullable)); Ok(()) } @@ -580,13 +583,16 @@ mod tests { NumericalAggregateOpts::default(), elem_dtype.clone(), )?; - acc.accumulate_list(groups, &mut array_session().create_execution_ctx())?; + acc.accumulate_list( + groups, + &mut default_session_builder().build().create_execution_ctx(), + )?; acc.finish() } #[test] fn grouped_sum_fixed_size_list() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6], Validity::NonNullable).into_array(); let groups = FixedSizeListArray::try_new(elements, 3, Validity::NonNullable, 2)?; @@ -601,7 +607,7 @@ mod tests { #[test] fn grouped_sum_with_null_elements() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5), Some(6)]) .into_array(); @@ -617,7 +623,7 @@ mod tests { #[test] fn grouped_sum_with_null_group() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9], Validity::NonNullable) .into_array(); @@ -635,7 +641,7 @@ mod tests { #[test] fn grouped_sum_all_null_elements_in_group() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([None::, None, Some(3), Some(4)]).into_array(); let groups = FixedSizeListArray::try_new(elements, 2, Validity::NonNullable, 2)?; @@ -650,7 +656,7 @@ mod tests { #[test] fn grouped_sum_bool() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements: BoolArray = [true, false, true, true, true, true].into_iter().collect(); let groups = FixedSizeListArray::try_new(elements.into_array(), 3, Validity::NonNullable, 2)?; @@ -665,7 +671,7 @@ mod tests { #[test] fn grouped_sum_finish_resets() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elem_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let mut acc = GroupedAccumulator::try_new(Sum, NumericalAggregateOpts::default(), elem_dtype)?; @@ -691,7 +697,7 @@ mod tests { #[test] fn grouped_sum_listview_out_of_order_offsets_with_null_group() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::new(buffer![100i32, 200, 300], Validity::NonNullable).into_array(); let offsets = PrimitiveArray::new(buffer![2i32, 0, 1], Validity::NonNullable).into_array(); @@ -729,7 +735,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result.as_primitive().as_::(), Some(20.8)); Ok(()) @@ -743,7 +749,7 @@ mod tests { let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?; let result = sum( &chunked.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) @@ -766,7 +772,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result.as_primitive().as_::(), Some(36.0)); Ok(()) @@ -781,7 +787,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert_eq!(result.as_primitive().as_::(), Some(1)); Ok(()) @@ -813,7 +819,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let decimal_result = result.as_decimal(); assert_eq!( @@ -849,7 +855,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let decimal_result = result.as_decimal(); assert_eq!( @@ -883,7 +889,7 @@ mod tests { let result = sum( &chunked.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let decimal_result = result.as_decimal(); assert_eq!( diff --git a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs index 87d8da4b143..747a694e04a 100644 --- a/vortex-array/src/aggregate_fn/fns/sum/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/sum/primitive.rs @@ -190,9 +190,9 @@ mod tests { use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::sum::sum; - use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::Nullability::Nullable; @@ -206,7 +206,10 @@ mod tests { #[test] fn sum_i32() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(10)); Ok(()) } @@ -214,7 +217,10 @@ mod tests { #[test] fn sum_u8() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![10u8, 20, 30], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(60)); Ok(()) } @@ -223,7 +229,10 @@ mod tests { fn sum_f64() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.5f64, 2.5, 3.0], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(7.0)); Ok(()) } @@ -231,7 +240,10 @@ mod tests { #[test] fn sum_with_nulls() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([Some(2i32), None, Some(4)]).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(6)); Ok(()) } @@ -251,7 +263,10 @@ mod tests { Some(6), ]) .into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(21)); Ok(()) } @@ -259,7 +274,10 @@ mod tests { #[test] fn sum_all_null() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([None::, None, None]).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(0)); Ok(()) } @@ -267,7 +285,10 @@ mod tests { #[test] fn sum_all_invalid_float() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter::([None, None, None]).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result, Scalar::primitive(0f64, Nullable)); Ok(()) } @@ -275,7 +296,10 @@ mod tests { #[test] fn sum_buffer_i32() -> VortexResult<()> { let arr = buffer![1, 1, 1, 1].into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().as_::(), Some(4)); Ok(()) } @@ -283,7 +307,10 @@ mod tests { #[test] fn sum_buffer_f64() -> VortexResult<()> { let arr = buffer![1., 1., 1., 1.].into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().as_::(), Some(4.)); Ok(()) } @@ -313,7 +340,10 @@ mod tests { Validity::NonNullable, ) .into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(6.0)); Ok(()) } @@ -322,7 +352,10 @@ mod tests { fn sum_f32_with_nan() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![1.0f32, f32::NAN, 4.0], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(5.0)); Ok(()) } @@ -331,7 +364,10 @@ mod tests { fn sum_f64_with_nan_and_nulls() -> VortexResult<()> { let arr = PrimitiveArray::from_option_iter([Some(1.0f64), None, Some(f64::NAN), Some(3.0)]) .into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(4.0)); Ok(()) } @@ -340,7 +376,10 @@ mod tests { fn sum_all_nan() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![f64::NAN, f64::NAN], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_eq!(result.as_primitive().typed_value::(), Some(0.0)); Ok(()) } @@ -351,7 +390,10 @@ mod tests { options: NumericalAggregateOpts, ) -> VortexResult { let mut acc = Accumulator::try_new(Sum, options, arr.dtype().clone())?; - acc.accumulate(arr, &mut array_session().create_execution_ctx())?; + acc.accumulate( + arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; acc.finish() } @@ -419,7 +461,10 @@ mod tests { Validity::NonNullable, ) .into_array(); - let acc = sum(&batch, &mut array_session().create_execution_ctx())?; + let acc = sum( + &batch, + &mut default_session_builder().build().create_execution_ctx(), + )?; // INFINITY + NEG_INFINITY = NaN, which is treated as saturated assert!(acc.as_primitive().typed_value::().unwrap().is_nan()); @@ -428,7 +473,10 @@ mod tests { NumericalAggregateOpts::default(), DType::Primitive(PType::F64, Nullability::NonNullable), )?; - acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; + acc.accumulate( + &batch, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert!(acc.is_saturated()); Ok(()) } @@ -436,7 +484,10 @@ mod tests { #[test] fn sum_checked_overflow() -> VortexResult<()> { let arr = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - let result = sum(&arr, &mut array_session().create_execution_ctx())?; + let result = sum( + &arr, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert!(result.is_null()); Ok(()) } @@ -449,7 +500,10 @@ mod tests { let batch = PrimitiveArray::new(buffer![i64::MAX, 1i64], Validity::NonNullable).into_array(); - acc.accumulate(&batch, &mut array_session().create_execution_ctx())?; + acc.accumulate( + &batch, + &mut default_session_builder().build().create_execution_ctx(), + )?; assert!(acc.is_saturated()); // finish resets state, clearing saturation diff --git a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs index 0b07c7cf98b..336c457e96c 100644 --- a/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/uncompressed_size_in_bytes/mod.rs @@ -323,7 +323,6 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -337,6 +336,7 @@ mod tests { use crate::arrays::VarBinViewArray; use crate::arrays::VariantArray; use crate::builders::builder_with_capacity; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::FieldNames; @@ -360,7 +360,7 @@ mod tests { } fn aggregate(array: &ArrayRef) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut acc = Accumulator::try_new(UncompressedSizeInBytes, EmptyOptions, array.dtype().clone())?; acc.accumulate(array, &mut ctx)?; @@ -545,7 +545,7 @@ mod tests { fn variant_stat_is_unsupported() -> VortexResult<()> { let child = ConstantArray::new(Scalar::variant(Scalar::from(42i32)), 3).into_array(); let array = VariantArray::try_new(child, None)?.into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!( array @@ -598,7 +598,7 @@ mod tests { #[test] fn helper_caches_result() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let size = uncompressed_size_in_bytes(&array, &mut ctx)?; diff --git a/vortex-array/src/aggregate_fn/proto.rs b/vortex-array/src/aggregate_fn/proto.rs index 3ca4df06d84..32cd0888cba 100644 --- a/vortex-array/src/aggregate_fn/proto.rs +++ b/vortex-array/src/aggregate_fn/proto.rs @@ -158,7 +158,7 @@ mod tests { #[test] fn aggregate_fn_serde() { - let session = crate::array_session(); + let session = crate::default_session_builder().build(); session.aggregate_fns().register(TestAgg); let agg_fn = TestAgg.bind(EmptyOptions); @@ -179,7 +179,7 @@ mod tests { fn numeric_aggregate_options_round_trip( #[case] options: NumericalAggregateOpts, ) -> VortexResult<()> { - let session = crate::array_session(); + let session = crate::default_session_builder().build(); let agg_fn = Sum.bind(options); let proto = agg_fn.serialize_proto()?; diff --git a/vortex-array/src/array/typed.rs b/vortex-array/src/array/typed.rs index 91225ac75d6..3b0b06deec6 100644 --- a/vortex-array/src/array/typed.rs +++ b/vortex-array/src/array/typed.rs @@ -528,15 +528,15 @@ mod tests { use super::Array; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::validity::Validity; #[test] fn typed_array_into_parts_roundtrips() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let expected = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); @@ -548,7 +548,7 @@ mod tests { #[test] fn typed_array_try_into_parts_requires_unique_owner() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable); let alias = array.clone(); diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 8b83d0f8128..c253393aa66 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -52,7 +52,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// ``` /// # fn main() -> vortex_error::VortexResult<()> { /// use vortex_array::arrays::BoolArray; -/// use vortex_array::{IntoArray, array_session, VortexSessionExecute}; +/// use vortex_array::{IntoArray, default_session_builder, VortexSessionExecute}; /// /// // Create from iterator using FromIterator impl /// let array: BoolArray = [true, false, true, false].into_iter().collect(); @@ -62,7 +62,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// assert_eq!(sliced.len(), 2); /// /// // Access individual values -/// let mut ctx = array_session().create_execution_ctx(); +/// let mut ctx = default_session_builder().build().create_execution_ctx(); /// let value = array.execute_scalar(0, &mut ctx).unwrap(); /// assert_eq!(value, true.into()); /// # Ok(()) @@ -372,17 +372,17 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::patches::Patches; use crate::validity::Validity; #[test] fn bool_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from_iter([true, false, true]); let scalar = bool::try_from(&arr.execute_scalar(0, &mut ctx).unwrap()).unwrap(); @@ -391,7 +391,7 @@ mod tests { #[test] fn test_all_some_iter() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(false)]); @@ -405,7 +405,7 @@ mod tests { #[test] fn test_bool_from_iter() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); let scalar = bool::try_from(&arr.execute_scalar(0, &mut ctx).unwrap()).unwrap(); @@ -426,7 +426,7 @@ mod tests { #[test] fn patch_sliced_bools() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.slice(4..12).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut ctx); @@ -464,7 +464,7 @@ mod tests { #[test] fn slice_array_in_middle() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let sliced = arr.slice(4..12).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut ctx); @@ -472,7 +472,7 @@ mod tests { #[test] fn patch_bools_owned() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(16)); let buf_ptr = arr.to_bit_buffer().inner().as_ptr(); @@ -495,7 +495,7 @@ mod tests { #[test] fn patch_sliced_bools_offset() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.slice(4..15).unwrap(); assert_arrays_eq!(sliced, BoolArray::from_iter([true; 11]), &mut ctx); diff --git a/vortex-array/src/arrays/bool/compute/cast.rs b/vortex-array/src/arrays/bool/compute/cast.rs index fe9332346ca..77187d4d9d8 100644 --- a/vortex-array/src/arrays/bool/compute/cast.rs +++ b/vortex-array/src/arrays/bool/compute/cast.rs @@ -68,7 +68,8 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn try_cast_bool_success() { diff --git a/vortex-array/src/arrays/bool/compute/filter.rs b/vortex-array/src/arrays/bool/compute/filter.rs index f463f27f437..2d14da25d71 100644 --- a/vortex-array/src/arrays/bool/compute/filter.rs +++ b/vortex-array/src/arrays/bool/compute/filter.rs @@ -304,14 +304,14 @@ mod tests { use super::*; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; + use crate::default_session_builder; #[test] fn filter_bool_test() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_iter([true, false, true]); @@ -321,7 +321,7 @@ mod tests { #[test] fn filter_bool_sparse_index_mask() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_indices(3, [0, 2]); @@ -331,7 +331,7 @@ mod tests { #[test] fn filter_bool_sparse_slice_mask() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_slices(3, vec![(0, 1), (2, 3)]); @@ -341,7 +341,7 @@ mod tests { #[test] fn filter_bool_sparse_buffer_mask() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from_iter([true, true, false]); let mask = Mask::from_buffer(BitBuffer::from_iter([true, false, true])); diff --git a/vortex-array/src/arrays/bool/compute/mod.rs b/vortex-array/src/arrays/bool/compute/mod.rs index 6037cac8c36..d84654b7391 100644 --- a/vortex-array/src/arrays/bool/compute/mod.rs +++ b/vortex-array/src/arrays/bool/compute/mod.rs @@ -17,9 +17,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::compute::conformance::consistency::test_array_consistency; + use crate::default_session_builder; #[rstest] // Basic bool arrays @@ -45,7 +45,7 @@ mod tests { fn test_bool_consistency(#[case] array: BoolArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/bool/compute/take.rs b/vortex-array/src/arrays/bool/compute/take.rs index abff3527bfe..c58f8d72ba0 100644 --- a/vortex-array/src/arrays/bool/compute/take.rs +++ b/vortex-array/src/arrays/bool/compute/take.rs @@ -91,17 +91,17 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::validity::Validity; #[test] fn take_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let reference = BoolArray::from_iter(vec![ Some(false), Some(true), @@ -127,7 +127,7 @@ mod test { #[test] fn test_bool_array_take_with_null_out_of_bounds_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -145,7 +145,7 @@ mod test { #[test] fn test_non_null_bool_array_take_with_null_out_of_bounds_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -162,7 +162,7 @@ mod test { #[test] fn test_bool_array_take_all_null_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = BoolArray::from_iter(vec![Some(false), Some(true), None, None, Some(false)]); let indices = PrimitiveArray::new( buffer![0, 3, 100], @@ -174,7 +174,7 @@ mod test { #[test] fn test_non_null_bool_array_take_all_null_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = BoolArray::from_iter(vec![false, true, false, true, false]); let indices = PrimitiveArray::new( buffer![0, 3, 100], diff --git a/vortex-array/src/arrays/bool/compute/zip.rs b/vortex-array/src/arrays/bool/compute/zip.rs index 80e7de80c82..131562707d1 100644 --- a/vortex-array/src/arrays/bool/compute/zip.rs +++ b/vortex-array/src/arrays/bool/compute/zip.rs @@ -90,11 +90,11 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; #[test] fn blend_value_bits_boundaries() { @@ -132,7 +132,7 @@ mod tests { let bits: Vec = (0..len).map(|i| i.is_multiple_of(5) || i == 64).collect(); let mask = Mask::from_iter(bits.iter().copied()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -167,7 +167,7 @@ mod tests { let bits: Vec = (0..len).map(|i| i.is_multiple_of(2)).collect(); let mask = Mask::from_iter(bits.iter().copied()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? diff --git a/vortex-array/src/arrays/bool/mod.rs b/vortex-array/src/arrays/bool/mod.rs index ea51e6b5f29..5e7b192588d 100644 --- a/vortex-array/src/arrays/bool/mod.rs +++ b/vortex-array/src/arrays/bool/mod.rs @@ -15,7 +15,7 @@ pub use compute::rules::BoolMaskedValidityRule; pub use vtable::Bool; pub use vtable::BoolArray; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/bool/patch.rs b/vortex-array/src/arrays/bool/patch.rs index babf5a0ae1d..6f91c24b0c2 100644 --- a/vortex-array/src/arrays/bool/patch.rs +++ b/vortex-array/src/arrays/bool/patch.rs @@ -48,13 +48,13 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; + use crate::default_session_builder; #[test] fn patch_sliced_bools() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(12)); let sliced = arr.into_array().slice(4..12).unwrap(); let expected = BoolArray::from_iter([true; 8]); @@ -63,7 +63,7 @@ mod tests { #[test] fn patch_sliced_bools_offset() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from(BitBuffer::new_set(15)); let sliced = arr.into_array().slice(4..15).unwrap(); let expected = BoolArray::from_iter([true; 11]); diff --git a/vortex-array/src/arrays/bool/vtable/kernel.rs b/vortex-array/src/arrays/bool/vtable/kernel.rs index c19d00d5118..41f1b12847a 100644 --- a/vortex-array/src/arrays/bool/vtable/kernel.rs +++ b/vortex-array/src/arrays/bool/vtable/kernel.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Bool; use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::BooleanExecuteAdaptor; @@ -18,8 +18,8 @@ use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Binary.id(), Bool, BooleanExecuteAdaptor(Bool)); kernels.register_execute_parent_kernel(Cast.id(), Bool, CastExecuteAdaptor(Bool)); kernels.register_execute_parent_kernel(FillNull.id(), Bool, FillNullExecuteAdaptor(Bool)); diff --git a/vortex-array/src/arrays/bool/vtable/mod.rs b/vortex-array/src/arrays/bool/vtable/mod.rs index 348a1f5f69a..36c8e1bad7b 100644 --- a/vortex-array/src/arrays/bool/vtable/mod.rs +++ b/vortex-array/src/arrays/bool/vtable/mod.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayRef; use crate::ExecutionCtx; @@ -44,7 +45,7 @@ use crate::hash::ArrayHash; /// A [`Bool`]-encoded Vortex array. pub type BoolArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } @@ -212,15 +213,15 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; #[test] fn test_nullable_bool_serde_roundtrip() { - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array = BoolArray::from_iter([Some(true), None, Some(false), None]); let dtype = array.dtype().clone(); diff --git a/vortex-array/src/arrays/bool/vtable/operations.rs b/vortex-array/src/arrays/bool/vtable/operations.rs index 9bae2b9e19c..52d872d4ee1 100644 --- a/vortex-array/src/arrays/bool/vtable/operations.rs +++ b/vortex-array/src/arrays/bool/vtable/operations.rs @@ -31,10 +31,10 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; + use crate::default_session_builder; #[test] fn test_slice_hundred_elements() { @@ -48,7 +48,7 @@ mod tests { #[test] fn test_slice() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = BoolArray::from_iter([Some(true), Some(true), None, Some(false), None]); #[expect(deprecated)] let sliced_arr = arr.into_array().slice(1..4).unwrap().to_bool(); diff --git a/vortex-array/src/arrays/chunked/array.rs b/vortex-array/src/arrays/chunked/array.rs index 54544ba1b4d..4cea42e3c34 100644 --- a/vortex-array/src/arrays/chunked/array.rs +++ b/vortex-array/src/arrays/chunked/array.rs @@ -279,11 +279,11 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -291,7 +291,7 @@ mod test { #[test] fn test_rechunk_one_chunk() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -305,7 +305,7 @@ mod test { #[test] fn test_rechunk_two_chunks() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![buffer![0u64].into_array(), buffer![5u64].into_array()], DType::Primitive(PType::U64, Nullability::NonNullable), @@ -320,7 +320,7 @@ mod test { #[test] fn test_rechunk_tiny_target_chunks() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2, 3].into_array(), @@ -339,7 +339,7 @@ mod test { #[test] fn test_rechunk_with_too_big_chunk() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let chunked = ChunkedArray::try_new( vec![ buffer![0u64, 1, 2].into_array(), @@ -373,11 +373,11 @@ mod test { ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be all_valid since all non-empty chunks are all_valid - assert!(chunked.all_valid(&mut array_session().create_execution_ctx())?); + assert!(chunked.all_valid(&mut default_session_builder().build().create_execution_ctx())?); assert!( !chunked .into_array() - .all_invalid(&mut array_session().create_execution_ctx())? + .all_invalid(&mut default_session_builder().build().create_execution_ctx())? ); Ok(()) @@ -397,11 +397,11 @@ mod test { ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be all_invalid since all non-empty chunks are all_invalid - assert!(!chunked.all_valid(&mut array_session().create_execution_ctx())?); + assert!(!chunked.all_valid(&mut default_session_builder().build().create_execution_ctx())?); assert!( chunked .into_array() - .all_invalid(&mut array_session().create_execution_ctx())? + .all_invalid(&mut default_session_builder().build().create_execution_ctx())? ); Ok(()) @@ -421,11 +421,11 @@ mod test { ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?; // Should be neither all_valid nor all_invalid - assert!(!chunked.all_valid(&mut array_session().create_execution_ctx())?); + assert!(!chunked.all_valid(&mut default_session_builder().build().create_execution_ctx())?); assert!( !chunked .into_array() - .all_invalid(&mut array_session().create_execution_ctx())? + .all_invalid(&mut default_session_builder().build().create_execution_ctx())? ); Ok(()) diff --git a/vortex-array/src/arrays/chunked/compute/aggregate.rs b/vortex-array/src/arrays/chunked/compute/aggregate.rs index 149c72f499e..b83310037f0 100644 --- a/vortex-array/src/arrays/chunked/compute/aggregate.rs +++ b/vortex-array/src/arrays/chunked/compute/aggregate.rs @@ -48,17 +48,17 @@ mod tests { use crate::aggregate_fn::DynAccumulator; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::sum::Sum; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::scalar::Scalar; fn run_sum(batch: &crate::ArrayRef) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut acc = Accumulator::try_new( Sum, NumericalAggregateOpts::default(), diff --git a/vortex-array/src/arrays/chunked/compute/cast.rs b/vortex-array/src/arrays/chunked/compute/cast.rs index 80c6848b973..3bd6a6a3db0 100644 --- a/vortex-array/src/arrays/chunked/compute/cast.rs +++ b/vortex-array/src/arrays/chunked/compute/cast.rs @@ -36,19 +36,19 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; #[test] fn test_cast_chunked() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr0 = buffer![0u32, 1].into_array(); let arr1 = buffer![2u32, 3].into_array(); diff --git a/vortex-array/src/arrays/chunked/compute/kernel.rs b/vortex-array/src/arrays/chunked/compute/kernel.rs index db0042105cd..2cf2fe9b2ab 100644 --- a/vortex-array/src/arrays/chunked/compute/kernel.rs +++ b/vortex-array/src/arrays/chunked/compute/kernel.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Chunked; @@ -11,15 +11,15 @@ use crate::arrays::Slice; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; use crate::arrays::slice::SliceExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::mask::MaskExecuteAdaptor; use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Filter.id(), Chunked, FilterExecuteAdaptor(Chunked)); kernels.register_execute_parent_kernel(Mask.id(), Chunked, MaskExecuteAdaptor(Chunked)); kernels.register_execute_parent_kernel(Slice.id(), Chunked, SliceExecuteAdaptor(Chunked)); diff --git a/vortex-array/src/arrays/chunked/compute/mod.rs b/vortex-array/src/arrays/chunked/compute/mod.rs index f7c9a4e9d61..d07ed75c571 100644 --- a/vortex-array/src/arrays/chunked/compute/mod.rs +++ b/vortex-array/src/arrays/chunked/compute/mod.rs @@ -19,11 +19,11 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::binary_numeric::test_binary_numeric_array; use crate::compute::conformance::consistency::test_array_consistency; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -83,7 +83,7 @@ mod tests { fn test_chunked_consistency(#[case] array: ChunkedArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } @@ -161,7 +161,7 @@ mod tests { fn test_chunked_binary_numeric(#[case] array: ChunkedArray) { test_binary_numeric_array( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) } } diff --git a/vortex-array/src/arrays/chunked/compute/take.rs b/vortex-array/src/arrays/chunked/compute/take.rs index 34f443f1ce1..a581278d1af 100644 --- a/vortex-array/src/arrays/chunked/compute/take.rs +++ b/vortex-array/src/arrays/chunked/compute/take.rs @@ -129,7 +129,6 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; @@ -137,13 +136,14 @@ mod test { use crate::arrays::chunked::ChunkedArrayExt; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::validity::Validity; #[test] fn test_take() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -157,7 +157,7 @@ mod test { #[test] fn test_take_nullable_values() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let a = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -175,7 +175,7 @@ mod test { #[test] fn test_take_nullable_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -196,7 +196,7 @@ mod test { #[test] fn test_take_nullable_struct() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let struct_array = StructArray::try_new(FieldNames::default(), vec![], 100, Validity::NonNullable) .unwrap(); @@ -222,7 +222,7 @@ mod test { #[test] fn test_empty_take() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let a = buffer![1i32, 2, 3].into_array(); let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone()) .unwrap(); @@ -243,7 +243,7 @@ mod test { #[test] fn test_take_shuffled_indices() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let c0 = buffer![0i32, 1, 2].into_array(); let c1 = buffer![3i32, 4, 5].into_array(); let c2 = buffer![6i32, 7, 8].into_array(); @@ -312,7 +312,7 @@ mod test { #[test] fn test_take_null_indices() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let c0 = buffer![10i32, 20, 30].into_array(); let c1 = buffer![40i32, 50, 60].into_array(); let arr = ChunkedArray::try_new( diff --git a/vortex-array/src/arrays/chunked/compute/zip.rs b/vortex-array/src/arrays/chunked/compute/zip.rs index 4037fd1dc49..20517e09b89 100644 --- a/vortex-array/src/arrays/chunked/compute/zip.rs +++ b/vortex-array/src/arrays/chunked/compute/zip.rs @@ -54,11 +54,11 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -94,7 +94,7 @@ mod tests { // One step of execution will push down the zip. let zipped = zipped .clone() - .execute::(&mut array_session().create_execution_ctx()) + .execute::(&mut default_session_builder().build().create_execution_ctx()) .unwrap(); let zipped = zipped .as_opt::() diff --git a/vortex-array/src/arrays/chunked/mod.rs b/vortex-array/src/arrays/chunked/mod.rs index 4bd6cb43b7e..872728b1915 100644 --- a/vortex-array/src/arrays/chunked/mod.rs +++ b/vortex-array/src/arrays/chunked/mod.rs @@ -12,7 +12,7 @@ pub(crate) mod paired_chunks; mod vtable; pub use vtable::Chunked; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { compute::kernel::initialize(session); } diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index d3cae8be8f0..2694e4ad91a 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -32,7 +32,8 @@ use crate::dtype::PType::I32; use crate::executor::execute_into_builder; use crate::validity::Validity; -static SESSION: LazyLock = LazyLock::new(crate::array_session); +static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); fn chunked_array() -> ChunkedArray { ChunkedArray::try_new( diff --git a/vortex-array/src/arrays/chunked/vtable/canonical.rs b/vortex-array/src/arrays/chunked/vtable/canonical.rs index 3cc9505d16b..93af23b4793 100644 --- a/vortex-array/src/arrays/chunked/vtable/canonical.rs +++ b/vortex-array/src/arrays/chunked/vtable/canonical.rs @@ -319,13 +319,14 @@ mod tests { use crate::dtype::PType::I32; use crate::memory::DefaultHostAllocator; use crate::memory::HostAllocator; - use crate::memory::MemorySessionExt; + use crate::memory::MemorySessionBuilderExt; use crate::memory::WritableHostBuffer; use crate::scalar::Scalar; use crate::validity::Validity; /// A shared session for these chunked-array tests, used to create execution contexts. - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[derive(Debug)] struct CountingAllocator { @@ -641,9 +642,11 @@ mod tests { #[test] fn list_canonicalize_uses_memory_session_allocator() { let allocations = Arc::new(AtomicUsize::new(0)); - let session = crate::array_session().with_allocator(Arc::new(CountingAllocator { - allocations: Arc::clone(&allocations), - })); + let session = crate::default_session_builder() + .with_allocator(Arc::new(CountingAllocator { + allocations: Arc::clone(&allocations), + })) + .build(); let mut ctx = session.create_execution_ctx(); let l1 = ListArray::try_new( diff --git a/vortex-array/src/arrays/chunked/vtable/operations.rs b/vortex-array/src/arrays/chunked/vtable/operations.rs index 8f9e0867a88..a727bf40fea 100644 --- a/vortex-array/src/arrays/chunked/vtable/operations.rs +++ b/vortex-array/src/arrays/chunked/vtable/operations.rs @@ -31,10 +31,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -59,7 +59,7 @@ mod tests { #[case::end(7..8, &[8u64])] #[case::exactly_end(6..9, &[7u64, 8, 9])] fn slice(#[case] range: Range, #[case] expected: &[u64]) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_arrays_eq!( chunked_array().slice(range).unwrap(), PrimitiveArray::from_iter(expected.iter().copied()), @@ -77,7 +77,7 @@ mod tests { #[test] fn scalar_at_empty_children_both_sides() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), @@ -94,7 +94,7 @@ mod tests { #[test] fn scalar_at_empty_children_trailing() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ buffer![1u64, 2].into_array(), @@ -110,7 +110,7 @@ mod tests { #[test] fn scalar_at_empty_children_leading() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = ChunkedArray::try_new( vec![ Buffer::::empty().into_array(), diff --git a/vortex-array/src/arrays/constant/compute/fill_null.rs b/vortex-array/src/arrays/constant/compute/fill_null.rs index 1c9c2aaf903..f052074fffb 100644 --- a/vortex-array/src/arrays/constant/compute/fill_null.rs +++ b/vortex-array/src/arrays/constant/compute/fill_null.rs @@ -23,10 +23,10 @@ impl FillNullReduce for Constant { mod test { use crate::IntoArray as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ConstantArray; use crate::arrow::ArrowSessionExt; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -34,7 +34,7 @@ mod test { #[test] fn test_null() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let actual = ConstantArray::new(Scalar::null_native::(), 3) .into_array() .fill_null(Scalar::from(1)) @@ -43,11 +43,13 @@ mod test { assert!(!actual.dtype().is_nullable()); - let actual_arrow = array_session() + let actual_arrow = default_session_builder() + .build() .arrow() .execute_arrow(actual.clone(), None, &mut ctx) .unwrap(); - let expected_arrow = array_session() + let expected_arrow = default_session_builder() + .build() .arrow() .execute_arrow(expected.clone(), None, &mut ctx) .unwrap(); @@ -62,7 +64,7 @@ mod test { #[test] fn test_non_null() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let actual = ConstantArray::new(Scalar::from(Some(1)), 3) .into_array() .fill_null(Scalar::from(1)) @@ -71,11 +73,13 @@ mod test { assert!(!actual.dtype().is_nullable()); - let actual_arrow = array_session() + let actual_arrow = default_session_builder() + .build() .arrow() .execute_arrow(actual.clone(), None, &mut ctx) .unwrap(); - let expected_arrow = array_session() + let expected_arrow = default_session_builder() + .build() .arrow() .execute_arrow(expected.clone(), None, &mut ctx) .unwrap(); @@ -90,7 +94,7 @@ mod test { #[test] fn test_non_nullable_with_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let actual = ConstantArray::new(Scalar::from(1), 3) .into_array() .fill_null(Scalar::new( @@ -104,11 +108,13 @@ mod test { assert!(actual.dtype().is_nullable()); - let actual_arrow = array_session() + let actual_arrow = default_session_builder() + .build() .arrow() .execute_arrow(actual.clone(), None, &mut ctx) .unwrap(); - let expected_arrow = array_session() + let expected_arrow = default_session_builder() + .build() .arrow() .execute_arrow(expected.clone(), None, &mut ctx) .unwrap(); diff --git a/vortex-array/src/arrays/constant/compute/mod.rs b/vortex-array/src/arrays/constant/compute/mod.rs index 80fab02dc88..c42506725bd 100644 --- a/vortex-array/src/arrays/constant/compute/mod.rs +++ b/vortex-array/src/arrays/constant/compute/mod.rs @@ -19,11 +19,11 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ConstantArray; use crate::compute::conformance::consistency::test_array_consistency; use crate::compute::conformance::filter::test_filter_conformance; use crate::compute::conformance::mask::test_mask_conformance; + use crate::default_session_builder; use crate::dtype::half::f16; use crate::scalar::Scalar; @@ -66,7 +66,7 @@ mod test { fn test_constant_consistency(#[case] array: ConstantArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/constant/compute/take.rs b/vortex-array/src/arrays/constant/compute/take.rs index 9fb1373a5cf..6702c5c8545 100644 --- a/vortex-array/src/arrays/constant/compute/take.rs +++ b/vortex-array/src/arrays/constant/compute/take.rs @@ -76,18 +76,18 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::scalar::Scalar; use crate::validity::Validity; #[test] fn take_nullable_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take( @@ -116,7 +116,10 @@ mod tests { taken .validity() .unwrap() - .execute_mask(taken.len(), &mut array_session().create_execution_ctx()) + .execute_mask( + taken.len(), + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .indices(), AllOr::Some(valid_indices) @@ -125,7 +128,7 @@ mod tests { #[test] fn take_all_valid_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = ConstantArray::new(42, 10).into_array(); let taken = array .take(PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array()) @@ -144,7 +147,10 @@ mod tests { taken .validity() .unwrap() - .execute_mask(taken.len(), &mut array_session().create_execution_ctx()) + .execute_mask( + taken.len(), + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .indices(), AllOr::All diff --git a/vortex-array/src/arrays/constant/compute/uncompressed_size.rs b/vortex-array/src/arrays/constant/compute/uncompressed_size.rs index f8a4d4f92a4..4e0c34e0b3a 100644 --- a/vortex-array/src/arrays/constant/compute/uncompressed_size.rs +++ b/vortex-array/src/arrays/constant/compute/uncompressed_size.rs @@ -16,10 +16,10 @@ mod tests { use crate::aggregate_fn::EmptyOptions; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::canonical_uncompressed_size_in_bytes; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -27,7 +27,7 @@ mod tests { use crate::validity::Validity; fn aggregate(array: &ArrayRef) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut acc = Accumulator::try_new(UncompressedSizeInBytes, EmptyOptions, array.dtype().clone())?; acc.accumulate(array, &mut ctx)?; @@ -82,7 +82,7 @@ mod tests { #[test] fn constant_utf8_matches_canonical_size() -> VortexResult<()> { let constant = ConstantArray::new("abcdefghijkl".to_string(), 10).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let canonical = constant.clone().execute::(&mut ctx)?; let expected = canonical_uncompressed_size_in_bytes(&canonical, &mut ctx)?; diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index e160bcc16a9..82c85d0665b 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/vortex-array/src/arrays/constant/vtable/canonical.rs @@ -364,7 +364,8 @@ mod tests { use crate::validity::Validity; /// A shared session for these constant-array tests, used to create execution contexts. - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn test_canonicalize_null() { diff --git a/vortex-array/src/arrays/constant/vtable/mod.rs b/vortex-array/src/arrays/constant/vtable/mod.rs index b2996e4985c..908dc04f010 100644 --- a/vortex-array/src/arrays/constant/vtable/mod.rs +++ b/vortex-array/src/arrays/constant/vtable/mod.rs @@ -285,7 +285,9 @@ mod tests { /// Appends `array` into a fresh builder and asserts the result matches `constant_canonicalize`. fn assert_append_matches_canonical(array: ConstantArray) -> VortexResult<()> { - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let expected = constant_canonicalize(array.as_view(), &mut ctx)?.into_array(); let mut builder = builder_with_capacity(array.dtype(), array.len()); diff --git a/vortex-array/src/arrays/datetime/test.rs b/vortex-array/src/arrays/datetime/test.rs index 1ae114874d9..31fc2373a2e 100644 --- a/vortex-array/src/arrays/datetime/test.rs +++ b/vortex-array/src/arrays/datetime/test.rs @@ -11,10 +11,10 @@ use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::datetime::TemporalData; use crate::assert_arrays_eq; +use crate::default_session_builder; use crate::expr::gt; use crate::expr::lit; use crate::expr::root; @@ -28,7 +28,7 @@ use crate::validity::Validity; macro_rules! test_temporal_roundtrip { ($prim:ty, $constructor:expr, $unit:expr) => {{ - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = buffer![100 as $prim].into_array(); let temporal: TemporalData = $constructor(array, $unit); @@ -149,7 +149,7 @@ test_fail_case!( // We test Timestamp explicitly to avoid the macro getting too complex. #[test] fn test_timestamp() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ts = buffer![100i64].into_array(); let ts_array = ts.into_array(); diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index 80fafde5059..2beb2e1cadf 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -388,12 +388,12 @@ mod tests { use super::upcast_decimal_values; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::DecimalArray; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] use crate::canonical::ToCanonical as _; use crate::compute::conformance::cast::test_cast_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::DecimalType; @@ -572,7 +572,7 @@ mod tests { .unwrap() .execute_mask( casted.as_ref().len(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap(); assert!(mask.value(0)); @@ -736,7 +736,7 @@ mod tests { .unwrap() .execute_mask( casted.as_ref().len(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap(); assert!(mask.value(0)); diff --git a/vortex-array/src/arrays/decimal/compute/fill_null.rs b/vortex-array/src/arrays/decimal/compute/fill_null.rs index 1c2904b7b6b..4d479f315b6 100644 --- a/vortex-array/src/arrays/decimal/compute/fill_null.rs +++ b/vortex-array/src/arrays/decimal/compute/fill_null.rs @@ -90,12 +90,12 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] use crate::canonical::ToCanonical as _; + use crate::default_session_builder; use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::scalar::DecimalValue; @@ -104,7 +104,7 @@ mod tests { #[test] fn fill_null_leading_none() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( [None, Some(800i128), None, Some(1000i128), None], @@ -135,7 +135,7 @@ mod tests { .unwrap() .execute_mask( p.as_ref().len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap() .all_true() @@ -144,7 +144,7 @@ mod tests { #[test] fn fill_null_all_none() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::from_option_iter( @@ -172,7 +172,7 @@ mod tests { /// fill_null with a value that overflows the array's storage type should upcast the array. #[test] fn fill_null_overflow_upcasts() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let decimal_dtype = DecimalDType::new(3, 0); let arr = DecimalArray::from_option_iter([None, Some(10i8), None], decimal_dtype); // i8 max is 127, so 200 doesn't fit — the array should be widened to i16. @@ -195,7 +195,7 @@ mod tests { #[test] fn fill_null_non_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let decimal_dtype = DecimalDType::new(19, 2); let arr = DecimalArray::new( diff --git a/vortex-array/src/arrays/decimal/compute/mod.rs b/vortex-array/src/arrays/decimal/compute/mod.rs index b2f21fa8398..b3f404cd6d3 100644 --- a/vortex-array/src/arrays/decimal/compute/mod.rs +++ b/vortex-array/src/arrays/decimal/compute/mod.rs @@ -15,9 +15,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::DecimalArray; use crate::compute::conformance::consistency::test_array_consistency; + use crate::default_session_builder; use crate::dtype::DecimalDType; use crate::validity::Validity; @@ -57,7 +57,7 @@ mod tests { fn test_decimal_consistency(#[case] array: DecimalArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index 762ac3d7365..037381fb5d9 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -54,17 +54,17 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::dtype::DecimalDType; use crate::validity::Validity; #[test] fn test_take() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![10i128, 11i128, 12i128, 13i128], @@ -81,7 +81,7 @@ mod tests { #[test] fn test_take_null_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ddtype = DecimalDType::new(19, 1); let array = DecimalArray::new( buffer![i128::MAX, 11i128, 12i128, 13i128], diff --git a/vortex-array/src/arrays/decimal/mod.rs b/vortex-array/src/arrays/decimal/mod.rs index e740a75ea91..cdf1a47f2e9 100644 --- a/vortex-array/src/arrays/decimal/mod.rs +++ b/vortex-array/src/arrays/decimal/mod.rs @@ -13,7 +13,7 @@ mod vtable; pub use compute::rules::DecimalMaskedValidityRule; pub use vtable::Decimal; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/decimal/vtable/kernel.rs b/vortex-array/src/arrays/decimal/vtable/kernel.rs index 19bb30d8c95..d70b682ccaa 100644 --- a/vortex-array/src/arrays/decimal/vtable/kernel.rs +++ b/vortex-array/src/arrays/decimal/vtable/kernel.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Decimal; use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; @@ -16,8 +16,8 @@ use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Between.id(), Decimal, BetweenExecuteAdaptor(Decimal)); kernels.register_execute_parent_kernel(Cast.id(), Decimal, CastExecuteAdaptor(Decimal)); kernels.register_execute_parent_kernel(FillNull.id(), Decimal, FillNullExecuteAdaptor(Decimal)); diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index a52f2882238..540ad557ca1 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -10,6 +10,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayRef; use crate::ExecutionCtx; @@ -42,7 +43,7 @@ use crate::hash::ArrayHash; /// A [`Decimal`]-encoded Vortex array. pub type DecimalArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } @@ -209,10 +210,10 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DecimalDType; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -220,7 +221,7 @@ mod tests { #[test] fn test_array_serde() { - let session = array_session(); + let session = default_session_builder().build(); let array = DecimalArray::new( buffer![100i128, 200i128, 300i128, 400i128, 500i128], DecimalDType::new(10, 2), @@ -250,7 +251,7 @@ mod tests { #[test] fn test_nullable_decimal_serde_roundtrip() { - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array = DecimalArray::new( buffer![1234567i32, 0i32, -9999999i32], diff --git a/vortex-array/src/arrays/decimal/vtable/operations.rs b/vortex-array/src/arrays/decimal/vtable/operations.rs index 257f26127ae..bc22bdd202d 100644 --- a/vortex-array/src/arrays/decimal/vtable/operations.rs +++ b/vortex-array/src/arrays/decimal/vtable/operations.rs @@ -33,9 +33,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Decimal; use crate::arrays::DecimalArray; + use crate::default_session_builder; use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::scalar::DecimalValue; @@ -81,7 +81,10 @@ mod tests { assert_eq!( array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::decimal( DecimalValue::I128(100), diff --git a/vortex-array/src/arrays/dict/array.rs b/vortex-array/src/arrays/dict/array.rs index f73f4965439..631fd38fb5a 100644 --- a/vortex-array/src/arrays/dict/array.rs +++ b/vortex-array/src/arrays/dict/array.rs @@ -302,12 +302,12 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::builder_with_capacity; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability::NonNullable; @@ -332,7 +332,7 @@ mod test { .unwrap() .execute_mask( dict.as_ref().len(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap(); let AllOr::Some(indices) = mask.indices() else { @@ -358,7 +358,7 @@ mod test { .unwrap() .execute_mask( dict.as_ref().len(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap(); let AllOr::Some(indices) = mask.indices() else { @@ -388,7 +388,7 @@ mod test { .unwrap() .execute_mask( dict.as_ref().len(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap(); let AllOr::Some(indices) = mask.indices() else { @@ -414,7 +414,7 @@ mod test { .unwrap() .execute_mask( dict.as_ref().len(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap(); let AllOr::Some(indices) = mask.indices() else { @@ -454,7 +454,7 @@ mod test { #[test] fn test_dict_array_from_primitive_chunks() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let len = 2; let chunk_count = 2; let array = make_dict_primitive_chunks::(len, 2, chunk_count); @@ -465,7 +465,7 @@ mod test { ); array.append_to_builder( builder.as_mut(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; #[expect(deprecated)] diff --git a/vortex-array/src/arrays/dict/compute/cast.rs b/vortex-array/src/arrays/dict/compute/cast.rs index 56a433627a2..46f55d71709 100644 --- a/vortex-array/src/arrays/dict/compute/cast.rs +++ b/vortex-array/src/arrays/dict/compute/cast.rs @@ -69,7 +69,8 @@ mod tests { use crate::dtype::PType; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn test_cast_dict_to_wider_type() { diff --git a/vortex-array/src/arrays/dict/compute/fill_null.rs b/vortex-array/src/arrays/dict/compute/fill_null.rs index 290fdd95bda..d74502c10a5 100644 --- a/vortex-array/src/arrays/dict/compute/fill_null.rs +++ b/vortex-array/src/arrays/dict/compute/fill_null.rs @@ -95,18 +95,18 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::scalar::Scalar; use crate::validity::Validity; #[test] fn nullable_codes_fill_in_values() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::new( buffer![0u32, 1, 2], diff --git a/vortex-array/src/arrays/dict/compute/like.rs b/vortex-array/src/arrays/dict/compute/like.rs index a4ca7c1f76e..729528fab98 100644 --- a/vortex-array/src/arrays/dict/compute/like.rs +++ b/vortex-array/src/arrays/dict/compute/like.rs @@ -57,20 +57,20 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::DictArray; use crate::arrays::VarBinArray; use crate::arrays::dict::compute::like::ConstantArray; use crate::arrays::scalar_fn::ScalarFnFactoryExt; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::optimizer::ArrayOptimizer; use crate::scalar_fn::fns::like::Like; use crate::scalar_fn::fns::like::LikeOptions; #[test] fn like_reduce_dict() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict = DictArray::try_new( buffer![0u8, 1, 0, 2].into_array(), VarBinArray::from(vec!["hello", "world", "help"]).into_array(), diff --git a/vortex-array/src/arrays/dict/compute/min_max.rs b/vortex-array/src/arrays/dict/compute/min_max.rs index ee38c0392d4..d80d919a65c 100644 --- a/vortex-array/src/arrays/dict/compute/min_max.rs +++ b/vortex-array/src/arrays/dict/compute/min_max.rs @@ -76,7 +76,8 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::builders::dict::dict_encode; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); fn assert_min_max(array: &ArrayRef, expected: Option<(i32, i32)>) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-array/src/arrays/dict/compute/mod.rs b/vortex-array/src/arrays/dict/compute/mod.rs index defa44a0976..b65e83b2562 100644 --- a/vortex-array/src/arrays/dict/compute/mod.rs +++ b/vortex-array/src/arrays/dict/compute/mod.rs @@ -78,7 +78,8 @@ mod test { use crate::dtype::PType::I32; use crate::scalar_fn::fns::operators::Operator; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn canonicalise_nullable_primitive() { @@ -341,7 +342,8 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[rstest] // Primitive arrays diff --git a/vortex-array/src/arrays/dict/compute/slice.rs b/vortex-array/src/arrays/dict/compute/slice.rs index 376621fba2f..9eb1abda478 100644 --- a/vortex-array/src/arrays/dict/compute/slice.rs +++ b/vortex-array/src/arrays/dict/compute/slice.rs @@ -92,11 +92,11 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::dict::compute::slice::ConstantArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; use crate::dtype::PType; @@ -104,7 +104,7 @@ mod tests { #[test] fn slice_constant_valid_code() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(1u8, 5).into_array(), buffer![10i32, 20, 30].into_array(), @@ -117,7 +117,7 @@ mod tests { #[test] fn slice_constant_null_code() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict = DictArray::new( ConstantArray::new(Scalar::null(DType::Primitive(PType::U8, Nullable)), 5).into_array(), buffer![10i32, 20, 30].into_array(), diff --git a/vortex-array/src/arrays/dict/mod.rs b/vortex-array/src/arrays/dict/mod.rs index e8a83aa88b3..e189a18ce2b 100644 --- a/vortex-array/src/arrays/dict/mod.rs +++ b/vortex-array/src/arrays/dict/mod.rs @@ -23,7 +23,7 @@ pub use take::*; pub mod vtable; pub use vtable::*; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/dict/tests.rs b/vortex-array/src/arrays/dict/tests.rs index 7961e2b7037..da4d72d3260 100644 --- a/vortex-array/src/arrays/dict/tests.rs +++ b/vortex-array/src/arrays/dict/tests.rs @@ -6,17 +6,17 @@ use vortex_buffer::buffer; use super::DictArray; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; +use crate::default_session_builder; use crate::validity::Validity; #[test] fn test_scalar_at_null_code() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict = DictArray::try_new( PrimitiveArray::from_option_iter(vec![None, Some(0u32), None]).into_array(), buffer![1i32].into_array(), diff --git a/vortex-array/src/arrays/dict/vtable/kernel.rs b/vortex-array/src/arrays/dict/vtable/kernel.rs index ab750f7d663..c979e8fda57 100644 --- a/vortex-array/src/arrays/dict/vtable/kernel.rs +++ b/vortex-array/src/arrays/dict/vtable/kernel.rs @@ -1,20 +1,20 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; use crate::scalar_fn::fns::fill_null::FillNull; use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Binary.id(), Dict, CompareExecuteAdaptor(Dict)); kernels.register_execute_parent_kernel(Dict.id(), Dict, TakeExecuteAdaptor(Dict)); kernels.register_execute_parent_kernel(FillNull.id(), Dict, FillNullExecuteAdaptor(Dict)); diff --git a/vortex-array/src/arrays/dict/vtable/mod.rs b/vortex-array/src/arrays/dict/vtable/mod.rs index cd2ef72a7ee..8f135c41bdd 100644 --- a/vortex-array/src/arrays/dict/vtable/mod.rs +++ b/vortex-array/src/arrays/dict/vtable/mod.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::CachedId; use super::DictData; @@ -55,7 +56,7 @@ mod validity; /// A [`Dict`]-encoded Vortex array. pub type DictArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } diff --git a/vortex-array/src/arrays/extension/compute/cast.rs b/vortex-array/src/arrays/extension/compute/cast.rs index dfce610f43a..999598c72b1 100644 --- a/vortex-array/src/arrays/extension/compute/cast.rs +++ b/vortex-array/src/arrays/extension/compute/cast.rs @@ -64,7 +64,8 @@ mod tests { use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn cast_same_ext_dtype() { diff --git a/vortex-array/src/arrays/extension/compute/mod.rs b/vortex-array/src/arrays/extension/compute/mod.rs index 8309ae43be2..f2b4a0131c0 100644 --- a/vortex-array/src/arrays/extension/compute/mod.rs +++ b/vortex-array/src/arrays/extension/compute/mod.rs @@ -88,10 +88,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ExtensionArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::consistency::test_array_consistency; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; @@ -118,7 +118,7 @@ mod tests { fn test_extension_consistency(#[case] array: ExtensionArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/extension/compute/rules.rs b/vortex-array/src/arrays/extension/compute/rules.rs index 611be1510a4..18d17483dca 100644 --- a/vortex-array/src/arrays/extension/compute/rules.rs +++ b/vortex-array/src/arrays/extension/compute/rules.rs @@ -112,7 +112,8 @@ mod tests { use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::operators::Operator; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[derive(Clone, Debug, Default, PartialEq, Eq, Hash)] struct TestExt; diff --git a/vortex-array/src/arrays/extension/mod.rs b/vortex-array/src/arrays/extension/mod.rs index 8dc5972b86f..e0485830216 100644 --- a/vortex-array/src/arrays/extension/mod.rs +++ b/vortex-array/src/arrays/extension/mod.rs @@ -10,6 +10,6 @@ pub(crate) mod compute; mod vtable; pub use vtable::Extension; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/extension/vtable/kernel.rs b/vortex-array/src/arrays/extension/vtable/kernel.rs index c90458f1d5f..cc9390bff9c 100644 --- a/vortex-array/src/arrays/extension/vtable/kernel.rs +++ b/vortex-array/src/arrays/extension/vtable/kernel.rs @@ -1,19 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::Extension; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel( Binary.id(), Extension, diff --git a/vortex-array/src/arrays/extension/vtable/mod.rs b/vortex-array/src/arrays/extension/vtable/mod.rs index a1fab1cd47a..96952a23e7d 100644 --- a/vortex-array/src/arrays/extension/vtable/mod.rs +++ b/vortex-array/src/arrays/extension/vtable/mod.rs @@ -9,6 +9,7 @@ use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::CachedId; use crate::ArrayRef; @@ -77,7 +78,7 @@ pub struct Extension; /// A [`Extension`]-encoded Vortex array. pub type ExtensionArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index 392a6373f76..d9109e1379b 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -125,11 +125,11 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::validity::Validity; @@ -151,7 +151,7 @@ mod test { #[test] fn filter_fixed_size_list_selects_correct_lists() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let array = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -168,7 +168,7 @@ mod test { #[test] fn filter_degenerate_list_size_zero() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::empty::(Nullability::NonNullable); let array = FixedSizeListArray::new(elements.into_array(), 0, Validity::NonNullable, 5); @@ -184,7 +184,7 @@ mod test { #[test] fn filter_nested_fixed_size_lists() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Inner lists of size 2, outer lists of size 2 (so 2 outer lists, each with 2 inner lists). let inner_elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); let inner_fsl = FixedSizeListArray::new(inner_elements, 2, Validity::NonNullable, 4); diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index 0507452c89d..1096c79c2e1 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -66,16 +66,16 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::filter::execute::ConstantArray; use crate::arrays::listview::ListViewArrayExt; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; + use crate::default_session_builder; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(array_session); + static SESSION: LazyLock = LazyLock::new(|| default_session_builder().build()); #[test] fn test_filter_listview_conformance() { @@ -135,7 +135,7 @@ mod test { #[test] fn filter_listview_selects_correct_lists() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // 3 lists: [10,20], [30,40], [50,60] let elements = PrimitiveArray::from_iter([10i32, 20, 30, 40, 50, 60]); let offsets = buffer![0u32, 2, 4].into_array(); @@ -163,7 +163,7 @@ mod test { #[test] fn test_filter_preserves_unreferenced_elements() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // ListView-specific: Test that filter preserves the entire elements array. // // Logical list: [[5,6,7], [2,3], [8,9], [0,1], [1,2,3,4]] @@ -198,7 +198,7 @@ mod test { #[test] fn test_filter_with_gaps() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // ListView-specific: Test filtering with gaps in elements array. // // Logical list: [[1,2,3], [7,8,9], [11,12], [2,3], [8,9]] diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index 8980c422ef3..c4da54efad2 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -50,13 +50,13 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability::Nullable; @@ -88,7 +88,7 @@ mod test { #[test] fn filter_struct_selects_correct_rows() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = StructArray::try_new( ["x", "y"].into(), vec![ @@ -119,7 +119,7 @@ mod test { #[test] fn filter_empty_struct() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let mask = Mask::from_iter([ @@ -134,7 +134,7 @@ mod test { #[test] fn filter_empty_struct_with_empty_filter() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 0, Validity::NonNullable).unwrap(); let filtered = struct_arr.filter(Mask::from_iter::<[bool; 0]>([])).unwrap(); diff --git a/vortex-array/src/arrays/filter/execute/take/tests.rs b/vortex-array/src/arrays/filter/execute/take/tests.rs index 36ab1a63dcb..1d861b9dd54 100644 --- a/vortex-array/src/arrays/filter/execute/take/tests.rs +++ b/vortex-array/src/arrays/filter/execute/take/tests.rs @@ -59,7 +59,9 @@ fn test_take_execute_kernel_maps_indices_through_filter() -> VortexResult<()> { filter.clone(), )? .into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -88,7 +90,9 @@ fn test_take_execute_kernel_nullable_fast_path_maps_indices_through_filter() -> filter.clone(), )? .into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -110,7 +114,9 @@ fn test_take_execute_kernel_fast_path_maps_indices_through_filter() -> VortexRes ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 3].into_array(), filter.clone())?.into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -131,7 +137,9 @@ fn assert_take_execute_rejects_out_of_bounds_rank( ) -> VortexResult<()> { let filter = FilterArray::new(child, filter_mask).into_array(); let parent = DictArray::try_new(codes, filter.clone())?.into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); if let Err(err) = execute_parent(&filter, &parent, 1, &mut ctx) { assert!( @@ -209,7 +217,9 @@ fn test_take_execute_kernel_handles_empty_sequential_take() -> VortexResult<()> filter.clone(), )? .into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -234,7 +244,9 @@ fn assert_take_execute_maps_child_dtype( let filter = FilterArray::new(child, Mask::from_iter([true, false, true, true, false])).into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -255,7 +267,9 @@ fn test_take_execute_kernel_skips_bool_filter_child() -> VortexResult<()> { ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)?; @@ -277,7 +291,9 @@ fn execute_primitive_take( .into_array(); let indices = PrimitiveArray::from_iter((0..take_len).map(|idx| (idx % filtered_len) as u64)); let parent = DictArray::try_new(indices.into_array(), filter.clone())?.into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); execute_parent(&filter, &parent, 1, &mut ctx) } @@ -344,7 +360,9 @@ fn test_take_execute_kernel_handles_nullable_primitive_filter_child() -> VortexR ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)?; @@ -367,7 +385,9 @@ fn test_take_execute_kernel_preserves_nullable_all_valid_fixed_width_child() -> ) .into_array(); let parent = DictArray::try_new(buffer![0u64, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its take parent"); @@ -394,7 +414,9 @@ fn test_take_execute_kernel_handles_nullable_decimal_filter_child() -> VortexRes ) .into_array(); let parent = DictArray::try_new(buffer![2u64, 0, 1].into_array(), filter.clone())?.into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)?; @@ -488,7 +510,9 @@ fn test_take_execute_kernel_preserves_nullable_indices_dtype_fast_path() -> Vort filter.clone(), )? .into_array(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let result = execute_parent(&filter, &parent, 1, &mut ctx)? .expect("filter child should execute its nullable take parent"); diff --git a/vortex-array/src/arrays/filter/kernel.rs b/vortex-array/src/arrays/filter/kernel.rs index 21bd225bf55..f7c8ac262cf 100644 --- a/vortex-array/src/arrays/filter/kernel.rs +++ b/vortex-array/src/arrays/filter/kernel.rs @@ -10,7 +10,7 @@ use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayRef; use crate::ArrayVTable; @@ -24,13 +24,15 @@ use crate::arrays::Filter; use crate::arrays::dict::TakeExecuteAdaptor; use crate::kernel::ExecuteParentKernel; use crate::matcher::Matcher; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::optimizer::rules::ArrayParentReduceRule; -pub(crate) fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(Dict.id(), Filter, TakeExecuteAdaptor(Filter)); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + builder_kernels(session).register_execute_parent_kernel( + Dict.id(), + Filter, + TakeExecuteAdaptor(Filter), + ); } pub trait FilterReduce: VTable { diff --git a/vortex-array/src/arrays/filter/mod.rs b/vortex-array/src/arrays/filter/mod.rs index 9510ca92f5b..86be43f85cb 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -20,6 +20,6 @@ mod rules; mod vtable; pub use vtable::Filter; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { kernel::initialize(session); } diff --git a/vortex-array/src/arrays/fixed_size_list/mod.rs b/vortex-array/src/arrays/fixed_size_list/mod.rs index 3d90d38e6e9..3d738eed3f9 100644 --- a/vortex-array/src/arrays/fixed_size_list/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/mod.rs @@ -12,7 +12,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::FixedSizeList; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs index 961801094d8..f48f9fa0213 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/basic.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/basic.rs @@ -7,9 +7,9 @@ use vortex_buffer::buffer; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -40,19 +40,28 @@ fn test_basic_fixed_size_list() { let first_list = fsl.fixed_size_list_elements_at(0).unwrap(); assert_eq!( first_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); assert_eq!( first_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 2i32.into() ); assert_eq!( first_list - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 3i32.into() ); @@ -60,19 +69,28 @@ fn test_basic_fixed_size_list() { let second_list = fsl.fixed_size_list_elements_at(1).unwrap(); assert_eq!( second_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 4i32.into() ); assert_eq!( second_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 5i32.into() ); assert_eq!( second_list - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 6i32.into() ); @@ -80,19 +98,28 @@ fn test_basic_fixed_size_list() { let third_list = fsl.fixed_size_list_elements_at(2).unwrap(); assert_eq!( third_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 7i32.into() ); assert_eq!( third_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 8i32.into() ); assert_eq!( third_list - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 9i32.into() ); @@ -100,19 +127,28 @@ fn test_basic_fixed_size_list() { let fourth_list = fsl.fixed_size_list_elements_at(3).unwrap(); assert_eq!( fourth_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 10i32.into() ); assert_eq!( fourth_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 11i32.into() ); assert_eq!( fourth_list - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 12i32.into() ); @@ -128,7 +164,10 @@ fn test_scalar_at() { // First list: [1, 2, 3]. let first = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert_eq!( first, @@ -143,26 +182,38 @@ fn test_scalar_at() { let first_list = fsl.fixed_size_list_elements_at(0).unwrap(); assert_eq!( first_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); assert_eq!( first_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 2i32.into() ); assert_eq!( first_list - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 3i32.into() ); // Second list: [4, 5, 6]. let second = fsl - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert_eq!( second, @@ -177,19 +228,28 @@ fn test_scalar_at() { let second_list = fsl.fixed_size_list_elements_at(1).unwrap(); assert_eq!( second_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 4i32.into() ); assert_eq!( second_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 5i32.into() ); assert_eq!( second_list - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 6i32.into() ); @@ -208,13 +268,19 @@ fn test_fixed_size_list_at() { assert_eq!(first_list.len(), list_size as usize); assert_eq!( first_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1.0f64.into() ); assert_eq!( first_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 2.0f64.into() ); @@ -224,13 +290,19 @@ fn test_fixed_size_list_at() { assert_eq!(third_list.len(), list_size as usize); assert_eq!( third_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 5.0f64.into() ); assert_eq!( third_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 6.0f64.into() ); diff --git a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs index f609fd72405..bdee3eb9428 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/degenerate.rs @@ -7,10 +7,10 @@ use vortex_buffer::buffer; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -57,7 +57,10 @@ fn test_fsl_size_0_length_1_non_nullable() { // Get the single empty list. let scalar = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!scalar.is_null()); assert_eq!( @@ -85,7 +88,10 @@ fn test_fsl_size_0_huge_length_non_nullable() { // Spot check a few lists. let scalar_first = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!scalar_first.is_null()); assert_eq!( @@ -98,7 +104,10 @@ fn test_fsl_size_0_huge_length_non_nullable() { ); let scalar_middle = fsl - .execute_scalar(500_000_000_000, &mut array_session().create_execution_ctx()) + .execute_scalar( + 500_000_000_000, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!scalar_middle.is_null()); assert_eq!( @@ -111,7 +120,10 @@ fn test_fsl_size_0_huge_length_non_nullable() { ); let scalar_end = fsl - .execute_scalar(999_999_999_999, &mut array_session().create_execution_ctx()) + .execute_scalar( + 999_999_999_999, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!scalar_end.is_null()); assert_eq!( @@ -163,7 +175,10 @@ fn test_fsl_size_0_length_1_nullable_valid() { // Get the single empty list (should be valid). let scalar = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!scalar.is_null()); assert_eq!( @@ -188,7 +203,10 @@ fn test_fsl_size_0_length_1_nullable_null() { // The single list should be null. let scalar = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(scalar.is_null()); } @@ -215,7 +233,10 @@ fn test_fsl_size_0_length_10_nullable_mixed() { ]; for i in 0..len { let scalar = fsl - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); if expected_valid[i] { assert!(!scalar.is_null()); @@ -256,7 +277,10 @@ fn test_fsl_size_0_nullable_elements() { // All lists should be empty but valid. for i in 0..len { let scalar = fsl - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!scalar.is_null()); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs index 6f609120901..567f697d746 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/filter.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/filter.rs @@ -9,7 +9,6 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; @@ -18,6 +17,7 @@ use crate::compute::conformance::filter::LARGE_SIZE; use crate::compute::conformance::filter::MEDIUM_SIZE; use crate::compute::conformance::filter::SMALL_SIZE; use crate::compute::conformance::filter::test_filter_conformance; +use crate::default_session_builder; use crate::dtype::Nullability; use crate::validity::Validity; @@ -53,7 +53,7 @@ fn test_filter_degenerate_list_size_zero( #[case] mask_values: Vec, #[case] expected_len: usize, ) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let new_validity = if matches!(validity, Validity::NonNullable) { Validity::NonNullable } else { @@ -82,7 +82,7 @@ fn test_filter_degenerate_list_size_zero( #[test] fn test_filter_with_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4), Some(5), Some(6)]); let validity = Validity::from_iter([true, false, true]); @@ -110,7 +110,7 @@ fn test_filter_with_nulls() { #[test] fn test_filter_all_null_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create an array where all elements are null. let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let validity = Validity::AllInvalid; @@ -127,7 +127,7 @@ fn test_filter_all_null_array() { #[test] fn test_filter_nested_fixed_size_lists() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create nested fixed-size lists: FSL>. // Inner lists are of size 2, outer lists are of size 3. // So we have 2 outer lists, each containing 3 inner lists, each containing 2 i32s. @@ -253,7 +253,7 @@ fn create_fsl_empty() -> ArrayRef { #[test] fn test_filter_all_null_various_list_sizes() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test filtering with all-null arrays of different list sizes. // The implementation returns ConstantArray only when validity_mask() is Mask::AllFalse. @@ -292,7 +292,7 @@ fn test_filter_all_null_various_list_sizes() { #[test] fn test_mask_expansion_threshold_boundary() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test with list_size == 8 (the FSL_SPARSE_MASK_LIST_SIZE_THRESHOLD). let list_size = 8u32; let num_lists = 100; @@ -369,7 +369,7 @@ fn test_mask_expansion_threshold_boundary() { // Test FSL-specific behavior with very large list sizes. #[test] fn test_filter_large_list_size() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test with list_size=100, which is significantly larger than typical use cases. let list_size = 100u32; let num_lists = 5; diff --git a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs index 44420d22cb3..3f2d815bd70 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nested.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nested.rs @@ -9,13 +9,13 @@ use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::builders::ArrayBuilder; use crate::builders::ListBuilder; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; @@ -76,7 +76,10 @@ fn test_fsl_of_fsl_basic() { // The first outer list should contain 3 inner lists. // We can check by slicing and examining scalars. let first_scalar = outer_fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!first_scalar.is_null()); @@ -92,13 +95,19 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_0 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); assert_eq!( inner_list_0 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 2i32.into() ); @@ -111,13 +120,19 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_1 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 3i32.into() ); assert_eq!( inner_list_1 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 4i32.into() ); @@ -130,13 +145,19 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_2 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 5i32.into() ); assert_eq!( inner_list_2 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 6i32.into() ); @@ -152,13 +173,19 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_0 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 7i32.into() ); assert_eq!( inner_list_0 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 8i32.into() ); @@ -171,13 +198,19 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_1 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 9i32.into() ); assert_eq!( inner_list_1 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 10i32.into() ); @@ -190,13 +223,19 @@ fn test_fsl_of_fsl_basic() { .unwrap(); assert_eq!( inner_list_2 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 11i32.into() ); assert_eq!( inner_list_2 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 12i32.into() ); @@ -246,7 +285,10 @@ fn test_fsl_of_fsl_with_nulls() { // First outer list is valid. assert!( !outer_fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); @@ -254,7 +296,10 @@ fn test_fsl_of_fsl_with_nulls() { // Second outer list is null. assert!( outer_fsl - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); @@ -262,7 +307,10 @@ fn test_fsl_of_fsl_with_nulls() { // Third outer list is valid. assert!( !outer_fsl - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); @@ -326,13 +374,19 @@ fn test_deeply_nested_fsl() { .unwrap(); assert_eq!( level1_0_0 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); assert_eq!( level1_0_0 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 2i32.into() ); @@ -344,13 +398,19 @@ fn test_deeply_nested_fsl() { .unwrap(); assert_eq!( level1_0_1 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 3i32.into() ); assert_eq!( level1_0_1 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 4i32.into() ); @@ -363,13 +423,19 @@ fn test_deeply_nested_fsl() { .unwrap(); assert_eq!( level1_1_0 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 5i32.into() ); assert_eq!( level1_1_0 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 6i32.into() ); @@ -381,13 +447,19 @@ fn test_deeply_nested_fsl() { .unwrap(); assert_eq!( level1_1_1 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 7i32.into() ); assert_eq!( level1_1_1 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 8i32.into() ); diff --git a/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs b/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs index f5e5339e756..7151263a5cf 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/nullability.rs @@ -7,11 +7,11 @@ use vortex_buffer::buffer; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -33,7 +33,10 @@ fn test_nullable_fsl_with_nulls() { // First list is valid: [1, 2]. let first = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!first.is_null()); assert_eq!( @@ -49,26 +52,38 @@ fn test_nullable_fsl_with_nulls() { let first_list = fsl.fixed_size_list_elements_at(0).unwrap(); assert_eq!( first_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); assert_eq!( first_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 2i32.into() ); // Second list is null. let second = fsl - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(second.is_null()); // Third list is valid: [5, 6]. let third = fsl - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!third.is_null()); assert_eq!( @@ -84,20 +99,29 @@ fn test_nullable_fsl_with_nulls() { let third_list = fsl.fixed_size_list_elements_at(2).unwrap(); assert_eq!( third_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 5i32.into() ); assert_eq!( third_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 6i32.into() ); // Fourth list is null. let fourth = fsl - .execute_scalar(3, &mut array_session().create_execution_ctx()) + .execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(fourth.is_null()); } @@ -123,7 +147,10 @@ fn test_nullable_elements_non_nullable_lists() { // First list: [Some(1), None, Some(3)]. let first = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!first.is_null()); assert_eq!( @@ -137,7 +164,10 @@ fn test_nullable_elements_non_nullable_lists() { // Second list: [Some(4), Some(5), None]. let second = fsl - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!second.is_null()); assert_eq!( @@ -165,7 +195,10 @@ fn test_nullable_elements_and_nullable_lists() { // First list is valid: [Some(10), None]. let first = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!first.is_null()); assert_eq!( @@ -181,26 +214,38 @@ fn test_nullable_elements_and_nullable_lists() { let first_list = fsl.fixed_size_list_elements_at(0).unwrap(); assert_eq!( first_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Some(10u16).into() ); assert_eq!( first_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), None::.into() ); // Second list is null (but elements would be [Some(20), Some(30)]). let second = fsl - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(second.is_null()); // Third list is valid: [None, None]. let third = fsl - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!third.is_null()); assert_eq!( @@ -216,13 +261,19 @@ fn test_nullable_elements_and_nullable_lists() { let third_list = fsl.fixed_size_list_elements_at(2).unwrap(); assert_eq!( third_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), None::.into() ); assert_eq!( third_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), None::.into() ); @@ -243,7 +294,10 @@ fn test_alternating_nulls() { // Check alternating pattern. for i in 0..len { let scalar = fsl - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); if i % 2 == 0 { assert!(!scalar.is_null()); @@ -275,9 +329,12 @@ fn test_validity_types() { let fsl = FixedSizeListArray::new(elements.clone(), list_size, Validity::AllInvalid, len); for i in 0..len { assert!( - fsl.execute_scalar(i, &mut array_session().create_execution_ctx()) - .unwrap() - .is_null() + fsl.execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap() + .is_null() ); } } @@ -293,24 +350,36 @@ fn test_validity_types() { ); assert!( - !fsl.execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap() - .is_null() + !fsl.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap() + .is_null() ); assert!( - !fsl.execute_scalar(1, &mut array_session().create_execution_ctx()) - .unwrap() - .is_null() + !fsl.execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap() + .is_null() ); assert!( - fsl.execute_scalar(2, &mut array_session().create_execution_ctx()) - .unwrap() - .is_null() + fsl.execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap() + .is_null() ); assert!( - !fsl.execute_scalar(3, &mut array_session().create_execution_ctx()) - .unwrap() - .is_null() + !fsl.execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap() + .is_null() ); } } @@ -338,31 +407,46 @@ fn test_mixed_nullability_patterns() { // List 0: valid with [Some(1), None]. let list0 = fsl - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!list0.is_null()); // List 1: null. let list1 = fsl - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(list1.is_null()); // List 2: valid with [Some(5), Some(6)]. let list2 = fsl - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!list2.is_null()); // List 3: valid with [Some(7), None]. let list3 = fsl - .execute_scalar(3, &mut array_session().create_execution_ctx()) + .execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!list3.is_null()); // List 4: valid with [None, Some(10)]. let list4 = fsl - .execute_scalar(4, &mut array_session().create_execution_ctx()) + .execute_scalar( + 4, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!list4.is_null()); } diff --git a/vortex-array/src/arrays/fixed_size_list/tests/take.rs b/vortex-array/src/arrays/fixed_size_list/tests/take.rs index e75139e5fc8..3bc874d8b74 100644 --- a/vortex-array/src/arrays/fixed_size_list/tests/take.rs +++ b/vortex-array/src/arrays/fixed_size_list/tests/take.rs @@ -12,13 +12,13 @@ use super::common::create_single_element_fsl; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::FixedSizeListArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::FixedSizeListBuilder; use crate::compute::conformance::take::test_take_conformance; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -40,7 +40,7 @@ fn test_take_fsl_conformance(#[case] fsl: FixedSizeListArray) { #[test] fn test_take_basic_smoke_test() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -98,7 +98,10 @@ fn test_take_degenerate_lists( for (i, expected_null) in expected_nulls.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null(), *expected_null @@ -108,7 +111,7 @@ fn test_take_degenerate_lists( #[test] fn test_take_large_list_size() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = buffer![0i32..300].into_array(); let fsl = FixedSizeListArray::new(elements, 100, Validity::NonNullable, 3); @@ -123,7 +126,7 @@ fn test_take_large_list_size() { #[test] fn test_take_fsl_with_null_indices_preserves_elements() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array(); let fsl = FixedSizeListArray::new(elements.into_array(), 2, Validity::NonNullable, 3); @@ -170,7 +173,7 @@ fn test_element_index_overflow( #[case] indices: ArrayRef, #[case] expected: FixedSizeListArray, ) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = fsl.take(indices).unwrap(); assert_arrays_eq!(result, expected, &mut ctx); } @@ -235,7 +238,10 @@ fn test_take_nullable_arrays_fsl_specific( for (i, expected_null) in expected_nulls.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null(), *expected_null diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs index 47dc3ec0651..e6da8cd0719 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/kernel.rs @@ -1,19 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::FixedSizeList; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel( Cast.id(), FixedSizeList, diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs index bfc172a5049..5cd7206daf1 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs +++ b/vortex-array/src/arrays/fixed_size_list/vtable/mod.rs @@ -11,6 +11,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::CachedId; use crate::ArrayEq; @@ -39,7 +40,7 @@ mod validity; /// A [`FixedSizeList`]-encoded Vortex array. pub type FixedSizeListArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index b480ac142fb..3464cdb26ab 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -450,10 +450,10 @@ mod tests { use super::*; use crate::Canonical; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::default_session_builder; /// Reference (oracle) implementation of the interleave spec, used only to validate the optimized /// [execute](super::execute) path. It is intentionally simple and slow: it pulls each output @@ -541,7 +541,7 @@ mod tests { InterleaveArray::try_new(values.clone(), array_indices.clone(), row_indices.clone())? .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let reference = interleave_reference(&values, &array_indices, &row_indices, &mut ctx)?; assert_arrays_eq!(interleaved, reference, &mut ctx); @@ -679,7 +679,7 @@ mod tests { InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)? .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let err = interleaved .execute::(&mut ctx) .err() @@ -700,7 +700,7 @@ mod tests { InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)? .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let err = interleaved .execute::(&mut ctx) .err() @@ -720,7 +720,7 @@ mod tests { let interleaved = InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices) .vortex_expect("primitive values should construct") .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); interleaved.execute::(&mut ctx).ok(); } } diff --git a/vortex-array/src/arrays/list/compute/cast.rs b/vortex-array/src/arrays/list/compute/cast.rs index 084d9613e75..67d09aca423 100644 --- a/vortex-array/src/arrays/list/compute/cast.rs +++ b/vortex-array/src/arrays/list/compute/cast.rs @@ -84,7 +84,8 @@ mod tests { use crate::dtype::PType; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn test_cast_list_success() { diff --git a/vortex-array/src/arrays/list/compute/kernels.rs b/vortex-array/src/arrays/list/compute/kernels.rs index 6af858dbf49..2b1f108aeed 100644 --- a/vortex-array/src/arrays/list/compute/kernels.rs +++ b/vortex-array/src/arrays/list/compute/kernels.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Dict; @@ -9,13 +9,13 @@ use crate::arrays::Filter; use crate::arrays::List; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Cast.id(), List, CastExecuteAdaptor(List)); kernels.register_execute_parent_kernel(Filter.id(), List, FilterExecuteAdaptor(List)); kernels.register_execute_parent_kernel(Dict.id(), List, TakeExecuteAdaptor(List)); diff --git a/vortex-array/src/arrays/list/compute/mod.rs b/vortex-array/src/arrays/list/compute/mod.rs index d8ddbd37583..4d77ac9373f 100644 --- a/vortex-array/src/arrays/list/compute/mod.rs +++ b/vortex-array/src/arrays/list/compute/mod.rs @@ -9,7 +9,7 @@ pub(crate) mod rules; mod slice; mod take; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { kernels::initialize(session); } @@ -20,13 +20,13 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::consistency::test_array_consistency; use crate::compute::conformance::filter::test_filter_conformance; use crate::compute::conformance::mask::test_mask_conformance; + use crate::default_session_builder; use crate::validity::Validity; #[test] @@ -82,7 +82,7 @@ mod tests { fn test_list_consistency(#[case] array: ListArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs index 2c6089e0c58..1e9e3478215 100644 --- a/vortex-array/src/arrays/list/compute/take.rs +++ b/vortex-array/src/arrays/list/compute/take.rs @@ -215,11 +215,11 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::PrimitiveArray; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType::I32; @@ -258,12 +258,18 @@ mod test { assert!( result - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::list( Arc::clone(&element_dtype), @@ -274,18 +280,27 @@ mod test { assert!( result - .is_invalid(1, &mut array_session().create_execution_ctx()) + .is_invalid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( result - .is_valid(2, &mut array_session().create_execution_ctx()) + .is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( result - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::list( Arc::clone(&element_dtype), @@ -296,12 +311,18 @@ mod test { assert!( result - .is_valid(3, &mut array_session().create_execution_ctx()) + .is_valid( + 3, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( result - .execute_scalar(3, &mut array_session().create_execution_ctx()) + .execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::list(element_dtype, vec![], Nullability::Nullable) ); @@ -361,12 +382,18 @@ mod test { assert!( result - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::list( Arc::clone(&element_dtype), @@ -377,12 +404,18 @@ mod test { assert!( result - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( result - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::list( Arc::clone(&element_dtype), @@ -393,12 +426,18 @@ mod test { assert!( result - .is_valid(2, &mut array_session().create_execution_ctx()) + .is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( result - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::list(element_dtype, vec![], Nullability::NonNullable) ); @@ -488,12 +527,18 @@ mod test { assert_eq!(result_view.len(), 2); assert!( result_view - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( result_view - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } @@ -518,17 +563,26 @@ mod test { assert_eq!(result_view.len(), 3); assert!( result_view - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( result_view - .is_invalid(1, &mut array_session().create_execution_ctx()) + .is_invalid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( result_view - .is_valid(2, &mut array_session().create_execution_ctx()) + .is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index 422a9523ff9..553375aee70 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -12,7 +12,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::List; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { compute::initialize(session); } diff --git a/vortex-array/src/arrays/list/tests.rs b/vortex-array/src/arrays/list/tests.rs index 906f4e03588..e3fe16862f5 100644 --- a/vortex-array/src/arrays/list/tests.rs +++ b/vortex-array/src/arrays/list/tests.rs @@ -15,13 +15,13 @@ use super::*; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::FilterArray; use crate::arrays::List; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::ListBuilder; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType::I32; @@ -29,7 +29,7 @@ use crate::scalar::Scalar; use crate::validity::Validity; /// A shared session for `List` tests, used to create execution contexts. -static SESSION: LazyLock = LazyLock::new(array_session); +static SESSION: LazyLock = LazyLock::new(|| default_session_builder().build()); #[test] fn test_empty_list_array() { diff --git a/vortex-array/src/arrays/listview/compute/zip.rs b/vortex-array/src/arrays/listview/compute/zip.rs index 46537f1fafe..34defd4f5e1 100644 --- a/vortex-array/src/arrays/listview/compute/zip.rs +++ b/vortex-array/src/arrays/listview/compute/zip.rs @@ -171,7 +171,6 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; @@ -181,6 +180,7 @@ mod tests { use crate::arrays::listview::ListViewArrayExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -216,7 +216,7 @@ mod tests { ); let mask = Mask::from_iter([true, false, true]); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -256,7 +256,7 @@ mod tests { // true -> if_true, false -> if_false let mask = Mask::from_iter([false, true, true]); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -293,7 +293,7 @@ mod tests { ); let mask = Mask::from_iter([true, true, false]); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -336,7 +336,7 @@ mod tests { ); let mask = Mask::from_iter([true, false]); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? diff --git a/vortex-array/src/arrays/listview/mod.rs b/vortex-array/src/arrays/listview/mod.rs index 80b2540a897..2f5a68be2ce 100644 --- a/vortex-array/src/arrays/listview/mod.rs +++ b/vortex-array/src/arrays/listview/mod.rs @@ -12,7 +12,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::ListView; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs index 3d4d6db284d..032745ec0c6 100644 --- a/vortex-array/src/arrays/listview/tests/basic.rs +++ b/vortex-array/src/arrays/listview/tests/basic.rs @@ -9,7 +9,6 @@ use vortex_error::VortexResult; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::ListArray; @@ -18,6 +17,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; use crate::arrays::listview::list_view_from_list; use crate::assert_arrays_eq; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -26,7 +26,7 @@ use crate::validity::Validity; #[test] fn test_basic_listview_comprehensive() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Comprehensive test for basic ListView functionality including scalar_at. // Logical lists: [[1,2,3], [4,5], [6,7,8,9]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -57,7 +57,10 @@ fn test_basic_listview_comprehensive() { // Test scalar_at which returns entire lists as Scalar values. let first_scalar = listview - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert_eq!( first_scalar, @@ -83,7 +86,7 @@ fn test_basic_listview_comprehensive() { #[test] fn test_out_of_order_offsets() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // ListView-specific: Tests that offsets can be non-sequential and out-of-order. // Logical lists: [[7,8,9], [1,2,3], [4,5,6]] let elements = buffer![1i32, 2, 3, 4, 5, 6, 7, 8, 9].into_array(); @@ -135,7 +138,7 @@ fn test_from_list_array() -> VortexResult<()> { let validity = Validity::from_iter([true, false, true]); let list_array = ListArray::try_new(elements, offsets, validity)?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let list_view = list_view_from_list(list_array, &mut ctx)?; assert_eq!(list_view.len(), 3); @@ -148,9 +151,18 @@ fn test_from_list_array() -> VortexResult<()> { ); // Check validity is preserved. - assert!(list_view.is_valid(0, &mut array_session().create_execution_ctx())?); - assert!(list_view.is_invalid(1, &mut array_session().create_execution_ctx())?); - assert!(list_view.is_valid(2, &mut array_session().create_execution_ctx())?); + assert!(list_view.is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + )?); + assert!(list_view.is_invalid( + 1, + &mut default_session_builder().build().create_execution_ctx() + )?); + assert!(list_view.is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + )?); // Check third list. assert_arrays_eq!( @@ -167,7 +179,7 @@ fn test_from_list_array() -> VortexResult<()> { #[case::constant_offsets(false, true)] // Varying sizes, constant offsets #[case::both_constant(true, true)] // Both constant fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_offsets: bool) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Logical lists vary by case: // - constant_sizes: [[1,2,3], [4,5,6], [7,8,9]] (size 3 each, varying offsets) // - constant_offsets: [[1,2,3], [1,2], [1]] (all start at 0, varying sizes) @@ -213,7 +225,10 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o listview .list_elements_at(0) .unwrap() - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); @@ -221,7 +236,10 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o listview .list_elements_at(1) .unwrap() - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); @@ -229,7 +247,10 @@ fn test_listview_with_constant_arrays(#[case] const_sizes: bool, #[case] const_o listview .list_elements_at(2) .unwrap() - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); diff --git a/vortex-array/src/arrays/listview/tests/common.rs b/vortex-array/src/arrays/listview/tests/common.rs index 1398a8defe6..f89045cef6e 100644 --- a/vortex-array/src/arrays/listview/tests/common.rs +++ b/vortex-array/src/arrays/listview/tests/common.rs @@ -16,7 +16,8 @@ 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(crate::array_session); +pub static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); /// Creates a basic ListView for testing: [[0,1,2], [3,4], [5,6], [7,8,9]] pub fn create_basic_listview() -> ListViewArray { diff --git a/vortex-array/src/arrays/listview/tests/density.rs b/vortex-array/src/arrays/listview/tests/density.rs index 8538a9fe20c..cc3b4526ca5 100644 --- a/vortex-array/src/arrays/listview/tests/density.rs +++ b/vortex-array/src/arrays/listview/tests/density.rs @@ -27,7 +27,7 @@ use crate::validity::Validity; const EPS: f32 = 1e-6; fn test_execution_ctx() -> ExecutionCtx { - let session = crate::array_session(); + let session = crate::default_session_builder().build(); session.create_execution_ctx() } diff --git a/vortex-array/src/arrays/listview/tests/filter.rs b/vortex-array/src/arrays/listview/tests/filter.rs index 25cc0de4aa8..5a2af47ff67 100644 --- a/vortex-array/src/arrays/listview/tests/filter.rs +++ b/vortex-array/src/arrays/listview/tests/filter.rs @@ -23,7 +23,8 @@ use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; use crate::validity::Validity; -static SESSION: LazyLock = LazyLock::new(crate::array_session); +static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); // Conformance tests for common filter scenarios. #[rstest] diff --git a/vortex-array/src/arrays/listview/tests/nested.rs b/vortex-array/src/arrays/listview/tests/nested.rs index 4b2676f4491..79e767cd6e3 100644 --- a/vortex-array/src/arrays/listview/tests/nested.rs +++ b/vortex-array/src/arrays/listview/tests/nested.rs @@ -5,12 +5,12 @@ use vortex_buffer::buffer; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::ListView; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::listview::ListViewArrayExt; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; @@ -74,7 +74,10 @@ fn test_listview_of_listview_with_overlapping() { // inner[0] should be [1, 2, 3]. assert_eq!( inner0 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_primitive() .as_::() @@ -83,7 +86,10 @@ fn test_listview_of_listview_with_overlapping() { ); assert_eq!( inner0 - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_primitive() .as_::() @@ -94,7 +100,10 @@ fn test_listview_of_listview_with_overlapping() { // inner[1] should be [3, 4, 5] - shares element 3 with inner[0]. assert_eq!( inner1 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_primitive() .as_::() @@ -103,7 +112,10 @@ fn test_listview_of_listview_with_overlapping() { ); assert_eq!( inner1 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_primitive() .as_::() @@ -287,7 +299,10 @@ fn test_listview_zero_and_overlapping() { assert_eq!(inner1.len(), 3); // [1, 2, 3] assert_eq!( inner1 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_primitive() .as_::() @@ -306,7 +321,10 @@ fn test_listview_zero_and_overlapping() { assert_eq!(inner3.len(), 3); // [2, 3, 4] assert_eq!( inner3 - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_primitive() .as_::() @@ -385,7 +403,10 @@ fn test_listview_of_struct_with_nulls() { // The middle element (struct[2]) should be null. assert!( list1 - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); diff --git a/vortex-array/src/arrays/listview/tests/nullability.rs b/vortex-array/src/arrays/listview/tests/nullability.rs index 4889b73de21..d8f46952a8b 100644 --- a/vortex-array/src/arrays/listview/tests/nullability.rs +++ b/vortex-array/src/arrays/listview/tests/nullability.rs @@ -8,11 +8,11 @@ use vortex_buffer::buffer; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::arrays::listview::ListViewArrayExt; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -38,17 +38,26 @@ fn test_nullable_listview_comprehensive() { // Check validity. assert!( listview - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( listview - .is_invalid(1, &mut array_session().create_execution_ctx()) + .is_invalid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( listview - .is_valid(2, &mut array_session().create_execution_ctx()) + .is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); @@ -60,7 +69,10 @@ fn test_nullable_listview_comprehensive() { // Test scalar_at with nulls. let first = listview - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!first.is_null()); assert_eq!( @@ -73,12 +85,18 @@ fn test_nullable_listview_comprehensive() { ); let second = listview - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(second.is_null()); let third = listview - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(!third.is_null()); assert_eq!( @@ -95,13 +113,19 @@ fn test_nullable_listview_comprehensive() { assert_eq!(null_list_data.len(), 2); assert_eq!( null_list_data - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 3i32.into() ); assert_eq!( null_list_data - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 4i32.into() ); @@ -123,7 +147,10 @@ fn test_nullable_patterns(#[case] validity: Validity, #[case] expected_validity: for (i, &expected) in expected_validity.iter().enumerate() { assert_eq!( listview - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), expected ); @@ -150,19 +177,28 @@ fn test_nullable_elements() { assert_eq!(first_list.len(), 2); assert!( !first_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert_eq!( first_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1i32.into() ); assert!( first_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); @@ -171,19 +207,28 @@ fn test_nullable_elements() { let second_list = listview.list_elements_at(1).unwrap(); assert!( !second_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert_eq!( second_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 3i32.into() ); assert!( second_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); @@ -192,25 +237,37 @@ fn test_nullable_elements() { let third_list = listview.list_elements_at(2).unwrap(); assert!( !third_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert_eq!( third_list - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 5i32.into() ); assert!( !third_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert_eq!( third_list - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 6i32.into() ); diff --git a/vortex-array/src/arrays/listview/tests/operations.rs b/vortex-array/src/arrays/listview/tests/operations.rs index f49680a854e..5ecc10db14b 100644 --- a/vortex-array/src/arrays/listview/tests/operations.rs +++ b/vortex-array/src/arrays/listview/tests/operations.rs @@ -17,7 +17,6 @@ use crate::IntoArray; use crate::ToCanonical as _; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::ListView; @@ -27,6 +26,7 @@ use crate::arrays::listview::ListViewArrayExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::mask::test_mask_conformance; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -64,10 +64,16 @@ fn test_slice_comprehensive() { assert_eq!( full_list .array() - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), listview - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), "Mismatch at index {}", i @@ -84,7 +90,7 @@ fn test_slice_comprehensive() { #[test] fn test_slice_out_of_order() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // ListView-specific: Test slicing with out-of-order offsets. // Logical lists: [[70,80], [10,20,30], [40,50,60], [90], [30]] let elements = buffer![10i32, 20, 30, 40, 50, 60, 70, 80, 90].into_array(); @@ -163,13 +169,19 @@ fn test_slice_with_nulls() { assert!( sliced_list .array() - .is_invalid(0, &mut array_session().create_execution_ctx()) + .is_invalid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Original index 1 was null. assert!( sliced_list .array() - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Original index 2 was valid. @@ -299,12 +311,18 @@ fn test_cast_with_nulls() { let result_list = result.to_listview(); assert!( result_list - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( result_list - .is_invalid(1, &mut array_session().create_execution_ctx()) + .is_invalid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } @@ -394,7 +412,7 @@ fn test_cast_large_dataset() { #[test] fn test_zip_widens_false_element_nullability() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // [[1, 2], [3], [4]] let if_true = ListViewArray::new( buffer![1i32, 2, 3, 4].into_array(), @@ -416,7 +434,7 @@ fn test_zip_widens_false_element_nullability() -> VortexResult<()> { let result = mask .into_array() .zip(if_true, if_false)? - .execute::(&mut array_session().create_execution_ctx())?; + .execute::(&mut default_session_builder().build().create_execution_ctx())?; assert!(result.is::()); assert_eq!( result.dtype(), @@ -440,7 +458,7 @@ fn test_zip_widens_false_element_nullability() -> VortexResult<()> { #[test] fn test_zip_widens_true_element_nullability() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // [[1, null], [3], [4]] let if_true = ListViewArray::new( PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4)]).into_array(), @@ -462,7 +480,7 @@ fn test_zip_widens_true_element_nullability() -> VortexResult<()> { let result = mask .into_array() .zip(if_true, if_false)? - .execute::(&mut array_session().create_execution_ctx())?; + .execute::(&mut default_session_builder().build().create_execution_ctx())?; assert!(result.is::()); assert_eq!( result.dtype(), @@ -540,7 +558,7 @@ fn test_is_constant_basic( ) .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(is_constant(&listview, &mut ctx).unwrap(), expected); } @@ -559,7 +577,7 @@ fn test_constant_with_constant_elements() { .into_array(); // All lists contain [42, 42] so should be constant. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!(is_constant(&listview, &mut ctx).unwrap()); } @@ -583,7 +601,7 @@ fn test_constant_with_nulls() { .with_zero_copy_to_list(true) } .into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!(!is_constant(&listview_mixed, &mut ctx).unwrap()); // Case 2: All nulls - should be constant. @@ -593,7 +611,7 @@ fn test_constant_with_nulls() { .with_zero_copy_to_list(true) } .into_array(); - let mut ctx2 = array_session().create_execution_ctx(); + let mut ctx2 = default_session_builder().build().create_execution_ctx(); assert!(is_constant(&listview_all_null, &mut ctx2).unwrap()); } @@ -608,7 +626,7 @@ fn test_constant_repeated_same_lists() { let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable).into_array(); // All lists are [10, 20, 30] so should be constant. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!(is_constant(&listview, &mut ctx).unwrap()); } @@ -650,22 +668,34 @@ fn test_mask_preserves_structure() { // Check validity: true in selection means null. assert!( !result_list - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Masked. assert!( result_list - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Not masked. assert!( !result_list - .is_valid(2, &mut array_session().create_execution_ctx()) + .is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Masked. assert!( !result_list - .is_valid(3, &mut array_session().create_execution_ctx()) + .is_valid( + 3, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Masked. @@ -704,17 +734,26 @@ fn test_mask_with_existing_nulls() { // Check combined validity: assert!( result_list - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Was valid, mask is false -> valid. assert!( !result_list - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Was invalid, mask is true -> invalid. assert!( !result_list - .is_valid(2, &mut array_session().create_execution_ctx()) + .is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Was valid, mask is true -> invalid. } @@ -737,17 +776,26 @@ fn test_mask_with_gaps() { assert_eq!(result_list.len(), 3); assert!( !result_list - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Masked assert!( result_list - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Not masked assert!( result_list - .is_valid(2, &mut array_session().create_execution_ctx()) + .is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Not masked @@ -782,17 +830,26 @@ fn test_mask_constant_arrays() { assert_eq!(result_list.len(), 3); assert!( result_list - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( !result_list - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); // Masked assert!( result_list - .is_valid(2, &mut array_session().create_execution_ctx()) + .is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); diff --git a/vortex-array/src/arrays/listview/tests/take.rs b/vortex-array/src/arrays/listview/tests/take.rs index 2a68f711d30..ae67380f7b2 100644 --- a/vortex-array/src/arrays/listview/tests/take.rs +++ b/vortex-array/src/arrays/listview/tests/take.rs @@ -22,7 +22,8 @@ use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; use crate::validity::Validity; -static SESSION: LazyLock = LazyLock::new(crate::array_session); +static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); // Conformance tests for common take scenarios. #[rstest] diff --git a/vortex-array/src/arrays/listview/vtable/kernel.rs b/vortex-array/src/arrays/listview/vtable/kernel.rs index 5621b23e79e..a7f89fc84de 100644 --- a/vortex-array/src/arrays/listview/vtable/kernel.rs +++ b/vortex-array/src/arrays/listview/vtable/kernel.rs @@ -1,18 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::arrays::ListView; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Cast.id(), ListView, CastExecuteAdaptor(ListView)); kernels.register_execute_parent_kernel(Zip.id(), ListView, ZipExecuteAdaptor(ListView)); } diff --git a/vortex-array/src/arrays/listview/vtable/mod.rs b/vortex-array/src/arrays/listview/vtable/mod.rs index ba5ae42b3ad..c5badcce005 100644 --- a/vortex-array/src/arrays/listview/vtable/mod.rs +++ b/vortex-array/src/arrays/listview/vtable/mod.rs @@ -12,6 +12,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::CachedId; use crate::ArrayEq; @@ -44,7 +45,7 @@ mod validity; /// A [`ListView`]-encoded Vortex array. pub type ListViewArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index 60b045d521d..6824a2d254f 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -11,9 +11,9 @@ use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::validity::Validity; @@ -51,7 +51,7 @@ fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> { let canonical = array .clone() .into_array() - .execute::(&mut array_session().create_execution_ctx())?; + .execute::(&mut default_session_builder().build().create_execution_ctx())?; assert_eq!(canonical.dtype(), array.dtype()); Ok(()) } @@ -67,35 +67,50 @@ fn test_masked_child_with_validity() { let prim = array.as_array().to_primitive(); // Positions where validity is false should be null in masked_child. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(prim.valid_count(&mut ctx).unwrap(), 3); assert!( - prim.is_valid(0, &mut array_session().create_execution_ctx()) - .unwrap() + prim.is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap() ); assert!( !prim - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( - prim.is_valid(2, &mut array_session().create_execution_ctx()) - .unwrap() + prim.is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap() ); assert!( !prim - .is_valid(3, &mut array_session().create_execution_ctx()) + .is_valid( + 3, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert!( - prim.is_valid(4, &mut array_session().create_execution_ctx()) - .unwrap() + prim.is_valid( + 4, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap() ); } #[test] fn test_masked_child_all_valid() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // When validity is AllValid, masked_child should invert to AllInvalid. let child = PrimitiveArray::from_iter([10i32, 20, 30]).into_array(); let array = MaskedArray::try_new(child, Validity::AllValid).unwrap(); @@ -103,7 +118,7 @@ fn test_masked_child_all_valid() { assert_eq!(array.len(), 3); assert_eq!( array - .valid_count(&mut array_session().create_execution_ctx()) + .valid_count(&mut default_session_builder().build().create_execution_ctx()) .unwrap(), 3 ); @@ -131,7 +146,7 @@ fn test_masked_child_preserves_length(#[case] validity: Validity) { assert_eq!(array.len(), len); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!( array .validity() diff --git a/vortex-array/src/arrays/masked/vtable/canonical.rs b/vortex-array/src/arrays/masked/vtable/canonical.rs index 32a5e334afa..e08bdc88277 100644 --- a/vortex-array/src/arrays/masked/vtable/canonical.rs +++ b/vortex-array/src/arrays/masked/vtable/canonical.rs @@ -9,9 +9,9 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::validity::Validity; @@ -41,7 +41,7 @@ mod tests { let canonical = array .clone() .into_array() - .execute::(&mut array_session().create_execution_ctx())?; + .execute::(&mut default_session_builder().build().create_execution_ctx())?; assert_eq!(canonical.dtype().nullability(), expected_nullability); assert_eq!(canonical.dtype(), array.dtype()); Ok(()) @@ -56,17 +56,32 @@ mod tests { let canonical = array .into_array() - .execute::(&mut array_session().create_execution_ctx())?; + .execute::(&mut default_session_builder().build().create_execution_ctx())?; let prim = canonical.into_primitive(); // Check that null positions match validity. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(prim.valid_count(&mut ctx)?, 3); - assert!(prim.is_valid(0, &mut array_session().create_execution_ctx())?); - assert!(!prim.is_valid(1, &mut array_session().create_execution_ctx())?); - assert!(prim.is_valid(2, &mut array_session().create_execution_ctx())?); - assert!(!prim.is_valid(3, &mut array_session().create_execution_ctx())?); - assert!(prim.is_valid(4, &mut array_session().create_execution_ctx())?); + assert!(prim.is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + )?); + assert!(!prim.is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + )?); + assert!(prim.is_valid( + 2, + &mut default_session_builder().build().create_execution_ctx() + )?); + assert!(!prim.is_valid( + 3, + &mut default_session_builder().build().create_execution_ctx() + )?); + assert!(prim.is_valid( + 4, + &mut default_session_builder().build().create_execution_ctx() + )?); Ok(()) } @@ -79,12 +94,12 @@ mod tests { let canonical = array .into_array() - .execute::(&mut array_session().create_execution_ctx())?; + .execute::(&mut default_session_builder().build().create_execution_ctx())?; assert_eq!(canonical.dtype().nullability(), Nullability::Nullable); assert_eq!( canonical .into_array() - .valid_count(&mut array_session().create_execution_ctx())?, + .valid_count(&mut default_session_builder().build().create_execution_ctx())?, 3 ); Ok(()) diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 70f7f566e41..68ce05929ee 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -207,10 +207,10 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Masked; use crate::arrays::MaskedArray; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -243,7 +243,11 @@ mod tests { let serialized = array .clone() .into_array() - .serialize(&ctx, &array_session(), &SerializeOptions::default()) + .serialize( + &ctx, + &default_session_builder().build(), + &SerializeOptions::default(), + ) .unwrap(); // Concat into a single buffer. @@ -259,7 +263,7 @@ mod tests { &dtype, len, &ReadContext::new(ctx.to_ids()), - &array_session(), + &default_session_builder().build(), ) .unwrap(); @@ -287,7 +291,7 @@ mod tests { assert_eq!(array.dtype().nullability(), Nullability::Nullable); // Execute the array. This should produce a Canonical with Nullable dtype. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result: Canonical = array.into_array().execute(&mut ctx)?; assert_eq!( diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index d93518eb859..9e5d45846a7 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -120,7 +120,7 @@ pub mod variant; pub use variant::Variant; pub use variant::VariantArray; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { bool::initialize(session); chunked::initialize(session); decimal::initialize(session); diff --git a/vortex-array/src/arrays/null/compute/cast.rs b/vortex-array/src/arrays/null/compute/cast.rs index a2527e5fb73..278f7f025cc 100644 --- a/vortex-array/src/arrays/null/compute/cast.rs +++ b/vortex-array/src/arrays/null/compute/cast.rs @@ -33,10 +33,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::NullArray; use crate::builtins::ArrayBuiltins; use crate::compute::conformance::cast::test_cast_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -68,7 +68,10 @@ mod tests { for i in 0..5 { assert!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); diff --git a/vortex-array/src/arrays/null/compute/mod.rs b/vortex-array/src/arrays/null/compute/mod.rs index 5e756da27d0..578f2e8a5cf 100644 --- a/vortex-array/src/arrays/null/compute/mod.rs +++ b/vortex-array/src/arrays/null/compute/mod.rs @@ -18,12 +18,12 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::NullArray; use crate::compute::conformance::consistency::test_array_consistency; use crate::compute::conformance::filter::test_filter_conformance; use crate::compute::conformance::mask::test_mask_conformance; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::dtype::DType; #[test] @@ -40,7 +40,7 @@ mod test { .unwrap() .execute_mask( sliced_arr.len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap(), Mask::AllFalse(4) @@ -62,7 +62,10 @@ mod test { taken_arr .validity() .unwrap() - .execute_mask(taken_arr.len(), &mut array_session().create_execution_ctx()) + .execute_mask( + taken_arr.len(), + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Mask::AllFalse(5) )); @@ -73,7 +76,10 @@ mod test { let nulls = NullArray::new(10); let scalar = nulls - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!(scalar.is_null()); assert_eq!(scalar.dtype().clone(), DType::Null); @@ -109,7 +115,7 @@ mod test { fn test_null_consistency(#[case] array: NullArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index 78fe191b6a0..05d08f8f1b6 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -123,7 +123,7 @@ impl VTable for Null { /// ``` /// # fn main() -> vortex_error::VortexResult<()> { /// use vortex_array::arrays::NullArray; -/// use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +/// use vortex_array::{IntoArray, VortexSessionExecute, default_session_builder}; /// /// // Create a null array with 5 elements /// let array = NullArray::new(5); @@ -133,7 +133,7 @@ impl VTable for Null { /// assert_eq!(sliced.len(), 2); /// /// // All elements are null -/// let mut ctx = array_session().create_execution_ctx(); +/// let mut ctx = default_session_builder().build().create_execution_ctx(); /// let scalar = array.execute_scalar(0, &mut ctx).unwrap(); /// assert!(scalar.is_null()); /// # Ok(()) diff --git a/vortex-array/src/arrays/patched/compute/compare.rs b/vortex-array/src/arrays/patched/compute/compare.rs index a9262b6621e..d67cfe1d878 100644 --- a/vortex-array/src/arrays/patched/compute/compare.rs +++ b/vortex-array/src/arrays/patched/compute/compare.rs @@ -188,7 +188,9 @@ mod tests { ) .unwrap(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx) .unwrap() @@ -221,7 +223,9 @@ mod tests { ) .unwrap(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx).unwrap(); // Slice the array so that the first patch should be skipped. @@ -258,7 +262,9 @@ mod tests { None, )?; - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx)? .into_array() .try_downcast::() @@ -292,7 +298,9 @@ mod tests { None, )?; - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let lhs = Patched::from_array_and_patches(lhs, &patches, &mut ctx)? .into_array() .try_downcast::() diff --git a/vortex-array/src/arrays/patched/compute/filter.rs b/vortex-array/src/arrays/patched/compute/filter.rs index 8cf0c186fb6..a52b62220c9 100644 --- a/vortex-array/src/arrays/patched/compute/filter.rs +++ b/vortex-array/src/arrays/patched/compute/filter.rs @@ -78,7 +78,9 @@ mod tests { #[test] fn test_filter_noop() -> VortexResult<()> { // Filter that doesn't prune any chunks (all data fits in one chunk). - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let array = buffer![u16::MIN; 5].into_array(); let patched_indices = buffer![3u8, 4].into_array(); @@ -105,7 +107,9 @@ mod tests { #[test] fn test_filter_with_offset() -> VortexResult<()> { // Test filtering where offset > 0. - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![5u16, 1030].into_array(); @@ -133,7 +137,9 @@ mod tests { #[test] fn test_filter_basic() -> VortexResult<()> { // Basic test: filter with mask that crosses boundaries. - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![1024u16, 1025].into_array(); @@ -160,7 +166,9 @@ mod tests { #[test] fn test_filter_complex() -> VortexResult<()> { // Filter with mask that crosses boundaries, with patches offset. - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let array = buffer![u16::MIN; 4096].into_array(); let patched_indices = buffer![1024u16, 1025].into_array(); @@ -187,7 +195,9 @@ mod tests { #[test] fn test_filter_sliced() -> VortexResult<()> { // Test filter on a sliced PatchedArray to exercise codepath where offset > 0. - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); // Create a larger array (6 chunks) so we can slice and still have room // for the filter to prune chunks. @@ -225,7 +235,9 @@ mod tests { fn test_filter_with_offset_nonuniform() -> VortexResult<()> { // Test filtering with offset > 0 using non-uniform base values. // This catches slice_chunks bugs where inner coordinates are miscalculated. - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); // Use non-uniform values so that incorrect slicing is detectable. let base_values: Vec = (0u16..4096).collect(); @@ -262,7 +274,9 @@ mod tests { fn test_filter_with_offset_last_chunk() -> VortexResult<()> { // Test filtering with offset > 0 where the mask touches the last chunk. // This ensures we don't accidentally slice past the end of the array or mask. - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); // Create a 6-chunk array (6144 elements). let array = buffer![u16::MIN; 6144].into_array(); diff --git a/vortex-array/src/arrays/patched/compute/take.rs b/vortex-array/src/arrays/patched/compute/take.rs index 993428808d4..f48528810a5 100644 --- a/vortex-array/src/arrays/patched/compute/take.rs +++ b/vortex-array/src/arrays/patched/compute/take.rs @@ -138,10 +138,10 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Patched; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::patches::Patches; fn make_patched_array( @@ -169,7 +169,7 @@ mod tests { #[test] fn test_take_basic() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -186,7 +186,7 @@ mod tests { #[test] fn test_take_sliced() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = make_patched_array(&[0; 10], &[1, 3], &[100, 200], 2..10)?; let indices = buffer![0u32, 1, 2, 3, 7].into_array(); @@ -201,7 +201,7 @@ mod tests { #[test] fn test_take_out_of_order() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at indices [1, 3] with values [10, 30] let array = make_patched_array(&[0; 5], &[1, 3], &[10, 30], 0..5)?; @@ -218,7 +218,7 @@ mod tests { #[test] fn test_take_duplicates() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Array with base values [0, 0, 0, 0, 0] patched at index [2] with value [99] let array = make_patched_array(&[0; 5], &[2], &[99], 0..5)?; @@ -239,7 +239,7 @@ mod tests { #[test] fn test_take_with_null_indices() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use crate::arrays::BoolArray; use crate::validity::Validity; diff --git a/vortex-array/src/arrays/patched/mod.rs b/vortex-array/src/arrays/patched/mod.rs index a88628e5587..0f5dfbbc3c9 100644 --- a/vortex-array/src/arrays/patched/mod.rs +++ b/vortex-array/src/arrays/patched/mod.rs @@ -78,7 +78,7 @@ pub use array::*; use vortex_buffer::ByteBuffer; pub use vtable::*; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/patched/vtable/kernels.rs b/vortex-array/src/arrays/patched/vtable/kernels.rs index 51f53d15625..9c23b7a76e3 100644 --- a/vortex-array/src/arrays/patched/vtable/kernels.rs +++ b/vortex-array/src/arrays/patched/vtable/kernels.rs @@ -1,19 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::Patched; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Binary.id(), Patched, CompareExecuteAdaptor(Patched)); kernels.register_execute_parent_kernel(Dict.id(), Patched, TakeExecuteAdaptor(Patched)); } diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index dbcb4c220dc..dbc9d1d6931 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -17,6 +17,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::CachedId; use crate::ArrayRef; @@ -54,7 +55,7 @@ use crate::serde::ArrayChildren; /// A [`Patched`]-encoded Vortex array. pub type PatchedArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernels::initialize(session); } @@ -355,7 +356,6 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Patched; use crate::arrays::PatchedArray; use crate::arrays::PrimitiveArray; @@ -364,6 +364,7 @@ mod tests { use crate::arrays::patched::PatchedSlotsView; use crate::assert_arrays_eq; use crate::builders::builder_with_capacity; + use crate::default_session_builder; use crate::patches::Patches; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -382,7 +383,7 @@ mod tests { ) .unwrap(); - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -415,7 +416,7 @@ mod tests { ) .unwrap(); - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -448,7 +449,7 @@ mod tests { ) .unwrap(); - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -481,7 +482,7 @@ mod tests { ) .unwrap(); - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -518,7 +519,7 @@ mod tests { ) .unwrap(); - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array = Patched::from_array_and_patches(values, &patches, &mut ctx) @@ -562,7 +563,7 @@ mod tests { let patches = Patches::new(len, 0, indices, patch_vals, None)?; - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); Patched::from_array_and_patches(array, &patches, &mut ctx) @@ -583,7 +584,7 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - let session = array_session(); + let session = default_session_builder().build(); session.arrays().register(Patched); let ctx = ArrayContext::empty().with_registry(session.arrays().registry().clone()); @@ -635,7 +636,7 @@ mod tests { assert_eq!(array_ref.dtype(), new_array.dtype()); // Execute both and compare results - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let original_executed = array_ref.execute::(&mut ctx)?.into_primitive(); let new_executed = new_array.execute::(&mut ctx)?.into_primitive(); @@ -661,7 +662,7 @@ mod tests { let new_array = array_ref.with_slots(slots.into_slots())?; // Execute and verify the inner values changed (except at patch positions) - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let executed = new_array.execute::(&mut ctx)?.into_primitive(); // Expected: all 5s except indices 1, 2, 3 which are patched to 10, 20, 30 diff --git a/vortex-array/src/arrays/patched/vtable/operations.rs b/vortex-array/src/arrays/patched/vtable/operations.rs index 51dd1fc9e3c..1d22e8de70f 100644 --- a/vortex-array/src/arrays/patched/vtable/operations.rs +++ b/vortex-array/src/arrays/patched/vtable/operations.rs @@ -62,8 +62,8 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Patched; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::optimizer::ArrayOptimizer; use crate::patches::Patches; @@ -90,25 +90,37 @@ mod tests { assert_eq!( array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::primitive(0u16, Nullability::NonNullable) ); assert_eq!( array - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::primitive(1u16, Nullability::NonNullable) ); assert_eq!( array - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::primitive(1u16, Nullability::NonNullable) ); assert_eq!( array - .execute_scalar(3, &mut array_session().create_execution_ctx()) + .execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::primitive(1u16, Nullability::NonNullable) ); @@ -135,7 +147,10 @@ mod tests { for index in 0..array.len() { let value = array - .execute_scalar(index, &mut array_session().create_execution_ctx()) + .execute_scalar( + index, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); if [1, 2, 3].contains(&index) { @@ -173,14 +188,20 @@ mod tests { assert_eq!( array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 1u16.into() ); for index in 1..array.len() { assert_eq!( array - .execute_scalar(index, &mut array_session().create_execution_ctx()) + .execute_scalar( + index, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), 0u16.into() ); diff --git a/vortex-array/src/arrays/patched/vtable/slice.rs b/vortex-array/src/arrays/patched/vtable/slice.rs index 1f59c2a854d..9de53870172 100644 --- a/vortex-array/src/arrays/patched/vtable/slice.rs +++ b/vortex-array/src/arrays/patched/vtable/slice.rs @@ -80,7 +80,9 @@ mod tests { let patch_values = buffer![u16::MAX; 3].into_array(); let patches = Patches::new(512, 0, patch_indices, patch_values, None)?; - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx)?; @@ -126,7 +128,9 @@ mod tests { ) .unwrap(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let patched_array = Patched::from_array_and_patches(inner.into_array(), &patches, &mut ctx) .unwrap() @@ -158,7 +162,9 @@ mod tests { let patched_values = buffer![0u64, 1, 2, 3, 4, 5].into_array(); let patches = Patches::new(10_000, 0, patched_indices, patched_values, None).unwrap(); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let patched_array = Patched::from_array_and_patches(values, &patches, &mut ctx) .unwrap() diff --git a/vortex-array/src/arrays/primitive/array/cast.rs b/vortex-array/src/arrays/primitive/array/cast.rs index bdab4f0beb8..bfc9e889b38 100644 --- a/vortex-array/src/arrays/primitive/array/cast.rs +++ b/vortex-array/src/arrays/primitive/array/cast.rs @@ -52,7 +52,8 @@ mod tests { use crate::dtype::PType; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn test_downcast_all_invalid() { diff --git a/vortex-array/src/arrays/primitive/array/mod.rs b/vortex-array/src/arrays/primitive/array/mod.rs index 65ef3ca0490..2653a924945 100644 --- a/vortex-array/src/arrays/primitive/array/mod.rs +++ b/vortex-array/src/arrays/primitive/array/mod.rs @@ -71,7 +71,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// ``` /// # fn main() -> vortex_error::VortexResult<()> { /// use vortex_array::arrays::PrimitiveArray; -/// use vortex_array::{VortexSessionExecute, array_session}; +/// use vortex_array::{VortexSessionExecute, default_session_builder}; /// /// // Create from iterator using FromIterator impl /// let array: PrimitiveArray = [1i32, 2, 3, 4, 5].into_iter().collect(); @@ -80,7 +80,7 @@ pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// let sliced = array.slice(1..3)?; /// /// // Access individual values -/// let mut ctx = array_session().create_execution_ctx(); +/// let mut ctx = default_session_builder().build().create_execution_ctx(); /// let value = sliced.execute_scalar(0, &mut ctx).unwrap(); /// assert_eq!(value, 2i32.into()); /// diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index d183552b66f..503b0715f47 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -138,8 +138,8 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::validity::Validity; /// Regression: patch_chunk must not OOB when chunk_offsets (chunk granularity) @@ -175,7 +175,7 @@ mod tests { #[test] fn patch_sliced() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let input = PrimitiveArray::new(buffer![2u32; 10], Validity::AllValid); let sliced = input.slice(2..8).unwrap(); #[expect(deprecated)] diff --git a/vortex-array/src/arrays/primitive/compute/cast.rs b/vortex-array/src/arrays/primitive/compute/cast.rs index 36b95b4c554..c0a6785b2b6 100644 --- a/vortex-array/src/arrays/primitive/compute/cast.rs +++ b/vortex-array/src/arrays/primitive/compute/cast.rs @@ -273,13 +273,13 @@ mod test { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] use crate::canonical::ToCanonical as _; use crate::compute::conformance::cast::test_cast_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -287,7 +287,7 @@ mod test { #[test] fn cast_u32_u8() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); // cast from u32 to u8 @@ -347,7 +347,7 @@ mod test { #[test] fn cast_u32_f32() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = buffer![0u32, 10, 200].into_array(); #[expect(deprecated)] let u8arr = arr.cast(PType::F32.into()).unwrap().to_primitive(); @@ -389,7 +389,7 @@ mod test { #[test] fn cast_with_invalid_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![-1i32, 0, 10], Validity::from_iter([false, true, true]), @@ -411,7 +411,7 @@ mod test { .unwrap() .execute_mask( p.as_ref().len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap(), Mask::from(BitBuffer::from(vec![false, true, true])) @@ -422,7 +422,7 @@ mod test { /// buffer without allocation (pointer identity). #[test] fn cast_same_width_int_reinterprets_buffer() -> vortex_error::VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let src = PrimitiveArray::from_iter([0u32, 10, 100]); let src_ptr = src.as_slice::().as_ptr(); @@ -468,7 +468,7 @@ mod test { /// not prevent the cast from succeeding. #[test] fn cast_same_width_int_nullable_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // The null position holds u32::MAX which doesn't fit in i32, but it's // masked as invalid so the cast should still succeed via reinterpret. let arr = PrimitiveArray::new( @@ -490,7 +490,7 @@ mod test { #[test] fn cast_u32_to_u8_with_out_of_range_nulls() -> vortex_error::VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![1000u32, 10u32, 42u32], Validity::from_iter([false, true, true]), diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index d73442a150e..c5557017112 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -52,19 +52,19 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::compute::fill_null::BoolArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] use crate::canonical::ToCanonical as _; + use crate::default_session_builder; use crate::scalar::Scalar; use crate::validity::Validity; #[test] fn fill_null_leading_none() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([None, Some(8u8), None, Some(10), None]); #[expect(deprecated)] let p = arr @@ -83,7 +83,7 @@ mod test { .unwrap() .execute_mask( p.as_ref().len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap() .all_true() @@ -92,7 +92,7 @@ mod test { #[test] fn fill_null_all_none() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::from_option_iter([Option::::None, None, None, None, None]); #[expect(deprecated)] @@ -112,7 +112,7 @@ mod test { .unwrap() .execute_mask( p.as_ref().len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap() .all_true() @@ -121,7 +121,7 @@ mod test { #[test] fn fill_null_nullable_non_null() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = PrimitiveArray::new( buffer![8u8, 10, 12, 14, 16], Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()), @@ -143,7 +143,7 @@ mod test { .unwrap() .execute_mask( p.as_ref().len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap() .all_true() @@ -152,7 +152,7 @@ mod test { #[test] fn fill_null_non_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = buffer![8u8, 10, 12, 14, 16].into_array(); #[expect(deprecated)] let p = arr.fill_null(Scalar::from(255u8)).unwrap().to_primitive(); @@ -167,7 +167,7 @@ mod test { .unwrap() .execute_mask( p.as_ref().len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap() .all_true() diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index 1ca4b17d3d5..5b9771ccc5b 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -16,10 +16,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::compute::conformance::binary_numeric::test_binary_numeric_array; use crate::compute::conformance::consistency::test_array_consistency; + use crate::default_session_builder; #[rstest] // Basic primitive arrays @@ -41,7 +41,7 @@ mod tests { fn test_primitive_consistency(#[case] array: PrimitiveArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } @@ -54,7 +54,7 @@ mod tests { fn test_primitive_binary_numeric(#[case] array: PrimitiveArray) { test_binary_numeric_array( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 9294659708e..36d7d2aa292 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -153,11 +153,11 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::compute::take::take_primitive_scalar; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::scalar::Scalar; use crate::validity::Validity; @@ -181,21 +181,30 @@ mod test { let actual = values.take(indices.into_array()).unwrap(); assert_eq!( actual - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("no fail"), Scalar::from(Some(1)) ); // position 3 is null assert_eq!( actual - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("no fail"), Scalar::null_native::() ); // the third index is null assert_eq!( actual - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("no fail"), Scalar::null_native::() ); @@ -223,15 +232,15 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::validity::Validity; #[test] fn take_null_index_skips_out_of_bounds_value() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = PrimitiveArray::from_iter([10i32, 20, 30]); let indices = PrimitiveArray::new( buffer![1u64, 3], diff --git a/vortex-array/src/arrays/primitive/compute/zip.rs b/vortex-array/src/arrays/primitive/compute/zip.rs index 83f03b10614..6380790de8a 100644 --- a/vortex-array/src/arrays/primitive/compute/zip.rs +++ b/vortex-array/src/arrays/primitive/compute/zip.rs @@ -153,11 +153,11 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; /// The branchless kernel must agree with the scalar reference across the chunk boundary (index /// 63/64) and the trailing remainder, for non-nullable inputs. @@ -170,7 +170,7 @@ mod tests { let bits: Vec = (0..len).map(|i| i.is_multiple_of(3) || i == 64).collect(); let mask = Mask::from_iter(bits.iter().copied()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? @@ -201,7 +201,7 @@ mod tests { let bits: Vec = (0..len).map(|i| i.is_multiple_of(2)).collect(); let mask = Mask::from_iter(bits.iter().copied()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask .into_array() .zip(if_true, if_false)? diff --git a/vortex-array/src/arrays/primitive/mod.rs b/vortex-array/src/arrays/primitive/mod.rs index e291ff521fb..1263dd87ae6 100644 --- a/vortex-array/src/arrays/primitive/mod.rs +++ b/vortex-array/src/arrays/primitive/mod.rs @@ -15,7 +15,7 @@ mod vtable; pub use compute::rules::PrimitiveMaskedValidityRule; pub use vtable::Primitive; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/primitive/vtable/kernel.rs b/vortex-array/src/arrays/primitive/vtable/kernel.rs index 6382ea73794..062e07ff4c0 100644 --- a/vortex-array/src/arrays/primitive/vtable/kernel.rs +++ b/vortex-array/src/arrays/primitive/vtable/kernel.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::Primitive; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::between::BetweenExecuteAdaptor; @@ -18,8 +18,8 @@ use crate::scalar_fn::fns::fill_null::FillNullExecuteAdaptor; use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel( Between.id(), Primitive, diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index b5eba2cb1a0..63fb70aa329 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -29,6 +29,7 @@ use std::hash::Hasher; use vortex_buffer::Alignment; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::CachedId; use crate::EqMode; @@ -41,7 +42,7 @@ use crate::hash::ArrayHash; /// A [`Primitive`]-encoded Vortex array. pub type PrimitiveArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } @@ -226,16 +227,16 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; use crate::validity::Validity; #[test] fn test_nullable_primitive_serde_roundtrip() { - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array = PrimitiveArray::new( buffer![1i32, 2, 3, 4], diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index c19e3e2288c..2aab1d91602 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -61,12 +61,12 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::ScalarFnArray; use crate::arrays::scalar_fn::ScalarFnArrayExt; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::scalar::Scalar; use crate::scalar_fn::TypedScalarFnInstance; use crate::scalar_fn::fns::binary::Binary; @@ -76,7 +76,7 @@ mod tests { #[test] fn test_scalar_fn_add() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = buffer![1i32, 2, 3].into_array(); let rhs = buffer![10i32, 20, 30].into_array(); @@ -87,7 +87,7 @@ mod tests { let result = scalar_fn_array .into_array() - .execute::(&mut array_session().create_execution_ctx())? + .execute::(&mut default_session_builder().build().create_execution_ctx())? .into_array(); let expected = buffer![11i32, 22, 33].into_array(); assert_arrays_eq!(result, expected, &mut ctx); @@ -131,7 +131,7 @@ mod tests { #[test] fn test_scalar_fn_mul() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = buffer![2i32, 3, 4].into_array(); let rhs = buffer![5i32, 6, 7].into_array(); @@ -140,7 +140,7 @@ mod tests { let result = scalar_fn_array .into_array() - .execute::(&mut array_session().create_execution_ctx())? + .execute::(&mut default_session_builder().build().create_execution_ctx())? .into_array(); let expected = buffer![10i32, 18, 28].into_array(); assert_arrays_eq!(result, expected, &mut ctx); @@ -150,7 +150,7 @@ mod tests { #[test] fn test_scalar_fn_with_nullable() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::AllValid).into_array(); let rhs = PrimitiveArray::new( buffer![10i32, 20, 30], @@ -163,7 +163,7 @@ mod tests { let result = scalar_fn_array .into_array() - .execute::(&mut array_session().create_execution_ctx())? + .execute::(&mut default_session_builder().build().create_execution_ctx())? .into_array(); let expected = PrimitiveArray::new( buffer![11i32, 0, 33], @@ -177,7 +177,7 @@ mod tests { #[test] fn test_scalar_fn_comparison() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = buffer![1i32, 5, 3].into_array(); let rhs = buffer![2i32, 5, 1].into_array(); @@ -186,7 +186,7 @@ mod tests { let result = scalar_fn_array .into_array() - .execute::(&mut array_session().create_execution_ctx())? + .execute::(&mut default_session_builder().build().create_execution_ctx())? .into_array(); let expected = BoolArray::from_iter([false, true, false]).into_array(); assert_arrays_eq!(result, expected, &mut ctx); diff --git a/vortex-array/src/arrays/shared/tests.rs b/vortex-array/src/arrays/shared/tests.rs index d1a33ed3e20..f13ec8810d6 100644 --- a/vortex-array/src/arrays/shared/tests.rs +++ b/vortex-array/src/arrays/shared/tests.rs @@ -19,7 +19,7 @@ fn shared_array_caches_on_canonicalize() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable).into_array(); let shared = SharedArray::new(array); - let session = crate::array_session(); + let session = crate::default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let first = shared.get_or_compute(|source| source.clone().execute::(&mut ctx))?; diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index ba716bee024..17cd802da6a 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -189,14 +189,14 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::assert_arrays_eq; + use crate::default_session_builder; #[test] fn test_slice_slice() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Slice(1..4, Slice(2..8, base)) combines to Slice(3..6, base) let arr = PrimitiveArray::from_iter(0i32..10).into_array(); let inner_slice = SliceArray::new(arr, 2..8).into_array(); diff --git a/vortex-array/src/arrays/struct_/array.rs b/vortex-array/src/arrays/struct_/array.rs index 4bea2e3641a..4141bc81244 100644 --- a/vortex-array/src/arrays/struct_/array.rs +++ b/vortex-array/src/arrays/struct_/array.rs @@ -57,7 +57,7 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// use vortex_array::arrays::{StructArray, BoolArray}; /// use vortex_array::validity::Validity; /// use vortex_array::dtype::FieldNames; -/// use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +/// use vortex_array::{IntoArray, VortexSessionExecute, default_session_builder}; /// use vortex_buffer::buffer; /// /// // Create struct with all non-null fields but struct-level nulls @@ -71,7 +71,7 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// 2, /// Validity::Array(BoolArray::from_iter([true, false]).into_array()), // row 1 is null /// ).unwrap(); -/// let mut ctx = array_session().create_execution_ctx(); +/// let mut ctx = default_session_builder().build().create_execution_ctx(); /// /// // Row 0 is valid - returns a struct scalar with field values /// let row0 = struct_array.execute_scalar(0, &mut ctx).unwrap(); @@ -92,7 +92,7 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// use vortex_array::arrays::struct_::StructArrayExt; /// use vortex_array::validity::Validity; /// use vortex_array::dtype::FieldNames; -/// use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +/// use vortex_array::{IntoArray, VortexSessionExecute, default_session_builder}; /// use vortex_buffer::buffer; /// /// // Create struct with duplicate "data" field names @@ -108,7 +108,7 @@ pub(super) const FIELDS_OFFSET: usize = 1; /// /// // field_by_name returns the FIRST "data" field /// let first_data = struct_array.unmasked_field_by_name("data").unwrap(); -/// let mut ctx = array_session().create_execution_ctx(); +/// let mut ctx = default_session_builder().build().create_execution_ctx(); /// assert_eq!(first_data.execute_scalar(0, &mut ctx).unwrap(), 1i32.into()); /// ``` /// diff --git a/vortex-array/src/arrays/struct_/compute/cast.rs b/vortex-array/src/arrays/struct_/compute/cast.rs index ac14983f0e2..a020eec6a13 100644 --- a/vortex-array/src/arrays/struct_/compute/cast.rs +++ b/vortex-array/src/arrays/struct_/compute/cast.rs @@ -4,7 +4,7 @@ use itertools::Itertools; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayRef; use crate::ArrayView; @@ -20,15 +20,13 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::StructFields; use crate::kernel::ExecuteParentKernel; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; -pub(crate) fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(Cast.id(), Struct, StructCastKernel); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + builder_kernels(session).register_execute_parent_kernel(Cast.id(), Struct, StructCastKernel); } #[derive(Debug)] @@ -158,7 +156,8 @@ mod tests { use crate::scalar_fn::fns::cast::Cast; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); fn null_struct_cast_execute_parent( child: &ArrayRef, @@ -214,7 +213,9 @@ mod tests { .try_new_array(source.len(), target.clone(), [source]) .unwrap(); let parent_id = cast.encoding_id(); - let session = VortexSession::empty().with_some(KernelSession::empty()); + let session = VortexSession::builder() + .with_some(KernelSession::empty()) + .build(); session.kernels().register_execute_parent( parent_id, child_id, diff --git a/vortex-array/src/arrays/struct_/compute/mod.rs b/vortex-array/src/arrays/struct_/compute/mod.rs index c1539a638bf..111dc47e2db 100644 --- a/vortex-array/src/arrays/struct_/compute/mod.rs +++ b/vortex-array/src/arrays/struct_/compute/mod.rs @@ -21,7 +21,6 @@ mod tests { use crate::IntoArray as _; use crate::VortexSessionExecute; use crate::aggregate_fn::fns::is_constant::is_constant; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; @@ -31,6 +30,7 @@ mod tests { use crate::compute::conformance::consistency::test_array_consistency; use crate::compute::conformance::mask::test_mask_conformance; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; @@ -40,7 +40,7 @@ mod tests { #[test] fn take_empty_struct() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let struct_arr = StructArray::try_new(FieldNames::empty(), vec![], 10, Validity::NonNullable).unwrap(); let indices = PrimitiveArray::from_option_iter([Some(1), None]); @@ -69,13 +69,13 @@ mod tests { let taken = struct_arr .take(indices.into_array()) .unwrap() - .execute::(&mut array_session().create_execution_ctx()) + .execute::(&mut default_session_builder().build().create_execution_ctx()) .unwrap(); assert_eq!(taken.len(), 1); assert!( taken .into_array() - .all_invalid(&mut array_session().create_execution_ctx()) + .all_invalid(&mut default_session_builder().build().create_execution_ctx()) .unwrap() ); } @@ -87,14 +87,14 @@ mod tests { let taken = arr .take(indices.into_array()) .unwrap() - .execute::(&mut array_session().create_execution_ctx()) + .execute::(&mut default_session_builder().build().create_execution_ctx()) .unwrap(); assert_eq!(taken.len(), 1); } #[test] fn take_field_struct() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let struct_arr = StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).into_array())]) .unwrap(); @@ -258,7 +258,7 @@ mod tests { #[test] fn test_empty_struct_is_constant() { let array = StructArray::new_fieldless_with_len(2); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = is_constant(&array.into_array(), &mut ctx) .vortex_expect("operation should succeed in test"); assert!(result); @@ -421,7 +421,7 @@ mod tests { fn test_struct_consistency(#[case] array: StructArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 01981803239..c2282895c8d 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -153,7 +153,8 @@ mod tests { use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); fn no_struct_cast_plugin( _child: &ArrayRef, @@ -264,7 +265,7 @@ mod tests { Struct.id(), &[no_struct_cast_plugin as ReduceParentFn], ); - let session = VortexSession::empty().with_some(kernels); + let session = VortexSession::builder().with_some(kernels).build(); let optimized = cast.optimize_ctx(&session).unwrap(); assert!(optimized.is::()); diff --git a/vortex-array/src/arrays/struct_/mod.rs b/vortex-array/src/arrays/struct_/mod.rs index da49411c3bc..1382f077234 100644 --- a/vortex-array/src/arrays/struct_/mod.rs +++ b/vortex-array/src/arrays/struct_/mod.rs @@ -10,7 +10,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::Struct; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { compute::cast::initialize(session); vtable::initialize(session); } diff --git a/vortex-array/src/arrays/struct_/tests.rs b/vortex-array/src/arrays/struct_/tests.rs index 3eb65e0c018..e5060a860bd 100644 --- a/vortex-array/src/arrays/struct_/tests.rs +++ b/vortex-array/src/arrays/struct_/tests.rs @@ -7,7 +7,6 @@ use vortex_error::VortexResult; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -15,6 +14,7 @@ use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; @@ -24,7 +24,7 @@ use crate::validity::Validity; #[test] fn test_project() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = VarBinArray::from_vec( vec!["a", "b", "c", "d", "e"], @@ -67,7 +67,7 @@ fn test_project() { #[test] fn test_remove_column() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::NonNullable); let ys = PrimitiveArray::new(buffer![4u64, 5, 6, 7, 8], Validity::NonNullable); @@ -113,7 +113,7 @@ fn test_remove_column() { #[test] fn test_duplicate_field_names() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test that StructArray allows duplicate field names and returns the first match let field1 = buffer![1i32, 2, 3].into_array(); let field2 = buffer![10i32, 20, 30].into_array(); @@ -161,12 +161,14 @@ fn test_uncompressed_size_in_bytes() -> VortexResult<()> { let canonical_size = struct_array .clone() .into_array() - .execute::(&mut array_session().create_execution_ctx())? + .execute::(&mut default_session_builder().build().create_execution_ctx())? .into_array() .nbytes(); let uncompressed_size = struct_array .statistics() - .compute_uncompressed_size_in_bytes(&mut array_session().create_execution_ctx()); + .compute_uncompressed_size_in_bytes( + &mut default_session_builder().build().create_execution_ctx(), + ); assert_eq!(canonical_size, 2); assert_eq!(uncompressed_size, Some(4000)); diff --git a/vortex-array/src/arrays/struct_/vtable/kernel.rs b/vortex-array/src/arrays/struct_/vtable/kernel.rs index aa7f648b935..546081380a9 100644 --- a/vortex-array/src/arrays/struct_/vtable/kernel.rs +++ b/vortex-array/src/arrays/struct_/vtable/kernel.rs @@ -1,16 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::arrays::Struct; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(Zip.id(), Struct, ZipExecuteAdaptor(Struct)); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + builder_kernels(session).register_execute_parent_kernel( + Zip.id(), + Struct, + ZipExecuteAdaptor(Struct), + ); } diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index df9b1170b02..0de4048c127 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -7,6 +7,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayRef; use crate::ExecutionCtx; @@ -36,7 +37,7 @@ use crate::array::ArrayId; /// A [`Struct`]-encoded Vortex array. pub type StructArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } diff --git a/vortex-array/src/arrays/varbin/builder.rs b/vortex-array/src/arrays/varbin/builder.rs index 404f574af52..cd1ba24ec2e 100644 --- a/vortex-array/src/arrays/varbin/builder.rs +++ b/vortex-array/src/arrays/varbin/builder.rs @@ -131,9 +131,9 @@ mod tests { use vortex_error::VortexResult; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::varbin::VarBinArrayExt; use crate::arrays::varbin::builder::VarBinBuilder; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; use crate::expr::stats::Precision; @@ -153,13 +153,19 @@ mod tests { assert_eq!(array.dtype().nullability(), Nullable); assert_eq!( array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::utf8("hello".to_string(), Nullable) ); assert!( array - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); diff --git a/vortex-array/src/arrays/varbin/compute/cast.rs b/vortex-array/src/arrays/varbin/compute/cast.rs index 5983886ea30..d225af8cd5b 100644 --- a/vortex-array/src/arrays/varbin/compute/cast.rs +++ b/vortex-array/src/arrays/varbin/compute/cast.rs @@ -82,7 +82,8 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[rstest] #[case( diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs index ede3c409d94..9f2de88409c 100644 --- a/vortex-array/src/arrays/varbin/compute/compare.rs +++ b/vortex-array/src/arrays/varbin/compute/compare.rs @@ -162,12 +162,12 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::bool::BoolArrayExt; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -200,7 +200,7 @@ mod test { .unwrap() .execute_mask( result.as_ref().len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap() .to_bit_buffer(), @@ -236,7 +236,7 @@ mod test { .unwrap() .execute_mask( result.as_ref().len(), - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .unwrap() .to_bit_buffer(), @@ -255,13 +255,13 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::VarBinArray; use crate::arrays::varbin::builder::VarBinBuilder; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -291,7 +291,7 @@ mod tests { /// [`CompareKernel`]: super::CompareKernel #[test] fn varbin_i64_offsets_compare_constant() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); @@ -312,7 +312,7 @@ mod tests { #[test] fn varbin_i64_offsets_compare_constant_binary() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = VarBinBuilder::::with_capacity(3); builder.append_value(b"abc"); builder.append_value(b"xyz"); diff --git a/vortex-array/src/arrays/varbin/compute/filter.rs b/vortex-array/src/arrays/varbin/compute/filter.rs index ab969cf3def..7f610ececa4 100644 --- a/vortex-array/src/arrays/varbin/compute/filter.rs +++ b/vortex-array/src/arrays/varbin/compute/filter.rs @@ -226,13 +226,13 @@ mod test { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::VarBinArray; use crate::arrays::varbin::compute::filter::filter_select_var_bin_by_index; use crate::arrays::varbin::compute::filter::filter_select_var_bin_by_slice; use crate::assert_arrays_eq; use crate::compute::conformance::filter::test_filter_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::Nullability::Nullable; @@ -240,7 +240,7 @@ mod test { #[test] fn filter_var_bin_test() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -257,7 +257,7 @@ mod test { #[test] fn filter_var_bin_slice_test() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arr = VarBinArray::from_vec( vec![ b"hello".as_slice(), @@ -282,7 +282,7 @@ mod test { #[test] fn filter_var_bin_slice_null() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bytes = [ b"one".as_slice(), b"two".as_slice(), @@ -319,7 +319,7 @@ mod test { #[test] fn filter_varbin_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bytes = [b"".as_slice(), b"two".as_slice(), b"two".as_slice()] .into_iter() .flat_map(|x| x.iter().cloned()) @@ -337,7 +337,7 @@ mod test { #[test] fn filter_varbin_all_null() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let offsets = buffer![0, 0, 0, 0].into_array(); let validity = Validity::Array(BoolArray::from_iter([false, false, false]).into_array()); let arr = VarBinArray::try_new( diff --git a/vortex-array/src/arrays/varbin/compute/mod.rs b/vortex-array/src/arrays/varbin/compute/mod.rs index 480c07a3031..45f8691971b 100644 --- a/vortex-array/src/arrays/varbin/compute/mod.rs +++ b/vortex-array/src/arrays/varbin/compute/mod.rs @@ -16,9 +16,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::VarBinArray; use crate::compute::conformance::consistency::test_array_consistency; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; @@ -64,7 +64,7 @@ mod tests { fn test_varbin_consistency(#[case] array: VarBinArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/varbin/compute/take.rs b/vortex-array/src/arrays/varbin/compute/take.rs index 74e0bd6af9d..fb265970b15 100644 --- a/vortex-array/src/arrays/varbin/compute/take.rs +++ b/vortex-array/src/arrays/varbin/compute/take.rs @@ -259,12 +259,12 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::compute::take::PrimitiveArray; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::validity::Validity; @@ -308,7 +308,7 @@ mod tests { #[test] fn test_take_overflow() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let scream = iter::once("a").cycle().take(128).collect::(); let bytes = ByteBuffer::copy_from(scream.as_bytes()); let offsets = buffer![0u8, 128u8].into_array(); diff --git a/vortex-array/src/arrays/varbin/mod.rs b/vortex-array/src/arrays/varbin/mod.rs index f0ccbe4585c..dde68e62ff7 100644 --- a/vortex-array/src/arrays/varbin/mod.rs +++ b/vortex-array/src/arrays/varbin/mod.rs @@ -12,7 +12,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::VarBin; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/varbin/tests.rs b/vortex-array/src/arrays/varbin/tests.rs index e2eb79fff54..f959ecbcd24 100644 --- a/vortex-array/src/arrays/varbin/tests.rs +++ b/vortex-array/src/arrays/varbin/tests.rs @@ -9,10 +9,10 @@ use vortex_buffer::buffer; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::validity::Validity; @@ -34,7 +34,7 @@ fn binary_array() -> ArrayRef { #[rstest] pub fn test_scalar_at(binary_array: ArrayRef) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_arrays_eq!( binary_array, VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]), @@ -44,7 +44,7 @@ pub fn test_scalar_at(binary_array: ArrayRef) { #[rstest] pub fn slice_array(binary_array: ArrayRef) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let binary_arr = binary_array.slice(1..2).unwrap(); assert_arrays_eq!( binary_arr, diff --git a/vortex-array/src/arrays/varbin/vtable/canonical.rs b/vortex-array/src/arrays/varbin/vtable/canonical.rs index 4865df5e1a3..154290ad87e 100644 --- a/vortex-array/src/arrays/varbin/vtable/canonical.rs +++ b/vortex-array/src/arrays/varbin/vtable/canonical.rs @@ -48,13 +48,13 @@ mod tests { use rstest::rstest; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::arrays::varbin::builder::VarBinBuilder; use crate::assert_arrays_eq; #[expect(deprecated)] use crate::canonical::ToCanonical as _; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; @@ -79,7 +79,10 @@ mod tests { assert!( !canonical - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); @@ -96,7 +99,7 @@ mod tests { #[case(DType::Utf8(Nullability::NonNullable))] #[case(DType::Binary(Nullability::NonNullable))] fn test_canonical_varbin_unsliced(#[case] dtype: DType) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let varbin = VarBinArray::from_iter_nonnull(["foo", "bar", "baz"], dtype.clone()); #[expect(deprecated)] let canonical = varbin.as_array().to_varbinview(); diff --git a/vortex-array/src/arrays/varbin/vtable/kernel.rs b/vortex-array/src/arrays/varbin/vtable/kernel.rs index 9e80abd1037..fe8fbc85894 100644 --- a/vortex-array/src/arrays/varbin/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbin/vtable/kernel.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Dict; @@ -9,15 +9,15 @@ use crate::arrays::Filter; use crate::arrays::VarBin; use crate::arrays::dict::TakeExecuteAdaptor; use crate::arrays::filter::FilterExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::binary::CompareExecuteAdaptor; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Cast.id(), VarBin, CastExecuteAdaptor(VarBin)); kernels.register_execute_parent_kernel(Binary.id(), VarBin, CompareExecuteAdaptor(VarBin)); kernels.register_execute_parent_kernel(Filter.id(), VarBin, FilterExecuteAdaptor(VarBin)); diff --git a/vortex-array/src/arrays/varbin/vtable/mod.rs b/vortex-array/src/arrays/varbin/vtable/mod.rs index 90dcf320427..00c43767e0a 100644 --- a/vortex-array/src/arrays/varbin/vtable/mod.rs +++ b/vortex-array/src/arrays/varbin/vtable/mod.rs @@ -37,6 +37,7 @@ mod validity; use canonical::varbin_to_canonical; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::EqMode; use crate::arrays::varbin::compute::rules::PARENT_RULES; @@ -46,7 +47,7 @@ use crate::hash::ArrayHash; /// A [`VarBin`]-encoded Vortex array. pub type VarBinArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } diff --git a/vortex-array/src/arrays/varbinview/compact.rs b/vortex-array/src/arrays/varbinview/compact.rs index fc3f3477ffc..e979a7e1a5c 100644 --- a/vortex-array/src/arrays/varbinview/compact.rs +++ b/vortex-array/src/arrays/varbinview/compact.rs @@ -202,15 +202,15 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; #[test] fn test_optimize_compacts_buffers() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create a VarBinViewArray with some long strings that will create multiple buffers let original = VarBinViewArray::from_iter_nullable_str([ Some("short"), @@ -228,7 +228,9 @@ mod tests { let indices = buffer![0u32, 4u32].into_array(); let taken = original.take(indices).unwrap(); let taken = taken - .execute::(&mut array_session().create_execution_ctx()) + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); // The taken array should still have the same number of buffers assert_eq!(taken.data_buffers().len(), original_buffers); @@ -251,7 +253,7 @@ mod tests { #[test] fn test_optimize_with_long_strings() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create strings that are definitely longer than 12 bytes let long_string_1 = "this is definitely a very long string that exceeds the inline limit"; let long_string_2 = "another extremely long string that also needs external buffer storage"; @@ -269,7 +271,9 @@ mod tests { let indices = buffer![0u32, 2u32].into_array(); let taken = original.take(indices).unwrap(); let taken_array = taken - .execute::(&mut array_session().create_execution_ctx()) + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); let optimized_array = taken_array.compact_with_threshold(1.0).unwrap(); @@ -287,7 +291,7 @@ mod tests { #[test] fn test_optimize_no_buffers() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create an array with only short strings (all inlined) let original = VarBinViewArray::from_iter_str(["a", "bb", "ccc", "dddd"]); @@ -304,7 +308,7 @@ mod tests { #[test] fn test_optimize_single_buffer() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create an array that naturally has only one buffer let str1 = "this is a long string that goes into a buffer"; let str2 = "another long string in the same buffer"; @@ -324,7 +328,7 @@ mod tests { #[test] fn test_selective_compaction_with_threshold_zero() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // threshold=0 should keep all buffers (no compaction) let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -338,7 +342,9 @@ mod tests { let indices = buffer![0u32].into_array(); let taken = original.take(indices).unwrap(); let taken = taken - .execute::(&mut array_session().create_execution_ctx()) + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); // Compact with threshold=0 (should not compact) let compacted = taken.compact_with_threshold(0.0).unwrap(); @@ -352,7 +358,7 @@ mod tests { #[test] fn test_selective_compaction_with_high_threshold() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // threshold=1.0 should compact any buffer with waste let original = VarBinViewArray::from_iter_str([ "this is a longer string that will be stored in a buffer", @@ -364,7 +370,9 @@ mod tests { let indices = buffer![0u32, 2u32].into_array(); let taken = original.take(indices).unwrap(); let taken = taken - .execute::(&mut array_session().create_execution_ctx()) + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); let original_buffers = taken.data_buffers().len(); @@ -381,7 +389,7 @@ mod tests { #[test] fn test_selective_compaction_preserves_well_utilized_buffers() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create an array with multiple strings in one buffer (well-utilized) let str1 = "first long string that needs external buffer storage"; let str2 = "second long string also in buffer"; @@ -404,7 +412,7 @@ mod tests { #[test] fn test_selective_compaction_with_mixed_utilization() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create array with some long strings let strings: Vec = (0..10) .map(|i| { @@ -421,7 +429,9 @@ mod tests { let indices_array = buffer![0u32, 2u32, 4u32, 6u32, 8u32].into_array(); let taken = original.take(indices_array).unwrap(); let taken = taken - .execute::(&mut array_session().create_execution_ctx()) + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); // Compact with moderate threshold @@ -436,7 +446,7 @@ mod tests { #[test] fn test_slice_strategy_with_contiguous_range() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create array with strings that will be in one buffer let strings: Vec = (0..20) .map(|i| format!("this is a long string number {} for slice test", i)) @@ -448,7 +458,9 @@ mod tests { let indices_array = buffer![0u32, 1u32, 2u32, 3u32, 4u32].into_array(); let taken = original.take(indices_array).unwrap(); let taken = taken - .execute::(&mut array_session().create_execution_ctx()) + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); // Get buffer stats before compaction let utils_before = taken.buffer_utilizations().unwrap(); diff --git a/vortex-array/src/arrays/varbinview/compute/cast.rs b/vortex-array/src/arrays/varbinview/compute/cast.rs index 476c74e7309..a0d3a31caca 100644 --- a/vortex-array/src/arrays/varbinview/compute/cast.rs +++ b/vortex-array/src/arrays/varbinview/compute/cast.rs @@ -86,7 +86,8 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[rstest] #[case( diff --git a/vortex-array/src/arrays/varbinview/compute/mod.rs b/vortex-array/src/arrays/varbinview/compute/mod.rs index 7a0d0b2d67f..3d9afdbc73a 100644 --- a/vortex-array/src/arrays/varbinview/compute/mod.rs +++ b/vortex-array/src/arrays/varbinview/compute/mod.rs @@ -42,8 +42,8 @@ mod tests { use rstest::rstest; use crate::VortexSessionExecute; - use crate::array_session; use crate::compute::conformance::consistency::test_array_consistency; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; @@ -83,7 +83,7 @@ mod tests { fn test_varbinview_consistency(#[case] array: VarBinViewArray) { test_array_consistency( &array.into_array(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ); } } diff --git a/vortex-array/src/arrays/varbinview/mod.rs b/vortex-array/src/arrays/varbinview/mod.rs index 6f3398884b3..13b3b1714b8 100644 --- a/vortex-array/src/arrays/varbinview/mod.rs +++ b/vortex-array/src/arrays/varbinview/mod.rs @@ -15,7 +15,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::VarBinView; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } diff --git a/vortex-array/src/arrays/varbinview/tests.rs b/vortex-array/src/arrays/varbinview/tests.rs index 981d4f602a1..0884c46d094 100644 --- a/vortex-array/src/arrays/varbinview/tests.rs +++ b/vortex-array/src/arrays/varbinview/tests.rs @@ -2,14 +2,14 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::VarBinViewArray; use crate::arrays::varbinview::BinaryView; use crate::assert_arrays_eq; +use crate::default_session_builder; #[test] pub fn varbin_view() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]); assert_arrays_eq!( @@ -21,7 +21,7 @@ pub fn varbin_view() { #[test] pub fn slice_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["hello world", "hello world this is a long string"]) .slice(1..2) @@ -35,7 +35,7 @@ pub fn slice_array() { #[test] pub fn flatten_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let binary_arr = VarBinViewArray::from_iter_str(["string1", "string2"]); assert_arrays_eq!( binary_arr, diff --git a/vortex-array/src/arrays/varbinview/vtable/kernel.rs b/vortex-array/src/arrays/varbinview/vtable/kernel.rs index e09b381d590..87fec4dfe46 100644 --- a/vortex-array/src/arrays/varbinview/vtable/kernel.rs +++ b/vortex-array/src/arrays/varbinview/vtable/kernel.rs @@ -1,21 +1,21 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::ArrayVTable; use crate::arrays::Dict; use crate::arrays::VarBinView; use crate::arrays::dict::TakeExecuteAdaptor; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::cast::CastExecuteAdaptor; use crate::scalar_fn::fns::zip::Zip; use crate::scalar_fn::fns::zip::ZipExecuteAdaptor; -pub(crate) fn initialize(session: &VortexSession) { - let kernels = session.kernels(); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + let kernels = builder_kernels(session); kernels.register_execute_parent_kernel(Cast.id(), VarBinView, CastExecuteAdaptor(VarBinView)); kernels.register_execute_parent_kernel(Dict.id(), VarBinView, TakeExecuteAdaptor(VarBinView)); kernels.register_execute_parent_kernel(Zip.id(), VarBinView, ZipExecuteAdaptor(VarBinView)); diff --git a/vortex-array/src/arrays/varbinview/vtable/mod.rs b/vortex-array/src/arrays/varbinview/vtable/mod.rs index 386625adcf1..d9b5c4b6950 100644 --- a/vortex-array/src/arrays/varbinview/vtable/mod.rs +++ b/vortex-array/src/arrays/varbinview/vtable/mod.rs @@ -12,6 +12,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::CachedId; use crate::ArrayRef; @@ -41,7 +42,7 @@ mod validity; /// A [`VarBinView`]-encoded Vortex array. pub type VarBinViewArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } @@ -249,8 +250,8 @@ mod tests { use crate::ArrayContext; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::serde::SerializeOptions; use crate::serde::SerializedArray; @@ -266,7 +267,7 @@ mod tests { let dtype = array.dtype().clone(); let len = array.len(); - let session = array_session(); + let session = default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let serialized = array diff --git a/vortex-array/src/arrays/variant/mod.rs b/vortex-array/src/arrays/variant/mod.rs index 1a5157e9aa1..d64faf28046 100644 --- a/vortex-array/src/arrays/variant/mod.rs +++ b/vortex-array/src/arrays/variant/mod.rs @@ -12,7 +12,7 @@ use vortex_error::vortex_ensure; pub use self::vtable::Variant; pub use self::vtable::VariantArray; -pub(crate) fn initialize(session: &vortex_session::VortexSession) { +pub(crate) fn initialize(session: &mut vortex_session::VortexSessionBuilder) { vtable::initialize(session); } @@ -89,7 +89,6 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -99,6 +98,7 @@ mod tests { use crate::arrays::variant::VariantArrayExt; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -141,7 +141,7 @@ mod tests { } fn execute_variant(array: ArrayRef) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let Canonical::Variant(variant) = array.execute::(&mut ctx)? else { return Err(vortex_err!("expected canonical variant array")); }; @@ -159,7 +159,7 @@ mod tests { let shredded = array .shredded() .ok_or_else(|| vortex_err!("expected shredded child"))?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let shredded = shredded.clone().execute::(&mut ctx)?; let expected_shredded_array = if let Some(values) = expected_shredded .iter() @@ -181,7 +181,7 @@ mod tests { ) -> VortexResult<()> { assert_eq!(array.len(), expected_core.len()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); for (idx, expected) in expected_core.iter().enumerate() { let scalar = array.core_storage().execute_scalar(idx, &mut ctx)?; let variant = scalar.as_variant(); @@ -273,7 +273,7 @@ mod tests { let shredded = PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]).into_array(); let variant = VariantArray::try_new(core_storage, Some(shredded))?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); for (idx, expected) in [Some(10i32), None, Some(3)].into_iter().enumerate() { let scalar = variant.execute_scalar(idx, &mut ctx)?; let variant = scalar.as_variant(); @@ -379,7 +379,7 @@ mod tests { #[test] fn variant_get_keeps_valid_shredded_rows_for_matching_dtype() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([( "a", @@ -395,7 +395,9 @@ mod tests { let result = variant .into_array() .apply(&expr)? - .execute::(&mut array_session().create_execution_ctx())?; + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_arrays_eq!( result, @@ -407,7 +409,7 @@ mod tests { #[test] fn variant_get_treats_value_and_typed_value_as_logical_field_names() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let core_storage = row_storage([1, 2, 3])?; let shredded = StructArray::try_from_iter([ ( @@ -430,7 +432,9 @@ mod tests { .clone() .into_array() .apply(&value_expr)? - .execute::(&mut array_session().create_execution_ctx())?; + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_arrays_eq!( value_result, PrimitiveArray::from_option_iter([Some(10i32), Some(20), Some(30)]), @@ -445,7 +449,9 @@ mod tests { let typed_value_result = variant .into_array() .apply(&typed_value_expr)? - .execute::(&mut array_session().create_execution_ctx())?; + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + )?; assert_arrays_eq!( typed_value_result, PrimitiveArray::from_option_iter([Some(40i32), Some(50), Some(60)]), diff --git a/vortex-array/src/arrays/variant/vtable/kernel.rs b/vortex-array/src/arrays/variant/vtable/kernel.rs index bfe89dea810..9dff3313991 100644 --- a/vortex-array/src/arrays/variant/vtable/kernel.rs +++ b/vortex-array/src/arrays/variant/vtable/kernel.rs @@ -3,7 +3,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use super::merge_typed_scalar_as_variant; use crate::ArrayRef; @@ -25,17 +25,19 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; use crate::kernel::ExecuteParentKernel; -use crate::optimizer::kernels::ArrayKernelsExt; +use crate::optimizer::kernels::builder_kernels; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::variant_get::VariantGet; use crate::scalar_fn::fns::variant_get::VariantGetOptions; use crate::scalar_fn::fns::variant_get::VariantPath; use crate::scalar_fn::fns::variant_get::VariantPathElement; -pub(crate) fn initialize(session: &VortexSession) { - session - .kernels() - .register_execute_parent_kernel(VariantGet.id(), Variant, VariantGetKernel); +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { + builder_kernels(session).register_execute_parent_kernel( + VariantGet.id(), + Variant, + VariantGetKernel, + ); } #[derive(Default, Debug)] diff --git a/vortex-array/src/arrays/variant/vtable/mod.rs b/vortex-array/src/arrays/variant/vtable/mod.rs index 19ade687875..27091e5aa4f 100644 --- a/vortex-array/src/arrays/variant/vtable/mod.rs +++ b/vortex-array/src/arrays/variant/vtable/mod.rs @@ -12,6 +12,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_proto::dtype as pb; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::CachedId; use vortex_utils::aliases::hash_set::HashSet; @@ -42,7 +43,7 @@ use crate::serde::ArrayChildren; /// A [`Variant`]-encoded Vortex array. pub type VariantArray = Array; -pub(crate) fn initialize(session: &VortexSession) { +pub(crate) fn initialize(session: &mut VortexSessionBuilder) { kernel::initialize(session); } diff --git a/vortex-array/src/arrow/executor/byte.rs b/vortex-array/src/arrow/executor/byte.rs index 2db9d3e6494..beba0dd6f45 100644 --- a/vortex-array/src/arrow/executor/byte.rs +++ b/vortex-array/src/arrow/executor/byte.rs @@ -86,9 +86,9 @@ mod tests { use crate::IntoArray; use crate::LEGACY_SESSION; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrow::ArrowArrayExecutor; use crate::arrow::executor::byte::VarBinViewArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; @@ -122,7 +122,7 @@ mod tests { #[case] vortex_array: VarBinViewArray, #[case] target_dtype: DataType, ) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arrow = vortex_array .into_array() .execute_arrow(Some(&target_dtype), &mut ctx) @@ -169,7 +169,7 @@ mod tests { vortex_dtype, ); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let arrow = vortex_array .into_array() .execute_arrow(Some(&target_dtype), &mut ctx) diff --git a/vortex-array/src/arrow/executor/decimal.rs b/vortex-array/src/arrow/executor/decimal.rs index a9aa9fdd2f5..596ae0ca3e8 100644 --- a/vortex-array/src/arrow/executor/decimal.rs +++ b/vortex-array/src/arrow/executor/decimal.rs @@ -224,18 +224,18 @@ mod tests { use crate::VortexSessionExecute; use crate::array::IntoArray; - use crate::array_session; use crate::arrow::ArrowArrayExecutor; use crate::arrow::executor::decimal::DecimalArray; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; + use crate::default_session_builder; use crate::dtype::DecimalDType; use crate::dtype::NativeDecimalType; use crate::validity::Validity; #[test] fn decimal_to_arrow() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Make a very simple i128 and i256 array. let decimal_vortex = DecimalArray::new( buffer![1i128, 2i128, 3i128, 4i128, 5i128], @@ -264,7 +264,7 @@ mod tests { fn test_to_arrow_decimal128( #[case] _decimal_type: T, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); decimal.append_value(10); decimal.append_value(11); @@ -292,7 +292,7 @@ mod tests { fn test_to_arrow_decimal32(#[case] _decimal_type: T) -> VortexResult<()> { use arrow_array::Decimal32Array; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); decimal.append_value(10); decimal.append_value(11); @@ -320,7 +320,7 @@ mod tests { fn test_to_arrow_decimal64(#[case] _decimal_type: T) -> VortexResult<()> { use arrow_array::Decimal64Array; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); decimal.append_value(10); decimal.append_value(11); @@ -348,7 +348,7 @@ mod tests { fn test_to_arrow_decimal256( #[case] _decimal_type: T, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut decimal = DecimalBuilder::new::(DecimalDType::new(2, 1), false.into()); decimal.append_value(10); decimal.append_value(11); diff --git a/vortex-array/src/arrow/executor/dictionary.rs b/vortex-array/src/arrow/executor/dictionary.rs index 97d25c3f7cc..cb37ea4c811 100644 --- a/vortex-array/src/arrow/executor/dictionary.rs +++ b/vortex-array/src/arrow/executor/dictionary.rs @@ -149,12 +149,12 @@ mod tests { use vortex_error::VortexResult; use crate::IntoArray; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinViewArray; use crate::arrow::ArrowArrayExecutor; use crate::arrow::executor::dictionary::ConstantArray; use crate::arrow::executor::dictionary::DictArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; use crate::executor::VortexSessionExecute; @@ -165,7 +165,10 @@ mod tests { } fn execute(array: crate::ArrayRef, dt: &DataType) -> VortexResult { - array.execute_arrow(Some(dt), &mut array_session().create_execution_ctx()) + array.execute_arrow( + Some(dt), + &mut default_session_builder().build().create_execution_ctx(), + ) } fn dict_basic_input() -> crate::ArrayRef { diff --git a/vortex-array/src/arrow/executor/list.rs b/vortex-array/src/arrow/executor/list.rs index 7bf7a16f496..ec856d3ccc5 100644 --- a/vortex-array/src/arrow/executor/list.rs +++ b/vortex-array/src/arrow/executor/list.rs @@ -225,7 +225,8 @@ mod tests { use crate::validity::Validity; /// A shared session for these list-executor tests, used to create execution contexts. - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn test_to_arrow_list_i32() -> VortexResult<()> { diff --git a/vortex-array/src/arrow/executor/list_view.rs b/vortex-array/src/arrow/executor/list_view.rs index 7d5666b3762..69508481581 100644 --- a/vortex-array/src/arrow/executor/list_view.rs +++ b/vortex-array/src/arrow/executor/list_view.rs @@ -127,7 +127,8 @@ mod tests { use crate::validity::Validity; /// A shared session for these list-view-executor tests, used to create execution contexts. - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn trims_zero_copy_with_significant_trailing_waste() -> VortexResult<()> { diff --git a/vortex-array/src/arrow/executor/run_end.rs b/vortex-array/src/arrow/executor/run_end.rs index 0632f871f29..007f79cbdb2 100644 --- a/vortex-array/src/arrow/executor/run_end.rs +++ b/vortex-array/src/arrow/executor/run_end.rs @@ -218,7 +218,8 @@ mod tests { use crate::executor::VortexSessionExecute; use crate::scalar::Scalar; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); fn ree_type(ends: DataType, values_dtype: DataType) -> DataType { DataType::RunEndEncoded( diff --git a/vortex-array/src/arrow/executor/struct_.rs b/vortex-array/src/arrow/executor/struct_.rs index d3b7bcda2dc..f01d12fa416 100644 --- a/vortex-array/src/arrow/executor/struct_.rs +++ b/vortex-array/src/arrow/executor/struct_.rs @@ -210,18 +210,18 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array; - use crate::array_session; use crate::arrays; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrow::ArrowArrayExecutor; use crate::arrow::FromArrowArray; + use crate::default_session_builder; use crate::dtype::FieldNames; use crate::validity::Validity; #[test] fn struct_nullable_non_null_to_arrow() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::AllValid); let struct_a = StructArray::try_new( @@ -242,7 +242,7 @@ mod tests { #[test] fn struct_nullable_with_nulls_to_arrow() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let xs = PrimitiveArray::from_option_iter(vec![Some(0_i64), Some(1), Some(2), None, Some(3)]); @@ -267,7 +267,7 @@ mod tests { #[test] fn struct_to_arrow_with_schema_mismatch() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let xs = PrimitiveArray::new(buffer![0i64, 1, 2, 3, 4], Validity::AllValid); let struct_a = StructArray::try_new( @@ -297,7 +297,7 @@ mod tests { #[test] fn test_to_arrow() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = StructArray::from_fields( vec![ ( @@ -340,7 +340,7 @@ mod tests { #[test] fn to_arrow_with_non_nullable_fields() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = StructArray::from_fields( vec![ ( diff --git a/vortex-array/src/arrow/session.rs b/vortex-array/src/arrow/session.rs index 0c4cc7b2a6d..ac6fbca8a4c 100644 --- a/vortex-array/src/arrow/session.rs +++ b/vortex-array/src/arrow/session.rs @@ -631,7 +631,7 @@ mod tests { use super::*; use crate::VortexSessionExecute; - use crate::array_session; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::Nullability; @@ -764,7 +764,7 @@ mod tests { #[test] fn execute_arrow_target_none_preserves_top_level_uuid_metadata() -> VortexResult<()> { - let vortex_session = array_session(); + let vortex_session = default_session_builder().build(); let mut ctx = vortex_session.create_execution_ctx(); let session = vortex_session.arrow(); diff --git a/vortex-array/src/builders/bool.rs b/vortex-array/src/builders/bool.rs index f26cbf04b57..ba23ef277cb 100644 --- a/vortex-array/src/builders/bool.rs +++ b/vortex-array/src/builders/bool.rs @@ -162,7 +162,6 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; @@ -172,6 +171,7 @@ mod tests { use crate::builders::builder_with_capacity; #[expect(deprecated)] use crate::canonical::ToCanonical as _; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -199,7 +199,7 @@ mod tests { let chunk_count = 10; let chunk = make_opt_bool_chunks(len, chunk_count); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = builder_with_capacity(chunk.dtype(), len * chunk_count); chunk .clone() @@ -221,7 +221,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = BoolBuilder::with_capacity(Nullability::Nullable, 10); // Test appending true value. diff --git a/vortex-array/src/builders/decimal.rs b/vortex-array/src/builders/decimal.rs index aa777051e39..497b3abed89 100644 --- a/vortex-array/src/builders/decimal.rs +++ b/vortex-array/src/builders/decimal.rs @@ -310,11 +310,11 @@ impl Default for DecimalBuffer { #[cfg(test)] mod tests { use crate::VortexSessionExecute; - use crate::array_session; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; use crate::builders::decimal::DecimalArray; + use crate::default_session_builder; use crate::dtype::DecimalDType; #[test] @@ -333,10 +333,16 @@ mod tests { for i in 0..i8s.len() { assert_eq!( - i8s.execute_scalar(i, &mut array_session().create_execution_ctx()) - .unwrap(), + i8s.execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), i128s - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } @@ -344,7 +350,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use crate::scalar::Scalar; // Simply test that the builder accepts its own finish output via scalar. @@ -364,7 +370,10 @@ mod tests { let mut builder2 = DecimalBuilder::new::(DecimalDType::new(10, 2), true.into()); for i in 0..array.len() { let scalar = array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); builder2.append_scalar(&scalar).unwrap(); } diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 00f14388aa2..2f24ba6233c 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -332,7 +332,8 @@ mod test { use crate::dtype::Nullability; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn encode_varbin() { diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index dea2932367b..53fc03f2961 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -209,7 +209,8 @@ mod test { use crate::builders::dict::dict_encode; use crate::builders::dict::primitive::PrimitiveArray; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn encode_primitive() { diff --git a/vortex-array/src/builders/extension.rs b/vortex-array/src/builders/extension.rs index 61d5fe8d7b0..d359dba2403 100644 --- a/vortex-array/src/builders/extension.rs +++ b/vortex-array/src/builders/extension.rs @@ -128,10 +128,10 @@ impl ArrayBuilder for ExtensionBuilder { mod tests { use super::*; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::extension::datetime::Date; use crate::extension::datetime::TimeUnit; @@ -139,7 +139,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::Nullable).erased(); let mut builder = ExtensionBuilder::new(ext_dtype.clone()); diff --git a/vortex-array/src/builders/fixed_size_list.rs b/vortex-array/src/builders/fixed_size_list.rs index 9129d8cb644..97a9b7afc5b 100644 --- a/vortex-array/src/builders/fixed_size_list.rs +++ b/vortex-array/src/builders/fixed_size_list.rs @@ -290,11 +290,11 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::fixed_size_list::FixedSizeListArrayExt; use crate::builders::ArrayBuilder; use crate::builders::fixed_size_list::FixedSizeListArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::Nullability::Nullable; @@ -448,7 +448,7 @@ mod tests { #[test] fn test_nullable_lists_non_nullable_elements() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype: Arc = Arc::new(DType::Primitive(I32, NonNullable)); let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 0); @@ -567,7 +567,7 @@ mod tests { #[test] fn test_append_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Elements must be nullable if we're going to append null lists let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); @@ -597,7 +597,7 @@ mod tests { #[test] fn test_append_scalar_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Elements must be nullable if we're going to append null lists let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); let mut builder = FixedSizeListBuilder::with_capacity(dtype, 2, Nullable, 0); @@ -670,7 +670,7 @@ mod tests { #[test] fn test_extend_from_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Create a source array. @@ -741,7 +741,7 @@ mod tests { #[test] fn test_extend_degenerate_arrays() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Create degenerate source arrays (size = 0). @@ -846,7 +846,7 @@ mod tests { #[test] fn test_mixed_operations() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Use nullable elements since we'll be appending nulls let dtype: Arc = Arc::new(DType::Primitive(I32, Nullable)); let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 0); @@ -932,7 +932,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 2, Nullable, 10); @@ -1002,7 +1002,7 @@ mod tests { #[test] fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = FixedSizeListBuilder::with_capacity(Arc::clone(&dtype), 3, NonNullable, 10); diff --git a/vortex-array/src/builders/list.rs b/vortex-array/src/builders/list.rs index 12c74ba53a4..1ed3f3b84a3 100644 --- a/vortex-array/src/builders/list.rs +++ b/vortex-array/src/builders/list.rs @@ -331,7 +331,6 @@ mod tests { use crate::IntoArray; #[expect(deprecated)] use crate::ToCanonical as _; - use crate::array_session; use crate::arrays::ChunkedArray; use crate::arrays::PrimitiveArray; use crate::arrays::list::ListArrayExt; @@ -340,6 +339,7 @@ mod tests { use crate::builders::ArrayBuilder; use crate::builders::list::ListArray; use crate::builders::list::ListBuilder; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::IntegerPType; use crate::dtype::Nullability; @@ -456,7 +456,7 @@ mod tests { .unwrap(); assert_eq!(list.len(), 3); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = ListBuilder::::with_capacity(Arc::new(I32.into()), Nullable, 18, 9); builder.extend_from_array(&list); @@ -544,18 +544,30 @@ mod tests { assert_eq!( one_trailing_unused_element - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), canon_values - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); assert_eq!( second_array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), canon_values - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } @@ -585,7 +597,7 @@ mod tests { let array = builder.finish_into_list(); assert_eq!(array.len(), 3); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Check actual values using scalar_at. @@ -642,7 +654,7 @@ mod tests { #[test] fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = ListBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10); diff --git a/vortex-array/src/builders/listview.rs b/vortex-array/src/builders/listview.rs index cb131048339..3540d138d66 100644 --- a/vortex-array/src/builders/listview.rs +++ b/vortex-array/src/builders/listview.rs @@ -440,13 +440,13 @@ mod tests { use super::ListViewBuilder; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ListArray; use crate::arrays::ListViewArray; use crate::arrays::listview::ListViewArrayExt; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::listview::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::Nullability::Nullable; @@ -465,7 +465,7 @@ mod tests { #[test] fn test_basic_append_and_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); let mut builder = ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); @@ -529,7 +529,7 @@ mod tests { #[test] fn test_different_offset_size_types() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test u32 offsets with u8 sizes. let dtype: Arc = Arc::new(I32.into()); let mut builder = @@ -602,7 +602,7 @@ mod tests { #[test] fn test_builder_trait_methods() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); let mut builder = ListViewBuilder::::with_capacity(Arc::clone(&dtype), Nullable, 0, 0); @@ -655,7 +655,7 @@ mod tests { #[test] fn test_extend_from_array() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Create a source ListArray. @@ -721,7 +721,7 @@ mod tests { #[test] fn test_extend_from_array_overlapping_listview() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype: Arc = Arc::new(I32.into()); // Non-ZCTL source: @@ -790,7 +790,7 @@ mod tests { #[test] fn test_append_array_as_list() { let dtype: Arc = Arc::new(I32.into()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = ListViewBuilder::::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10); diff --git a/vortex-array/src/builders/mod.rs b/vortex-array/src/builders/mod.rs index 48a347aa3d5..d4c3c4e2b78 100644 --- a/vortex-array/src/builders/mod.rs +++ b/vortex-array/src/builders/mod.rs @@ -11,7 +11,7 @@ //! ``` //! use vortex_array::builders::{builder_with_capacity, ArrayBuilder}; //! use vortex_array::dtype::{DType, Nullability}; -//! use vortex_array::{VortexSessionExecute, array_session}; +//! use vortex_array::{VortexSessionExecute, default_session_builder}; //! //! // Create a new builder for string data. //! let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4); @@ -22,7 +22,7 @@ //! builder.append_scalar(&"d".into()).unwrap(); //! //! let strings = builder.finish(); -//! let mut ctx = array_session().create_execution_ctx(); +//! let mut ctx = default_session_builder().build().create_execution_ctx(); //! //! assert_eq!(strings.execute_scalar(0, &mut ctx).unwrap(), "a".into()); //! assert_eq!(strings.execute_scalar(1, &mut ctx).unwrap(), "b".into()); @@ -225,7 +225,7 @@ pub trait ArrayBuilder: Send { /// ``` /// use vortex_array::builders::{builder_with_capacity, ArrayBuilder}; /// use vortex_array::dtype::{DType, Nullability}; -/// use vortex_array::{VortexSessionExecute, array_session}; +/// use vortex_array::{VortexSessionExecute, default_session_builder}; /// /// // Create a new builder for string data. /// let mut builder = builder_with_capacity(&DType::Utf8(Nullability::NonNullable), 4); @@ -236,7 +236,7 @@ pub trait ArrayBuilder: Send { /// builder.append_scalar(&"d".into()).unwrap(); /// /// let strings = builder.finish(); -/// let mut ctx = array_session().create_execution_ctx(); +/// let mut ctx = default_session_builder().build().create_execution_ctx(); /// /// assert_eq!(strings.execute_scalar(0, &mut ctx).unwrap(), "a".into()); /// assert_eq!(strings.execute_scalar(1, &mut ctx).unwrap(), "b".into()); diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index e7e8d3d7512..449592bbf00 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -377,8 +377,8 @@ mod tests { use super::*; use crate::VortexSessionExecute; - use crate::array_session; use crate::assert_arrays_eq; + use crate::default_session_builder; /// REGRESSION TEST: This test verifies that multiple sequential ranges have correct offsets. /// @@ -386,7 +386,7 @@ mod tests { /// buffer. #[test] fn test_multiple_uninit_ranges_correct_offsets() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); // First range. @@ -454,19 +454,28 @@ mod tests { // Check validity using scalar_at - nulls will return is_null() = true. assert!( !array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert!( array - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert!( !array - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); @@ -497,7 +506,7 @@ mod tests { /// This verifies the new simplified API without the redundant `len` parameter. #[test] fn test_copy_from_slice_with_offsets() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 10); let mut range = builder.uninit_range(6); @@ -564,13 +573,19 @@ mod tests { // Check validity - the first two should be valid (from append_value). assert!( !array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); // initial value 100 assert!( !array - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); // initial value 200 @@ -578,19 +593,28 @@ mod tests { // Check the range items with modified validity. assert!( !array - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); // range index 0 - set to valid assert!( array - .execute_scalar(3, &mut array_session().create_execution_ctx()) + .execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); // range index 1 - left as null assert!( !array - .execute_scalar(4, &mut array_session().create_execution_ctx()) + .execute_scalar( + 4, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); // range index 2 - set to valid @@ -686,7 +710,7 @@ mod tests { // values[2] might be any value since it's null. // Check validity - first two should be valid, third should be null. - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!( array .validity() diff --git a/vortex-array/src/builders/struct_.rs b/vortex-array/src/builders/struct_.rs index 016d42508ee..1a4c6c5470b 100644 --- a/vortex-array/src/builders/struct_.rs +++ b/vortex-array/src/builders/struct_.rs @@ -212,13 +212,13 @@ impl ArrayBuilder for StructBuilder { mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::struct_::StructArray; use crate::builders::struct_::StructBuilder; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType::I32; @@ -258,7 +258,7 @@ mod tests { assert_eq!(struct_.dtype(), &dtype); assert_eq!( struct_ - .valid_count(&mut array_session().create_execution_ctx()) + .valid_count(&mut default_session_builder().build().create_execution_ctx()) .unwrap(), 1 ); @@ -266,7 +266,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use crate::scalar::Scalar; let dtype = DType::Struct( diff --git a/vortex-array/src/builders/tests.rs b/vortex-array/src/builders/tests.rs index 02babab4a7d..403aff4f118 100644 --- a/vortex-array/src/builders/tests.rs +++ b/vortex-array/src/builders/tests.rs @@ -9,9 +9,9 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use crate::VortexSessionExecute; -use crate::array_session; use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; @@ -95,10 +95,16 @@ fn test_append_zeros_matches_default_value(#[case] dtype: DType) { // Compare each element. for i in 0..num_elements { let scalar_zeros = array_zeros - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); let scalar_manual = array_manual - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert_eq!( @@ -197,7 +203,10 @@ fn test_append_defaults_behavior(#[case] dtype: DType, #[case] should_be_null: b for i in 0..3 { let scalar = array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); if should_be_null { assert!(scalar.is_null(), "Element at index {} should be null", i); @@ -255,10 +264,16 @@ where // Compare each element. for i in 0..array_direct.len() { let scalar_direct = array_direct - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); let scalar_indirect = array_indirect - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert_eq!( @@ -549,7 +564,10 @@ fn test_append_scalar_comprehensive(#[case] dtype: DType) { // Verify each scalar matches. for (i, expected_scalar) in scalars.iter().enumerate() { let actual_scalar = array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert_scalars_equal(&actual_scalar, expected_scalar, &dtype, i); } @@ -557,7 +575,10 @@ fn test_append_scalar_comprehensive(#[case] dtype: DType) { // If nullable, verify the last element is null. if dtype.is_nullable() { let null_scalar = array - .execute_scalar(num_elements, &mut array_session().create_execution_ctx()) + .execute_scalar( + num_elements, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert!( null_scalar.is_null(), @@ -707,31 +728,46 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { // Check the pattern. assert!( !array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert!( array - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert!( !array - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert!( array - .execute_scalar(3, &mut array_session().create_execution_ctx()) + .execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); assert!( !array - .execute_scalar(4, &mut array_session().create_execution_ctx()) + .execute_scalar( + 4, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .is_null() ); @@ -739,7 +775,10 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { // Verify non-null values match. assert_scalars_equal( &array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(), &test_scalars[0], &dtype, @@ -747,7 +786,10 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { ); assert_scalars_equal( &array - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(), &test_scalars[1], &dtype, @@ -755,7 +797,10 @@ fn test_append_scalar_mixed_nulls(#[case] dtype: DType) { ); assert_scalars_equal( &array - .execute_scalar(4, &mut array_session().create_execution_ctx()) + .execute_scalar( + 4, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(), &test_scalars[2], &dtype, @@ -812,7 +857,10 @@ fn test_append_scalar_repeated_same_instance() { // All values should be 42. for i in 0..5 { let actual = array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert_eq!( actual.as_primitive().typed_value::(), @@ -855,7 +903,7 @@ fn test_set_validity_overrides_validity( builder.set_validity(mask); let validity = builder.finish().validity()?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); for (i, &valid) in expected.iter().enumerate() { assert_eq!( validity.execute_is_valid(i, &mut ctx)?, @@ -877,7 +925,7 @@ fn test_set_validity_noop_when_non_nullable() -> VortexResult<()> { builder.set_validity(Mask::new_false(4)); let validity = builder.finish().validity()?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); for i in 0..4 { assert!( validity.execute_is_valid(i, &mut ctx)?, diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index 3aabfd87c31..21f299dfc71 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -874,17 +874,17 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::VarBinViewBuilder; use crate::builders::varbinview::VarBinViewArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; #[test] fn test_utf8_builder() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); builder.append_value("Hello"); @@ -912,7 +912,7 @@ mod tests { #[test] fn test_utf8_builder_with_extend() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = { let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10); @@ -953,7 +953,7 @@ mod tests { let mut builder = VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullability::Nullable), 10); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); array.append_to_builder(&mut builder, &mut ctx)?; assert_eq!(builder.completed_block_count(), 1); @@ -988,7 +988,7 @@ mod tests { #[test] fn test_append_scalar() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use crate::scalar::Scalar; // Test with Utf8 builder. @@ -1081,7 +1081,10 @@ mod tests { // Verify the value was stored correctly let retrieved = array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_binary() .value() diff --git a/vortex-array/src/canonical.rs b/vortex-array/src/canonical.rs index 28953a93e4e..32adc6b2c07 100644 --- a/vortex-array/src/canonical.rs +++ b/vortex-array/src/canonical.rs @@ -1159,7 +1159,8 @@ mod test { use crate::scalar::Scalar; /// A shared session for these canonical tests, used to create execution contexts. - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); fn variant_core_storage(len: usize) -> ArrayRef { ConstantArray::new( diff --git a/vortex-array/src/compute/conformance/cast.rs b/vortex-array/src/compute/conformance/cast.rs index c823b2a778a..1c8f241c19f 100644 --- a/vortex-array/src/compute/conformance/cast.rs +++ b/vortex-array/src/compute/conformance/cast.rs @@ -12,8 +12,8 @@ use crate::VortexSessionExecute; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::MinMaxResult; use crate::aggregate_fn::fns::min_max::min_max; -use crate::array_session; use crate::builtins::ArrayBuiltins; +use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -23,7 +23,9 @@ use crate::scalar::Scalar; fn cast_and_execute(array: &ArrayRef, dtype: DType) -> VortexResult { Ok(array .cast(dtype)? - .execute::(&mut array_session().create_execution_ctx())? + .execute::( + &mut default_session_builder().build().create_execution_ctx(), + )? .0 .into_array()) } @@ -75,10 +77,16 @@ fn test_cast_identity(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -110,7 +118,10 @@ fn test_cast_from_null(array: &ArrayRef) { for i in 0..array.len().min(10) { assert!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") .is_null() ); @@ -130,7 +141,7 @@ fn test_cast_from_null(array: &ArrayRef) { fn test_cast_to_non_nullable(array: &ArrayRef) { if array - .invalid_count(&mut array_session().create_execution_ctx()) + .invalid_count(&mut default_session_builder().build().create_execution_ctx()) .vortex_expect("invalid_count should succeed in conformance test") == 0 { @@ -142,10 +153,16 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), non_nullable - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -158,10 +175,16 @@ fn test_cast_to_non_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), back_to_nullable - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -191,10 +214,16 @@ fn test_cast_to_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), nullable - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -207,10 +236,16 @@ fn test_cast_to_nullable(array: &ArrayRef) { for i in 0..array.len().min(10) { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), - back.execute_scalar(i, &mut array_session().create_execution_ctx()) - .vortex_expect("scalar_at should succeed in conformance test") + back.execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) + .vortex_expect("scalar_at should succeed in conformance test") ); } } @@ -248,7 +283,7 @@ fn fits(value: &Scalar, ptype: PType) -> bool { } fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip: bool) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let maybe_min_max = min_max(array, &mut ctx, NumericalAggregateOpts::default()) .vortex_expect("cast should succeed in conformance test"); @@ -289,20 +324,32 @@ fn test_cast_to_primitive(array: &ArrayRef, target_ptype: PType, test_round_trip array .validity() .vortex_expect("validity_mask should succeed in conformance test") - .execute_mask(array.len(), &mut array_session().create_execution_ctx()) + .execute_mask( + array.len(), + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("Failed to compute validity mask"), casted .validity() .vortex_expect("validity_mask should succeed in conformance test") - .execute_mask(casted.len(), &mut array_session().create_execution_ctx()) + .execute_mask( + casted.len(), + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("Failed to compute validity mask") ); for i in 0..array.len().min(10) { let original = array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .vortex_expect("scalar_at should succeed in conformance test"); let casted = casted - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!( original diff --git a/vortex-array/src/compute/conformance/filter.rs b/vortex-array/src/compute/conformance/filter.rs index 01bfb3bb2fa..79e6f813f11 100644 --- a/vortex-array/src/compute/conformance/filter.rs +++ b/vortex-array/src/compute/conformance/filter.rs @@ -7,8 +7,8 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::assert_arrays_eq; +use crate::default_session_builder; use crate::dtype::DType; // Standard test array sizes @@ -62,7 +62,7 @@ pub fn create_runs_pattern(len: usize, run_length: usize) -> Vec { /// Tests that filtering with an all-true mask returns all elements unchanged fn test_all_filter(array: &ArrayRef) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let len = array.len(); let mask = Mask::new_true(len); let filtered = array @@ -101,10 +101,16 @@ fn test_selective_filter(array: &ArrayRef) { for (filtered_idx, i) in (0..len).step_by(2).enumerate() { assert_eq!( filtered - .execute_scalar(filtered_idx, &mut array_session().create_execution_ctx()) + .execute_scalar( + filtered_idx, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -121,18 +127,30 @@ fn test_selective_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 2); assert_eq!( filtered - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( filtered - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) + .execute_scalar( + len - 1, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -154,10 +172,16 @@ fn test_single_element_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 1); assert_eq!( filtered - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); @@ -172,10 +196,16 @@ fn test_single_element_filter(array: &ArrayRef) { assert_eq!(filtered.len(), 1); assert_eq!( filtered - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) + .execute_scalar( + len - 1, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -238,10 +268,16 @@ fn test_alternating_pattern_filter(array: &ArrayRef) { if keep { assert_eq!( filtered - .execute_scalar(filtered_idx, &mut array_session().create_execution_ctx()) + .execute_scalar( + filtered_idx, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); filtered_idx += 1; diff --git a/vortex-array/src/compute/conformance/mask.rs b/vortex-array/src/compute/conformance/mask.rs index 7a002662d38..c61409275d9 100644 --- a/vortex-array/src/compute/conformance/mask.rs +++ b/vortex-array/src/compute/conformance/mask.rs @@ -7,10 +7,10 @@ use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::builtins::ArrayBuiltins; +use crate::default_session_builder; /// Test mask compute function with various array sizes and patterns. /// The mask operation sets elements to null where the mask is true. @@ -54,16 +54,25 @@ fn test_heterogenous_mask(array: &ArrayRef) { if masked_out { assert!( !masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -87,10 +96,16 @@ fn test_empty_mask(array: &ArrayRef) { for i in 0..len { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -113,7 +128,10 @@ fn test_full_mask(array: &ArrayRef) { for i in 0..len { assert!( !masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("is_valid should succeed in conformance test") ); } @@ -135,16 +153,25 @@ fn test_alternating_mask(array: &ArrayRef) { if i % 2 == 0 { assert!( !masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -173,7 +200,10 @@ fn test_sparse_mask(array: &ArrayRef) { let valid_count = (0..len) .filter(|&i| { masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .vortex_expect("is_valid should succeed in conformance test") }) .count(); @@ -185,7 +215,10 @@ fn test_sparse_mask(array: &ArrayRef) { .filter(|&i| { pattern[i] || !array - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .vortex_expect("is_valid should succeed in conformance test") }) .count(); @@ -208,17 +241,26 @@ fn test_single_element_mask(array: &ArrayRef) { .vortex_expect("mask should succeed in conformance test"); assert!( !masked - .is_valid(0, &mut array_session().create_execution_ctx()) + .is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("is_valid should succeed in conformance test") ); for i in 1..len { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -249,16 +291,25 @@ fn test_double_mask(array: &ArrayRef) { if mask1_pattern[i] || mask2_pattern[i] { assert!( !double_masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( double_masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); @@ -281,8 +332,8 @@ fn test_nullable_mask_input(array: &ArrayRef) { let validity = crate::validity::Validity::from_iter(validity_values.clone()); let nullable_mask = BoolArray::new(bool_array.to_bit_buffer(), validity); - let mask_array = - nullable_mask.to_mask_fill_null_false(&mut array_session().create_execution_ctx()); + let mask_array = nullable_mask + .to_mask_fill_null_false(&mut default_session_builder().build().create_execution_ctx()); let masked = array .clone() .mask((!&mask_array).into_array()) @@ -293,16 +344,25 @@ fn test_nullable_mask_input(array: &ArrayRef) { if bool_values[i] && validity_values[i] { assert!( !masked - .is_valid(i, &mut array_session().create_execution_ctx()) + .is_valid( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("is_valid should succeed in conformance test") ); } else { assert_eq!( masked - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") .into_nullable() ); diff --git a/vortex-array/src/compute/conformance/search_sorted.rs b/vortex-array/src/compute/conformance/search_sorted.rs index 0d0cf9d8cc8..b5089ec027d 100644 --- a/vortex-array/src/compute/conformance/search_sorted.rs +++ b/vortex-array/src/compute/conformance/search_sorted.rs @@ -12,8 +12,8 @@ use vortex_buffer::buffer; use crate::ArrayRef; use crate::VortexSessionExecute; use crate::array::IntoArray; -use crate::array_session; use crate::arrays::PrimitiveArray; +use crate::default_session_builder; use crate::patches::Patches; use crate::validity::Validity; @@ -32,7 +32,7 @@ pub fn sparse_high_null_fill() -> ArrayRef { None, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() .into_array() @@ -49,7 +49,7 @@ pub fn sparse_high_non_null_fill() -> ArrayRef { None, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() .into_array() @@ -66,7 +66,7 @@ pub fn sparse_low() -> ArrayRef { None, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() .into_array() @@ -83,7 +83,7 @@ pub fn sparse_low_high() -> ArrayRef { None, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() .into_array() @@ -100,7 +100,7 @@ pub fn sparse_edge_patch_high() -> ArrayRef { None, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() .into_array() @@ -117,7 +117,7 @@ pub fn sparse_edge_patch_low() -> ArrayRef { None, ) .unwrap(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .unwrap() .into_array() diff --git a/vortex-array/src/compute/conformance/take.rs b/vortex-array/src/compute/conformance/take.rs index babaf392eb9..0e9cb52817b 100644 --- a/vortex-array/src/compute/conformance/take.rs +++ b/vortex-array/src/compute/conformance/take.rs @@ -8,8 +8,8 @@ use crate::ArrayRef; use crate::Canonical; use crate::IntoArray as _; use crate::VortexSessionExecute; -use crate::array_session; use crate::arrays::PrimitiveArray; +use crate::default_session_builder; use crate::dtype::Nullability; /// Test conformance of the take compute function for an array. @@ -84,10 +84,16 @@ fn test_take_all(array: &ArrayRef) { for i in 0..len { assert_eq!( array - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -125,11 +131,14 @@ fn test_take_selective(array: &ArrayRef) { array .execute_scalar( original_idx as usize, - &mut array_session().create_execution_ctx() + &mut default_session_builder().build().create_execution_ctx() ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(result_idx, &mut array_session().create_execution_ctx()) + .execute_scalar( + result_idx, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -145,18 +154,30 @@ fn test_take_first_and_last(array: &ArrayRef) { assert_eq!(result.len(), 2); assert_eq!( array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); assert_eq!( array - .execute_scalar(len - 1, &mut array_session().create_execution_ctx()) + .execute_scalar( + len - 1, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -190,17 +211,26 @@ fn test_take_with_nullable_indices(array: &ArrayRef) { match idx_opt { Some(idx) => { let expected = array - .execute_scalar(*idx as usize, &mut array_session().create_execution_ctx()) + .execute_scalar( + *idx as usize, + &mut default_session_builder().build().create_execution_ctx(), + ) .vortex_expect("scalar_at should succeed in conformance test"); let actual = result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx(), + ) .vortex_expect("scalar_at should succeed in conformance test"); assert_eq!(expected, actual); } None => { assert!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") .is_null() ); @@ -222,12 +252,18 @@ fn test_take_repeated_indices(array: &ArrayRef) { assert_eq!(result.len(), 3); let first_elem = array - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .vortex_expect("scalar_at should succeed in conformance test"); for i in 0..3 { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), first_elem ); @@ -258,10 +294,16 @@ fn test_take_reverse(array: &ArrayRef) { for i in 0..len { assert_eq!( array - .execute_scalar(len - 1 - i, &mut array_session().create_execution_ctx()) + .execute_scalar( + len - 1 - i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -279,10 +321,16 @@ fn test_take_single_middle(array: &ArrayRef) { assert_eq!(result.len(), 1); assert_eq!( array - .execute_scalar(middle_idx, &mut array_session().create_execution_ctx()) + .execute_scalar( + middle_idx, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -310,10 +358,16 @@ fn test_take_random_unsorted(array: &ArrayRef) { for (i, &idx) in indices.iter().enumerate() { assert_eq!( array - .execute_scalar(idx as usize, &mut array_session().create_execution_ctx()) + .execute_scalar( + idx as usize, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -336,10 +390,16 @@ fn test_take_contiguous_range(array: &ArrayRef) { for i in 0..(end - start) { assert_eq!( array - .execute_scalar(start + i, &mut array_session().create_execution_ctx()) + .execute_scalar( + start + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -372,10 +432,16 @@ fn test_take_mixed_repeated(array: &ArrayRef) { for (i, &idx) in indices.iter().enumerate() { assert_eq!( array - .execute_scalar(idx as usize, &mut array_session().create_execution_ctx()) + .execute_scalar( + idx as usize, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } @@ -404,10 +470,16 @@ fn test_take_large_indices(array: &ArrayRef) { let expected_idx = indices[i] as usize; assert_eq!( array - .execute_scalar(expected_idx, &mut array_session().create_execution_ctx()) + .execute_scalar( + expected_idx, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test"), result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("scalar_at should succeed in conformance test") ); } diff --git a/vortex-array/src/dtype/mod.rs b/vortex-array/src/dtype/mod.rs index 009c2610c73..0d009c04b79 100644 --- a/vortex-array/src/dtype/mod.rs +++ b/vortex-array/src/dtype/mod.rs @@ -203,5 +203,6 @@ mod test { use vortex_session::VortexSession; - pub(crate) static SESSION: LazyLock = LazyLock::new(crate::array_session); + pub(crate) static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); } diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 78d4c62f412..72d2218761f 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -501,7 +501,7 @@ mod tests { #[test] fn test_unknown_extension_allow_unknown() { - let session = crate::array_session().allow_unknown(); + let session = crate::default_session_builder().allow_unknown().build(); let proto = pb::DType { dtype_type: Some(DtypeType::Extension(Box::new(pb::Extension { id: "vortex.test.foreign_ext".to_string(), diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 0c083f18ee5..3b21a9befb3 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -876,7 +876,9 @@ mod tests { #[test] fn execution_ctx_snapshots_execute_parent_kernels_at_creation() { - let session = VortexSession::empty().with_some(KernelSession::empty()); + let session = VortexSession::builder() + .with_some(KernelSession::empty()) + .build(); let key = execute_parent_key(Bool.id(), Primitive.id()); let before_registration = session.create_execution_ctx(); diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index 303424f56c0..dda34fdbc1e 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -422,13 +422,13 @@ where /// ``` /// # use vortex_array::IntoArray; /// # use vortex_array::arrow::ArrowArrayExecutor; -/// # use vortex_array::{VortexSessionExecute, array_session}; +/// # use vortex_array::{VortexSessionExecute, default_session_builder}; /// # use vortex_buffer::buffer; /// # use vortex_array::expr::{checked_add, lit, root}; /// let xs = buffer![1, 2, 3].into_array(); /// let result = xs.apply(&checked_add(root(), lit(5))).unwrap(); /// -/// let mut ctx = array_session().create_execution_ctx(); +/// let mut ctx = default_session_builder().build().create_execution_ctx(); /// assert_eq!( /// &result.execute_arrow(None, &mut ctx).unwrap(), /// &buffer![6, 7, 8] diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index ebf1f877f15..df1fcdc90ce 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -74,7 +74,7 @@ mod tests { use vortex_session::VortexSession; use super::ExprSerializeProtoExt; - use crate::array_session; + use crate::default_session_builder; use crate::expr::Expression; use crate::expr::and; use crate::expr::between; @@ -108,16 +108,18 @@ mod tests { let s_expr = expr.serialize_proto().unwrap(); let buf = s_expr.encode_to_vec(); let s_expr = pb::Expr::decode(buf.as_slice()).unwrap(); - let deser_expr = Expression::from_proto(&s_expr, &array_session()).unwrap(); + let deser_expr = + Expression::from_proto(&s_expr, &default_session_builder().build()).unwrap(); assert_eq!(&deser_expr, &expr); } #[test] fn unknown_expression_id_allow_unknown() { - let session = VortexSession::empty() + let session = VortexSession::builder() .with::() - .allow_unknown(); + .allow_unknown() + .build(); let expr_proto = pb::Expr { id: "vortex.test.foreign_scalar_fn".to_string(), diff --git a/vortex-array/src/expr/stats/mod.rs b/vortex-array/src/expr/stats/mod.rs index 57e11135ef3..473668d6ad1 100644 --- a/vortex-array/src/expr/stats/mod.rs +++ b/vortex-array/src/expr/stats/mod.rs @@ -269,15 +269,18 @@ mod test { use enum_iterator::all; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::expr::stats::Stat; #[test] fn min_of_nulls_is_not_panic() { let min = PrimitiveArray::from_option_iter::([None, None, None, None]) .statistics() - .compute_as::(Stat::Min, &mut array_session().create_execution_ctx()); + .compute_as::( + Stat::Min, + &mut default_session_builder().build().create_execution_ctx(), + ); assert_eq!(min, None); } diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index df5ca44c555..94c75b3532f 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -91,13 +91,13 @@ pub use metadata::*; pub use smallvec; pub use vortex_array_macros::array_slots; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::Context; use crate::aggregate_fn::session::AggregateFnSession; use crate::arrow::ArrowSession; use crate::dtype::session::DTypeSession; use crate::memory::MemorySession; -use crate::optimizer::kernels::ArrayKernelsExt; use crate::optimizer::kernels::KernelSession; use crate::scalar_fn::session::ScalarFnSession; use crate::session::ArraySession; @@ -157,20 +157,20 @@ pub mod flatbuffers { /// /// If the session contains a [`KernelSession`], this registers into its registry. Sessions that use /// [`KernelSession::default`] already receive these built-in kernels. -pub fn initialize(session: &VortexSession) { - if session.kernels_opt().is_some() { +pub fn initialize(session: &mut VortexSessionBuilder) { + if session.get_mut_opt::().is_some() { arrays::initialize(session); } } -/// Builds a fresh [`VortexSession`] registered with all of vortex-array's built-in session +/// Builds a fresh [`VortexSessionBuilder`] registered with all of vortex-array's built-in session /// variables: arrays, dtypes, scalar functions, stats, optimizer kernels, aggregate functions, /// Arrow conversion, and memory. /// /// Each call returns an independent session (with its own registries), so callers may register /// additional encodings or kernels into it without affecting any other session. This does not /// register file, layout, or runtime state — those live in higher-level crates. -pub fn array_session() -> VortexSession { +pub fn default_session_builder() -> VortexSessionBuilder { VortexSession::builder() .with::() .with::() @@ -180,12 +180,12 @@ pub fn array_session() -> VortexSession { .with::() .with::() .with::() - .build() } // TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot invoke execute // from the new array encodings to support back-compat for legacy encodings. So we hold a session // here... -pub static LEGACY_SESSION: LazyLock = LazyLock::new(array_session); +pub static LEGACY_SESSION: LazyLock = + LazyLock::new(|| default_session_builder().build()); pub type ArrayContext = Context; diff --git a/vortex-array/src/mask.rs b/vortex-array/src/mask.rs index ee55abc9875..8e83375e121 100644 --- a/vortex-array/src/mask.rs +++ b/vortex-array/src/mask.rs @@ -104,16 +104,16 @@ mod tests { use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; fn ctx() -> ExecutionCtx { - array_session().create_execution_ctx() + default_session_builder().build().create_execution_ctx() } #[test] diff --git a/vortex-array/src/memory.rs b/vortex-array/src/memory.rs index 78713dfbde7..11fd090fa4d 100644 --- a/vortex-array/src/memory.rs +++ b/vortex-array/src/memory.rs @@ -18,7 +18,7 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_session::SessionExt; use vortex_session::SessionVar; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; /// Mutable host buffer contract used by [`WritableHostBuffer`]. pub trait HostBufferMut: Send + 'static { @@ -221,17 +221,23 @@ pub trait MemorySessionExt: SessionExt { fn allocator(&self) -> HostAllocatorRef { self.memory().allocator() } - - /// Returns a new session configured to use `allocator` as its host allocator. - fn with_allocator(self, allocator: HostAllocatorRef) -> VortexSession { - let mut builder = self.session().to_builder(); - builder.get_mut::().set_allocator(allocator); - builder.build() - } } impl MemorySessionExt for S {} +/// Extension trait for configuring session-scoped memory before a session is built. +pub trait MemorySessionBuilderExt { + /// Configure the host allocator. + fn with_allocator(self, allocator: HostAllocatorRef) -> Self; +} + +impl MemorySessionBuilderExt for VortexSessionBuilder { + fn with_allocator(mut self, allocator: HostAllocatorRef) -> Self { + self.get_mut::().set_allocator(allocator); + self + } +} + /// Default host allocator. #[derive(Debug, Default)] pub struct DefaultHostAllocator; diff --git a/vortex-array/src/normalize.rs b/vortex-array/src/normalize.rs index 7957adc5297..d45bff96bbc 100644 --- a/vortex-array/src/normalize.rs +++ b/vortex-array/src/normalize.rs @@ -127,7 +127,9 @@ mod tests { )? .into_array(); let allowed = HashSet::from_iter([array.encoding_id(), field.encoding_id()]); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let normalized = array.clone().normalize(&mut NormalizeOptions { allowed: &allowed, @@ -172,7 +174,9 @@ mod tests { )? .into_array(); let allowed = HashSet::from_iter([array.encoding_id(), unchanged.encoding_id()]); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let normalized = array.clone().normalize(&mut NormalizeOptions { allowed: &allowed, @@ -211,7 +215,9 @@ mod tests { assert_eq!(sliced.encoding_id(), Slice.id()); let allowed = HashSet::from_iter([Dict.id(), Primitive.id()]); - let mut ctx = crate::array_session().create_execution_ctx(); + let mut ctx = crate::default_session_builder() + .build() + .create_execution_ctx(); let normalized = sliced.normalize(&mut NormalizeOptions { allowed: &allowed, diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index 4e98e896b96..c67f866f528 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -34,7 +34,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_session::SessionExt; use vortex_session::SessionVar; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use vortex_session::registry::Id; use vortex_utils::aliases::DefaultHashBuilder; use vortex_utils::aliases::hash_map::HashMap; @@ -320,13 +320,13 @@ impl Default for KernelSession { fn default() -> Self { // `ArrayKernels::default` installs the built-in parent-reduce kernels. The execute-parent // kernels are registered by the per-encoding `initialize` functions, which operate on a - // session. `KernelSession` clones share their registry storage, so kernels registered into - // the temporary session land in `this.kernels`. + // session builder. `KernelSession` clones share their registry storage, so kernels + // registered into the temporary builder land in `this.kernels`. let this = Self { kernels: ArrayKernels::default(), }; - let session = VortexSession::empty().with_some(this.clone()); - crate::arrays::initialize(&session); + let mut session = VortexSessionBuilder::default().with_some(this.clone()); + crate::arrays::initialize(&mut session); this } } @@ -341,7 +341,13 @@ impl SessionVar for KernelSession { } } -/// Extension trait for accessing the optimizer kernel registry from a [`VortexSession`]. +/// Returns the active [`ArrayKernels`] registry for a session builder. +pub fn builder_kernels(session: &mut VortexSessionBuilder) -> &ArrayKernels { + session.get_mut::().kernels() +} + +/// Extension trait for accessing the optimizer kernel registry from a +/// [`VortexSession`][vortex_session::VortexSession]. pub trait ArrayKernelsExt: SessionExt { /// Returns the active [`ArrayKernels`] registry if the session contains a [`KernelSession`]. fn kernels_opt(&self) -> Option<&ArrayKernels> { @@ -372,18 +378,22 @@ mod tests { #[test] fn kernel_session_default_registers_builtin_kernels() { - let session = VortexSession::empty().with::(); + let session = VortexSession::builder().with::().build(); assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); } #[test] fn initialize_registers_builtin_kernels_into_empty_kernel_session() { - let session = VortexSession::empty().with_some(KernelSession::empty()); + let session = VortexSession::builder() + .with_some(KernelSession::empty()) + .build(); assert!(!session.kernels().has_execute_parent(Binary.id(), Bool.id())); - crate::initialize(&session); + let mut builder = VortexSession::builder().with_some(KernelSession::empty()); + crate::initialize(&mut builder); + let session = builder.build(); assert!(session.kernels().has_execute_parent(Binary.id(), Bool.id())); } diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index a9952a7e99b..12eef3634fa 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -1213,8 +1213,8 @@ mod test { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::patches::Patches; use crate::patches::PrimitiveArray; use crate::search_sorted::SearchResult; @@ -1222,7 +1222,7 @@ mod test { #[test] fn test_filter() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 100, 0, @@ -1251,7 +1251,7 @@ mod test { #[test] fn take_with_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1297,7 +1297,7 @@ mod test { #[test] fn take_search_with_nulls_chunked() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1343,7 +1343,7 @@ mod test { #[test] fn take_search_chunked_multiple_chunks() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 2048, 0, @@ -1372,7 +1372,7 @@ mod test { #[test] fn take_search_chunked_indices_with_no_patches() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 20, 0, @@ -1395,7 +1395,7 @@ mod test { #[test] fn take_search_chunked_interleaved() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 30, 0, @@ -1424,7 +1424,7 @@ mod test { #[test] fn test_take_search_multiple_chunk_offsets() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 1500, 0, @@ -1448,7 +1448,7 @@ mod test { #[test] fn test_slice() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1465,7 +1465,7 @@ mod test { #[test] fn doubly_sliced() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = buffer![15_u32, 135, 13531, 42].into_array(); let indices = buffer![10_u64, 11, 50, 100].into_array(); @@ -1489,7 +1489,7 @@ mod test { #[test] fn test_mask_all_true() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1506,7 +1506,7 @@ mod test { #[test] fn test_mask_all_false() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1539,7 +1539,7 @@ mod test { #[test] fn test_mask_partial() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1573,7 +1573,7 @@ mod test { #[test] fn test_mask_with_offset() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1605,7 +1605,7 @@ mod test { #[test] fn test_mask_nullable_values() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1642,7 +1642,7 @@ mod test { #[test] fn test_filter_keep_all() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1670,7 +1670,7 @@ mod test { #[test] fn test_filter_none() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1688,7 +1688,7 @@ mod test { #[test] fn test_filter_with_indices() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1716,7 +1716,7 @@ mod test { #[test] fn test_slice_full_range() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1742,7 +1742,7 @@ mod test { #[test] fn test_slice_partial() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1787,7 +1787,7 @@ mod test { #[test] fn test_slice_with_offset() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 5, // offset @@ -1815,7 +1815,7 @@ mod test { #[test] fn test_patch_values() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, @@ -1881,7 +1881,7 @@ mod test { #[test] fn test_mask_boundary_patches() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test masking patches at array boundaries let patches = Patches::new( 10, @@ -1912,7 +1912,7 @@ mod test { #[test] fn test_mask_all_patches_removed() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test when all patches are masked out let patches = Patches::new( 10, @@ -1933,7 +1933,7 @@ mod test { #[test] fn test_mask_no_patches_removed() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test when no patches are masked let patches = Patches::new( 10, @@ -1964,7 +1964,7 @@ mod test { #[test] fn test_mask_single_patch() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test with a single patch let patches = Patches::new( 5, @@ -1992,7 +1992,7 @@ mod test { #[test] fn test_mask_contiguous_patches() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test with contiguous patches let patches = Patches::new( 10, @@ -2023,7 +2023,7 @@ mod test { #[test] fn test_mask_with_large_offset() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test with a large offset that shifts all indices let patches = Patches::new( 20, @@ -2056,7 +2056,7 @@ mod test { #[test] #[should_panic(expected = "Filter mask length 5 does not match array length 10")] fn test_mask_wrong_length() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let patches = Patches::new( 10, 0, diff --git a/vortex-array/src/scalar/tests/mod.rs b/vortex-array/src/scalar/tests/mod.rs index 676363c853c..d991da6dd35 100644 --- a/vortex-array/src/scalar/tests/mod.rs +++ b/vortex-array/src/scalar/tests/mod.rs @@ -14,4 +14,5 @@ use std::sync::LazyLock; use vortex_session::VortexSession; -pub(crate) static SESSION: LazyLock = LazyLock::new(crate::array_session); +pub(crate) static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index 2151a310dde..7d300bd6a33 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -341,7 +341,8 @@ mod tests { use crate::test_harness::to_int_indices; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); #[test] fn test_display() { diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index b7f838c6d13..a5f626c65a4 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -758,13 +758,13 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; #[expect(deprecated)] use crate::canonical::ToCanonical as _; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::scalar::Scalar; @@ -772,7 +772,7 @@ mod tests { #[test] fn test_kleene_truth_table() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = BoolArray::from_iter([ Some(true), Some(true), @@ -835,7 +835,7 @@ mod tests { #[test] fn test_null_constant_kleene() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = BoolArray::from_iter([Some(false), Some(true), None]).into_array(); let null = ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), lhs.len()) .into_array(); @@ -869,22 +869,34 @@ mod tests { let r = r.to_bool().into_array(); let v0 = r - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_bool() .value(); let v1 = r - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_bool() .value(); let v2 = r - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_bool() .value(); let v3 = r - .execute_scalar(3, &mut array_session().create_execution_ctx()) + .execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_bool() .value(); @@ -913,22 +925,34 @@ mod tests { .into_array(); let v0 = r - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_bool() .value(); let v1 = r - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_bool() .value(); let v2 = r - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_bool() .value(); let v3 = r - .execute_scalar(3, &mut array_session().create_execution_ctx()) + .execute_scalar( + 3, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap() .as_bool() .value(); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare.rs b/vortex-array/src/scalar_fn/fns/binary/compare.rs index b7a53c21d62..cf0eae52147 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare.rs @@ -252,7 +252,6 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ListArray; use crate::arrays::ListViewArray; @@ -262,6 +261,7 @@ mod tests { use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; @@ -280,7 +280,7 @@ mod tests { #[test] fn test_bool_basic_comparisons() { - let ctx = &mut array_session().create_execution_ctx(); + let ctx = &mut default_session_builder().build().create_execution_ctx(); let arr = BoolArray::new( BitBuffer::from_iter([true, true, false, true, false]), Validity::from_iter([false, true, true, true, true]), @@ -357,7 +357,10 @@ mod tests { .unwrap(); assert_eq!(result.len(), 10); let scalar = result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); assert_eq!(scalar.as_bool().value(), Some(false)); } @@ -368,7 +371,7 @@ mod tests { #[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())] #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())] fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let res = left.binary(right, Operator::Eq).unwrap(); let expected = BoolArray::from_iter([true, true]); assert_arrays_eq!(res, expected, &mut ctx); @@ -376,7 +379,7 @@ mod tests { #[test] fn test_list_array_comparison() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values1 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets1 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list1 = ListArray::try_new( @@ -421,7 +424,7 @@ mod tests { #[test] fn test_list_array_constant_comparison() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); let offsets = PrimitiveArray::from_iter([0i32, 2, 4, 6]); let list = ListArray::try_new( @@ -448,7 +451,7 @@ mod tests { #[test] fn test_struct_array_comparison() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bool_field1 = BoolArray::from_iter([Some(true), Some(false), Some(true)]); let int_field1 = PrimitiveArray::from_iter([1i32, 2, 3]); @@ -485,7 +488,7 @@ mod tests { #[test] fn test_empty_struct_compare() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let empty1 = StructArray::try_new( FieldNames::from(Vec::::new()), Vec::new(), @@ -514,7 +517,7 @@ mod tests { /// different Vortex encodings (VarBinArray vs VarBinViewArray) must not panic. #[test] fn struct_compare_mixed_binary_encodings() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // LHS: struct with a VarBinArray (offset-based) binary field let bin_field1 = VarBinArray::from(vec![ "apple".as_bytes(), @@ -570,7 +573,7 @@ mod tests { #[test] fn test_empty_list() { - let ctx = &mut array_session().create_execution_ctx(); + let ctx = &mut default_session_builder().build().create_execution_ctx(); let list = ListViewArray::new( BoolArray::from_iter(Vec::::new()).into_array(), buffer![0i32, 0i32, 0i32].into_array(), diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 156540a2820..ecc3d2cb1e0 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -282,9 +282,9 @@ mod tests { use super::*; use crate::VortexSessionExecute; - use crate::array_session; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -499,7 +499,10 @@ mod tests { let result_equal = lhs_struct.binary(rhs_struct_equal, Operator::Eq).unwrap(); assert_eq!( result_equal - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("value"), Scalar::bool(true, Nullability::NonNullable), "Equal structs should be equal" @@ -510,7 +513,10 @@ mod tests { .unwrap(); assert_eq!( result_different - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .vortex_expect("value"), Scalar::bool(false, Nullability::NonNullable), "Different structs should not be equal" @@ -519,7 +525,7 @@ mod tests { #[test] fn test_or_kleene_validity() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use crate::IntoArray; use crate::arrays::BoolArray; use crate::arrays::StructArray; @@ -547,7 +553,7 @@ mod tests { #[test] fn test_scalar_subtract_unsigned() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -562,7 +568,7 @@ mod tests { #[test] fn test_scalar_subtract_signed() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; @@ -577,7 +583,7 @@ mod tests { #[test] fn test_scalar_subtract_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use crate::IntoArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; @@ -594,7 +600,7 @@ mod tests { #[test] fn test_scalar_subtract_float() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); use vortex_buffer::buffer; use crate::IntoArray; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index e95627155e9..8a12d57ffb6 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -928,11 +928,11 @@ mod test { use crate::IntoArray; use crate::RecursiveCanonical; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::builtins::ArrayBuiltins; + use crate::default_session_builder; use crate::scalar::Scalar; use crate::scalar_fn::fns::operators::Operator; use crate::validity::Validity; @@ -944,14 +944,16 @@ mod test { Operator::Sub, ) .and_then(|a| { - a.execute::(&mut array_session().create_execution_ctx()) + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) }) .map(|a| a.0.into_array()) } #[test] fn test_scalar_subtract_unsigned() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = buffer![1u16, 2, 3].into_array(); let result = sub_scalar(&values, 1u16).unwrap(); assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]), &mut ctx); @@ -959,7 +961,7 @@ mod test { #[test] fn test_scalar_subtract_signed() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = buffer![1i64, 2, 3].into_array(); let result = sub_scalar(&values, -1i64).unwrap(); assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]), &mut ctx); @@ -967,7 +969,7 @@ mod test { #[test] fn test_scalar_subtract_nullable() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]); let result = sub_scalar(&values.into_array(), Some(1u16)).unwrap(); assert_arrays_eq!( @@ -979,7 +981,7 @@ mod test { #[test] fn test_scalar_subtract_float() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = buffer![1.0f64, 2.0, 3.0].into_array(); let result = sub_scalar(&values, -1f64).unwrap(); assert_arrays_eq!( @@ -998,14 +1000,18 @@ mod test { #[test] fn test_float_divide_by_zero_is_ok() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = buffer![1.0f64, -1.0].into_array(); let result = values .binary( ConstantArray::new(0.0f64, values.len()).into_array(), Operator::Div, ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .and_then(|a| { + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) + }) .unwrap(); assert_arrays_eq!( @@ -1023,7 +1029,11 @@ mod test { ConstantArray::new(1u8, values.len()).into_array(), Operator::Add, ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + .and_then(|a| { + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) + }); assert!(result.is_err()); } @@ -1036,7 +1046,11 @@ mod test { ConstantArray::new(0i32, values.len()).into_array(), Operator::Div, ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + .and_then(|a| { + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) + }); assert!(result.is_err()); } @@ -1049,21 +1063,27 @@ mod test { ConstantArray::new(-1i64, values.len()).into_array(), Operator::Div, ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + .and_then(|a| { + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) + }); assert!(result.is_err()); } #[test] fn test_integer_divide_errors_ignore_null_lanes() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let lhs = PrimitiveArray::new(buffer![10i32, 10], Validity::from_iter([false, true])) .into_array(); let rhs = buffer![0i32, 2].into_array(); let result = lhs .binary(rhs, Operator::Div) .and_then(|a| { - a.execute::(&mut array_session().create_execution_ctx()) + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) }) .map(|a| a.0.into_array()) .unwrap(); @@ -1077,7 +1097,7 @@ mod test { #[test] fn test_integer_errors_ignore_null_lanes() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true])) .into_array(); let result = values @@ -1086,7 +1106,9 @@ mod test { Operator::Add, ) .and_then(|a| { - a.execute::(&mut array_session().create_execution_ctx()) + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) }) .map(|a| a.0.into_array()) .unwrap(); @@ -1106,23 +1128,29 @@ mod test { ) .into_array(); let rhs = buffer![1u8, 1, 1].into_array(); - let result = lhs - .binary(rhs, Operator::Add) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())); + let result = lhs.binary(rhs, Operator::Add).and_then(|a| { + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) + }); assert!(result.is_err()); } #[test] fn test_present_nullable_constant_preserves_nullable_output() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = buffer![1u8, 2].into_array(); let result = values .binary( ConstantArray::new(Some(1u8), values.len()).into_array(), Operator::Add, ) - .and_then(|a| a.execute::(&mut array_session().create_execution_ctx())) + .and_then(|a| { + a.execute::( + &mut default_session_builder().build().create_execution_ctx(), + ) + }) .unwrap(); assert_arrays_eq!( diff --git a/vortex-array/src/scalar_fn/fns/byte_length.rs b/vortex-array/src/scalar_fn/fns/byte_length.rs index cfab5290a8f..1cb7f89d3d2 100644 --- a/vortex-array/src/scalar_fn/fns/byte_length.rs +++ b/vortex-array/src/scalar_fn/fns/byte_length.rs @@ -181,12 +181,12 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::byte_length; @@ -202,7 +202,7 @@ mod tests { #[case] array: ArrayRef, #[case] expected_lens: Vec, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(expected_lens); assert_arrays_eq!(result, expected, &mut ctx); @@ -211,7 +211,7 @@ mod tests { #[test] fn test_varbinview_byte_length() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = VarBinViewArray::from_iter_str(["short", "a longer string here"]).into_array(); let result = array.apply(&byte_length(root()))?; let expected = PrimitiveArray::from_iter(vec![5u64, 20]); @@ -225,16 +225,22 @@ mod tests { .into_array(); let result = array.apply(&byte_length(root()))?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!(result.is_valid(0, &mut ctx)?); assert!(!result.is_valid(1, &mut ctx)?); assert!(result.is_valid(2, &mut ctx)?); assert_eq!( - result.execute_scalar(0, &mut array_session().create_execution_ctx())?, + result.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + )?, Scalar::primitive(5u64, Nullability::Nullable), ); assert_eq!( - result.execute_scalar(2, &mut array_session().create_execution_ctx())?, + result.execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + )?, Scalar::primitive(18u64, Nullability::Nullable), ); Ok(()) @@ -245,7 +251,7 @@ mod tests { let null_scalar = Scalar::null(DType::Utf8(Nullability::Nullable)); let array = ConstantArray::new(null_scalar, 2).into_array(); let result = array.apply(&byte_length(root()))?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!(!result.is_valid(0, &mut ctx)?); assert!(!result.is_valid(1, &mut ctx)?); Ok(()) diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index d57026e00f5..b13a6ce0cc6 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -473,7 +473,8 @@ mod tests { use crate::expr::test_harness; use crate::scalar::Scalar; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); /// Helper to evaluate an expression using the apply+execute pattern fn evaluate_expr(expr: &Expression, array: &ArrayRef) -> ArrayRef { diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 74058aa19f2..c6477e8b8de 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -278,9 +278,9 @@ mod tests { use super::*; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -305,7 +305,7 @@ mod tests { #[test] fn execute_with_value() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -321,7 +321,7 @@ mod tests { #[test] fn execute_without_value_default_true() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -337,7 +337,7 @@ mod tests { #[test] fn execute_without_value_default_false() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let input = buffer![1i32, 5, 10].into_array(); let expr = dynamic( CompareOperator::Lt, @@ -357,7 +357,7 @@ mod tests { #[test] fn execute_value_flips() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let threshold = Arc::new(AtomicI32::new(5)); let threshold_clone = Arc::clone(&threshold); let expr = dynamic( diff --git a/vortex-array/src/scalar_fn/fns/ext_storage.rs b/vortex-array/src/scalar_fn/fns/ext_storage.rs index 4a014a3325b..419bb27f0fe 100644 --- a/vortex-array/src/scalar_fn/fns/ext_storage.rs +++ b/vortex-array/src/scalar_fn/fns/ext_storage.rs @@ -142,7 +142,9 @@ mod tests { assert_arrays_eq!( result, storage, - &mut crate::array_session().create_execution_ctx() + &mut crate::default_session_builder() + .build() + .create_execution_ctx() ); Ok(()) } @@ -162,7 +164,9 @@ mod tests { assert_arrays_eq!( result, storage, - &mut crate::array_session().create_execution_ctx() + &mut crate::default_session_builder() + .build() + .create_execution_ctx() ); Ok(()) } @@ -183,7 +187,9 @@ mod tests { assert_arrays_eq!( result, ConstantArray::new(storage_scalar, 3), - &mut crate::array_session().create_execution_ctx() + &mut crate::default_session_builder() + .build() + .create_execution_ctx() ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index c9eae72ad08..aa399d9da5e 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -180,10 +180,10 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -210,7 +210,7 @@ mod tests { #[test] fn evaluate() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let test_array = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), None, Some(5)]) .into_array(); @@ -231,7 +231,7 @@ mod tests { #[test] fn evaluate_struct_field() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let test_array = StructArray::from_fields(&[( "a", PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), @@ -251,7 +251,7 @@ mod tests { #[test] fn evaluate_non_nullable_input() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let test_array = buffer![1i32, 2, 3].into_array(); let expr = fill_null(root(), lit(0i32)); let result = test_array.apply(&expr).unwrap(); diff --git a/vortex-array/src/scalar_fn/fns/is_not_null.rs b/vortex-array/src/scalar_fn/fns/is_not_null.rs index 75d86db5ee8..e0f411b56b4 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -108,9 +108,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::col; @@ -129,7 +129,7 @@ mod tests { use crate::stats::null_count; static STATS_SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + LazyLock::new(|| VortexSession::builder().with::().build()); #[test] fn dtype() { @@ -162,7 +162,10 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); @@ -179,7 +182,10 @@ mod tests { for i in 0..result.len() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(true, Nullability::NonNullable) ); @@ -198,7 +204,10 @@ mod tests { for i in 0..result.len() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(false, Nullability::NonNullable) ); @@ -227,7 +236,10 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/fns/is_null.rs index 5e605ca6ee0..4d47bf84397 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -99,9 +99,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::col; @@ -118,7 +118,7 @@ mod tests { use crate::stats::null_count; static STATS_SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + LazyLock::new(|| VortexSession::builder().with::().build()); #[test] fn dtype() { @@ -151,7 +151,10 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); @@ -169,7 +172,10 @@ mod tests { for i in 0..result.len() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(false, Nullability::NonNullable) ); @@ -189,7 +195,10 @@ mod tests { for i in 0..result.len() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(true, Nullability::NonNullable) ); @@ -218,7 +227,10 @@ mod tests { for (i, expected_value) in expected.iter().enumerate() { assert_eq!( result - .execute_scalar(i, &mut array_session().create_execution_ctx()) + .execute_scalar( + i, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(*expected_value, Nullability::NonNullable) ); diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index 63ced07c8e7..fb17158bdbf 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -245,9 +245,9 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::BoolArray; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::get_item; @@ -262,7 +262,7 @@ mod tests { fn invert_booleans() { let not_expr = not(root()); let bools = BoolArray::from_iter([false, true, false, false, true, true]); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_arrays_eq!( bools.into_array().apply(¬_expr).unwrap(), BoolArray::from_iter([true, false, true, true, false, false]), diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 5185e513e70..192c767ead8 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -406,12 +406,12 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ListArray; use crate::arrays::VarBinArray; use crate::assert_arrays_eq; #[expect(deprecated)] use crate::canonical::ToCanonical as _; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType::I32; @@ -437,7 +437,7 @@ mod tests { use crate::validity::Validity; static STATS_SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + LazyLock::new(|| VortexSession::builder().with::().build()); fn stat(expr: Expression, stat: Stat) -> Expression { stat_expr(expr, stat.aggregate_fn().unwrap()) @@ -461,13 +461,19 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(true, Nullability::Nullable) ); assert_eq!( - item.execute_scalar(1, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(false, Nullability::Nullable) ); } @@ -480,13 +486,19 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(true, Nullability::Nullable) ); assert_eq!( - item.execute_scalar(1, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(true, Nullability::Nullable) ); } @@ -499,13 +511,19 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(false, Nullability::Nullable) ); assert_eq!( - item.execute_scalar(1, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(false, Nullability::Nullable) ); } @@ -524,13 +542,19 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(true, Nullability::Nullable) ); assert_eq!( - item.execute_scalar(1, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(false, Nullability::Nullable) ); } @@ -549,13 +573,19 @@ mod tests { let item = arr.apply(&expr).unwrap(); assert_eq!( - item.execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(), + item.execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) + .unwrap(), Scalar::bool(true, Nullability::Nullable) ); assert!( !item - .is_valid(1, &mut array_session().create_execution_ctx()) + .is_valid( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } @@ -647,7 +677,10 @@ mod tests { let result = arr.clone().apply(&expr).unwrap(); assert_eq!( result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(true, Nullability::NonNullable) ); @@ -657,7 +690,10 @@ mod tests { let result = arr.apply(&expr).unwrap(); assert_eq!( result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::bool(false, Nullability::NonNullable) ); @@ -752,7 +788,7 @@ mod tests { #[case] value: Option<&str>, #[case] expected: BoolArray, ) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let element_nullability = list_array .dtype() .as_list_element_opt() @@ -770,7 +806,7 @@ mod tests { #[test] fn test_constant_list() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let list_array = ConstantArray::new( Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -789,7 +825,7 @@ mod tests { #[test] fn test_all_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let list_array = ConstantArray::new( Scalar::null(DType::List( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), @@ -811,7 +847,7 @@ mod tests { #[test] fn test_list_array_element() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let list_scalar = Scalar::list( Arc::new(DType::Primitive(I32, Nullability::NonNullable)), vec![1.into(), 3.into(), 6.into()], @@ -828,7 +864,7 @@ mod tests { #[test] fn test_list_contains_empty_listview() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let empty_elements = PrimitiveArray::empty::(Nullability::NonNullable); let offsets = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); let sizes = Buffer::from_iter([0u32, 0, 0, 0]).into_array(); @@ -852,7 +888,7 @@ mod tests { #[test] fn test_list_contains_all_null_elements() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = PrimitiveArray::from_option_iter::([None, None, None, None, None]); let offsets = Buffer::from_iter([0u32, 2, 4]).into_array(); let sizes = Buffer::from_iter([2u32, 2, 1]).into_array(); @@ -888,7 +924,7 @@ mod tests { #[test] fn test_list_contains_large_offsets() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = Buffer::from_iter([1i32, 2, 3, 4, 5]).into_array(); let offsets = Buffer::from_iter([0u32, 1, 4, 0]).into_array(); @@ -912,7 +948,7 @@ mod tests { #[test] fn test_list_contains_offset_size_boundary() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = Buffer::from_iter(0..256).into_array(); let offsets = Buffer::from_iter([0u8, 100, 200, 254]).into_array(); let sizes = Buffer::from_iter([50u8, 50, 54, 2]).into_array(); diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index c9f5783de27..cb472f2cc3a 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -280,10 +280,10 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType::I32; @@ -320,7 +320,7 @@ mod tests { #[test] pub fn test_merge_right_most() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let expr = merge_opts( vec![ get_item("0", root()), diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 81019e22d6f..970460e9acb 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -173,10 +173,10 @@ mod tests { #[expect(deprecated)] use crate::ToCanonical as _; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; use crate::arrays::struct_::StructArrayExt; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::expr::col; use crate::expr::pack; @@ -232,7 +232,7 @@ mod tests { #[test] pub fn test_simple_pack() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), @@ -266,7 +266,7 @@ mod tests { #[test] pub fn test_nested_pack() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let expr = Pack.new_expr( PackOptions { names: ["one", "two", "three"].into(), diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index 1a4854a6a4a..48ccc2cf693 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -402,7 +402,6 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; @@ -411,6 +410,7 @@ mod tests { use crate::arrays::variant::VariantArrayExt; use crate::assert_arrays_eq; use crate::assert_nth_scalar_is_null; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; @@ -545,7 +545,7 @@ mod tests { let expr = variant_get(root(), parse_path(path)?, dtype); array .apply(&expr)? - .execute::(&mut array_session().create_execution_ctx()) + .execute::(&mut default_session_builder().build().create_execution_ctx()) } #[test] @@ -636,7 +636,7 @@ mod tests { Some(DType::Primitive(PType::I32, Nullability::NonNullable)), ); let proto = expr.serialize_proto().unwrap(); - let actual = Expression::from_proto(&proto, &array_session()).unwrap(); + let actual = Expression::from_proto(&proto, &default_session_builder().build()).unwrap(); assert_eq!(actual, expr); } @@ -678,7 +678,7 @@ mod tests { "$.items[1]", Some(DType::Primitive(PType::I32, Nullability::NonNullable)), )?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_arrays_eq!( result, @@ -712,7 +712,7 @@ mod tests { "$.a", Some(DType::Primitive(PType::I32, Nullability::NonNullable)), )?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_arrays_eq!( result, @@ -746,7 +746,7 @@ mod tests { )?; assert!(!result.is::()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_arrays_eq!( result, PrimitiveArray::from_option_iter([Some(10i32), Some(20), None]), @@ -772,7 +772,7 @@ mod tests { let result = execute_variant_get(array, "$.a", None)?; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let row0 = result.execute_scalar(0, &mut ctx)?; assert_eq!( row0.as_variant() @@ -807,10 +807,11 @@ mod tests { ])?; let result = execute_variant_get(array, "$.a", None)?; - let variant = result - .clone() - .execute::(&mut array_session().create_execution_ctx())?; - let canonical = result.execute::(&mut array_session().create_execution_ctx())?; + let variant = result.clone().execute::( + &mut default_session_builder().build().create_execution_ctx(), + )?; + let canonical = result + .execute::(&mut default_session_builder().build().create_execution_ctx())?; let Canonical::Variant(canonical_variant) = canonical else { vortex_bail!("expected Variant canonical array"); }; @@ -820,7 +821,7 @@ mod tests { assert_eq!(variant.core_storage().dtype(), variant.dtype()); assert_eq!(variant.core_storage().len(), variant.len()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); for (idx, expected) in [10i32, 20].into_iter().enumerate() { let scalar = variant.execute_scalar(idx, &mut ctx)?; let actual = scalar diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 48449bc8428..7c1fcd811a2 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -277,7 +277,6 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::Struct; @@ -290,6 +289,7 @@ mod tests { use crate::builders::VarBinViewBuilder; use crate::builtins::ArrayBuiltins; use crate::columnar::Columnar; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -317,7 +317,7 @@ mod tests { #[test] fn test_zip_basic() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mask = Mask::from_iter([true, false, false, true, false]); let if_true = buffer![10, 20, 30, 40, 50].into_array(); let if_false = buffer![1, 2, 3, 4, 5].into_array(); @@ -330,7 +330,7 @@ mod tests { #[test] fn test_zip_all_true() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mask = Mask::new_true(4); let if_true = buffer![10, 20, 30, 40].into_array(); let if_false = @@ -346,7 +346,7 @@ mod tests { #[test] fn test_zip_all_false_widens_nullability() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mask = Mask::new_false(4); let if_true = PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array(); @@ -367,7 +367,7 @@ mod tests { let if_false = PrimitiveArray::from_option_iter([Some(1), Some(2), Some(3), None]).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?; assert_arrays_eq!( result, @@ -386,7 +386,7 @@ mod tests { PrimitiveArray::from_option_iter([Some(10), Some(20), Some(30), None]).into_array(); let if_false = buffer![1i32, 2, 3, 4].into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = zip_impl(&if_true, &if_false, &mask, &mut ctx)?; assert_arrays_eq!( result, @@ -427,7 +427,7 @@ mod tests { let mask = Mask::from_indices(len, indices); let mask_array = mask.into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let result = mask_array .zip(const1.clone(), const2.clone())? .execute::(&mut ctx)? @@ -487,7 +487,7 @@ mod tests { let mask = Mask::from_indices(200, (0..100).filter(|i| i % 3 != 0)); let mask_array = mask.clone().into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let zipped = mask_array .zip(if_true.clone(), if_false.clone()) .unwrap() @@ -496,25 +496,29 @@ mod tests { let zipped = zipped.as_opt::().unwrap(); assert_eq!(zipped.data_buffers().len(), 2); - let mut arrow_ctx = array_session().create_execution_ctx(); + let mut arrow_ctx = default_session_builder().build().create_execution_ctx(); let expected = arrow_zip( - array_session() + default_session_builder() + .build() .arrow() .execute_arrow(mask.into_array(), None, &mut arrow_ctx) .unwrap() .as_boolean(), - &array_session() + &default_session_builder() + .build() .arrow() .execute_arrow(if_true, None, &mut arrow_ctx) .unwrap(), - &array_session() + &default_session_builder() + .build() .arrow() .execute_arrow(if_false, None, &mut arrow_ctx) .unwrap(), ) .unwrap(); - let actual = array_session() + let actual = default_session_builder() + .build() .arrow() .execute_arrow(zipped.array().clone(), None, &mut arrow_ctx) .unwrap(); diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 6e143eaa4dd..b9b654ad36b 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -129,14 +129,16 @@ mod tests { #[test] fn array_session_default_registers_encodings() { - let session = VortexSession::empty().with::(); + let session = VortexSession::builder().with::().build(); assert!(session.arrays().registry().find(&Bool.id()).is_some()); } #[test] fn empty_array_session_registers_no_encodings() { - let session = VortexSession::empty().with_some(ArraySession::empty()); + let session = VortexSession::builder() + .with_some(ArraySession::empty()) + .build(); assert!(session.arrays().registry().find(&Bool.id()).is_none()); } diff --git a/vortex-array/src/stats/expr.rs b/vortex-array/src/stats/expr.rs index 2038df0b613..d475e1a46dc 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -89,13 +89,13 @@ mod tests { use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::Chunked; use crate::arrays::ChunkedArray; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; use crate::arrays::chunked::ChunkedArrayExt; use crate::assert_arrays_eq; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -106,7 +106,7 @@ mod tests { use crate::scalar::ScalarValue; use crate::validity::Validity; - static SESSION: LazyLock = LazyLock::new(array_session); + static SESSION: LazyLock = LazyLock::new(|| default_session_builder().build()); #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { diff --git a/vortex-array/src/stats/rewrite.rs b/vortex-array/src/stats/rewrite.rs index 2bbeeb00022..16183eea79f 100644 --- a/vortex-array/src/stats/rewrite.rs +++ b/vortex-array/src/stats/rewrite.rs @@ -200,7 +200,7 @@ mod tests { #[test] fn combines_multiple_falsifiers_with_or() -> VortexResult<()> { - let session = crate::array_session(); + let session = crate::default_session_builder().build(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); session.stats().register_rewrite(StaticLiteralRule { falsifier: Some(lit(false)), @@ -220,7 +220,7 @@ mod tests { #[test] fn combines_multiple_satisfiers_with_or() -> VortexResult<()> { - let session = crate::array_session(); + let session = crate::default_session_builder().build(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); session.stats().register_rewrite(StaticLiteralRule { falsifier: None, @@ -240,7 +240,7 @@ mod tests { #[test] fn unregistered_expression_has_no_rewrite() -> VortexResult<()> { - let session = crate::array_session(); + let session = crate::default_session_builder().build(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); assert_eq!(lit(true).falsify(&dtype, &session)?, None); @@ -250,7 +250,7 @@ mod tests { #[test] fn non_predicate_expression_errors() { - let session = crate::array_session(); + let session = crate::default_session_builder().build(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); assert!(lit(7).falsify(&dtype, &session).is_err()); diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index ea656f24546..16efa4743e4 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -738,7 +738,8 @@ mod tests { use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::internal::row_count::RowCount; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = + LazyLock::new(|| crate::default_session_builder().build()); fn stat(expr: Expression, stat: Stat) -> Expression { let aggregate_fn = stat.aggregate_fn().expect("stat should have aggregate fn"); diff --git a/vortex-array/src/stats/stats_set.rs b/vortex-array/src/stats/stats_set.rs index 6025afed2c6..0ef82247b15 100644 --- a/vortex-array/src/stats/stats_set.rs +++ b/vortex-array/src/stats/stats_set.rs @@ -566,8 +566,8 @@ mod test { use smallvec::smallvec; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -877,7 +877,10 @@ mod test { .collect_vec(); array .statistics() - .compute_all(&all_stats, &mut array_session().create_execution_ctx()) + .compute_all( + &all_stats, + &mut default_session_builder().build().create_execution_ctx(), + ) .unwrap(); let stats = array.statistics().to_owned(); diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs index 9c80e34267e..05241cfb290 100644 --- a/vortex-array/src/validity.rs +++ b/vortex-array/src/validity.rs @@ -651,8 +651,8 @@ mod tests { use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; - use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::default_session_builder; use crate::dtype::Nullability; use crate::validity::BoolArray; use crate::validity::Validity; @@ -720,7 +720,7 @@ mod tests { let indices = PrimitiveArray::new(Buffer::copy_from(positions), Validity::NonNullable).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!( validity @@ -734,7 +734,7 @@ mod tests { #[test] #[should_panic] fn out_of_bounds_patch() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); Validity::NonNullable .patch( 2, @@ -786,7 +786,7 @@ mod tests { #[case] indices: ArrayRef, #[case] expected: Validity, ) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert!( validity .take(&indices) @@ -833,7 +833,7 @@ mod tests { #[case] rhs: Validity, #[case] expected: bool, ) -> vortex_error::VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!(lhs.mask_eq(&rhs, 3, &mut ctx)?, expected); Ok(()) } diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 8981b4859cf..08c5253b5de 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -66,7 +66,7 @@ pub use output::BenchmarkOutput; pub use output::create_output_writer; use vortex::VortexSessionDefault; pub use vortex::error::vortex_panic; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; // All benchmarks run with mimalloc for consistency. @@ -74,7 +74,7 @@ use vortex::session::VortexSession; static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; pub static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_tokio()); + LazyLock::new(|| VortexSession::default_builder().with_tokio().build()); #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] pub struct Target { diff --git a/vortex-btrblocks/benches/compress.rs b/vortex-btrblocks/benches/compress.rs index 6898b4a9ab4..0741ae93287 100644 --- a/vortex-btrblocks/benches/compress.rs +++ b/vortex-btrblocks/benches/compress.rs @@ -22,7 +22,8 @@ mod benchmarks { use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn make_clickbench_window_name() -> ArrayRef { // A test that's meant to mirror the WindowName column from ClickBench. diff --git a/vortex-btrblocks/benches/compress_listview.rs b/vortex-btrblocks/benches/compress_listview.rs index 502acf79e0f..d1863373dba 100644 --- a/vortex-btrblocks/benches/compress_listview.rs +++ b/vortex-btrblocks/benches/compress_listview.rs @@ -29,7 +29,8 @@ mod benchmarks { const NUM_ROWS: usize = 8192; const SEED: u64 = 42; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); const SHORT_STRINGS: &[&str] = &[ "alpha_one", diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index bddb0be113a..6e78a8b1518 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -85,7 +85,8 @@ mod tests { #[cfg(feature = "zstd")] use crate::BtrBlocksCompressorBuilder; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[rstest] #[case::zctl( diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 69ddb59738c..ff8cb90d958 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -40,7 +40,7 @@ //! # Example //! //! ```rust -//! use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +//! use vortex_array::{IntoArray, VortexSessionExecute, default_session_builder}; //! use vortex_array::arrays::PrimitiveArray; //! use vortex_array::validity::Validity; //! use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt}; @@ -48,7 +48,7 @@ //! use vortex_buffer::buffer; //! //! # fn example() -> vortex_error::VortexResult<()> { -//! let session = array_session(); +//! let session = default_session_builder().build(); //! let array = PrimitiveArray::new(buffer![42u64; 1024], Validity::NonNullable).into_array(); //! //! let compressor = BtrBlocksCompressor::default(); diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index 4c7a5f85fa6..d4aa758d0a3 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -21,7 +21,8 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_constant_compressed() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index 2d0a04542cc..56bf822ef89 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -22,7 +22,8 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::float::FloatRLEScheme; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_empty() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index e4227a472ec..6ab3a79373b 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -28,7 +28,8 @@ use vortex_session::VortexSession; use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_constant_compressed() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index a9ef24dc0e0..3214e563efc 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -27,7 +27,8 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_empty() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index c1473146607..889698e3ac4 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -18,7 +18,8 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_constant_compressed() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/string/tests.rs b/vortex-btrblocks/src/schemes/string/tests.rs index ec9c4060f64..08f0a0e2e14 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -16,7 +16,8 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); #[test] fn test_strings() -> VortexResult<()> { diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index 31fb700a2d6..7fda9887313 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -24,7 +24,8 @@ use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); /// Helper: synthetic short-string corpus that the cascading compressor should /// route through OnPair. diff --git a/vortex-compressor/benches/dict_encode.rs b/vortex-compressor/benches/dict_encode.rs index 2c4e24108a7..e82cc445363 100644 --- a/vortex-compressor/benches/dict_encode.rs +++ b/vortex-compressor/benches/dict_encode.rs @@ -15,7 +15,8 @@ use vortex_array::validity::Validity; use vortex_buffer::BufferMut; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn make_array() -> PrimitiveArray { let values: BufferMut = (0..50).cycle().take(64_000).collect(); diff --git a/vortex-compressor/benches/stats_calc.rs b/vortex-compressor/benches/stats_calc.rs index abbc34817f1..64f4e7cd68d 100644 --- a/vortex-compressor/benches/stats_calc.rs +++ b/vortex-compressor/benches/stats_calc.rs @@ -16,7 +16,8 @@ mod benchmarks { use vortex_compressor::stats::IntegerStats; use vortex_session::VortexSession; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn generate_dataset(max_run: u32, distinct: u32) -> Buffer { let mut output = BufferMut::with_capacity(64_000); diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index 5288637a609..684ec5c8e7d 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -263,7 +263,9 @@ mod tests { #[test] fn test_float_dict_encode() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let values = buffer![1f32, 2f32, 2f32, 0f32, 1f32]; let validity = Validity::Array(BoolArray::from_iter([true, true, true, false, true]).into_array()); diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 8ec3bf53345..e120177c40f 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -269,7 +269,9 @@ mod tests { #[test] fn test_dict_encode_integer_stats() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let data = buffer![100i32, 200, 100, 0, 100]; let validity = Validity::Array(BoolArray::from_iter([true, true, true, false, true]).into_array()); diff --git a/vortex-compressor/src/compressor.rs b/vortex-compressor/src/compressor.rs index e66ffb053e7..8523a40c672 100644 --- a/vortex-compressor/src/compressor.rs +++ b/vortex-compressor/src/compressor.rs @@ -615,7 +615,8 @@ mod tests { use crate::estimate::WinnerEstimate; use crate::scheme::SchemeExt; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn compressor() -> CascadingCompressor { CascadingCompressor::new(vec![&IntDictScheme, &FloatDictScheme, &StringDictScheme]) diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 7e6854eacbe..66926dcd5d7 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -23,14 +23,14 @@ //! compression is selected. //! //! ```rust -//! use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +//! use vortex_array::{IntoArray, VortexSessionExecute, default_session_builder}; //! use vortex_array::arrays::PrimitiveArray; //! use vortex_array::validity::Validity; //! use vortex_buffer::buffer; //! use vortex_compressor::CascadingCompressor; //! //! # fn example() -> vortex_error::VortexResult<()> { -//! let session = array_session(); +//! let session = default_session_builder().build(); //! let array = PrimitiveArray::new(buffer![1i32, 2, 3], Validity::NonNullable).into_array(); //! let compressor = CascadingCompressor::new(Vec::new()); //! diff --git a/vortex-compressor/src/sample.rs b/vortex-compressor/src/sample.rs index deb5bbe6f37..6eb547a45e2 100644 --- a/vortex-compressor/src/sample.rs +++ b/vortex-compressor/src/sample.rs @@ -134,9 +134,9 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexResult; @@ -145,7 +145,7 @@ mod tests { #[test] fn sample_is_deterministic() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Create a deterministic array with linear-with-noise pattern let values: Vec = (0i64..100_000).map(|i| i + (i * 7 + 3) % 11).collect(); diff --git a/vortex-compressor/src/stats/bool.rs b/vortex-compressor/src/stats/bool.rs index fbbc4001b90..0f279d8666c 100644 --- a/vortex-compressor/src/stats/bool.rs +++ b/vortex-compressor/src/stats/bool.rs @@ -93,8 +93,8 @@ impl BoolStats { #[cfg(test)] mod tests { use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; + use vortex_array::default_session_builder; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_error::VortexResult; @@ -103,7 +103,7 @@ mod tests { #[test] fn test_all_true() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, true, true]), Validity::NonNullable, @@ -118,7 +118,7 @@ mod tests { #[test] fn test_all_false() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![false, false, false]), Validity::NonNullable, @@ -133,7 +133,7 @@ mod tests { #[test] fn test_mixed() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true]), Validity::NonNullable, @@ -148,7 +148,7 @@ mod tests { #[test] fn test_with_nulls() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = BoolArray::new( BitBuffer::from(vec![true, false, true]), Validity::from_iter([true, false, true]), diff --git a/vortex-compressor/src/stats/float.rs b/vortex-compressor/src/stats/float.rs index c505993e9e5..5650a46319a 100644 --- a/vortex-compressor/src/stats/float.rs +++ b/vortex-compressor/src/stats/float.rs @@ -277,8 +277,8 @@ where mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::default_session_builder; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -288,7 +288,7 @@ mod tests { #[test] fn test_float_stats() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let floats = buffer![0.0f32, 1.0f32, 2.0f32].into_array(); let floats = floats.execute::(&mut ctx)?; @@ -309,7 +309,7 @@ mod tests { #[test] fn test_float_stats_leading_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let floats = PrimitiveArray::new( buffer![0.0f32, 1.0f32, 2.0f32], Validity::from_iter([false, true, true]), diff --git a/vortex-compressor/src/stats/integer.rs b/vortex-compressor/src/stats/integer.rs index a399c1060f9..05c610d0a57 100644 --- a/vortex-compressor/src/stats/integer.rs +++ b/vortex-compressor/src/stats/integer.rs @@ -562,8 +562,8 @@ mod tests { use std::iter; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::default_session_builder; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; @@ -575,7 +575,7 @@ mod tests { #[test] fn test_naive_count_distinct_values() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new(buffer![217u8, 0], Validity::NonNullable); let stats = typed_int_stats::(&array, true, &mut ctx)?; assert_eq!(stats.distinct_count().unwrap(), 2); @@ -584,7 +584,7 @@ mod tests { #[test] fn test_naive_count_distinct_values_nullable() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new( buffer![217u8, 0], Validity::from(BitBuffer::from(vec![true, false])), @@ -596,7 +596,7 @@ mod tests { #[test] fn test_count_distinct_values() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new((0..128u8).collect::>(), Validity::NonNullable); let stats = typed_int_stats::(&array, true, &mut ctx)?; assert_eq!(stats.distinct_count().unwrap(), 128); @@ -605,7 +605,7 @@ mod tests { #[test] fn test_count_distinct_values_nullable() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new( (0..128u8).collect::>(), Validity::from(BitBuffer::from_iter( @@ -619,7 +619,7 @@ mod tests { #[test] fn test_integer_stats_leading_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ints = PrimitiveArray::new(buffer![0, 1, 2], Validity::from_iter([false, true, true])); let stats = IntegerStats::generate_opts( diff --git a/vortex-cuda/benches/alp_cuda.rs b/vortex-cuda/benches/alp_cuda.rs index 5a1fadb7570..b959b0e185f 100644 --- a/vortex-cuda/benches/alp_cuda.rs +++ b/vortex-cuda/benches/alp_cuda.rs @@ -22,8 +22,8 @@ use cudarc::driver::DeviceRepr; use futures::executor::block_on; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; -use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; +use vortex::array::default_session_builder; use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; use vortex::dtype::NativePType; @@ -74,7 +74,7 @@ where let encoded = alp_encode( primitive_array.as_view(), None, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), ) .vortex_expect("failed to ALP-encode array"); diff --git a/vortex-cuda/benches/arrow_validity_cuda.rs b/vortex-cuda/benches/arrow_validity_cuda.rs index d1142e361f4..ea745149622 100644 --- a/vortex-cuda/benches/arrow_validity_cuda.rs +++ b/vortex-cuda/benches/arrow_validity_cuda.rs @@ -26,7 +26,6 @@ use vortex::buffer::Buffer; use vortex::dtype::PType; use vortex::error::VortexExpect; use vortex::error::VortexResult; -use vortex::session::VortexSession; use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaSession; use vortex_cuda::arrow::ArrowDeviceArray; @@ -45,6 +44,11 @@ fn validity_bitmap_byte_len(len: usize, bit_offset: usize) -> usize { (bit_offset + len).div_ceil(8) } +fn create_cuda_execution_ctx() -> CudaExecutionCtx { + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context") +} + unsafe fn release_arrow_device_array(array: &mut ArrowDeviceArray) { unsafe { if let Some(release) = array.array.release { @@ -105,9 +109,7 @@ fn benchmark_arrow_validity_export(c: &mut Criterion) { &len, |b, &len| { b.iter_custom(|iters| { - let mut cuda_ctx = - CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context"); + let mut cuda_ctx = create_cuda_execution_ctx(); let array = block_on(primitive_with_device_bool_validity( len, validity_offset, @@ -157,9 +159,8 @@ fn benchmark_arrow_validity_repack(c: &mut Criterion) { b.iter_custom(|iters| { let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_launch_strategy(Arc::new(timed)); + let mut cuda_ctx = + create_cuda_execution_ctx().with_launch_strategy(Arc::new(timed)); let (input_offset, input_buffer) = block_on(device_validity_buffer(len, INPUT_OFFSET, &mut cuda_ctx)) .vortex_expect("failed to create validity fixture"); @@ -199,9 +200,8 @@ fn benchmark_arrow_validity_count_nulls(c: &mut Criterion) { b.iter_custom(|iters| { let timed = TimedLaunchStrategy::default(); let timer = timed.timer(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) - .vortex_expect("failed to create execution context") - .with_launch_strategy(Arc::new(timed)); + let mut cuda_ctx = + create_cuda_execution_ctx().with_launch_strategy(Arc::new(timed)); let (_, input_buffer) = block_on(device_validity_buffer(len, ARROW_OFFSET, &mut cuda_ctx)) .vortex_expect("failed to create validity fixture"); diff --git a/vortex-cuda/benches/bitpacked_cuda.rs b/vortex-cuda/benches/bitpacked_cuda.rs index f216c33e2dd..6a438a30c89 100644 --- a/vortex-cuda/benches/bitpacked_cuda.rs +++ b/vortex-cuda/benches/bitpacked_cuda.rs @@ -22,8 +22,8 @@ use cudarc::driver::DeviceRepr; use futures::executor::block_on; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; -use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; +use vortex::array::default_session_builder; use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; use vortex::dtype::NativePType; @@ -57,7 +57,7 @@ where .collect(); let primitive_array = PrimitiveArray::new(Buffer::from(values), NonNullable); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); BitPackedData::encode(&primitive_array.into_array(), bit_width, &mut ctx) .vortex_expect("failed to create BitPacked array") } @@ -98,7 +98,7 @@ where .collect(); let primitive_array = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); BitPackedData::encode(&primitive_array, bit_width, &mut ctx) .vortex_expect("failed to create BitPacked array with patches") } diff --git a/vortex-cuda/benches/dynamic_dispatch_cuda.rs b/vortex-cuda/benches/dynamic_dispatch_cuda.rs index 9945b016190..971b58a0bd0 100644 --- a/vortex-cuda/benches/dynamic_dispatch_cuda.rs +++ b/vortex-cuda/benches/dynamic_dispatch_cuda.rs @@ -28,10 +28,10 @@ use futures::executor::block_on; use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; -use vortex::array::array_session; use vortex::array::arrays::DictArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::buffer::BufferHandle; +use vortex::array::default_session_builder; use vortex::array::scalar::Scalar; use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; @@ -210,7 +210,7 @@ fn bench_for_bitpacked(c: &mut Criterion) { .map(|i| (i as u64 % (max_val + 1)) as u32) .collect(); let prim = PrimitiveArray::new(Buffer::from(residuals), NonNullable); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bp = BitPackedData::encode(&prim.into_array(), bit_width, &mut ctx).vortex_expect("bitpack"); let array = FoR::try_new(bp.into_array(), Scalar::from(reference)) @@ -255,7 +255,7 @@ fn bench_dict_bp_codes(c: &mut Criterion) { let codes: Vec = (0..*len).map(|i| (i % dict_size) as u32).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let codes_bp = BitPackedData::encode(&codes_prim.into_array(), dict_bit_width, &mut ctx) .vortex_expect("bitpack codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); @@ -300,7 +300,7 @@ fn bench_runend(c: &mut Criterion) { let ends: Vec = (1..=num_runs).map(|i| (i * run_len) as u32).collect(); let values: Vec = (0..num_runs).map(|i| (i * 7 + 42) as u32).collect(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array(); let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); let re = RunEnd::new(ends_arr, values_arr, &mut ctx); @@ -347,7 +347,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_dynanmic_dispatch(c: &mut Criterion) { .map(|i| ::decode_single(10 + i as i32, exponents)) .collect(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // values: ALP → FoR → BitPacked. let float_prim = PrimitiveArray::new(Buffer::from(dict_floats), NonNullable); @@ -663,7 +663,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_composed_standalone(c: &mut Criterion) .map(|i| ::decode_single(10 + (i as i32 % 64), exponents)) .collect(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let float_prim = PrimitiveArray::new(Buffer::from(dict_floats), NonNullable); let alp = alp_encode(float_prim.as_view(), Some(exponents), &mut ctx).vortex_expect("alp_encode"); @@ -735,7 +735,7 @@ fn bench_dict_bp_codes_alp_for_bp_values_composed_standalone(c: &mut Criterion) // Benchmark: ALP(FoR(BitPacked)) — f64 // --------------------------------------------------------------------------- fn bench_alp_for_bitpacked_f64(c: &mut Criterion) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut group = c.benchmark_group("cuda"); let exponents = Exponents { e: 2, f: 0 }; @@ -810,7 +810,7 @@ fn bench_dict_bp_codes_bp_for_values(c: &mut Criterion) { // Dict values: residuals 0..63 bitpacked, FoR adds 1_000_000 let dict_residuals: Vec = (0..dict_size as u32).collect(); let dict_prim = PrimitiveArray::new(Buffer::from(dict_residuals), NonNullable); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict_bp = BitPackedData::encode(&dict_prim.into_array(), dict_bit_width, &mut ctx) .vortex_expect("bitpack dict"); let dict_for = @@ -857,7 +857,7 @@ fn bench_dict_bp_codes_bp_for_values(c: &mut Criterion) { // Benchmark: ALP(FoR(BitPacked)) for f32 // --------------------------------------------------------------------------- fn bench_alp_for_bitpacked(c: &mut Criterion) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut group = c.benchmark_group("cuda"); let exponents = Exponents { e: 2, f: 0 }; @@ -935,7 +935,7 @@ fn bench_dict_bp_u8_codes_u32_values(c: &mut Criterion) { let codes: Vec = (0..*len).map(|i| (i % dict_size) as u8).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let codes_bp = BitPackedData::encode(&codes_prim.into_array(), bit_width, &mut ctx) .vortex_expect("bitpack u8 codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); @@ -978,7 +978,7 @@ fn bench_dict_bp_u16_codes_u32_values(c: &mut Criterion) { let codes: Vec = (0..*len).map(|i| (i % dict_size) as u16).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let codes_bp = BitPackedData::encode(&codes_prim.into_array(), bit_width, &mut ctx) .vortex_expect("bitpack u16 codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); @@ -1021,7 +1021,7 @@ fn bench_dict_bp_u32_codes_u32_values(c: &mut Criterion) { let codes: Vec = (0..*len).map(|i| (i % dict_size) as u32).collect(); let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let codes_bp = BitPackedData::encode(&codes_prim.into_array(), bit_width, &mut ctx) .vortex_expect("bitpack u32 codes"); let values_prim = PrimitiveArray::new(Buffer::from(dict_values.clone()), NonNullable); diff --git a/vortex-cuda/benches/for_cuda.rs b/vortex-cuda/benches/for_cuda.rs index 31acc764141..948dee87e32 100644 --- a/vortex-cuda/benches/for_cuda.rs +++ b/vortex-cuda/benches/for_cuda.rs @@ -22,8 +22,8 @@ use cudarc::driver::DeviceRepr; use futures::executor::block_on; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; -use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; +use vortex::array::default_session_builder; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::dtype::NativePType; @@ -58,7 +58,7 @@ where PrimitiveArray::new(Buffer::from(data), Validity::NonNullable).into_array(); if bp && T::PTYPE != PType::U8 { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let child = BitPackedData::encode(&primitive_array, 8, &mut ctx).vortex_expect("failed to bitpack"); FoR::try_new(child.into_array(), reference.into()) diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index b5718d3d5fa..f96e9fdb4cc 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -14,6 +14,7 @@ use std::ptr; use arrow_schema::ffi::FFI_ArrowSchema; use vortex::error::VortexResult; use vortex::error::vortex_ensure; +use vortex::error::vortex_err; use vortex::session::VortexSession; use vortex_cuda::CudaSession; use vortex_cuda::arrow::ArrowDeviceArray; @@ -34,16 +35,14 @@ use vortex_ffi::vx_session_ref; const VX_CUDA_OK: c_int = 0; const VX_CUDA_ERR: c_int = 1; -/// Return a Vortex session with a [`CudaSession`] session variable. -/// -/// If `session` already has CUDA support, this returns a clone of it. Otherwise it -/// returns a new session cloned from `session` with a default [`CudaSession`] attached. -fn session_with_cuda(session: &VortexSession) -> VortexResult { - if session.get_opt::().is_some() { - return Ok(session.clone()); +fn session_with_cuda(session: &VortexSession) -> VortexResult<&VortexSession> { + if session.get_opt::().is_none() { + return Err(vortex_err!( + InvalidArgument: "CUDA session state is not configured; create the session with vx_cuda_session_new" + )); } - Ok(session.clone().with_some(CudaSession::try_default()?)) + Ok(session) } /// Create a CUDA Vortex session. @@ -97,7 +96,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_export_arrow_device( let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; let array = unsafe { vx_array_ref(array) }?.clone(); - let mut ctx = CudaSession::create_execution_ctx(&session)?; + let mut ctx = CudaSession::create_execution_ctx(session)?; let exported = futures::executor::block_on(array.export_device_array_with_schema(&mut ctx))?; @@ -142,7 +141,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_partition_scan_arrow_device_stream( let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; // Drive the stream on the same runtime the partition's scan spawned its work onto. - let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?; + let device_stream = array_stream.export_device_array_stream(session, ffi_runtime())?; unsafe { ptr::write(out_stream, device_stream) }; Ok(VX_CUDA_OK) @@ -156,7 +155,6 @@ mod tests { use arrow_schema::Field; use arrow_schema::Schema; - use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::IntoArray; use vortex::array::arrays::PrimitiveArray; @@ -173,6 +171,12 @@ mod tests { Box::into_raw(Box::new(session)).cast::() } + fn test_cuda_session() -> VortexResult { + Ok(vortex::array::default_session_builder() + .with_some(CudaSession::try_default()?) + .build()) + } + unsafe fn free_test_session(session: *mut vx_session) { unsafe { drop(Box::from_raw(session.cast::())) }; } @@ -212,9 +216,9 @@ mod tests { } #[cuda_test] - fn test_export_primitive_arrow_device() { + fn test_export_primitive_arrow_device() -> VortexResult<()> { let mut error = ptr::null_mut(); - let session = test_session(VortexSession::default()); + let session = test_session(test_cuda_session()?); let array = test_array(PrimitiveArray::from_iter(0u32..5)); let mut schema = FFI_ArrowSchema::empty(); let mut device_array = empty_device_array(); @@ -231,7 +235,7 @@ mod tests { assert_eq!(status, VX_CUDA_OK); assert!(error.is_null()); - let field = Field::try_from(&schema).expect("schema should be a field"); + let field = Field::try_from(&schema)?; assert_eq!(field.name(), ""); assert_eq!(device_array.array.length, 5); assert_eq!(device_array.array.n_buffers, 2); @@ -245,12 +249,13 @@ mod tests { free_test_array(array); free_test_session(session); } + Ok(()) } #[cuda_test] fn test_export_struct_arrow_device_table() -> VortexResult<()> { let mut error = ptr::null_mut(); - let session = test_session(VortexSession::default()); + let session = test_session(test_cuda_session()?); let array = test_array(StructArray::try_new( ["ids", "values"].into(), vec![ @@ -342,7 +347,7 @@ mod tests { #[cuda_not_available] #[test] fn test_export_reports_cuda_initialization_error() { - let session = test_session(VortexSession::default()); + let session = test_session(vortex::array::default_session_builder().build()); let array = test_array(PrimitiveArray::from_iter(0u32..5)); let mut schema = FFI_ArrowSchema::empty(); let mut device_array = empty_device_array(); diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 9078e30c691..c66cb1319a2 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -100,6 +100,15 @@ fn cuda_write_strategy() -> Arc { .build() } +#[cuda_available] +fn cuda_scan_session() -> VortexResult { + let session = VortexSession::default_builder() + .with_some(CudaSession::try_default()?) + .build(); + register_cuda_layout(&session); + Ok(session) +} + /// Convert an input Vortex file to CUDA-compatible encodings and write to disk. #[cuda_available] async fn cmd_convert(input: PathBuf, output: PathBuf) -> VortexResult<()> { @@ -149,8 +158,7 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes .init(); } - let session = VortexSession::default(); - register_cuda_layout(&session); + let session = cuda_scan_session()?; let mut cuda_ctx = CudaSession::create_execution_ctx(&session)? .with_launch_strategy(Arc::new(TracingLaunchStrategy)); diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index fb475ab2b24..25f6bc66692 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -952,7 +952,9 @@ mod tests { #[cuda_test] fn test_export_device_array_stream_schema_next_eos_release() -> VortexResult<()> { let runtime = CurrentThreadRuntime::new(); - let session = VortexSession::default().with_some(CudaSession::try_default()?); + let session = VortexSession::default_builder() + .with_some(CudaSession::try_default()?) + .build(); let array = PrimitiveArray::from_iter(0u32..5).into_array(); let stream = array.to_array_stream().boxed(); let mut device_stream = stream.export_device_array_stream(&session, &runtime)?; @@ -1004,7 +1006,9 @@ mod tests { #[cuda_test] fn test_export_device_array_stream_empty_stream_schema_and_eos() -> VortexResult<()> { let runtime = CurrentThreadRuntime::new(); - let session = VortexSession::default().with_some(CudaSession::try_default()?); + let session = VortexSession::default_builder() + .with_some(CudaSession::try_default()?) + .build(); let dtype = DType::Primitive(PType::U32, Nullability::NonNullable); let stream = ArrayStreamAdapter::new(dtype, stream::empty::>()).boxed(); @@ -1045,7 +1049,9 @@ mod tests { #[cuda_test] fn test_export_device_array_stream_null_outputs_report_error() -> VortexResult<()> { let runtime = CurrentThreadRuntime::new(); - let session = VortexSession::default().with_some(CudaSession::try_default()?); + let session = VortexSession::default_builder() + .with_some(CudaSession::try_default()?) + .build(); let array = PrimitiveArray::from_iter(0u32..5).into_array(); let stream = array.to_array_stream().boxed(); let mut device_stream = stream.export_device_array_stream(&session, &runtime)?; diff --git a/vortex-cuda/src/dynamic_dispatch/mod.rs b/vortex-cuda/src/dynamic_dispatch/mod.rs index f4d374f7b00..413840cb9d7 100644 --- a/vortex-cuda/src/dynamic_dispatch/mod.rs +++ b/vortex-cuda/src/dynamic_dispatch/mod.rs @@ -568,9 +568,9 @@ mod tests { use vortex::mask::Mask; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::FilterArray; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use vortex_array::patches::Patches; use super::*; @@ -627,7 +627,7 @@ mod tests { #[crate::test] fn test_max_scalar_ops() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 6; let len = 2050; let references: [u32; 4] = [1, 2, 4, 8]; @@ -848,7 +848,7 @@ mod tests { #[crate::test] async fn test_bitpacked() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let max_val = (1u64 << bit_width).saturating_sub(1); @@ -869,7 +869,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 6; let len = 3000; let reference = 42u32; @@ -921,7 +921,7 @@ mod tests { #[crate::test] async fn test_dict_for_bp_values_bp_codes() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Dict where both codes and values are BitPacked+FoR. let dict_reference = 1_000_000u32; let dict_residuals: Vec = (0..64).collect(); @@ -955,7 +955,7 @@ mod tests { #[crate::test] async fn test_alp_for_bitpacked() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // ALP(FoR(BitPacked)): encode each layer, then reassemble the tree // bottom-up because encode() methods produce flat outputs. let len = 3000; @@ -988,7 +988,7 @@ mod tests { #[crate::test] async fn test_zigzag_bitpacked() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // ZigZag(BitPacked): unpack then zigzag-decode. let bit_width: u8 = 4; let len = 3000; @@ -1076,7 +1076,7 @@ mod tests { #[crate::test] async fn test_dict_for_bp_codes() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Dict(codes=FoR(BitPacked), values=primitive) let dict_values: Vec = (0..8).map(|i| i * 1000 + 7).collect(); let dict_size = dict_values.len(); @@ -1105,7 +1105,7 @@ mod tests { #[crate::test] async fn test_dict_primitive_values_bp_codes() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let dict_size = dict_values.len(); let len = 3000; @@ -1131,7 +1131,7 @@ mod tests { #[crate::test] async fn test_dict_mixed_width_u8_codes_u32_values() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict_values: Vec = vec![100, 200, 300, 400]; let len = 3000; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u8).collect(); @@ -1164,7 +1164,7 @@ mod tests { #[crate::test] async fn test_dict_mixed_width_u16_codes_u32_values() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let dict_values: Vec = vec![1000, 2000, 3000, 4000, 5000]; let len = 2048; let codes: Vec = (0..len).map(|i| (i % dict_values.len()) as u16).collect(); @@ -1197,7 +1197,7 @@ mod tests { #[crate::test] async fn test_runend_mixed_width_u64_ends_u32_values() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let ends: Vec = vec![1000, 2000, 3000]; let values: Vec = vec![10, 20, 30]; let len = 3000; @@ -1300,7 +1300,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; let len = 5000; @@ -1402,7 +1402,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; let len = 5000; @@ -1448,7 +1448,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let reference = 100u32; let bit_width = 10u8; let max_val = (1u32 << bit_width) - 1; @@ -1498,7 +1498,7 @@ mod tests { #[case] slice_start: usize, #[case] slice_end: usize, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict_reference = 1_000_000u32; let dict_residuals: Vec = (0..64).collect(); let dict_expected: Vec = dict_residuals.iter().map(|&r| r + dict_reference).collect(); @@ -1604,7 +1604,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u8() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 4; let len = 3000; let reference = 100u8; @@ -1635,7 +1635,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u16() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 10; let len = 3000; let reference = 1000u16; @@ -1668,7 +1668,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_u64() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 20; let len = 3000; let reference = 100_000u64; @@ -1710,7 +1710,7 @@ mod tests { #[crate::test] async fn test_single_element() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let values: Vec = vec![42]; let primitive = PrimitiveArray::new(Buffer::from(values.clone()), NonNullable); let bp = @@ -1732,7 +1732,7 @@ mod tests { #[crate::test] async fn test_exactly_elements_per_block() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Exactly 2048 elements — one full block, no remainder let bit_width: u8 = 6; let len = 2048; @@ -1767,7 +1767,7 @@ mod tests { #[crate::test] fn test_has_standalone_kernel_true_cases() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // BitPacked — leaf encoding, no children. let bp = bitpacked_array_u32(6, 2048, &mut ctx); let bp_arr = bp.into_array(); @@ -1819,7 +1819,7 @@ mod tests { /// a kernel registered in the CUDA session. #[crate::test] fn test_has_standalone_kernel_implies_registered_kernel() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let session = CudaSession::create_execution_ctx(&cuda_session())?; let cuda_session = session.cuda_session(); @@ -1886,7 +1886,7 @@ mod tests { #[crate::test] async fn test_alp_f64_for_bitpacked() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // ALP(FoR(BitPacked)) with f64: same structure as the f32 test. let len = 3000; let exponents = Exponents { e: 2, f: 0 }; @@ -1931,7 +1931,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut values: Vec = (0..len).map(|i| (i as f64) * 1.1).collect(); // Insert exception values that ALP can't encode. values[0] = 99.9; @@ -1976,7 +1976,7 @@ mod tests { #[crate::test] async fn alp_slice_device_patches() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); // Regression test for https://github.com/vortex-data/vortex/issues/7838#issuecomment-4452796116. let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 4096; @@ -2039,7 +2039,7 @@ mod tests { #[crate::test] async fn test_runend_u32_ends_u16_values() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); // RunEnd with u32 ends, u16 values. Output type = u16. // Ends (u32) differ from output (u16) → pending subtree. let ends: Vec = vec![500, 1000, 1500, 2000]; @@ -2084,7 +2084,7 @@ mod tests { #[crate::test] async fn test_dict_bitpacked_u8_codes_u32_values() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); // Dict with BitPacked u8 codes (narrower than u32 output) and u32 values. // The kernel's bitunpack_typed decodes at the source's native width and // widens to T, so this fuses into a single kernel launch. @@ -2130,7 +2130,7 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); @@ -2170,7 +2170,7 @@ mod tests { #[case] len: usize, #[case] dict_values: Vec, ) -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let dict_size = dict_values.len(); let codes: Vec = (0..len).map(|i| (i % dict_size) as u8).collect(); @@ -2220,7 +2220,7 @@ mod tests { #[case] ends: Vec, #[case] values: Vec, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ends_u64: Vec = ends.iter().map(|e| (*e).into()).collect(); let len = *ends_u64.last().unwrap() as usize; let bit_width = 64 - ends_u64.iter().max().unwrap().leading_zeros() as u8; @@ -2260,7 +2260,7 @@ mod tests { /// at native width and widens to T, fusing everything. #[crate::test] async fn test_runend_mixed_width_for_bp_u16_ends_u32_values() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ends: Vec = vec![500, 1000, 1500, 2000]; let values: Vec = vec![100, 200, 300, 400]; let len = 2000usize; @@ -2298,7 +2298,7 @@ mod tests { #[crate::test] async fn test_sliced_dict_mixed_width() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Sliced Dict with u8 codes and u32 values — combines PartiallyFused + slice handling. let dict_values: Vec = vec![100, 200, 300, 400]; let full_len = 4096; @@ -2409,7 +2409,7 @@ mod tests { /// Nullable Primitive array — LOAD source with validity propagated. #[crate::test] async fn test_nullable_primitive() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let array = PrimitiveArray::from_option_iter( @@ -2432,7 +2432,7 @@ mod tests { /// validity, so this produces a real nullable FoR(BitPacked) tree. #[crate::test] async fn test_nullable_for_bitpacked() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 2048; @@ -2500,7 +2500,7 @@ mod tests { /// AllValid nullable array — should fuse and produce AllValid output. #[crate::test] async fn test_all_valid_nullable() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let values: Vec = (0..2048).collect(); @@ -2536,7 +2536,7 @@ mod tests { /// Dict with non-nullable codes but nullable values should still fuse. #[crate::test] async fn test_dict_nullable_values_fuses() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let codes = PrimitiveArray::new(buffer![0u32, 1, 2, 2, 1, 0], NonNullable); @@ -2560,7 +2560,7 @@ mod tests { /// Validity must survive through fused dispatch and into the filter. #[crate::test] async fn test_nullable_fused_then_filter() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 2048usize; @@ -2611,7 +2611,7 @@ mod tests { #[crate::test] async fn test_bitpacked_with_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let len = 3000; let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; @@ -2652,7 +2652,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; let values: Vec = (0..len) @@ -2690,7 +2690,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_with_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let len = 3000; let bit_width: u8 = 6; let reference = 42u32; @@ -2727,7 +2727,7 @@ mod tests { #[crate::test] async fn test_for_bitpacked_with_patches_sliced() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let len = 5000; let bit_width: u8 = 6; let reference = 42u32; @@ -2773,7 +2773,7 @@ mod tests { #[case] len: usize, #[case] slice_range: Option>, ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut values: Vec = (0..len).map(|i| (i as f32) * 1.1).collect(); // Insert exception values that ALP can't encode. values[0] = 99.9; @@ -2826,7 +2826,7 @@ mod tests { /// u8 BitPacked with patches (bit_width=3, patch values > 7). #[crate::test] async fn test_bitpacked_with_patches_u8() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 3; let len = 3000usize; let max_val = (1u8 << bit_width) - 1; @@ -2856,7 +2856,7 @@ mod tests { /// u16 BitPacked with patches (bit_width=6, patch values > 63). #[crate::test] async fn test_bitpacked_with_patches_u16() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 6; let len = 3000usize; let max_val = (1u16 << bit_width) - 1; @@ -2886,7 +2886,7 @@ mod tests { /// u64 BitPacked with patches (bit_width=4, patch values > 15). #[crate::test] async fn test_bitpacked_with_patches_u64() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 4; let len = 3000usize; let max_val = (1u64 << bit_width) - 1; @@ -2916,7 +2916,7 @@ mod tests { /// Dict where codes are BitPacked u32 with patches exceeding the bit width. #[crate::test] async fn test_dict_bitpacked_codes_with_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dict_values: Vec = (0..256).map(|i| i * 1000 + 42).collect(); let len = 3000; let bit_width: u8 = 4; @@ -2951,7 +2951,7 @@ mod tests { /// Patches placed exactly at FastLanes chunk boundaries (1024-element chunks). #[crate::test] async fn test_bitpacked_patches_at_chunk_boundaries() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let len = 4096usize; let bit_width: u8 = 4; let max_val = (1u32 << bit_width) - 1; @@ -2977,7 +2977,7 @@ mod tests { /// Large array (100k elements) spanning many blocks with sparse patches. #[crate::test] async fn test_bitpacked_large_array_with_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let len = 100_000usize; let bit_width: u8 = 6; let max_val = (1u32 << bit_width) - 1; @@ -3007,7 +3007,7 @@ mod tests { /// dispatch alongside patch application. #[crate::test] async fn test_nullable_bitpacked_with_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; let len = 3000usize; @@ -3045,7 +3045,7 @@ mod tests { /// Extreme case: ALL values are patches (bit_width=1, every value > 1). #[crate::test] async fn test_bitpacked_all_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bit_width: u8 = 1; let len = 2000usize; // All values >= 2, so every single element exceeds max storable (1) and diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index b3a0151a727..8202b6ece3f 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -172,14 +172,14 @@ mod tests { use vortex::mask::Mask; use vortex::scalar::Scalar; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; + use vortex_array::default_session_builder; use crate::CanonicalCudaExt; use crate::executor::CudaArrayExt; use crate::session::CudaSession; fn for_bp>(values: Vec, reference: T) -> ArrayRef { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bp = BitPacked::encode( &PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(), 7, @@ -194,7 +194,7 @@ mod tests { /// FoR(BitPacked) u32 — entire tree compiles into a single fused plan. #[crate::test] async fn test_fused() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); let values: Vec = (0..2048).map(|i| (i % 128) as u32).collect(); @@ -222,7 +222,7 @@ mod tests { /// Exercises the unsigned type reinterpretation in CudaDispatchPlan::execute. #[crate::test] async fn test_fused_f32() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); use vortex::encodings::alp::ALP; use vortex::encodings::alp::Exponents; @@ -258,7 +258,7 @@ mod tests { /// ALP with patches — plan builder rejects it, falls back to ALPExecutor. #[crate::test] async fn test_fallback() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); use vortex::array::patches::Patches; use vortex::array::validity::Validity::NonNullable as NN; use vortex::buffer::buffer; @@ -300,15 +300,16 @@ mod tests { #[cfg(feature = "unstable_encodings")] #[crate::test] async fn test_partial_fusion() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); use vortex::array::arrays::DictArray; - use vortex::array::session::ArraySessionExt; + use vortex::array::session::ArraySession; use vortex::encodings::fastlanes; use vortex::encodings::zstd::ZstdBuffers; - let session = crate::cuda_session(); - fastlanes::initialize(&session); - session.arrays().register(ZstdBuffers); + let mut builder = crate::cuda_session_builder(); + fastlanes::initialize(&mut builder); + builder.get_mut::().register(ZstdBuffers); + let session = builder.build(); let mut ctx = CudaSession::create_execution_ctx(&session).vortex_expect("ctx"); let num_values: u32 = 64; @@ -364,7 +365,7 @@ mod tests { /// Filter(FoR(BP), mask) — FoR+BP fuses via dyn dispatch, then CUB filters the result. #[crate::test] async fn test_filter_fused_child() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -408,7 +409,7 @@ mod tests { ))] #[crate::test] async fn test_ext_storage_gpu_decode(#[case] ext: ExtensionArray) -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); @@ -432,7 +433,7 @@ mod tests { /// Extension over already-canonical storage executes unchanged. #[crate::test] async fn test_ext_canonical_storage() -> VortexResult<()> { - let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cpu_ctx = default_session_builder().build().create_execution_ctx(); let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).vortex_expect("ctx"); diff --git a/vortex-cuda/src/kernel/arrays/constant.rs b/vortex-cuda/src/kernel/arrays/constant.rs index d385333a024..4fdcca1cc12 100644 --- a/vortex-cuda/src/kernel/arrays/constant.rs +++ b/vortex-cuda/src/kernel/arrays/constant.rs @@ -224,7 +224,9 @@ mod tests { async fn test_cuda_constant_materialization( #[case] constant_array: ConstantArray, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -243,7 +245,9 @@ mod tests { #[crate::test] async fn test_cuda_constant_empty_array() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -263,7 +267,9 @@ mod tests { #[crate::test] async fn test_cuda_constant_small_array() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index 394de699f49..40c4302851c 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -325,7 +325,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_u32_values_u8_codes() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -355,7 +357,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_u64_values_u16_codes() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -388,7 +392,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_i32_values_u32_codes() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -417,7 +423,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_large_array() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -447,7 +455,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_values_with_validity() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -477,7 +487,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_codes_with_validity() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -512,7 +524,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_both_with_validity() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -554,7 +568,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_i64_values_with_validity() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -597,7 +613,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_all_valid_matches_baseline() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -638,7 +656,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i8_values() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -665,7 +685,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i16_values() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -692,7 +714,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i32_values() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -719,7 +743,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i64_values() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -749,7 +775,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i128_values() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -792,7 +820,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_u8_codes() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -816,7 +846,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_u16_codes() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -840,7 +872,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_max_inlined_12_bytes() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -866,7 +900,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_outlined_views() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -895,7 +931,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_empty_strings() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -919,7 +957,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_values_with_validity() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -944,7 +984,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_string_outlined_with_validity() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -976,7 +1018,9 @@ mod tests { #[crate::test] async fn test_cuda_dict_decimal_i256_values() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/alp.rs b/vortex-cuda/src/kernel/encodings/alp.rs index d027aa0bbd6..a4ab9e63da5 100644 --- a/vortex-cuda/src/kernel/encodings/alp.rs +++ b/vortex-cuda/src/kernel/encodings/alp.rs @@ -141,9 +141,9 @@ mod tests { use std::f64; use vortex::array::IntoArray; - use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; use vortex::array::assert_arrays_eq; + use vortex::array::default_session_builder; use vortex::array::patches::Patches; use vortex::array::validity::Validity; use vortex::buffer::Buffer; @@ -169,7 +169,7 @@ mod tests { /// Patches must carry `chunk_offsets` — the fused kernel requires them. #[crate::test] async fn test_cuda_alp_decompression_f32() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -215,7 +215,7 @@ mod tests { /// preserved through the standalone ALP GPU executor. #[crate::test] async fn test_cuda_alp_nullable_with_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -236,7 +236,7 @@ mod tests { let alp_array = alp_encode( prim.as_view(), None, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let gpu_result = alp_array @@ -256,7 +256,7 @@ mod tests { /// elements are actually null. #[crate::test] async fn test_cuda_alp_all_valid_nullable() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -267,7 +267,7 @@ mod tests { let alp_array = alp_encode( values.as_view(), None, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let gpu_result = alp_array @@ -289,7 +289,7 @@ mod tests { /// (zero patches) via the offset math rather than the NULL sentinel. #[crate::test] async fn test_cuda_alp_multi_chunk_sparse_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -309,7 +309,7 @@ mod tests { let alp_array = alp_encode( prim.as_view(), None, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert!( alp_array.patches().is_some(), @@ -334,7 +334,7 @@ mod tests { /// so this guards the fast-path for the (i64, f64) kernel variant. #[crate::test] async fn test_cuda_alp_f64_multi_chunk_with_patches() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -352,7 +352,7 @@ mod tests { let alp_array = alp_encode( prim.as_view(), None, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert!( alp_array.patches().is_some(), @@ -378,7 +378,7 @@ mod tests { /// (existing tests have ≤ 6 patches per chunk). #[crate::test] async fn test_cuda_alp_dense_patches_single_chunk() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -425,7 +425,7 @@ mod tests { /// loop. Includes a patch in the tail. #[crate::test] async fn test_cuda_alp_partial_tail_chunk() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -436,7 +436,7 @@ mod tests { let alp_array = alp_encode( prim.as_view(), None, - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; assert!( alp_array.patches().is_some(), diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 6820050b72f..10eba26c004 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -249,7 +249,9 @@ mod tests { #[case] iter: impl Iterator, #[case] bw: u8, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -276,7 +278,9 @@ mod tests { #[crate::test] fn test_patches() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -314,7 +318,9 @@ mod tests { #[case::bw_7(7)] #[crate::test] fn test_cuda_bitunpack_u8(#[case] bit_width: u8) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -362,7 +368,9 @@ mod tests { #[case::bw_15(15)] #[crate::test] fn test_cuda_bitunpack_u16(#[case] bit_width: u8) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -426,7 +434,9 @@ mod tests { #[case::bw_31(31)] #[crate::test] fn test_cuda_bitunpack_u32(#[case] bit_width: u8) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -522,7 +532,9 @@ mod tests { #[case::bw_63(63)] #[crate::test] fn test_cuda_bitunpack_u64(#[case] bit_width: u8) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -554,7 +566,9 @@ mod tests { #[crate::test] fn test_cuda_bitunpack_sliced() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let bit_width = 32; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -599,7 +613,9 @@ mod tests { #[case] expected_offset: u16, #[case] expected_packed_len: usize, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let values = PrimitiveArray::new( (0u16..4096) .map(|i| if i % 1000 == 0 { 600 } else { i % 512 }) @@ -631,7 +647,9 @@ mod tests { /// offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_sliced_patches_offset_within_chunk() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -667,7 +685,9 @@ mod tests { /// Test slicing a bitpacked array multiple times, accumulating offset_within_chunk. #[crate::test] fn test_cuda_bitunpack_double_sliced_patches() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -715,7 +735,9 @@ mod tests { /// Test slicing to skip an entire chunk's worth of patches. #[crate::test] fn test_cuda_bitunpack_sliced_skip_first_chunk_patches() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/date_time_parts.rs b/vortex-cuda/src/kernel/encodings/date_time_parts.rs index dcb1fd1f69e..3ce15b8652d 100644 --- a/vortex-cuda/src/kernel/encodings/date_time_parts.rs +++ b/vortex-cuda/src/kernel/encodings/date_time_parts.rs @@ -280,7 +280,9 @@ mod tests { #[case] subseconds: Vec, #[case] time_unit: TimeUnit, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -300,7 +302,9 @@ mod tests { #[crate::test] async fn test_cuda_datetimeparts_large_array() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -325,7 +329,9 @@ mod tests { #[crate::test] async fn test_cuda_datetimeparts_with_nulls() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index 5f7e639b8a1..8e6b018b854 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -86,7 +86,9 @@ mod tests { #[case] precision: u8, #[case] scale: i8, ) { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/for_.rs b/vortex-cuda/src/kernel/encodings/for_.rs index 6348f8479b7..827aaf4daae 100644 --- a/vortex-cuda/src/kernel/encodings/for_.rs +++ b/vortex-cuda/src/kernel/encodings/for_.rs @@ -139,7 +139,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::scalar::Scalar; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; + use vortex_array::default_session_builder; use super::*; use crate::CanonicalCudaExt; @@ -160,7 +160,7 @@ mod tests { #[case::u64(make_for_array((0..2050).map(|i| (i % 2050) as u64).collect(), 1000000u64))] #[crate::test] async fn test_cuda_for_decompression(#[case] for_array: FoRArray) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -179,7 +179,7 @@ mod tests { #[crate::test] async fn test_signed_ffor() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -188,9 +188,13 @@ mod tests { .take(1024) .collect::>() .into_array(); - let packed = BitPacked::encode(&values, 3, &mut array_session().create_execution_ctx()) - .unwrap() - .into_array(); + let packed = BitPacked::encode( + &values, + 3, + &mut default_session_builder().build().create_execution_ctx(), + ) + .unwrap() + .into_array(); let for_array = FoR::try_new(packed, (-8i8).into()).unwrap(); let gpu_result = FoRExecutor diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index eef0db049b8..71532c1dc6d 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -236,7 +236,9 @@ mod tests { #[case] strings: Vec>, #[case] nullability: Nullability, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -260,7 +262,9 @@ mod tests { /// Exercises the multi-block grid-stride path on a larger dataset. #[crate::test] async fn test_cuda_fsst_decompression_roundtrip_large() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); use vortex_fsst::test_utils::make_fsst_clickbench_urls; let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index 4ac042e2972..2a3177eb597 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -205,7 +205,9 @@ mod tests { #[case::u64_ends_i32_values(|ctx: &mut ExecutionCtx| make_runend_array(vec![2u64, 5, 10], vec![1i32, 2, 3], ctx))] #[crate::test] async fn test_cuda_runend_types(#[case] build: RunEndBuilder) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -225,7 +227,9 @@ mod tests { #[crate::test] async fn test_cuda_runend_large_array() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -254,7 +258,9 @@ mod tests { #[crate::test] async fn test_cuda_runend_single_run() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -275,7 +281,9 @@ mod tests { #[crate::test] async fn test_cuda_runend_many_small_runs() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -301,7 +309,9 @@ mod tests { #[crate::test] async fn test_cuda_runend_nullable_values_falls_back_to_cpu() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/sequence.rs b/vortex-cuda/src/kernel/encodings/sequence.rs index e52a706c8b7..9a0c0d4a0ba 100644 --- a/vortex-cuda/src/kernel/encodings/sequence.rs +++ b/vortex-cuda/src/kernel/encodings/sequence.rs @@ -125,7 +125,9 @@ mod tests { len: usize, nullability: Nullability, ) { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).unwrap(); let array = Sequence::try_new_typed(base, multiplier, nullability, len) diff --git a/vortex-cuda/src/kernel/encodings/zigzag.rs b/vortex-cuda/src/kernel/encodings/zigzag.rs index 107ad2b36a1..b83a51d89f0 100644 --- a/vortex-cuda/src/kernel/encodings/zigzag.rs +++ b/vortex-cuda/src/kernel/encodings/zigzag.rs @@ -113,7 +113,9 @@ mod tests { #[crate::test] async fn test_cuda_zigzag_decompression_u32() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/zstd.rs b/vortex-cuda/src/kernel/encodings/zstd.rs index 6d14cc46e86..cf042faaba3 100644 --- a/vortex-cuda/src/kernel/encodings/zstd.rs +++ b/vortex-cuda/src/kernel/encodings/zstd.rs @@ -365,7 +365,9 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_utf8() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -392,7 +394,9 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_multiple_frames() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -429,7 +433,9 @@ mod tests { #[crate::test] async fn test_cuda_zstd_decompression_sliced() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -464,7 +470,9 @@ mod tests { /// correct results instead of panicking. #[crate::test] async fn test_cuda_zstd_nullable_falls_back_to_cpu() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs index 7de2158ac48..80e28f17a6d 100644 --- a/vortex-cuda/src/kernel/encodings/zstd_buffers.rs +++ b/vortex-cuda/src/kernel/encodings/zstd_buffers.rs @@ -227,7 +227,7 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; + use vortex_array::default_session_builder; use super::*; use crate::CanonicalCudaExt; @@ -235,7 +235,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_buffers_decompression_primitive() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -254,7 +254,7 @@ mod tests { #[crate::test] async fn test_cuda_zstd_buffers_decompression_varbinview() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); diff --git a/vortex-cuda/src/kernel/filter/decimal.rs b/vortex-cuda/src/kernel/filter/decimal.rs index dc10b908461..d196ead7726 100644 --- a/vortex-cuda/src/kernel/filter/decimal.rs +++ b/vortex-cuda/src/kernel/filter/decimal.rs @@ -89,7 +89,9 @@ mod tests { #[case] input: DecimalArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -110,7 +112,9 @@ mod tests { #[crate::test] async fn test_gpu_filter_decimal_large_array() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); diff --git a/vortex-cuda/src/kernel/filter/primitive.rs b/vortex-cuda/src/kernel/filter/primitive.rs index c2a21ccc7d7..b1cc41abffe 100644 --- a/vortex-cuda/src/kernel/filter/primitive.rs +++ b/vortex-cuda/src/kernel/filter/primitive.rs @@ -83,7 +83,9 @@ mod tests { #[case] input: PrimitiveArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); @@ -104,7 +106,9 @@ mod tests { #[crate::test] async fn test_gpu_filter_large_array() -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); diff --git a/vortex-cuda/src/kernel/filter/varbinview.rs b/vortex-cuda/src/kernel/filter/varbinview.rs index 7bbafff824c..4de6bba5064 100644 --- a/vortex-cuda/src/kernel/filter/varbinview.rs +++ b/vortex-cuda/src/kernel/filter/varbinview.rs @@ -69,7 +69,9 @@ mod tests { #[case] input: VarBinViewArray, #[case] mask: Mask, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create CUDA execution context"); diff --git a/vortex-cuda/src/kernel/patches/mod.rs b/vortex-cuda/src/kernel/patches/mod.rs index bd6d2b737db..a2c55c7c95a 100644 --- a/vortex-cuda/src/kernel/patches/mod.rs +++ b/vortex-cuda/src/kernel/patches/mod.rs @@ -176,12 +176,12 @@ mod tests { use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; - use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::assert_arrays_eq; use vortex::array::buffer::BufferHandle; use vortex::array::builtins::ArrayBuiltins; + use vortex::array::default_session_builder; use vortex::array::patches::Patches; use vortex::array::validity::Validity; use vortex::buffer::buffer; @@ -217,7 +217,7 @@ mod tests { async fn full_test_case() { let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()).unwrap(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = PrimitiveArray::from_iter(0..128); let values = force_cast::(values, &mut ctx); diff --git a/vortex-cuda/src/kernel/patches/types.rs b/vortex-cuda/src/kernel/patches/types.rs index 3bfe2270b66..3c689972af4 100644 --- a/vortex-cuda/src/kernel/patches/types.rs +++ b/vortex-cuda/src/kernel/patches/types.rs @@ -318,7 +318,9 @@ mod tests { #[case] chunk_offsets: ArrayRef, #[case] expected: ArrayRef, ) -> VortexResult<()> { - let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut ctx = vortex_array::default_session_builder() + .build() + .create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let indices = PrimitiveArray::from_iter([100u32, 1100, 2100, 3100, 4100]); let values = PrimitiveArray::from_iter([10u32, 11, 12, 13, 14]); diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 07cfe40b06f..4d030612dde 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -134,7 +134,12 @@ pub fn initialize_cuda(session: &CudaSession) { /// # Panics /// /// Panics if CUDA device 0 cannot be initialized (the same contract as [`CudaSession::default`]). +#[cfg(any(test, feature = "_test-harness"))] +pub fn cuda_session_builder() -> vortex::session::VortexSessionBuilder { + vortex::array::default_session_builder().with::() +} + #[cfg(any(test, feature = "_test-harness"))] pub fn cuda_session() -> vortex::session::VortexSession { - vortex::array::array_session().with::() + cuda_session_builder().build() } diff --git a/vortex-cxx/src/lib.rs b/vortex-cxx/src/lib.rs index b6d002513f6..ce2464b6311 100644 --- a/vortex-cxx/src/lib.rs +++ b/vortex-cxx/src/lib.rs @@ -18,7 +18,7 @@ use scalar::*; use vortex::VortexSessionDefault; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; use write::*; @@ -28,8 +28,11 @@ use write::*; // this runtime. pub(crate) static RUNTIME: LazyLock = LazyLock::new(CurrentThreadRuntime::new); -pub(crate) static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle())); +pub(crate) static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::default_builder() + .with_handle(RUNTIME.handle()) + .build() +}); #[cxx::bridge(namespace = "vortex::ffi")] #[allow(let_underscore_drop)] diff --git a/vortex-datafusion/examples/vortex_table.rs b/vortex-datafusion/examples/vortex_table.rs index 2570fc47afd..30a3df7d35a 100644 --- a/vortex-datafusion/examples/vortex_table.rs +++ b/vortex-datafusion/examples/vortex_table.rs @@ -19,13 +19,13 @@ use vortex::array::validity::Validity; use vortex::buffer::buffer; use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; use vortex_datafusion::VortexFormat; #[tokio::main] async fn main() -> anyhow::Result<()> { - let session = VortexSession::default().with_tokio(); + let session = VortexSession::default_builder().with_tokio().build(); let temp_dir = tempdir()?; let strings = ChunkedArray::from_iter([ diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 1d12dab9dbe..1dbff816381 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -402,7 +402,7 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::error::VortexExpect; use vortex::mask::Mask; - use vortex_array::array_session; + use vortex_array::default_session_builder; use vortex_geo::extension::WellKnownBinaryData; use wkb::writer::WriteOptions; use wkb::writer::write_point; @@ -415,7 +415,7 @@ mod tests { #[test] fn test_integer_vector_conversion() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1i32, 2, 3, 4, 5]; let len = values.len(); @@ -437,7 +437,7 @@ mod tests { #[test] fn test_timestamp_vector_conversion() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1_703_980_800_000_000_i64, 0i64, -86_400_000_000_i64]; // microseconds let len = values.len(); @@ -465,7 +465,7 @@ mod tests { #[test] fn test_timestamp_seconds_vector_conversion() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1_703_980_800_i64, 0i64, -86_400_i64]; // seconds let len = values.len(); @@ -493,7 +493,7 @@ mod tests { #[test] fn test_timestamp_milliseconds_vector_conversion() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1_703_980_800_000_i64, 0i64, -86_400_000_i64]; // milliseconds let len = values.len(); @@ -521,7 +521,7 @@ mod tests { #[test] fn test_timestamp_with_nulls_conversion() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1_703_980_800_000_000_i64, 0i64, -86_400_000_000_i64]; let len = values.len(); @@ -585,7 +585,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&vector, len).unwrap(); let vortex_array = TemporalArray::try_from(result).unwrap(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let vortex_values = vortex_array .temporal_values() .clone() @@ -598,7 +598,7 @@ mod tests { #[test] fn test_timestamp_single_value() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1_703_980_800_000_000_i64]; // Single microsecond timestamp let len = values.len(); @@ -626,7 +626,7 @@ mod tests { #[test] fn test_boolean_vector_conversion() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![true, false, true, false]; let len = values.len(); @@ -648,7 +648,7 @@ mod tests { #[test] fn test_vector_with_nulls() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1i32, 2, 3]; let len = values.len(); @@ -707,7 +707,7 @@ mod tests { // Test conversion let result = flat_vector_to_vortex(&vector, len).unwrap(); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let vortex_array = result.execute::(&mut ctx).unwrap(); assert_eq!(vortex_array.len(), len); @@ -720,7 +720,7 @@ mod tests { #[test] fn test_fixed_sized_list() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![1i32, 2, 3, 4]; let len = 1; @@ -750,7 +750,7 @@ mod tests { #[test] fn test_empty_struct() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let len = 4; let logical_type = LogicalType::struct_type([], []) .vortex_expect("LogicalTypeRef creation should succeed for test data"); @@ -766,7 +766,7 @@ mod tests { #[test] fn test_struct() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values1 = vec![1i32, 2, 3, 4]; let values2 = vec![5i32, 6, 7, 8]; let len = values1.len(); @@ -812,7 +812,7 @@ mod tests { #[test] fn test_list_with_trailing_null() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Regression test: when the last list entry is null, its offset/length may be 0/0, // so we can't use the last entry to compute child vector length. let child_values = vec![1i32, 2, 3, 4]; @@ -867,7 +867,7 @@ mod tests { #[test] fn test_list_out_of_order() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Regression test: list views can be out of order in DuckDB. The child vector length // must be computed as the maximum end offset, not just the last entry's end offset. let child_values = vec![1i32, 2, 3, 4]; @@ -915,7 +915,7 @@ mod tests { #[test] fn test_list_null_garbage_data() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Test that null list entries with garbage offset/size values don't cause issues. // DuckDB doesn't guarantee valid offset/size for null list views, so we must check // validity before reading the offset/size values. @@ -999,7 +999,7 @@ mod tests { #[test] fn test_geometry_vector_conversion() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut wkb_a: Vec = Vec::new(); write_point( diff --git a/vortex-duckdb/src/duckdb/vector.rs b/vortex-duckdb/src/duckdb/vector.rs index 5ac3ad01aa9..0919913f6c1 100644 --- a/vortex-duckdb/src/duckdb/vector.rs +++ b/vortex-duckdb/src/duckdb/vector.rs @@ -383,7 +383,7 @@ impl ValidityRef<'_> { #[cfg(test)] mod tests { - use vortex::array::array_session; + use vortex::array::default_session_builder; use vortex::mask::Mask; use vortex_array::VortexSessionExecute; @@ -420,7 +420,7 @@ mod tests { let validity = validity.to_validity(); assert_eq!(validity.maybe_len(), Some(len)); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); assert_eq!( validity.execute_mask(len, &mut ctx).unwrap(), Mask::from_indices(len, vec![0, 2, 4, 5, 6, 8, 9]) @@ -440,7 +440,10 @@ mod tests { let validity = validity.to_validity(); assert!( validity - .execute_is_null(0, &mut array_session().create_execution_ctx()) + .execute_is_null( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } @@ -457,7 +460,10 @@ mod tests { let validity = validity.to_validity(); assert!( validity - .execute_is_valid(0, &mut array_session().create_execution_ctx()) + .execute_is_valid( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() ); } diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 6cc28483571..2c57c1298d7 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -20,7 +20,6 @@ use num_traits::AsPrimitive; use tempfile::NamedTempFile; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; -use vortex::array::array_session; use vortex::array::arrays::BoolArray; use vortex::array::arrays::ConstantArray; use vortex::array::arrays::DictArray; @@ -30,6 +29,7 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinArray; use vortex::array::arrays::VarBinViewArray; +use vortex::array::default_session_builder; use vortex::array::validity::Validity; use vortex::buffer::buffer; use vortex::dtype::Nullability; @@ -779,7 +779,7 @@ async fn write_vortex_file_with_encodings() -> NamedTempFile { // 4. Run-End let run_ends = buffer![3u32, 5]; let run_values = buffer![100i32, 200]; - let mut rle_ctx = array_session().create_execution_ctx(); + let mut rle_ctx = default_session_builder().build().create_execution_ctx(); let rle_array = RunEnd::try_new(run_ends.into_array(), run_values.into_array(), &mut rle_ctx).unwrap(); diff --git a/vortex-duckdb/src/exporter/decimal.rs b/vortex-duckdb/src/exporter/decimal.rs index 361f4477474..af30ba8d804 100644 --- a/vortex-duckdb/src/exporter/decimal.rs +++ b/vortex-duckdb/src/exporter/decimal.rs @@ -135,8 +135,8 @@ pub fn precision_to_duckdb_storage_size(decimal_dtype: &DecimalDType) -> VortexR #[cfg(test)] mod tests { use vortex::array::VortexSessionExecute; - use vortex::array::array_session; use vortex::array::arrays::DecimalArray; + use vortex::array::default_session_builder; use vortex::dtype::DecimalDType; use vortex::error::VortexExpect; @@ -150,7 +150,7 @@ mod tests { ) -> VortexResult> { let validity = array.as_ref().validity()?.execute_mask( array.as_ref().len(), - &mut array_session().create_execution_ctx(), + &mut default_session_builder().build().create_execution_ctx(), )?; let dest_values_type = precision_to_duckdb_storage_size(&array.decimal_dtype())?; diff --git a/vortex-duckdb/src/lib.rs b/vortex-duckdb/src/lib.rs index 55039b88737..0bdd07d7eec 100644 --- a/vortex-duckdb/src/lib.rs +++ b/vortex-duckdb/src/lib.rs @@ -13,7 +13,7 @@ use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; use crate::duckdb::Database; @@ -42,9 +42,9 @@ mod e2e_test; // A global runtime for Vortex operations within DuckDB. static RUNTIME: LazyLock = LazyLock::new(CurrentThreadRuntime::new); static SESSION: LazyLock = LazyLock::new(|| { - let session = VortexSession::default().with_handle(RUNTIME.handle()); - vortex_geo::initialize(&session); - session + let mut builder = VortexSession::default_builder(); + vortex_geo::initialize(&mut builder); + builder.with_handle(RUNTIME.handle()).build() }); // Duckdb's logger requires a *Context as first argument which diff --git a/vortex-ffi/examples/hello_vortex.rs b/vortex-ffi/examples/hello_vortex.rs index 93ed17772f1..f878f4280a9 100644 --- a/vortex-ffi/examples/hello_vortex.rs +++ b/vortex-ffi/examples/hello_vortex.rs @@ -29,12 +29,15 @@ use vortex::file::WriteOptionsSessionExt; use vortex::io::VortexWrite; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; static RUNTIME: LazyLock = LazyLock::new(CurrentThreadRuntime::new); -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle())); +static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::default_builder() + .with_handle(RUNTIME.handle()) + .build() +}); const BIN_NAME: &str = "hello_vortex"; diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index a87178446de..94e5c19eb17 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -497,12 +497,12 @@ mod tests { use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; - use vortex::array::array_session; use vortex::array::arrays::BoolArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::bool::BoolArrayExt; + use vortex::array::default_session_builder; use vortex::array::validity::Validity; use vortex::buffer::buffer; #[cfg(not(miri))] @@ -827,7 +827,7 @@ mod tests { assert!(!res.is_null()); { let res = vx_array::as_ref(res); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let bool_array = res.clone().execute::(&mut ctx).unwrap(); let buffer = bool_array.to_bit_buffer(); let expected = BoolArray::from_iter(vec![false, false, true, true]); diff --git a/vortex-ffi/src/expression.rs b/vortex-ffi/src/expression.rs index 269e1e5d941..9f94d102ee6 100644 --- a/vortex-ffi/src/expression.rs +++ b/vortex-ffi/src/expression.rs @@ -337,13 +337,13 @@ mod tests { use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; - use vortex::array::array_session; use vortex::array::arrays::BoolArray; use vortex::array::arrays::ListArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::bool::BoolArrayExt; + use vortex::array::default_session_builder; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::buffer::buffer; @@ -393,7 +393,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_get_item() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let (array, names_array, ages_array) = struct_array(); unsafe { let root = vx_expression_root(); @@ -444,7 +444,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_literal() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable).into_array(); @@ -528,7 +528,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_and_or() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let col1 = BoolArray::from_iter([true, false, true, true]); let col2 = BoolArray::from_iter([false, true, true, false]); let col3 = BoolArray::from_iter([false, true, true, true]); @@ -620,7 +620,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_list_contains() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let elements = buffer![1i32, 2, 3, 4, 5].into_array(); let offsets = buffer![0u32, 2, 5, 5].into_array(); let array = ListArray::try_new(elements, offsets, Validity::NonNullable).unwrap(); diff --git a/vortex-ffi/src/scan.rs b/vortex-ffi/src/scan.rs index 619f82e3f3e..11d8402d4ff 100644 --- a/vortex-ffi/src/scan.rs +++ b/vortex-ffi/src/scan.rs @@ -454,9 +454,9 @@ mod tests { use vortex::array::arrays::StructArray; use vortex::session::VortexSession; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use crate::array::vx_array; use crate::array::vx_array_free; @@ -533,7 +533,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_no_options() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let (array, struct_array) = scan(ptr::null()); assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut ctx); unsafe { vx_array_free(array) }; @@ -542,7 +542,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_project_all() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let opts = vx_scan_options::default(); let (array, struct_array) = scan(&raw const opts); assert_arrays_eq!(vx_array::as_ref(array), struct_array, &mut ctx); @@ -553,7 +553,7 @@ mod tests { #[cfg_attr(miri, ignore)] fn test_project_single_field() { unsafe { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let root = vx_expression_root(); let mut opts = vx_scan_options::default(); @@ -715,7 +715,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_ordered() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let opts = vx_scan_options { ordered: true, ..Default::default() diff --git a/vortex-ffi/src/session.rs b/vortex-ffi/src/session.rs index 2af3d997444..7e3d1fc2cef 100644 --- a/vortex-ffi/src/session.rs +++ b/vortex-ffi/src/session.rs @@ -5,8 +5,9 @@ use vortex::VortexSessionDefault; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::io::runtime::BlockingRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; +use vortex::session::VortexSessionBuilder; use crate::RUNTIME; use crate::box_wrapper; @@ -19,11 +20,11 @@ box_wrapper!( /// Create an FFI session from a configured default session. pub fn vx_session_new_with( - configure: impl FnOnce(VortexSession) -> VortexSession, + configure: impl FnOnce(VortexSessionBuilder) -> VortexSessionBuilder, ) -> *mut vx_session { - vx_session::new(configure( - VortexSession::default().with_handle(RUNTIME.handle()), - )) + vx_session::new( + configure(VortexSession::default_builder().with_handle(RUNTIME.handle())).build(), + ) } /// Create a new Vortex session. @@ -31,7 +32,7 @@ pub fn vx_session_new_with( /// The caller is responsible for freeing the session with [`vx_session_free`]. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_session_new() -> *mut vx_session { - vx_session_new_with(|session| session) + vx_session_new_with(|builder| builder) } /// Clone a Vortex session, returning an owned copy. diff --git a/vortex-ffi/src/struct_array.rs b/vortex-ffi/src/struct_array.rs index 765864a6176..4ae87918d40 100644 --- a/vortex-ffi/src/struct_array.rs +++ b/vortex-ffi/src/struct_array.rs @@ -133,11 +133,11 @@ mod tests { use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; - use vortex::array::array_session; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::assert_arrays_eq; + use vortex::array::default_session_builder; use vortex::array::validity::Validity; use vortex::buffer::buffer; @@ -170,7 +170,7 @@ mod tests { #[test] #[cfg_attr(miri, ignore)] fn test_many() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let names = ["age", "name"]; let age_field = PrimitiveArray::new(buffer![30u8, 25u8, 35u8], Validity::NonNullable); let name_field = VarBinViewArray::from_iter_str(["Alice", "Bob", "Charlie"]); diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index e9c7741af93..7be43135779 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -118,11 +118,12 @@ pub use footer::*; pub use forever_constant::*; pub use open::*; pub use strategy::*; +use vortex_array::arrays::Dict; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; -use vortex_array::session::ArraySessionExt; +use vortex_array::session::ArraySession; use vortex_pco::Pco; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; pub use writer::*; /// The current version of the Vortex file format @@ -161,7 +162,7 @@ mod forever_constant { /// /// NOTE: this function will be changed in the future to encapsulate logic for using different /// Vortex "Editions" that may support different sets of encodings. -pub fn register_default_encodings(session: &VortexSession) { +pub fn register_default_encodings(session: &mut VortexSessionBuilder) { vortex_bytebool::initialize(session); vortex_fsst::initialize(session); #[cfg(feature = "unstable_encodings")] @@ -169,7 +170,8 @@ pub fn register_default_encodings(session: &VortexSession) { vortex_zigzag::initialize(session); { - let arrays = session.arrays(); + let arrays = session.get_mut::(); + arrays.register(Dict); arrays.register(Pco); #[cfg(feature = "zstd")] arrays.register(vortex_zstd::Zstd); @@ -195,8 +197,8 @@ pub fn register_default_encodings(session: &VortexSession) { #[cfg(test)] mod default_encoding_tests { use vortex_array::VTable as _; - use vortex_array::array_session; use vortex_array::arrays::Filter; + use vortex_array::default_session_builder; use vortex_array::optimizer::kernels::ArrayKernelsExt as _; use vortex_array::session::ArraySessionExt as _; use vortex_fsst::FSST; @@ -205,12 +207,14 @@ mod default_encoding_tests { #[test] fn register_default_encodings_registers_external_execute_parent_kernels() { - let session = array_session(); + let session = default_session_builder().build(); assert!(session.arrays().registry().find(&FSST.id()).is_none()); assert!(!session.kernels().has_execute_parent(Filter.id(), FSST.id())); - register_default_encodings(&session); + let mut builder = default_session_builder(); + register_default_encodings(&mut builder); + let session = builder.build(); assert!(session.arrays().registry().find(&FSST.id()).is_some()); assert!(session.kernels().has_execute_parent(Filter.id(), FSST.id())); diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index 1bd42c7be41..1c2f3d67c36 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -390,7 +390,7 @@ mod tests { use vortex_array::buffer::BufferHandle; use vortex_array::memory::DefaultHostAllocator; use vortex_array::memory::HostAllocator; - use vortex_array::memory::MemorySessionExt; + use vortex_array::memory::MemorySessionBuilderExt; use vortex_array::memory::WritableHostBuffer; use vortex_buffer::Alignment; use vortex_buffer::Buffer; @@ -450,11 +450,12 @@ mod tests { #[tokio::test] async fn test_initial_read_size() { - let session = vortex_array::array_session() + let mut builder = vortex_array::default_session_builder() .with::() .with::(); - crate::register_default_encodings(&session); + crate::register_default_encodings(&mut builder); + let session = builder.build(); // Create a large file (> 1MB) let mut buf = ByteBufferMut::empty(); @@ -511,11 +512,16 @@ mod tests { #[cfg(not(target_arch = "wasm32"))] #[tokio::test] async fn test_open_path_uses_memory_session_allocator() { - let session = vortex_array::array_session() + let allocations = Arc::new(AtomicUsize::new(0)); + let mut builder = vortex_array::default_session_builder() .with::() - .with::(); + .with::() + .with_allocator(Arc::new(CountingAllocator { + allocations: Arc::clone(&allocations), + })); - crate::register_default_encodings(&session); + crate::register_default_encodings(&mut builder); + let session = builder.build(); let mut buf = ByteBufferMut::empty(); let array = Buffer::from((0i32..16_384).collect::>()).into_array(); @@ -531,11 +537,7 @@ mod tests { )); std::fs::write(&file_path, ByteBuffer::from(buf).as_slice()).unwrap(); - let allocations = Arc::new(AtomicUsize::new(0)); - let session = session.with_allocator(Arc::new(CountingAllocator { - allocations: Arc::clone(&allocations), - })); - + allocations.store(0, Ordering::Relaxed); let _file = session.open_options().open_path(&file_path).await.unwrap(); std::fs::remove_file(&file_path).unwrap(); diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 32ec3c75f8d..3293bc237f7 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -15,7 +15,6 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::accessor::ArrayAccessor; -use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; @@ -29,6 +28,7 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; +use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; @@ -80,13 +80,13 @@ use crate::VortexFile; use crate::WriteOptionsSessionExt; use crate::footer::SegmentSpec; static SESSION: LazyLock = LazyLock::new(|| { - let session = array_session() + let mut session = default_session_builder() .with::() .with::(); - crate::register_default_encodings(&session); + crate::register_default_encodings(&mut session); - session + session.build() }); #[tokio::test] diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index fc14eaf57d7..12f2f91574f 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -185,9 +185,10 @@ mod tests { use vortex_array::stats::StatsSet; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_io::runtime::Handle; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSession; - use vortex_io::session::RuntimeSessionExt; + use vortex_io::session::RuntimeSessionBuilderExt; use vortex_layout::LayoutReader; use vortex_layout::LayoutStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; @@ -203,11 +204,20 @@ mod tests { use super::*; static SESSION: LazyLock = LazyLock::new(|| { - vortex_array::array_session() + vortex_array::default_session_builder() .with::() .with::() + .build() }); + fn session_with_handle(handle: Handle) -> VortexSession { + vortex_array::default_session_builder() + .with::() + .with::() + .with_handle(handle) + .build() + } + fn test_file_stats(min: i32, max: i32) -> FileStatistics { let mut stats = StatsSet::default(); stats.set(Stat::Min, Precision::exact(ScalarValue::from(min))); @@ -233,7 +243,7 @@ mod tests { #[test] fn pruning_when_filter_out_of_range() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -272,7 +282,7 @@ mod tests { #[test] fn no_pruning_when_filter_in_range() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -312,7 +322,7 @@ mod tests { #[test] fn no_pruning_for_computed_expression_stats() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -352,7 +362,7 @@ mod tests { #[test] fn is_null_pruning_on_nullable_timestamp_column() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -404,7 +414,7 @@ mod tests { #[test] fn pruning_is_not_null_when_file_is_all_null() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 2afd61f7b23..09b0b47b940 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -196,8 +196,7 @@ impl VortexWriteOptions { // buffer stream, so we don't need to poll it until all buffers have been drained. let ctx2 = ctx.clone(); let session = self.session.clone(); - let layout_fut = self.session.handle().spawn_nested(move |h| async move { - let session = session.with_handle(h); + let layout_fut = self.session.handle().spawn_nested(move |_h| async move { let layout = self .strategy .write_stream( diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index bd664cc24ed..c2f70c54bfe 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -30,13 +30,13 @@ use vortex_layout::layouts::table::TableStrategy; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session() + let mut builder = vortex_array::default_session_builder() .with::() .with::(); - vortex_file::register_default_encodings(&session); + vortex_file::register_default_encodings(&mut builder); - session + builder.build() }); #[tokio::test] diff --git a/vortex-geo/src/extension/point.rs b/vortex-geo/src/extension/point.rs index 19e33c212f5..e48ab00d9c9 100644 --- a/vortex-geo/src/extension/point.rs +++ b/vortex-geo/src/extension/point.rs @@ -312,7 +312,7 @@ mod tests { /// A `Point` column round-trips through scalar execution back to the original coordinates. #[test] fn point_unpacks_coordinates() -> VortexResult<()> { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let points = point_column(vec![1.0, -111.7610], vec![2.0, 34.8697])?; diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 951d93b7b4f..b24ffdcc2dd 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -3,10 +3,10 @@ use std::sync::Arc; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::dtype::session::DTypeSessionExt; -use vortex_array::scalar_fn::session::ScalarFnSessionExt; -use vortex_session::VortexSession; +use vortex_array::arrow::ArrowSession; +use vortex_array::dtype::session::DTypeSession; +use vortex_array::scalar_fn::session::ScalarFnSession; +use vortex_session::VortexSessionBuilder; use crate::extension::Point; use crate::extension::Polygon; @@ -21,18 +21,24 @@ mod test_harness; mod tests; /// Set up a session with support for geospatial extension types, encodings and layouts. -pub fn initialize(session: &VortexSession) { +pub fn initialize(session: &mut VortexSessionBuilder) { // Register the geospatial extension types. - session.dtypes().register(WellKnownBinary); - session.arrow().register_exporter(Arc::new(WellKnownBinary)); - session.arrow().register_importer(Arc::new(WellKnownBinary)); - session.dtypes().register(Point); - session.arrow().register_exporter(Arc::new(Point)); - session.arrow().register_importer(Arc::new(Point)); - session.dtypes().register(Polygon); - session.arrow().register_exporter(Arc::new(Polygon)); - session.arrow().register_importer(Arc::new(Polygon)); + { + let dtypes = session.get_mut::(); + dtypes.register(WellKnownBinary); + dtypes.register(Point); + dtypes.register(Polygon); + } + { + let arrow = session.get_mut::(); + arrow.register_exporter(Arc::new(WellKnownBinary)); + arrow.register_importer(Arc::new(WellKnownBinary)); + arrow.register_exporter(Arc::new(Point)); + arrow.register_importer(Arc::new(Point)); + arrow.register_exporter(Arc::new(Polygon)); + arrow.register_importer(Arc::new(Polygon)); + } // Register the geometry scalar functions. - session.scalar_fns().register(GeoDistance); + session.get_mut::().register(GeoDistance); } diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index feb7ea833aa..bfd4577d0b7 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -164,7 +164,7 @@ mod tests { /// (3–4–5 triangles), computed via the geo crate. #[test] fn distance_over_points() -> VortexResult<()> { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let a = point_column(vec![0.0, 3.0, 0.0, 3.0], vec![0.0, 0.0, 4.0, 4.0])?; @@ -178,7 +178,7 @@ mod tests { /// Column-to-column distance pairs corresponding rows of the two columns. #[test] fn distance_between_columns() -> VortexResult<()> { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let a = point_column(vec![0.0, 1.0], vec![0.0, 1.0])?; @@ -192,7 +192,7 @@ mod tests { /// The constant query point may be either operand; distance is symmetric. #[test] fn distance_with_constant_first_operand() -> VortexResult<()> { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let a = point_constant(0.0, 0.0, 4, &mut ctx)?; @@ -206,7 +206,7 @@ mod tests { /// Two constant operands: every row has the same distance. #[test] fn distance_between_two_constants() -> VortexResult<()> { - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); let mut ctx = session.create_execution_ctx(); let a = point_constant(0.0, 0.0, 3, &mut ctx)?; diff --git a/vortex-geo/src/tests/mod.rs b/vortex-geo/src/tests/mod.rs index 546de758eba..92acbb3c303 100644 --- a/vortex-geo/src/tests/mod.rs +++ b/vortex-geo/src/tests/mod.rs @@ -13,7 +13,7 @@ use vortex_session::VortexSession; /// A session with the geospatial types and functions registered. static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + crate::initialize(&mut builder); + builder.build() }); diff --git a/vortex-io/src/session.rs b/vortex-io/src/session.rs index 52acb71339d..5e642fd5ca6 100644 --- a/vortex-io/src/session.rs +++ b/vortex-io/src/session.rs @@ -7,7 +7,7 @@ use std::fmt::Debug; use vortex_error::VortexExpect; use vortex_session::SessionExt; use vortex_session::SessionVar; -use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; use crate::runtime::Handle; @@ -50,23 +50,31 @@ pub trait RuntimeSessionExt: SessionExt { .vortex_expect("Runtime handle not configured in Vortex session. Please setup a `CurrentThreadRuntime`, or configure the session for `with_tokio`.") .clone() } +} +impl RuntimeSessionExt for S {} +/// Extension trait for configuring runtime session data before a session is built. +pub trait RuntimeSessionBuilderExt { /// Configure the runtime session to use the application's Tokio runtime. /// /// For example, if the application is launched using `#[tokio::main]`. #[cfg(feature = "tokio")] - fn with_tokio(self) -> VortexSession { + fn with_tokio(self) -> Self; + + /// Configure the runtime session to use a specific Vortex runtime handle. + fn with_handle(self, handle: Handle) -> Self; +} + +impl RuntimeSessionBuilderExt for VortexSessionBuilder { + #[cfg(feature = "tokio")] + fn with_tokio(mut self) -> Self { use crate::runtime::tokio::TokioRuntime; - let mut builder = self.session().to_builder(); - builder.get_mut::().handle = Some(TokioRuntime::current()); - builder.build() + self.get_mut::().handle = Some(TokioRuntime::current()); + self } - /// Configure the runtime session to use a specific Vortex runtime handle. - fn with_handle(self, handle: Handle) -> VortexSession { - let mut builder = self.session().to_builder(); - builder.get_mut::().handle = Some(handle); - builder.build() + fn with_handle(mut self, handle: Handle) -> Self { + self.get_mut::().handle = Some(handle); + self } } -impl RuntimeSessionExt for S {} diff --git a/vortex-jni/src/session.rs b/vortex-jni/src/session.rs index 9adaf544431..14b571f7a7d 100644 --- a/vortex-jni/src/session.rs +++ b/vortex-jni/src/session.rs @@ -8,7 +8,7 @@ use jni::objects::JClass; use jni::sys::jlong; use vortex::VortexSessionDefault; use vortex::io::runtime::BlockingRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; use crate::RUNTIME; @@ -16,8 +16,9 @@ use crate::RUNTIME; /// Constructs a fresh [`VortexSession`] bound to the JNI-shared tokio runtime and returns /// an opaque pointer that Java must pass to [`Java_dev_vortex_jni_NativeSession_free`]. pub(crate) fn new_session() -> Box { - let session = VortexSession::default().with_handle(RUNTIME.handle()); - vortex_parquet_variant::initialize(&session); + let mut builder = VortexSession::default_builder(); + vortex_parquet_variant::initialize(&mut builder); + let session = builder.with_handle(RUNTIME.handle()).build(); Box::new(session) } diff --git a/vortex-json/src/arrow.rs b/vortex-json/src/arrow.rs index 9db94d86743..a860dc3feb8 100644 --- a/vortex-json/src/arrow.rs +++ b/vortex-json/src/arrow.rs @@ -168,8 +168,9 @@ mod tests { /// Export a JSON extension array to Arrow's canonical JSON extension. #[test] fn exports_json_extension_array_as_arrow_json() -> VortexResult<()> { - let session = vortex_array::array_session(); - initialize(&session); + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + let session = builder.build(); let storage = VarBinArray::from_iter( [Some("{\"id\":1}"), Some("{\"id\":2}")], @@ -206,8 +207,9 @@ mod tests { /// Import Arrow's canonical JSON extension as a Vortex JSON extension array. #[test] fn imports_arrow_json_extension_array_as_vortex_json() -> VortexResult<()> { - let session = vortex_array::array_session(); - initialize(&session); + let mut builder = vortex_array::default_session_builder(); + initialize(&mut builder); + let session = builder.build(); let mut field = Field::new("data", DataType::Utf8, false); field.try_with_extension_type(ArrowJson::default())?; diff --git a/vortex-json/src/json_to_variant.rs b/vortex-json/src/json_to_variant.rs index d0febf4e646..b324f68a599 100644 --- a/vortex-json/src/json_to_variant.rs +++ b/vortex-json/src/json_to_variant.rs @@ -306,7 +306,6 @@ mod tests { use vortex_array::expr::proto::ExprSerializeProtoExt; use vortex_array::expr::root; use vortex_array::scalar_fn::session::ScalarFnSession; - use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_array::session::ArraySession; use super::*; @@ -318,12 +317,12 @@ mod tests { /// A session that knows the `JsonToVariant` definition but has no Variant encoding registered, /// so executing the function exercises the fallback that errors when no kernel is present. fn session() -> VortexSession { - let session = VortexSession::empty() + let mut session = VortexSession::builder() .with::() .with::() .with::(); - session.scalar_fns().register(JsonToVariant); - session + session.get_mut::().register(JsonToVariant); + session.build() } #[test] diff --git a/vortex-json/src/lib.rs b/vortex-json/src/lib.rs index 21f86ca9a7c..2d647a83eb6 100644 --- a/vortex-json/src/lib.rs +++ b/vortex-json/src/lib.rs @@ -20,15 +20,18 @@ pub use json_to_variant::JsonToVariant; pub use json_to_variant::JsonToVariantOptions; pub use json_to_variant::ShreddingSpec; pub use json_to_variant::json_to_variant; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::dtype::session::DTypeSessionExt; -use vortex_array::scalar_fn::session::ScalarFnSessionExt; -use vortex_session::VortexSession; +use vortex_array::arrow::ArrowSession; +use vortex_array::dtype::session::DTypeSession; +use vortex_array::scalar_fn::session::ScalarFnSession; +use vortex_session::VortexSessionBuilder; /// Register JSON extension support with a session. -pub fn initialize(session: &VortexSession) { - session.dtypes().register(Json); - session.arrow().register_exporter(Arc::new(Json)); - session.arrow().register_importer(Arc::new(Json)); - session.scalar_fns().register(JsonToVariant); +pub fn initialize(session: &mut VortexSessionBuilder) { + session.get_mut::().register(Json); + + let arrow = session.get_mut::(); + arrow.register_exporter(Arc::new(Json)); + arrow.register_importer(Arc::new(Json)); + + session.get_mut::().register(JsonToVariant); } diff --git a/vortex-layout/src/display.rs b/vortex-layout/src/display.rs index 5acf685e946..8ca68fb15e2 100644 --- a/vortex-layout/src/display.rs +++ b/vortex-layout/src/display.rs @@ -230,7 +230,6 @@ mod tests { use vortex_buffer::BitBufferMut; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use crate::IntoLayout; use crate::OwnedLayoutChildren; @@ -242,7 +241,6 @@ mod tests { use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::strategy::LayoutStrategy; - use crate::test::SESSION; /// Test display_tree with inline array_tree metadata (no segment source needed). #[test] @@ -251,7 +249,7 @@ mod tests { if std::env::var("NEXTEST_RUN_ID").is_ok() { temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", Some("1"), || { block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -349,7 +347,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5 if std::env::var("NEXTEST_RUN_ID").is_ok() { temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", None::<&str>, || { block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -419,7 +417,7 @@ vortex.chunked, dtype: i32, children: 2, rows: 10 // Create a simple primitive array let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::AllValid); let layout = block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); FlatLayoutStrategy::default() .write_stream( ctx.clone(), @@ -464,7 +462,7 @@ vortex.flat, dtype: i32?, segment 0, buffers=[20B], total=20B // Create a simple primitive array let array = PrimitiveArray::new(buffer![10i64, 20, 30], Validity::NonNullable); let layout = block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); FlatLayoutStrategy::default() .write_stream( ctx, diff --git a/vortex-layout/src/flatbuffers.rs b/vortex-layout/src/flatbuffers.rs index 92d34cd5de3..7b0db82718a 100644 --- a/vortex-layout/src/flatbuffers.rs +++ b/vortex-layout/src/flatbuffers.rs @@ -272,7 +272,9 @@ mod tests { LayoutEncodingId::new("vortex.test.foreign_child_layout"), ]); let array_ctx = ReadContext::new([]); - let session = vortex_array::array_session().with::(); + let session = vortex_array::default_session_builder() + .with::() + .build(); let layout = layout_from_flatbuffer_with_options( layout_buffer, diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index a9870229792..2cc70fc8f9f 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -362,7 +362,6 @@ mod test { use vortex_array::expr::root; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use vortex_session::registry::ReadContext; use crate::IntoLayout; @@ -392,7 +391,7 @@ mod test { let (mut sequence_id, eof) = SequenceId::root().split(); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); strategy .write_stream( ctx, diff --git a/vortex-layout/src/layouts/chunked/writer.rs b/vortex-layout/src/layouts/chunked/writer.rs index 29596fadff0..65865f759ed 100644 --- a/vortex-layout/src/layouts/chunked/writer.rs +++ b/vortex-layout/src/layouts/chunked/writer.rs @@ -66,8 +66,7 @@ impl LayoutStrategy for ChunkedLayoutStrategy { let dtype = dtype2.clone(); let session = session.clone(); - yield handle.spawn_nested(move |handle| async move { - let session = session.with_handle(handle); + yield handle.spawn_nested(move |_handle| async move { chunk_strategy .write_stream( ctx, diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index a4b3e20cbff..bbe7b8fcbbd 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -339,11 +339,11 @@ mod tests { use vortex_array::IntoArray as _; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinArray; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; @@ -365,7 +365,7 @@ mod tests { use vortex_io::runtime::Handle; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSession; - use vortex_io::session::RuntimeSessionExt; + use vortex_io::session::RuntimeSessionBuilderExt; use vortex_session::VortexSession; use crate::LayoutId; @@ -386,10 +386,11 @@ mod tests { // FIXME(ngates): Deprecate the global `runtime::single::block_on` helper and require tests // to call `block_on` on an explicit runtime instance. fn session_with_handle(handle: Handle) -> VortexSession { - array_session() + default_session_builder() .with::() .with::() .with_handle(handle) + .build() } async fn write_dict_layout( @@ -704,7 +705,7 @@ mod tests { } fn test_apply(original: Expression, outer: Expression, inner: Expression) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = VarBinArray::from_iter( [Some("abc"), Some("def"), None], DType::Utf8(Nullability::Nullable), diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 1a376452984..f56ac071af4 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -191,8 +191,7 @@ impl LayoutStrategy for DictStrategy { let ctx2 = ctx.clone(); let segment_sink2 = Arc::clone(&segment_sink); let session2 = session.clone(); - let codes_fut = handle.spawn_nested(move |h| async move { - let session2 = session2.with_handle(h); + let codes_fut = handle.spawn_nested(move |_h| async move { codes.write_stream( ctx2, segment_sink2, @@ -208,8 +207,7 @@ impl LayoutStrategy for DictStrategy { let segment_sink2 = Arc::clone(&segment_sink); let dtype2 = dtype2.clone(); let session2 = session.clone(); - let values_layout = handle.spawn_nested(move |h| async move { - let session2 = session2.with_handle(h); + let values_layout = handle.spawn_nested(move |_h| async move { values.write_stream( ctx2, segment_sink2, @@ -608,7 +606,7 @@ mod tests { use crate::sequence::SequentialStreamExt; static SESSION: LazyLock = - LazyLock::new(|| VortexSession::empty().with::()); + LazyLock::new(|| VortexSession::builder().with::().build()); /// Regression test for a bug where the codes stream dtype was hardcoded to U16 instead of /// using the actual codes dtype from the array. When `max_len <= 255`, the dict encoder diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 983f9e59a58..002b1e0ee3e 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -498,10 +498,10 @@ impl FileStatsAccumulator { #[cfg(test)] mod tests { use rstest::rstest; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::builders::VarBinViewBuilder; + use vortex_array::default_session_builder; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; @@ -511,7 +511,7 @@ mod tests { #[case(DType::Utf8(Nullability::NonNullable))] #[case(DType::Binary(Nullability::NonNullable))] fn truncates_accumulated_stats(#[case] dtype: DType) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let mut builder = VarBinViewBuilder::with_capacity(dtype.clone(), 2); builder.append_value("Value to be truncated"); builder.append_value("untruncated"); @@ -556,7 +556,7 @@ mod tests { #[test] fn always_adds_is_truncated_column() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = buffer![0, 1, 2].into_array(); let mut acc = StatsAccumulator::new(array.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12); acc.push_chunk(&array, &mut ctx) diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 50ec634a61a..71c7fdf8a19 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -244,19 +244,17 @@ mod test { use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use crate::LayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; - use crate::test::SESSION; #[test] fn flat_identity() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -296,7 +294,7 @@ mod test { #[test] fn flat_expr() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); @@ -336,7 +334,7 @@ mod test { #[test] fn flat_unaligned_row_mask() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index f6b6d5acf02..0851cecd56b 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -209,7 +209,6 @@ mod tests { use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; @@ -218,6 +217,7 @@ mod tests { use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinViewBuilder; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; @@ -233,7 +233,6 @@ mod tests { use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_utils::aliases::hash_set::HashSet; @@ -251,7 +250,7 @@ mod tests { #[test] fn flat_stats() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -289,7 +288,7 @@ mod tests { #[test] fn truncates_variable_size_stats() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -350,8 +349,8 @@ mod tests { #[test] fn struct_array_round_trip() { block_on(|handle| async { - let mut ctx_exec = array_session().create_execution_ctx(); - let session = SESSION.clone().with_handle(handle); + let mut ctx_exec = default_session_builder().build().create_execution_ctx(); + let session = crate::test::session_with_handle(handle); let mut validity_builder = BitBufferMut::with_capacity(2); validity_builder.append(true); validity_builder.append(false); @@ -437,7 +436,7 @@ mod tests { #[test] fn flat_invalid_array_fails() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let prim: PrimitiveArray = (0..10).collect(); let filter = prim.filter(Mask::from_indices(10, vec![2, 3]))?; @@ -477,7 +476,7 @@ mod tests { #[test] fn flat_valid_array_writes() -> VortexResult<()> { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let codes: PrimitiveArray = (0u32..10).collect(); let values: PrimitiveArray = (0..10).collect(); let dict = DictArray::new(codes.into_array(), values.into_array()); diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 89a1317d240..fbf0055714a 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -277,18 +277,17 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SharedArray; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use super::*; use crate::LayoutStrategy; @@ -297,7 +296,6 @@ mod tests { use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; - use crate::test::SESSION; const ONE_MEG: u64 = 1 << 20; @@ -398,7 +396,7 @@ mod tests { let stream = fsl.into_array().to_array_stream().sequenced(ptr); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); strategy .write_stream( ctx, @@ -463,7 +461,7 @@ mod tests { let stream = elements.into_array().to_array_stream().sequenced(ptr); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); strategy .write_stream( ctx, @@ -490,7 +488,7 @@ mod tests { /// `pop_front` subtracted the larger Cached-era values. #[test] fn chunks_buffer_pop_front_no_panic_after_shared_execution() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let n = 20_000usize; let block_len = 10_000usize; diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 4e5e8e494fa..f9378acf212 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -337,7 +337,6 @@ mod tests { use vortex_array::expr::root; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use crate::LayoutReader; use crate::LayoutStrategy; @@ -347,12 +346,11 @@ mod tests { use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; - use crate::test::SESSION; #[test] fn flat_expr_no_row_id() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -397,7 +395,7 @@ mod tests { #[test] fn flat_expr_row_id() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); @@ -442,7 +440,7 @@ mod tests { #[test] fn flat_expr_or() { block_on(|handle| async { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); let mut ctx = session.create_execution_ctx(); let array_ctx = ArrayContext::empty(); let segments = Arc::new(TestSegments::default()); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index ab03a4624fa..eab20ba987f 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -396,13 +396,13 @@ mod tests { use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; use vortex_array::assert_nth_scalar; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::Nullability; @@ -422,7 +422,6 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; use crate::LayoutRef; @@ -447,7 +446,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); strategy .write_stream( ctx, @@ -484,7 +483,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); strategy .write_stream( ctx, @@ -524,7 +523,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); strategy .write_stream( ctx, @@ -569,7 +568,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); strategy .write_stream( ctx, @@ -677,7 +676,7 @@ mod tests { fn test_struct_layout_select( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let reader = layout .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap(); @@ -740,7 +739,10 @@ mod tests { // ...and the result is masked with the validity of the parent StructArray assert_eq!( result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap(), Scalar::null(result.dtype().clone()), ); @@ -785,7 +787,10 @@ mod tests { // Row 0: struct is valid, field "c" is 4. assert_eq!( result - .execute_scalar(0, &mut array_session().create_execution_ctx()) + .execute_scalar( + 0, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_struct() .field_by_idx(0) @@ -796,7 +801,10 @@ mod tests { // Row 1: struct is null (because root.a.b was null at this row). assert!( result - .execute_scalar(1, &mut array_session().create_execution_ctx()) + .execute_scalar( + 1, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_struct() .is_null() @@ -805,7 +813,10 @@ mod tests { // Row 2: struct is valid, field "c" is 6. assert_eq!( result - .execute_scalar(2, &mut array_session().create_execution_ctx()) + .execute_scalar( + 2, + &mut default_session_builder().build().create_execution_ctx() + ) .unwrap() .as_struct() .field_by_idx(0) @@ -848,7 +859,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = crate::test::session_with_handle(handle); strategy .write_stream( ctx, diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 28b3af52c87..14d137c709e 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -322,7 +322,7 @@ impl LayoutStrategy for TableStrategy { let session = session.clone(); let ctx = ctx.clone(); let segment_sink = Arc::clone(&segment_sink); - handle.spawn_nested(move |h| { + handle.spawn_nested(move |_h| { let validity = Arc::clone(&self.validity); // descend further and try with new fields let writer = self @@ -338,8 +338,6 @@ impl LayoutStrategy for TableStrategy { Arc::clone(&self.fallback) } }); - let session = session.with_handle(h); - async move { // If we have a matching writer, we use it. // Otherwise, we descend into a new modified one. diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 3d82fa27b34..33ca7fd450c 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -548,7 +548,9 @@ mod tests { }; let deserialized = ZonedMetadata::deserialize(&metadata.serialize())?; - let session = VortexSession::empty().with::(); + let session = VortexSession::builder() + .with::() + .build(); let aggregate_fns = aggregate_fns_from_specs(&deserialized.aggregate_specs, &session)?; assert_eq!(aggregate_fns.as_ref(), std::slice::from_ref(&aggregate_fn)); @@ -619,7 +621,7 @@ mod tests { ) .into_layout(), ]); - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); let build_read_ctx = ReadContext::new([]); let build_ctx = LayoutBuildContext { session: &session, @@ -650,7 +652,7 @@ mod tests { aggregate_specs: Arc::new([]), }; let children = OwnedLayoutChildren::layout_children(vec![]); - let session = vortex_array::array_session(); + let session = vortex_array::default_session_builder().build(); let build_read_ctx = ReadContext::new([]); let build_ctx = LayoutBuildContext { session: &session, diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 8feaf46fa03..70486c38578 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -227,10 +227,10 @@ mod test { use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use vortex_array::expr::gt; use vortex_array::expr::is_not_null; use vortex_array::expr::lit; @@ -241,7 +241,7 @@ mod test { use vortex_io::runtime::Handle; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSession; - use vortex_io::session::RuntimeSessionExt; + use vortex_io::session::RuntimeSessionBuilderExt; use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; @@ -267,10 +267,11 @@ mod test { use crate::session::LayoutSession; fn session_with_handle(handle: Handle) -> VortexSession { - array_session() + default_session_builder() .with::() .with::() .with_handle(handle) + .build() } #[fixture] @@ -311,7 +312,7 @@ mod test { #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), ) { block_on(|handle| async { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let session = session_with_handle(handle); let result = layout .new_reader("".into(), segments, &session, &Default::default()) @@ -432,7 +433,7 @@ mod test { let zoned_layout = layout.as_::(); let children = OwnedLayoutChildren::layout_children(vec![layout.child(0)?, layout.child(1)?]); - let session = array_session(); + let session = default_session_builder().build(); let read_ctx = ReadContext::new([]); let build_ctx = LayoutBuildContext { session: &session, diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 11fd5c7b882..12013aad9eb 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -469,8 +469,8 @@ mod test { use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; + use vortex_array::default_session_builder; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldPath; @@ -730,7 +730,7 @@ mod test { #[test] fn into_stream_executes_after_prepare() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let calls = Arc::new(AtomicUsize::new(0)); let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls))); @@ -874,7 +874,7 @@ mod test { #[test] fn into_stream_with_row_range() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let calls = Arc::new(AtomicUsize::new(0)); let reader = Arc::new(SplittingLayoutReader::new(Arc::clone(&calls))); diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 2af312f7537..6c2cde11b72 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -136,7 +136,6 @@ mod test { use vortex_array::expr::Expression; use vortex_buffer::buffer; use vortex_io::runtime::single::block_on; - use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; use super::*; @@ -145,6 +144,7 @@ mod test { use crate::RowSplits; use crate::layouts::flat::writer::FlatLayoutStrategy; use crate::scan::test::SCAN_SESSION; + use crate::scan::test::session_with_handle; use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; @@ -154,7 +154,7 @@ mod test { let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let layout = block_on(|handle| async { - let session = SCAN_SESSION.clone().with_handle(handle); + let session = session_with_handle(handle); FlatLayoutStrategy::default() .write_stream( ctx, diff --git a/vortex-layout/src/scan/test.rs b/vortex-layout/src/scan/test.rs index 809639d169d..cf63394a851 100644 --- a/vortex-layout/src/scan/test.rs +++ b/vortex-layout/src/scan/test.rs @@ -5,19 +5,24 @@ use std::sync::LazyLock; use vortex_io::runtime::Handle; use vortex_io::session::RuntimeSession; -use vortex_io::session::RuntimeSessionExt; +use vortex_io::session::RuntimeSessionBuilderExt; use vortex_session::VortexSession; use crate::session::LayoutSession; pub fn new_session() -> VortexSession { - vortex_array::array_session() + vortex_array::default_session_builder() .with::() .with::() + .build() } pub fn session_with_handle(handle: Handle) -> VortexSession { - new_session().with_handle(handle) + vortex_array::default_session_builder() + .with::() + .with::() + .with_handle(handle) + .build() } pub static SCAN_SESSION: LazyLock = LazyLock::new(new_session); diff --git a/vortex-layout/src/test.rs b/vortex-layout/src/test.rs index 72258a1161e..07a709562cf 100644 --- a/vortex-layout/src/test.rs +++ b/vortex-layout/src/test.rs @@ -3,13 +3,24 @@ use std::sync::LazyLock; +use vortex_io::runtime::Handle; use vortex_io::session::RuntimeSession; +use vortex_io::session::RuntimeSessionBuilderExt; use vortex_session::VortexSession; use crate::session::LayoutSession; pub static SESSION: LazyLock = LazyLock::new(|| { - vortex_array::array_session() + vortex_array::default_session_builder() .with::() .with::() + .build() }); + +pub fn session_with_handle(handle: Handle) -> VortexSession { + vortex_array::default_session_builder() + .with::() + .with::() + .with_handle(handle) + .build() +} diff --git a/vortex-python/src/arrays/range_to_sequence.rs b/vortex-python/src/arrays/range_to_sequence.rs index eed03c64d13..54a25320556 100644 --- a/vortex-python/src/arrays/range_to_sequence.rs +++ b/vortex-python/src/arrays/range_to_sequence.rs @@ -79,7 +79,7 @@ mod test { use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; + use vortex_array::default_session_builder; use crate::arrays::range_to_sequence::range_len; use crate::arrays::range_to_sequence::sequence_array_from_range; @@ -100,7 +100,7 @@ mod test { #[test] fn test_sequence_array_from_len() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let dtype = DType::Primitive(PType::U16, Nullability::NonNullable); let arr = sequence_array_from_range::(0, 10, 1, dtype).unwrap(); assert_arrays_eq!( diff --git a/vortex-python/src/cli.rs b/vortex-python/src/cli.rs index 4033da1ad3b..498dd8da0ed 100644 --- a/vortex-python/src/cli.rs +++ b/vortex-python/src/cli.rs @@ -11,7 +11,7 @@ use vortex::error::VortexError; use vortex::error::VortexExpect as _; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::tokio::TokioRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; use crate::install_module; @@ -23,8 +23,11 @@ static TOKIO_RUNTIME: LazyLock = LazyLock::new(|| { }); static TUI_RUNTIME: LazyLock = LazyLock::new(|| TokioRuntime::new(TOKIO_RUNTIME.handle().clone())); -static TUI_SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_handle(TUI_RUNTIME.handle())); +static TUI_SESSION: LazyLock = LazyLock::new(|| { + VortexSession::default_builder() + .with_handle(TUI_RUNTIME.handle()) + .build() +}); pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { let m = PyModule::new(py, "cli")?; diff --git a/vortex-python/src/session.rs b/vortex-python/src/session.rs index 5598241d2ef..04b9344b46a 100644 --- a/vortex-python/src/session.rs +++ b/vortex-python/src/session.rs @@ -5,13 +5,16 @@ use std::sync::LazyLock; use vortex::VortexSessionDefault; use vortex::io::runtime::BlockingRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; use crate::RUNTIME; -static SESSION: LazyLock = - LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle())); +static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::default_builder() + .with_handle(RUNTIME.handle()) + .build() +}); pub(crate) fn session() -> &'static VortexSession { &SESSION diff --git a/vortex-row/benches/row_encode.rs b/vortex-row/benches/row_encode.rs index 0747c5facac..6be751c75b7 100644 --- a/vortex-row/benches/row_encode.rs +++ b/vortex-row/benches/row_encode.rs @@ -41,7 +41,8 @@ static GLOBAL: MiMalloc = MiMalloc; const N: usize = 100_000; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::default_session_builder().build()); fn main() { divan::main(); diff --git a/vortex-row/src/encoder.rs b/vortex-row/src/encoder.rs index c1d7a8f2a54..ddd2c28cc18 100644 --- a/vortex-row/src/encoder.rs +++ b/vortex-row/src/encoder.rs @@ -28,12 +28,12 @@ use crate::size::RowSize; /// # Example /// /// ```rust -/// use vortex_array::{IntoArray, VortexSessionExecute, array_session}; +/// use vortex_array::{IntoArray, VortexSessionExecute, default_session_builder}; /// use vortex_array::arrays::PrimitiveArray; /// use vortex_row::{RowEncoder, RowSortField}; /// /// # fn example() -> vortex_error::VortexResult<()> { -/// let mut ctx = array_session().create_execution_ctx(); +/// let mut ctx = default_session_builder().build().create_execution_ctx(); /// let ids = PrimitiveArray::from_iter([3i32, 1, 2]).into_array(); /// let scores = PrimitiveArray::from_iter([10i32, 20, 10]).into_array(); /// diff --git a/vortex-row/src/lib.rs b/vortex-row/src/lib.rs index f33fa8a4fcb..0b57aa91e02 100644 --- a/vortex-row/src/lib.rs +++ b/vortex-row/src/lib.rs @@ -15,7 +15,7 @@ //! - [`convert_columns`] and [`compute_row_sizes`], compatibility helpers around //! [`RowEncoder`]. //! - [`initialize`], which registers the [`RowSize`] and [`RowEncode`] scalar functions on a -//! [`VortexSession`]. +//! [`VortexSession`][vortex_session::VortexSession]. //! //! Internally, encoding is split into two scalar functions. [`RowSize`] performs the sizing //! pass and classifies fixed-width versus variable-width input columns. [`RowEncode`] uses @@ -48,16 +48,17 @@ pub use encoder::convert_columns_with_options; pub use options::RowEncodingOptions; pub use options::RowSortField; pub use size::RowSize; -use vortex_array::scalar_fn::session::ScalarFnSessionExt; -use vortex_session::VortexSession; +use vortex_array::scalar_fn::session::ScalarFnSession; +use vortex_session::VortexSessionBuilder; /// Register the row-encoding scalar functions ([`RowSize`] and [`RowEncode`]) on the given -/// session. +/// session builder. /// /// Call this during session construction when row encoding must be available through the /// expression layer. The direct [`RowEncoder`] API constructs the scalar-function calls /// itself and does not require global registration. -pub fn initialize(session: &VortexSession) { - session.scalar_fns().register(RowSize); - session.scalar_fns().register(RowEncode); +pub fn initialize(session: &mut VortexSessionBuilder) { + let scalar_fns = session.get_mut::(); + scalar_fns.register(RowSize); + scalar_fns.register(RowEncode); } diff --git a/vortex-row/src/tests.rs b/vortex-row/src/tests.rs index c7f9e8a2b49..9a57dc4037a 100644 --- a/vortex-row/src/tests.rs +++ b/vortex-row/src/tests.rs @@ -8,7 +8,6 @@ use std::f64::consts::PI; use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::ListViewArray; @@ -16,6 +15,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::listview::ListViewArrayExt; +use vortex_array::default_session_builder; use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::Date; use vortex_array::extension::datetime::TimeUnit; @@ -29,7 +29,7 @@ use crate::convert_columns; use crate::convert_columns_with_options; fn collect_row_bytes(array: &ListViewArray) -> Vec> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let nrows = array.len(); (0..nrows) .map(|i| { @@ -43,7 +43,7 @@ fn collect_row_bytes(array: &ListViewArray) -> Vec> { /// Encode each column independently, sort the resulting row bytes, and check the permutation /// matches the natural sort order of `values`. fn assert_sort_order_i64(values: Vec, descending: bool) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let col = PrimitiveArray::from_iter(values.clone()).into_array(); let field = RowSortField::new(descending, true); let encoded = convert_columns(&[col], &[field], &mut ctx)?; @@ -77,7 +77,7 @@ fn primitive_i64_roundtrip(#[case] descending: bool) -> VortexResult<()> { #[test] fn primitive_u32_sort_order() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values: Vec = vec![0, 1, 100, u32::MAX, 42, 17]; let col = PrimitiveArray::from_iter(values.clone()).into_array(); let encoded = convert_columns(&[col], &[RowSortField::default()], &mut ctx)?; @@ -95,7 +95,7 @@ fn primitive_u32_sort_order() -> VortexResult<()> { #[test] fn reject_temporal_extension_dtype_early() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let storage = PrimitiveArray::from_iter([2i32, -1, 0, 7]).into_array(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); let col = ExtensionArray::new(ext_dtype, storage).into_array(); @@ -111,7 +111,7 @@ fn reject_temporal_extension_dtype_early() -> VortexResult<()> { #[test] fn reject_nested_temporal_extension_dtype_early() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let storage = PrimitiveArray::from_iter([2i32, -1, 0, 7]).into_array(); let ext_dtype = Date::new(TimeUnit::Days, Nullability::NonNullable).erased(); let date_col = ExtensionArray::new(ext_dtype, storage).into_array(); @@ -130,7 +130,7 @@ fn reject_nested_temporal_extension_dtype_early() -> VortexResult<()> { #[test] fn primitive_f64_sort_order() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // We use IEEE total-ordering semantics: -0.0 < +0.0 in the byte encoding (matches // `arrow-row`). Avoid -0.0 in the natural-order baseline since partial_cmp says // -0.0 == 0.0. @@ -151,7 +151,7 @@ fn primitive_f64_sort_order() -> VortexResult<()> { #[test] fn bool_sort_order() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let col = BoolArray::from_iter([true, false, true, false]).into_array(); let encoded = convert_columns(&[col], &[RowSortField::default()], &mut ctx)?; let rows = collect_row_bytes(&encoded); @@ -168,7 +168,7 @@ fn bool_sort_order() -> VortexResult<()> { #[test] fn utf8_sort_order() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values = vec![ "banana", "apple", @@ -193,7 +193,7 @@ fn utf8_sort_order() -> VortexResult<()> { #[test] fn multi_column_sort() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ints: Vec = vec![1, 2, 1, 2, 1, 3]; let strs = vec!["b", "a", "a", "b", "c", "z"]; let col0 = PrimitiveArray::from_iter(ints.clone()).into_array(); @@ -216,7 +216,7 @@ fn multi_column_sort() -> VortexResult<()> { #[test] fn nulls_first_and_last() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values: Vec> = vec![Some(5), None, Some(1), None, Some(3)]; let col = PrimitiveArray::from_option_iter(values.clone()).into_array(); @@ -250,7 +250,7 @@ fn nulls_first_and_last() -> VortexResult<()> { #[test] fn reusable_options_helpers() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let options = RowEncodingOptions::new([RowSortField::descending().nulls_last()]); assert_eq!(options.len(), 1); assert!(!options.is_empty()); @@ -282,7 +282,7 @@ fn reusable_options_helpers() -> VortexResult<()> { #[test] fn row_encoder_new_accepts_sort_fields() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let encoder = RowEncoder::new([RowSortField::ascending()]); let col = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); @@ -293,7 +293,7 @@ fn row_encoder_new_accepts_sort_fields() -> VortexResult<()> { #[test] fn default_row_encoder_uses_default_fields() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let col0 = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); let col1 = PrimitiveArray::from_iter([4i32, 5, 6]).into_array(); @@ -305,7 +305,7 @@ fn default_row_encoder_uses_default_fields() -> VortexResult<()> { #[test] fn struct_sort_order() -> VortexResult<()> { use vortex_array::arrays::StructArray; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ids: Vec = vec![3, 1, 3, 1, 2]; let names = vec!["b", "a", "a", "b", "z"]; let id_arr = PrimitiveArray::from_iter(ids.clone()).into_array(); @@ -332,7 +332,7 @@ fn row_size_struct_shape() -> VortexResult<()> { use crate::compute_row_sizes; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let ints: Vec = vec![1, 2, 3, 4, 5]; let strs = vec!["a", "bb", "ccc", "", "eeeee"]; let col0 = PrimitiveArray::from_iter(ints).into_array(); @@ -372,7 +372,7 @@ fn row_size_struct_shape() -> VortexResult<()> { #[test] fn single_buffer_invariant() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Encoded rows here are all > 12 bytes, forcing the Ref-view path that points back into // the shared data buffer. let nrows = 64usize; @@ -411,7 +411,7 @@ fn single_buffer_invariant() -> VortexResult<()> { /// boundaries always align. #[test] fn multi_column_varlen_empty_vs_nul_byte_string() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // col1: empty vs single 0-byte. col2: same int for all rows. let col1 = VarBinViewArray::from_iter_str(["", "\0", "a", "ab"]).into_array(); let col2 = PrimitiveArray::from_iter([1i32, 1, 1, 1]).into_array(); @@ -442,7 +442,7 @@ fn multi_column_varlen_empty_vs_nul_byte_string() -> VortexResult<()> { /// the 3-sentinel scheme null=0x00, empty=0x01 differ at byte 0. #[test] fn multi_column_varlen_null_vs_empty() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let col1 = VarBinViewArray::from_iter_nullable_str([ None::<&str>, Some(""), @@ -505,7 +505,7 @@ fn multi_column_varlen_null_vs_empty() -> VortexResult<()> { /// non-empty's first byte is smaller than empty's first byte. #[test] fn varlen_descending_empty_vs_non_empty() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let col = VarBinViewArray::from_iter_str(["a", "", "abc"]).into_array(); let encoded = convert_columns(&[col], &[RowSortField::descending()], &mut ctx)?; let rows = collect_row_bytes(&encoded); @@ -531,7 +531,7 @@ fn null_struct_rows_with_varying_child_lengths_are_byte_equal() -> VortexResult< use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); // Build a nullable struct{name: utf8} where rows 0 and 2 are null but the underlying // child has different length data ("short" vs "much longer text data"). let names = @@ -556,7 +556,7 @@ fn null_struct_rows_with_varying_child_lengths_are_byte_equal() -> VortexResult< #[test] fn primitive_f32_sort_order() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values: Vec = vec![-1.5, 0.0, 1.5, f32::INFINITY, f32::NEG_INFINITY]; let col = PrimitiveArray::from_iter(values.clone()).into_array(); let encoded = convert_columns(&[col], &[RowSortField::default()], &mut ctx)?; @@ -573,7 +573,7 @@ fn primitive_f32_sort_order() -> VortexResult<()> { #[test] fn primitive_f16_sort_order() -> VortexResult<()> { use vortex_array::dtype::half::f16; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let values: Vec = vec![ f16::from_f32(-1.5), f16::from_f32(0.0), @@ -600,7 +600,7 @@ fn reject_list_dtype_early() { use vortex_array::validity::Validity; use vortex_buffer::buffer; - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let offsets = PrimitiveArray::new(buffer![0u32, 1, 2], Validity::NonNullable).into_array(); let elements = PrimitiveArray::from_iter([10i32, 20]).into_array(); let list: ArrayRef = ListArray::try_new(elements, offsets, Validity::NonNullable) diff --git a/vortex-session/src/immutable.rs b/vortex-session/src/immutable.rs index 8647914b274..336ceed7c05 100644 --- a/vortex-session/src/immutable.rs +++ b/vortex-session/src/immutable.rs @@ -1,13 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The immutable [`VortexSession`] and its [`SessionBuilder`]. +//! The immutable [`VortexSession`] and its [`VortexSessionBuilder`]. //! //! [`VortexSession`] is backed by an immutable [`HashMap`]: once built, its set of session -//! variables can never change, so reads are lock-free and return plain references. To modify -//! a session, clone its state into a mutable [`SessionBuilder`] with -//! [`VortexSession::to_builder`], apply changes, and call [`SessionBuilder::build`] to freeze it -//! back into a new [`VortexSession`]. +//! variables can never change, so reads are lock-free and return plain references. Build all +//! mutable state in a [`VortexSessionBuilder`], then call [`VortexSessionBuilder::build`] to +//! freeze it into a [`VortexSession`]. use std::any::TypeId; use std::any::type_name; @@ -24,24 +23,9 @@ use crate::SessionVar; use crate::UnknownPluginPolicy; /// A [`SessionVar`] that can be stored in a [`VortexSession`]. -/// -/// This trait is implemented automatically for every [`SessionVar`] that is also [`Clone`], -/// so types opt in by implementing [`Clone`] rather than implementing this trait directly. -/// -/// Cloning must produce an *independent* copy: [`SessionBuilder`] relies on it to detach a -/// session's state before mutating it. A variable whose `Clone` shares interior-mutable state -/// (e.g. an `Arc` around a mutable registry) will leak builder mutations back into the session -/// it was cloned from. -pub trait VortexSessionVar: SessionVar { - /// Clones this variable into a new boxed instance. - fn clone_var(&self) -> Box; -} +pub trait VortexSessionVar: SessionVar {} -impl VortexSessionVar for V { - fn clone_var(&self) -> Box { - Box::new(self.clone()) - } -} +impl VortexSessionVar for V {} /// The type map backing a [`VortexSession`]. Immutable once wrapped in an `Arc`. type VortexSessionVars = HashMap, BuildHasherDefault>; @@ -52,8 +36,7 @@ type VortexSessionVars = HashMap, BuildHasherD /// It is also the entry-point passed to dynamic libraries to initialize Vortex plugins. /// /// The set of session variables is fixed at construction time, so lookups take no locks and -/// hand out plain references. To change a session's state, clone it into a mutable -/// [`SessionBuilder`] with [`VortexSession::to_builder`] and build a new session from it. +/// hand out plain references. #[derive(Clone, Debug)] pub struct VortexSession(Arc); @@ -63,23 +46,9 @@ impl VortexSession { Self(Default::default()) } - /// Create an empty [`SessionBuilder`]. - pub fn builder() -> SessionBuilder { - SessionBuilder::default() - } - - /// Clone this session's state into a mutable [`SessionBuilder`]. - /// - /// The builder owns an independent copy of every variable, so mutations made through it - /// never affect this session. - pub fn to_builder(&self) -> SessionBuilder { - SessionBuilder { - vars: self - .0 - .iter() - .map(|(id, var)| (*id, var.clone_var())) - .collect(), - } + /// Create an empty [`VortexSessionBuilder`]. + pub fn builder() -> VortexSessionBuilder { + VortexSessionBuilder::default() } /// Returns the session variable of type `V`. @@ -116,41 +85,20 @@ impl VortexSession { self.get_opt::() .is_some_and(|p| p.allow_unknown) } - - /// Inserts a new session variable of type `V` with its default value, returning the - /// updated session. - /// - /// # Panics - /// - /// If a variable of that type already exists. - pub fn with(self) -> Self { - self.to_builder().with::().build() - } - - /// Inserts a new session variable of type `V`, returning the updated session. - /// - /// # Panics - /// - /// If a variable of that type already exists. - pub fn with_some(self, var: V) -> Self { - self.to_builder().with_some(var).build() - } - - /// Allow deserializing unknown plugin IDs as non-executable foreign placeholders. - pub fn allow_unknown(self) -> Self { - self.to_builder().allow_unknown().build() - } } /// A mutable builder for [`VortexSession`]. /// /// Holds a private, mutable copy of a session's state. Modify it freely, then call -/// [`SessionBuilder::build`] to freeze it into a new immutable [`VortexSession`]. +/// [`VortexSessionBuilder::build`] to freeze it into a new immutable [`VortexSession`]. #[derive(Debug, Default)] pub struct SessionBuilder { vars: VortexSessionVars, } +/// Public builder type for constructing [`VortexSession`] values. +pub type VortexSessionBuilder = SessionBuilder; + impl SessionBuilder { /// Inserts a new session variable of type `V` with its default value. /// @@ -258,23 +206,9 @@ mod tests { assert!(session.get_opt::().is_none()); } - #[test] - fn to_builder_does_not_mutate_source() { - let source = VortexSession::builder() - .with_some(Counter { count: 1 }) - .build(); - - let mut builder = source.to_builder(); - builder.get_mut::().count = 2; - let modified = builder.build(); - - assert_eq!(source.get::().count, 1); - assert_eq!(modified.get::().count, 2); - } - #[test] fn get_mut_inserts_default() { - let mut builder = VortexSession::empty().to_builder(); + let mut builder = VortexSession::builder(); assert!(builder.get_mut_opt::().is_none()); builder.get_mut::().count = 42; @@ -302,7 +236,7 @@ mod tests { let session = VortexSession::empty(); assert!(!session.allows_unknown()); - let session = session.to_builder().allow_unknown().build(); + let session = VortexSession::builder().allow_unknown().build(); assert!(session.allows_unknown()); } } diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index f933e831d38..f467018b06b 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -10,6 +10,7 @@ use std::hash::Hasher; pub use immutable::SessionBuilder; pub use immutable::VortexSession; +pub use immutable::VortexSessionBuilder; pub use immutable::VortexSessionVar; #[derive(Debug, Clone, Copy, Default)] @@ -101,7 +102,7 @@ mod tests { let session = VortexSession::empty(); assert!(!session.allows_unknown()); - let session = session.allow_unknown(); + let session = VortexSession::builder().allow_unknown().build(); assert!(session.allows_unknown()); } } diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index 17b13b1ba1e..4583f27aca5 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -13,11 +13,11 @@ use std::sync::Arc; use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin; -use vortex_array::arrow::ArrowSessionExt; -use vortex_array::dtype::session::DTypeSessionExt; -use vortex_array::scalar_fn::session::ScalarFnSessionExt; -use vortex_array::session::ArraySessionExt; -use vortex_session::VortexSession; +use vortex_array::arrow::ArrowSession; +use vortex_array::dtype::session::DTypeSession; +use vortex_array::scalar_fn::session::ScalarFnSession; +use vortex_array::session::ArraySession; +use vortex_session::VortexSessionBuilder; use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::scalar_fns::inner_product::InnerProduct; @@ -47,33 +47,39 @@ mod utils; /// deserialize. Opt-in by setting the variable to any non-empty value. pub const SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str = "VX_SCALAR_FN_ARRAY_TENSOR_PLUGIN"; -/// Initialize the Vortex tensor library with a Vortex session. -pub fn initialize(session: &VortexSession) { - session.dtypes().register(Vector); - session.dtypes().register(FixedShapeTensor); - - let arrow_session = session.arrow(); - arrow_session.register_exporter(Arc::new(Vector)); - arrow_session.register_importer(Arc::new(Vector)); +/// Initialize the Vortex tensor library with a Vortex session builder. +pub fn initialize(session: &mut VortexSessionBuilder) { + { + let dtypes = session.get_mut::(); + dtypes.register(Vector); + dtypes.register(FixedShapeTensor); + } - let session_fns = session.scalar_fns(); + { + let arrow_session = session.get_mut::(); + arrow_session.register_exporter(Arc::new(Vector)); + arrow_session.register_importer(Arc::new(Vector)); + } - session_fns.register(CosineSimilarity); - session_fns.register(InnerProduct); - session_fns.register(L2Denorm); - session_fns.register(L2Norm); + { + let scalar_fns = session.get_mut::(); + scalar_fns.register(CosineSimilarity); + scalar_fns.register(InnerProduct); + scalar_fns.register(L2Denorm); + scalar_fns.register(L2Norm); + } // Registering the scalar-fn array plugins lets the tensor scalar fns be serialized as array // encodings inside Vortex files. Gate this on an env var so applications that do not intend // to persist these encodings do not pay the registry cost or widen their stable-encoding // surface unintentionally. if std::env::var_os(SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV).is_some_and(|v| !v.is_empty()) { - let session_arrays = session.arrays(); + let arrays = session.get_mut::(); - session_arrays.register(ScalarFnArrayPlugin::new(CosineSimilarity)); - session_arrays.register(ScalarFnArrayPlugin::new(InnerProduct)); - session_arrays.register(ScalarFnArrayPlugin::new(L2Denorm)); - session_arrays.register(ScalarFnArrayPlugin::new(L2Norm)); + arrays.register(ScalarFnArrayPlugin::new(CosineSimilarity)); + arrays.register(ScalarFnArrayPlugin::new(InnerProduct)); + arrays.register(ScalarFnArrayPlugin::new(L2Denorm)); + arrays.register(ScalarFnArrayPlugin::new(L2Norm)); } } @@ -84,8 +90,8 @@ mod tests { use vortex_session::VortexSession; pub static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session + let mut builder = vortex_array::default_session_builder(); + crate::initialize(&mut builder); + builder.build() }); } diff --git a/vortex-test/compat-gen/src/adapter.rs b/vortex-test/compat-gen/src/adapter.rs index 1c8ee9bdc6a..e7048a9077e 100644 --- a/vortex-test/compat-gen/src/adapter.rs +++ b/vortex-test/compat-gen/src/adapter.rs @@ -17,7 +17,7 @@ use vortex::array::MaskFuture; use vortex::array::expr::root; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::layout::LayoutStrategy; use vortex::layout::layouts::flat::Flat; use vortex::layout::layouts::flat::writer::FlatLayoutStrategy; @@ -70,7 +70,7 @@ pub fn write_compressed( let stream = ArrayStreamAdapter::new(chunk.dtype().clone(), stream::iter([Ok(chunk)])); runtime()?.block_on(async { - let session = VortexSession::default().with_tokio(); + let session = VortexSession::default_builder().with_tokio().build(); let mut file = tokio::fs::File::create(path) .await .map_err(|e| vortex_err!("failed to create {}: {e}", path.display()))?; @@ -91,7 +91,7 @@ pub fn write_compressed_to_bytes( let stream = ArrayStreamAdapter::new(chunk.dtype().clone(), stream::iter([Ok(chunk)])); runtime()?.block_on(async { - let session = VortexSession::default().with_tokio(); + let session = VortexSession::default_builder().with_tokio().build(); let mut bytes = Vec::new(); let _summary = session .write_options() @@ -105,7 +105,7 @@ pub fn write_compressed_to_bytes( /// Read a `.vortex` file from bytes, returning the arrays. pub fn read_file(bytes: ByteBuffer) -> VortexResult { runtime()?.block_on(async { - let session = VortexSession::default().with_tokio(); + let session = VortexSession::default_builder().with_tokio().build(); let file = session.open_options().open_buffer(bytes)?; file.scan()?.into_array_stream()?.read_all().await }) @@ -120,7 +120,7 @@ pub fn read_file(bytes: ByteBuffer) -> VortexResult { /// If any segment is corrupt or any array fails to decode, this will error. pub fn read_layout_tree(bytes: ByteBuffer) -> VortexResult<()> { runtime()?.block_on(async { - let session = VortexSession::default().with_tokio(); + let session = VortexSession::default_builder().with_tokio().build(); let file = session.open_options().open_buffer(bytes)?; let root_layout = Arc::clone(file.footer().layout()); let segment_source = file.segment_source(); diff --git a/vortex-test/compat-gen/src/check.rs b/vortex-test/compat-gen/src/check.rs index ce8b29078ef..1c94e9924d0 100644 --- a/vortex-test/compat-gen/src/check.rs +++ b/vortex-test/compat-gen/src/check.rs @@ -6,8 +6,8 @@ use std::path::Path; use clap::ValueEnum; use serde::Serialize; use vortex_array::VortexSessionExecute; -use vortex_array::array_session; use vortex_array::assert_arrays_eq; +use vortex_array::default_session_builder; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -48,7 +48,7 @@ struct FailedFixture { /// Prints JSON result to stdout, human-readable progress to stderr. /// Returns error if any fixture failed or if mode constraints are violated. pub fn check(dir: &Path, mode: Mode, exclude: &[String]) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let fixtures = all_fixtures(); let fixtures: Vec<_> = fixtures .into_iter() diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs index eb453478017..60c08e490ba 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/mod.rs @@ -18,8 +18,8 @@ pub fn fixtures() -> Vec> { #[cfg(test)] mod tests { use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::assert_arrays_eq; + use vortex_array::default_session_builder; use super::fixtures; use crate::adapter; @@ -28,7 +28,7 @@ mod tests { #[test] fn roundtrip_fixtures_to_bytes() { for fixture in fixtures() { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let array = fixture.build(&mut ctx).unwrap(); check_expected_encodings(&array, fixture.as_ref()).unwrap(); let bytes = adapter::write_file_to_bytes(array.clone()).unwrap(); diff --git a/vortex-test/compat-gen/src/generate.rs b/vortex-test/compat-gen/src/generate.rs index 706c064ed70..0a52e644389 100644 --- a/vortex-test/compat-gen/src/generate.rs +++ b/vortex-test/compat-gen/src/generate.rs @@ -9,7 +9,7 @@ use sha2::Digest; use sha2::Sha256; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; -use vortex_array::array_session; +use vortex_array::default_session_builder; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -90,7 +90,7 @@ pub fn write_manifest(output_dir: &Path, infos: Vec) -> VortexResul /// Generate all fixtures into `output_dir` and write the manifest. pub fn generate(output_dir: &Path, exclude: &[String]) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let infos = write_fixtures(output_dir, exclude, &mut ctx)?; write_manifest(output_dir, infos) } diff --git a/vortex-test/e2e-cuda/src/lib.rs b/vortex-test/e2e-cuda/src/lib.rs index 493d7c7d23c..7f51024bbdb 100644 --- a/vortex-test/e2e-cuda/src/lib.rs +++ b/vortex-test/e2e-cuda/src/lib.rs @@ -69,10 +69,11 @@ use vortex_cuda::arrow::DeviceArrayStreamExt; const PRIMITIVE_DTYPE_ENV: &str = "VORTEX_CUDF_PRIMITIVE_DTYPE"; static SESSION: LazyLock = LazyLock::new(|| { - vortex::array::array_session() + vortex::array::default_session_builder() .with::() .with::() .with::() + .build() }); fn primitive_dtype_case() -> String { diff --git a/vortex-tui/src/lib.rs b/vortex-tui/src/lib.rs index 98a3343f24d..b67bdff6658 100644 --- a/vortex-tui/src/lib.rs +++ b/vortex-tui/src/lib.rs @@ -10,12 +10,12 @@ //! //! ```ignore //! use vortex::session::VortexSession; -//! use vortex::io::session::RuntimeSessionExt; +//! use vortex::io::session::RuntimeSessionBuilderExt; //! use vortex_tui::browse; //! //! #[tokio::main] //! async fn main() -> anyhow::Result<()> { -//! let session = VortexSession::default().with_tokio(); +//! let session = VortexSession::default_builder().with_tokio().build(); //! browse::exec_tui(&session, "my_file.vortex").await?; //! Ok(()) //! } diff --git a/vortex-tui/src/main.rs b/vortex-tui/src/main.rs index cc1e740eb83..03f9d923fba 100644 --- a/vortex-tui/src/main.rs +++ b/vortex-tui/src/main.rs @@ -2,13 +2,16 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex::VortexSessionDefault; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; use vortex_tui::launch; #[tokio::main] async fn main() -> anyhow::Result<()> { - let session = VortexSession::default().with_tokio().allow_unknown(); + let session = VortexSession::default_builder() + .with_tokio() + .allow_unknown() + .build(); if let Err(err) = launch(&session).await { // Defer help/version/usage errors back to clap so their formatting // and exit codes match the standalone-binary convention exactly. diff --git a/vortex-tui/src/wasm.rs b/vortex-tui/src/wasm.rs index 5fd7dd9054b..f2689fb822d 100644 --- a/vortex-tui/src/wasm.rs +++ b/vortex-tui/src/wasm.rs @@ -122,9 +122,11 @@ pub fn open_vortex_file(data: &[u8]) -> Result<(), JsValue> { use vortex::VortexSessionDefault; use vortex::buffer::ByteBuffer; use vortex::io::runtime::wasm::WasmRuntime; - use vortex::io::session::RuntimeSessionExt; + use vortex::io::session::RuntimeSessionBuilderExt; - let session = VortexSession::default().with_handle(WasmRuntime::handle()); + let session = VortexSession::default_builder() + .with_handle(WasmRuntime::handle()) + .build(); let buffer = ByteBuffer::from(data.to_vec()); let app = Rc::new(RefCell::new( AppState::from_buffer(session, buffer).map_err(|e| JsValue::from_str(&e.to_string()))?, diff --git a/vortex-web/crate/src/lib.rs b/vortex-web/crate/src/lib.rs index 03b36ee4515..4b77a114c73 100644 --- a/vortex-web/crate/src/lib.rs +++ b/vortex-web/crate/src/lib.rs @@ -11,13 +11,14 @@ use std::sync::LazyLock; use vortex::VortexSessionDefault; use vortex::io::runtime::wasm::WasmRuntime; -use vortex::io::session::RuntimeSessionExt; +use vortex::io::session::RuntimeSessionBuilderExt; use vortex::session::VortexSession; mod wasm; static SESSION: LazyLock = LazyLock::new(|| { - VortexSession::default() + VortexSession::default_builder() .with_handle(WasmRuntime::handle()) .allow_unknown() + .build() }); diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 1d27f841b5d..70dee4ec728 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -114,6 +114,7 @@ use vortex_array::stats::session::StatsSession; use vortex_io::session::RuntimeSession; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; +use vortex_session::VortexSessionBuilder; // We re-export like so in order to allow users to search inside subcrates when using the Rust docs. @@ -282,14 +283,20 @@ pub mod encodings { /// Extension trait to create a default Vortex session. pub trait VortexSessionDefault { + /// Creates a default Vortex session builder with standard arrays, layouts, scalar functions, + /// optimizer kernels, expressions, aggregate functions, and runtime support. + fn default_builder() -> VortexSessionBuilder; + /// Creates a default Vortex session with standard arrays, layouts, scalar functions, /// optimizer kernels, expressions, aggregate functions, and runtime support. - fn default() -> VortexSession; + fn default() -> VortexSession { + Self::default_builder().build() + } } impl VortexSessionDefault for VortexSession { - fn default() -> VortexSession { - let session = VortexSession::empty() + fn default_builder() -> VortexSessionBuilder { + let builder = VortexSession::builder() .with::() .with::() .with::() @@ -302,13 +309,14 @@ impl VortexSessionDefault for VortexSession { .with::(); #[cfg(feature = "files")] - let session = { - let session = session.with::(); - file::register_default_encodings(&session); - session - }; - - session + { + let mut builder = builder.with::(); + file::register_default_encodings(&mut builder); + builder + } + + #[cfg(not(feature = "files"))] + builder } } @@ -322,9 +330,9 @@ mod test { use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::array_session; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; + use vortex_array::default_session_builder; use vortex_array::dtype::FieldNames; use vortex_array::expr::gt; use vortex_array::expr::lit; @@ -472,7 +480,7 @@ mod test { assert_eq!(recovered_array.len(), array.len()); - let mut ctx = array_session().create_execution_ctx(); + let mut ctx = default_session_builder().build().create_execution_ctx(); let recovered_primitive = recovered_array.execute::(&mut ctx)?; assert!(recovered_primitive.validity()?.mask_eq(