From b9329d9655f1f3eca19c75b238f22679e9b10004 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 15:18:43 -0400 Subject: [PATCH 1/2] 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 | 10 +- benchmarks/string-bench/src/serialized.rs | 36 +- encodings/parquet-variant/src/vtable.rs | 11 +- fuzz/fuzz_targets/file_io.rs | 15 +- fuzz/src/array/mod.rs | 24 +- vortex-bench/src/conversions.rs | 13 +- vortex-bench/src/lib.rs | 35 +- vortex-btrblocks/Cargo.toml | 2 +- vortex-btrblocks/benches/compress.rs | 3 +- vortex-btrblocks/benches/compress_listview.rs | 3 +- vortex-btrblocks/src/builder.rs | 321 ------------------ vortex-btrblocks/src/canonical_compressor.rs | 295 ++-------------- vortex-btrblocks/src/lib.rs | 110 +++++- .../schemes/float/scheme_selection_tests.rs | 9 +- vortex-btrblocks/src/schemes/float/tests.rs | 7 +- .../schemes/integer/scheme_selection_tests.rs | 42 ++- vortex-btrblocks/src/schemes/integer/tests.rs | 11 +- vortex-btrblocks/src/schemes/string/fsst.rs | 2 +- .../schemes/string/scheme_selection_tests.rs | 30 +- vortex-btrblocks/src/schemes/string/tests.rs | 5 +- vortex-btrblocks/src/tests.rs | 255 ++++++++++++++ vortex-btrblocks/src/trace_tests.rs | 16 +- vortex-btrblocks/tests/golden.rs | 62 ++-- vortex-btrblocks/tests/onpair_roundtrip.rs | 11 +- vortex-btrblocks/tests/varbin_scheme.rs | 25 +- vortex-compressor/Cargo.toml | 3 +- vortex-compressor/src/compressor/mod.rs | 11 + vortex-compressor/src/lib.rs | 1 + vortex-compressor/src/session.rs | 133 ++++++++ vortex-cuda/src/layout.rs | 62 +++- vortex-ffi/src/sink.rs | 2 +- vortex-file/benches/split_collection.rs | 7 +- vortex-file/src/lib.rs | 2 + vortex-file/src/strategy.rs | 45 +-- vortex-file/src/tests.rs | 34 +- vortex-file/src/writer.rs | 13 +- vortex-file/tests/test_write_table.rs | 2 +- vortex-layout/src/layouts/dict/reader.rs | 9 +- vortex-layout/src/layouts/table.rs | 6 +- 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 | 28 +- wasm-test/src/main.rs | 3 +- 52 files changed, 951 insertions(+), 907 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..1d79728682d 100644 --- a/benchmarks/compress-bench/src/gpu/vortex.rs +++ b/benchmarks/compress-bench/src/gpu/vortex.rs @@ -22,7 +22,7 @@ use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::BtrBlocksCompressor; use vortex::error::VortexResult; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; @@ -37,7 +37,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 +45,7 @@ use vortex_cuda::CudaSession; use vortex_cuda::PooledFileReadAtOptions; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::CudaFlatLayoutStrategy; +use vortex_cuda::layout::cuda_compatible_schemes; use vortex_cuda::layout::register_cuda_layout; use crate::gpu::writer::GPU_ROW_GROUP_SIZE; @@ -100,11 +100,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(), + BtrBlocksCompressor::new(cuda_compatible_schemes(&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..b54bbe88e6f 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -18,7 +18,8 @@ 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_btrblocks::COMPACT_SCHEMES; +use vortex_btrblocks::CompressionSessionExt; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_file::OpenOptionsSessionExt; @@ -65,8 +66,16 @@ 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()) + WriteStrategyBuilder::from_session(&SESSION) + .with_schemes( + SESSION.permit( + SESSION + .registered_schemes() + .into_iter() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ), + ) .build(), ), }; diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index e513c5daf81..94bc1ea61f7 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -66,7 +66,9 @@ 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_btrblocks::COMPACT_SCHEMES; +#[cfg(feature = "zstd")] +use vortex_btrblocks::DEFAULT_SCHEMES; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -249,7 +251,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,14 +563,18 @@ 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() - .compress(array, ctx) - .vortex_expect("Compact compress should succeed in fuzz test"), + CompressorStrategy::Compact => BtrBlocksCompressor::new( + DEFAULT_SCHEMES + .iter() + .copied() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ) + .compress(array, ctx) + .vortex_expect("Compact compress should succeed in fuzz test"), } } @@ -579,7 +585,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/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..d6e1f0ccdf5 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::DEFAULT_SCHEMES; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; @@ -51,7 +52,7 @@ mod benchmarks { let array = make_clickbench_window_name() .execute::(&mut ctx) .unwrap(); - let compressor = BtrBlocksCompressor::default(); + let compressor = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .input_counter(|(array, _)| ItemsCount::new(array.len())) diff --git a/vortex-btrblocks/benches/compress_listview.rs b/vortex-btrblocks/benches/compress_listview.rs index 881f0f0a0eb..2bb96d0c750 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::DEFAULT_SCHEMES; use vortex_buffer::buffer_mut; use vortex_session::VortexSession; @@ -183,7 +184,7 @@ mod benchmarks { fn compress_listview(bencher: Bencher, layout: OffsetLayout) { let array = build_nested_listview(NUM_ROWS, layout); let nbytes = array.nbytes(); - let compressor = BtrBlocksCompressor::default(); + let compressor = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .input_counter(|_| ItemsCount::new(NUM_ROWS)) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs deleted file mode 100644 index 3bcda909227..00000000000 --- a/vortex-btrblocks/src/builder.rs +++ /dev/null @@ -1,321 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Builder for configuring `BtrBlocksCompressor` instances. - -use vortex_array::ArrayId; -use vortex_utils::aliases::hash_set::HashSet; - -use crate::BtrBlocksCompressor; -use crate::CascadingCompressor; -use crate::Scheme; -use crate::SchemeExt; -use crate::SchemeId; -use crate::schemes::binary; -use crate::schemes::decimal; -use crate::schemes::float; -use crate::schemes::integer; -use crate::schemes::string; -use crate::schemes::temporal; - -/// All available compression schemes. -/// -/// This list is order-sensitive: the builder preserves this order when constructing -/// the final scheme list, so that tie-breaking is deterministic. -pub const ALL_SCHEMES: &[&dyn Scheme] = &[ - //////////////////////////////////////////////////////////////////////////////////////////////// - // Integer schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - // NOTE: FoR must precede BitPacking to avoid unnecessary patches. - &integer::FoRScheme, - // NOTE: ZigZag should precede BitPacking because we don't want negative numbers. - &integer::ZigZagScheme, - &integer::BitPackingScheme, - &integer::SparseScheme, - &integer::IntDictScheme, - &integer::RunEndScheme, - &integer::SequenceScheme, - &integer::IntRLEScheme, - // Delta is omitted here: see [`DELTA_SCHEME`]. - //////////////////////////////////////////////////////////////////////////////////////////////// - // Float schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &float::ALPScheme, - &float::ALPRDScheme, - &float::FloatDictScheme, - &float::NullDominatedSparseScheme, - &float::FloatRLEScheme, - //////////////////////////////////////////////////////////////////////////////////////////////// - // String schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &string::StringDictScheme, - // Both string-fragmentation schemes are registered; the sample-based - // selector keeps whichever is smaller per column. - &string::FSSTScheme, - &string::OnPairScheme, - &string::NullDominatedSparseScheme, - //////////////////////////////////////////////////////////////////////////////////////////////// - // Binary schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &binary::BinaryDictScheme, - &binary::VarBinScheme, - // Decimal schemes. - &decimal::DecimalScheme, - // Temporal schemes. - &temporal::TemporalScheme, -]; - -/// Delta, kept out of [`ALL_SCHEMES`] because it is slower to decompress than the schemes that -/// would otherwise win. Callers that want it opt in with -/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme). -/// -/// TODO(robert): Return it to [`ALL_SCHEMES`] once we have scheme filtering. -pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); - -/// Builder for creating configured [`BtrBlocksCompressor`] instances. -/// -/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated -/// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via -/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the -/// `zstd` feature is enabled. -/// -/// # Examples -/// -/// ```rust -/// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt}; -/// use vortex_btrblocks::schemes::integer::IntDictScheme; -/// -/// // Default compressor with all schemes in ALL_SCHEMES. -/// let compressor = BtrBlocksCompressorBuilder::default().build(); -/// -/// // Remove specific schemes. -/// let compressor = BtrBlocksCompressorBuilder::default() -/// .exclude_schemes([IntDictScheme.id()]) -/// .build(); -/// ``` -#[derive(Debug, Clone)] -pub struct BtrBlocksCompressorBuilder { - schemes: Vec<&'static dyn Scheme>, -} - -impl Default for BtrBlocksCompressorBuilder { - fn default() -> Self { - Self { - schemes: ALL_SCHEMES.to_vec(), - } - } -} - -impl BtrBlocksCompressorBuilder { - /// Creates a builder with no schemes registered. - /// - /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor. - pub fn empty() -> Self { - Self { - schemes: Vec::new(), - } - } - - /// Adds an external compression scheme not in [`ALL_SCHEMES`]. - /// - /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes - /// with the compressor. - /// - /// # Panics - /// - /// Panics if a scheme with the same [`SchemeId`] is already present. - pub fn with_new_scheme(mut self, scheme: &'static dyn Scheme) -> Self { - assert!( - !self.schemes.iter().any(|s| s.id() == scheme.id()), - "scheme {:?} is already present in the builder", - scheme.id(), - ); - - self.schemes.push(scheme); - self - } - - /// Adds compact encoding schemes (Zstd for strings and binary, Pco for numerics). - /// - /// This provides better compression ratios than the default, especially for floating-point - /// heavy datasets. Requires the `zstd` feature. When the `pco` feature is also enabled, - /// Pco schemes for integers and floats are included. - /// - /// # Panics - /// - /// Panics if any of the compact schemes are already present. - #[cfg(feature = "zstd")] - pub fn with_compact(self) -> Self { - let builder = self - .with_new_scheme(&string::ZstdScheme) - .with_new_scheme(&binary::ZstdScheme); - - #[cfg(feature = "pco")] - let builder = builder - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme); - - builder - } - - /// Excludes schemes without CUDA kernel support, keeps FSST for string compression, - /// and adds Zstd for binary compression. - /// - /// Both the array-level and the buffer-level Zstd schemes are added. Buffer-level - /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression, - /// but belongs to the opt-in `zstd` edition, so callers filter the two through - /// [`retain_allowed_encodings`](Self::retain_allowed_encodings). - /// - /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a - /// larger encoded representation than the default compressor. - pub fn only_cuda_compatible(self) -> Self { - // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other - // string fragmentation and dictionary schemes still require unsupported decode paths. - #[cfg_attr(not(any(feature = "pco", feature = "zstd")), allow(unused_mut))] - let mut excluded: Vec = vec![ - integer::SparseScheme.id(), - integer::IntRLEScheme.id(), - float::ALPRDScheme.id(), - float::FloatRLEScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - string::StringDictScheme.id(), - binary::BinaryDictScheme.id(), - ]; - // Delta now has a CUDA decode kernel, so arrays that reach the GPU already encoded with - // it — the Delta children OnPair emits, for instance — decode there. It stays excluded - // from this preset until GPU delta decode is benchmarked against the schemes it would - // displace, since the preset picks encodings rather than merely decoding them. - excluded.push(integer::DeltaScheme::default().id()); - #[cfg(feature = "pco")] - excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]); - let builder = self.exclude_schemes(excluded); - - #[cfg(feature = "zstd")] - let builder = builder - .with_new_scheme(&binary::ZstdScheme) - .with_new_scheme(&binary::ZstdBuffersScheme); - - builder - } - - /// Removes the specified compression schemes by their [`SchemeId`]. - pub fn exclude_schemes(mut self, ids: impl IntoIterator) -> Self { - let ids: HashSet<_> = ids.into_iter().collect(); - self.schemes.retain(|s| !ids.contains(&s.id())); - self - } - - /// Retains only schemes whose produced serialized IDs all belong to `allowed`. - /// - /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions - /// permit. - pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes - .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); - self - } - - /// Builds the configured [`BtrBlocksCompressor`]. - pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) - } -} - -#[cfg(test)] -mod tests { - use vortex_array::VTable; - use vortex_fastlanes::FoR; - - use super::*; - - #[test] - fn empty_starts_with_no_schemes() { - let builder = BtrBlocksCompressorBuilder::empty(); - assert!(builder.schemes.is_empty()); - } - - #[test] - fn default_includes_all_schemes() { - let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn retain_allowed_encodings_filters_schemes() { - let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), 1); - assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); - - let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); - assert!(none.schemes.is_empty()); - } - - #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn cuda_compatible_excludes_alprd() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - !builder - .schemes - .iter() - .any(|s| s.id() == float::ALPRDScheme.id()) - ); - } - - /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. - #[test] - fn cuda_compatible_excludes_every_sparse_scheme() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - for excluded in [ - integer::SparseScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - ] { - assert!( - !builder.schemes.iter().any(|s| s.id() == excluded), - "{excluded} should be excluded" - ); - } - } - - #[test] - fn cuda_compatible_uses_fsst_for_strings() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::FSSTScheme.id()) - ); - #[cfg(feature = "zstd")] - assert!( - !builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::ZstdScheme.id()) - ); - } - - #[test] - #[cfg(feature = "pco")] - fn cuda_compatible_excludes_pco() { - let builder = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme) - .only_cuda_compatible(); - for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { - assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); - } - } -} diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index d93be365550..7580a013e06 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -8,28 +8,37 @@ 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::Scheme; -/// 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`]. [`from_session`](Self::from_session) +/// compresses with the schemes registered on a session that its enabled editions permit; +/// [`new`](Self::new) takes an explicit scheme list, for example [`DEFAULT_SCHEMES`] for in-memory +/// compression where no edition applies. +/// +/// [`DEFAULT_SCHEMES`]: crate::DEFAULT_SCHEMES /// /// # Examples /// /// ```rust -/// use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt}; +/// use vortex_btrblocks::{BtrBlocksCompressor, DEFAULT_SCHEMES, SchemeExt}; /// use vortex_btrblocks::schemes::integer::IntDictScheme; /// -/// // Default compressor - all schemes allowed. -/// let compressor = BtrBlocksCompressor::default(); +/// // Every default scheme, for in-memory compression. +/// let compressor = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); /// -/// // Remove specific schemes using the builder. -/// let compressor = BtrBlocksCompressorBuilder::default() -/// .exclude_schemes([IntDictScheme.id()]) -/// .build(); +/// // Every default scheme except one. +/// let compressor = BtrBlocksCompressor::new( +/// DEFAULT_SCHEMES +/// .iter() +/// .copied() +/// .filter(|scheme| scheme.id() != IntDictScheme.id()) +/// .collect(), +/// ); /// ``` #[derive(Clone)] pub struct BtrBlocksCompressor( @@ -38,6 +47,17 @@ pub struct BtrBlocksCompressor( ); impl BtrBlocksCompressor { + /// Creates a compressor over exactly `schemes`, in order. + pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { + Self(CascadingCompressor::new(schemes)) + } + + /// 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::from_session(session)) + } + /// Compresses an array using BtrBlocks-inspired compression. pub fn compress(&self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.0.compress(array, ctx) @@ -51,256 +71,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..4fad6d3a94a 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -30,8 +30,9 @@ //! //! 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. //! //! 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 +44,34 @@ //! 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, DEFAULT_SCHEMES}; //! use vortex_buffer::buffer; //! //! # fn example() -> vortex_error::VortexResult<()> { //! let session = array_session(); //! let array = PrimitiveArray::new(buffer![42u64; 1024], Validity::NonNullable).into_array(); //! -//! let compressor = BtrBlocksCompressor::default(); +//! // In memory, with no editions to respect, compress with the default schemes directly. +//! let compressor = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); //! 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 +80,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..a219d98062e 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -20,6 +20,7 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::DEFAULT_SCHEMES; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -27,7 +28,7 @@ static SESSION: LazyLock = LazyLock::new(vortex_array::array_sess fn test_constant_compressed() -> VortexResult<()> { let values: Vec = vec![42.5; 100]; let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -37,7 +38,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::new(DEFAULT_SCHEMES.to_vec()); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -50,7 +51,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::new(DEFAULT_SCHEMES.to_vec()); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert!(compressed.children()[0].is::()); @@ -69,7 +70,7 @@ fn test_null_dominated_compressed() -> VortexResult<()> { } builder.append_nulls(95); let array = builder.finish_into_primitive(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); 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..5ed7818384f 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -21,12 +21,13 @@ use vortex_fastlanes::RLE; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::DEFAULT_SCHEMES; use crate::schemes::float::FloatRLEScheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); let array = PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable).into_array(); let result = btr.compress(&array, &mut SESSION.create_execution_ctx())?; @@ -42,7 +43,7 @@ fn test_compress() -> VortexResult<()> { } let array = values.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); let compressed = btr.compress(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 1024); @@ -92,7 +93,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::new(DEFAULT_SCHEMES.to_vec()); 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..8e44ad0d5cb 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::DEFAULT_SCHEMES; use crate::DELTA_SCHEME; +use crate::Scheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +/// The default schemes plus opt-in Delta. +fn with_delta() -> BtrBlocksCompressor { + BtrBlocksCompressor::new( + DEFAULT_SCHEMES + .iter() + .copied() + .chain([&DELTA_SCHEME as &dyn Scheme]) + .collect(), + ) +} + #[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::new(DEFAULT_SCHEMES.to_vec()); 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::new(DEFAULT_SCHEMES.to_vec()); 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::new(DEFAULT_SCHEMES.to_vec()); 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::new(DEFAULT_SCHEMES.to_vec()); 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::new(DEFAULT_SCHEMES.to_vec()); 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::new(DEFAULT_SCHEMES.to_vec()); 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::new(DEFAULT_SCHEMES.to_vec()); 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::new(DEFAULT_SCHEMES.to_vec()); 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..ffddea757df 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -26,13 +26,14 @@ use vortex_sequence::Sequence; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::DEFAULT_SCHEMES; use crate::schemes::integer::IntRLEScheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); #[test] fn test_empty() -> VortexResult<()> { // Make sure empty array compression does not fail. - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); let array = PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable); let result = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; @@ -60,7 +61,7 @@ fn test_dict_encodable() -> VortexResult<()> { } } - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); let compressed = btr.compress( &codes.freeze().into_array(), &mut SESSION.create_execution_ctx(), @@ -80,7 +81,7 @@ fn constant_mostly_nulls() -> VortexResult<()> { ); let validity = array.validity()?; - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -99,7 +100,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::new(DEFAULT_SCHEMES.to_vec()); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -141,7 +142,7 @@ fn compress_large_int() -> VortexResult<()> { .collect::() .into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()); 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..51a28c8fbe0 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -17,6 +17,7 @@ use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::DEFAULT_SCHEMES; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -25,8 +26,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 = - BtrBlocksCompressor::default().compress(&array_ref, &mut SESSION.create_execution_ctx())?; + let compressed = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -40,8 +41,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::new(DEFAULT_SCHEMES.to_vec()) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -51,10 +52,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 +72,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::new(DEFAULT_SCHEMES.to_vec()) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), "expected OnPair, got {}", @@ -81,21 +82,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 +106,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 compressor = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&FSSTScheme) - .build(); + let compressor = BtrBlocksCompressor::new(vec![&FSSTScheme]); 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..96992a48c3e 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -15,6 +15,7 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::DEFAULT_SCHEMES; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -30,7 +31,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::new(DEFAULT_SCHEMES.to_vec()); let compressed = btr.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 2048); @@ -57,7 +58,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::new(DEFAULT_SCHEMES.to_vec()); 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..84c6dd66304 --- /dev/null +++ b/vortex-btrblocks/src/tests.rs @@ -0,0 +1,255 @@ +// 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::DEFAULT_SCHEMES; +#[cfg(feature = "zstd")] +use crate::Scheme; +#[cfg(feature = "zstd")] +use crate::schemes::binary; + +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::new(DEFAULT_SCHEMES.to_vec()) + .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::new(DEFAULT_SCHEMES.to_vec()); + 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::new(DEFAULT_SCHEMES.to_vec()); + 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::new(DEFAULT_SCHEMES.to_vec()); + 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::new(DEFAULT_SCHEMES.to_vec()); + 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::new(DEFAULT_SCHEMES.to_vec()); + 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::new(DEFAULT_SCHEMES.to_vec()); + 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::new(DEFAULT_SCHEMES.to_vec()); + 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 = BtrBlocksCompressor::new( + DEFAULT_SCHEMES + .iter() + .copied() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(), + ); + 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 compressor = BtrBlocksCompressor::new(vec![scheme]); + 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..9a5f911906e 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -53,8 +53,10 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use crate::BtrBlocksCompressorBuilder; +use crate::BtrBlocksCompressor; +use crate::DEFAULT_SCHEMES; use crate::DELTA_SCHEME; +use crate::Scheme; /// A session with the default Vortex encodings registered. /// @@ -127,10 +129,14 @@ fn lineitem() -> VortexResult { /// Delta is opt-in, and these traces cover the delta-encoded FSST offsets, so enable it here. fn compressed_lineitem() -> VortexResult { - BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build() - .compress(&lineitem()?, &mut execution_ctx()) + BtrBlocksCompressor::new( + DEFAULT_SCHEMES + .iter() + .copied() + .chain([&DELTA_SCHEME as &dyn Scheme]) + .collect(), + ) + .compress(&lineitem()?, &mut execution_ctx()) } fn field(array: &ArrayRef, name: &str) -> VortexResult { diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index fc636252c27..4e805cb0a27 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,13 @@ 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::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 +397,16 @@ 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 { let session = vortex_array::array_session().with::(); + vortex_btrblocks::initialize(&session); for family in EDITION_FAMILIES { session.editions().declare_family(family)?; } @@ -414,35 +419,10 @@ 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 compressor = BtrBlocksCompressor::new(without_onpair(session.permitted_schemes())); golden_corpus_snapshots("regular", &compressor) } @@ -450,7 +430,7 @@ fn golden_regular() -> VortexResult<()> { #[test] fn golden_onpair() -> VortexResult<()> { let session = edition_session(&[CORE_2026_08_3])?; - let compressor = compressor_with_onpair(&session, BtrBlocksCompressorBuilder::default()); + let compressor = BtrBlocksCompressor::from_session(&session); golden_snapshots( "onpair", &compressor, @@ -464,9 +444,11 @@ fn golden_compact() -> VortexResult<()> { let session = edition_session(&[CORE_2026_08_3])?; vortex_zstd::initialize(&session); session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; - let compressor = compressor_for_session( - &session, - BtrBlocksCompressorBuilder::default().with_compact(), - ); + let schemes = session + .registered_schemes() + .into_iter() + .chain(COMPACT_SCHEMES.iter().copied()) + .collect(); + let compressor = BtrBlocksCompressor::new(without_onpair(session.permit(schemes))); golden_corpus_snapshots("compact", &compressor) } diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index 31734d6a60e..34726082a62 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -19,6 +19,7 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::DEFAULT_SCHEMES; use vortex_session::VortexSession; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -58,7 +59,7 @@ fn nonnullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert a specific scheme — both OnPair and FSST are registered and @@ -101,7 +102,7 @@ fn nullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert OnPair specifically here — the sample-based selector may @@ -137,7 +138,7 @@ fn large_unique_short_strings_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); @@ -166,7 +167,7 @@ fn empty_and_short_string_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); let decoded = compressed @@ -211,7 +212,7 @@ fn delta_dict_offsets_roundtrip() { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()) .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..dc31728296a 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -14,8 +14,10 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::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; @@ -25,6 +27,17 @@ static SESSION: LazyLock = LazyLock::new(vortex_array::array_sess const N: usize = 100_000; +/// The default schemes minus `excluded`. +fn default_without(excluded: SchemeId) -> BtrBlocksCompressor { + BtrBlocksCompressor::new( + DEFAULT_SCHEMES + .iter() + .copied() + .filter(|scheme| scheme.id() != excluded) + .collect(), + ) +} + fn lcg(state: &mut u64) -> u64 { *state = state .wrapping_mul(6364136223846793005) @@ -70,10 +83,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::new(DEFAULT_SCHEMES.to_vec()); + let without = default_without(VarBinScheme.id()); println!( "{:<20}{:>12}{:>14}{:>14}{:>9}", @@ -116,9 +127,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/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..da742b84add 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,12 +9,15 @@ mod sample; mod select; mod structural; +use vortex_session::VortexSession; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::scheme::SchemeId; +use crate::session::CompressionSessionExt; /// Synthetic scheme ID used for the compressor's own root-level cascading. pub(crate) const ROOT_SCHEME_ID: SchemeId = SchemeId { @@ -65,6 +68,14 @@ impl CascadingCompressor { root_exclusions, } } + + /// Creates a compressor over the schemes registered on `session` whose serialized IDs the + /// session's enabled editions permit. + /// + /// See [`CompressionSessionExt::permitted_schemes`](crate::session::CompressionSessionExt). + pub fn from_session(session: &VortexSession) -> Self { + Self::new(session.permitted_schemes()) + } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. 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..5b076a6459a 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -29,10 +29,14 @@ 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::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 +557,59 @@ 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::new(schemes)) + .with_probe_compressor(BtrBlocksCompressor::new(Vec::new())) .with_row_block_size(block_rows) .with_data_block_target_bytes(None) .build() } } +/// 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 +694,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..1591d2191dc 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::DEFAULT_SCHEMES; 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(DEFAULT_SCHEMES.to_vec()) .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(DEFAULT_SCHEMES.to_vec()); 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..f11b63d5db4 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -7,7 +7,9 @@ use std::num::NonZeroUsize; use std::sync::Arc; use vortex_array::dtype::FieldPath; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::CompressionSessionExt; +use vortex_btrblocks::Scheme; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::integer::IntDictScheme; use vortex_error::VortexExpect; @@ -26,16 +28,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 +67,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 +136,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 +181,13 @@ 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::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 +211,7 @@ 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::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..09091ea942d 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -65,7 +65,8 @@ 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::CompressionSessionExt; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::string::StringDictScheme; use vortex_buffer::Buffer; @@ -1874,7 +1875,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 +2254,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 +2579,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 +2587,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 +2618,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::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..8994cf32599 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::DEFAULT_SCHEMES; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::runtime::Handle; @@ -426,7 +427,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec())), ); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -456,7 +457,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec())), ); let array = VarBinArray::from_iter( @@ -559,7 +560,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec())), ); let array = @@ -614,7 +615,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec())), ); let array = VarBinArray::from_iter( diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..55ee4174e82 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -111,14 +111,14 @@ impl TableStrategy { /// ```ignore /// # use std::sync::Arc; /// # use vortex_array::dtype::{field_path, Field, FieldPath}; - /// # use vortex_btrblocks::BtrBlocksCompressor; + /// # use vortex_btrblocks::{BtrBlocksCompressor, DEFAULT_SCHEMES}; /// # use vortex_layout::layouts::compressed::CompressingStrategy; /// # use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; /// # 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::new(DEFAULT_SCHEMES.to_vec()); + /// 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..75e50968390 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::DEFAULT_SCHEMES; 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(DEFAULT_SCHEMES.to_vec()) + .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( + DEFAULT_SCHEMES + .iter() + .copied() + .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..8b224a5f0d5 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::DEFAULT_SCHEMES; 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( + DEFAULT_SCHEMES + .iter() + .copied() + .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(DEFAULT_SCHEMES.to_vec()) + .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..0c01d73e8e2 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,16 @@ 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::schemes; } /// Vortex editions: versioned sets of serialized components. @@ -366,7 +373,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 +430,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 +497,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..50e7d87df4f 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::DEFAULT_SCHEMES; use vortex::session::VortexSession; use vortex::VortexSessionDefault; @@ -17,7 +18,7 @@ pub fn main() { let array = PrimitiveArray::new(buffer![1i32; 1024], Validity::AllValid).into_array(); let session = VortexSession::default(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::new(DEFAULT_SCHEMES.to_vec()) .compress(&array, &mut session.create_execution_ctx()) .unwrap(); println!("Compressed size: {}", compressed.len()); From 17af1335df77a3f588c23397347258bf3dde3b5f Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 24 Sep 2026 18:10:17 -0400 Subject: [PATCH 2/2] Upgrade DecimalScheme to the v2 wire format when the writer permits it Signed-off-by: Matt Katz --- vortex-btrblocks/src/lib.rs | 5 +- vortex-btrblocks/src/schemes/decimal.rs | 121 ++++++++-- vortex-btrblocks/tests/decimal_config.rs | 273 +++++++++++++++++++++++ vortex-compressor/src/scheme/mod.rs | 10 + vortex-compressor/src/session.rs | 36 ++- vortex-cuda/src/layout.rs | 15 +- vortex-file/src/tests.rs | 90 ++++++++ vortex-file/src/writer.rs | 13 +- vortex/src/lib.rs | 1 + 9 files changed, 523 insertions(+), 41 deletions(-) create mode 100644 vortex-btrblocks/tests/decimal_config.rs diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 4fad6d3a94a..3f5e71ef9ed 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -82,6 +82,7 @@ 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::session::permit_schemes; pub use vortex_compressor::stats::ArrayAndStats; pub use vortex_compressor::stats::BoolStats; pub use vortex_compressor::stats::FloatStats; @@ -138,8 +139,8 @@ pub const DEFAULT_SCHEMES: &[&dyn Scheme] = &[ //////////////////////////////////////////////////////////////////////////////////////////////// &binary::BinaryDictScheme, &binary::VarBinScheme, - // Decimal schemes. - &decimal::DecimalScheme, + // Decimal schemes. Permitting the v2 format upgrades this to the multi-part variant. + &decimal::DecimalScheme::v1(), // Temporal schemes. &temporal::TemporalScheme, ]; diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index f77a77d8c50..2e13bb3b6f6 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -9,14 +9,17 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::DecimalArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; use vortex_array::dtype::DecimalType; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsSlots; use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_decimal_byte_parts::split_decimal; use vortex_error::VortexResult; +use vortex_utils::aliases::hash_set::HashSet; use crate::ArrayAndStats; use crate::CascadingCompressor; @@ -24,12 +27,53 @@ use crate::CompressorContext; use crate::Scheme; use crate::SchemeExt; +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +enum DecimalSchemeMode { + V1, + V2, +} + +static DECIMAL_V2: DecimalScheme = DecimalScheme::v2(); + /// Compression scheme for decimal arrays via byte-part decomposition. /// -/// Narrows the decimal to the smallest integer type, compresses the underlying primitive, and wraps -/// the result in a `DecimalBytePartsArray`. +/// Narrows the decimal to the smallest integer type and compresses its byte parts independently. +/// The v1 mode leaves values wider than `i64` canonical; v2 splits them into a signed most +/// significant part and up to three unsigned lower parts. Single-part arrays serialize as v1 +/// in either mode, while arrays with lower parts serialize as v2. +/// +/// The default uses v1. Permitting both decimal IDs upgrades a registered v1 scheme to v2, see +/// [`Scheme::try_upgrade`]. A v2 scheme is dropped if either ID is not permitted. #[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct DecimalScheme; +pub struct DecimalScheme { + mode: DecimalSchemeMode, +} + +impl DecimalScheme { + /// Creates a decimal scheme configured for v1, disallowing splitting of wide decimals. + /// + /// Values that remain wider than `i64` after narrowing stay canonical. Permitting both + /// serialized IDs upgrades this scheme to v2. + pub const fn v1() -> Self { + Self { + mode: DecimalSchemeMode::V1, + } + } + + /// Creates a decimal scheme configured for v2, allowing splitting of wide decimals. + /// It is dropped if either decimal serialized ID is not permitted. + pub const fn v2() -> Self { + Self { + mode: DecimalSchemeMode::V2, + } + } +} + +impl Default for DecimalScheme { + fn default() -> Self { + Self::v1() + } +} impl Scheme for DecimalScheme { fn scheme_name(&self) -> &'static str { @@ -41,14 +85,27 @@ impl Scheme for DecimalScheme { } fn produced_encodings(&self) -> Vec { - // This scheme only builds single-part arrays, which serialize under the frozen v1 ID. - // The in-memory ID is the v2 wire ID, which no edition permits yet. - vec![decimal_byte_parts_v1_id()] + match self.mode { + DecimalSchemeMode::V1 => vec![decimal_byte_parts_v1_id()], + DecimalSchemeMode::V2 => { + vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()] + } + } } - /// Children: primitive=0. + fn try_upgrade(&self, allowed: &HashSet) -> Option<&'static dyn Scheme> { + (self.mode == DecimalSchemeMode::V1 + && allowed.contains(&decimal_byte_parts_v1_id()) + && allowed.contains(&decimal_byte_parts_v2_id())) + .then_some(&DECIMAL_V2 as &dyn Scheme) + } + + /// Children: msp=0, then up to three lower parts in v2 mode. fn num_children(&self) -> usize { - 1 + match self.mode { + DecimalSchemeMode::V1 => 1, + DecimalSchemeMode::V2 => 4, + } } fn expected_compression_ratio( @@ -68,22 +125,38 @@ impl Scheme for DecimalScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): add support splitting i128/256 buffers into chunks of primitive values - // for compression. 2 for i128 and 4 for i256. let decimal = data.array().clone().execute::(exec_ctx)?; let decimal = narrowed_decimal(decimal); - let validity = decimal.validity()?; - let prim = match decimal.values_type() { - DecimalType::I8 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I16 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I32 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I64 => PrimitiveArray::new(decimal.buffer::(), validity), - _ => return Ok(decimal.into_array()), - }; - - let compressed = - compressor.compress_child(&prim.into_array(), &compress_ctx, self.id(), 0, exec_ctx)?; - - DecimalByteParts::try_new(compressed, decimal.decimal_dtype()).map(|d| d.into_array()) + if self.mode == DecimalSchemeMode::V1 + && matches!(decimal.values_type(), DecimalType::I128 | DecimalType::I256) + { + return Ok(decimal.into_array()); + } + + let parts = split_decimal(&decimal, exec_ctx)?; + let msp = compressor.compress_child( + &parts.msp, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::MSP, + exec_ctx, + )?; + let lower_parts = parts + .lower_parts + .iter() + .enumerate() + .map(|(idx, part)| { + compressor.compress_child( + part, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx, + exec_ctx, + ) + }) + .collect::>>()?; + + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal.decimal_dtype()) + .map(IntoArray::into_array) } } diff --git a/vortex-btrblocks/tests/decimal_config.rs b/vortex-btrblocks/tests/decimal_config.rs new file mode 100644 index 00000000000..47e62ef18b9 --- /dev/null +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -0,0 +1,273 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal mode selection, serialized permissions, and compression of wide decimal parts. + +#![cfg(test)] + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Decimal; +use vortex_array::arrays::DecimalArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_array::serde::SerializeOptions; +use vortex_array::serde::SerializedArray; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_btrblocks::BtrBlocksCompressor; +use vortex_btrblocks::CompressionSessionExt; +use vortex_btrblocks::DEFAULT_SCHEMES; +use vortex_btrblocks::Scheme; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::permit_schemes; +use vortex_btrblocks::schemes::decimal::DecimalScheme; +use vortex_btrblocks::schemes::integer::BitPackingScheme; +use vortex_btrblocks::schemes::integer::FoRScheme; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_edition::ComponentKind; +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::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; +use vortex_utils::aliases::hash_set::HashSet; + +static DECIMAL_V1: DecimalScheme = DecimalScheme::v1(); +static DECIMAL_V2: DecimalScheme = DecimalScheme::v2(); + +/// Encodings registered for serialization checks; schemes and permissions come from each test. +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_decimal_byte_parts::initialize(&session); + vortex_fastlanes::initialize(&session); + session +}); + +/// Like [`SESSION`], with the latest core edition enabled: it permits decimal v1 but not v2. +static CORE_SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session().with::(); + for family in EDITION_FAMILIES { + session + .editions() + .declare_family(family) + .expect("first-party edition family"); + } + for declaration in EDITION_DECLARATIONS { + session + .register_edition(declaration) + .expect("first-party edition"); + } + session + .enable_edition(CORE_2026_08_3) + .expect("core edition is registered"); + vortex_decimal_byte_parts::initialize(&session); + vortex_fastlanes::initialize(&session); + vortex_btrblocks::initialize(&session); + session +}); + +/// The default schemes permitting exactly `ids`. +fn permitting(ids: impl IntoIterator) -> BtrBlocksCompressor { + BtrBlocksCompressor::new(permit_schemes( + DEFAULT_SCHEMES.to_vec(), + &ids.into_iter().collect(), + )) +} + +fn decimal_array(wide: bool) -> ArrayRef { + let base = if wide { 1i128 << 70 } else { 0 }; + DecimalArray::new( + (0..128i128).map(|i| base + i).collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + .into_array() +} + +fn assert_decimal_output( + compressor: BtrBlocksCompressor, + wide: bool, + expected_id: Option, +) -> VortexResult<()> { + let array = decimal_array(wide); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + if let Some(expected_id) = expected_id { + assert!(compressed.is::()); + let serialized = SESSION + .array_serialize(&compressed)? + .ok_or_else(|| vortex_err!("expected serializable decimal byte parts"))?; + assert_eq!(serialized.serialized_id, expected_id); + } else { + assert!(compressed.is::()); + } + assert_arrays_eq!(array, compressed, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::core_edition(None, Some(false))] +#[case::neither(Some(vec![]), None)] +#[case::v1(Some(vec![decimal_byte_parts_v1_id()]), Some(false))] +#[case::v2_only(Some(vec![decimal_byte_parts_v2_id()]), None)] +#[case::both(Some(vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]), Some(true))] +fn decimal_mode_follows_permissions( + #[case] ids: Option>, + #[case] v2: Option, + #[values(false, true)] wide: bool, +) -> VortexResult<()> { + let compressor = ids + .map(permitting) + .unwrap_or_else(|| BtrBlocksCompressor::from_session(&CORE_SESSION)); + let expected = match v2 { + Some(true) if wide => Some(decimal_byte_parts_v2_id()), + Some(_) if !wide => Some(decimal_byte_parts_v1_id()), + _ => None, + }; + assert_decimal_output(compressor, wide, expected) +} + +/// Permitting v2 on top of the core edition upgrades the registered v1 scheme. +#[rstest] +fn core_edition_plus_v2_upgrades(#[values(false, true)] wide: bool) -> VortexResult<()> { + let mut allowed: HashSet = CORE_SESSION + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .collect(); + allowed.insert(decimal_byte_parts_v2_id()); + let compressor = + BtrBlocksCompressor::new(permit_schemes(CORE_SESSION.registered_schemes(), &allowed)); + let expected = if wide { + decimal_byte_parts_v2_id() + } else { + decimal_byte_parts_v1_id() + }; + assert_decimal_output(compressor, wide, Some(expected)) +} + +#[rstest] +fn explicit_decimal_modes_only_upgrade( + #[values(false, true)] allowed_v2: bool, + #[values(false, true)] initial_v2: bool, + #[values(false, true)] wide: bool, +) -> VortexResult<()> { + let scheme: &'static dyn Scheme = if initial_v2 { &DECIMAL_V2 } else { &DECIMAL_V1 }; + let mut allowed = HashSet::from([decimal_byte_parts_v1_id()]); + if allowed_v2 { + allowed.insert(decimal_byte_parts_v2_id()); + } + let compressor = BtrBlocksCompressor::new(permit_schemes(vec![scheme], &allowed)); + let mode = (!initial_v2 || allowed_v2).then_some(allowed_v2); + let expected = match mode { + Some(true) if wide => Some(decimal_byte_parts_v2_id()), + Some(_) if !wide => Some(decimal_byte_parts_v1_id()), + _ => None, + }; + assert_decimal_output(compressor, wide, expected) +} + +#[test] +fn upgrades_do_not_restore_excluded_decimal() -> VortexResult<()> { + let allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); + let schemes = DEFAULT_SCHEMES + .iter() + .copied() + .filter(|scheme| scheme.id() != DecimalScheme::default().id()) + .collect(); + let compressor = BtrBlocksCompressor::new(permit_schemes(schemes, &allowed)); + assert_decimal_output(compressor, false, None) +} + +#[rstest] +#[case::i128(false, 1)] +#[case::i256(true, 3)] +fn wide_decimal_parts_roundtrip( + #[case] use_i256: bool, + #[case] lower_part_count: usize, + #[values(false, true)] negative: bool, + #[values(false, true)] nullable: bool, + #[values(false, true)] compress_children: bool, +) -> VortexResult<()> { + let validity = if nullable { + Validity::from_iter((0..2048).map(|i| i % 7 != 0)) + } else { + Validity::NonNullable + }; + let array = if use_i256 { + let values = (1..=2048u32) + .map(|i| { + let msp = if negative { + -i128::from(i) + } else { + i128::from(i) + }; + i256::from_parts( + (u128::from(i) << 64) | u128::from(i * 131 + 17), + (msp << 64) | i128::from(i * 3 + 1), + ) + }) + .collect::>(); + DecimalArray::new(values, DecimalDType::new(76, 2), validity) + } else { + let values = (1..=2048i128) + .map(|i| { + let msp = if negative { -i } else { i }; + (msp << 70) + i * 131 + 17 + }) + .collect::>(); + DecimalArray::new(values, DecimalDType::new(38, 2), validity) + } + .into_array(); + let mut allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); + if compress_children { + allowed.extend(FoRScheme.produced_encodings()); + allowed.extend(BitPackingScheme.produced_encodings()); + } + let compressor = BtrBlocksCompressor::new(permit_schemes( + vec![&DECIMAL_V2, &FoRScheme, &BitPackingScheme], + &allowed, + )); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + // Every part varies, so without child compression splitting alone cannot save space. + assert_eq!(compressed.is::(), compress_children); + if compress_children { + let parts = compressed + .as_opt::() + .ok_or_else(|| vortex_err!("expected decimal byte parts"))?; + assert_eq!(parts.lower_parts().len(), lower_part_count); + assert!(!parts.msp().is_canonical()); + assert!(parts.lower_parts().iter().all(|part| !part.is_canonical())); + } + + let array_ctx = ArrayContext::empty(); + let mut bytes = ByteBufferMut::empty(); + for buffer in compressed.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { + bytes.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + array.dtype(), + array.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + )?; + assert_arrays_eq!(array, decoded, &mut ctx); + Ok(()) +} diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 0ba1c90202a..57a9c893160 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -28,6 +28,7 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_error::VortexResult; +use vortex_utils::aliases::hash_set::HashSet; use crate::CascadingCompressor; use crate::stats::ArrayAndStats; @@ -135,6 +136,15 @@ pub trait Scheme: Debug + Send + Sync { /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. fn produced_encodings(&self) -> Vec; + /// Returns a newer variant of this scheme that the permitted serialized IDs support. + /// + /// `None` keeps this scheme. A variant must share this scheme's [`SchemeId`] and must not + /// downgrade what was registered. Every ID the variant declares in + /// [`produced_encodings`](Self::produced_encodings) must still be permitted for it to be used. + fn try_upgrade(&self, _allowed: &HashSet) -> Option<&'static dyn Scheme> { + None + } + /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies /// every scheme. diff --git a/vortex-compressor/src/session.rs b/vortex-compressor/src/session.rs index 1e0bf028014..b03aa1fee52 100644 --- a/vortex-compressor/src/session.rs +++ b/vortex-compressor/src/session.rs @@ -9,6 +9,7 @@ use std::any::Any; +use vortex_array::ArrayId; use vortex_edition::ComponentKind; use vortex_edition::EditionSessionExt; use vortex_session::SessionExt; @@ -78,26 +79,39 @@ pub trait CompressionSessionExt: SessionExt { self.permit(self.registered_schemes()) } - /// Keeps the schemes in `schemes` whose serialized IDs the enabled editions all permit. + /// Keeps the schemes in `schemes` whose serialized IDs the enabled editions all permit, + /// see [`permit_schemes`]. fn permit(&self, schemes: Vec<&'static dyn Scheme>) -> Vec<&'static dyn Scheme> { - let allowed: HashSet<_> = self + let allowed = self .enabled_component_ids(ComponentKind::Array) .into_iter() .collect(); - schemes - .into_iter() - .filter(|scheme| { - scheme - .produced_encodings() - .iter() - .all(|id| allowed.contains(id)) - }) - .collect() + permit_schemes(schemes, &allowed) } } impl CompressionSessionExt for S {} +/// Keeps the schemes in `schemes` whose serialized IDs are all in `allowed`. +/// +/// A scheme with several wire formats is first replaced by the variant it offers for `allowed`, +/// see [`Scheme::try_upgrade`], so a newer format is used exactly when the writer permits it. +pub fn permit_schemes( + schemes: Vec<&'static dyn Scheme>, + allowed: &HashSet, +) -> Vec<&'static dyn Scheme> { + schemes + .into_iter() + .map(|scheme| scheme.try_upgrade(allowed).unwrap_or(scheme)) + .filter(|scheme| { + scheme + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + }) + .collect() +} + #[cfg(test)] mod tests { use vortex_array::array_session; diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 5b076a6459a..4f7a02b8e45 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -607,7 +607,20 @@ pub fn cuda_compatible_schemes(session: &VortexSession) -> Vec<&'static dyn Sche .collect(); cuda.push(&schemes::binary::ZstdScheme); cuda.push(&schemes::binary::ZstdBuffersScheme); - session.permit(cuda) + // Multi-part DecimalByteParts arrays have no CUDA decode kernel, so wide decimals stay + // canonical: keep the single-part variant even where the editions would permit v2. + static DECIMAL_V1: schemes::decimal::DecimalScheme = schemes::decimal::DecimalScheme::v1(); + session + .permit(cuda) + .into_iter() + .map(|scheme| { + if scheme.id() == DECIMAL_V1.id() { + &DECIMAL_V1 as &dyn Scheme + } else { + scheme + } + }) + .collect() } #[derive(Clone, Debug)] diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 09091ea942d..5c312900d57 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -38,6 +38,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; +use vortex_array::dtype::i256; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::and; @@ -73,7 +74,12 @@ use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use vortex_edition::EDITION_DECLARATIONS; 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; @@ -174,6 +180,90 @@ async fn test_read_simple() { assert_eq!(row_count, 8); } +/// Wide decimals split into multi-part arrays only when the writer permits the v2 format. The +/// default writer permits every registered encoding once editions are disabled; a strategy built +/// from the session keeps the enabled editions' permissions regardless. +#[rstest] +#[case::default_writer(false, false)] +#[case::custom_layout(true, false)] +#[case::explicit_compressor(true, true)] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn decimal_writer_uses_default_or_session_permissions( + #[values(false, true)] use_i256: bool, + #[case] custom_strategy: bool, + #[case] explicit_compressor: bool, +) -> VortexResult<()> { + let session = array_session() + .with::() + .with::() + .with::(); + crate::register_default_encodings(&session); + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(CORE_2026_08_3)?; + + let array = if use_i256 { + DecimalArray::new( + (0..1024u128) + .map(|i| i256::from_parts(i * 17, 1i128 << 70)) + .collect::>(), + DecimalDType::new(76, 2), + Validity::NonNullable, + ) + } else { + DecimalArray::new( + (0..1024i128) + .map(|i| (1i128 << 70) + i * 17) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + } + .into_array(); + let strategy = crate::strategy::WriteStrategyBuilder::from_session(&session) + .with_row_block_size(256) + .with_data_block_target_bytes(None); + let strategy = if explicit_compressor { + strategy.with_schemes(session.permitted_schemes()) + } else { + strategy + } + .build(); + + for disable_editions in [false, true] { + let mut options = session.write_options(); + if custom_strategy { + options = options.with_strategy(Arc::clone(&strategy)); + } + let options = if disable_editions { + options.disable_editions() + } else { + options + }; + let mut buffer = ByteBufferMut::empty(); + options + .write(&mut buffer, array.clone().to_array_stream()) + .await?; + let actual = session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + let uses_v2 = actual.depth_first_traversal().any(|array| { + array + .as_opt::() + .is_some_and(|parts| !parts.lower_parts().is_empty()) + }); + assert_eq!(uses_v2, disable_editions && !custom_strategy); + assert_arrays_eq!(array, actual, &mut session.create_execution_ctx()); + } + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_round_trip_many_types() { diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 17711da378a..10737001c1c 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -31,6 +31,7 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; use vortex_btrblocks::CompressionSessionExt; +use vortex_btrblocks::permit_schemes; use vortex_buffer::ByteBuffer; use vortex_edition::ComponentKind; use vortex_edition::EditionSessionExt; @@ -251,9 +252,15 @@ 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 => { + let registered: HashSet = ctx.array_ctx().to_ids().into_iter().collect(); + WriteStrategyBuilder::from_session(&self.session) + .with_schemes(permit_schemes( + self.session.registered_schemes(), + ®istered, + )) + .build() + } }; let dtype = stream.dtype().clone(); if enforce_editions { diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 0c01d73e8e2..1a5fbc7b713 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -153,6 +153,7 @@ pub mod compressor { pub use vortex_btrblocks::Scheme; pub use vortex_btrblocks::SchemeExt; pub use vortex_btrblocks::SchemeId; + pub use vortex_btrblocks::permit_schemes; pub use vortex_btrblocks::schemes; }