From bca8e159eb515bf7c04096203e880059f330540a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 15:18:43 -0400 Subject: [PATCH 01/11] Register compression schemes on the session Signed-off-by: Matt Katz --- Cargo.lock | 1 + benchmarks/compress-bench/README.md | 2 +- benchmarks/compress-bench/src/gpu/vortex.rs | 9 +- benchmarks/string-bench/src/serialized.rs | 36 +- encodings/parquet-variant/src/vtable.rs | 11 +- fuzz/fuzz_targets/file_io.rs | 9 +- fuzz/src/array/mod.rs | 14 +- fuzz/src/lib.rs | 17 + vortex-bench/src/conversions.rs | 13 +- vortex-bench/src/lib.rs | 35 +- vortex-btrblocks/Cargo.toml | 2 +- vortex-btrblocks/benches/compress.rs | 8 +- vortex-btrblocks/benches/compress_listview.rs | 8 +- vortex-btrblocks/src/builder.rs | 321 ------------------ vortex-btrblocks/src/canonical_compressor.rs | 290 ++-------------- vortex-btrblocks/src/lib.rs | 114 ++++++- .../schemes/float/scheme_selection_tests.rs | 14 +- vortex-btrblocks/src/schemes/float/tests.rs | 12 +- .../schemes/integer/scheme_selection_tests.rs | 44 +-- vortex-btrblocks/src/schemes/integer/tests.rs | 16 +- vortex-btrblocks/src/schemes/string/fsst.rs | 2 +- .../schemes/string/scheme_selection_tests.rs | 39 ++- vortex-btrblocks/src/schemes/string/tests.rs | 10 +- vortex-btrblocks/src/tests.rs | 260 ++++++++++++++ vortex-btrblocks/src/trace_tests.rs | 10 +- vortex-btrblocks/tests/golden.rs | 77 ++--- vortex-btrblocks/tests/onpair_roundtrip.rs | 16 +- vortex-btrblocks/tests/varbin_scheme.rs | 33 +- vortex-compressor/Cargo.toml | 3 +- vortex-compressor/src/lib.rs | 1 + vortex-compressor/src/session.rs | 133 ++++++++ vortex-cuda/src/layout.rs | 68 +++- vortex-ffi/src/sink.rs | 2 +- vortex-file/benches/split_collection.rs | 7 +- vortex-file/src/lib.rs | 2 + vortex-file/src/strategy.rs | 50 +-- vortex-file/src/tests.rs | 35 +- vortex-file/src/writer.rs | 13 +- vortex-file/tests/test_write_table.rs | 2 +- vortex-layout/src/layouts/dict/reader.rs | 14 +- vortex-layout/src/layouts/table.rs | 4 +- vortex-python/src/compress.rs | 3 +- vortex-python/src/io.rs | 23 +- .../src/fixtures/arrays/datasets/mod.rs | 17 +- vortex-test/compat-gen/src/fixtures/mod.rs | 22 +- vortex-tui/src/convert.rs | 23 +- vortex/Cargo.toml | 2 +- vortex/examples/compression_showcase.rs | 12 +- vortex/examples/tracing_vortex.rs | 15 +- vortex/src/editions/tests.rs | 24 +- vortex/src/lib.rs | 29 +- wasm-test/src/main.rs | 4 +- 52 files changed, 1006 insertions(+), 925 deletions(-) delete mode 100644 vortex-btrblocks/src/builder.rs create mode 100644 vortex-btrblocks/src/tests.rs create mode 100644 vortex-compressor/src/session.rs diff --git a/Cargo.lock b/Cargo.lock index 041f9a9e2f6..e8969b748f2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10923,6 +10923,7 @@ dependencies = [ "tracing", "vortex-array", "vortex-buffer", + "vortex-edition", "vortex-error", "vortex-mask", "vortex-session", diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index d309a2122f1..541c5382537 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 + (`cuda_compatible_schemes`) 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..b9214df5294 100644 --- a/benchmarks/compress-bench/src/gpu/vortex.rs +++ b/benchmarks/compress-bench/src/gpu/vortex.rs @@ -22,7 +22,6 @@ 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::error::VortexResult; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; @@ -37,7 +36,6 @@ 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; @@ -46,6 +44,7 @@ use vortex_cuda::CudaSession; use vortex_cuda::PooledFileReadAtOptions; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::CudaFlatLayoutStrategy; +use vortex_cuda::layout::cuda_compressor; use vortex_cuda::layout::register_cuda_layout; use crate::gpu::writer::GPU_ROW_GROUP_SIZE; @@ -100,11 +99,7 @@ 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(), + cuda_compressor(&SESSION), ))); let start = Instant::now(); SESSION diff --git a/benchmarks/string-bench/src/serialized.rs b/benchmarks/string-bench/src/serialized.rs index 62de3dd9a51..825746ee1be 100644 --- a/benchmarks/string-bench/src/serialized.rs +++ b/benchmarks/string-bench/src/serialized.rs @@ -31,7 +31,8 @@ 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::CompressionSessionExt; +use vortex::compressor::Scheme; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; @@ -57,7 +58,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,16 +171,23 @@ 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) + let excluded: Vec = default_string_scheme_ids() + .into_iter() + .filter(|&id| id != forced) + .chain([DeltaScheme::default().id()]) + .collect(); + let schemes: Vec<&'static dyn Scheme> = session + .permitted_schemes() + .into_iter() + .filter(|scheme| !excluded.contains(&scheme.id())) + .collect(); + WriteStrategyBuilder::from_session(session) + .with_schemes(schemes) .build() } @@ -251,7 +259,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 +363,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::DEFAULT_SCHEMES; use vortex_btrblocks::SchemeExt; use super::*; @@ -365,7 +373,7 @@ 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 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..35027856fea 100644 --- a/fuzz/src/lib.rs +++ b/fuzz/src/lib.rs @@ -56,6 +56,10 @@ mod native_runtime { use std::sync::LazyLock; use vortex::VortexSessionDefault; + #[cfg(feature = "zstd")] + use vortex::compressor::COMPACT_SCHEMES; + #[cfg(feature = "zstd")] + use vortex::compressor::CompressionSessionExt; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::current::CurrentThreadRuntime; use vortex_io::session::RuntimeSessionExt; @@ -76,8 +80,21 @@ 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); + for scheme in COMPACT_SCHEMES { + session.register_scheme(*scheme); + } + 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..3892f2ad3fd 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,7 @@ use wkb::writer::write_geometry; use crate::CompactionStrategy; use crate::Format; use crate::SESSION; -use crate::retain_edition_encodings; +use crate::compact_schemes; use crate::utils::file::idempotent_async; /// Memory budget per concurrent conversion stream in GB. This is somewhat arbitary. @@ -246,12 +246,9 @@ fn write_options_for( return compaction.apply_options(SESSION.write_options()); } - let mut builder = WriteStrategyBuilder::default(); + let mut builder = WriteStrategyBuilder::from_session(&SESSION); if matches!(compaction, CompactionStrategy::Compact) { - builder = builder.with_btrblocks_builder(retain_edition_encodings( - &SESSION, - BtrBlocksCompressorBuilder::default().with_compact(), - )); + builder = builder.with_schemes(compact_schemes()); } for name in binary_fields { builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); @@ -263,7 +260,7 @@ fn write_options_for( fn no_dict_layout() -> Arc { Arc::new(CompressingStrategy::new( ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), - retain_edition_encodings(&SESSION, BtrBlocksCompressorBuilder::default()).build(), + BtrBlocksCompressor::from_session(&SESSION), )) } diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 569dfc74a4d..18594fdd03e 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -28,7 +28,9 @@ use tpcds::TpcDsBenchmark; use tpch::benchmark::TpcHBenchmark; pub use utils::file::*; pub use utils::logging::*; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::COMPACT_SCHEMES; +use vortex::compressor::CompressionSessionExt; +use vortex::compressor::Scheme; use vortex::error::VortexExpect; use vortex::error::vortex_err; use vortex::file::VortexWriteOptions; @@ -70,8 +72,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; @@ -254,11 +254,8 @@ impl CompactionStrategy { pub fn apply_options(&self, options: VortexWriteOptions) -> VortexWriteOptions { match self { CompactionStrategy::Compact => options.with_strategy( - WriteStrategyBuilder::default() - .with_btrblocks_builder(retain_edition_encodings( - &SESSION, - BtrBlocksCompressorBuilder::default().with_compact(), - )) + WriteStrategyBuilder::from_session(&SESSION) + .with_schemes(compact_schemes()) .build(), ), CompactionStrategy::Default => options, @@ -266,19 +263,15 @@ impl CompactionStrategy { } } -/// 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) +/// The schemes [`SESSION`] permits plus the compact ones, for [`CompactionStrategy::Compact`]. +pub fn compact_schemes() -> Vec<&'static dyn Scheme> { + SESSION.permit( + SESSION + .registered_schemes() + .into_iter() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ) } /// Verify that local data has already been prepared for the requested benchmark formats. diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 24a03337768..9c499b81eaa 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -33,6 +33,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 } @@ -51,7 +52,6 @@ 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..5f8801ddfbb 100644 --- a/vortex-btrblocks/benches/compress.rs +++ b/vortex-btrblocks/benches/compress.rs @@ -24,7 +24,11 @@ mod benchmarks { use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session + }); fn make_clickbench_window_name() -> ArrayRef { // A test that's meant to mirror the WindowName column from ClickBench. @@ -51,7 +55,7 @@ mod benchmarks { let array = make_clickbench_window_name() .execute::(&mut ctx) .unwrap(); - let compressor = BtrBlocksCompressor::default(); + let compressor = BtrBlocksCompressor::from_session_no_editions(&SESSION); 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..aed620db9fc 100644 --- a/vortex-btrblocks/benches/compress_listview.rs +++ b/vortex-btrblocks/benches/compress_listview.rs @@ -31,7 +31,11 @@ mod benchmarks { const NUM_ROWS: usize = 8192; const SEED: u64 = 42; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session + }); const SHORT_STRINGS: &[&str] = &[ "alpha_one", @@ -183,7 +187,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::from_session_no_editions(&SESSION); 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..d131feb0c2c 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -8,28 +8,28 @@ use std::ops::Deref; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_error::VortexResult; +use vortex_session::VortexSession; -use crate::BtrBlocksCompressorBuilder; use crate::CascadingCompressor; +use crate::CompressionSessionExt; -/// 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_no_editions`](Self::from_session_no_editions) +/// keeps every registered scheme, for in-memory compression where no edition applies. /// /// # Examples /// /// ```rust -/// use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt}; -/// use vortex_btrblocks::schemes::integer::IntDictScheme; +/// use vortex_btrblocks::BtrBlocksCompressor; /// -/// // Default compressor - all schemes allowed. -/// let compressor = BtrBlocksCompressor::default(); +/// let session = vortex_array::array_session(); +/// vortex_btrblocks::initialize(&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_no_editions(&session); /// ``` #[derive(Clone)] pub struct BtrBlocksCompressor( @@ -38,6 +38,19 @@ pub struct BtrBlocksCompressor( ); impl BtrBlocksCompressor { + /// 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(CascadingCompressor::new(session.permitted_schemes())) + } + + /// Creates a compressor over every scheme registered on `session`, ignoring editions. + /// + /// Use for in-memory compression, where no edition restricts what a file may contain. + pub fn from_session_no_editions(session: &VortexSession) -> Self { + Self(CascadingCompressor::new(session.registered_schemes())) + } + /// Compresses an array using BtrBlocks-inspired compression. pub fn compress(&self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.0.compress(array, ctx) @@ -51,256 +64,3 @@ impl Deref for BtrBlocksCompressor { &self.0 } } - -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 crate::BtrBlocksCompressor; - #[cfg(feature = "zstd")] - use crate::BtrBlocksCompressorBuilder; - - 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(()) - } - - #[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(()) - } - - #[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(()) - } -} diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 2e8ae484f90..9cb8e9aae65 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -30,8 +30,12 @@ //! //! 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). Schemes are registered on a session: [`initialize`] registers [`DEFAULT_SCHEMES`], +//! and [`BtrBlocksCompressor::from_session`] compresses with the registered schemes whose +//! serialized IDs the session's enabled editions permit ([`from_session_no_editions`] ignores +//! them). +//! +//! [`from_session_no_editions`]: BtrBlocksCompressor::from_session_no_editions //! //! 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,42 +47,35 @@ //! 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_buffer::buffer; //! //! # fn example() -> vortex_error::VortexResult<()> { //! let session = array_session(); +//! vortex_btrblocks::initialize(&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_no_editions(&session); //! 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; #[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 schemes::patches::compress_patches; pub use vortex_compressor::CascadingCompressor; @@ -87,9 +84,96 @@ pub use vortex_compressor::scheme::MAX_CASCADE; pub use vortex_compressor::scheme::Scheme; pub use vortex_compressor::scheme::SchemeExt; pub use vortex_compressor::scheme::SchemeId; +pub use vortex_compressor::session::CompressionSession; +pub use vortex_compressor::session::CompressionSessionExt; pub use vortex_compressor::stats::ArrayAndStats; pub use vortex_compressor::stats::BoolStats; pub use vortex_compressor::stats::FloatStats; pub use vortex_compressor::stats::GenerateStatsOptions; pub use vortex_compressor::stats::IntegerStats; pub use vortex_compressor::stats::StringStats; +use vortex_session::VortexSession; + +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. +/// +/// This list is order-sensitive: [`initialize`] registers it in this order and the compressor +/// preserves registration order, so that tie-breaking is deterministic. +pub 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, +]; + +/// Compact schemes (Zstd for strings and binary, Pco for numerics when the `pco` feature is on). +/// +/// Not part of [`DEFAULT_SCHEMES`]: they trade decode speed for compression ratio, so callers add +/// them to a compressor's scheme list explicitly. +#[cfg(feature = "zstd")] +pub const COMPACT_SCHEMES: &[&dyn Scheme] = &[ + &string::ZstdScheme, + &binary::ZstdScheme, + #[cfg(feature = "pco")] + &integer::PcoScheme, + #[cfg(feature = "pco")] + &float::PcoScheme, +]; + +/// Delta, kept out of [`DEFAULT_SCHEMES`] because it is slower to decompress than the schemes that +/// would otherwise win. Callers that want it add it to their scheme list and permit +/// `fastlanes.delta`. +/// +/// TODO(robert): Return it to [`DEFAULT_SCHEMES`] once we have scheme filtering. +pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); + +/// Registers [`DEFAULT_SCHEMES`] on `session`, in order. +/// +/// Registration is idempotent, so this may run more than once. +pub fn initialize(session: &VortexSession) { + for scheme in DEFAULT_SCHEMES { + session.register_scheme(*scheme); + } +} diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index c6f54a9dafe..96a186adc41 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -21,13 +21,17 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] 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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -37,7 +41,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -50,7 +54,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert!(compressed.children()[0].is::()); @@ -69,7 +73,7 @@ fn test_null_dominated_compressed() -> VortexResult<()> { } builder.append_nulls(95); let array = builder.finish_into_primitive(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); 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..e21b031ad70 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -22,11 +22,15 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::float::FloatRLEScheme; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_empty() -> VortexResult<()> { - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); let array = PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable).into_array(); let result = btr.compress(&array, &mut SESSION.create_execution_ctx())?; @@ -42,7 +46,7 @@ fn test_compress() -> VortexResult<()> { } let array = values.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 1024); @@ -92,7 +96,7 @@ fn test_sparse_compression() -> VortexResult<()> { array.append_nulls(90); let array = array.finish_into_primitive().into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); 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..d1bbbfe4e4d 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -28,15 +28,27 @@ use vortex_session::VortexSession; use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; -use crate::BtrBlocksCompressorBuilder; +use crate::CompressionSessionExt; use crate::DELTA_SCHEME; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + +/// The default schemes plus opt-in Delta. +fn with_delta() -> BtrBlocksCompressor { + let session = vortex_array::array_session(); + crate::initialize(&session); + session.register_scheme(&DELTA_SCHEME); + BtrBlocksCompressor::from_session_no_editions(&session) +} #[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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -46,7 +58,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -56,7 +68,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert_eq!( @@ -85,7 +97,7 @@ fn test_sparse_compressed() -> VortexResult<()> { } } let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -109,7 +121,7 @@ fn test_dict_compressed() -> VortexResult<()> { } let array = PrimitiveArray::new(Buffer::copy_from(&codes), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -122,7 +134,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -132,7 +144,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -149,7 +161,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; eprintln!("{}", compressed.display_tree()); assert!(compressed.is::()); @@ -176,9 +188,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 +226,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 +253,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..b331dece461 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -27,12 +27,16 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_empty() -> VortexResult<()> { // Make sure empty array compression does not fail. - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); let array = PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable); let result = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; @@ -60,7 +64,7 @@ fn test_dict_encodable() -> VortexResult<()> { } } - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress( &codes.freeze().into_array(), &mut SESSION.create_execution_ctx(), @@ -80,7 +84,7 @@ fn constant_mostly_nulls() -> VortexResult<()> { ); let validity = array.validity()?; - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -99,7 +103,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -141,7 +145,7 @@ fn compress_large_int() -> VortexResult<()> { .collect::() .into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); 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..e0c9c7f44b9 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 +/// [`crate::DEFAULT_SCHEMES`] (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..e0e72a1e4dc 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -17,16 +17,22 @@ use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::CompressionSessionExt; +use crate::DEFAULT_SCHEMES; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_constant_compressed() -> VortexResult<()> { let strings: Vec> = vec![Some("constant_value"); 100]; 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())?; + let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -40,8 +46,8 @@ 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())?; + let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -51,10 +57,10 @@ 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 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 DEFAULT_SCHEMES" ); } @@ -71,8 +77,8 @@ 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())?; + let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), "expected OnPair, got {}", @@ -81,21 +87,20 @@ 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. 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 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 +111,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 session = vortex_array::array_session(); + session.register_scheme(&FSSTScheme); + let compressor = BtrBlocksCompressor::from_session_no_editions(&session); 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..e76a46b36fb 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -16,7 +16,11 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_strings() -> VortexResult<()> { @@ -30,7 +34,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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 2048); @@ -57,7 +61,7 @@ fn test_sparse_nulls() -> VortexResult<()> { let strings = strings.finish_into_varbinview(); let array_ref = strings.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); let compressed = btr.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 100); diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs new file mode 100644 index 00000000000..b3225c0e789 --- /dev/null +++ b/vortex-btrblocks/src/tests.rs @@ -0,0 +1,260 @@ +// 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; +#[cfg(feature = "zstd")] +use crate::COMPACT_SCHEMES; +use crate::CompressionSessionExt; +#[cfg(feature = "zstd")] +use crate::Scheme; +#[cfg(feature = "zstd")] +use crate::schemes::binary; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + 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::from_session_no_editions(&SESSION) + .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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); + 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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); + 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::from_session_no_editions(&SESSION); + 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::from_session_no_editions(&SESSION); + 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::from_session_no_editions(&SESSION); + 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::from_session_no_editions(&SESSION); + 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::from_session_no_editions(&SESSION); + 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 session = vortex_array::array_session(); + crate::initialize(&session); + for scheme in COMPACT_SCHEMES { + session.register_scheme(*scheme); + } + let compressor = BtrBlocksCompressor::from_session_no_editions(&session); + 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 session = vortex_array::array_session(); + session.register_scheme(scheme); + let compressor = BtrBlocksCompressor::from_session_no_editions(&session); + 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..62c4720d6e2 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -53,7 +53,8 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use crate::BtrBlocksCompressorBuilder; +use crate::BtrBlocksCompressor; +use crate::CompressionSessionExt; use crate::DELTA_SCHEME; /// A session with the default Vortex encodings registered. @@ -127,9 +128,10 @@ 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() + let session = trace_session(); + crate::initialize(&session); + session.register_scheme(&DELTA_SCHEME); + BtrBlocksCompressor::from_session_no_editions(&session) .compress(&lineitem()?, &mut execution_ctx()) } diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index fc636252c27..3fc43a81877 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 [`COMPACT_SCHEMES`](vortex_btrblocks::COMPACT_SCHEMES) +//! — 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. @@ -49,9 +49,14 @@ use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +#[cfg(all(feature = "zstd", feature = "pco"))] +use vortex_btrblocks::COMPACT_SCHEMES; +use vortex_btrblocks::CompressionSessionExt; +use vortex_btrblocks::DEFAULT_SCHEMES; +use vortex_btrblocks::Scheme; +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,22 @@ 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(schemes: Vec<&'static dyn Scheme>) -> Vec<&'static dyn Scheme> { + schemes + .into_iter() + .filter(|scheme| scheme.id() != OnPairScheme.id()) + .collect() } -fn edition_session(editions: &[EditionId]) -> VortexResult { +/// A session registering `schemes` and enabling `editions`. +fn edition_session( + editions: &[EditionId], + schemes: Vec<&'static dyn Scheme>, +) -> VortexResult { let session = vortex_array::array_session().with::(); + for scheme in schemes { + session.register_scheme(scheme); + } for family in EDITION_FAMILIES { session.editions().declare_family(family)?; } @@ -414,43 +426,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], without_onpair(DEFAULT_SCHEMES.to_vec()))?; + let compressor = BtrBlocksCompressor::from_session(&session); 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], DEFAULT_SCHEMES.to_vec())?; + let compressor = BtrBlocksCompressor::from_session(&session); golden_snapshots( "onpair", &compressor, @@ -461,12 +448,14 @@ fn golden_onpair() -> VortexResult<()> { #[cfg(all(feature = "zstd", feature = "pco"))] #[test] fn golden_compact() -> VortexResult<()> { - let session = edition_session(&[CORE_2026_08_3])?; + let schemes = DEFAULT_SCHEMES + .iter() + .chain(COMPACT_SCHEMES.iter()) + .copied() + .collect(); + let session = edition_session(&[CORE_2026_08_3], without_onpair(schemes))?; vortex_zstd::initialize(&session); session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; - let compressor = compressor_for_session( - &session, - BtrBlocksCompressorBuilder::default().with_compact(), - ); + let compressor = BtrBlocksCompressor::from_session(&session); golden_corpus_snapshots("compact", &compressor) } diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index 31734d6a60e..cd03d9f1524 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -21,7 +21,11 @@ use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session +}); /// Helper: synthetic short-string corpus that the cascading compressor should /// route through OnPair. @@ -58,7 +62,7 @@ fn nonnullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert a specific scheme — both OnPair and FSST are registered and @@ -101,7 +105,7 @@ fn nullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert OnPair specifically here — the sample-based selector may @@ -137,7 +141,7 @@ fn large_unique_short_strings_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); @@ -166,7 +170,7 @@ fn empty_and_short_string_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); let decoded = compressed @@ -211,7 +215,7 @@ fn delta_dict_offsets_roundtrip() { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) .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..0a0ac1ded2f 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -14,17 +14,36 @@ 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::CompressionSessionExt; +use vortex_btrblocks::DEFAULT_SCHEMES; use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::SchemeId; use vortex_btrblocks::schemes::binary::VarBinScheme; use vortex_btrblocks::schemes::string::OnPairScheme; use vortex_error::VortexResult; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session +}); const N: usize = 100_000; +/// The default schemes minus `excluded`. +fn default_without(excluded: SchemeId) -> BtrBlocksCompressor { + let session = vortex_array::array_session(); + for scheme in DEFAULT_SCHEMES + .iter() + .filter(|scheme| scheme.id() != excluded) + { + session.register_scheme(*scheme); + } + BtrBlocksCompressor::from_session_no_editions(&session) +} + 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 = BtrBlocksCompressor::from_session_no_editions(&SESSION); + 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-compressor/Cargo.toml b/vortex-compressor/Cargo.toml index 5977ac0227f..89bdf7a760e 100644 --- a/vortex-compressor/Cargo.toml +++ b/vortex-compressor/Cargo.toml @@ -22,15 +22,16 @@ rustc-hash = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } +vortex-session = { workspace = true } vortex-utils = { workspace = true } [dev-dependencies] divan = { workspace = true } mimalloc = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } -vortex-session = { workspace = true } [lints] workspace = true diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 55bb9b188f6..1c8f6452d26 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -64,6 +64,7 @@ pub mod builtins; pub mod scheme; +pub mod session; pub mod stats; mod compressor; diff --git a/vortex-compressor/src/session.rs b/vortex-compressor/src/session.rs new file mode 100644 index 00000000000..1e0bf028014 --- /dev/null +++ b/vortex-compressor/src/session.rs @@ -0,0 +1,133 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Session registry of compression schemes. +//! +//! Registering a scheme makes it available to compressors built from the session with +//! [`CascadingCompressor::from_session`](crate::CascadingCompressor::from_session). Whether a +//! registered scheme may write its encodings is decided by the session's enabled editions. + +use std::any::Any; + +use vortex_edition::ComponentKind; +use vortex_edition::EditionSessionExt; +use vortex_session::SessionExt; +use vortex_session::SessionGuard; +use vortex_session::SessionVar; +use vortex_utils::aliases::hash_set::HashSet; + +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; + +/// 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 crates in the same order compress identically. +#[derive(Clone, Debug, Default)] +pub struct CompressionSession { + /// Registered schemes in registration order. + schemes: Vec<&'static dyn Scheme>, +} + +impl CompressionSession { + /// Registers a scheme. + /// + /// Registering a [`SchemeId`](crate::scheme::SchemeId) that is already present is a no-op, so + /// initializers may run more than once. + 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 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); + } + + /// The registered compression schemes in registration order. + fn registered_schemes(&self) -> Vec<&'static dyn Scheme> { + self.compression().schemes().to_vec() + } + + /// The registered schemes whose serialized IDs the enabled editions all permit. + fn permitted_schemes(&self) -> Vec<&'static dyn Scheme> { + self.permit(self.registered_schemes()) + } + + /// Keeps the schemes in `schemes` whose serialized IDs the enabled editions all permit. + fn permit(&self, schemes: Vec<&'static dyn Scheme>) -> Vec<&'static dyn Scheme> { + let allowed: HashSet<_> = self + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + schemes + .into_iter() + .filter(|scheme| { + scheme + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + }) + .collect() + } +} + +impl CompressionSessionExt for S {} + +#[cfg(test)] +mod tests { + use vortex_array::array_session; + + use super::*; + use crate::builtins::FloatDictScheme; + use crate::builtins::IntDictScheme; + + fn ids(schemes: &[&'static dyn Scheme]) -> Vec { + schemes.iter().map(|scheme| scheme.id()).collect() + } + + #[test] + fn registration_keeps_order_and_is_idempotent() { + let session = array_session(); + assert!(session.registered_schemes().is_empty()); + session.register_scheme(&IntDictScheme); + session.register_scheme(&FloatDictScheme); + session.register_scheme(&IntDictScheme); + assert_eq!( + ids(&session.registered_schemes()), + vec![IntDictScheme.id(), FloatDictScheme.id()] + ); + } + + /// Without enabled editions no serialized ID is permitted, so nothing survives. + #[test] + fn no_editions_permit_nothing() { + let session = array_session(); + session.register_scheme(&IntDictScheme); + assert!(session.permitted_schemes().is_empty()); + } +} diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 7cc7d0322da..9b4d5416487 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -29,10 +29,15 @@ 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::compressor::CascadingCompressor; +use vortex::compressor::CompressionSessionExt; +use vortex::compressor::Scheme; +use vortex::compressor::SchemeExt; +use vortex::compressor::SchemeId; +use vortex::compressor::schemes; use vortex::dtype::DType; use vortex::dtype::FieldMask; -use vortex::editions::ComponentKind; use vortex::editions::Edition; use vortex::editions::EditionDeclaration; use vortex::editions::EditionFamily; @@ -553,28 +558,64 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { /// nonzero sets row blocks without outer dictionaries or byte coalescing, retaining per-block /// dictionary compression. 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 schemes = cuda_compatible_schemes(session); + let strategy = WriteStrategyBuilder::from_session(session) .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); if block_rows == 0 { - strategy.with_btrblocks_builder(builder).build() + strategy.with_schemes(schemes).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(CascadingCompressor::new(schemes))) + .with_probe_compressor(BtrBlocksCompressor(CascadingCompressor::new(Vec::new()))) .with_row_block_size(block_rows) .with_data_block_target_bytes(None) .build() } } +/// A compressor over [`cuda_compatible_schemes`]. +pub fn cuda_compressor(session: &VortexSession) -> BtrBlocksCompressor { + BtrBlocksCompressor(CascadingCompressor::new(cuda_compatible_schemes(session))) +} + +/// The schemes registered on `session` that CUDA kernels can decode, keeping FSST for string +/// compression and adding 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 the session's enabled editions decide which of the two survives. +/// +/// Files written with these schemes may be larger than with the default compressor: the list +/// picks encodings the GPU decodes, not the smallest ones. +pub fn cuda_compatible_schemes(session: &VortexSession) -> Vec<&'static dyn Scheme> { + // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other string + // fragmentation and dictionary schemes still require unsupported decode paths. + let excluded: Vec = vec![ + schemes::integer::SparseScheme.id(), + schemes::integer::IntRLEScheme.id(), + schemes::float::ALPRDScheme.id(), + schemes::float::FloatRLEScheme.id(), + schemes::float::NullDominatedSparseScheme.id(), + schemes::string::NullDominatedSparseScheme.id(), + schemes::string::StringDictScheme.id(), + schemes::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 + // until GPU delta decode is benchmarked against the schemes it would displace, since this + // list picks encodings rather than merely decoding them. + schemes::integer::DeltaScheme::default().id(), + ]; + let mut cuda: Vec<&'static dyn Scheme> = session + .registered_schemes() + .into_iter() + .filter(|scheme| !excluded.contains(&scheme.id())) + .collect(); + cuda.push(&schemes::binary::ZstdScheme); + cuda.push(&schemes::binary::ZstdBuffersScheme); + session.permit(cuda) +} + #[derive(Clone, Debug)] struct CudaLayoutRegistration(Arc); @@ -659,6 +700,7 @@ mod tests { use vortex::array::assert_arrays_eq; use vortex::buffer::ByteBufferMut; use vortex::editions::CORE_2025_05_0; + use vortex::editions::ComponentKind; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; use vortex::file::WriteOptionsSessionExt; 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..d855a55fde5 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::CompressionSessionExt; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; use vortex_file::OpenOptionsSessionExt; @@ -85,7 +86,8 @@ fn make_file(columns: usize, chunks: usize) -> VortexFile { .collect::>(); let array = ChunkedArray::from_iter(struct_chunks).into_array(); - let strategy = vortex_file::WriteStrategyBuilder::default() + let strategy = vortex_file::WriteStrategyBuilder::from_session(&SESSION) + .with_schemes(SESSION.registered_schemes()) .with_row_block_size(ROWS_PER_CHUNK) .with_data_block_target_bytes(None) .build(); @@ -143,7 +145,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(&SESSION) + .with_schemes(SESSION.registered_schemes()); for (c, (name, _)) in fields.iter().enumerate() { let field_strategy = RepartitionStrategy::new( ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 707de8bb47e..1f2eb1a3b0c 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -194,6 +194,8 @@ pub fn register_default_encodings(session: &VortexSession) { #[cfg(feature = "tensor")] vortex_tensor::initialize(session); + + vortex_btrblocks::initialize(session); } #[cfg(test)] diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index c8110fe88c3..e70f9cdc9d2 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -7,7 +7,10 @@ use std::num::NonZeroUsize; use std::sync::Arc; use vortex_array::dtype::FieldPath; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::CascadingCompressor; +use vortex_btrblocks::CompressionSessionExt; +use vortex_btrblocks::Scheme; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::integer::IntDictScheme; use vortex_error::VortexExpect; @@ -26,16 +29,17 @@ 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. + /// Schemes for the [`BtrBlocksCompressor`]s that [`WriteStrategyBuilder::build`] creates. /// `IntDictScheme` is automatically excluded from the data compressor to prevent recursive /// dictionary encoding. - BtrBlocks(BtrBlocksCompressorBuilder), + Schemes(Vec<&'static dyn Scheme>), /// An opaque compressor used as-is for both data and stats compression. Opaque(Arc), } @@ -64,12 +68,13 @@ 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::Schemes(session.permitted_schemes()), row_block_size: 8192, data_block_target_bytes: Some(ONE_MEG), field_writers: HashMap::new(), @@ -132,12 +137,12 @@ impl WriteStrategyBuilder { self } - /// Override the default [`BtrBlocksCompressorBuilder`] used for compression. + /// Override the compression schemes. /// - /// 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); + /// The strategy builds two compressors from them: one for data, without `IntDictScheme`, and + /// one for stats. The list is used as given; it is not filtered by the session's editions. + pub fn with_schemes(mut self, schemes: Vec<&'static dyn Scheme>) -> Self { + self.compressor = CompressorConfig::Schemes(schemes); self } @@ -177,12 +182,15 @@ impl WriteStrategyBuilder { // 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::Schemes(schemes) => { + Arc::new(BtrBlocksCompressor(CascadingCompressor::new( + schemes + .iter() + .copied() + .filter(|scheme| scheme.id() != IntDictScheme.id()) + .collect(), + ))) + } CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; let compressing = CompressingStrategy::new(buffered, data_compressor); @@ -206,7 +214,9 @@ 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::Schemes(schemes) => { + Arc::new(BtrBlocksCompressor(CascadingCompressor::new(schemes))) + } CompressorConfig::Opaque(compressor) => compressor, }; let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 640de874d2b..ec98e06cc1b 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -65,7 +65,9 @@ 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::CascadingCompressor; +use vortex_btrblocks::CompressionSessionExt; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::string::StringDictScheme; use vortex_buffer::Buffer; @@ -1874,7 +1876,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 +2255,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 +2580,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,14 +2588,17 @@ 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 no_string_dict: Vec<_> = SESSION + .permitted_schemes() + .into_iter() + .filter(|scheme| scheme.id() != StringDictScheme.id()) + .collect(); let mut buf = ByteBufferMut::empty(); let summary = SESSION .write_options() .with_strategy( - crate::strategy::WriteStrategyBuilder::default() - .with_btrblocks_builder(no_string_dict) + crate::strategy::WriteStrategyBuilder::from_session(&SESSION) + .with_schemes(no_string_dict) .build(), ) .write(&mut buf, strings.to_array_stream()) @@ -2614,15 +2619,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(CascadingCompressor::new( + SESSION + .permitted_schemes() + .into_iter() + .filter(|scheme| scheme.id() != StringDictScheme.id()) + .collect(), + )); 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..17711da378a 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::CompressionSessionExt; use vortex_buffer::ByteBuffer; use vortex_edition::ComponentKind; use vortex_edition::EditionSessionExt; @@ -247,15 +247,12 @@ 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), - ) + None if enforce_editions => WriteStrategyBuilder::from_session(&self.session).build(), + // With editions disabled every registered encoding may be written. + None => WriteStrategyBuilder::from_session(&self.session) + .with_schemes(self.session.registered_schemes()) .build(), }; let dtype = stream.dtype().clone(); 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..34aa95d4295 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -411,10 +411,12 @@ mod tests { // FIXME(ngates): Deprecate the global `runtime::single::block_on` helper and require tests // to call `block_on` on an explicit runtime instance. fn session_with_handle(handle: Handle) -> VortexSession { - array_session() + let session = array_session() .with::() .with::() - .with_handle(handle) + .with_handle(handle); + vortex_btrblocks::initialize(&session); + session } async fn write_dict_layout( @@ -426,7 +428,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::from_session_no_editions(session)), ); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -456,7 +458,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::from_session_no_editions(&session)), ); let array = VarBinArray::from_iter( @@ -559,7 +561,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::from_session_no_editions(&session)), ); let array = @@ -614,7 +616,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::from_session_no_editions(&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..034887e060f 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -16,9 +16,8 @@ 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::compressor::COMPACT_SCHEMES; +use vortex::compressor::CompressionSessionExt; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::file::WriteOptionsSessionExt; @@ -380,16 +379,18 @@ impl PyVortexWriteOptions { ) -> PyVortexResult<()> { let session = session(); py.detach(|| { - let allowed_encodings = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - let mut compressor = BtrBlocksCompressorBuilder::default(); + let mut strategy = WriteStrategyBuilder::from_session(session); if self.use_compact_encodings { - compressor = compressor.with_compact(); + strategy = strategy.with_schemes( + session.permit( + session + .registered_schemes() + .into_iter() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ), + ); } - 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()))? { 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..9b60ffc20dc 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::COMPACT_SCHEMES; + use vortex::compressor::CompressionSessionExt; use vortex::editions::CORE_2026_08_3; use vortex::editions::EditionSessionExt; use vortex::file::WriteStrategyBuilder; @@ -44,15 +45,23 @@ mod tests { let regular_bytes = adapter::write_compressed_to_bytes_with_session( &session, array.clone(), - WriteStrategyBuilder::default().build(), + WriteStrategyBuilder::from_session(&session) + .with_schemes(session.registered_schemes()) + .build(), )?; let _regular = adapter::read_file(regular_bytes)?; let compact_bytes = adapter::write_compressed_to_bytes_with_session( &session, array, - WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) + WriteStrategyBuilder::from_session(&session) + .with_schemes( + session + .registered_schemes() + .into_iter() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ) .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..6749ca5728f 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::COMPACT_SCHEMES; +use vortex::compressor::CompressionSessionExt; use vortex::file::WriteStrategyBuilder; +use vortex::session::VortexSession; use vortex_array::ExecutionCtx; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; @@ -137,13 +140,24 @@ 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 execution context's session registers no compression schemes, so build the + // strategy from the same default session the adapter writes with. + let session = VortexSession::default(); if self.compact { - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) + let strategy = WriteStrategyBuilder::from_session(&session) + .with_schemes( + session + .registered_schemes() + .into_iter() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ) .build(); adapter::write_compressed(&path, array, strategy)?; } else { - let strategy = WriteStrategyBuilder::default().build(); + let strategy = WriteStrategyBuilder::from_session(&session) + .with_schemes(session.registered_schemes()) + .build(); adapter::write_compressed(&path, array, strategy)?; } Ok(vec![FixtureEntry { diff --git a/vortex-tui/src/convert.rs b/vortex-tui/src/convert.rs index ab316982b27..3c08dc0e392 100644 --- a/vortex-tui/src/convert.rs +++ b/vortex-tui/src/convert.rs @@ -13,9 +13,8 @@ use parquet::arrow::ParquetRecordBatchStreamBuilder; use tokio::fs::File; use tokio::io::AsyncWriteExt; use vortex::array::stream::ArrayStreamAdapter; -use vortex::compressor::BtrBlocksCompressorBuilder; -use vortex::editions::ComponentKind; -use vortex::editions::EditionSessionExt; +use vortex::compressor::COMPACT_SCHEMES; +use vortex::compressor::CompressionSessionExt; use vortex::error::VortexExpect; use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; @@ -98,16 +97,18 @@ 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(); + let mut strategy = WriteStrategyBuilder::from_session(session); if matches!(flags.strategy, Strategy::Compact) { - compressor = compressor.with_compact(); + strategy = strategy.with_schemes( + session.permit( + session + .registered_schemes() + .into_iter() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ), + ); } - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); let mut file = File::create(output_path).await?; session diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 49fc0a44ddf..0bc6a351bd8 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -83,7 +83,7 @@ tokio = [ "vortex-io/tokio", "vortex-layout/tokio", ] -zstd = ["dep:vortex-zstd", "vortex-file?/zstd"] +zstd = ["dep:vortex-zstd", "vortex-btrblocks/pco", "vortex-btrblocks/zstd", "vortex-file?/zstd"] tensor = ["dep:vortex-tensor", "vortex-file?/tensor"] wasm-bindgen = [ "vortex-file?/wasm-bindgen", 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 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..3fee1c71617 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()); @@ -144,9 +144,17 @@ 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; + #[cfg(feature = "zstd")] + pub use vortex_btrblocks::COMPACT_SCHEMES; + pub use vortex_btrblocks::CascadingCompressor; + pub use vortex_btrblocks::CompressionSession; + pub use vortex_btrblocks::CompressionSessionExt; + pub use vortex_btrblocks::DEFAULT_SCHEMES; pub use vortex_btrblocks::Scheme; + pub use vortex_btrblocks::SchemeExt; pub use vortex_btrblocks::SchemeId; + pub use vortex_btrblocks::initialize; + pub use vortex_btrblocks::schemes; } /// Vortex editions: versioned sets of serialized components. @@ -366,7 +374,8 @@ mod test { use vortex_array::expr::select; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; - use vortex_btrblocks::BtrBlocksCompressorBuilder; + use vortex_btrblocks::COMPACT_SCHEMES; + use vortex_btrblocks::CompressionSessionExt; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_file::OpenOptionsSessionExt; @@ -422,7 +431,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(), )?; @@ -489,8 +498,16 @@ mod test { session .write_options() .with_strategy( - WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) + WriteStrategyBuilder::from_session(&session) + .with_schemes( + session.permit( + session + .registered_schemes() + .into_iter() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ), + ) .build(), ) .write( diff --git a/wasm-test/src/main.rs b/wasm-test/src/main.rs index 964d3a36c9a..647aa6baf15 100644 --- a/wasm-test/src/main.rs +++ b/wasm-test/src/main.rs @@ -7,6 +7,7 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::validity::Validity; use vortex::buffer::buffer; use vortex::compressor::BtrBlocksCompressor; +use vortex::compressor::initialize; use vortex::session::VortexSession; use vortex::VortexSessionDefault; @@ -17,7 +18,8 @@ pub fn main() { let array = PrimitiveArray::new(buffer![1i32; 1024], Validity::AllValid).into_array(); let session = VortexSession::default(); - let compressed = BtrBlocksCompressor::default() + initialize(&session); + let compressed = BtrBlocksCompressor::from_session(&session) .compress(&array, &mut session.create_execution_ctx()) .unwrap(); println!("Compressed size: {}", compressed.len()); From d18f5433c040e0feededb2fe41e1863345fb5db5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 19:44:04 -0400 Subject: [PATCH 02/11] Default the scheme registry, drop initialize, configure CUDA sessions Signed-off-by: Matt Katz --- Cargo.lock | 1 - benchmarks/compress-bench/README.md | 2 +- benchmarks/compress-bench/src/gpu/vortex.rs | 41 +++- vortex-btrblocks/Cargo.toml | 2 +- vortex-btrblocks/benches/compress.rs | 6 +- vortex-btrblocks/benches/compress_listview.rs | 6 +- vortex-btrblocks/src/canonical_compressor.rs | 1 - vortex-btrblocks/src/lib.rs | 104 +------- .../schemes/float/scheme_selection_tests.rs | 6 +- vortex-btrblocks/src/schemes/float/tests.rs | 6 +- .../schemes/integer/scheme_selection_tests.rs | 7 +- vortex-btrblocks/src/schemes/integer/tests.rs | 6 +- .../schemes/string/scheme_selection_tests.rs | 9 +- vortex-btrblocks/src/schemes/string/tests.rs | 6 +- vortex-btrblocks/src/session.rs | 228 ++++++++++++++++++ vortex-btrblocks/src/tests.rs | 11 +- vortex-btrblocks/src/trace_tests.rs | 1 - vortex-btrblocks/tests/golden.rs | 5 +- vortex-btrblocks/tests/onpair_roundtrip.rs | 6 +- vortex-btrblocks/tests/varbin_scheme.rs | 9 +- vortex-compressor/Cargo.toml | 3 +- vortex-compressor/src/lib.rs | 1 - vortex-compressor/src/session.rs | 133 ---------- vortex-cuda/ffi/cinclude/vortex_cuda.h | 3 +- vortex-cuda/ffi/src/lib.rs | 7 +- vortex-cuda/gpu-scan-cli/src/main.rs | 3 + vortex-cuda/src/layout.rs | 56 ++--- vortex-file/src/lib.rs | 2 - vortex-layout/src/layouts/dict/reader.rs | 6 +- vortex/src/lib.rs | 3 +- wasm-test/src/main.rs | 2 - 31 files changed, 336 insertions(+), 346 deletions(-) create mode 100644 vortex-btrblocks/src/session.rs delete mode 100644 vortex-compressor/src/session.rs diff --git a/Cargo.lock b/Cargo.lock index e8969b748f2..041f9a9e2f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10923,7 +10923,6 @@ dependencies = [ "tracing", "vortex-array", "vortex-buffer", - "vortex-edition", "vortex-error", "vortex-mask", "vortex-session", diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 541c5382537..44cd3c36b8d 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 - (`cuda_compatible_schemes`) and a CUDA flat layout, then decoded on the device all the way to + (`use_cuda_schemes`) 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 b9214df5294..380283d5996 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,21 +17,24 @@ 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::BtrBlocksCompressor; 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; @@ -44,11 +48,21 @@ use vortex_cuda::CudaSession; use vortex_cuda::PooledFileReadAtOptions; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::CudaFlatLayoutStrategy; -use vortex_cuda::layout::cuda_compressor; use vortex_cuda::layout::register_cuda_layout; +use vortex_cuda::layout::use_cuda_schemes; 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); + use_cuda_schemes(&session); + session +}); + /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. pub struct GpuVortexCompressor { verify: bool, @@ -90,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?; @@ -99,10 +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(), - cuda_compressor(&SESSION), + BtrBlocksCompressor::from_session(&GPU_SESSION), ))); let start = Instant::now(); - SESSION + GPU_SESSION .write_options() .with_strategy(strategy) .write(&mut output, array.to_array_stream()) @@ -130,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 @@ -165,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 { @@ -192,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 @@ -209,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)) @@ -298,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/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 9c499b81eaa..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 } @@ -50,7 +51,6 @@ 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 } [features] diff --git a/vortex-btrblocks/benches/compress.rs b/vortex-btrblocks/benches/compress.rs index 5f8801ddfbb..bc5c00987d2 100644 --- a/vortex-btrblocks/benches/compress.rs +++ b/vortex-btrblocks/benches/compress.rs @@ -24,11 +24,7 @@ mod benchmarks { use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; - static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_btrblocks::initialize(&session); - session - }); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn make_clickbench_window_name() -> ArrayRef { // A test that's meant to mirror the WindowName column from ClickBench. diff --git a/vortex-btrblocks/benches/compress_listview.rs b/vortex-btrblocks/benches/compress_listview.rs index aed620db9fc..b6d01c86dd2 100644 --- a/vortex-btrblocks/benches/compress_listview.rs +++ b/vortex-btrblocks/benches/compress_listview.rs @@ -31,11 +31,7 @@ mod benchmarks { const NUM_ROWS: usize = 8192; const SEED: u64 = 42; - static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_btrblocks::initialize(&session); - session - }); + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); const SHORT_STRINGS: &[&str] = &[ "alpha_one", diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index d131feb0c2c..b57515342b1 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -26,7 +26,6 @@ use crate::CompressionSessionExt; /// use vortex_btrblocks::BtrBlocksCompressor; /// /// let session = vortex_array::array_session(); -/// vortex_btrblocks::initialize(&session); /// /// // Every registered scheme; this session enables no editions. /// let compressor = BtrBlocksCompressor::from_session_no_editions(&session); diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 9cb8e9aae65..fa4a44099a1 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -30,10 +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). Schemes are registered on a session: [`initialize`] registers [`DEFAULT_SCHEMES`], -//! and [`BtrBlocksCompressor::from_session`] compresses with the registered schemes whose -//! serialized IDs the session's enabled editions permit ([`from_session_no_editions`] ignores -//! them). +//! sample). The schemes available to a compressor are those registered on its session's +//! [`CompressionSession`], which starts with [`DEFAULT_SCHEMES`]. +//! [`BtrBlocksCompressor::from_session`] keeps the registered schemes whose serialized IDs the +//! session's enabled editions permit; [`from_session_no_editions`] keeps them all. //! //! [`from_session_no_editions`]: BtrBlocksCompressor::from_session_no_editions //! @@ -52,7 +52,6 @@ //! //! # fn example() -> vortex_error::VortexResult<()> { //! let session = array_session(); -//! vortex_btrblocks::initialize(&session); //! let array = PrimitiveArray::new(buffer![42u64; 1024], Validity::NonNullable).into_array(); //! //! // In memory, with no editions to respect, compress with every registered scheme. @@ -68,6 +67,8 @@ mod canonical_compressor; /// Compression scheme implementations. pub mod schemes; +/// Session registry of compression schemes. +pub mod session; #[cfg(test)] mod tests; #[cfg(test)] @@ -78,102 +79,21 @@ mod trace_tests; // Btrblocks-specific exports. pub use canonical_compressor::BtrBlocksCompressor; pub use schemes::patches::compress_patches; +#[cfg(feature = "zstd")] +pub use session::COMPACT_SCHEMES; +pub use session::CompressionSession; +pub use session::CompressionSessionExt; +pub use session::DEFAULT_SCHEMES; +pub use session::DELTA_SCHEME; pub use vortex_compressor::CascadingCompressor; pub use vortex_compressor::scheme::CompressorContext; pub use vortex_compressor::scheme::MAX_CASCADE; pub use vortex_compressor::scheme::Scheme; pub use vortex_compressor::scheme::SchemeExt; pub use vortex_compressor::scheme::SchemeId; -pub use vortex_compressor::session::CompressionSession; -pub use vortex_compressor::session::CompressionSessionExt; pub use vortex_compressor::stats::ArrayAndStats; pub use vortex_compressor::stats::BoolStats; pub use vortex_compressor::stats::FloatStats; pub use vortex_compressor::stats::GenerateStatsOptions; pub use vortex_compressor::stats::IntegerStats; pub use vortex_compressor::stats::StringStats; -use vortex_session::VortexSession; - -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. -/// -/// This list is order-sensitive: [`initialize`] registers it in this order and the compressor -/// preserves registration order, so that tie-breaking is deterministic. -pub 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, -]; - -/// Compact schemes (Zstd for strings and binary, Pco for numerics when the `pco` feature is on). -/// -/// Not part of [`DEFAULT_SCHEMES`]: they trade decode speed for compression ratio, so callers add -/// them to a compressor's scheme list explicitly. -#[cfg(feature = "zstd")] -pub const COMPACT_SCHEMES: &[&dyn Scheme] = &[ - &string::ZstdScheme, - &binary::ZstdScheme, - #[cfg(feature = "pco")] - &integer::PcoScheme, - #[cfg(feature = "pco")] - &float::PcoScheme, -]; - -/// Delta, kept out of [`DEFAULT_SCHEMES`] because it is slower to decompress than the schemes that -/// would otherwise win. Callers that want it add it to their scheme list and permit -/// `fastlanes.delta`. -/// -/// TODO(robert): Return it to [`DEFAULT_SCHEMES`] once we have scheme filtering. -pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); - -/// Registers [`DEFAULT_SCHEMES`] on `session`, in order. -/// -/// Registration is idempotent, so this may run more than once. -pub fn initialize(session: &VortexSession) { - for scheme in DEFAULT_SCHEMES { - session.register_scheme(*scheme); - } -} diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index 96a186adc41..dbea853803f 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -21,11 +21,7 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_constant_compressed() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index e21b031ad70..7f1888e92b0 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -22,11 +22,7 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::float::FloatRLEScheme; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index d1bbbfe4e4d..1e8ae27e454 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -30,16 +30,11 @@ use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; use crate::CompressionSessionExt; use crate::DELTA_SCHEME; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); /// The default schemes plus opt-in Delta. fn with_delta() -> BtrBlocksCompressor { let session = vortex_array::array_session(); - crate::initialize(&session); session.register_scheme(&DELTA_SCHEME); BtrBlocksCompressor::from_session_no_editions(&session) } diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index b331dece461..f641d2335ad 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -27,11 +27,7 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index e0e72a1e4dc..22b052fcbcd 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -17,14 +17,11 @@ use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::CompressionSession; use crate::CompressionSessionExt; use crate::DEFAULT_SCHEMES; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_constant_compressed() -> VortexResult<()> { @@ -111,7 +108,7 @@ 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 session = vortex_array::array_session(); + let session = vortex_array::array_session().with_some(CompressionSession::empty()); session.register_scheme(&FSSTScheme); let compressor = BtrBlocksCompressor::from_session_no_editions(&session); let compressed = compressor.compress(&array_ref, &mut SESSION.create_execution_ctx())?; diff --git a/vortex-btrblocks/src/schemes/string/tests.rs b/vortex-btrblocks/src/schemes/string/tests.rs index e76a46b36fb..57398bb1e3b 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -16,11 +16,7 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_strings() -> VortexResult<()> { diff --git a/vortex-btrblocks/src/session.rs b/vortex-btrblocks/src/session.rs new file mode 100644 index 00000000000..fba52af561d --- /dev/null +++ b/vortex-btrblocks/src/session.rs @@ -0,0 +1,228 @@ +// 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 [`DEFAULT_SCHEMES`]. Whether a registered scheme may write its encodings is +//! decided by the session's enabled editions. + +use std::any::Any; + +use vortex_edition::ComponentKind; +use vortex_edition::EditionSessionExt; +use vortex_session::SessionExt; +use vortex_session::SessionGuard; +use vortex_session::SessionVar; +use vortex_utils::aliases::hash_set::HashSet; + +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. +pub 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, +]; + +/// Compact schemes (Zstd for strings and binary, Pco for numerics when the `pco` feature is on). +/// +/// Not part of [`DEFAULT_SCHEMES`]: they trade decode speed for compression ratio, so callers add +/// them to a compressor's scheme list explicitly. +#[cfg(feature = "zstd")] +pub const COMPACT_SCHEMES: &[&dyn Scheme] = &[ + &string::ZstdScheme, + &binary::ZstdScheme, + #[cfg(feature = "pco")] + &integer::PcoScheme, + #[cfg(feature = "pco")] + &float::PcoScheme, +]; + +/// Delta, kept out of [`DEFAULT_SCHEMES`] because it is slower to decompress than the schemes that +/// would otherwise win. Callers that want it add it to their scheme list and permit +/// `fastlanes.delta`. +/// +/// TODO(robert): Return it to [`DEFAULT_SCHEMES`] 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 [`DEFAULT_SCHEMES`]; [`empty`](Self::empty) registers 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(), + } + } + + /// 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); + } + + /// The registered compression schemes in registration order. + fn registered_schemes(&self) -> Vec<&'static dyn Scheme> { + self.compression().schemes().to_vec() + } + + /// The registered schemes whose serialized IDs the enabled editions all permit. + fn permitted_schemes(&self) -> Vec<&'static dyn Scheme> { + self.permit(self.registered_schemes()) + } + + /// Keeps the schemes in `schemes` whose serialized IDs the enabled editions all permit. + fn permit(&self, schemes: Vec<&'static dyn Scheme>) -> Vec<&'static dyn Scheme> { + let allowed: HashSet<_> = self + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + schemes + .into_iter() + .filter(|scheme| { + scheme + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + }) + .collect() + } +} + +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.registered_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.registered_schemes()), + vec![IntDictScheme.id(), FloatDictScheme.id()] + ); + } + + /// Without enabled editions no serialized ID is permitted, so nothing survives. + #[test] + fn no_editions_permit_nothing() { + let session = array_session(); + assert!(session.permitted_schemes().is_empty()); + } +} diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index b3225c0e789..f64ef7f9900 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -31,17 +31,15 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; #[cfg(feature = "zstd")] use crate::COMPACT_SCHEMES; +#[cfg(feature = "zstd")] +use crate::CompressionSession; use crate::CompressionSessionExt; #[cfg(feature = "zstd")] use crate::Scheme; #[cfg(feature = "zstd")] use crate::schemes::binary; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - crate::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[rstest] #[case::zctl( @@ -208,7 +206,6 @@ fn test_compact_binary_zstd_compressed() -> VortexResult<()> { ); let session = vortex_array::array_session(); - crate::initialize(&session); for scheme in COMPACT_SCHEMES { session.register_scheme(*scheme); } @@ -248,7 +245,7 @@ fn test_binary_zstd_scheme_encoding( DType::Binary(Nullability::NonNullable), ); - let session = vortex_array::array_session(); + let session = vortex_array::array_session().with_some(CompressionSession::empty()); session.register_scheme(scheme); let compressor = BtrBlocksCompressor::from_session_no_editions(&session); let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 62c4720d6e2..610fc81bb42 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -129,7 +129,6 @@ fn lineitem() -> VortexResult { /// Delta is opt-in, and these traces cover the delta-encoded FSST offsets, so enable it here. fn compressed_lineitem() -> VortexResult { let session = trace_session(); - crate::initialize(&session); session.register_scheme(&DELTA_SCHEME); BtrBlocksCompressor::from_session_no_editions(&session) .compress(&lineitem()?, &mut execution_ctx()) diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index 3fc43a81877..b656a657556 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -51,6 +51,7 @@ use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; #[cfg(all(feature = "zstd", feature = "pco"))] use vortex_btrblocks::COMPACT_SCHEMES; +use vortex_btrblocks::CompressionSession; use vortex_btrblocks::CompressionSessionExt; use vortex_btrblocks::DEFAULT_SCHEMES; use vortex_btrblocks::Scheme; @@ -410,7 +411,9 @@ fn edition_session( editions: &[EditionId], schemes: Vec<&'static dyn Scheme>, ) -> VortexResult { - let session = vortex_array::array_session().with::(); + let session = vortex_array::array_session() + .with_some(CompressionSession::empty()) + .with::(); for scheme in schemes { session.register_scheme(scheme); } diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index cd03d9f1524..45d6d67c64b 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -21,11 +21,7 @@ use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_btrblocks::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); /// Helper: synthetic short-string corpus that the cascading compressor should /// route through OnPair. diff --git a/vortex-btrblocks/tests/varbin_scheme.rs b/vortex-btrblocks/tests/varbin_scheme.rs index 0a0ac1ded2f..b312a06a9f7 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -15,6 +15,7 @@ use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::CompressionSession; use vortex_btrblocks::CompressionSessionExt; use vortex_btrblocks::DEFAULT_SCHEMES; use vortex_btrblocks::SchemeExt; @@ -24,17 +25,13 @@ use vortex_btrblocks::schemes::string::OnPairScheme; use vortex_error::VortexResult; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(|| { - let session = vortex_array::array_session(); - vortex_btrblocks::initialize(&session); - session -}); +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); const N: usize = 100_000; /// The default schemes minus `excluded`. fn default_without(excluded: SchemeId) -> BtrBlocksCompressor { - let session = vortex_array::array_session(); + let session = vortex_array::array_session().with_some(CompressionSession::empty()); for scheme in DEFAULT_SCHEMES .iter() .filter(|scheme| scheme.id() != excluded) diff --git a/vortex-compressor/Cargo.toml b/vortex-compressor/Cargo.toml index 89bdf7a760e..5977ac0227f 100644 --- a/vortex-compressor/Cargo.toml +++ b/vortex-compressor/Cargo.toml @@ -22,16 +22,15 @@ rustc-hash = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } vortex-array = { workspace = true } vortex-buffer = { workspace = true } -vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } -vortex-session = { workspace = true } vortex-utils = { workspace = true } [dev-dependencies] divan = { workspace = true } mimalloc = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-session = { workspace = true } [lints] workspace = true diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 1c8f6452d26..55bb9b188f6 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -64,7 +64,6 @@ pub mod builtins; pub mod scheme; -pub mod session; pub mod stats; mod compressor; diff --git a/vortex-compressor/src/session.rs b/vortex-compressor/src/session.rs deleted file mode 100644 index 1e0bf028014..00000000000 --- a/vortex-compressor/src/session.rs +++ /dev/null @@ -1,133 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Session registry of compression schemes. -//! -//! Registering a scheme makes it available to compressors built from the session with -//! [`CascadingCompressor::from_session`](crate::CascadingCompressor::from_session). Whether a -//! registered scheme may write its encodings is decided by the session's enabled editions. - -use std::any::Any; - -use vortex_edition::ComponentKind; -use vortex_edition::EditionSessionExt; -use vortex_session::SessionExt; -use vortex_session::SessionGuard; -use vortex_session::SessionVar; -use vortex_utils::aliases::hash_set::HashSet; - -use crate::scheme::Scheme; -use crate::scheme::SchemeExt; - -/// 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 crates in the same order compress identically. -#[derive(Clone, Debug, Default)] -pub struct CompressionSession { - /// Registered schemes in registration order. - schemes: Vec<&'static dyn Scheme>, -} - -impl CompressionSession { - /// Registers a scheme. - /// - /// Registering a [`SchemeId`](crate::scheme::SchemeId) that is already present is a no-op, so - /// initializers may run more than once. - 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 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); - } - - /// The registered compression schemes in registration order. - fn registered_schemes(&self) -> Vec<&'static dyn Scheme> { - self.compression().schemes().to_vec() - } - - /// The registered schemes whose serialized IDs the enabled editions all permit. - fn permitted_schemes(&self) -> Vec<&'static dyn Scheme> { - self.permit(self.registered_schemes()) - } - - /// Keeps the schemes in `schemes` whose serialized IDs the enabled editions all permit. - fn permit(&self, schemes: Vec<&'static dyn Scheme>) -> Vec<&'static dyn Scheme> { - let allowed: HashSet<_> = self - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - schemes - .into_iter() - .filter(|scheme| { - scheme - .produced_encodings() - .iter() - .all(|id| allowed.contains(id)) - }) - .collect() - } -} - -impl CompressionSessionExt for S {} - -#[cfg(test)] -mod tests { - use vortex_array::array_session; - - use super::*; - use crate::builtins::FloatDictScheme; - use crate::builtins::IntDictScheme; - - fn ids(schemes: &[&'static dyn Scheme]) -> Vec { - schemes.iter().map(|scheme| scheme.id()).collect() - } - - #[test] - fn registration_keeps_order_and_is_idempotent() { - let session = array_session(); - assert!(session.registered_schemes().is_empty()); - session.register_scheme(&IntDictScheme); - session.register_scheme(&FloatDictScheme); - session.register_scheme(&IntDictScheme); - assert_eq!( - ids(&session.registered_schemes()), - vec![IntDictScheme.id(), FloatDictScheme.id()] - ); - } - - /// Without enabled editions no serialized ID is permitted, so nothing survives. - #[test] - fn no_editions_permit_nothing() { - let session = array_session(); - session.register_scheme(&IntDictScheme); - assert!(session.permitted_schemes().is_empty()); - } -} diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 2fb8ae8054e..970f628bd4e 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. Opening a sink restricts the session's + * compression schemes to those the GPU decodes, as `vx_cuda_session_new` already does. * * # Safety * diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 0647b2936a7..6bf843dc8d1 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -40,6 +40,7 @@ use vortex_cuda::arrow::DeviceArrayExt; use vortex_cuda::arrow::DeviceArrayStreamExt; use vortex_cuda::layout::cuda_write_strategy; use vortex_cuda::layout::register_cuda_layout; +use vortex_cuda::layout::use_cuda_schemes; use vortex_ffi::ffi_runtime; use vortex_ffi::try_or; use vortex_ffi::vx_array; @@ -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); + use_cuda_schemes(&session); 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. Opening a sink restricts the session's +/// compression schemes to those the GPU decodes, as `vx_cuda_session_new` already does. /// /// # Safety /// @@ -145,6 +148,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( try_or(error_out, ptr::null_mut(), || { // SAFETY: The caller supplies a live borrowed session handle. let vortex_session = session_with_cuda(unsafe { vx_session_ref(session) }?); + use_cuda_schemes(vortex_session); // SAFETY: All borrowed inputs satisfy the underlying sink's requirements. unsafe { vx_array_sink_open_file_with_strategy( @@ -1033,6 +1037,7 @@ mod tests { fn test_projection_gpu_values_and_validity() -> VortexResult<()> { let session = session().with_some(CudaSession::try_default()?); register_cuda_layout(&session); + use_cuda_schemes(&session); 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..04cff42a5e2 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -36,6 +36,7 @@ use vortex_cuda::TracingLaunchStrategy; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::cuda_write_strategy; use vortex_cuda::layout::register_cuda_layout; +use vortex_cuda::layout::use_cuda_schemes; use vortex_cuda_macros::cuda_available; use vortex_cuda_macros::cuda_not_available; @@ -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); + use_cuda_schemes(&session); 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); + use_cuda_schemes(&session); 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 9b4d5416487..91eb32b4268 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -31,10 +31,9 @@ use vortex::buffer::BufferString; use vortex::buffer::ByteBuffer; use vortex::compressor::BtrBlocksCompressor; use vortex::compressor::CascadingCompressor; -use vortex::compressor::CompressionSessionExt; -use vortex::compressor::Scheme; +use vortex::compressor::CompressionSession; +use vortex::compressor::DEFAULT_SCHEMES; use vortex::compressor::SchemeExt; -use vortex::compressor::SchemeId; use vortex::compressor::schemes; use vortex::dtype::DType; use vortex::dtype::FieldMask; @@ -552,21 +551,21 @@ 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; +/// Requires [`register_cuda_layout`], and [`use_cuda_schemes`] for the file to use only +/// encodings the GPU decodes. Zero `block_rows` uses default sizing and dictionary policy; /// nonzero sets row blocks without outer dictionaries or byte coalescing, retaining per-block /// dictionary compression. pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc { - let schemes = cuda_compatible_schemes(session); let strategy = WriteStrategyBuilder::from_session(session) .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); if block_rows == 0 { - strategy.with_schemes(schemes).build() + strategy.build() } else { // An opaque compressor keeps IntDict; disabling the probe avoids u16-sized outer blocks. strategy - .with_compressor(BtrBlocksCompressor(CascadingCompressor::new(schemes))) + .with_compressor(BtrBlocksCompressor::from_session(session)) .with_probe_compressor(BtrBlocksCompressor(CascadingCompressor::new(Vec::new()))) .with_row_block_size(block_rows) .with_data_block_target_bytes(None) @@ -574,24 +573,21 @@ pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(cuda_compatible_schemes(session))) -} - -/// The schemes registered on `session` that CUDA kernels can decode, keeping FSST for string -/// compression and adding Zstd for binary compression. +/// Replace the compression schemes registered on `session` with those CUDA kernels decode: the +/// defaults minus the schemes the GPU cannot decode, keeping FSST for string compression, plus +/// 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 the session's enabled editions decide which of the two survives. +/// opt-in `zstd` edition, so the session's enabled editions decide which of the two a writer uses. /// -/// Files written with these schemes may be larger than with the default compressor: the list -/// picks encodings the GPU decodes, not the smallest ones. -pub fn cuda_compatible_schemes(session: &VortexSession) -> Vec<&'static dyn Scheme> { +/// Call it, alongside [`register_cuda_layout`], on sessions that write CUDA-readable files; +/// sessions that only read them need not. Files written from such a session may be larger than +/// with the default schemes: the set picks encodings the GPU decodes, not the smallest ones. +pub fn use_cuda_schemes(session: &VortexSession) { // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other string // fragmentation and dictionary schemes still require unsupported decode paths. - let excluded: Vec = vec![ + let excluded = [ schemes::integer::SparseScheme.id(), schemes::integer::IntRLEScheme.id(), schemes::float::ALPRDScheme.id(), @@ -603,17 +599,19 @@ pub fn cuda_compatible_schemes(session: &VortexSession) -> Vec<&'static dyn Sche // 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 // until GPU delta decode is benchmarked against the schemes it would displace, since this - // list picks encodings rather than merely decoding them. + // set picks encodings rather than merely decoding them. schemes::integer::DeltaScheme::default().id(), ]; - let mut cuda: Vec<&'static dyn Scheme> = session - .registered_schemes() - .into_iter() + let mut registry = CompressionSession::empty(); + for scheme in DEFAULT_SCHEMES + .iter() .filter(|scheme| !excluded.contains(&scheme.id())) - .collect(); - cuda.push(&schemes::binary::ZstdScheme); - cuda.push(&schemes::binary::ZstdBuffersScheme); - session.permit(cuda) + { + registry.register(*scheme); + } + registry.register(&schemes::binary::ZstdScheme); + registry.register(&schemes::binary::ZstdBuffersScheme); + session.register(registry); } #[derive(Clone, Debug)] @@ -756,6 +754,7 @@ mod tests { let runtime = CurrentThreadRuntime::new(); let session = VortexSession::default().with_handle(runtime.handle()); register_cuda_layout(&session); + use_cuda_schemes(&session); runtime.block_on(async { let input = repeated_ids(8, 2 * block_rows + 137)?; let file = write_file(&session, input.clone(), block_rows).await?; @@ -790,6 +789,7 @@ mod tests { let runtime = CurrentThreadRuntime::new(); let session = VortexSession::default().with_handle(runtime.handle()); register_cuda_layout(&session); + use_cuda_schemes(&session); runtime.block_on(async { // Exceed u16 cardinality while remaining eligible for outer dictionaries. let block_rows = 70_000 * 8; diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 1f2eb1a3b0c..707de8bb47e 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -194,8 +194,6 @@ pub fn register_default_encodings(session: &VortexSession) { #[cfg(feature = "tensor")] vortex_tensor::initialize(session); - - vortex_btrblocks::initialize(session); } #[cfg(test)] diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 34aa95d4295..27080e35cfc 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -411,12 +411,10 @@ mod tests { // FIXME(ngates): Deprecate the global `runtime::single::block_on` helper and require tests // to call `block_on` on an explicit runtime instance. fn session_with_handle(handle: Handle) -> VortexSession { - let session = array_session() + array_session() .with::() .with::() - .with_handle(handle); - vortex_btrblocks::initialize(&session); - session + .with_handle(handle) } async fn write_dict_layout( diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 3fee1c71617..53e1bb1a40a 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -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; @@ -153,7 +154,6 @@ pub mod compressor { pub use vortex_btrblocks::Scheme; pub use vortex_btrblocks::SchemeExt; pub use vortex_btrblocks::SchemeId; - pub use vortex_btrblocks::initialize; pub use vortex_btrblocks::schemes; } @@ -328,6 +328,7 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::() + .with::() .with::() .with::() .with::() diff --git a/wasm-test/src/main.rs b/wasm-test/src/main.rs index 647aa6baf15..508ae4e24b9 100644 --- a/wasm-test/src/main.rs +++ b/wasm-test/src/main.rs @@ -7,7 +7,6 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::validity::Validity; use vortex::buffer::buffer; use vortex::compressor::BtrBlocksCompressor; -use vortex::compressor::initialize; use vortex::session::VortexSession; use vortex::VortexSessionDefault; @@ -18,7 +17,6 @@ pub fn main() { let array = PrimitiveArray::new(buffer![1i32; 1024], Validity::AllValid).into_array(); let session = VortexSession::default(); - initialize(&session); let compressed = BtrBlocksCompressor::from_session(&session) .compress(&array, &mut session.create_execution_ctx()) .unwrap(); From 17805da49cd818f50eb1d29a6e5cd7edd01ec771 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 19:58:52 -0400 Subject: [PATCH 03/11] Remove with_schemes: compact and no-editions writes configure the session Signed-off-by: Matt Katz --- benchmarks/string-bench/src/serialized.rs | 21 +++-- vortex-bench/src/conversions.rs | 13 ++- vortex-bench/src/datasets/taxi_data.rs | 2 +- vortex-bench/src/lib.rs | 40 ++++------ vortex-bench/src/tpch/tpchgen.rs | 3 +- vortex-file/benches/split_collection.rs | 7 +- vortex-file/src/strategy.rs | 18 ++--- vortex-file/src/tests.rs | 26 +++--- vortex-file/src/writer.rs | 5 +- vortex-python/src/io.rs | 25 ++---- vortex-python/src/session.rs | 80 +++++++++++++------ .../src/fixtures/arrays/datasets/mod.rs | 21 ++--- vortex-test/compat-gen/src/fixtures/mod.rs | 24 ++---- vortex-tui/src/convert.rs | 27 +++---- vortex/examples/tracing_vortex.rs | 20 ++--- vortex/src/lib.rs | 17 +--- 16 files changed, 165 insertions(+), 184 deletions(-) diff --git a/benchmarks/string-bench/src/serialized.rs b/benchmarks/string-bench/src/serialized.rs index 825746ee1be..95991c250c4 100644 --- a/benchmarks/string-bench/src/serialized.rs +++ b/benchmarks/string-bench/src/serialized.rs @@ -25,17 +25,19 @@ use anyhow::Result; use anyhow::bail; use bytes::Bytes; use futures::TryStreamExt; +use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::VarBinViewArray; +use vortex::compressor::CompressionSession; use vortex::compressor::CompressionSessionExt; -use vortex::compressor::Scheme; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; +use vortex::io::session::RuntimeSessionExt; use vortex::layout::LayoutStrategy; use vortex::session::VortexSession; use vortex_bench::Format; @@ -181,14 +183,19 @@ fn serialized_write_strategy( .filter(|&id| id != forced) .chain([DeltaScheme::default().id()]) .collect(); - let schemes: Vec<&'static dyn Scheme> = session - .permitted_schemes() + // A session like `session` registering only the remaining schemes; its editions still decide + // which of them may write. + let mut registry = CompressionSession::empty(); + for scheme in session + .registered_schemes() .into_iter() .filter(|scheme| !excluded.contains(&scheme.id())) - .collect(); - WriteStrategyBuilder::from_session(session) - .with_schemes(schemes) - .build() + { + registry.register(scheme); + } + let forced_session = VortexSession::default().with_handle(session.handle()); + forced_session.register(registry); + WriteStrategyBuilder::from_session(&forced_session).build() } /// Write one canonical string column to an in-memory Vortex file, forcing the diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 3892f2ad3fd..49ea7dd15a9 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -66,7 +66,6 @@ use wkb::writer::write_geometry; use crate::CompactionStrategy; use crate::Format; use crate::SESSION; -use crate::compact_schemes; use crate::utils::file::idempotent_async; /// Memory budget per concurrent conversion stream in GB. This is somewhat arbitary. @@ -243,17 +242,15 @@ 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::from_session(&SESSION); - if matches!(compaction, CompactionStrategy::Compact) { - builder = builder.with_schemes(compact_schemes()); - } + 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. @@ -340,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 18594fdd03e..2273db8ef99 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -30,11 +30,8 @@ pub use utils::file::*; pub use utils::logging::*; use vortex::compressor::COMPACT_SCHEMES; use vortex::compressor::CompressionSessionExt; -use vortex::compressor::Scheme; 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; @@ -80,11 +77,22 @@ 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(); + for scheme in COMPACT_SCHEMES { + session.register_scheme(*scheme); + } + 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,29 +259,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::from_session(&SESSION) - .with_schemes(compact_schemes()) - .build(), - ), - CompactionStrategy::Default => options, + CompactionStrategy::Compact => &COMPACT_SESSION, + CompactionStrategy::Default => &SESSION, } } } -/// The schemes [`SESSION`] permits plus the compact ones, for [`CompactionStrategy::Compact`]. -pub fn compact_schemes() -> Vec<&'static dyn Scheme> { - SESSION.permit( - SESSION - .registered_schemes() - .into_iter() - .chain(COMPACT_SCHEMES.iter().copied()) - .collect(), - ) -} - /// 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-file/benches/split_collection.rs b/vortex-file/benches/split_collection.rs index d855a55fde5..1af9421f5c9 100644 --- a/vortex-file/benches/split_collection.rs +++ b/vortex-file/benches/split_collection.rs @@ -22,7 +22,6 @@ use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::Field; use vortex_array::dtype::FieldMask; -use vortex_btrblocks::CompressionSessionExt; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; use vortex_file::OpenOptionsSessionExt; @@ -86,8 +85,7 @@ fn make_file(columns: usize, chunks: usize) -> VortexFile { .collect::>(); let array = ChunkedArray::from_iter(struct_chunks).into_array(); - let strategy = vortex_file::WriteStrategyBuilder::from_session(&SESSION) - .with_schemes(SESSION.registered_schemes()) + let strategy = vortex_file::WriteStrategyBuilder::from_session_no_editions(&SESSION) .with_row_block_size(ROWS_PER_CHUNK) .with_data_block_target_bytes(None) .build(); @@ -145,8 +143,7 @@ fn make_misaligned_file(columns: usize, chunks: usize) -> VortexFile { .unwrap() .into_array(); - let mut strategy = vortex_file::WriteStrategyBuilder::from_session(&SESSION) - .with_schemes(SESSION.registered_schemes()); + let mut strategy = vortex_file::WriteStrategyBuilder::from_session_no_editions(&SESSION); 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 e70f9cdc9d2..2a8a4d2dd7d 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -83,6 +83,15 @@ impl WriteStrategyBuilder { use_list_layout: use_experimental_list_layout(), } } + + /// Create a new builder whose compressor uses every scheme registered on `session`, ignoring + /// its editions. The writer defaults to this when editions are disabled. + pub fn from_session_no_editions(session: &VortexSession) -> Self { + Self { + compressor: CompressorConfig::Schemes(session.registered_schemes()), + ..Self::from_session(session) + } + } } impl WriteStrategyBuilder { @@ -137,15 +146,6 @@ impl WriteStrategyBuilder { self } - /// Override the compression schemes. - /// - /// The strategy builds two compressors from them: one for data, without `IntDictScheme`, and - /// one for stats. The list is used as given; it is not filtered by the session's editions. - pub fn with_schemes(mut self, schemes: Vec<&'static dyn Scheme>) -> Self { - self.compressor = CompressorConfig::Schemes(schemes); - self - } - /// Set the compressor to an opaque [`CompressorPlugin`]. /// /// The compressor is used as-is for both data and stats compression. Use this when the diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index ec98e06cc1b..8d110f21cd4 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -67,6 +67,7 @@ use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; use vortex_btrblocks::CascadingCompressor; +use vortex_btrblocks::CompressionSession; use vortex_btrblocks::CompressionSessionExt; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::string::StringDictScheme; @@ -103,7 +104,9 @@ use crate::VortexFile; use crate::WriteOptionsSessionExt; use crate::flatbuffers::footer as fb; use crate::footer::SegmentSpec; -static SESSION: LazyLock = LazyLock::new(|| { +static SESSION: LazyLock = LazyLock::new(new_session); + +fn new_session() -> VortexSession { let session = array_session() .with::() .with::(); @@ -112,7 +115,7 @@ static SESSION: LazyLock = LazyLock::new(|| { crate::enable_all_registered_array_encodings(&session); session -}); +} fn strict_sorted(indices: Buffer) -> StrictSortedBuffer { StrictSortedBuffer::try_new(indices).expect("test indices should be strictly increasing") @@ -2588,19 +2591,20 @@ async fn dict_probe_honours_configured_compressor() -> VortexResult<()> { "default builder should produce a dict layout for low-cardinality strings" ); - let no_string_dict: Vec<_> = SESSION - .permitted_schemes() + let mut no_string_dict = CompressionSession::empty(); + for scheme in SESSION + .registered_schemes() .into_iter() .filter(|scheme| scheme.id() != StringDictScheme.id()) - .collect(); + { + no_string_dict.register(scheme); + } + let session = new_session(); + session.register(no_string_dict); let mut buf = ByteBufferMut::empty(); - let summary = SESSION + let summary = session .write_options() - .with_strategy( - crate::strategy::WriteStrategyBuilder::from_session(&SESSION) - .with_schemes(no_string_dict) - .build(), - ) + .with_strategy(crate::strategy::WriteStrategyBuilder::from_session(&session).build()) .write(&mut buf, strings.to_array_stream()) .await?; assert!( diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 17711da378a..61376fd7344 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -30,7 +30,6 @@ use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; -use vortex_btrblocks::CompressionSessionExt; use vortex_buffer::ByteBuffer; use vortex_edition::ComponentKind; use vortex_edition::EditionSessionExt; @@ -251,9 +250,7 @@ impl VortexWriteOptions { Some(strategy) => strategy, None if enforce_editions => WriteStrategyBuilder::from_session(&self.session).build(), // With editions disabled every registered encoding may be written. - None => WriteStrategyBuilder::from_session(&self.session) - .with_schemes(self.session.registered_schemes()) - .build(), + None => WriteStrategyBuilder::from_session_no_editions(&self.session).build(), }; let dtype = stream.dtype().clone(); if enforce_editions { diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 034887e060f..b81595aab4c 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -16,12 +16,9 @@ use vortex::array::IntoArray; use vortex::array::iter::ArrayIterator; use vortex::array::iter::ArrayIteratorAdapter; use vortex::array::iter::ArrayIteratorExt; -use vortex::compressor::COMPACT_SCHEMES; -use vortex::compressor::CompressionSessionExt; 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; @@ -45,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<()> { @@ -377,28 +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 mut strategy = WriteStrategyBuilder::from_session(session); - if self.use_compact_encodings { - strategy = strategy.with_schemes( - session.permit( - session - .registered_schemes() - .into_iter() - .chain(COMPACT_SCHEMES.iter().copied()) - .collect(), - ), - ); - } - 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?; @@ -408,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..d2ac1f60dcd 100644 --- a/vortex-python/src/session.rs +++ b/vortex-python/src/session.rs @@ -11,8 +11,9 @@ use std::sync::atomic::AtomicPtr; use std::sync::atomic::Ordering; use vortex::VortexSessionDefault; +use vortex::compressor::COMPACT_SCHEMES; +use vortex::compressor::CompressionSessionExt; use vortex::io::runtime::BlockingRuntime; -#[cfg(unix)] use vortex::io::runtime::Handle; use vortex::io::session::RuntimeSessionExt; #[cfg(unix)] @@ -22,34 +23,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 +87,35 @@ 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); + for scheme in COMPACT_SCHEMES { + session.register_scheme(*scheme); + } + 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 9b60ffc20dc..7ca9c2306bb 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs @@ -37,6 +37,11 @@ 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)?; + for scheme in COMPACT_SCHEMES { + compact_session.register_scheme(*scheme); + } for dataset in fixtures() .into_iter() .filter(|fixture| !is_clickbench_fixture(fixture.name())) @@ -45,24 +50,14 @@ mod tests { let regular_bytes = adapter::write_compressed_to_bytes_with_session( &session, array.clone(), - WriteStrategyBuilder::from_session(&session) - .with_schemes(session.registered_schemes()) - .build(), + WriteStrategyBuilder::from_session_no_editions(&session).build(), )?; let _regular = adapter::read_file(regular_bytes)?; let compact_bytes = adapter::write_compressed_to_bytes_with_session( - &session, + &compact_session, array, - WriteStrategyBuilder::from_session(&session) - .with_schemes( - session - .registered_schemes() - .into_iter() - .chain(COMPACT_SCHEMES.iter().copied()) - .collect(), - ) - .build(), + WriteStrategyBuilder::from_session_no_editions(&compact_session).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 6749ca5728f..10b41d235d6 100644 --- a/vortex-test/compat-gen/src/fixtures/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/mod.rs @@ -140,26 +140,16 @@ 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 execution context's session registers no compression schemes, so build the - // strategy from the same default session the adapter writes with. + // 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::from_session(&session) - .with_schemes( - session - .registered_schemes() - .into_iter() - .chain(COMPACT_SCHEMES.iter().copied()) - .collect(), - ) - .build(); - adapter::write_compressed(&path, array, strategy)?; - } else { - let strategy = WriteStrategyBuilder::from_session(&session) - .with_schemes(session.registered_schemes()) - .build(); - adapter::write_compressed(&path, array, strategy)?; + for scheme in COMPACT_SCHEMES { + session.register_scheme(*scheme); + } } + let strategy = WriteStrategyBuilder::from_session_no_editions(&session).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 3c08dc0e392..37f06fbf25e 100644 --- a/vortex-tui/src/convert.rs +++ b/vortex-tui/src/convert.rs @@ -12,13 +12,14 @@ 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::COMPACT_SCHEMES; use vortex::compressor::CompressionSessionExt; 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; @@ -97,23 +98,21 @@ pub async fn exec_convert(session: &VortexSession, flags: ConvertArgs) -> anyhow .boxed(); } - let mut strategy = WriteStrategyBuilder::from_session(session); - if matches!(flags.strategy, Strategy::Compact) { - strategy = strategy.with_schemes( - session.permit( - session - .registered_schemes() - .into_iter() - .chain(COMPACT_SCHEMES.iter().copied()) - .collect(), - ), - ); - } + // 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()); + for scheme in COMPACT_SCHEMES { + compact_session.register_scheme(*scheme); + } + &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/tracing_vortex.rs b/vortex/examples/tracing_vortex.rs index c2308e0d84c..9d14753cf4a 100644 --- a/vortex/examples/tracing_vortex.rs +++ b/vortex/examples/tracing_vortex.rs @@ -43,7 +43,6 @@ use vortex::compressor::COMPACT_SCHEMES; use vortex::compressor::CompressionSessionExt; use vortex::dtype::DType; use vortex::dtype::Nullability; -use vortex::file::WriteStrategyBuilder; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_session::VortexSession; @@ -54,6 +53,10 @@ async fn main() -> 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. + for scheme in COMPACT_SCHEMES { + session.register_scheme(*scheme); + } // Create output directory let output_dir: PathBuf = "vortex-traces/".into(); @@ -390,20 +393,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::from_session(&session) - .with_schemes( - session.permit( - session - .registered_schemes() - .into_iter() - .chain(COMPACT_SCHEMES.iter().copied()) - .collect(), - ), - ) - .build(), - ); + let write_opts = session.write_options(); write_opts .write(&mut file, struct_array.into_array().to_array_stream()) diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 53e1bb1a40a..b58813b91f8 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -381,7 +381,6 @@ mod test { use vortex_error::VortexResult; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; - use vortex_file::WriteStrategyBuilder; use vortex_session::VortexSession; use crate as vortex; @@ -496,21 +495,11 @@ mod test { let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("example_compact.vortex"); + for scheme in COMPACT_SCHEMES { + session.register_scheme(*scheme); + } session .write_options() - .with_strategy( - WriteStrategyBuilder::from_session(&session) - .with_schemes( - session.permit( - session - .registered_schemes() - .into_iter() - .chain(COMPACT_SCHEMES.iter().copied()) - .collect(), - ), - ) - .build(), - ) .write( &mut tokio::fs::File::create(&path).await?, array.clone().into_array().to_array_stream(), From 2f5d6cc3254f595a453fdadd3bffefdbdfdd0050 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 20:17:15 -0400 Subject: [PATCH 04/11] Add CompressionSession::compact and ::cuda, hide the scheme lists Signed-off-by: Matt Katz --- benchmarks/compress-bench/README.md | 2 +- benchmarks/compress-bench/src/gpu/vortex.rs | 4 +- benchmarks/string-bench/src/serialized.rs | 6 +- fuzz/src/lib.rs | 8 +- vortex-bench/src/lib.rs | 7 +- vortex-btrblocks/src/lib.rs | 5 +- vortex-btrblocks/src/schemes/string/fsst.rs | 2 +- .../schemes/string/scheme_selection_tests.rs | 11 +-- vortex-btrblocks/src/session.rs | 89 ++++++++++++++++--- vortex-btrblocks/src/tests.rs | 7 +- vortex-btrblocks/tests/golden.rs | 46 +++++----- vortex-btrblocks/tests/varbin_scheme.rs | 5 +- vortex-cuda/ffi/src/lib.rs | 8 +- vortex-cuda/gpu-scan-cli/src/main.rs | 6 +- vortex-cuda/src/layout.rs | 54 ++--------- vortex-python/src/session.rs | 7 +- .../src/fixtures/arrays/datasets/mod.rs | 7 +- vortex-test/compat-gen/src/fixtures/mod.rs | 7 +- vortex-tui/src/convert.rs | 7 +- vortex/examples/tracing_vortex.rs | 7 +- vortex/src/lib.rs | 10 +-- 21 files changed, 145 insertions(+), 160 deletions(-) diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 44cd3c36b8d..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 - (`use_cuda_schemes`) 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 380283d5996..5588aabb9b8 100644 --- a/benchmarks/compress-bench/src/gpu/vortex.rs +++ b/benchmarks/compress-bench/src/gpu/vortex.rs @@ -25,6 +25,7 @@ use vortex::array::VortexSessionExecute; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; use vortex::compressor::BtrBlocksCompressor; +use vortex::compressor::CompressionSession; use vortex::error::VortexResult; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; @@ -49,7 +50,6 @@ use vortex_cuda::PooledFileReadAtOptions; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::CudaFlatLayoutStrategy; use vortex_cuda::layout::register_cuda_layout; -use vortex_cuda::layout::use_cuda_schemes; use crate::gpu::writer::GPU_ROW_GROUP_SIZE; @@ -59,7 +59,7 @@ use crate::gpu::writer::GPU_ROW_GROUP_SIZE; static GPU_SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::default().with_tokio(); register_cuda_layout(&session); - use_cuda_schemes(&session); + session.register(CompressionSession::cuda()); session }); diff --git a/benchmarks/string-bench/src/serialized.rs b/benchmarks/string-bench/src/serialized.rs index 95991c250c4..ff229e44659 100644 --- a/benchmarks/string-bench/src/serialized.rs +++ b/benchmarks/string-bench/src/serialized.rs @@ -370,7 +370,7 @@ mod tests { use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::current::CurrentThreadRuntime; use vortex::io::session::RuntimeSessionExt; - use vortex_btrblocks::DEFAULT_SCHEMES; + use vortex_btrblocks::CompressionSession; use vortex_btrblocks::SchemeExt; use super::*; @@ -380,7 +380,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 = DEFAULT_SCHEMES + let default = CompressionSession::default(); + let mut actual = default + .schemes() .iter() .filter(|scheme| scheme.matches(&canonical)) .map(|scheme| scheme.id()) diff --git a/fuzz/src/lib.rs b/fuzz/src/lib.rs index 35027856fea..e4025001c0e 100644 --- a/fuzz/src/lib.rs +++ b/fuzz/src/lib.rs @@ -57,9 +57,7 @@ mod native_runtime { use vortex::VortexSessionDefault; #[cfg(feature = "zstd")] - use vortex::compressor::COMPACT_SCHEMES; - #[cfg(feature = "zstd")] - use vortex::compressor::CompressionSessionExt; + use vortex::compressor::CompressionSession; use vortex_io::runtime::BlockingRuntime; use vortex_io::runtime::current::CurrentThreadRuntime; use vortex_io::session::RuntimeSessionExt; @@ -86,9 +84,7 @@ mod native_runtime { pub static COMPACT_SESSION: LazyLock = LazyLock::new(|| { let session = VortexSession::default().with_handle(RUNTIME.handle()); super::enable_latest_core_edition(&session); - for scheme in COMPACT_SCHEMES { - session.register_scheme(*scheme); - } + session.register(CompressionSession::compact()); session }); } diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 2273db8ef99..11b13fa715c 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -28,8 +28,7 @@ use tpcds::TpcDsBenchmark; use tpch::benchmark::TpcHBenchmark; pub use utils::file::*; pub use utils::logging::*; -use vortex::compressor::COMPACT_SCHEMES; -use vortex::compressor::CompressionSessionExt; +use vortex::compressor::CompressionSession; use vortex::error::VortexExpect; use vortex::error::vortex_err; use vortex::utils::aliases::hash_map::HashMap; @@ -82,9 +81,7 @@ 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(); - for scheme in COMPACT_SCHEMES { - session.register_scheme(*scheme); - } + session.register(CompressionSession::compact()); session }); diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index fa4a44099a1..09b6b53a8ff 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -31,7 +31,7 @@ //! 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). The schemes available to a compressor are those registered on its session's -//! [`CompressionSession`], which starts with [`DEFAULT_SCHEMES`]. +//! [`CompressionSession`], which starts with the default schemes. //! [`BtrBlocksCompressor::from_session`] keeps the registered schemes whose serialized IDs the //! session's enabled editions permit; [`from_session_no_editions`] keeps them all. //! @@ -79,11 +79,8 @@ mod trace_tests; // Btrblocks-specific exports. pub use canonical_compressor::BtrBlocksCompressor; pub use schemes::patches::compress_patches; -#[cfg(feature = "zstd")] -pub use session::COMPACT_SCHEMES; pub use session::CompressionSession; pub use session::CompressionSessionExt; -pub use session::DEFAULT_SCHEMES; pub use session::DELTA_SCHEME; pub use vortex_compressor::CascadingCompressor; pub use vortex_compressor::scheme::CompressorContext; diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index e0c9c7f44b9..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::DEFAULT_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 22b052fcbcd..017b83a08d1 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -19,7 +19,6 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::CompressionSession; use crate::CompressionSessionExt; -use crate::DEFAULT_SCHEMES; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -54,10 +53,11 @@ fn test_all_schemes_includes_onpair() { use crate::SchemeExt; use crate::schemes::string::onpair::OnPairScheme; - let ids: Vec<_> = DEFAULT_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 DEFAULT_SCHEMES" + "OnPairScheme not registered in the default schemes" ); } @@ -92,9 +92,10 @@ fn test_fsst_in_default_scheme_list() -> VortexResult<()> { use crate::schemes::string::FSSTScheme; // FSST is registered by default. + let default = CompressionSession::default(); assert!( - DEFAULT_SCHEMES.iter().any(|s| s.id() == FSSTScheme.id()), - "FSSTScheme should be in DEFAULT_SCHEMES", + default.schemes().iter().any(|s| s.id() == FSSTScheme.id()), + "FSSTScheme should be in the default schemes", ); // An FSST-only compressor still produces an FSST array for FSST-favourable diff --git a/vortex-btrblocks/src/session.rs b/vortex-btrblocks/src/session.rs index fba52af561d..5ddcc18169a 100644 --- a/vortex-btrblocks/src/session.rs +++ b/vortex-btrblocks/src/session.rs @@ -5,8 +5,9 @@ //! //! A session's [`CompressionSession`] holds the schemes available to compressors built from it //! with [`BtrBlocksCompressor::from_session`](crate::BtrBlocksCompressor::from_session). It -//! starts with [`DEFAULT_SCHEMES`]. Whether a registered scheme may write its encodings is -//! decided by the session's enabled editions. +//! 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 the session's enabled editions. use std::any::Any; @@ -30,7 +31,7 @@ use crate::schemes::temporal; /// /// This list is order-sensitive: the compressor preserves registration order, so that /// tie-breaking is deterministic. -pub const DEFAULT_SCHEMES: &[&dyn Scheme] = &[ +const DEFAULT_SCHEMES: &[&dyn Scheme] = &[ //////////////////////////////////////////////////////////////////////////////////////////////// // Integer schemes. //////////////////////////////////////////////////////////////////////////////////////////////// @@ -73,12 +74,10 @@ pub const DEFAULT_SCHEMES: &[&dyn Scheme] = &[ &temporal::TemporalScheme, ]; -/// Compact schemes (Zstd for strings and binary, Pco for numerics when the `pco` feature is on). -/// -/// Not part of [`DEFAULT_SCHEMES`]: they trade decode speed for compression ratio, so callers add -/// them to a compressor's scheme list explicitly. +/// The schemes [`CompressionSession::compact`] adds to the defaults: Zstd for strings and binary, +/// and Pco for numerics when the `pco` feature is on. #[cfg(feature = "zstd")] -pub const COMPACT_SCHEMES: &[&dyn Scheme] = &[ +const COMPACT_SCHEMES: &[&dyn Scheme] = &[ &string::ZstdScheme, &binary::ZstdScheme, #[cfg(feature = "pco")] @@ -87,18 +86,18 @@ pub const COMPACT_SCHEMES: &[&dyn Scheme] = &[ &float::PcoScheme, ]; -/// Delta, kept out of [`DEFAULT_SCHEMES`] because it is slower to decompress than the schemes that -/// would otherwise win. Callers that want it add it to their scheme list and permit -/// `fastlanes.delta`. +/// 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 [`DEFAULT_SCHEMES`] once we have scheme filtering. +/// 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 [`DEFAULT_SCHEMES`]; [`empty`](Self::empty) registers none. +/// [`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. @@ -113,6 +112,55 @@ impl CompressionSession { } } + /// The default schemes plus the compact ones: Zstd for strings and binary, and Pco for + /// numerics when the `pco` feature is on. They trade decode speed for compression ratio. + #[cfg(feature = "zstd")] + 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. @@ -219,6 +267,21 @@ mod tests { ); } + #[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())); + } + + #[cfg(feature = "zstd")] + #[test] + fn compact_extends_the_defaults() { + let compact = ids(CompressionSession::compact().schemes()); + assert_eq!(&compact[..DEFAULT_SCHEMES.len()], &ids(DEFAULT_SCHEMES)[..]); + assert!(compact.contains(&string::ZstdScheme.id())); + } + /// Without enabled editions no serialized ID is permitted, so nothing survives. #[test] fn no_editions_permit_nothing() { diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index f64ef7f9900..481d7b20300 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -30,8 +30,6 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; #[cfg(feature = "zstd")] -use crate::COMPACT_SCHEMES; -#[cfg(feature = "zstd")] use crate::CompressionSession; use crate::CompressionSessionExt; #[cfg(feature = "zstd")] @@ -205,10 +203,7 @@ fn test_compact_binary_zstd_compressed() -> VortexResult<()> { DType::Binary(Nullability::NonNullable), ); - let session = vortex_array::array_session(); - for scheme in COMPACT_SCHEMES { - session.register_scheme(*scheme); - } + let session = vortex_array::array_session().with_some(CompressionSession::compact()); let compressor = BtrBlocksCompressor::from_session_no_editions(&session); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index b656a657556..dde3d4a12b1 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -14,7 +14,7 @@ //! - `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 [`COMPACT_SCHEMES`](vortex_btrblocks::COMPACT_SCHEMES) +//! 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 @@ -49,12 +49,7 @@ use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; -#[cfg(all(feature = "zstd", feature = "pco"))] -use vortex_btrblocks::COMPACT_SCHEMES; use vortex_btrblocks::CompressionSession; -use vortex_btrblocks::CompressionSessionExt; -use vortex_btrblocks::DEFAULT_SCHEMES; -use vortex_btrblocks::Scheme; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::string::OnPairScheme; use vortex_buffer::Buffer; @@ -399,24 +394,26 @@ 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(schemes: Vec<&'static dyn Scheme>) -> Vec<&'static dyn Scheme> { - schemes - .into_iter() +fn without_onpair(registry: &CompressionSession) -> CompressionSession { + let mut filtered = CompressionSession::empty(); + for scheme in registry + .schemes() + .iter() .filter(|scheme| scheme.id() != OnPairScheme.id()) - .collect() + { + filtered.register(*scheme); + } + filtered } -/// A session registering `schemes` and enabling `editions`. +/// A session with the schemes in `registry` and `editions` enabled. fn edition_session( editions: &[EditionId], - schemes: Vec<&'static dyn Scheme>, + registry: CompressionSession, ) -> VortexResult { let session = vortex_array::array_session() - .with_some(CompressionSession::empty()) + .with_some(registry) .with::(); - for scheme in schemes { - session.register_scheme(scheme); - } for family in EDITION_FAMILIES { session.editions().declare_family(family)?; } @@ -431,7 +428,10 @@ fn edition_session( #[test] fn golden_regular() -> VortexResult<()> { - let session = edition_session(&[CORE_2026_08_3], without_onpair(DEFAULT_SCHEMES.to_vec()))?; + let session = edition_session( + &[CORE_2026_08_3], + without_onpair(&CompressionSession::default()), + )?; let compressor = BtrBlocksCompressor::from_session(&session); golden_corpus_snapshots("regular", &compressor) } @@ -439,7 +439,7 @@ fn golden_regular() -> VortexResult<()> { /// Pins OnPair's selection over FSST on the structured-string entry. #[test] fn golden_onpair() -> VortexResult<()> { - let session = edition_session(&[CORE_2026_08_3], DEFAULT_SCHEMES.to_vec())?; + let session = edition_session(&[CORE_2026_08_3], CompressionSession::default())?; let compressor = BtrBlocksCompressor::from_session(&session); golden_snapshots( "onpair", @@ -451,12 +451,10 @@ fn golden_onpair() -> VortexResult<()> { #[cfg(all(feature = "zstd", feature = "pco"))] #[test] fn golden_compact() -> VortexResult<()> { - let schemes = DEFAULT_SCHEMES - .iter() - .chain(COMPACT_SCHEMES.iter()) - .copied() - .collect(); - let session = edition_session(&[CORE_2026_08_3], without_onpair(schemes))?; + let session = edition_session( + &[CORE_2026_08_3], + without_onpair(&CompressionSession::compact()), + )?; vortex_zstd::initialize(&session); session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; let compressor = BtrBlocksCompressor::from_session(&session); diff --git a/vortex-btrblocks/tests/varbin_scheme.rs b/vortex-btrblocks/tests/varbin_scheme.rs index b312a06a9f7..be4c1b8b7de 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -17,7 +17,6 @@ use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; use vortex_btrblocks::CompressionSession; use vortex_btrblocks::CompressionSessionExt; -use vortex_btrblocks::DEFAULT_SCHEMES; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::SchemeId; use vortex_btrblocks::schemes::binary::VarBinScheme; @@ -32,7 +31,9 @@ const N: usize = 100_000; /// The default schemes minus `excluded`. fn default_without(excluded: SchemeId) -> BtrBlocksCompressor { let session = vortex_array::array_session().with_some(CompressionSession::empty()); - for scheme in DEFAULT_SCHEMES + let default = CompressionSession::default(); + for scheme in default + .schemes() .iter() .filter(|scheme| scheme.id() != excluded) { diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 6bf843dc8d1..35d965697b5 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; @@ -40,7 +41,6 @@ use vortex_cuda::arrow::DeviceArrayExt; use vortex_cuda::arrow::DeviceArrayStreamExt; use vortex_cuda::layout::cuda_write_strategy; use vortex_cuda::layout::register_cuda_layout; -use vortex_cuda::layout::use_cuda_schemes; use vortex_ffi::ffi_runtime; use vortex_ffi::try_or; use vortex_ffi::vx_array; @@ -101,7 +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); - use_cuda_schemes(&session); + session.register(CompressionSession::cuda()); session })) }) @@ -148,7 +148,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( try_or(error_out, ptr::null_mut(), || { // SAFETY: The caller supplies a live borrowed session handle. let vortex_session = session_with_cuda(unsafe { vx_session_ref(session) }?); - use_cuda_schemes(vortex_session); + vortex_session.register(CompressionSession::cuda()); // SAFETY: All borrowed inputs satisfy the underlying sink's requirements. unsafe { vx_array_sink_open_file_with_strategy( @@ -1037,7 +1037,7 @@ mod tests { fn test_projection_gpu_values_and_validity() -> VortexResult<()> { let session = session().with_some(CudaSession::try_default()?); register_cuda_layout(&session); - use_cuda_schemes(&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 04cff42a5e2..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; @@ -36,7 +37,6 @@ use vortex_cuda::TracingLaunchStrategy; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::cuda_write_strategy; use vortex_cuda::layout::register_cuda_layout; -use vortex_cuda::layout::use_cuda_schemes; use vortex_cuda_macros::cuda_available; use vortex_cuda_macros::cuda_not_available; @@ -94,7 +94,7 @@ async fn main() -> VortexResult<()> { async fn cmd_convert(input: PathBuf, output: PathBuf) -> VortexResult<()> { let session = VortexSession::default(); register_cuda_layout(&session); - use_cuda_schemes(&session); + session.register(CompressionSession::cuda()); let input_file = session.open_options().open_path(&input).await?; let scan = input_file.scan()?.into_array_stream()?; @@ -141,7 +141,7 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes let session = VortexSession::default(); register_cuda_layout(&session); - use_cuda_schemes(&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 91eb32b4268..4b8abc4c8df 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -31,10 +31,6 @@ use vortex::buffer::BufferString; use vortex::buffer::ByteBuffer; use vortex::compressor::BtrBlocksCompressor; use vortex::compressor::CascadingCompressor; -use vortex::compressor::CompressionSession; -use vortex::compressor::DEFAULT_SCHEMES; -use vortex::compressor::SchemeExt; -use vortex::compressor::schemes; use vortex::dtype::DType; use vortex::dtype::FieldMask; use vortex::editions::Edition; @@ -553,8 +549,8 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { /// Build a CUDA-flat writer from the schemes registered on `session` that its editions permit. /// -/// Requires [`register_cuda_layout`], and [`use_cuda_schemes`] for the file to use only -/// encodings the GPU decodes. Zero `block_rows` uses default sizing and dictionary policy; +/// Requires [`register_cuda_layout`], and a [`CompressionSession::cuda`] registry for the file +/// to use only encodings the GPU decodes. Zero `block_rows` uses default sizing and dictionary policy; /// nonzero sets row blocks without outer dictionaries or byte coalescing, retaining per-block /// dictionary compression. pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc { @@ -573,47 +569,6 @@ pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc); @@ -697,6 +652,7 @@ 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; @@ -754,7 +710,7 @@ mod tests { let runtime = CurrentThreadRuntime::new(); let session = VortexSession::default().with_handle(runtime.handle()); register_cuda_layout(&session); - use_cuda_schemes(&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?; @@ -789,7 +745,7 @@ mod tests { let runtime = CurrentThreadRuntime::new(); let session = VortexSession::default().with_handle(runtime.handle()); register_cuda_layout(&session); - use_cuda_schemes(&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-python/src/session.rs b/vortex-python/src/session.rs index d2ac1f60dcd..9e33dfbd6ea 100644 --- a/vortex-python/src/session.rs +++ b/vortex-python/src/session.rs @@ -11,8 +11,7 @@ use std::sync::atomic::AtomicPtr; use std::sync::atomic::Ordering; use vortex::VortexSessionDefault; -use vortex::compressor::COMPACT_SCHEMES; -use vortex::compressor::CompressionSessionExt; +use vortex::compressor::CompressionSession; use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Handle; use vortex::io::session::RuntimeSessionExt; @@ -93,9 +92,7 @@ fn new_session(handle: Handle) -> VortexSession { fn new_compact_session(handle: Handle) -> VortexSession { let session = new_session(handle); - for scheme in COMPACT_SCHEMES { - session.register_scheme(*scheme); - } + session.register(CompressionSession::compact()); session } 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 7ca9c2306bb..7acb397bb26 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/datasets/mod.rs @@ -17,8 +17,7 @@ pub fn fixtures() -> Vec> { #[cfg(test)] mod tests { use vortex::VortexSessionDefault; - use vortex::compressor::COMPACT_SCHEMES; - use vortex::compressor::CompressionSessionExt; + use vortex::compressor::CompressionSession; use vortex::editions::CORE_2026_08_3; use vortex::editions::EditionSessionExt; use vortex::file::WriteStrategyBuilder; @@ -39,9 +38,7 @@ mod tests { session.enable_edition(CORE_2026_08_3)?; let compact_session = VortexSession::default(); compact_session.enable_edition(CORE_2026_08_3)?; - for scheme in COMPACT_SCHEMES { - compact_session.register_scheme(*scheme); - } + compact_session.register(CompressionSession::compact()); for dataset in fixtures() .into_iter() .filter(|fixture| !is_clickbench_fixture(fixture.name())) diff --git a/vortex-test/compat-gen/src/fixtures/mod.rs b/vortex-test/compat-gen/src/fixtures/mod.rs index 10b41d235d6..955b013d88d 100644 --- a/vortex-test/compat-gen/src/fixtures/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/mod.rs @@ -9,8 +9,7 @@ use std::sync::Arc; use vortex::VortexSessionDefault; use vortex::array::ArrayId; use vortex::array::ArrayRef; -use vortex::compressor::COMPACT_SCHEMES; -use vortex::compressor::CompressionSessionExt; +use vortex::compressor::CompressionSession; use vortex::file::WriteStrategyBuilder; use vortex::session::VortexSession; use vortex_array::ExecutionCtx; @@ -144,9 +143,7 @@ impl Fixture for DatasetFixtureAdapter { // session may be used. let session = VortexSession::default(); if self.compact { - for scheme in COMPACT_SCHEMES { - session.register_scheme(*scheme); - } + session.register(CompressionSession::compact()); } let strategy = WriteStrategyBuilder::from_session_no_editions(&session).build(); adapter::write_compressed(&path, array, strategy)?; diff --git a/vortex-tui/src/convert.rs b/vortex-tui/src/convert.rs index 37f06fbf25e..73f0948110d 100644 --- a/vortex-tui/src/convert.rs +++ b/vortex-tui/src/convert.rs @@ -14,8 +14,7 @@ use tokio::fs::File; use tokio::io::AsyncWriteExt; use vortex::VortexSessionDefault; use vortex::array::stream::ArrayStreamAdapter; -use vortex::compressor::COMPACT_SCHEMES; -use vortex::compressor::CompressionSessionExt; +use vortex::compressor::CompressionSession; use vortex::error::VortexExpect; use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; @@ -102,9 +101,7 @@ pub async fn exec_convert(session: &VortexSession, flags: ConvertArgs) -> anyhow let compact_session; let session = if matches!(flags.strategy, Strategy::Compact) { compact_session = VortexSession::default().with_handle(session.handle()); - for scheme in COMPACT_SCHEMES { - compact_session.register_scheme(*scheme); - } + compact_session.register(CompressionSession::compact()); &compact_session } else { session diff --git a/vortex/examples/tracing_vortex.rs b/vortex/examples/tracing_vortex.rs index 9d14753cf4a..5d74935303e 100644 --- a/vortex/examples/tracing_vortex.rs +++ b/vortex/examples/tracing_vortex.rs @@ -39,8 +39,7 @@ use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinArray; use vortex::array::stream::ArrayStreamExt; use vortex::array::validity::Validity; -use vortex::compressor::COMPACT_SCHEMES; -use vortex::compressor::CompressionSessionExt; +use vortex::compressor::CompressionSession; use vortex::dtype::DType; use vortex::dtype::Nullability; use vortex_file::OpenOptionsSessionExt; @@ -54,9 +53,7 @@ async fn main() -> Result<(), Box> { let session = VortexSession::default(); // Use compact encodings (Pco + Zstd) for the telemetry files. - for scheme in COMPACT_SCHEMES { - session.register_scheme(*scheme); - } + session.register(CompressionSession::compact()); // Create output directory let output_dir: PathBuf = "vortex-traces/".into(); diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index b58813b91f8..69a1e38864f 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -145,12 +145,9 @@ pub mod buffer { /// Default adaptive compression APIs based on the maintained BtrBlocks-style compressor. pub mod compressor { pub use vortex_btrblocks::BtrBlocksCompressor; - #[cfg(feature = "zstd")] - pub use vortex_btrblocks::COMPACT_SCHEMES; pub use vortex_btrblocks::CascadingCompressor; pub use vortex_btrblocks::CompressionSession; pub use vortex_btrblocks::CompressionSessionExt; - pub use vortex_btrblocks::DEFAULT_SCHEMES; pub use vortex_btrblocks::Scheme; pub use vortex_btrblocks::SchemeExt; pub use vortex_btrblocks::SchemeId; @@ -375,8 +372,7 @@ mod test { use vortex_array::expr::select; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; - use vortex_btrblocks::COMPACT_SCHEMES; - use vortex_btrblocks::CompressionSessionExt; + use vortex_btrblocks::CompressionSession; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_file::OpenOptionsSessionExt; @@ -495,9 +491,7 @@ mod test { let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("example_compact.vortex"); - for scheme in COMPACT_SCHEMES { - session.register_scheme(*scheme); - } + session.register(CompressionSession::compact()); session .write_options() .write( From 54f3d9bf5b1297ee3cef9451fada7c0557568854 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 21:23:15 -0400 Subject: [PATCH 05/11] Add BtrBlocksOptions, build strategy compressors from the session at build time Signed-off-by: Matt Katz --- vortex-btrblocks/benches/compress.rs | 9 ++- vortex-btrblocks/benches/compress_listview.rs | 9 ++- vortex-btrblocks/src/canonical_compressor.rs | 56 +++++++++++--- vortex-btrblocks/src/lib.rs | 14 +++- .../schemes/float/scheme_selection_tests.rs | 10 +-- vortex-btrblocks/src/schemes/float/tests.rs | 8 +- .../schemes/integer/scheme_selection_tests.rs | 19 ++--- vortex-btrblocks/src/schemes/integer/tests.rs | 12 +-- .../schemes/string/scheme_selection_tests.rs | 10 +-- vortex-btrblocks/src/schemes/string/tests.rs | 6 +- vortex-btrblocks/src/session.rs | 12 +-- vortex-btrblocks/src/tests.rs | 34 ++++++--- vortex-btrblocks/src/trace_tests.rs | 5 +- vortex-btrblocks/tests/onpair_roundtrip.rs | 23 ++++-- vortex-btrblocks/tests/varbin_scheme.rs | 17 ++++- vortex-cuda/src/layout.rs | 3 +- vortex-file/benches/split_collection.rs | 21 ++++-- vortex-file/src/strategy.rs | 75 ++++++++++--------- vortex-file/src/tests.rs | 16 ++-- vortex-file/src/writer.rs | 10 ++- vortex-layout/src/layouts/dict/reader.rs | 21 +++++- .../src/fixtures/arrays/datasets/mod.rs | 14 +++- vortex-test/compat-gen/src/fixtures/mod.rs | 10 ++- vortex/src/lib.rs | 4 +- 24 files changed, 285 insertions(+), 133 deletions(-) diff --git a/vortex-btrblocks/benches/compress.rs b/vortex-btrblocks/benches/compress.rs index bc5c00987d2..4b7c528813d 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::BtrBlocksOptions; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; @@ -51,7 +52,13 @@ mod benchmarks { let array = make_clickbench_window_name() .execute::(&mut ctx) .unwrap(); - let compressor = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let compressor = BtrBlocksCompressor::from_session_with_options( + &SESSION, + &BtrBlocksOptions { + enforce_editions: false, + ..Default::default() + }, + ); 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 b6d01c86dd2..31f2714363a 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::BtrBlocksOptions; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; @@ -183,7 +184,13 @@ mod benchmarks { fn compress_listview(bencher: Bencher, layout: OffsetLayout) { let array = build_nested_listview(NUM_ROWS, layout); let nbytes = array.nbytes(); - let compressor = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let compressor = BtrBlocksCompressor::from_session_with_options( + &SESSION, + &BtrBlocksOptions { + enforce_editions: false, + ..Default::default() + }, + ); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .input_counter(|_| ItemsCount::new(NUM_ROWS)) diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index b57515342b1..ac81310cf7b 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -12,23 +12,52 @@ use vortex_session::VortexSession; use crate::CascadingCompressor; use crate::CompressionSessionExt; +use crate::SchemeExt; +use crate::SchemeId; + +/// Options for building a [`BtrBlocksCompressor`] from a session. +#[derive(Clone, Debug)] +pub struct BtrBlocksOptions { + /// Keep only the registered schemes whose serialized IDs the session's enabled editions + /// permit, which a file writer requires. Off, every registered scheme is used, for in-memory + /// compression where no edition applies. + 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. /// /// 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_no_editions`](Self::from_session_no_editions) -/// keeps every registered scheme, for in-memory compression where no edition applies. +/// 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; +/// use vortex_btrblocks::BtrBlocksOptions; /// /// let session = vortex_array::array_session(); /// /// // Every registered scheme; this session enables no editions. -/// let compressor = BtrBlocksCompressor::from_session_no_editions(&session); +/// let compressor = BtrBlocksCompressor::from_session_with_options( +/// &session, +/// &BtrBlocksOptions { +/// enforce_editions: false, +/// ..Default::default() +/// }, +/// ); /// ``` #[derive(Clone)] pub struct BtrBlocksCompressor( @@ -37,17 +66,26 @@ pub struct BtrBlocksCompressor( ); impl BtrBlocksCompressor { + /// A compressor with no schemes, which leaves every array as it is. + pub fn empty() -> Self { + Self(CascadingCompressor::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(CascadingCompressor::new(session.permitted_schemes())) + Self::from_session_with_options(session, &BtrBlocksOptions::default()) } - /// Creates a compressor over every scheme registered on `session`, ignoring editions. - /// - /// Use for in-memory compression, where no edition restricts what a file may contain. - pub fn from_session_no_editions(session: &VortexSession) -> Self { - Self(CascadingCompressor::new(session.registered_schemes())) + /// Creates a compressor over the schemes registered on `session`, per `options`. + pub fn from_session_with_options(session: &VortexSession, options: &BtrBlocksOptions) -> Self { + let mut schemes = if options.enforce_editions { + session.permitted_schemes() + } else { + session.registered_schemes() + }; + schemes.retain(|scheme| !options.exclude_schemes.contains(&scheme.id())); + Self(CascadingCompressor::new(schemes)) } /// Compresses an array using BtrBlocks-inspired compression. diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 09b6b53a8ff..533a1382481 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -33,9 +33,7 @@ //! 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; [`from_session_no_editions`] keeps them all. -//! -//! [`from_session_no_editions`]: BtrBlocksCompressor::from_session_no_editions +//! 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 @@ -48,6 +46,7 @@ //! use vortex_array::arrays::PrimitiveArray; //! use vortex_array::validity::Validity; //! use vortex_btrblocks::BtrBlocksCompressor; +//! use vortex_btrblocks::BtrBlocksOptions; //! use vortex_buffer::buffer; //! //! # fn example() -> vortex_error::VortexResult<()> { @@ -55,7 +54,13 @@ //! let array = PrimitiveArray::new(buffer![42u64; 1024], Validity::NonNullable).into_array(); //! //! // In memory, with no editions to respect, compress with every registered scheme. -//! let compressor = BtrBlocksCompressor::from_session_no_editions(&session); +//! 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()); //! # Ok(()) @@ -78,6 +83,7 @@ mod trace_tests; // Re-export framework types from vortex-compressor for backwards compatibility. // Btrblocks-specific exports. pub use canonical_compressor::BtrBlocksCompressor; +pub use canonical_compressor::BtrBlocksOptions; pub use schemes::patches::compress_patches; pub use session::CompressionSession; pub use session::CompressionSessionExt; diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index dbea853803f..fa384f28e53 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::no_editions_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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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 7f1888e92b0..bf4e75fec89 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -20,13 +20,13 @@ use vortex_error::VortexResult; use vortex_fastlanes::RLE; use vortex_session::VortexSession; -use crate::BtrBlocksCompressor; use crate::schemes::float::FloatRLEScheme; +use crate::tests::no_editions_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { - let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 1024); @@ -92,7 +92,7 @@ fn test_sparse_compression() -> VortexResult<()> { array.append_nulls(90); let array = array.finish_into_primitive().into_array(); - let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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 1e8ae27e454..42288a0347b 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -30,20 +30,21 @@ use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; use crate::CompressionSessionExt; use crate::DELTA_SCHEME; +use crate::tests::no_editions_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); /// The default schemes plus opt-in Delta. fn with_delta() -> BtrBlocksCompressor { let session = vortex_array::array_session(); session.register_scheme(&DELTA_SCHEME); - BtrBlocksCompressor::from_session_no_editions(&session) + no_editions_compressor(&session) } #[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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -53,7 +54,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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -63,7 +64,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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert_eq!( @@ -92,7 +93,7 @@ fn test_sparse_compressed() -> VortexResult<()> { } } let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -116,7 +117,7 @@ fn test_dict_compressed() -> VortexResult<()> { } let array = PrimitiveArray::new(Buffer::copy_from(&codes), Validity::NonNullable); - let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -129,7 +130,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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -139,7 +140,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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -156,7 +157,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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; eprintln!("{}", compressed.display_tree()); assert!(compressed.is::()); diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index f641d2335ad..a7fdf71f4e4 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -25,14 +25,14 @@ use vortex_fastlanes::RLE; use vortex_sequence::Sequence; use vortex_session::VortexSession; -use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; +use crate::tests::no_editions_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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); btr.compress(&prim, &mut SESSION.create_execution_ctx())?; Ok(()) diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index 017b83a08d1..6beb6a251c5 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -16,9 +16,9 @@ use vortex_error::VortexResult; use vortex_fsst::FSST; use vortex_session::VortexSession; -use crate::BtrBlocksCompressor; use crate::CompressionSession; use crate::CompressionSessionExt; +use crate::tests::no_editions_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -27,7 +27,7 @@ fn test_constant_compressed() -> VortexResult<()> { let strings: Vec> = vec![Some("constant_value"); 100]; let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + let compressed = no_editions_compressor(&SESSION) .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -42,7 +42,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::from_session_no_editions(&SESSION) + let compressed = no_editions_compressor(&SESSION) .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -74,7 +74,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::from_session_no_editions(&SESSION) + let compressed = no_editions_compressor(&SESSION) .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), @@ -111,7 +111,7 @@ fn test_fsst_in_default_scheme_list() -> VortexResult<()> { let session = vortex_array::array_session().with_some(CompressionSession::empty()); session.register_scheme(&FSSTScheme); - let compressor = BtrBlocksCompressor::from_session_no_editions(&session); + let compressor = no_editions_compressor(&session); 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 57398bb1e3b..1b79c011eb8 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::no_editions_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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); 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 index 5ddcc18169a..f5760155592 100644 --- a/vortex-btrblocks/src/session.rs +++ b/vortex-btrblocks/src/session.rs @@ -74,11 +74,12 @@ const DEFAULT_SCHEMES: &[&dyn Scheme] = &[ &temporal::TemporalScheme, ]; -/// The schemes [`CompressionSession::compact`] adds to the defaults: Zstd for strings and binary, -/// and Pco for numerics when the `pco` feature is on. -#[cfg(feature = "zstd")] +/// 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, @@ -113,8 +114,7 @@ impl CompressionSession { } /// The default schemes plus the compact ones: Zstd for strings and binary, and Pco for - /// numerics when the `pco` feature is on. They trade decode speed for compression ratio. - #[cfg(feature = "zstd")] + /// 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 { @@ -274,11 +274,11 @@ mod tests { assert!(!cuda.contains(&string::StringDictScheme.id())); } - #[cfg(feature = "zstd")] #[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 index 481d7b20300..b2b88f030fd 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -29,8 +29,10 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::BtrBlocksOptions; #[cfg(feature = "zstd")] use crate::CompressionSession; +#[cfg(feature = "zstd")] use crate::CompressionSessionExt; #[cfg(feature = "zstd")] use crate::Scheme; @@ -39,6 +41,18 @@ use crate::schemes::binary; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +/// A compressor over every scheme registered on `session`: these tests compress in memory, where +/// no edition applies. +pub(crate) fn no_editions_compressor(session: &VortexSession) -> BtrBlocksCompressor { + BtrBlocksCompressor::from_session_with_options( + session, + &BtrBlocksOptions { + enforce_editions: false, + ..Default::default() + }, + ) +} + #[rstest] #[case::zctl( unsafe { @@ -66,7 +80,7 @@ fn listview_compress_roundtrip( ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let array_ref = input.clone().into_array(); - let result = BtrBlocksCompressor::from_session_no_editions(&SESSION) + let result = no_editions_compressor(&SESSION) .compress(&array_ref, &mut SESSION.create_execution_ctx())?; if expect_list { assert!(result.as_opt::().is_some()); @@ -81,7 +95,7 @@ fn listview_compress_roundtrip( 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -95,7 +109,7 @@ fn test_constant_all_true() -> VortexResult<()> { 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -112,7 +126,7 @@ fn test_nullable_all_valid_compressed() -> VortexResult<()> { BitBuffer::from(vec![true; 100]), Validity::from(BitBuffer::from(vec![true; 100])), ); - let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -127,7 +141,7 @@ 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -144,7 +158,7 @@ fn test_mixed_not_constant() -> VortexResult<()> { BitBuffer::from(vec![true, false, true, false, true]), Validity::NonNullable, ); - let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -159,7 +173,7 @@ 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::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -177,7 +191,7 @@ fn test_binary_dict_compressed() -> VortexResult<()> { .map(|idx| Some(distinct_values[idx % distinct_values.len()])) .collect::>(); let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); - let btr = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let btr = no_editions_compressor(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -204,7 +218,7 @@ fn test_compact_binary_zstd_compressed() -> VortexResult<()> { ); let session = vortex_array::array_session().with_some(CompressionSession::compact()); - let compressor = BtrBlocksCompressor::from_session_no_editions(&session); + let compressor = no_editions_compressor(&session); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; @@ -242,7 +256,7 @@ fn test_binary_zstd_scheme_encoding( let session = vortex_array::array_session().with_some(CompressionSession::empty()); session.register_scheme(scheme); - let compressor = BtrBlocksCompressor::from_session_no_editions(&session); + let compressor = no_editions_compressor(&session); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 610fc81bb42..10ad1bab745 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -53,9 +53,9 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use crate::BtrBlocksCompressor; use crate::CompressionSessionExt; use crate::DELTA_SCHEME; +use crate::tests::no_editions_compressor; /// A session with the default Vortex encodings registered. /// @@ -130,8 +130,7 @@ fn lineitem() -> VortexResult { fn compressed_lineitem() -> VortexResult { let session = trace_session(); session.register_scheme(&DELTA_SCHEME); - BtrBlocksCompressor::from_session_no_editions(&session) - .compress(&lineitem()?, &mut execution_ctx()) + no_editions_compressor(&session).compress(&lineitem()?, &mut execution_ctx()) } fn field(array: &ArrayRef, name: &str) -> VortexResult { diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index 45d6d67c64b..16aeff20138 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -19,10 +19,23 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::BtrBlocksOptions; use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +/// 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() + }, + ) +} + /// Helper: synthetic short-string corpus that the cascading compressor should /// route through OnPair. fn corpus(n: usize) -> Vec { @@ -58,7 +71,7 @@ fn nonnullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + let compressed = no_editions_compressor(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert a specific scheme — both OnPair and FSST are registered and @@ -101,7 +114,7 @@ fn nullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + let compressed = no_editions_compressor(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert OnPair specifically here — the sample-based selector may @@ -137,7 +150,7 @@ fn large_unique_short_strings_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + let compressed = no_editions_compressor(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); @@ -166,7 +179,7 @@ fn empty_and_short_string_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + let compressed = no_editions_compressor(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); let decoded = compressed @@ -211,7 +224,7 @@ fn delta_dict_offsets_roundtrip() { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let compressed = BtrBlocksCompressor::from_session_no_editions(&SESSION) + let compressed = no_editions_compressor(&SESSION) .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 be4c1b8b7de..f318ccf937d 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -15,6 +15,7 @@ use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::BtrBlocksOptions; use vortex_btrblocks::CompressionSession; use vortex_btrblocks::CompressionSessionExt; use vortex_btrblocks::SchemeExt; @@ -26,6 +27,18 @@ use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +/// 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() + }, + ) +} + const N: usize = 100_000; /// The default schemes minus `excluded`. @@ -39,7 +52,7 @@ fn default_without(excluded: SchemeId) -> BtrBlocksCompressor { { session.register_scheme(*scheme); } - BtrBlocksCompressor::from_session_no_editions(&session) + no_editions_compressor(&session) } fn lcg(state: &mut u64) -> u64 { @@ -87,7 +100,7 @@ fn cases() -> Vec<(&'static str, ArrayRef)> { #[test] fn varbin_scheme_shrinks_binary() -> VortexResult<()> { - let with = BtrBlocksCompressor::from_session_no_editions(&SESSION); + let with = no_editions_compressor(&SESSION); let without = default_without(VarBinScheme.id()); println!( diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 4b8abc4c8df..607e5b44f09 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -30,7 +30,6 @@ use vortex::array::stats::StatsSetRef; use vortex::buffer::BufferString; use vortex::buffer::ByteBuffer; use vortex::compressor::BtrBlocksCompressor; -use vortex::compressor::CascadingCompressor; use vortex::dtype::DType; use vortex::dtype::FieldMask; use vortex::editions::Edition; @@ -562,7 +561,7 @@ pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc = 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::from_session_no_editions(&SESSION) - .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::from_session_no_editions(&SESSION); + 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 2a8a4d2dd7d..0110c082106 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -8,9 +8,7 @@ use std::sync::Arc; use vortex_array::dtype::FieldPath; use vortex_btrblocks::BtrBlocksCompressor; -use vortex_btrblocks::CascadingCompressor; -use vortex_btrblocks::CompressionSessionExt; -use vortex_btrblocks::Scheme; +use vortex_btrblocks::BtrBlocksOptions; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::integer::IntDictScheme; use vortex_error::VortexExpect; @@ -36,10 +34,13 @@ const ONE_MEG: u64 = 1 << 20; /// How the compressor was configured on [`WriteStrategyBuilder`]. enum CompressorConfig { - /// Schemes for the [`BtrBlocksCompressor`]s that [`WriteStrategyBuilder::build`] creates. - /// `IntDictScheme` is automatically excluded from the data compressor to prevent recursive - /// dictionary encoding. - Schemes(Vec<&'static dyn Scheme>), + /// [`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), } @@ -74,7 +75,10 @@ impl WriteStrategyBuilder { /// [`LayoutStrategy`]. pub fn from_session(session: &VortexSession) -> Self { Self { - compressor: CompressorConfig::Schemes(session.permitted_schemes()), + compressor: CompressorConfig::BtrBlocks { + session: session.clone(), + options: BtrBlocksOptions::default(), + }, row_block_size: 8192, data_block_target_bytes: Some(ONE_MEG), field_writers: HashMap::new(), @@ -84,11 +88,13 @@ impl WriteStrategyBuilder { } } - /// Create a new builder whose compressor uses every scheme registered on `session`, ignoring - /// its editions. The writer defaults to this when editions are disabled. - pub fn from_session_no_editions(session: &VortexSession) -> Self { + /// 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::Schemes(session.registered_schemes()), + compressor: CompressorConfig::BtrBlocks { + session: session.clone(), + options, + }, ..Self::from_session(session) } } @@ -170,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)); @@ -178,21 +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::Schemes(schemes) => { - Arc::new(BtrBlocksCompressor(CascadingCompressor::new( - schemes - .iter() - .copied() - .filter(|scheme| scheme.id() != IntDictScheme.id()) - .collect(), - ))) - } - CompressorConfig::Opaque(compressor) => Arc::clone(compressor), - }; let compressing = CompressingStrategy::new(buffered, data_compressor); // 4. prior to compression, coalesce up to a minimum size @@ -213,12 +226,6 @@ impl WriteStrategyBuilder { ); // 2.1. | 3.1. compress stats tables and dict values. - let stats_compressor: Arc = match compressor { - CompressorConfig::Schemes(schemes) => { - Arc::new(BtrBlocksCompressor(CascadingCompressor::new(schemes))) - } - 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 8d110f21cd4..b5e2fb43958 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -66,7 +66,7 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; -use vortex_btrblocks::CascadingCompressor; +use vortex_btrblocks::BtrBlocksOptions; use vortex_btrblocks::CompressionSession; use vortex_btrblocks::CompressionSessionExt; use vortex_btrblocks::SchemeExt; @@ -2623,13 +2623,13 @@ 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 = BtrBlocksCompressor(CascadingCompressor::new( - SESSION - .permitted_schemes() - .into_iter() - .filter(|scheme| scheme.id() != StringDictScheme.id()) - .collect(), - )); + 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 diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 61376fd7344..a76ab47a6d2 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -30,6 +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::BtrBlocksOptions; use vortex_buffer::ByteBuffer; use vortex_edition::ComponentKind; use vortex_edition::EditionSessionExt; @@ -250,7 +251,14 @@ impl VortexWriteOptions { Some(strategy) => strategy, None if enforce_editions => WriteStrategyBuilder::from_session(&self.session).build(), // With editions disabled every registered encoding may be written. - None => WriteStrategyBuilder::from_session_no_editions(&self.session).build(), + None => WriteStrategyBuilder::from_session_with_options( + &self.session, + BtrBlocksOptions { + enforce_editions: false, + ..Default::default() + }, + ) + .build(), }; let dtype = stream.dtype().clone(); if enforce_editions { diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 27080e35cfc..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::from_session_no_editions(session)), + 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::from_session_no_editions(&session)), + 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::from_session_no_editions(&session)), + Arc::new(no_editions_compressor(&session)), ); let array = @@ -614,7 +627,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::from_session_no_editions(&session)), + Arc::new(no_editions_compressor(&session)), ); let array = VarBinArray::from_iter( 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 7acb397bb26..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,6 +17,7 @@ pub fn fixtures() -> Vec> { #[cfg(test)] mod tests { use vortex::VortexSessionDefault; + use vortex::compressor::BtrBlocksOptions; use vortex::compressor::CompressionSession; use vortex::editions::CORE_2026_08_3; use vortex::editions::EditionSessionExt; @@ -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") } @@ -47,14 +56,15 @@ mod tests { let regular_bytes = adapter::write_compressed_to_bytes_with_session( &session, array.clone(), - WriteStrategyBuilder::from_session_no_editions(&session).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( &compact_session, array, - WriteStrategyBuilder::from_session_no_editions(&compact_session).build(), + 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 955b013d88d..74e027d6a71 100644 --- a/vortex-test/compat-gen/src/fixtures/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/mod.rs @@ -9,6 +9,7 @@ use std::sync::Arc; use vortex::VortexSessionDefault; use vortex::array::ArrayId; use vortex::array::ArrayRef; +use vortex::compressor::BtrBlocksOptions; use vortex::compressor::CompressionSession; use vortex::file::WriteStrategyBuilder; use vortex::session::VortexSession; @@ -145,7 +146,14 @@ impl Fixture for DatasetFixtureAdapter { if self.compact { session.register(CompressionSession::compact()); } - let strategy = WriteStrategyBuilder::from_session_no_editions(&session).build(); + 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(), diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 69a1e38864f..0a483fbbff1 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -145,13 +145,11 @@ 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::CascadingCompressor; + pub use vortex_btrblocks::BtrBlocksOptions; pub use vortex_btrblocks::CompressionSession; pub use vortex_btrblocks::CompressionSessionExt; pub use vortex_btrblocks::Scheme; - pub use vortex_btrblocks::SchemeExt; pub use vortex_btrblocks::SchemeId; - pub use vortex_btrblocks::schemes; } /// Vortex editions: versioned sets of serialized components. From ae0e6b592e687812f067f3f474b2c88f3d9da657 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 21:31:09 -0400 Subject: [PATCH 06/11] Restore the facade zstd feature to develop's definition Signed-off-by: Matt Katz --- vortex/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 0bc6a351bd8..49fc0a44ddf 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -83,7 +83,7 @@ tokio = [ "vortex-io/tokio", "vortex-layout/tokio", ] -zstd = ["dep:vortex-zstd", "vortex-btrblocks/pco", "vortex-btrblocks/zstd", "vortex-file?/zstd"] +zstd = ["dep:vortex-zstd", "vortex-file?/zstd"] tensor = ["dep:vortex-tensor", "vortex-file?/tensor"] wasm-bindgen = [ "vortex-file?/wasm-bindgen", From 795d3050f6f5a63a99098b67d0ec3bae6d042bc2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 21:45:15 -0400 Subject: [PATCH 07/11] Fold the writer's default strategy arms into one Signed-off-by: Matt Katz --- vortex-file/src/writer.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index a76ab47a6d2..6248d012deb 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -249,12 +249,11 @@ impl VortexWriteOptions { }; let strategy = match self.strategy { Some(strategy) => strategy, - None if enforce_editions => WriteStrategyBuilder::from_session(&self.session).build(), // With editions disabled every registered encoding may be written. None => WriteStrategyBuilder::from_session_with_options( &self.session, BtrBlocksOptions { - enforce_editions: false, + enforce_editions, ..Default::default() }, ) From 09a98446c6e10bc17c3ff63994b38aa9b7be3a52 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 21:55:07 -0400 Subject: [PATCH 08/11] Use exclude_schemes instead of rebuilding registries minus one scheme Signed-off-by: Matt Katz --- benchmarks/string-bench/src/serialized.rs | 26 +++++++------------ vortex-btrblocks/tests/golden.rs | 28 +++++++------------- vortex-btrblocks/tests/varbin_scheme.rs | 19 +++++--------- vortex-file/src/tests.rs | 31 ++++++++++------------- 4 files changed, 38 insertions(+), 66 deletions(-) diff --git a/benchmarks/string-bench/src/serialized.rs b/benchmarks/string-bench/src/serialized.rs index ff229e44659..d18682d8747 100644 --- a/benchmarks/string-bench/src/serialized.rs +++ b/benchmarks/string-bench/src/serialized.rs @@ -25,19 +25,16 @@ use anyhow::Result; use anyhow::bail; use bytes::Bytes; use futures::TryStreamExt; -use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::VarBinViewArray; -use vortex::compressor::CompressionSession; -use vortex::compressor::CompressionSessionExt; +use vortex::compressor::BtrBlocksOptions; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; -use vortex::io::session::RuntimeSessionExt; use vortex::layout::LayoutStrategy; use vortex::session::VortexSession; use vortex_bench::Format; @@ -183,19 +180,14 @@ fn serialized_write_strategy( .filter(|&id| id != forced) .chain([DeltaScheme::default().id()]) .collect(); - // A session like `session` registering only the remaining schemes; its editions still decide - // which of them may write. - let mut registry = CompressionSession::empty(); - for scheme in session - .registered_schemes() - .into_iter() - .filter(|scheme| !excluded.contains(&scheme.id())) - { - registry.register(scheme); - } - let forced_session = VortexSession::default().with_handle(session.handle()); - forced_session.register(registry); - WriteStrategyBuilder::from_session(&forced_session).build() + 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 diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index dde3d4a12b1..73cc20c9bee 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -49,6 +49,7 @@ use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::BtrBlocksOptions; use vortex_btrblocks::CompressionSession; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::string::OnPairScheme; @@ -394,16 +395,11 @@ 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(registry: &CompressionSession) -> CompressionSession { - let mut filtered = CompressionSession::empty(); - for scheme in registry - .schemes() - .iter() - .filter(|scheme| scheme.id() != OnPairScheme.id()) - { - filtered.register(*scheme); +fn without_onpair() -> BtrBlocksOptions { + BtrBlocksOptions { + exclude_schemes: vec![OnPairScheme.id()], + ..Default::default() } - filtered } /// A session with the schemes in `registry` and `editions` enabled. @@ -428,11 +424,8 @@ fn edition_session( #[test] fn golden_regular() -> VortexResult<()> { - let session = edition_session( - &[CORE_2026_08_3], - without_onpair(&CompressionSession::default()), - )?; - let compressor = BtrBlocksCompressor::from_session(&session); + 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) } @@ -451,12 +444,9 @@ fn golden_onpair() -> VortexResult<()> { #[cfg(all(feature = "zstd", feature = "pco"))] #[test] fn golden_compact() -> VortexResult<()> { - let session = edition_session( - &[CORE_2026_08_3], - without_onpair(&CompressionSession::compact()), - )?; + let session = edition_session(&[CORE_2026_08_3], CompressionSession::compact())?; vortex_zstd::initialize(&session); session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; - let compressor = BtrBlocksCompressor::from_session(&session); + let compressor = BtrBlocksCompressor::from_session_with_options(&session, &without_onpair()); golden_corpus_snapshots("compact", &compressor) } diff --git a/vortex-btrblocks/tests/varbin_scheme.rs b/vortex-btrblocks/tests/varbin_scheme.rs index f318ccf937d..0c8d771ba4d 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -16,8 +16,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; use vortex_btrblocks::BtrBlocksOptions; -use vortex_btrblocks::CompressionSession; -use vortex_btrblocks::CompressionSessionExt; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::SchemeId; use vortex_btrblocks::schemes::binary::VarBinScheme; @@ -43,16 +41,13 @@ const N: usize = 100_000; /// The default schemes minus `excluded`. fn default_without(excluded: SchemeId) -> BtrBlocksCompressor { - let session = vortex_array::array_session().with_some(CompressionSession::empty()); - let default = CompressionSession::default(); - for scheme in default - .schemes() - .iter() - .filter(|scheme| scheme.id() != excluded) - { - session.register_scheme(*scheme); - } - no_editions_compressor(&session) + BtrBlocksCompressor::from_session_with_options( + &SESSION, + &BtrBlocksOptions { + enforce_editions: false, + exclude_schemes: vec![excluded], + }, + ) } fn lcg(state: &mut u64) -> u64 { diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index b5e2fb43958..180d596133b 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -67,8 +67,6 @@ use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; use vortex_btrblocks::BtrBlocksOptions; -use vortex_btrblocks::CompressionSession; -use vortex_btrblocks::CompressionSessionExt; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::string::StringDictScheme; use vortex_buffer::Buffer; @@ -104,9 +102,7 @@ use crate::VortexFile; use crate::WriteOptionsSessionExt; use crate::flatbuffers::footer as fb; use crate::footer::SegmentSpec; -static SESSION: LazyLock = LazyLock::new(new_session); - -fn new_session() -> VortexSession { +static SESSION: LazyLock = LazyLock::new(|| { let session = array_session() .with::() .with::(); @@ -115,7 +111,7 @@ fn new_session() -> VortexSession { crate::enable_all_registered_array_encodings(&session); session -} +}); fn strict_sorted(indices: Buffer) -> StrictSortedBuffer { StrictSortedBuffer::try_new(indices).expect("test indices should be strictly increasing") @@ -2591,20 +2587,19 @@ async fn dict_probe_honours_configured_compressor() -> VortexResult<()> { "default builder should produce a dict layout for low-cardinality strings" ); - let mut no_string_dict = CompressionSession::empty(); - for scheme in SESSION - .registered_schemes() - .into_iter() - .filter(|scheme| scheme.id() != StringDictScheme.id()) - { - no_string_dict.register(scheme); - } - let session = new_session(); - session.register(no_string_dict); let mut buf = ByteBufferMut::empty(); - let summary = session + let summary = SESSION .write_options() - .with_strategy(crate::strategy::WriteStrategyBuilder::from_session(&session).build()) + .with_strategy( + 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?; assert!( From 3f2859627dcdfa6c16eae26d26bf46580e4849ee Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 22:07:28 -0400 Subject: [PATCH 09/11] Fix the CompressionSession::cuda doc link in vortex-cuda Signed-off-by: Matt Katz --- vortex-cuda/src/layout.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 607e5b44f09..cd46c718307 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -552,6 +552,8 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { /// to use only encodings the GPU decodes. 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 strategy = WriteStrategyBuilder::from_session(session) .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); From 033636e139ab59f7d9bd0470b9cbd28ae96a805e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 22:23:03 -0400 Subject: [PATCH 10/11] don't replace the caller's schemes in the cuda sink Signed-off-by: Matt Katz --- vortex-cuda/ffi/cinclude/vortex_cuda.h | 4 ++-- vortex-cuda/ffi/src/lib.rs | 5 ++--- vortex-cuda/src/layout.rs | 10 ++++++---- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 970f628bd4e..4601478b70c 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -101,8 +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. Opening a sink restricts the session's - * compression schemes to those the GPU decodes, as `vx_cuda_session_new` already does. + * 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 35d965697b5..b99ed624d60 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -110,8 +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. Opening a sink restricts the session's -/// compression schemes to those the GPU decodes, as `vx_cuda_session_new` already does. +/// 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 /// @@ -148,7 +148,6 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( try_or(error_out, ptr::null_mut(), || { // SAFETY: The caller supplies a live borrowed session handle. let vortex_session = session_with_cuda(unsafe { vx_session_ref(session) }?); - vortex_session.register(CompressionSession::cuda()); // SAFETY: All borrowed inputs satisfy the underlying sink's requirements. unsafe { vx_array_sink_open_file_with_strategy( diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index cd46c718307..2435708b684 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -548,10 +548,12 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { /// Build a CUDA-flat writer from the schemes registered on `session` that its editions permit. /// -/// Requires [`register_cuda_layout`], and a [`CompressionSession::cuda`] registry for the file -/// to use only encodings the GPU decodes. 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 { From c6fc22e9ae381cd1a9801b65754b851c234340b2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 25 Sep 2026 00:27:26 -0400 Subject: [PATCH 11/11] Only compress with encodings the session can write; add BtrBlocksCompressor::new Signed-off-by: Matt Katz --- vortex-btrblocks/benches/compress.rs | 10 +- vortex-btrblocks/benches/compress_listview.rs | 10 +- vortex-btrblocks/src/canonical_compressor.rs | 92 +++++++++++++++++-- .../schemes/float/scheme_selection_tests.rs | 10 +- vortex-btrblocks/src/schemes/float/tests.rs | 12 +-- .../schemes/integer/scheme_selection_tests.rs | 27 +++--- vortex-btrblocks/src/schemes/integer/tests.rs | 16 ++-- .../schemes/string/scheme_selection_tests.rs | 22 ++--- vortex-btrblocks/src/schemes/string/tests.rs | 6 +- vortex-btrblocks/src/session.rs | 44 +-------- vortex-btrblocks/src/tests.rs | 49 +++++----- vortex-btrblocks/src/trace_tests.rs | 10 +- vortex-btrblocks/tests/golden.rs | 22 ++++- vortex-btrblocks/tests/onpair_roundtrip.rs | 25 ++--- vortex-btrblocks/tests/varbin_scheme.rs | 30 +++--- vortex-file/src/tests.rs | 38 ++++++++ 16 files changed, 243 insertions(+), 180 deletions(-) diff --git a/vortex-btrblocks/benches/compress.rs b/vortex-btrblocks/benches/compress.rs index 4b7c528813d..f95d37b241d 100644 --- a/vortex-btrblocks/benches/compress.rs +++ b/vortex-btrblocks/benches/compress.rs @@ -20,7 +20,7 @@ mod benchmarks { use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_btrblocks::BtrBlocksCompressor; - use vortex_btrblocks::BtrBlocksOptions; + use vortex_btrblocks::CompressionSession; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; @@ -52,13 +52,7 @@ mod benchmarks { let array = make_clickbench_window_name() .execute::(&mut ctx) .unwrap(); - let compressor = BtrBlocksCompressor::from_session_with_options( - &SESSION, - &BtrBlocksOptions { - enforce_editions: false, - ..Default::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 31f2714363a..5e12d5c2bf2 100644 --- a/vortex-btrblocks/benches/compress_listview.rs +++ b/vortex-btrblocks/benches/compress_listview.rs @@ -25,7 +25,7 @@ mod benchmarks { use vortex_array::dtype::FieldNames; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; - use vortex_btrblocks::BtrBlocksOptions; + use vortex_btrblocks::CompressionSession; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; @@ -184,13 +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::from_session_with_options( - &SESSION, - &BtrBlocksOptions { - enforce_editions: false, - ..Default::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/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index ac81310cf7b..526c2608100 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -5,22 +5,28 @@ 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::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 registered schemes whose serialized IDs the session's enabled editions - /// permit, which a file writer requires. Off, every registered scheme is used, for in-memory - /// compression where no edition applies. + /// 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, @@ -66,9 +72,17 @@ 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(CascadingCompressor::new(Vec::new())) + Self::new(Vec::new()) } /// Creates a compressor over the schemes registered on `session` whose serialized IDs the @@ -78,14 +92,14 @@ impl BtrBlocksCompressor { } /// 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 = if options.enforce_editions { - session.permitted_schemes() - } else { - session.registered_schemes() - }; + let mut schemes = writable_schemes(session, options.enforce_editions); schemes.retain(|scheme| !options.exclude_schemes.contains(&scheme.id())); - Self(CascadingCompressor::new(schemes)) + Self::new(schemes) } /// Compresses an array using BtrBlocks-inspired compression. @@ -94,6 +108,36 @@ impl BtrBlocksCompressor { } } +/// 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; @@ -101,3 +145,31 @@ impl Deref for BtrBlocksCompressor { &self.0 } } + +#[cfg(test)] +mod tests { + use vortex_array::array_session; + + use super::*; + use crate::SchemeId; + use crate::schemes::integer::FoRScheme; + use crate::schemes::integer::IntDictScheme; + + 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 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 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/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index fa384f28e53..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::tests::no_editions_compressor; +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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 bf4e75fec89..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::no_editions_compressor; +use crate::tests::default_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { - let btr = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 42288a0347b..852b05cf590 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -28,23 +28,24 @@ use vortex_session::VortexSession; use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; -use crate::CompressionSessionExt; +use crate::CompressionSession; use crate::DELTA_SCHEME; -use crate::tests::no_editions_compressor; +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 session = vortex_array::array_session(); - session.register_scheme(&DELTA_SCHEME); - no_editions_compressor(&session) + 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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -54,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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -64,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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert_eq!( @@ -93,7 +94,7 @@ fn test_sparse_compressed() -> VortexResult<()> { } } let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -117,7 +118,7 @@ fn test_dict_compressed() -> VortexResult<()> { } let array = PrimitiveArray::new(Buffer::copy_from(&codes), Validity::NonNullable); - let btr = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -130,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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -140,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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -157,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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; eprintln!("{}", compressed.display_tree()); assert!(compressed.is::()); diff --git a/vortex-btrblocks/src/schemes/integer/tests.rs b/vortex-btrblocks/src/schemes/integer/tests.rs index a7fdf71f4e4..23bfd97afe9 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -19,20 +19,20 @@ 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; use vortex_session::VortexSession; +use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; -use crate::tests::no_editions_compressor; +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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); btr.compress(&prim, &mut SESSION.create_execution_ctx())?; Ok(()) diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index 6beb6a251c5..98b2eccb8f1 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -17,8 +17,8 @@ use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::CompressionSession; -use crate::CompressionSessionExt; -use crate::tests::no_editions_compressor; +use crate::tests::compressor; +use crate::tests::default_compressor; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -27,8 +27,8 @@ fn test_constant_compressed() -> VortexResult<()> { let strings: Vec> = vec![Some("constant_value"); 100]; let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressed = no_editions_compressor(&SESSION) - .compress(&array_ref, &mut SESSION.create_execution_ctx())?; + let compressed = + default_compressor().compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -42,8 +42,8 @@ fn test_dict_compressed() -> VortexResult<()> { } let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressed = no_editions_compressor(&SESSION) - .compress(&array_ref, &mut SESSION.create_execution_ctx())?; + let compressed = + default_compressor().compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -74,8 +74,8 @@ 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 = no_editions_compressor(&SESSION) - .compress(&array_ref, &mut SESSION.create_execution_ctx())?; + let compressed = + default_compressor().compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), "expected OnPair, got {}", @@ -109,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 session = vortex_array::array_session().with_some(CompressionSession::empty()); - session.register_scheme(&FSSTScheme); - let compressor = no_editions_compressor(&session); + 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 1b79c011eb8..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::tests::no_editions_compressor; +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 = no_editions_compressor(&SESSION); + 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 = no_editions_compressor(&SESSION); + 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 index f5760155592..1ac372e1d5a 100644 --- a/vortex-btrblocks/src/session.rs +++ b/vortex-btrblocks/src/session.rs @@ -7,16 +7,14 @@ //! 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 the session's enabled editions. +//! 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_edition::ComponentKind; -use vortex_edition::EditionSessionExt; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; -use vortex_utils::aliases::hash_set::HashSet; use crate::Scheme; use crate::SchemeExt; @@ -205,33 +203,6 @@ pub trait CompressionSessionExt: SessionExt { fn register_scheme(&self, scheme: &'static dyn Scheme) { self.get_mut::().register(scheme); } - - /// The registered compression schemes in registration order. - fn registered_schemes(&self) -> Vec<&'static dyn Scheme> { - self.compression().schemes().to_vec() - } - - /// The registered schemes whose serialized IDs the enabled editions all permit. - fn permitted_schemes(&self) -> Vec<&'static dyn Scheme> { - self.permit(self.registered_schemes()) - } - - /// Keeps the schemes in `schemes` whose serialized IDs the enabled editions all permit. - fn permit(&self, schemes: Vec<&'static dyn Scheme>) -> Vec<&'static dyn Scheme> { - let allowed: HashSet<_> = self - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - schemes - .into_iter() - .filter(|scheme| { - scheme - .produced_encodings() - .iter() - .all(|id| allowed.contains(id)) - }) - .collect() - } } impl CompressionSessionExt for S {} @@ -252,7 +223,7 @@ mod tests { #[test] fn default_registers_default_schemes() { let session = array_session(); - assert_eq!(ids(&session.registered_schemes()), ids(DEFAULT_SCHEMES)); + assert_eq!(ids(session.compression().schemes()), ids(DEFAULT_SCHEMES)); } #[test] @@ -262,7 +233,7 @@ mod tests { session.register_scheme(&FloatDictScheme); session.register_scheme(&IntDictScheme); assert_eq!( - ids(&session.registered_schemes()), + ids(session.compression().schemes()), vec![IntDictScheme.id(), FloatDictScheme.id()] ); } @@ -281,11 +252,4 @@ mod tests { #[cfg(feature = "zstd")] assert!(compact.contains(&string::ZstdScheme.id())); } - - /// Without enabled editions no serialized ID is permitted, so nothing survives. - #[test] - fn no_editions_permit_nothing() { - let session = array_session(); - assert!(session.permitted_schemes().is_empty()); - } } diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index b2b88f030fd..430bc0a960c 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -29,28 +29,23 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -use crate::BtrBlocksOptions; -#[cfg(feature = "zstd")] use crate::CompressionSession; #[cfg(feature = "zstd")] -use crate::CompressionSessionExt; -#[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 every scheme registered on `session`: these tests compress in memory, where -/// no edition applies. -pub(crate) fn no_editions_compressor(session: &VortexSession) -> BtrBlocksCompressor { - BtrBlocksCompressor::from_session_with_options( - session, - &BtrBlocksOptions { - enforce_editions: false, - ..Default::default() - }, - ) +/// 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] @@ -80,8 +75,7 @@ fn listview_compress_roundtrip( ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let array_ref = input.clone().into_array(); - let result = no_editions_compressor(&SESSION) - .compress(&array_ref, &mut SESSION.create_execution_ctx())?; + let result = default_compressor().compress(&array_ref, &mut SESSION.create_execution_ctx())?; if expect_list { assert!(result.as_opt::().is_some()); } else { @@ -95,7 +89,7 @@ fn listview_compress_roundtrip( 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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -109,7 +103,7 @@ fn test_constant_all_true() -> VortexResult<()> { 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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -126,7 +120,7 @@ fn test_nullable_all_valid_compressed() -> VortexResult<()> { BitBuffer::from(vec![true; 100]), Validity::from(BitBuffer::from(vec![true; 100])), ); - let btr = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -141,7 +135,7 @@ 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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -158,7 +152,7 @@ fn test_mixed_not_constant() -> VortexResult<()> { BitBuffer::from(vec![true, false, true, false, true]), Validity::NonNullable, ); - let btr = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -173,7 +167,7 @@ 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 = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -191,7 +185,7 @@ fn test_binary_dict_compressed() -> VortexResult<()> { .map(|idx| Some(distinct_values[idx % distinct_values.len()])) .collect::>(); let array = VarBinViewArray::from_iter(values, DType::Binary(Nullability::NonNullable)); - let btr = no_editions_compressor(&SESSION); + let btr = default_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -217,8 +211,7 @@ fn test_compact_binary_zstd_compressed() -> VortexResult<()> { DType::Binary(Nullability::NonNullable), ); - let session = vortex_array::array_session().with_some(CompressionSession::compact()); - let compressor = no_editions_compressor(&session); + let compressor = compressor(&CompressionSession::compact()); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; @@ -254,9 +247,9 @@ fn test_binary_zstd_scheme_encoding( DType::Binary(Nullability::NonNullable), ); - let session = vortex_array::array_session().with_some(CompressionSession::empty()); - session.register_scheme(scheme); - let compressor = no_editions_compressor(&session); + 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)?; diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 10ad1bab745..2dc52e35cdb 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -53,9 +53,9 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use crate::CompressionSessionExt; +use crate::CompressionSession; use crate::DELTA_SCHEME; -use crate::tests::no_editions_compressor; +use crate::tests::compressor; /// A session with the default Vortex encodings registered. /// @@ -128,9 +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 { - let session = trace_session(); - session.register_scheme(&DELTA_SCHEME); - no_editions_compressor(&session).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 73cc20c9bee..5deab2f03b9 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -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,6 +49,7 @@ 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::BtrBlocksOptions; @@ -410,6 +413,24 @@ fn edition_session( 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)?; } @@ -445,7 +466,6 @@ fn golden_onpair() -> VortexResult<()> { #[test] fn golden_compact() -> VortexResult<()> { let session = edition_session(&[CORE_2026_08_3], CompressionSession::compact())?; - vortex_zstd::initialize(&session); session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; 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 16aeff20138..2600ce1bdc3 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -19,21 +19,14 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; -use vortex_btrblocks::BtrBlocksOptions; +use vortex_btrblocks::CompressionSession; use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); -/// 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() - }, - ) +/// 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 @@ -71,7 +64,7 @@ fn nonnullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = no_editions_compressor(&SESSION) + 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 @@ -114,7 +107,7 @@ fn nullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = no_editions_compressor(&SESSION) + let compressed = default_compressor() .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert OnPair specifically here — the sample-based selector may @@ -150,7 +143,7 @@ fn large_unique_short_strings_roundtrip() { ) .into_array(); - let compressed = no_editions_compressor(&SESSION) + let compressed = default_compressor() .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); @@ -179,7 +172,7 @@ fn empty_and_short_string_roundtrip() { ) .into_array(); - let compressed = no_editions_compressor(&SESSION) + let compressed = default_compressor() .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); let decoded = compressed @@ -224,7 +217,7 @@ fn delta_dict_offsets_roundtrip() { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let compressed = no_editions_compressor(&SESSION) + 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 0c8d771ba4d..33a03fdca5c 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -15,7 +15,7 @@ use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; -use vortex_btrblocks::BtrBlocksOptions; +use vortex_btrblocks::CompressionSession; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::SchemeId; use vortex_btrblocks::schemes::binary::VarBinScheme; @@ -25,28 +25,22 @@ use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); -/// 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() - }, - ) +/// 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::from_session_with_options( - &SESSION, - &BtrBlocksOptions { - enforce_editions: false, - exclude_schemes: vec![excluded], - }, + BtrBlocksCompressor::new( + CompressionSession::default() + .schemes() + .iter() + .copied() + .filter(|scheme| scheme.id() != excluded) + .collect(), ) } @@ -95,7 +89,7 @@ fn cases() -> Vec<(&'static str, ArrayRef)> { #[test] fn varbin_scheme_shrinks_binary() -> VortexResult<()> { - let with = no_editions_compressor(&SESSION); + let with = default_compressor(); let without = default_without(VarBinScheme.id()); println!( diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 180d596133b..f18c089db25 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -73,7 +73,11 @@ 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; @@ -1694,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.