Pass ExecutionCtx through the compressor (#7578)

## Summary

Tracking issue: #7216

In order to have a "less breaking change" before, I just put an
execution context in an arc'd mutex in the compressor itself. This is
really stupid, and the caller should pass an execution context whenever
they want to compress something (because we need to look at buffers to
canonicalize and do compute in the compressor).

## Testing

N/A since this is just moving where things are coming from.

---------

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

Cargo.lock

@@ -237,7 +237,7 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction {
                     };
 
                     let compressed = BtrBlocksCompressor::default()
-                        .compress(&indices_array)
+                        .compress(&indices_array, &mut ctx)
                         .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test");
                     (
                         Action::Take(compressed),
@@ -541,14 +541,15 @@ fn random_action_from_list(
 /// Compress an array using the given strategy.
 #[cfg(feature = "zstd")]
 pub fn compress_array(array: &ArrayRef, strategy: CompressorStrategy) -> ArrayRef {
+    let mut ctx = SESSION.create_execution_ctx();
     match strategy {
         CompressorStrategy::Default => BtrBlocksCompressor::default()
-            .compress(array)
+            .compress(array, &mut ctx)
             .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test"),
         CompressorStrategy::Compact => BtrBlocksCompressorBuilder::default()
             .with_compact()
             .build()
-            .compress(array)
+            .compress(array, &mut ctx)
             .vortex_expect("Compact compress should succeed in fuzz test"),
     }
 }
@@ -557,7 +558,7 @@ pub fn compress_array(array: &ArrayRef, strategy: CompressorStrategy) -> ArrayRe
 #[cfg(not(feature = "zstd"))]
 pub fn compress_array(array: &ArrayRef, _strategy: CompressorStrategy) -> ArrayRef {
     BtrBlocksCompressor::default()
-        .compress(array)
+        .compress(array, &mut SESSION.create_execution_ctx())
         .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test")
 }
 

fuzz/src/array/mod.rs

@@ -44,6 +44,7 @@ divan = { workspace = true }
 rstest = { workspace = true }
 test-with = { workspace = true }
 vortex-array = { workspace = true, features = ["_test-harness"] }
+vortex-session = { workspace = true }
 
 [features]
 # This feature enabled unstable encodings for which we don't guarantee stability.

vortex-btrblocks/Cargo.toml

@@ -5,6 +5,8 @@
 
 #[cfg(not(codspeed))]
 mod benchmarks {
+    use std::sync::LazyLock;
+
     use divan::Bencher;
     use divan::counter::BytesCount;
     use divan::counter::ItemsCount;
@@ -13,13 +15,17 @@ mod benchmarks {
     use rand::prelude::StdRng;
     use vortex_array::ArrayRef;
     use vortex_array::IntoArray;
-    use vortex_array::LEGACY_SESSION;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::session::ArraySession;
     use vortex_btrblocks::BtrBlocksCompressor;
     use vortex_buffer::buffer_mut;
+    use vortex_session::VortexSession;
     use vortex_utils::aliases::hash_set::HashSet;
 
+    static SESSION: LazyLock<VortexSession> =
+        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
+
     fn make_clickbench_window_name() -> ArrayRef {
         // A test that's meant to mirror the WindowName column from ClickBench.
         let mut values = buffer_mut![-1i32; 65_536];
@@ -41,16 +47,20 @@ mod benchmarks {
 
     #[divan::bench]
     fn btrblocks(bencher: Bencher) {
-        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+        let mut ctx = SESSION.create_execution_ctx();
         let array = make_clickbench_window_name()
             .execute::<PrimitiveArray>(&mut ctx)
             .unwrap();
         let compressor = BtrBlocksCompressor::default();
         bencher
-            .with_inputs(|| &array)
-            .input_counter(|array| ItemsCount::new(array.len()))
-            .input_counter(|array| BytesCount::of_many::<i32>(array.len()))
-            .bench_refs(|array| compressor.compress(&array.clone().into_array()).unwrap());
+            .with_inputs(|| (&array, SESSION.create_execution_ctx()))
+            .input_counter(|(array, _)| ItemsCount::new(array.len()))
+            .input_counter(|(array, _)| BytesCount::of_many::<i32>(array.len()))
+            .bench_refs(|(array, ctx)| {
+                compressor
+                    .compress(&array.clone().into_array(), ctx)
+                    .unwrap()
+            });
     }
 }
 

vortex-btrblocks/benches/compress.rs

@@ -6,6 +6,8 @@
 
 #[cfg(not(codspeed))]
 mod benchmarks {
+    use std::sync::LazyLock;
+
     use divan::Bencher;
     use divan::counter::BytesCount;
     use divan::counter::ItemsCount;
@@ -14,17 +16,23 @@ mod benchmarks {
     use rand::prelude::StdRng;
     use vortex_array::ArrayRef;
     use vortex_array::IntoArray;
+    use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::ListViewArray;
     use vortex_array::arrays::StructArray;
     use vortex_array::arrays::VarBinViewArray;
     use vortex_array::dtype::FieldNames;
+    use vortex_array::session::ArraySession;
     use vortex_array::validity::Validity;
     use vortex_btrblocks::BtrBlocksCompressor;
     use vortex_buffer::buffer_mut;
+    use vortex_session::VortexSession;
 
     const NUM_ROWS: usize = 8192;
     const SEED: u64 = 42;
 
+    static SESSION: LazyLock<VortexSession> =
+        LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
+
     const SHORT_STRINGS: &[&str] = &[
         "alpha_one",
         "bravo_two",
@@ -177,10 +185,10 @@ mod benchmarks {
         let nbytes = array.nbytes();
         let compressor = BtrBlocksCompressor::default();
         bencher
-            .with_inputs(|| &array)
+            .with_inputs(|| (&array, SESSION.create_execution_ctx()))
             .input_counter(|_| ItemsCount::new(NUM_ROWS))
             .input_counter(move |_| BytesCount::new(nbytes as usize))
-            .bench_refs(|array| compressor.compress(array).unwrap());
+            .bench_refs(|(array, ctx)| compressor.compress(array, ctx).unwrap());
     }
 }