add execution context to compressor

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

Author: connortsui20

Cargo.lock

@@ -172,6 +172,7 @@ impl dyn DynArray + '_ {
 ///
 /// Accumulates a trace of execution steps. Individual steps are logged at TRACE level for
 /// real-time following, and the full trace is dumped at DEBUG level when the context is dropped.
+#[derive(Debug, Clone)]
 pub struct ExecutionCtx {
     id: usize,
     session: VortexSession,

vortex-array/src/executor.rs

@@ -4,9 +4,7 @@
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::IntoArray;
-use vortex_array::LEGACY_SESSION;
 use vortex_array::ToCanonical;
-use vortex_array::VortexSessionExecute;
 use vortex_array::aggregate_fn::fns::is_constant::is_constant;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::DictArray;
@@ -253,7 +251,7 @@ impl Scheme for ConstantScheme {
 
     fn expected_compression_ratio(
         &self,
-        _compressor: &CascadingCompressor,
+        compressor: &CascadingCompressor,
         data: &mut ArrayAndStats,
         ctx: CompressorContext,
         _excludes: &[SchemeId],
@@ -264,10 +262,11 @@ impl Scheme for ConstantScheme {
 
         let stats = data.string_stats();
 
-        // TODO(connor): Put the execution context somewhere!
-        let mut ctx = LEGACY_SESSION.create_execution_ctx();
         if stats.estimated_distinct_count().is_none_or(|c| c > 1)
-            || !is_constant(&stats.source().clone().into_array(), &mut ctx)?
+            || !is_constant(
+                &stats.source().clone().into_array(),
+                &mut compressor.execution_ctx(),
+            )?
         {
             return Ok(0.0);
         }

vortex-btrblocks/src/compressor/string.rs

@@ -6,9 +6,7 @@
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::IntoArray;
-use vortex_array::LEGACY_SESSION;
 use vortex_array::ToCanonical;
-use vortex_array::VortexSessionExecute;
 use vortex_array::aggregate_fn::fns::is_constant::is_constant;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::TemporalArray;
@@ -78,11 +76,11 @@ impl Scheme for TemporalScheme {
         let ext_array = array.to_extension();
         let temporal_array = TemporalArray::try_from(ext_array.clone().into_array())?;
 
-        // TODO(connor): Put the execution context somewhere!
-        let mut ctx = LEGACY_SESSION.create_execution_ctx();
-
         // Check for constant array and return early if so.
-        let is_constant = is_constant(&ext_array.clone().into_array(), &mut ctx)?;
+        let is_constant = is_constant(
+            &ext_array.clone().into_array(),
+            &mut compressor.execution_ctx(),
+        )?;
 
         if is_constant {
             return Ok(

vortex-btrblocks/src/compressor/temporal.rs

@@ -16,6 +16,7 @@ version = { workspace = true }
 [dependencies]
 itertools = { workspace = true }
 num-traits = { workspace = true }
+parking_lot = { workspace = true }
 rand = { workspace = true }
 rustc-hash = { workspace = true }
 tracing = { workspace = true }

vortex-compressor/Cargo.toml

@@ -3,10 +3,15 @@
 
 //! Cascading array compression implementation.
 
+use std::sync::Arc;
+
+use parking_lot::Mutex;
+use parking_lot::MutexGuard;
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::CanonicalValidity;
 use vortex_array::DynArray;
+use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 use vortex_array::LEGACY_SESSION;
 use vortex_array::ToCanonical;
@@ -42,16 +47,30 @@ use crate::stats::GenerateStatsOptions;
 /// 2. Pre-filtering schemes by [`Scheme::matches`] and excludes.
 /// 3. Evaluating each matching scheme's compression ratio on a sample.
 /// 4. Compressing with the best scheme and verifying the result is smaller.
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 pub struct CascadingCompressor {
     /// The enabled compression schemes.
     pub schemes: Vec<&'static dyn Scheme>,
+
+    /// Shared execution context for array operations during compression.
+    ///
+    /// This should have low contention as we only execute arrays one at a time during compression.
+    ctx: Arc<Mutex<ExecutionCtx>>,
 }
 
 impl CascadingCompressor {
     /// Creates a new compressor with the given schemes.
     pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self {
-        Self { schemes }
+        Self {
+            schemes,
+            // TODO(connor): The caller should probably pass this in.
+            ctx: Arc::new(Mutex::new(LEGACY_SESSION.create_execution_ctx())),
+        }
+    }
+
+    /// Returns a mutable borrow of the execution context.
+    pub fn execution_ctx(&self) -> MutexGuard<'_, ExecutionCtx> {
+        self.ctx.lock()
     }
 
     /// Compresses an array using cascading adaptive compression.
@@ -64,7 +83,7 @@ impl CascadingCompressor {
     pub fn compress(&self, array: &ArrayRef) -> VortexResult<ArrayRef> {
         let canonical = array
             .clone()
-            .execute::<CanonicalValidity>(&mut LEGACY_SESSION.create_execution_ctx())?
+            .execute::<CanonicalValidity>(&mut self.execution_ctx())?
             .0;
 
         // Compact it, removing any wasted space before we attempt to compress it.