diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index d309a2122f1..c27b6db05d3 100644 --- a/benchmarks/compress-bench/README.md +++ b/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 diff --git a/benchmarks/compress-bench/src/gpu/vortex.rs b/benchmarks/compress-bench/src/gpu/vortex.rs index bb48b461e87..5588aabb9b8 100644 --- a/benchmarks/compress-bench/src/gpu/vortex.rs +++ b/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 = 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,8 +104,6 @@ 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 { - register_cuda_layout(&SESSION); - let array = input.vortex()?; let gpu_file = NamedTempFile::new()?; let mut output = tokio::fs::File::create(gpu_file.path()).await?; @@ -100,14 +111,10 @@ impl Compressor for GpuVortexCompressor { // 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 { - 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 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(()); diff --git a/benchmarks/string-bench/src/serialized.rs b/benchmarks/string-bench/src/serialized.rs index 62de3dd9a51..d18682d8747 100644 --- a/benchmarks/string-bench/src/serialized.rs +++ b/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 { +fn serialized_write_strategy( + session: &VortexSession, + encoder: StringEncoder, +) -> Arc { 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 = 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 { - 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()) diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index f4e90d758fa..14abf9b2521 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -466,8 +466,11 @@ mod tests { } #[fixture] - fn write_strategy() -> Arc { - vortex_file::WriteStrategyBuilder::default().build() + fn write_strategy( + parquet_variant_file_session: VortexResult, + ) -> VortexResult> { + let session = parquet_variant_file_session?; + Ok(vortex_file::WriteStrategyBuilder::from_session(&session).build()) } #[test] @@ -544,7 +547,7 @@ mod tests { async fn test_file_roundtrip_typed_value_variant_with_zoned_strategy( #[from(typed_value_variant_array)] expected: VortexResult, parquet_variant_file_session: VortexResult, - write_strategy: Arc, + write_strategy: VortexResult>, ) -> VortexResult<()> { let expected = expected?; let parquet_variant_file_session = parquet_variant_file_session?; @@ -552,7 +555,7 @@ mod tests { 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?; diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index 6d9c8906fc9..f96ac2e90c8 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/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(); diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index e513c5daf81..bc2ce09f576 100644 --- a/fuzz/src/array/mod.rs +++ b/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") } diff --git a/fuzz/src/lib.rs b/fuzz/src/lib.rs index ca6cd49535e..e4025001c0e 100644 --- a/fuzz/src/lib.rs +++ b/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 = 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"))] diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 6b4ed871f2e..49ea7dd15a9 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -37,7 +37,7 @@ use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::builders::builder_with_capacity_in; use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::BtrBlocksCompressor; use vortex::dtype::DType; use vortex::dtype::FieldPath; use vortex::dtype::StructFields; @@ -66,7 +66,6 @@ use wkb::writer::write_geometry; use crate::CompactionStrategy; use crate::Format; use crate::SESSION; -use crate::retain_edition_encodings; use crate::utils::file::idempotent_async; /// Memory budget per concurrent conversion stream in GB. This is somewhat arbitary. @@ -243,27 +242,22 @@ fn write_options_for( _ => Vec::new(), }; if binary_fields.is_empty() { - return compaction.apply_options(SESSION.write_options()); + return compaction.session().write_options(); } - let mut builder = WriteStrategyBuilder::default(); - if matches!(compaction, CompactionStrategy::Compact) { - builder = builder.with_btrblocks_builder(retain_edition_encodings( - &SESSION, - BtrBlocksCompressorBuilder::default().with_compact(), - )); - } + let session = compaction.session(); + let mut builder = WriteStrategyBuilder::from_session(session); for name in binary_fields { builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); } - SESSION.write_options().with_strategy(builder.build()) + session.write_options().with_strategy(builder.build()) } /// A chunked + compressed layout that skips dictionary encoding for opaque `Binary` blobs. fn no_dict_layout() -> Arc { Arc::new(CompressingStrategy::new( ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), - retain_edition_encodings(&SESSION, BtrBlocksCompressorBuilder::default()).build(), + BtrBlocksCompressor::from_session(&SESSION), )) } @@ -343,7 +337,7 @@ pub async fn write_parquet_as_vortex( idempotent_async(vortex_path, |output_fname| async move { let mut output_file = File::create(&output_fname).await?; let data = parquet_to_vortex_chunks(parquet_path).await?; - let write_options = compaction.apply_options(SESSION.write_options()); + let write_options = compaction.session().write_options(); write_options .write(&mut output_file, data.into_array().to_array_stream()) .await?; diff --git a/vortex-bench/src/datasets/taxi_data.rs b/vortex-bench/src/datasets/taxi_data.rs index e592d84aec9..27b8096cc63 100644 --- a/vortex-bench/src/datasets/taxi_data.rs +++ b/vortex-bench/src/datasets/taxi_data.rs @@ -115,7 +115,7 @@ pub async fn taxi_data_vortex_compact() -> Result { let mut output_file = TokioFile::create(output_fname).await?; // This is the only difference to `taxi_data_vortex`. - let write_options = CompactionStrategy::Compact.apply_options(SESSION.write_options()); + let write_options = CompactionStrategy::Compact.session().write_options(); let data = parquet_to_vortex_chunks(taxi_data_parquet().await?).await?; diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 569dfc74a4d..11b13fa715c 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -28,11 +28,9 @@ use tpcds::TpcDsBenchmark; use tpch::benchmark::TpcHBenchmark; pub use utils::file::*; pub use utils::logging::*; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::CompressionSession; use vortex::error::VortexExpect; use vortex::error::vortex_err; -use vortex::file::VortexWriteOptions; -use vortex::file::WriteStrategyBuilder; use vortex::utils::aliases::hash_map::HashMap; use crate::spatialbench::SpatialBenchBenchmark; @@ -70,8 +68,6 @@ pub use datasets::BenchmarkDataset; pub use output::BenchmarkOutput; pub use output::create_output_writer; use vortex::VortexSessionDefault; -use vortex::editions::ComponentKind; -use vortex::editions::EditionSessionExt; pub use vortex::error::vortex_panic; use vortex::io::session::RuntimeSessionExt; use vortex::session::VortexSession; @@ -80,11 +76,20 @@ use vortex::session::VortexSession; #[global_allocator] static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; -pub static SESSION: LazyLock = LazyLock::new(|| { +pub static SESSION: LazyLock = LazyLock::new(new_session); + +/// [`SESSION`] plus the compact (Zstd and Pco) schemes, for [`CompactionStrategy::Compact`]. +pub static COMPACT_SESSION: LazyLock = LazyLock::new(|| { + let session = new_session(); + session.register(CompressionSession::compact()); + session +}); + +fn new_session() -> VortexSession { let session = VortexSession::default().with_tokio(); vortex_spatial::initialize(&session); session -}); +} #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)] pub struct Target { @@ -251,36 +256,15 @@ pub enum CompactionStrategy { } impl CompactionStrategy { - pub fn apply_options(&self, options: VortexWriteOptions) -> VortexWriteOptions { + /// The session to write with: [`COMPACT_SESSION`] for compact files, else [`SESSION`]. + pub fn session(&self) -> &'static VortexSession { match self { - CompactionStrategy::Compact => options.with_strategy( - WriteStrategyBuilder::default() - .with_btrblocks_builder(retain_edition_encodings( - &SESSION, - BtrBlocksCompressorBuilder::default().with_compact(), - )) - .build(), - ), - CompactionStrategy::Default => options, + CompactionStrategy::Compact => &COMPACT_SESSION, + CompactionStrategy::Default => &SESSION, } } } -/// Restrict `builder` to the encodings permitted by the session's enabled editions. -/// -/// The default writer applies this filter itself. An explicit strategy bypasses it, so a -/// benchmark that builds its own compressor applies it here to stay within editions. -pub fn retain_edition_encodings( - session: &VortexSession, - builder: BtrBlocksCompressorBuilder, -) -> BtrBlocksCompressorBuilder { - let allowed = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - builder.retain_allowed_encodings(&allowed) -} - /// Verify that local data has already been prepared for the requested benchmark formats. /// /// Engine-specific benchmark binaries call this before running queries. Data generation itself diff --git a/vortex-bench/src/tpch/tpchgen.rs b/vortex-bench/src/tpch/tpchgen.rs index 2c4a26fd527..ce7b6377e3f 100644 --- a/vortex-bench/src/tpch/tpchgen.rs +++ b/vortex-bench/src/tpch/tpchgen.rs @@ -343,7 +343,8 @@ impl VortexWriter { let mut file = TokioFile::create(&file_path).await?; compaction_strategy - .apply_options(SESSION.write_options()) + .session() + .write_options() .write(&mut file, stream) .await .map_err(|e| anyhow!("Vortex write failed: {}", e))?; diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 24a03337768..f59f05bb5ed 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -26,6 +26,7 @@ vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } +vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-fsst = { workspace = true } @@ -33,6 +34,7 @@ vortex-onpair = { workspace = true } vortex-pco = { workspace = true, optional = true } vortex-runend = { workspace = true } vortex-sequence = { workspace = true } +vortex-session = { workspace = true } vortex-sparse = { workspace = true } vortex-utils = { workspace = true } vortex-zigzag = { workspace = true } @@ -49,9 +51,7 @@ tpchgen = { workspace = true } tpchgen-arrow = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-arrow = { workspace = true } -vortex-edition = { workspace = true } vortex-mask = { workspace = true } -vortex-session = { workspace = true } [features] pco = ["dep:pco", "dep:vortex-pco"] diff --git a/vortex-btrblocks/benches/compress.rs b/vortex-btrblocks/benches/compress.rs index b088c257489..f95d37b241d 100644 --- a/vortex-btrblocks/benches/compress.rs +++ b/vortex-btrblocks/benches/compress.rs @@ -20,6 +20,7 @@ mod benchmarks { use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_btrblocks::BtrBlocksCompressor; + use vortex_btrblocks::CompressionSession; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; @@ -51,7 +52,7 @@ mod benchmarks { let array = make_clickbench_window_name() .execute::(&mut ctx) .unwrap(); - let compressor = BtrBlocksCompressor::default(); + let compressor = BtrBlocksCompressor::new(CompressionSession::default().schemes().to_vec()); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .input_counter(|(array, _)| ItemsCount::new(array.len())) diff --git a/vortex-btrblocks/benches/compress_listview.rs b/vortex-btrblocks/benches/compress_listview.rs index 881f0f0a0eb..5e12d5c2bf2 100644 --- a/vortex-btrblocks/benches/compress_listview.rs +++ b/vortex-btrblocks/benches/compress_listview.rs @@ -25,6 +25,7 @@ mod benchmarks { use vortex_array::dtype::FieldNames; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; + use vortex_btrblocks::CompressionSession; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; @@ -183,7 +184,7 @@ mod benchmarks { fn compress_listview(bencher: Bencher, layout: OffsetLayout) { let array = build_nested_listview(NUM_ROWS, layout); let nbytes = array.nbytes(); - let compressor = BtrBlocksCompressor::default(); + let compressor = BtrBlocksCompressor::new(CompressionSession::default().schemes().to_vec()); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .input_counter(|_| ItemsCount::new(NUM_ROWS)) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs deleted file mode 100644 index 3bcda909227..00000000000 --- a/vortex-btrblocks/src/builder.rs +++ /dev/null @@ -1,321 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Builder for configuring `BtrBlocksCompressor` instances. - -use vortex_array::ArrayId; -use vortex_utils::aliases::hash_set::HashSet; - -use crate::BtrBlocksCompressor; -use crate::CascadingCompressor; -use crate::Scheme; -use crate::SchemeExt; -use crate::SchemeId; -use crate::schemes::binary; -use crate::schemes::decimal; -use crate::schemes::float; -use crate::schemes::integer; -use crate::schemes::string; -use crate::schemes::temporal; - -/// All available compression schemes. -/// -/// This list is order-sensitive: the builder preserves this order when constructing -/// the final scheme list, so that tie-breaking is deterministic. -pub const ALL_SCHEMES: &[&dyn Scheme] = &[ - //////////////////////////////////////////////////////////////////////////////////////////////// - // Integer schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - // NOTE: FoR must precede BitPacking to avoid unnecessary patches. - &integer::FoRScheme, - // NOTE: ZigZag should precede BitPacking because we don't want negative numbers. - &integer::ZigZagScheme, - &integer::BitPackingScheme, - &integer::SparseScheme, - &integer::IntDictScheme, - &integer::RunEndScheme, - &integer::SequenceScheme, - &integer::IntRLEScheme, - // Delta is omitted here: see [`DELTA_SCHEME`]. - //////////////////////////////////////////////////////////////////////////////////////////////// - // Float schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &float::ALPScheme, - &float::ALPRDScheme, - &float::FloatDictScheme, - &float::NullDominatedSparseScheme, - &float::FloatRLEScheme, - //////////////////////////////////////////////////////////////////////////////////////////////// - // String schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &string::StringDictScheme, - // Both string-fragmentation schemes are registered; the sample-based - // selector keeps whichever is smaller per column. - &string::FSSTScheme, - &string::OnPairScheme, - &string::NullDominatedSparseScheme, - //////////////////////////////////////////////////////////////////////////////////////////////// - // Binary schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &binary::BinaryDictScheme, - &binary::VarBinScheme, - // Decimal schemes. - &decimal::DecimalScheme, - // Temporal schemes. - &temporal::TemporalScheme, -]; - -/// Delta, kept out of [`ALL_SCHEMES`] because it is slower to decompress than the schemes that -/// would otherwise win. Callers that want it opt in with -/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme). -/// -/// TODO(robert): Return it to [`ALL_SCHEMES`] once we have scheme filtering. -pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); - -/// Builder for creating configured [`BtrBlocksCompressor`] instances. -/// -/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated -/// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via -/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the -/// `zstd` feature is enabled. -/// -/// # Examples -/// -/// ```rust -/// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt}; -/// use vortex_btrblocks::schemes::integer::IntDictScheme; -/// -/// // Default compressor with all schemes in ALL_SCHEMES. -/// let compressor = BtrBlocksCompressorBuilder::default().build(); -/// -/// // Remove specific schemes. -/// let compressor = BtrBlocksCompressorBuilder::default() -/// .exclude_schemes([IntDictScheme.id()]) -/// .build(); -/// ``` -#[derive(Debug, Clone)] -pub struct BtrBlocksCompressorBuilder { - schemes: Vec<&'static dyn Scheme>, -} - -impl Default for BtrBlocksCompressorBuilder { - fn default() -> Self { - Self { - schemes: ALL_SCHEMES.to_vec(), - } - } -} - -impl BtrBlocksCompressorBuilder { - /// Creates a builder with no schemes registered. - /// - /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor. - pub fn empty() -> Self { - Self { - schemes: Vec::new(), - } - } - - /// Adds an external compression scheme not in [`ALL_SCHEMES`]. - /// - /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes - /// with the compressor. - /// - /// # Panics - /// - /// Panics if a scheme with the same [`SchemeId`] is already present. - pub fn with_new_scheme(mut self, scheme: &'static dyn Scheme) -> Self { - assert!( - !self.schemes.iter().any(|s| s.id() == scheme.id()), - "scheme {:?} is already present in the builder", - scheme.id(), - ); - - self.schemes.push(scheme); - self - } - - /// Adds compact encoding schemes (Zstd for strings and binary, Pco for numerics). - /// - /// This provides better compression ratios than the default, especially for floating-point - /// heavy datasets. Requires the `zstd` feature. When the `pco` feature is also enabled, - /// Pco schemes for integers and floats are included. - /// - /// # Panics - /// - /// Panics if any of the compact schemes are already present. - #[cfg(feature = "zstd")] - pub fn with_compact(self) -> Self { - let builder = self - .with_new_scheme(&string::ZstdScheme) - .with_new_scheme(&binary::ZstdScheme); - - #[cfg(feature = "pco")] - let builder = builder - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme); - - builder - } - - /// Excludes schemes without CUDA kernel support, keeps FSST for string compression, - /// and adds Zstd for binary compression. - /// - /// Both the array-level and the buffer-level Zstd schemes are added. Buffer-level - /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression, - /// but belongs to the opt-in `zstd` edition, so callers filter the two through - /// [`retain_allowed_encodings`](Self::retain_allowed_encodings). - /// - /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a - /// larger encoded representation than the default compressor. - pub fn only_cuda_compatible(self) -> Self { - // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other - // string fragmentation and dictionary schemes still require unsupported decode paths. - #[cfg_attr(not(any(feature = "pco", feature = "zstd")), allow(unused_mut))] - let mut excluded: Vec = vec![ - integer::SparseScheme.id(), - integer::IntRLEScheme.id(), - float::ALPRDScheme.id(), - float::FloatRLEScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - string::StringDictScheme.id(), - binary::BinaryDictScheme.id(), - ]; - // Delta now has a CUDA decode kernel, so arrays that reach the GPU already encoded with - // it — the Delta children OnPair emits, for instance — decode there. It stays excluded - // from this preset until GPU delta decode is benchmarked against the schemes it would - // displace, since the preset picks encodings rather than merely decoding them. - excluded.push(integer::DeltaScheme::default().id()); - #[cfg(feature = "pco")] - excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]); - let builder = self.exclude_schemes(excluded); - - #[cfg(feature = "zstd")] - let builder = builder - .with_new_scheme(&binary::ZstdScheme) - .with_new_scheme(&binary::ZstdBuffersScheme); - - builder - } - - /// Removes the specified compression schemes by their [`SchemeId`]. - pub fn exclude_schemes(mut self, ids: impl IntoIterator) -> Self { - let ids: HashSet<_> = ids.into_iter().collect(); - self.schemes.retain(|s| !ids.contains(&s.id())); - self - } - - /// Retains only schemes whose produced serialized IDs all belong to `allowed`. - /// - /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions - /// permit. - pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes - .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); - self - } - - /// Builds the configured [`BtrBlocksCompressor`]. - pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) - } -} - -#[cfg(test)] -mod tests { - use vortex_array::VTable; - use vortex_fastlanes::FoR; - - use super::*; - - #[test] - fn empty_starts_with_no_schemes() { - let builder = BtrBlocksCompressorBuilder::empty(); - assert!(builder.schemes.is_empty()); - } - - #[test] - fn default_includes_all_schemes() { - let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn retain_allowed_encodings_filters_schemes() { - let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), 1); - assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); - - let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); - assert!(none.schemes.is_empty()); - } - - #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn cuda_compatible_excludes_alprd() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - !builder - .schemes - .iter() - .any(|s| s.id() == float::ALPRDScheme.id()) - ); - } - - /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. - #[test] - fn cuda_compatible_excludes_every_sparse_scheme() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - for excluded in [ - integer::SparseScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - ] { - assert!( - !builder.schemes.iter().any(|s| s.id() == excluded), - "{excluded} should be excluded" - ); - } - } - - #[test] - fn cuda_compatible_uses_fsst_for_strings() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::FSSTScheme.id()) - ); - #[cfg(feature = "zstd")] - assert!( - !builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::ZstdScheme.id()) - ); - } - - #[test] - #[cfg(feature = "pco")] - fn cuda_compatible_excludes_pco() { - let builder = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme) - .only_cuda_compatible(); - for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { - assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); - } - } -} diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index d93be365550..526c2608100 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -5,31 +5,65 @@ use std::ops::Deref; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::session::ArraySessionExt; +use vortex_edition::ComponentKind; +use vortex_edition::EditionSessionExt; use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; -use crate::BtrBlocksCompressorBuilder; use crate::CascadingCompressor; +use crate::CompressionSessionExt; +use crate::Scheme; +use crate::SchemeExt; +use crate::SchemeId; + +/// Options for building a [`BtrBlocksCompressor`] from a session. +#[derive(Clone, Debug)] +pub struct BtrBlocksOptions { + /// Keep only the schemes whose serialized IDs the session's enabled editions permit, which a + /// file writer requires. Off, for in-memory compression where no edition applies, every + /// scheme whose encodings have a registered plugin is used. + pub enforce_editions: bool, + /// Schemes to leave out. + pub exclude_schemes: Vec, +} + +impl Default for BtrBlocksOptions { + fn default() -> Self { + Self { + enforce_editions: true, + exclude_schemes: Vec::new(), + } + } +} -/// The BtrBlocks-style compressor with all built-in schemes pre-registered. +/// The BtrBlocks-style compressor. /// -/// This is a thin wrapper around [`CascadingCompressor`] that provides a default set of -/// compression schemes via [`BtrBlocksCompressorBuilder`]. +/// This is a thin wrapper around [`CascadingCompressor`] built from the schemes registered on a +/// session. [`from_session`](Self::from_session) keeps the schemes whose serialized IDs the +/// session's enabled editions permit; [`from_session_with_options`](Self::from_session_with_options) +/// takes [`BtrBlocksOptions`] to ignore editions or leave schemes out. /// /// # Examples /// /// ```rust -/// use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt}; -/// use vortex_btrblocks::schemes::integer::IntDictScheme; +/// use vortex_btrblocks::BtrBlocksCompressor; +/// use vortex_btrblocks::BtrBlocksOptions; /// -/// // Default compressor - all schemes allowed. -/// let compressor = BtrBlocksCompressor::default(); +/// let session = vortex_array::array_session(); /// -/// // Remove specific schemes using the builder. -/// let compressor = BtrBlocksCompressorBuilder::default() -/// .exclude_schemes([IntDictScheme.id()]) -/// .build(); +/// // Every registered scheme; this session enables no editions. +/// let compressor = BtrBlocksCompressor::from_session_with_options( +/// &session, +/// &BtrBlocksOptions { +/// enforce_editions: false, +/// ..Default::default() +/// }, +/// ); /// ``` #[derive(Clone)] pub struct BtrBlocksCompressor( @@ -38,12 +72,72 @@ pub struct BtrBlocksCompressor( ); impl BtrBlocksCompressor { + /// Creates a compressor over `schemes` as given. + /// + /// Prefer [`from_session`](Self::from_session) for writing files: it only keeps schemes the + /// session can serialize and its editions permit, which this constructor does not check. + pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { + Self(CascadingCompressor::new(schemes)) + } + + /// A compressor with no schemes, which leaves every array as it is. + pub fn empty() -> Self { + Self::new(Vec::new()) + } + + /// Creates a compressor over the schemes registered on `session` whose serialized IDs the + /// session's enabled editions permit. + pub fn from_session(session: &VortexSession) -> Self { + Self::from_session_with_options(session, &BtrBlocksOptions::default()) + } + + /// Creates a compressor over the schemes registered on `session`, per `options`. + /// + /// Only schemes the session can write are kept: every encoding they produce has a registered + /// array plugin and, with `enforce_editions`, is permitted by the enabled editions. This is + /// the rule the file writer applies to the arrays it serializes. + pub fn from_session_with_options(session: &VortexSession, options: &BtrBlocksOptions) -> Self { + let mut schemes = writable_schemes(session, options.enforce_editions); + schemes.retain(|scheme| !options.exclude_schemes.contains(&scheme.id())); + Self::new(schemes) + } + /// Compresses an array using BtrBlocks-inspired compression. pub fn compress(&self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.0.compress(array, ctx) } } +/// The registered schemes whose produced encodings all have a registered array plugin and, with +/// `enforce_editions`, are permitted by the enabled editions. +fn writable_schemes(session: &VortexSession, enforce_editions: bool) -> Vec<&'static dyn Scheme> { + let registered: HashSet = session + .arrays() + .registry() + .read(|registry| registry.keys().copied().collect()); + let allowed: HashSet = if enforce_editions { + session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .filter(|id| registered.contains(id)) + .collect() + } else { + registered + }; + session + .compression() + .schemes() + .iter() + .copied() + .filter(|scheme| { + scheme + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + }) + .collect() +} + impl Deref for BtrBlocksCompressor { type Target = CascadingCompressor; @@ -52,255 +146,30 @@ impl Deref for BtrBlocksCompressor { } } -impl Default for BtrBlocksCompressor { - fn default() -> Self { - BtrBlocksCompressorBuilder::default().build() - } -} - #[cfg(test)] mod tests { - use std::sync::LazyLock; - - use rstest::rstest; - #[cfg(feature = "zstd")] - use vortex_array::ArrayId; - #[cfg(feature = "zstd")] - use vortex_array::ArrayPlugin; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::Constant; - use vortex_array::arrays::Dict; - use vortex_array::arrays::List; - use vortex_array::arrays::ListView; - use vortex_array::arrays::ListViewArray; - use vortex_array::arrays::VarBinViewArray; - use vortex_array::assert_arrays_eq; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::validity::Validity; - use vortex_buffer::BitBuffer; - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - #[cfg(feature = "zstd")] - use vortex_utils::aliases::hash_set::HashSet; + use vortex_array::array_session; - use crate::BtrBlocksCompressor; - #[cfg(feature = "zstd")] - use crate::BtrBlocksCompressorBuilder; + use super::*; + use crate::SchemeId; + use crate::schemes::integer::FoRScheme; + use crate::schemes::integer::IntDictScheme; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); - - #[rstest] - #[case::zctl( - unsafe { - ListViewArray::new_unchecked( - buffer![1i32, 2, 3, 4, 5].into_array(), - buffer![0i32, 3].into_array(), - buffer![3i32, 2].into_array(), - Validity::NonNullable, - ).with_zero_copy_to_list(true) - }, - true, - )] - #[case::overlapping( - ListViewArray::new( - buffer![1i32, 2, 3].into_array(), - buffer![0i32, 0, 0].into_array(), - buffer![3i32, 3, 3].into_array(), - Validity::NonNullable, - ), - false, - )] - fn listview_compress_roundtrip( - #[case] input: ListViewArray, - #[case] expect_list: bool, - ) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let array_ref = input.clone().into_array(); - let result = BtrBlocksCompressor::default() - .compress(&array_ref, &mut SESSION.create_execution_ctx())?; - if expect_list { - assert!(result.as_opt::().is_some()); - } else { - assert!(result.as_opt::().is_some()); - } - assert_arrays_eq!(result, input, &mut ctx); - Ok(()) + fn ids(schemes: &[&'static dyn Scheme]) -> Vec { + schemes.iter().map(|scheme| scheme.id()).collect() } + /// Without enabled editions no serialized ID is permitted, so nothing survives. #[test] - fn test_constant_all_true() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; - assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) + fn no_editions_permit_nothing() { + assert!(writable_schemes(&array_session(), true).is_empty()); } + /// A scheme whose encoding has no registered plugin can never be written, editions or not. #[test] - fn test_constant_all_false() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; - assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) - } - - #[test] - fn test_nullable_all_valid_compressed() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let array = BoolArray::new( - BitBuffer::from(vec![true; 100]), - Validity::from(BitBuffer::from(vec![true; 100])), - ); - let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; - assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) - } - - #[test] - fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0))); - let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity); - let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; - assert!(!compressed.is::()); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) - } - - #[test] - fn test_mixed_not_constant() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let array = BoolArray::new( - BitBuffer::from(vec![true, false, true, false, true]), - Validity::NonNullable, - ); - let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; - assert!(!compressed.is::()); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) - } - - #[test] - fn test_binary_constant_compressed() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values = vec![Some(b"constant-bytes".as_slice()); 100]; - let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); - let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; - assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) - } - - #[test] - fn test_binary_dict_compressed() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"]; - let values = (0..1000) - .map(|idx| Some(distinct_values[idx % distinct_values.len()])) - .collect::>(); - let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); - let btr = BtrBlocksCompressor::default(); - let compressed = btr.compress( - &array.clone().into_array(), - &mut SESSION.create_execution_ctx(), - )?; - assert!(compressed.is::()); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) - } - - #[cfg(feature = "zstd")] - #[test] - fn test_compact_binary_zstd_compressed() -> VortexResult<()> { - let values = (0..1024) - .map(|idx| { - let mut value = Vec::from(&b"common binary payload prefix "[..]); - value.extend_from_slice(&(idx as u32).to_le_bytes()); - value.extend_from_slice(&[b'x'; 96]); - value - }) - .collect::>(); - let array = VarBinViewArray::from_iter( - values.iter().map(|value| Some(value.as_slice())), - DType::Binary(Nullability::NonNullable), - ); - - let compressor = BtrBlocksCompressorBuilder::default().with_compact().build(); - let mut ctx = SESSION.create_execution_ctx(); - let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; - - assert!( - compressed.is::(), - "expected Zstd, got {}", - compressed.encoding_id() - ); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) - } - - #[cfg(feature = "zstd")] - #[rstest] - #[case::array_level(vortex_zstd::Zstd.id())] - #[case::buffer_level(vortex_zstd::ZstdBuffers.id())] - fn test_cuda_compatible_binary_zstd_follows_editions( - #[case] allowed: ArrayId, - ) -> VortexResult<()> { - let values = (0..1024) - .map(|idx| { - let mut value = Vec::from(&b"common binary payload prefix "[..]); - value.extend_from_slice(&(idx as u32).to_le_bytes()); - value.extend_from_slice(&[b'x'; 96]); - value - }) - .collect::>(); - let array = VarBinViewArray::from_iter( - values.iter().map(|value| Some(value.as_slice())), - DType::Binary(Nullability::NonNullable), - ); - - // The CUDA preset carries both Zstd schemes; the edition filter decides which one - // survives. - let compressor = BtrBlocksCompressorBuilder::default() - .only_cuda_compatible() - .retain_allowed_encodings(&HashSet::from([allowed])) - .build(); - let mut ctx = SESSION.create_execution_ctx(); - let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; - - assert_eq!(compressed.encoding_id(), allowed); - assert_arrays_eq!(compressed, array, &mut ctx); - Ok(()) + fn unregistered_encodings_are_never_writable() { + let writable = ids(&writable_schemes(&array_session(), false)); + assert!(writable.contains(&IntDictScheme.id())); + assert!(!writable.contains(&FoRScheme.id())); } } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 2e8ae484f90..533a1382481 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -30,8 +30,10 @@ //! //! Each `Scheme` implementation declares whether it [`matches`](Scheme::matches) a given //! canonical form and, if so, estimates the compression ratio (often by compressing a ~1% -//! sample). There is no dynamic registry — the set of schemes is fixed at build time via -//! [`ALL_SCHEMES`]. +//! sample). The schemes available to a compressor are those registered on its session's +//! [`CompressionSession`], which starts with the default schemes. +//! [`BtrBlocksCompressor::from_session`] keeps the registered schemes whose serialized IDs the +//! session's enabled editions permit; [`BtrBlocksOptions`] turns that off or leaves schemes out. //! //! Schemes can produce arrays that are themselves further compressed (e.g. FoR then BitPacking), //! up to [`MAX_CASCADE`] (3) layers deep. Descendant exclusion rules for of [`SchemeId`] prevents @@ -43,44 +45,49 @@ //! use vortex_array::{IntoArray, VortexSessionExecute, array_session}; //! use vortex_array::arrays::PrimitiveArray; //! use vortex_array::validity::Validity; -//! use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt}; -//! use vortex_btrblocks::schemes::integer::IntDictScheme; +//! use vortex_btrblocks::BtrBlocksCompressor; +//! use vortex_btrblocks::BtrBlocksOptions; //! use vortex_buffer::buffer; //! //! # fn example() -> vortex_error::VortexResult<()> { //! let session = array_session(); //! let array = PrimitiveArray::new(buffer![42u64; 1024], Validity::NonNullable).into_array(); //! -//! let compressor = BtrBlocksCompressor::default(); +//! // In memory, with no editions to respect, compress with every registered scheme. +//! let compressor = BtrBlocksCompressor::from_session_with_options( +//! &session, +//! &BtrBlocksOptions { +//! enforce_editions: false, +//! ..Default::default() +//! }, +//! ); //! let compressed = compressor.compress(&array, &mut session.create_execution_ctx())?; //! assert_eq!(compressed.dtype(), array.dtype()); -//! -//! // Remove specific schemes using the builder. -//! let compressor = BtrBlocksCompressorBuilder::default() -//! .exclude_schemes([IntDictScheme.id()]) -//! .build(); -//! # let _ = compressor; //! # Ok(()) //! # } //! ``` //! //! [BtrBlocks]: https://www.cs.cit.tum.de/fileadmin/w00cfj/dis/papers/btrblocks.pdf -mod builder; mod canonical_compressor; /// Compression scheme implementations. pub mod schemes; +/// Session registry of compression schemes. +pub mod session; +#[cfg(test)] +mod tests; #[cfg(test)] #[cfg(not(codspeed))] mod trace_tests; // Re-export framework types from vortex-compressor for backwards compatibility. // Btrblocks-specific exports. -pub use builder::ALL_SCHEMES; -pub use builder::BtrBlocksCompressorBuilder; -pub use builder::DELTA_SCHEME; pub use canonical_compressor::BtrBlocksCompressor; +pub use canonical_compressor::BtrBlocksOptions; pub use schemes::patches::compress_patches; +pub use session::CompressionSession; +pub use session::CompressionSessionExt; +pub use session::DELTA_SCHEME; pub use vortex_compressor::CascadingCompressor; pub use vortex_compressor::scheme::CompressorContext; pub use vortex_compressor::scheme::MAX_CASCADE; diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index c6f54a9dafe..9c560d4cadd 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -19,7 +19,7 @@ use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::BtrBlocksCompressor; +use crate::tests::default_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -27,7 +27,7 @@ static SESSION: LazyLock = LazyLock::new(vortex_array::array_sess fn test_constant_compressed() -> VortexResult<()> { let values: Vec = vec![42.5; 100]; let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -37,7 +37,7 @@ fn test_constant_compressed() -> VortexResult<()> { fn test_alp_compressed() -> VortexResult<()> { let values: Vec = (0..1000).map(|i| (i as f64) * 0.01).collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -50,7 +50,7 @@ fn test_dict_compressed() -> VortexResult<()> { .map(|i| distinct_values[i % distinct_values.len()]) .collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert!(compressed.children()[0].is::()); @@ -69,7 +69,7 @@ fn test_null_dominated_compressed() -> VortexResult<()> { } builder.append_nulls(95); let array = builder.finish_into_primitive(); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; // Verify the compressed array preserves values. assert_eq!(compressed.len(), 100); diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index bb5301b807b..d0831bbcf3a 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -15,18 +15,18 @@ use vortex_array::dtype::Nullability; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer_mut; -use vortex_compressor::CascadingCompressor; use vortex_error::VortexResult; use vortex_fastlanes::RLE; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::float::FloatRLEScheme; +use crate::tests::default_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let array = PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable).into_array(); let result = btr.compress(&array, &mut SESSION.create_execution_ctx())?; @@ -42,7 +42,7 @@ fn test_compress() -> VortexResult<()> { } let array = values.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 1024); @@ -66,7 +66,7 @@ fn test_rle_compression() -> VortexResult<()> { let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let compressor = CascadingCompressor::new(vec![&FloatRLEScheme]); + let compressor = BtrBlocksCompressor::new(vec![&FloatRLEScheme]); let compressed = compressor.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -92,7 +92,7 @@ fn test_sparse_compression() -> VortexResult<()> { array.append_nulls(90); let array = array.finish_into_primitive().into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 96); diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index b4726dab9b5..852b05cf590 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -28,15 +28,24 @@ use vortex_session::VortexSession; use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; -use crate::BtrBlocksCompressorBuilder; +use crate::CompressionSession; use crate::DELTA_SCHEME; +use crate::tests::compressor; +use crate::tests::default_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +/// The default schemes plus opt-in Delta. +fn with_delta() -> BtrBlocksCompressor { + let mut registry = CompressionSession::default(); + registry.register(&DELTA_SCHEME); + compressor(®istry) +} + #[test] fn test_constant_compressed() -> VortexResult<()> { let values: Vec = iter::repeat_n(42, 100).collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -46,7 +55,7 @@ fn test_constant_compressed() -> VortexResult<()> { fn test_for_compressed() -> VortexResult<()> { let values: Vec = (0..1000).map(|i| 1_000_000 + ((i * 37) % 100)).collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -56,7 +65,7 @@ fn test_for_compressed() -> VortexResult<()> { fn test_bitpacking_compressed() -> VortexResult<()> { let values: Vec = (0..1000).map(|i| i % 16).collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert_eq!( @@ -85,7 +94,7 @@ fn test_sparse_compressed() -> VortexResult<()> { } } let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -109,7 +118,7 @@ fn test_dict_compressed() -> VortexResult<()> { } let array = PrimitiveArray::new(Buffer::copy_from(&codes), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -122,7 +131,7 @@ fn test_runend_compressed() -> VortexResult<()> { values.extend(iter::repeat_n((i32::MAX - 50).wrapping_add(i), 10)); } let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -132,7 +141,7 @@ fn test_runend_compressed() -> VortexResult<()> { fn test_sequence_compressed() -> VortexResult<()> { let values: Vec = (0..1000).map(|i| i * 7).collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -149,7 +158,7 @@ fn test_rle_compressed() -> VortexResult<()> { values.extend(iter::repeat_n(v, 10)); } let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; eprintln!("{}", compressed.display_tree()); assert!(compressed.is::()); @@ -176,9 +185,7 @@ fn test_delta_compressed() -> VortexResult<()> { .collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build(); + let btr = with_delta(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -216,9 +223,7 @@ fn test_delta_compressed_unaligned_length() -> VortexResult<()> { .collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build(); + let btr = with_delta(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -245,9 +250,7 @@ fn test_delta_nullable_unaligned_sum() -> VortexResult<()> { let array = PrimitiveArray::from_option_iter(iter::once(None).chain((1i32..=100_000).map(Some))); - let btr = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build(); + let btr = with_delta(); let compressed = btr.compress(&array.clone().into_array(), &mut ctx)?; assert!( compressed.is::(), diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index a9ef24dc0e0..23bfd97afe9 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -19,7 +19,6 @@ use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::buffer; -use vortex_compressor::CascadingCompressor; use vortex_error::VortexResult; use vortex_fastlanes::RLE; use vortex_sequence::Sequence; @@ -27,12 +26,13 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; +use crate::tests::default_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { // Make sure empty array compression does not fail. - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let array = PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable); let result = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; @@ -60,7 +60,7 @@ fn test_dict_encodable() -> VortexResult<()> { } } - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress( &codes.freeze().into_array(), &mut SESSION.create_execution_ctx(), @@ -80,7 +80,7 @@ fn constant_mostly_nulls() -> VortexResult<()> { ); let validity = array.validity()?; - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -99,7 +99,7 @@ fn nullable_sequence() -> VortexResult<()> { let values = (0i32..20).step_by(7).collect_vec(); let array = PrimitiveArray::from_option_iter(values.clone().into_iter().map(Some)); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -118,7 +118,7 @@ fn test_rle_compression() -> VortexResult<()> { values.extend(iter::repeat_n(987i32, 150)); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let compressor = CascadingCompressor::new(vec![&IntRLEScheme]); + let compressor = BtrBlocksCompressor::new(vec![&IntRLEScheme]); let compressed = compressor.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -141,7 +141,7 @@ fn compress_large_int() -> VortexResult<()> { .collect::() .into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); btr.compress(&prim, &mut SESSION.create_execution_ctx())?; Ok(()) diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index fd3fd28696a..a42060becf0 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -35,7 +35,7 @@ use crate::SchemeExt; /// FSST (Fast Static Symbol Table) compression. /// /// One of the two string-fragmentation schemes in the default -/// [`crate::ALL_SCHEMES`] (alongside `OnPairScheme`); the sample-based selector +/// [`CompressionSession`](crate::CompressionSession) (alongside `OnPairScheme`); the sample-based selector /// keeps whichever is smaller per column. FSST compresses faster, OnPair /// usually wins on ratio. #[derive(Debug, Copy, Clone, PartialEq, Eq)] diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index aac0b4de4de..98b2eccb8f1 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -16,7 +16,9 @@ use vortex_error::VortexResult; use vortex_fsst::FSST; use vortex_session::VortexSession; -use crate::BtrBlocksCompressor; +use crate::CompressionSession; +use crate::tests::compressor; +use crate::tests::default_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -26,7 +28,7 @@ fn test_constant_compressed() -> VortexResult<()> { let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); let compressed = - BtrBlocksCompressor::default().compress(&array_ref, &mut SESSION.create_execution_ctx())?; + default_compressor().compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -41,7 +43,7 @@ fn test_dict_compressed() -> VortexResult<()> { let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); let compressed = - BtrBlocksCompressor::default().compress(&array_ref, &mut SESSION.create_execution_ctx())?; + default_compressor().compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -51,10 +53,11 @@ fn test_all_schemes_includes_onpair() { use crate::SchemeExt; use crate::schemes::string::onpair::OnPairScheme; - let ids: Vec<_> = crate::ALL_SCHEMES.iter().map(|s| s.id()).collect(); + let default = CompressionSession::default(); + let ids: Vec<_> = default.schemes().iter().map(|s| s.id()).collect(); assert!( ids.contains(&OnPairScheme.id()), - "OnPairScheme not registered in ALL_SCHEMES" + "OnPairScheme not registered in the default schemes" ); } @@ -72,7 +75,7 @@ fn test_default_btrblocks_compressor_selects_onpair() -> VortexResult<()> { let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); let compressed = - BtrBlocksCompressor::default().compress(&array_ref, &mut SESSION.create_execution_ctx())?; + default_compressor().compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), "expected OnPair, got {}", @@ -81,21 +84,21 @@ fn test_default_btrblocks_compressor_selects_onpair() -> VortexResult<()> { Ok(()) } -/// FSST is registered in the default scheme list, and an FSST-only builder +/// FSST is registered in the default scheme list, and an FSST-only compressor /// still produces an FSST array. #[test] fn test_fsst_in_default_scheme_list() -> VortexResult<()> { - use crate::BtrBlocksCompressorBuilder; use crate::SchemeExt; use crate::schemes::string::FSSTScheme; // FSST is registered by default. + let default = CompressionSession::default(); assert!( - crate::ALL_SCHEMES.iter().any(|s| s.id() == FSSTScheme.id()), - "FSSTScheme should be in ALL_SCHEMES", + default.schemes().iter().any(|s| s.id() == FSSTScheme.id()), + "FSSTScheme should be in the default schemes", ); - // An FSST-only builder still produces an FSST array for FSST-favourable + // An FSST-only compressor still produces an FSST array for FSST-favourable // input. let mut strings = Vec::with_capacity(1000); for i in 0..1000 { @@ -106,9 +109,9 @@ fn test_fsst_in_default_scheme_list() -> VortexResult<()> { let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressor = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&FSSTScheme) - .build(); + let mut registry = CompressionSession::empty(); + registry.register(&FSSTScheme); + let compressor = compressor(®istry); let compressed = compressor.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), diff --git a/vortex-btrblocks/src/schemes/string/tests.rs b/vortex-btrblocks/src/schemes/string/tests.rs index 1928f0065a2..f6e6b47f680 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -14,7 +14,7 @@ use vortex_array::dtype::Nullability; use vortex_error::VortexResult; use vortex_session::VortexSession; -use crate::BtrBlocksCompressor; +use crate::tests::default_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -30,7 +30,7 @@ fn test_strings() -> VortexResult<()> { let strings = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = strings.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 2048); @@ -57,7 +57,7 @@ fn test_sparse_nulls() -> VortexResult<()> { let strings = strings.finish_into_varbinview(); let array_ref = strings.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = default_compressor(); let compressed = btr.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 100); diff --git a/vortex-btrblocks/src/session.rs b/vortex-btrblocks/src/session.rs new file mode 100644 index 00000000000..1ac372e1d5a --- /dev/null +++ b/vortex-btrblocks/src/session.rs @@ -0,0 +1,255 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Session registry of compression schemes. +//! +//! A session's [`CompressionSession`] holds the schemes available to compressors built from it +//! with [`BtrBlocksCompressor::from_session`](crate::BtrBlocksCompressor::from_session). It +//! starts with the default schemes; [`CompressionSession::compact`] and +//! [`CompressionSession::cuda`] build the other standard registries. Whether a registered scheme +//! may write its encodings is decided by [`BtrBlocksCompressor::from_session`] from the array +//! plugins registered on the session and its enabled editions. + +use std::any::Any; + +use vortex_session::SessionExt; +use vortex_session::SessionGuard; +use vortex_session::SessionVar; + +use crate::Scheme; +use crate::SchemeExt; +use crate::schemes::binary; +use crate::schemes::decimal; +use crate::schemes::float; +use crate::schemes::integer; +use crate::schemes::string; +use crate::schemes::temporal; + +/// The default compression schemes, registered by [`CompressionSession::default`]. +/// +/// This list is order-sensitive: the compressor preserves registration order, so that +/// tie-breaking is deterministic. +const DEFAULT_SCHEMES: &[&dyn Scheme] = &[ + //////////////////////////////////////////////////////////////////////////////////////////////// + // Integer schemes. + //////////////////////////////////////////////////////////////////////////////////////////////// + // NOTE: FoR must precede BitPacking to avoid unnecessary patches. + &integer::FoRScheme, + // NOTE: ZigZag should precede BitPacking because we don't want negative numbers. + &integer::ZigZagScheme, + &integer::BitPackingScheme, + &integer::SparseScheme, + &integer::IntDictScheme, + &integer::RunEndScheme, + &integer::SequenceScheme, + &integer::IntRLEScheme, + // Delta is omitted here: see [`DELTA_SCHEME`]. + //////////////////////////////////////////////////////////////////////////////////////////////// + // Float schemes. + //////////////////////////////////////////////////////////////////////////////////////////////// + &float::ALPScheme, + &float::ALPRDScheme, + &float::FloatDictScheme, + &float::NullDominatedSparseScheme, + &float::FloatRLEScheme, + //////////////////////////////////////////////////////////////////////////////////////////////// + // String schemes. + //////////////////////////////////////////////////////////////////////////////////////////////// + &string::StringDictScheme, + // Both string-fragmentation schemes are registered; the sample-based + // selector keeps whichever is smaller per column. + &string::FSSTScheme, + &string::OnPairScheme, + &string::NullDominatedSparseScheme, + //////////////////////////////////////////////////////////////////////////////////////////////// + // Binary schemes. + //////////////////////////////////////////////////////////////////////////////////////////////// + &binary::BinaryDictScheme, + &binary::VarBinScheme, + // Decimal schemes. + &decimal::DecimalScheme, + // Temporal schemes. + &temporal::TemporalScheme, +]; + +/// The schemes [`CompressionSession::compact`] adds to the defaults: Zstd for strings and binary +/// when the `zstd` feature is on, and Pco for numerics when the `pco` feature is on. +const COMPACT_SCHEMES: &[&dyn Scheme] = &[ + #[cfg(feature = "zstd")] + &string::ZstdScheme, + #[cfg(feature = "zstd")] + &binary::ZstdScheme, + #[cfg(feature = "pco")] + &integer::PcoScheme, + #[cfg(feature = "pco")] + &float::PcoScheme, +]; + +/// Delta, kept out of the default schemes because it is slower to decompress than the schemes +/// that would otherwise win. Callers that want it register it and permit `fastlanes.delta`. +/// +/// TODO(robert): Return it to the defaults once we have scheme filtering. +pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); + +/// The compression schemes registered on a session, in registration order. +/// +/// Registration order is the compressor's tie-break order between equally good schemes, so +/// sessions that register the same schemes in the same order compress identically. +/// [`Default`] registers the default schemes, [`compact`](Self::compact) and [`cuda`](Self::cuda) +/// their variants, and [`empty`](Self::empty) none. +#[derive(Clone, Debug)] +pub struct CompressionSession { + /// Registered schemes in registration order. + schemes: Vec<&'static dyn Scheme>, +} + +impl CompressionSession { + /// A registry with no schemes. + pub fn empty() -> Self { + Self { + schemes: Vec::new(), + } + } + + /// The default schemes plus the compact ones: Zstd for strings and binary, and Pco for + /// numerics, each when its feature is on. They trade decode speed for compression ratio. + pub fn compact() -> Self { + let mut this = Self::default(); + for scheme in COMPACT_SCHEMES { + this.register(*scheme); + } + this + } + + /// The default schemes that CUDA kernels decode, keeping FSST for string compression, plus + /// Zstd for binary compression when the `zstd` feature is on. + /// + /// Both the array-level and the buffer-level Zstd schemes are added. Buffer-level compression + /// preserves binary arrays' buffer layout for zero-conversion GPU decompression, but belongs + /// to the opt-in `zstd` edition, so a session's enabled editions decide which of the two a + /// writer uses. Files written with these schemes may be larger than with the defaults: the + /// set picks encodings the GPU decodes, not the smallest ones. + pub fn cuda() -> Self { + // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other string + // fragmentation and dictionary schemes still require unsupported decode paths. Delta is + // not a default either: it has a CUDA decode kernel, but GPU delta decode has not been + // benchmarked against the schemes it would displace. + let excluded = [ + integer::SparseScheme.id(), + integer::IntRLEScheme.id(), + float::ALPRDScheme.id(), + float::FloatRLEScheme.id(), + float::NullDominatedSparseScheme.id(), + string::NullDominatedSparseScheme.id(), + string::StringDictScheme.id(), + binary::BinaryDictScheme.id(), + ]; + let mut this = Self::empty(); + for scheme in DEFAULT_SCHEMES + .iter() + .filter(|scheme| !excluded.contains(&scheme.id())) + { + this.register(*scheme); + } + #[cfg(feature = "zstd")] + { + this.register(&binary::ZstdScheme); + this.register(&binary::ZstdBuffersScheme); + } + this + } + + /// Registers a scheme. + /// + /// Registering a [`SchemeId`](crate::SchemeId) that is already present is a no-op. + pub fn register(&mut self, scheme: &'static dyn Scheme) { + if !self.schemes.iter().any(|s| s.id() == scheme.id()) { + self.schemes.push(scheme); + } + } + + /// The registered schemes in registration order. + pub fn schemes(&self) -> &[&'static dyn Scheme] { + &self.schemes + } +} + +impl Default for CompressionSession { + fn default() -> Self { + Self { + schemes: DEFAULT_SCHEMES.to_vec(), + } + } +} + +impl SessionVar for CompressionSession { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +/// Session access to the compression scheme registry. +pub trait CompressionSessionExt: SessionExt { + /// Returns the compression scheme registry. + fn compression(&self) -> SessionGuard<'_, CompressionSession> { + self.get::() + } + + /// Registers a compression scheme, see [`CompressionSession::register`]. + fn register_scheme(&self, scheme: &'static dyn Scheme) { + self.get_mut::().register(scheme); + } +} + +impl CompressionSessionExt for S {} + +#[cfg(test)] +mod tests { + use vortex_array::array_session; + + use super::*; + use crate::SchemeId; + use crate::schemes::float::FloatDictScheme; + use crate::schemes::integer::IntDictScheme; + + fn ids(schemes: &[&'static dyn Scheme]) -> Vec { + schemes.iter().map(|scheme| scheme.id()).collect() + } + + #[test] + fn default_registers_default_schemes() { + let session = array_session(); + assert_eq!(ids(session.compression().schemes()), ids(DEFAULT_SCHEMES)); + } + + #[test] + fn registration_keeps_order_and_is_idempotent() { + let session = array_session().with_some(CompressionSession::empty()); + session.register_scheme(&IntDictScheme); + session.register_scheme(&FloatDictScheme); + session.register_scheme(&IntDictScheme); + assert_eq!( + ids(session.compression().schemes()), + vec![IntDictScheme.id(), FloatDictScheme.id()] + ); + } + + #[test] + fn cuda_keeps_fsst_and_drops_string_dict() { + let cuda = ids(CompressionSession::cuda().schemes()); + assert!(cuda.contains(&string::FSSTScheme.id())); + assert!(!cuda.contains(&string::StringDictScheme.id())); + } + + #[test] + fn compact_extends_the_defaults() { + let compact = ids(CompressionSession::compact().schemes()); + assert_eq!(&compact[..DEFAULT_SCHEMES.len()], &ids(DEFAULT_SCHEMES)[..]); + #[cfg(feature = "zstd")] + assert!(compact.contains(&string::ZstdScheme.id())); + } +} diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs new file mode 100644 index 00000000000..430bc0a960c --- /dev/null +++ b/vortex-btrblocks/src/tests.rs @@ -0,0 +1,259 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compressor behaviour over the default scheme set. + +use std::sync::LazyLock; + +use rstest::rstest; +#[cfg(feature = "zstd")] +use vortex_array::ArrayId; +#[cfg(feature = "zstd")] +use vortex_array::ArrayPlugin; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::Dict; +use vortex_array::arrays::List; +use vortex_array::arrays::ListView; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::VarBinViewArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::BtrBlocksCompressor; +use crate::CompressionSession; +#[cfg(feature = "zstd")] +use crate::Scheme; +#[cfg(feature = "zstd")] +use crate::schemes::binary; + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +/// A compressor over `registry`'s schemes as given. These tests exercise scheme selection, so +/// they need neither the session's array plugins nor its editions. +pub(crate) fn compressor(registry: &CompressionSession) -> BtrBlocksCompressor { + BtrBlocksCompressor::new(registry.schemes().to_vec()) +} + +/// [`compressor`] over the default schemes. +pub(crate) fn default_compressor() -> BtrBlocksCompressor { + compressor(&CompressionSession::default()) +} + +#[rstest] +#[case::zctl( + unsafe { + ListViewArray::new_unchecked( + buffer![1i32, 2, 3, 4, 5].into_array(), + buffer![0i32, 3].into_array(), + buffer![3i32, 2].into_array(), + Validity::NonNullable, + ).with_zero_copy_to_list(true) + }, + true, +)] +#[case::overlapping( + ListViewArray::new( + buffer![1i32, 2, 3].into_array(), + buffer![0i32, 0, 0].into_array(), + buffer![3i32, 3, 3].into_array(), + Validity::NonNullable, + ), + false, +)] +fn listview_compress_roundtrip( + #[case] input: ListViewArray, + #[case] expect_list: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array_ref = input.clone().into_array(); + let result = default_compressor().compress(&array_ref, &mut SESSION.create_execution_ctx())?; + if expect_list { + assert!(result.as_opt::().is_some()); + } else { + assert!(result.as_opt::().is_some()); + } + assert_arrays_eq!(result, input, &mut ctx); + Ok(()) +} + +#[test] +fn test_constant_all_true() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = BoolArray::new(BitBuffer::from(vec![true; 100]), Validity::NonNullable); + let btr = default_compressor(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!(compressed.is::()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_constant_all_false() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = BoolArray::new(BitBuffer::from(vec![false; 100]), Validity::NonNullable); + let btr = default_compressor(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!(compressed.is::()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_nullable_all_valid_compressed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = BoolArray::new( + BitBuffer::from(vec![true; 100]), + Validity::from(BitBuffer::from(vec![true; 100])), + ); + let btr = default_compressor(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!(compressed.is::()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_nullable_with_nulls_not_compressed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let validity = Validity::from(BitBuffer::from_iter((0..100).map(|i| i % 3 != 0))); + let array = BoolArray::new(BitBuffer::from(vec![true; 100]), validity); + let btr = default_compressor(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!(!compressed.is::()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_mixed_not_constant() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = BoolArray::new( + BitBuffer::from(vec![true, false, true, false, true]), + Validity::NonNullable, + ); + let btr = default_compressor(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!(!compressed.is::()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_binary_constant_compressed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = vec![Some(b"constant-bytes".as_slice()); 100]; + let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); + let btr = default_compressor(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!(compressed.is::()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[test] +fn test_binary_dict_compressed() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let distinct_values: [&[u8]; 3] = [b"alpha", b"beta", b"gamma"]; + let values = (0..1000) + .map(|idx| Some(distinct_values[idx % distinct_values.len()])) + .collect::>(); + let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); + let btr = default_compressor(); + let compressed = btr.compress( + &array.clone().into_array(), + &mut SESSION.create_execution_ctx(), + )?; + assert!(compressed.is::()); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +#[cfg(feature = "zstd")] +#[test] +fn test_compact_binary_zstd_compressed() -> VortexResult<()> { + let values = (0..1024) + .map(|idx| { + let mut value = Vec::from(&b"common binary payload prefix "[..]); + value.extend_from_slice(&(idx as u32).to_le_bytes()); + value.extend_from_slice(&[b'x'; 96]); + value + }) + .collect::>(); + let array = VarBinViewArray::from_iter( + values.iter().map(|value| Some(value.as_slice())), + DType::Binary(Nullability::NonNullable), + ); + + let compressor = compressor(&CompressionSession::compact()); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; + + assert!( + compressed.is::(), + "expected Zstd, got {}", + compressed.encoding_id() + ); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} + +/// Each binary Zstd scheme writes exactly its own encoding, so a compressor over just that +/// scheme emits it. +#[cfg(feature = "zstd")] +#[rstest] +#[case::array_level(&binary::ZstdScheme, vortex_zstd::Zstd.id())] +#[case::buffer_level(&binary::ZstdBuffersScheme, vortex_zstd::ZstdBuffers.id())] +fn test_binary_zstd_scheme_encoding( + #[case] scheme: &'static dyn Scheme, + #[case] expected: ArrayId, +) -> VortexResult<()> { + let values = (0..1024) + .map(|idx| { + let mut value = Vec::from(&b"common binary payload prefix "[..]); + value.extend_from_slice(&(idx as u32).to_le_bytes()); + value.extend_from_slice(&[b'x'; 96]); + value + }) + .collect::>(); + let array = VarBinViewArray::from_iter( + values.iter().map(|value| Some(value.as_slice())), + DType::Binary(Nullability::NonNullable), + ); + + let mut registry = CompressionSession::empty(); + registry.register(scheme); + let compressor = compressor(®istry); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; + + assert_eq!(compressed.encoding_id(), expected); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) +} diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index e23e4ef0244..2dc52e35cdb 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -53,8 +53,9 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use crate::BtrBlocksCompressorBuilder; +use crate::CompressionSession; use crate::DELTA_SCHEME; +use crate::tests::compressor; /// A session with the default Vortex encodings registered. /// @@ -127,10 +128,9 @@ fn lineitem() -> VortexResult { /// Delta is opt-in, and these traces cover the delta-encoded FSST offsets, so enable it here. fn compressed_lineitem() -> VortexResult { - BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build() - .compress(&lineitem()?, &mut execution_ctx()) + let mut registry = CompressionSession::default(); + registry.register(&DELTA_SCHEME); + compressor(®istry).compress(&lineitem()?, &mut execution_ctx()) } fn field(array: &ArrayRef, name: &str) -> VortexResult { diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index fc636252c27..5deab2f03b9 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -14,8 +14,8 @@ //! - `regular`: the schemes permitted by the default `core` edition, minus OnPair. //! - `onpair`: the structured-string entry with OnPair enabled — pins OnPair selection. //! - `compact`: the schemes permitted by the default `core` and opt-in `zstd` editions, with -//! the `zstd` + `pco` features and -//! [`BtrBlocksCompressorBuilder::with_compact`] — pins Zstd / Pco selection. +//! the `zstd` + `pco` features and [`CompressionSession::compact`] +//! — pins Zstd / Pco selection. //! //! Every corpus entry is longer than 1024 values so the sampling-based estimation path is //! exercised, and each entry is compressed twice per run to assert determinism directly. @@ -35,10 +35,12 @@ use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::ListArray; +use vortex_array::arrays::Patched; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::display::EncodingSummaryExtractor; use vortex_array::display::MetadataExtractor; use vortex_array::display::TreeContext; @@ -47,11 +49,14 @@ use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; +use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::BtrBlocksOptions; +use vortex_btrblocks::CompressionSession; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::string::OnPairScheme; use vortex_buffer::Buffer; -use vortex_edition::ComponentKind; use vortex_edition::EDITION_DECLARATIONS; use vortex_edition::EDITION_FAMILIES; use vortex_edition::EditionId; @@ -393,15 +398,39 @@ fn list_of_int_runs() -> VortexResult { /// Excludes OnPair from the `regular` and `compact` variants: it beats FSST on /// `string_fsst_structured`, and those variants pin the FSST selection. OnPair's own decisions /// are pinned by [`golden_onpair`]. -fn without_onpair(builder: BtrBlocksCompressorBuilder) -> BtrBlocksCompressorBuilder { - use vortex_btrblocks::SchemeExt; - use vortex_btrblocks::schemes::string::OnPairScheme; - - builder.exclude_schemes([OnPairScheme.id()]) +fn without_onpair() -> BtrBlocksOptions { + BtrBlocksOptions { + exclude_schemes: vec![OnPairScheme.id()], + ..Default::default() + } } -fn edition_session(editions: &[EditionId]) -> VortexResult { - let session = vortex_array::array_session().with::(); +/// A session with the schemes in `registry` and `editions` enabled. +fn edition_session( + editions: &[EditionId], + registry: CompressionSession, +) -> VortexResult { + let session = vortex_array::array_session() + .with_some(registry) + .with::(); + // The compressor only uses schemes whose encodings this session can serialize. + vortex_alp::initialize(&session); + vortex_datetime_parts::initialize(&session); + vortex_decimal_byte_parts::initialize(&session); + vortex_fastlanes::initialize(&session); + vortex_fsst::initialize(&session); + vortex_onpair::initialize(&session); + vortex_runend::initialize(&session); + vortex_sequence::initialize(&session); + vortex_sparse::initialize(&session); + vortex_zigzag::initialize(&session); + #[cfg(feature = "zstd")] + vortex_zstd::initialize(&session); + #[cfg(feature = "pco")] + session.arrays().register(vortex_pco::Pco); + if use_experimental_patches() { + session.arrays().register(Patched); + } for family in EDITION_FAMILIES { session.editions().declare_family(family)?; } @@ -414,43 +443,18 @@ fn edition_session(editions: &[EditionId]) -> VortexResult { Ok(session) } -fn compressor_for_session( - session: &VortexSession, - builder: BtrBlocksCompressorBuilder, -) -> BtrBlocksCompressor { - let allowed = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - without_onpair(builder) - .retain_allowed_encodings(&allowed) - .build() -} - -/// Like [`compressor_for_session`] but keeps OnPair in the scheme pool. -fn compressor_with_onpair( - session: &VortexSession, - builder: BtrBlocksCompressorBuilder, -) -> BtrBlocksCompressor { - let allowed = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - builder.retain_allowed_encodings(&allowed).build() -} - #[test] fn golden_regular() -> VortexResult<()> { - let session = edition_session(&[CORE_2026_08_3])?; - let compressor = compressor_for_session(&session, BtrBlocksCompressorBuilder::default()); + let session = edition_session(&[CORE_2026_08_3], CompressionSession::default())?; + let compressor = BtrBlocksCompressor::from_session_with_options(&session, &without_onpair()); golden_corpus_snapshots("regular", &compressor) } /// Pins OnPair's selection over FSST on the structured-string entry. #[test] fn golden_onpair() -> VortexResult<()> { - let session = edition_session(&[CORE_2026_08_3])?; - let compressor = compressor_with_onpair(&session, BtrBlocksCompressorBuilder::default()); + let session = edition_session(&[CORE_2026_08_3], CompressionSession::default())?; + let compressor = BtrBlocksCompressor::from_session(&session); golden_snapshots( "onpair", &compressor, @@ -461,12 +465,8 @@ fn golden_onpair() -> VortexResult<()> { #[cfg(all(feature = "zstd", feature = "pco"))] #[test] fn golden_compact() -> VortexResult<()> { - let session = edition_session(&[CORE_2026_08_3])?; - vortex_zstd::initialize(&session); + let session = edition_session(&[CORE_2026_08_3], CompressionSession::compact())?; session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; - let compressor = compressor_for_session( - &session, - BtrBlocksCompressorBuilder::default().with_compact(), - ); + let compressor = BtrBlocksCompressor::from_session_with_options(&session, &without_onpair()); golden_corpus_snapshots("compact", &compressor) } diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index 31734d6a60e..2600ce1bdc3 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -19,10 +19,16 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::CompressionSession; use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +/// A compressor over the default schemes as given, independent of the session's plugins. +fn default_compressor() -> BtrBlocksCompressor { + BtrBlocksCompressor::new(CompressionSession::default().schemes().to_vec()) +} + /// Helper: synthetic short-string corpus that the cascading compressor should /// route through OnPair. fn corpus(n: usize) -> Vec { @@ -58,7 +64,7 @@ fn nonnullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = default_compressor() .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert a specific scheme — both OnPair and FSST are registered and @@ -101,7 +107,7 @@ fn nullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = default_compressor() .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert OnPair specifically here — the sample-based selector may @@ -137,7 +143,7 @@ fn large_unique_short_strings_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = default_compressor() .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); @@ -166,7 +172,7 @@ fn empty_and_short_string_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = default_compressor() .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); let decoded = compressed @@ -211,7 +217,7 @@ fn delta_dict_offsets_roundtrip() { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = default_compressor() .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); let decoded = compressed diff --git a/vortex-btrblocks/tests/varbin_scheme.rs b/vortex-btrblocks/tests/varbin_scheme.rs index d47c8280af1..33a03fdca5c 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -14,8 +14,10 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::CompressionSession; use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::SchemeId; use vortex_btrblocks::schemes::binary::VarBinScheme; use vortex_btrblocks::schemes::string::OnPairScheme; use vortex_error::VortexResult; @@ -23,8 +25,25 @@ use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +/// A compressor over the default schemes as given, independent of the session's plugins. +fn default_compressor() -> BtrBlocksCompressor { + BtrBlocksCompressor::new(CompressionSession::default().schemes().to_vec()) +} + const N: usize = 100_000; +/// The default schemes minus `excluded`. +fn default_without(excluded: SchemeId) -> BtrBlocksCompressor { + BtrBlocksCompressor::new( + CompressionSession::default() + .schemes() + .iter() + .copied() + .filter(|scheme| scheme.id() != excluded) + .collect(), + ) +} + fn lcg(state: &mut u64) -> u64 { *state = state .wrapping_mul(6364136223846793005) @@ -70,10 +89,8 @@ fn cases() -> Vec<(&'static str, ArrayRef)> { #[test] fn varbin_scheme_shrinks_binary() -> VortexResult<()> { - let with = BtrBlocksCompressorBuilder::default().build(); - let without = BtrBlocksCompressorBuilder::default() - .exclude_schemes([VarBinScheme.id()]) - .build(); + let with = default_compressor(); + let without = default_without(VarBinScheme.id()); println!( "{:<20}{:>12}{:>14}{:>14}{:>9}", @@ -116,9 +133,7 @@ fn varbin_scheme_shrinks_binary() -> VortexResult<()> { /// change the result. `OnPairScheme` only matches utf8 and would otherwise win the utf8 column. #[test] fn fsst_versus_varbin_on_identical_bytes() -> VortexResult<()> { - let builder = BtrBlocksCompressorBuilder::default(); - let builder = builder.exclude_schemes([OnPairScheme.id()]); - let compressor = builder.build(); + let compressor = default_without(OnPairScheme.id()); let mut seed = 99u64; let shared_prefix: Vec = (0..N).map(|i| format!("PREFIX_{i:09}")).collect(); diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 2fb8ae8054e..4601478b70c 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -101,7 +101,8 @@ vx_session *vx_cuda_session_new(vx_error **error_out); * Open a Vortex file sink configured to produce CUDA-readable files. * * Push host arrays and close/abort with `vx_array_sink_*`. Only on-disk encodings and layouts - * change; writing does not move arrays to the GPU. + * change; writing does not move arrays to the GPU. The sink compresses with the session's + * schemes, so pass a session from `vx_cuda_session_new` to use only those the GPU decodes. * * # Safety * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 0647b2936a7..b99ed624d60 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::array::ArrayRef; use vortex::array::stream::ArrayStreamExt; +use vortex::compressor::CompressionSession; use vortex::dtype::FieldName; use vortex::dtype::FieldNames; use vortex::error::VortexResult; @@ -100,6 +101,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_session_new( Ok(vx_session_new_with(|session| { let session = session.with_some(cuda_session); register_cuda_layout(&session); + session.register(CompressionSession::cuda()); session })) }) @@ -108,7 +110,8 @@ pub unsafe extern "C-unwind" fn vx_cuda_session_new( /// Open a Vortex file sink configured to produce CUDA-readable files. /// /// Push host arrays and close/abort with `vx_array_sink_*`. Only on-disk encodings and layouts -/// change; writing does not move arrays to the GPU. +/// change; writing does not move arrays to the GPU. The sink compresses with the session's +/// schemes, so pass a session from `vx_cuda_session_new` to use only those the GPU decodes. /// /// # Safety /// @@ -1033,6 +1036,7 @@ mod tests { fn test_projection_gpu_values_and_validity() -> VortexResult<()> { let session = session().with_some(CudaSession::try_default()?); register_cuda_layout(&session); + session.register(CompressionSession::cuda()); let input = table()?; let columns = ["値.x", "ids"]; let expected = input.project(names(&columns)?.as_ref())?.into_array(); diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 74ff881a0ea..d21395243cf 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -23,6 +23,7 @@ use vortex::array::arrays::Dict; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; use vortex::buffer::ByteBufferMut; +use vortex::compressor::CompressionSession; use vortex::error::VortexResult; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; @@ -93,6 +94,7 @@ async fn main() -> VortexResult<()> { async fn cmd_convert(input: PathBuf, output: PathBuf) -> VortexResult<()> { let session = VortexSession::default(); register_cuda_layout(&session); + session.register(CompressionSession::cuda()); let input_file = session.open_options().open_path(&input).await?; let scan = input_file.scan()?.into_array_stream()?; @@ -139,6 +141,7 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes let session = VortexSession::default(); register_cuda_layout(&session); + session.register(CompressionSession::cuda()); let mut cuda_ctx = CudaSession::create_execution_ctx(&session)? .with_launch_strategy(Arc::new(TracingLaunchStrategy)); diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 7cc7d0322da..2435708b684 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -29,10 +29,9 @@ use vortex::array::serde::SerializedArray; use vortex::array::stats::StatsSetRef; use vortex::buffer::BufferString; use vortex::buffer::ByteBuffer; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::BtrBlocksCompressor; use vortex::dtype::DType; use vortex::dtype::FieldMask; -use vortex::editions::ComponentKind; use vortex::editions::Edition; use vortex::editions::EditionDeclaration; use vortex::editions::EditionFamily; @@ -547,28 +546,26 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { result } -/// Build a CUDA-flat writer using only CUDA-compatible, session-enabled array encodings. +/// Build a CUDA-flat writer from the schemes registered on `session` that its editions permit. /// -/// Requires [`register_cuda_layout`]. Zero `block_rows` uses default sizing and dictionary policy; -/// nonzero sets row blocks without outer dictionaries or byte coalescing, retaining per-block -/// dictionary compression. +/// `session` must have [`register_cuda_layout`] applied and a [`CompressionSession::cuda`] +/// registry: the writer uses whatever schemes are registered, so with any other registry the file +/// may contain encodings the GPU does not decode. +/// +/// Zero `block_rows` uses default sizing and dictionary policy; nonzero sets row blocks without +/// outer dictionaries or byte coalescing, retaining per-block dictionary compression. +/// +/// [`CompressionSession::cuda`]: vortex::compressor::CompressionSession::cuda pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc { - let allowed_encodings = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - let builder = BtrBlocksCompressorBuilder::default() - .only_cuda_compatible() - .retain_allowed_encodings(&allowed_encodings); - let strategy = WriteStrategyBuilder::default() + let strategy = WriteStrategyBuilder::from_session(session) .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); if block_rows == 0 { - strategy.with_btrblocks_builder(builder).build() + strategy.build() } else { // An opaque compressor keeps IntDict; disabling the probe avoids u16-sized outer blocks. strategy - .with_compressor(builder.build()) - .with_probe_compressor(BtrBlocksCompressorBuilder::empty().build()) + .with_compressor(BtrBlocksCompressor::from_session(session)) + .with_probe_compressor(BtrBlocksCompressor::empty()) .with_row_block_size(block_rows) .with_data_block_target_bytes(None) .build() @@ -658,7 +655,9 @@ mod tests { use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::assert_arrays_eq; use vortex::buffer::ByteBufferMut; + use vortex::compressor::CompressionSession; use vortex::editions::CORE_2025_05_0; + use vortex::editions::ComponentKind; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; use vortex::file::WriteOptionsSessionExt; @@ -714,6 +713,7 @@ mod tests { let runtime = CurrentThreadRuntime::new(); let session = VortexSession::default().with_handle(runtime.handle()); register_cuda_layout(&session); + session.register(CompressionSession::cuda()); runtime.block_on(async { let input = repeated_ids(8, 2 * block_rows + 137)?; let file = write_file(&session, input.clone(), block_rows).await?; @@ -748,6 +748,7 @@ mod tests { let runtime = CurrentThreadRuntime::new(); let session = VortexSession::default().with_handle(runtime.handle()); register_cuda_layout(&session); + session.register(CompressionSession::cuda()); runtime.block_on(async { // Exceed u16 cardinality while remaining eligible for outer dictionaries. let block_rows = 70_000 * 8; diff --git a/vortex-ffi/src/sink.rs b/vortex-ffi/src/sink.rs index 12d5bfb1909..2ac790bed85 100644 --- a/vortex-ffi/src/sink.rs +++ b/vortex-ffi/src/sink.rs @@ -106,7 +106,7 @@ pub unsafe extern "C-unwind" fn vx_array_sink_open_file( error_out: *mut *mut vx_error, ) -> *mut vx_array_sink { try_or_default(error_out, || { - let strategy = WriteStrategyBuilder::default().build(); + let strategy = WriteStrategyBuilder::from_session(vx_session::as_ref(session)).build(); unsafe { vx_array_sink_open_file_with_strategy(session, path, dtype, strategy) } }) } diff --git a/vortex-file/benches/split_collection.rs b/vortex-file/benches/split_collection.rs index 36c8e6b9391..3d2c9c34c82 100644 --- a/vortex-file/benches/split_collection.rs +++ b/vortex-file/benches/split_collection.rs @@ -22,6 +22,7 @@ use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::Field; use vortex_array::dtype::FieldMask; +use vortex_btrblocks::BtrBlocksOptions; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; use vortex_file::OpenOptionsSessionExt; @@ -67,6 +68,14 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +/// Options ignoring editions: this session enables none. +fn no_editions() -> BtrBlocksOptions { + BtrBlocksOptions { + enforce_editions: false, + ..Default::default() + } +} + fn make_file(columns: usize, chunks: usize) -> VortexFile { let field_names = (0..columns).map(|c| format!("col_{c}")).collect::>(); let struct_chunks = (0..chunks) @@ -85,10 +94,11 @@ fn make_file(columns: usize, chunks: usize) -> VortexFile { .collect::>(); let array = ChunkedArray::from_iter(struct_chunks).into_array(); - let strategy = vortex_file::WriteStrategyBuilder::default() - .with_row_block_size(ROWS_PER_CHUNK) - .with_data_block_target_bytes(None) - .build(); + let strategy = + vortex_file::WriteStrategyBuilder::from_session_with_options(&SESSION, no_editions()) + .with_row_block_size(ROWS_PER_CHUNK) + .with_data_block_target_bytes(None) + .build(); let mut buf = ByteBufferMut::empty(); RUNTIME @@ -143,7 +153,8 @@ fn make_misaligned_file(columns: usize, chunks: usize) -> VortexFile { .unwrap() .into_array(); - let mut strategy = vortex_file::WriteStrategyBuilder::default(); + let mut strategy = + vortex_file::WriteStrategyBuilder::from_session_with_options(&SESSION, no_editions()); for (c, (name, _)) in fields.iter().enumerate() { let field_strategy = RepartitionStrategy::new( ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index c8110fe88c3..0110c082106 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -7,7 +7,8 @@ use std::num::NonZeroUsize; use std::sync::Arc; use vortex_array::dtype::FieldPath; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::BtrBlocksOptions; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::integer::IntDictScheme; use vortex_error::VortexExpect; @@ -26,16 +27,20 @@ use vortex_layout::layouts::table::TableStrategy; use vortex_layout::layouts::table::use_experimental_list_layout; use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; use vortex_layout::layouts::zoned::writer::ZonedStrategy; +use vortex_session::VortexSession; use vortex_utils::aliases::hash_map::HashMap; const ONE_MEG: u64 = 1 << 20; /// How the compressor was configured on [`WriteStrategyBuilder`]. enum CompressorConfig { - /// A [`BtrBlocksCompressorBuilder`] that [`WriteStrategyBuilder::build`] will finalize. - /// `IntDictScheme` is automatically excluded from the data compressor to prevent recursive - /// dictionary encoding. - BtrBlocks(BtrBlocksCompressorBuilder), + /// [`BtrBlocksCompressor`]s that [`WriteStrategyBuilder::build`] creates from the session: the + /// data compressor without `IntDictScheme`, to prevent recursive dictionary encoding, and the + /// stats compressor with every scheme `options` allows. + BtrBlocks { + session: VortexSession, + options: BtrBlocksOptions, + }, /// An opaque compressor used as-is for both data and stats compression. Opaque(Arc), } @@ -64,12 +69,16 @@ pub struct WriteStrategyBuilder { use_list_layout: bool, } -impl Default for WriteStrategyBuilder { - /// Create a new empty builder. It can be further configured, - /// and then finally built yielding the [`LayoutStrategy`]. - fn default() -> Self { +impl WriteStrategyBuilder { + /// Create a new builder whose compressor uses the schemes registered on `session` that its + /// enabled editions permit. It can be further configured, and then finally built yielding the + /// [`LayoutStrategy`]. + pub fn from_session(session: &VortexSession) -> Self { Self { - compressor: CompressorConfig::BtrBlocks(BtrBlocksCompressorBuilder::default()), + compressor: CompressorConfig::BtrBlocks { + session: session.clone(), + options: BtrBlocksOptions::default(), + }, row_block_size: 8192, data_block_target_bytes: Some(ONE_MEG), field_writers: HashMap::new(), @@ -78,6 +87,17 @@ impl Default for WriteStrategyBuilder { use_list_layout: use_experimental_list_layout(), } } + + /// Create a new builder whose compressor is built from `session` per `options`. + pub fn from_session_with_options(session: &VortexSession, options: BtrBlocksOptions) -> Self { + Self { + compressor: CompressorConfig::BtrBlocks { + session: session.clone(), + options, + }, + ..Self::from_session(session) + } + } } impl WriteStrategyBuilder { @@ -132,15 +152,6 @@ impl WriteStrategyBuilder { self } - /// Override the default [`BtrBlocksCompressorBuilder`] used for compression. - /// - /// The builder produces two compressors: one for data and one for stats. - /// An explicitly built compressor is used as configured. - pub fn with_btrblocks_builder(mut self, builder: BtrBlocksCompressorBuilder) -> Self { - self.compressor = CompressorConfig::BtrBlocks(builder); - self - } - /// Set the compressor to an opaque [`CompressorPlugin`]. /// /// The compressor is used as-is for both data and stats compression. Use this when the @@ -165,7 +176,29 @@ impl WriteStrategyBuilder { Arc::new(FlatLayoutStrategy::default()) }; - let compressor = self.compressor; + // The data compressor (step 5) excludes IntDictScheme because DictStrategy (step 3) + // already dictionary-encodes columns; allowing it there would redundantly + // dictionary-encode the integer codes produced by that earlier step. Stats tables and + // dict values (steps 2.1 and 3.1) use every scheme. + let (data_compressor, stats_compressor): ( + Arc, + Arc, + ) = match self.compressor { + CompressorConfig::BtrBlocks { session, options } => { + let mut data_options = options.clone(); + data_options.exclude_schemes.push(IntDictScheme.id()); + ( + Arc::new(BtrBlocksCompressor::from_session_with_options( + &session, + &data_options, + )), + Arc::new(BtrBlocksCompressor::from_session_with_options( + &session, &options, + )), + ) + } + CompressorConfig::Opaque(compressor) => (Arc::clone(&compressor), compressor), + }; // 7. for each chunk create a flat layout let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); @@ -173,18 +206,6 @@ impl WriteStrategyBuilder { let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB // 5. compress each chunk. - // Exclude IntDictScheme from the data compressor because DictStrategy (step 3) already - // dictionary-encodes columns. Allowing IntDictScheme here would redundantly - // dictionary-encode the integer codes produced by that earlier step. - let data_compressor: Arc = match &compressor { - CompressorConfig::BtrBlocks(builder) => Arc::new( - builder - .clone() - .exclude_schemes([IntDictScheme.id()]) - .build(), - ), - CompressorConfig::Opaque(compressor) => Arc::clone(compressor), - }; let compressing = CompressingStrategy::new(buffered, data_compressor); // 4. prior to compression, coalesce up to a minimum size @@ -205,10 +226,6 @@ impl WriteStrategyBuilder { ); // 2.1. | 3.1. compress stats tables and dict values. - let stats_compressor: Arc = match compressor { - CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), - CompressorConfig::Opaque(compressor) => compressor, - }; let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); // 3. apply dict encoding or fallback diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 640de874d2b..f18c089db25 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -65,14 +65,19 @@ use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::BtrBlocksOptions; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::string::StringDictScheme; use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; +use vortex_edition::EDITION_DECLARATIONS; +use vortex_edition::EDITION_FAMILIES; use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::session::RuntimeSession; @@ -1693,6 +1698,40 @@ async fn test_buffered_bytes_are_writer_scoped() -> VortexResult<()> { Ok(()) } +/// The core edition permits encodings whose plugins this session never registers. The compressor +/// must skip those schemes, as the writer would reject their output. +#[tokio::test] +async fn write_uses_only_registered_encodings() -> VortexResult<()> { + let session = array_session() + .with::() + .with::() + .with::(); + for family in EDITION_FAMILIES { + session.editions().declare_family(family)?; + } + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(CORE_2026_08_3)?; + + // Sorted integers, which FoR and BitPacking would otherwise compress. + let array = PrimitiveArray::from_iter(0..4096i64).into_array(); + let mut buf = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buf, array.to_array_stream()) + .await?; + let read = session + .open_options() + .open_buffer(buf.freeze())? + .scan()? + .into_array_stream()? + .read_all() + .await?; + assert_arrays_eq!(array, read, &mut session.create_execution_ctx()); + Ok(()) +} + #[tokio::test] async fn test_encoding_registered_after_write_options() -> VortexResult<()> { // A session that does not know about ZigZag yet. @@ -1874,7 +1913,7 @@ async fn write_read_roundtrip_with_layout( array: ArrayRef, use_list_layout: bool, ) -> VortexResult { - let strategy = crate::strategy::WriteStrategyBuilder::default() + let strategy = crate::strategy::WriteStrategyBuilder::from_session(&SESSION) .with_list_layout() .build(); let mut buf = ByteBufferMut::empty(); @@ -2253,14 +2292,14 @@ async fn timestamp_unit_mismatch() -> Result<(), Box> { #[tokio::test] async fn timestamp_unit_mismatch_errors_with_constant_children() -> Result<(), Box> { - let compressor = vortex_btrblocks::BtrBlocksCompressor::default(); + let compressor = BtrBlocksCompressor::from_session(&SESSION); // Write file with MILLISECONDS timestamps using this compressor. let ts_array = PrimitiveArray::from_iter(vec![1704067200000i64, 1704153600000, 1704240000000]) .into_array(); let temporal = TemporalArray::new_timestamp(ts_array, TimeUnit::Milliseconds, None); - let strategy = crate::strategy::WriteStrategyBuilder::default() + let strategy = crate::strategy::WriteStrategyBuilder::from_session(&SESSION) .with_compressor(compressor) .build(); @@ -2578,7 +2617,7 @@ async fn dict_probe_honours_configured_compressor() -> VortexResult<()> { let mut buf = ByteBufferMut::empty(); let summary = SESSION .write_options() - .with_strategy(crate::strategy::WriteStrategyBuilder::default().build()) + .with_strategy(crate::strategy::WriteStrategyBuilder::from_session(&SESSION).build()) .write(&mut buf, strings.clone().to_array_stream()) .await?; assert!( @@ -2586,15 +2625,18 @@ async fn dict_probe_honours_configured_compressor() -> VortexResult<()> { "default builder should produce a dict layout for low-cardinality strings" ); - let no_string_dict = - BtrBlocksCompressorBuilder::default().exclude_schemes([StringDictScheme.id()]); let mut buf = ByteBufferMut::empty(); let summary = SESSION .write_options() .with_strategy( - crate::strategy::WriteStrategyBuilder::default() - .with_btrblocks_builder(no_string_dict) - .build(), + crate::strategy::WriteStrategyBuilder::from_session_with_options( + &SESSION, + BtrBlocksOptions { + exclude_schemes: vec![StringDictScheme.id()], + ..Default::default() + }, + ) + .build(), ) .write(&mut buf, strings.to_array_stream()) .await?; @@ -2614,15 +2656,19 @@ async fn probe_compressor_override_is_independent() -> VortexResult<()> { let values: Vec<&str> = (0..n).map(|i| ["alpha", "beta", "gamma"][i % 3]).collect(); let strings = VarBinArray::from(values).into_array(); - let probe_without_dict = BtrBlocksCompressorBuilder::default() - .exclude_schemes([StringDictScheme.id()]) - .build(); + let probe_without_dict = BtrBlocksCompressor::from_session_with_options( + &SESSION, + &BtrBlocksOptions { + exclude_schemes: vec![StringDictScheme.id()], + ..Default::default() + }, + ); let mut buf = ByteBufferMut::empty(); let summary = SESSION .write_options() .with_strategy( - crate::strategy::WriteStrategyBuilder::default() + crate::strategy::WriteStrategyBuilder::from_session(&SESSION) .with_probe_compressor(probe_without_dict) .build(), ) diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 874d08be306..6248d012deb 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -30,7 +30,7 @@ use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::BtrBlocksOptions; use vortex_buffer::ByteBuffer; use vortex_edition::ComponentKind; use vortex_edition::EditionSessionExt; @@ -247,16 +247,17 @@ impl VortexWriteOptions { } else { ctx }; - let allowed_serialized_ids: HashSet = - ctx.array_ctx().to_ids().into_iter().collect(); let strategy = match self.strategy { Some(strategy) => strategy, - None => WriteStrategyBuilder::default() - .with_btrblocks_builder( - BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_serialized_ids), - ) - .build(), + // With editions disabled every registered encoding may be written. + None => WriteStrategyBuilder::from_session_with_options( + &self.session, + BtrBlocksOptions { + enforce_editions, + ..Default::default() + }, + ) + .build(), }; let dtype = stream.dtype().clone(); if enforce_editions { diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index 3f69de67394..56d3394f13b 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -73,7 +73,7 @@ async fn test_file_roundtrip() { // the b and the a.raw columns uncompressed. let default_strategy = Arc::new(CompressingStrategy::new( FlatLayoutStrategy::default(), - BtrBlocksCompressor::default(), + BtrBlocksCompressor::from_session(&SESSION), )); let writer = Arc::new( diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 10984c9d9e7..16c6fa06138 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -385,6 +385,7 @@ mod tests { use vortex_array::expr::root; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; + use vortex_btrblocks::BtrBlocksOptions; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::runtime::Handle; @@ -417,6 +418,18 @@ mod tests { .with_handle(handle) } + /// A compressor over every scheme registered on `session`: these tests compress in memory, where + /// no edition applies. + fn no_editions_compressor(session: &VortexSession) -> BtrBlocksCompressor { + BtrBlocksCompressor::from_session_with_options( + session, + &BtrBlocksOptions { + enforce_editions: false, + ..Default::default() + }, + ) + } + async fn write_dict_layout( array: ArrayRef, session: &VortexSession, @@ -426,7 +439,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(no_editions_compressor(session)), ); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -456,7 +469,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(no_editions_compressor(&session)), ); let array = VarBinArray::from_iter( @@ -559,7 +572,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(no_editions_compressor(&session)), ); let array = @@ -614,7 +627,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(no_editions_compressor(&session)), ); let array = VarBinArray::from_iter( diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..7334e7b76fd 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -117,8 +117,8 @@ impl TableStrategy { /// # use vortex_layout::layouts::table::TableStrategy; /// /// // A strategy for compressing data using the balanced BtrBlocks compressor. - /// let compress = - /// CompressingStrategy::new(FlatLayoutStrategy::default(), BtrBlocksCompressor::default()); + /// let compressor = BtrBlocksCompressor::from_session(&session); + /// let compress = CompressingStrategy::new(FlatLayoutStrategy::default(), compressor); /// /// // Our combined strategy uses no compression for validity buffers, BtrBlocks compression /// // for most columns, and stores a nested binary column uncompressed (flat) because it diff --git a/vortex-python/src/compress.rs b/vortex-python/src/compress.rs index 8688fbabf35..d3f9c2d7f74 100644 --- a/vortex-python/src/compress.rs +++ b/vortex-python/src/compress.rs @@ -56,7 +56,8 @@ pub fn compress(py: Python, array: PyArrayRef) -> PyVortexResult { let session = session(); let array = array.into_inner(); let compressed = py.detach(move || { - BtrBlocksCompressor::default().compress(&array, &mut session.create_execution_ctx()) + BtrBlocksCompressor::from_session(session) + .compress(&array, &mut session.create_execution_ctx()) })?; Ok(PyArrayRef::from(compressed)) } diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 7288fc5d1e3..b81595aab4c 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -16,13 +16,9 @@ use vortex::array::IntoArray; use vortex::array::iter::ArrayIterator; use vortex::array::iter::ArrayIteratorAdapter; use vortex::array::iter::ArrayIteratorExt; -use vortex::compressor::BtrBlocksCompressorBuilder; -use vortex::editions::ComponentKind; -use vortex::editions::EditionSessionExt; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; use vortex::io::VortexWrite; use vortex::io::object_store::ObjectStoreWrite; use vortex::io::runtime::BlockingRuntime; @@ -46,6 +42,7 @@ use crate::object_store::resolve::ResolvedStore; use crate::object_store::resolve::resolve_store; use crate::opendal_store::CosStore; use crate::opendal_store::GoosefsStore; +use crate::session::compact_session; use crate::session::session; pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { @@ -378,26 +375,18 @@ impl PyVortexWriteOptions { path: &str, store: Option, ) -> PyVortexResult<()> { - let session = session(); + let session = if self.use_compact_encodings { + compact_session() + } else { + session() + }; py.detach(|| { - let allowed_encodings = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - let mut compressor = BtrBlocksCompressorBuilder::default(); - if self.use_compact_encodings { - compressor = compressor.with_compact(); - } - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); - let strategy = strategy.build(); current_runtime().block_on(async move { match resolve_store(path, store.map(|x| x.into_inner()))? { ResolvedStore::ObjectStore(store, path) => { let mut store = ObjectStoreWrite::new(store, &path).await?; session .write_options() - .with_strategy(strategy) .write(&mut store, iter.into_inner().into_array_stream()) .await?; store.shutdown().await?; @@ -407,7 +396,6 @@ impl PyVortexWriteOptions { let mut w = FileWrite::create(path, current_runtime().handle()).await?; session .write_options() - .with_strategy(strategy) .write(&mut w, iter.into_inner().into_array_stream()) .await?; w.shutdown().await?; diff --git a/vortex-python/src/session.rs b/vortex-python/src/session.rs index 651addb5457..9e33dfbd6ea 100644 --- a/vortex-python/src/session.rs +++ b/vortex-python/src/session.rs @@ -11,8 +11,8 @@ use std::sync::atomic::AtomicPtr; use std::sync::atomic::Ordering; use vortex::VortexSessionDefault; +use vortex::compressor::CompressionSession; use vortex::io::runtime::BlockingRuntime; -#[cfg(unix)] use vortex::io::runtime::Handle; use vortex::io::session::RuntimeSessionExt; #[cfg(unix)] @@ -22,34 +22,58 @@ use vortex::session::VortexSession; use crate::current_runtime; #[cfg(not(unix))] -static SESSION: LazyLock = LazyLock::new(new_session); +static SESSION: LazyLock = LazyLock::new(|| new_session(current_runtime().handle())); +#[cfg(not(unix))] +static COMPACT_SESSION: LazyLock = + LazyLock::new(|| new_compact_session(current_runtime().handle())); -/// The shared session is published without an initialization lock so a forked child cannot inherit -/// it in a permanently initializing state. +/// The shared sessions are published without an initialization lock so a forked child cannot +/// inherit them in a permanently initializing state. #[cfg(unix)] static SESSION: AtomicPtr = AtomicPtr::new(ptr::null_mut()); +#[cfg(unix)] +static COMPACT_SESSION: AtomicPtr = AtomicPtr::new(ptr::null_mut()); #[cfg(not(unix))] pub(crate) fn session() -> &'static VortexSession { &SESSION } +/// The shared session plus the compact (Zstd and Pco) schemes, for compact writes. +#[cfg(not(unix))] +pub(crate) fn compact_session() -> &'static VortexSession { + &COMPACT_SESSION +} + #[cfg(unix)] pub(crate) fn session() -> &'static VortexSession { + published(&SESSION, new_session) +} + +/// The shared session plus the compact (Zstd and Pco) schemes, for compact writes. +#[cfg(unix)] +pub(crate) fn compact_session() -> &'static VortexSession { + published(&COMPACT_SESSION, new_compact_session) +} + +/// Returns the session published in `slot`, building it with `new` on first use. +#[cfg(unix)] +fn published( + slot: &AtomicPtr, + new: fn(Handle) -> VortexSession, +) -> &'static VortexSession { // Ensure a forked child has published its new runtime and repointed an existing session before // returning that session to a caller. let runtime = current_runtime(); loop { - let current = SESSION.load(Ordering::Acquire); + let current = slot.load(Ordering::Acquire); if !current.is_null() { // SAFETY: A published session is never freed or replaced. return unsafe { &*current }; } - let fresh = Box::into_raw(Box::new( - VortexSession::default().with_handle(runtime.handle()), - )); - match SESSION.compare_exchange(current, fresh, Ordering::AcqRel, Ordering::Acquire) { + let fresh = Box::into_raw(Box::new(new(runtime.handle()))); + match slot.compare_exchange(current, fresh, Ordering::AcqRel, Ordering::Acquire) { Ok(_) => { // SAFETY: `fresh` was just published and published sessions are never freed. return unsafe { &*fresh }; @@ -62,26 +86,33 @@ pub(crate) fn session() -> &'static VortexSession { } } -#[cfg(not(unix))] -fn new_session() -> VortexSession { - VortexSession::default().with_handle(current_runtime().handle()) +fn new_session(handle: Handle) -> VortexSession { + VortexSession::default().with_handle(handle) } -/// Point the shared session at `handle`, replacing any previously configured runtime handle. +fn new_compact_session(handle: Handle) -> VortexSession { + let session = new_session(handle); + session.register(CompressionSession::compact()); + session +} + +/// Point the shared sessions at `handle`, replacing any previously configured runtime handle. /// /// A [`Handle`] is a weak reference to its executor, so after the runtime is rebuilt in a forked -/// child (see [`crate::current_runtime`]) the session must be repointed, or every spawn would panic -/// on a dropped runtime. `VortexSession` has interior mutability, so the change is visible through -/// every existing clone of the session. +/// child (see [`crate::current_runtime`]) the sessions must be repointed, or every spawn would +/// panic on a dropped runtime. `VortexSession` has interior mutability, so the change is visible +/// through every existing clone of a session. /// -/// Does nothing if the session has not been built yet — in that case it will pick up the current -/// runtime's handle when it is first built. +/// Does nothing for a session that has not been built yet — it will pick up the current runtime's +/// handle when it is first built. #[cfg(unix)] pub(crate) fn reset_session_handle(handle: Handle) { - let current = SESSION.load(Ordering::Acquire); - if !current.is_null() { - // SAFETY: A published session is never freed or replaced. - let session = unsafe { &*current }; - session.session().with_handle(handle); + for slot in [&SESSION, &COMPACT_SESSION] { + let current = slot.load(Ordering::Acquire); + if !current.is_null() { + // SAFETY: A published session is never freed or replaced. + let session = unsafe { &*current }; + session.session().with_handle(handle.clone()); + } } } diff --git a/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs index 55aac492306..a2ac0dd6a5d 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs @@ -17,7 +17,8 @@ pub fn fixtures() -> Vec> { #[cfg(test)] mod tests { use vortex::VortexSessionDefault; - use vortex::compressor::BtrBlocksCompressorBuilder; + use vortex::compressor::BtrBlocksOptions; + use vortex::compressor::CompressionSession; use vortex::editions::CORE_2026_08_3; use vortex::editions::EditionSessionExt; use vortex::file::WriteStrategyBuilder; @@ -28,6 +29,14 @@ mod tests { use super::fixtures; use crate::adapter; + /// The adapter writes with editions disabled, so every registered scheme may be used. + fn no_editions() -> BtrBlocksOptions { + BtrBlocksOptions { + enforce_editions: false, + ..Default::default() + } + } + fn is_clickbench_fixture(name: &str) -> bool { name.contains("clickbench") } @@ -36,6 +45,9 @@ mod tests { fn roundtrip_non_clickbench_fixtures_to_bytes() -> VortexResult<()> { let session = VortexSession::default(); session.enable_edition(CORE_2026_08_3)?; + let compact_session = VortexSession::default(); + compact_session.enable_edition(CORE_2026_08_3)?; + compact_session.register(CompressionSession::compact()); for dataset in fixtures() .into_iter() .filter(|fixture| !is_clickbench_fixture(fixture.name())) @@ -44,15 +56,14 @@ mod tests { let regular_bytes = adapter::write_compressed_to_bytes_with_session( &session, array.clone(), - WriteStrategyBuilder::default().build(), + WriteStrategyBuilder::from_session_with_options(&session, no_editions()).build(), )?; let _regular = adapter::read_file(regular_bytes)?; let compact_bytes = adapter::write_compressed_to_bytes_with_session( - &session, + &compact_session, array, - WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) + WriteStrategyBuilder::from_session_with_options(&compact_session, no_editions()) .build(), )?; let _compact = adapter::read_file(compact_bytes)?; diff --git a/vortex-test/compat-gen/src/fixtures/mod.rs b/vortex-test/compat-gen/src/fixtures/mod.rs index 291d9a5f5ff..74e027d6a71 100644 --- a/vortex-test/compat-gen/src/fixtures/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/mod.rs @@ -6,10 +6,13 @@ mod arrays; use std::path::Path; use std::sync::Arc; +use vortex::VortexSessionDefault; use vortex::array::ArrayId; use vortex::array::ArrayRef; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::BtrBlocksOptions; +use vortex::compressor::CompressionSession; use vortex::file::WriteStrategyBuilder; +use vortex::session::VortexSession; use vortex_array::ExecutionCtx; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; @@ -137,15 +140,21 @@ impl Fixture for DatasetFixtureAdapter { fn write(&self, dir: &Path, ctx: &mut ExecutionCtx) -> VortexResult> { let array = self.inner.build(&ctx.session().arrow())?; let path = dir.join(self.name()); + // The adapter writes with editions disabled, so every scheme registered on this default + // session may be used. + let session = VortexSession::default(); if self.compact { - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) - .build(); - adapter::write_compressed(&path, array, strategy)?; - } else { - let strategy = WriteStrategyBuilder::default().build(); - adapter::write_compressed(&path, array, strategy)?; + session.register(CompressionSession::compact()); } + let strategy = WriteStrategyBuilder::from_session_with_options( + &session, + BtrBlocksOptions { + enforce_editions: false, + ..Default::default() + }, + ) + .build(); + adapter::write_compressed(&path, array, strategy)?; Ok(vec![FixtureEntry { name: self.name().to_string(), description: self.description().to_string(), diff --git a/vortex-tui/src/convert.rs b/vortex-tui/src/convert.rs index ab316982b27..73f0948110d 100644 --- a/vortex-tui/src/convert.rs +++ b/vortex-tui/src/convert.rs @@ -12,14 +12,13 @@ use indicatif::ProgressBar; use parquet::arrow::ParquetRecordBatchStreamBuilder; use tokio::fs::File; use tokio::io::AsyncWriteExt; +use vortex::VortexSessionDefault; use vortex::array::stream::ArrayStreamAdapter; -use vortex::compressor::BtrBlocksCompressorBuilder; -use vortex::editions::ComponentKind; -use vortex::editions::EditionSessionExt; +use vortex::compressor::CompressionSession; use vortex::error::VortexExpect; use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; +use vortex::io::session::RuntimeSessionExt; use vortex::session::VortexSession; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; @@ -98,21 +97,19 @@ pub async fn exec_convert(session: &VortexSession, flags: ConvertArgs) -> anyhow .boxed(); } - let allowed_encodings = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - let mut compressor = BtrBlocksCompressorBuilder::default(); - if matches!(flags.strategy, Strategy::Compact) { - compressor = compressor.with_compact(); - } - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); + // Compact files come from a session like `session` that also registers the compact schemes. + let compact_session; + let session = if matches!(flags.strategy, Strategy::Compact) { + compact_session = VortexSession::default().with_handle(session.handle()); + compact_session.register(CompressionSession::compact()); + &compact_session + } else { + session + }; let mut file = File::create(output_path).await?; session .write_options() - .with_strategy(strategy.build()) .write(&mut file, ArrayStreamAdapter::new(dtype, vortex_stream)) .await?; file.shutdown().await?; diff --git a/vortex/examples/compression_showcase.rs b/vortex/examples/compression_showcase.rs index 6aa9cf08218..b8954f881cc 100644 --- a/vortex/examples/compression_showcase.rs +++ b/vortex/examples/compression_showcase.rs @@ -67,7 +67,7 @@ fn compress_sequential_data(session: &VortexSession) -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box> { println!("This example demonstrates using Vortex as a backend for structured logging.\n"); let session = VortexSession::default(); + // Use compact encodings (Pco + Zstd) for the telemetry files. + session.register(CompressionSession::compact()); // Create output directory let output_dir: PathBuf = "vortex-traces/".into(); @@ -389,12 +390,7 @@ async fn write_batch_to_vortex( let file_path = output_dir.join(format!("traces_{:04}.vortex", file_index)); let mut file = tokio::fs::File::create(&file_path).await?; - // Use compact encodings (Pco + Zstd) for the telemetry files. - let write_opts = session.write_options().with_strategy( - WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) - .build(), - ); + let write_opts = session.write_options(); write_opts .write(&mut file, struct_array.into_array().to_array_stream()) diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index 98a4f4a7d30..ccdf6204228 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -567,30 +567,22 @@ async fn btrblocks_respects_enabled_array_encodings() -> VortexResult<()> { Ok(()) } -/// An explicitly supplied strategy is not reconfigured by the writer. Its unsupported output is -/// still caught by the serialization context. +/// An explicit default strategy is built from the session, so it only emits the encodings the +/// enabled editions permit. #[tokio::test] -async fn explicit_btrblocks_strategy_is_not_reconfigured() -> VortexResult<()> { +async fn explicit_default_strategy_respects_enabled_editions() -> VortexResult<()> { let session = writer_test_session()?; - let strategy = WriteStrategyBuilder::default().build(); + let strategy = WriteStrategyBuilder::from_session(&session).build(); let mut buffer = ByteBufferMut::empty(); - let error = session + session .write_options() .with_strategy(strategy) .write( &mut buffer, sequential_integers().into_array().to_array_stream(), ) - .await - .err() - .ok_or_else(|| vortex_err!("explicit BtrBlocks strategy was unexpectedly reconfigured"))?; - assert!( - error - .to_string() - .contains("Serialized array ID vortex.sequence not permitted by ctx"), - "unexpected error: {error}" - ); + .await?; Ok(()) } @@ -601,7 +593,7 @@ async fn explicit_btrblocks_strategy_is_not_reconfigured() -> VortexResult<()> { #[tokio::test] async fn serialization_context_rejects_unsupported_compressor_output() -> VortexResult<()> { let session = writer_test_session()?; - let strategy = WriteStrategyBuilder::default() + let strategy = WriteStrategyBuilder::from_session(&session) .with_compressor(forbidden_sequence_compressor) .build(); let mut buffer = ByteBufferMut::empty(); @@ -633,7 +625,7 @@ async fn serialization_context_accepts_supported_compressor_output() -> VortexRe use crate::VortexSessionDefault; let session = VortexSession::default(); - let strategy = WriteStrategyBuilder::default() + let strategy = WriteStrategyBuilder::from_session(&session) .with_compressor(forbidden_sequence_compressor) .build(); let mut buffer = ByteBufferMut::empty(); diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index f171522faee..0a483fbbff1 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -48,7 +48,7 @@ //! # fn example() -> vortex::error::VortexResult<()> { //! let session = VortexSession::default(); //! let array = PrimitiveArray::new(buffer![42u64; 1024], Validity::NonNullable).into_array(); -//! let compressed = BtrBlocksCompressor::default() +//! let compressed = BtrBlocksCompressor::from_session(&session) //! .compress(&array, &mut session.create_execution_ctx())?; //! //! assert_eq!(compressed.dtype(), array.dtype()); @@ -114,6 +114,7 @@ pub use vortex_array::scalar_fn; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; use vortex_array::stats::session::StatsSession; +use vortex_btrblocks::CompressionSession; use vortex_io::session::RuntimeSession; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; @@ -144,7 +145,9 @@ pub mod buffer { /// Default adaptive compression APIs based on the maintained BtrBlocks-style compressor. pub mod compressor { pub use vortex_btrblocks::BtrBlocksCompressor; - pub use vortex_btrblocks::BtrBlocksCompressorBuilder; + pub use vortex_btrblocks::BtrBlocksOptions; + pub use vortex_btrblocks::CompressionSession; + pub use vortex_btrblocks::CompressionSessionExt; pub use vortex_btrblocks::Scheme; pub use vortex_btrblocks::SchemeId; } @@ -320,6 +323,7 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::() + .with::() .with::() .with::() .with::() @@ -366,12 +370,11 @@ mod test { use vortex_array::expr::select; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; - use vortex_btrblocks::BtrBlocksCompressorBuilder; + use vortex_btrblocks::CompressionSession; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; - use vortex_file::WriteStrategyBuilder; use vortex_session::VortexSession; use crate as vortex; @@ -422,7 +425,7 @@ mod test { // You can compress an array in-memory with the BtrBlocks compressor let session = VortexSession::default(); - let compressed = BtrBlocksCompressor::default().compress( + let compressed = BtrBlocksCompressor::from_session(&session).compress( &array.clone().into_array(), &mut session.create_execution_ctx(), )?; @@ -486,13 +489,9 @@ mod test { let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("example_compact.vortex"); + session.register(CompressionSession::compact()); session .write_options() - .with_strategy( - WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) - .build(), - ) .write( &mut tokio::fs::File::create(&path).await?, array.clone().into_array().to_array_stream(), diff --git a/wasm-test/src/main.rs b/wasm-test/src/main.rs index 964d3a36c9a..508ae4e24b9 100644 --- a/wasm-test/src/main.rs +++ b/wasm-test/src/main.rs @@ -17,7 +17,7 @@ pub fn main() { let array = PrimitiveArray::new(buffer![1i32; 1024], Validity::AllValid).into_array(); let session = VortexSession::default(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::from_session(&session) .compress(&array, &mut session.create_execution_ctx()) .unwrap(); println!("Compressed size: {}", compressed.len());