Register compression schemes on the session

benchmarks/compress-bench/README.md

@@ -29,7 +29,7 @@ cargo run -p compress-bench --profile release_debug
 GPU dataset list in `src/main.rs`. It measures decompression only, for two backends:
 
 - **Vortex** — the file is written with CUDA-compatible BtrBlocks encodings only
-  (`only_cuda_compatible`) and a CUDA flat layout, then decoded on the device all the way to
+  (`CompressionSession::cuda()`) and a CUDA flat layout, then decoded on the device all the way to
   canonical arrays.
 - **Parquet** — the file is rewritten with GPU-friendly writer settings (see below) and read
   back with [cuDF](https://github.com/rapidsai/cudf)'s `read_parquet`, which performs the

benchmarks/compress-bench/src/gpu/vortex.rs

@@ -4,6 +4,7 @@
 use std::hint::black_box;
 use std::path::Path;
 use std::sync::Arc;
+use std::sync::LazyLock;
 use std::time::Duration;
 use std::time::Instant;
 
@@ -16,28 +17,30 @@ use async_trait::async_trait;
 use futures::Stream;
 use futures::StreamExt;
 use tempfile::NamedTempFile;
+use vortex::VortexSessionDefault;
 use vortex::array::ArrayRef;
 use vortex::array::ExecutionCtx;
 use vortex::array::IntoArray;
 use vortex::array::VortexSessionExecute;
 use vortex::array::arrays::StructArray;
 use vortex::array::arrays::struct_::StructArrayExt;
-use vortex::compressor::BtrBlocksCompressorBuilder;
+use vortex::compressor::BtrBlocksCompressor;
+use vortex::compressor::CompressionSession;
 use vortex::error::VortexResult;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::WriteOptionsSessionExt;
+use vortex::io::session::RuntimeSessionExt;
 use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy;
 use vortex::layout::layouts::compressed::CompressingStrategy;
 use vortex::layout::scan::split_by::SplitBy;
+use vortex::session::VortexSession;
 use vortex_arrow::ArrowSessionExt;
 use vortex_bench::Format;
-use vortex_bench::SESSION;
 use vortex_bench::compress::Compressed;
 use vortex_bench::compress::CompressedData;
 use vortex_bench::compress::Compressor;
 use vortex_bench::compress::Uncompressed;
 use vortex_bench::conversions::parquet_to_vortex_chunks_with_batch_size;
-use vortex_bench::retain_edition_encodings;
 use vortex_cuda::CanonicalCudaExt;
 use vortex_cuda::CudaExecutionCtx;
 use vortex_cuda::CudaOpenOptionsExt;
@@ -50,6 +53,16 @@ use vortex_cuda::layout::register_cuda_layout;
 
 use crate::gpu::writer::GPU_ROW_GROUP_SIZE;
 
+/// The session that writes and reads the CUDA-compatible files: the CUDA flat layout plus only
+/// the compression schemes the GPU decodes. Separate from vortex-bench's `SESSION`, so the host
+/// backends keep the default schemes.
+static GPU_SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = VortexSession::default().with_tokio();
+    register_cuda_layout(&session);
+    session.register(CompressionSession::cuda());
+    session
+});
+
 /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU.
 pub struct GpuVortexCompressor {
     verify: bool,
@@ -91,23 +104,17 @@ impl Compressor for GpuVortexCompressor {
     /// GPU mode never publishes this timing: `--gpu-decompress` restricts the suite to
     /// decompression, and the write runs on the host anyway.
     async fn compress(&self, input: &Uncompressed) -> Result<Compressed> {
-        register_cuda_layout(&SESSION);
-
         let array = input.vortex()?;
         let gpu_file = NamedTempFile::new()?;
         let mut output = tokio::fs::File::create(gpu_file.path()).await?;
         // Write those batches straight through as root chunks, so a chunk on disk is one
         // partition rather than whatever the default strategy would regroup them into.
         let strategy = Arc::new(ChunkedLayoutStrategy::new(CompressingStrategy::new(
             CudaFlatLayoutStrategy::default(),
-            retain_edition_encodings(
-                &SESSION,
-                BtrBlocksCompressorBuilder::default().only_cuda_compatible(),
-            )
-            .build(),
+            BtrBlocksCompressor::from_session(&GPU_SESSION),
         )));
         let start = Instant::now();
-        SESSION
+        GPU_SESSION
             .write_options()
             .with_strategy(strategy)
             .write(&mut output, array.to_array_stream())
@@ -135,7 +142,7 @@ impl Compressor for GpuVortexCompressor {
             verify_against_host_scan(gpu_file.path(), self.direct_io).await?;
         }
 
-        let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?;
+        let mut cuda_ctx = CudaSession::create_execution_ctx(&GPU_SESSION)?;
         let start = Instant::now();
         let file = open_gpu(gpu_file.path(), self.direct_io).await?;
         // Split reads on the same boundary the file was written with, so a scan batch is one
@@ -170,7 +177,7 @@ impl Compressor for GpuVortexCompressor {
 /// for `--gpu-direct-io` where it cannot be honoured is an error rather than a silent no-op,
 /// because the flag changes what the resulting number means.
 async fn open_gpu(path: &Path, direct_io: bool) -> Result<vortex::file::VortexFile> {
-    let open_options = SESSION.open_options().with_cuda();
+    let open_options = GPU_SESSION.open_options().with_cuda();
 
     #[cfg(target_os = "linux")]
     let open_options = if direct_io {
@@ -197,11 +204,11 @@ async fn open_gpu(path: &Path, direct_io: bool) -> Result<vortex::file::VortexFi
 /// This times nothing and reports no measurement: the caller runs its own timed scan afterwards,
 /// so a verifying run publishes the same kind of number as a plain one.
 async fn verify_against_host_scan(path: &Path, direct_io: bool) -> Result<()> {
-    let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?;
+    let mut cuda_ctx = CudaSession::create_execution_ctx(&GPU_SESSION)?;
     // Everything on the reference side — the host scan and both Arrow conversions — has to run
     // through a plain host context. A CUDA context allocates its outputs in device memory, and
     // the Arrow conversion then reads those buffers on the host.
-    let mut host_ctx = SESSION.create_execution_ctx();
+    let mut host_ctx = GPU_SESSION.create_execution_ctx();
 
     // The host scan reads a copy rather than the same path. The session's segment cache is
     // keyed by URI, and the CUDA reader deliberately bypasses it because its buffers are
@@ -214,7 +221,10 @@ async fn verify_against_host_scan(path: &Path, direct_io: bool) -> Result<()> {
         .scan()?
         .with_split_by(SplitBy::RowCount(GPU_ROW_GROUP_SIZE))
         .into_array_stream()?;
-    let host_file = SESSION.open_options().open_path(host_path.path()).await?;
+    let host_file = GPU_SESSION
+        .open_options()
+        .open_path(host_path.path())
+        .await?;
     let mut host_batches = host_file
         .scan()?
         .with_split_by(SplitBy::RowCount(GPU_ROW_GROUP_SIZE))
@@ -303,11 +313,11 @@ fn verify_field(
     batch_index: usize,
     field_index: usize,
 ) -> Result<()> {
-    let expected = SESSION.arrow().execute_arrow(host.clone(), None, ctx)?;
+    let expected = GPU_SESSION.arrow().execute_arrow(host.clone(), None, ctx)?;
     // Pin the Arrow target type so the two sides cannot land on different but equivalent
     // encodings of the same logical values.
     let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable());
-    let actual = SESSION.arrow().execute_arrow(gpu, Some(&target), ctx)?;
+    let actual = GPU_SESSION.arrow().execute_arrow(gpu, Some(&target), ctx)?;
 
     if expected.to_data() == actual.to_data() {
         return Ok(());

benchmarks/string-bench/src/serialized.rs

@@ -31,7 +31,7 @@ use vortex::array::IntoArray;
 use vortex::array::VortexSessionExecute;
 use vortex::array::arrays::ChunkedArray;
 use vortex::array::arrays::VarBinViewArray;
-use vortex::compressor::BtrBlocksCompressorBuilder;
+use vortex::compressor::BtrBlocksOptions;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::WriteOptionsSessionExt;
 use vortex::file::WriteStrategyBuilder;
@@ -57,7 +57,7 @@ use crate::prepare_column;
 use crate::throughput;
 use crate::verify_canonicalized;
 
-/// The btrblocks string schemes that `BtrBlocksCompressorBuilder::default()` can
+/// The btrblocks string schemes that `BtrBlocksCompressor::from_session` can
 /// choose between. Forcing one encoder excludes every entry except its own
 /// scheme, so this list must track the default scheme set: add a row whenever a
 /// new string encoder becomes selectable by default (e.g. Zstd).
@@ -170,17 +170,24 @@ impl SerializedResult {
 
 /// Build the file writer strategy that forces one selected string scheme while
 /// leaving editioned non-string child compression enabled.
-fn serialized_write_strategy(encoder: StringEncoder) -> Arc<dyn LayoutStrategy> {
+fn serialized_write_strategy(
+    session: &VortexSession,
+    encoder: StringEncoder,
+) -> Arc<dyn LayoutStrategy> {
     let forced = encoder.scheme_id();
-    let compressor = BtrBlocksCompressorBuilder::default().exclude_schemes(
-        default_string_scheme_ids()
-            .into_iter()
-            .filter(|&id| id != forced)
-            .chain([DeltaScheme::default().id()]),
-    );
-    WriteStrategyBuilder::default()
-        .with_btrblocks_builder(compressor)
-        .build()
+    let excluded: Vec<SchemeId> = default_string_scheme_ids()
+        .into_iter()
+        .filter(|&id| id != forced)
+        .chain([DeltaScheme::default().id()])
+        .collect();
+    WriteStrategyBuilder::from_session_with_options(
+        session,
+        BtrBlocksOptions {
+            exclude_schemes: excluded,
+            ..Default::default()
+        },
+    )
+    .build()
 }
 
 /// Write one canonical string column to an in-memory Vortex file, forcing the
@@ -251,7 +258,7 @@ async fn prepare_serialized_file(
     verify: bool,
     ctx: &mut ExecutionCtx,
 ) -> Result<SerializedFile> {
-    let strategy = serialized_write_strategy(encoder);
+    let strategy = serialized_write_strategy(session, encoder);
     let data = write_serialized_file(session, input, &strategy).await?;
     let file_bytes = data.len() as u64;
 
@@ -355,7 +362,7 @@ mod tests {
     use vortex::io::runtime::BlockingRuntime;
     use vortex::io::runtime::current::CurrentThreadRuntime;
     use vortex::io::session::RuntimeSessionExt;
-    use vortex_btrblocks::ALL_SCHEMES;
+    use vortex_btrblocks::CompressionSession;
     use vortex_btrblocks::SchemeExt;
 
     use super::*;
@@ -365,7 +372,9 @@ mod tests {
         // Every default scheme whose dtype gate accepts canonical Utf8 must be
         // excluded when another root string encoding is forced.
         let canonical = Canonical::VarBinView(VarBinViewArray::from_iter_str(["value"]));
-        let mut actual = ALL_SCHEMES
+        let default = CompressionSession::default();
+        let mut actual = default
+            .schemes()
             .iter()
             .filter(|scheme| scheme.matches(&canonical))
             .map(|scheme| scheme.id())

encodings/parquet-variant/src/vtable.rs

@@ -466,8 +466,11 @@ mod tests {
     }
 
     #[fixture]
-    fn write_strategy() -> Arc<dyn LayoutStrategy> {
-        vortex_file::WriteStrategyBuilder::default().build()
+    fn write_strategy(
+        parquet_variant_file_session: VortexResult<VortexSession>,
+    ) -> VortexResult<Arc<dyn LayoutStrategy>> {
+        let session = parquet_variant_file_session?;
+        Ok(vortex_file::WriteStrategyBuilder::from_session(&session).build())
     }
 
     #[test]
@@ -544,15 +547,15 @@ mod tests {
     async fn test_file_roundtrip_typed_value_variant_with_zoned_strategy(
         #[from(typed_value_variant_array)] expected: VortexResult<ArrayRef>,
         parquet_variant_file_session: VortexResult<VortexSession>,
-        write_strategy: Arc<dyn LayoutStrategy>,
+        write_strategy: VortexResult<Arc<dyn LayoutStrategy>>,
     ) -> VortexResult<()> {
         let expected = expected?;
         let parquet_variant_file_session = parquet_variant_file_session?;
 
         let mut bytes = ByteBufferMut::empty();
         parquet_variant_file_session
             .write_options()
-            .with_strategy(write_strategy)
+            .with_strategy(write_strategy?)
             .write(&mut bytes, expected.to_array_stream())
             .await?;
 

fuzz/fuzz_targets/file_io.rs

@@ -18,12 +18,11 @@ use vortex_array::dtype::StructFields;
 use vortex_array::expr::lit;
 use vortex_array::expr::root;
 use vortex_array::scalar_fn::fns::operators::Operator;
-use vortex_btrblocks::BtrBlocksCompressorBuilder;
 use vortex_error::VortexExpect;
 use vortex_error::vortex_panic;
 use vortex_file::OpenOptionsSessionExt;
 use vortex_file::WriteOptionsSessionExt;
-use vortex_file::WriteStrategyBuilder;
+use vortex_fuzz::COMPACT_SESSION;
 use vortex_fuzz::CompressorStrategy;
 use vortex_fuzz::FuzzFileAction;
 use vortex_fuzz::RUNTIME;
@@ -64,11 +63,7 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus {
 
     let write_options = match compressor_strategy {
         CompressorStrategy::Default => SESSION.write_options(),
-        CompressorStrategy::Compact => SESSION.write_options().with_strategy(
-            WriteStrategyBuilder::default()
-                .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact())
-                .build(),
-        ),
+        CompressorStrategy::Compact => COMPACT_SESSION.write_options(),
     };
 
     let mut full_buff = Vec::new();

fuzz/src/array/mod.rs

@@ -65,13 +65,13 @@ use vortex_array::search_sorted::SearchResult;
 use vortex_array::search_sorted::SearchSorted;
 use vortex_array::search_sorted::SearchSortedSide;
 use vortex_btrblocks::BtrBlocksCompressor;
-#[cfg(feature = "zstd")]
-use vortex_btrblocks::BtrBlocksCompressorBuilder;
 use vortex_error::VortexExpect;
 use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 use vortex_utils::aliases::hash_set::HashSet;
 
+#[cfg(feature = "zstd")]
+use crate::COMPACT_SESSION;
 use crate::FUZZ_ARRAY_MAX_LEN;
 use crate::SESSION;
 use crate::error::Backtrace;
@@ -249,7 +249,7 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction {
                         .into_array()
                     };
 
-                    let compressed = BtrBlocksCompressor::default()
+                    let compressed = BtrBlocksCompressor::from_session(&SESSION)
                         .compress(&indices_array, &mut ctx)
                         .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test");
                     (
@@ -561,12 +561,10 @@ pub fn compress_array(
     ctx: &mut ExecutionCtx,
 ) -> ArrayRef {
     match strategy {
-        CompressorStrategy::Default => BtrBlocksCompressor::default()
+        CompressorStrategy::Default => BtrBlocksCompressor::from_session(&SESSION)
             .compress(array, ctx)
             .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test"),
-        CompressorStrategy::Compact => BtrBlocksCompressorBuilder::default()
-            .with_compact()
-            .build()
+        CompressorStrategy::Compact => BtrBlocksCompressor::from_session(&COMPACT_SESSION)
             .compress(array, ctx)
             .vortex_expect("Compact compress should succeed in fuzz test"),
     }
@@ -579,7 +577,7 @@ pub fn compress_array(
     _strategy: CompressorStrategy,
     ctx: &mut ExecutionCtx,
 ) -> ArrayRef {
-    BtrBlocksCompressor::default()
+    BtrBlocksCompressor::from_session(&SESSION)
         .compress(array, ctx)
         .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test")
 }

fuzz/src/lib.rs

@@ -56,6 +56,8 @@ mod native_runtime {
     use std::sync::LazyLock;
 
     use vortex::VortexSessionDefault;
+    #[cfg(feature = "zstd")]
+    use vortex::compressor::CompressionSession;
     use vortex_io::runtime::BlockingRuntime;
     use vortex_io::runtime::current::CurrentThreadRuntime;
     use vortex_io::session::RuntimeSessionExt;
@@ -76,8 +78,19 @@ mod native_runtime {
         super::enable_latest_core_edition(&session);
         session
     });
+
+    /// A default session that also registers the compact (Zstd and Pco) schemes.
+    #[cfg(feature = "zstd")]
+    pub static COMPACT_SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = VortexSession::default().with_handle(RUNTIME.handle());
+        super::enable_latest_core_edition(&session);
+        session.register(CompressionSession::compact());
+        session
+    });
 }
 
+#[cfg(all(feature = "zstd", not(target_arch = "wasm32")))]
+pub use native_runtime::COMPACT_SESSION;
 #[cfg(not(target_arch = "wasm32"))]
 pub use native_runtime::RUNTIME;
 #[cfg(not(target_arch = "wasm32"))]