diff --git a/Cargo.lock b/Cargo.lock index 041f9a9e2f6..a4be7981264 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10605,6 +10605,7 @@ dependencies = [ "rstest", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-fastlanes", "vortex-mask", @@ -11063,6 +11064,7 @@ dependencies = [ "rstest", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11082,6 +11084,7 @@ dependencies = [ "vortex-array", "vortex-bench-support", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11164,6 +11167,7 @@ dependencies = [ "vortex-alp", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-fastlanes", "vortex-mask", @@ -11258,6 +11262,7 @@ dependencies = [ "test-with", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11476,6 +11481,7 @@ dependencies = [ "rstest", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11519,6 +11525,7 @@ dependencies = [ "vortex-array", "vortex-arrow", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11603,6 +11610,7 @@ dependencies = [ "rstest", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11633,6 +11641,7 @@ dependencies = [ "smallvec", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11662,6 +11671,7 @@ dependencies = [ "vortex-array", "vortex-arrow", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11819,6 +11829,7 @@ dependencies = [ "rstest", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-error", "vortex-mask", "vortex-session", @@ -11835,6 +11846,7 @@ dependencies = [ "rstest", "vortex-array", "vortex-buffer", + "vortex-compressor", "vortex-edition", "vortex-error", "vortex-mask", diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index d309a2122f1..08a2ae509fc 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -28,9 +28,8 @@ cargo run -p compress-bench --profile release_debug `--gpu-decompress` is opt-in, requires the `cuda` feature, and restricts the suite to the 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 - canonical arrays. +- **Vortex**: the file is written with the CUDA output edition 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 whole read on the device: page header decode, codec decompression, dictionary/RLE/plain diff --git a/benchmarks/compress-bench/src/gpu/vortex.rs b/benchmarks/compress-bench/src/gpu/vortex.rs index bb48b461e87..9e377f43c7e 100644 --- a/benchmarks/compress-bench/src/gpu/vortex.rs +++ b/benchmarks/compress-bench/src/gpu/vortex.rs @@ -22,9 +22,14 @@ use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::BtrBlocksCompressor; +use vortex::compressor::CompressionSessionExt; +use vortex::editions::EditionSessionExt; +use vortex::editions::EnabledEditions; +use vortex::editions::cuda::CUDA_2026_09_0; use vortex::error::VortexResult; use vortex::file::OpenOptionsSessionExt; +use vortex::file::VortexFile; use vortex::file::WriteOptionsSessionExt; use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; use vortex::layout::layouts::compressed::CompressingStrategy; @@ -37,7 +42,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; @@ -100,11 +104,13 @@ 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(), + { + let compression_session = CompressionSessionExt::fork_compression(&*SESSION); + compression_session.register(EnabledEditions::default()); + EditionSessionExt::set_enabled_editions(&compression_session, [CUDA_2026_09_0]) + .expect("CUDA edition must be registered"); + BtrBlocksCompressor::from_session(&compression_session) + }, ))); let start = Instant::now(); SESSION @@ -169,7 +175,7 @@ impl Compressor for GpuVortexCompressor { /// Windows too, and the whole crate still has to compile on a developer's macOS machine. Asking /// for `--gpu-direct-io` where it cannot be honoured is an error rather than a silent no-op, /// because the flag changes what the resulting number means. -async fn open_gpu(path: &Path, direct_io: bool) -> Result { +async fn open_gpu(path: &Path, direct_io: bool) -> Result { let open_options = SESSION.open_options().with_cuda(); #[cfg(target_os = "linux")] diff --git a/benchmarks/string-bench/src/serialized.rs b/benchmarks/string-bench/src/serialized.rs index 62de3dd9a51..adfdd0833cc 100644 --- a/benchmarks/string-bench/src/serialized.rs +++ b/benchmarks/string-bench/src/serialized.rs @@ -31,7 +31,7 @@ use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; use vortex::array::arrays::VarBinViewArray; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::BtrBlocksCompressor; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; @@ -57,7 +57,7 @@ use crate::prepare_column; use crate::throughput; use crate::verify_canonicalized; -/// The btrblocks string schemes that `BtrBlocksCompressorBuilder::default()` can +/// The btrblocks string schemes that `BtrBlocksCompressor::from_session` can /// choose between. Forcing one encoder excludes every entry except its own /// scheme, so this list must track the default scheme set: add a row whenever a /// new string encoder becomes selectable by default (e.g. Zstd). @@ -170,16 +170,26 @@ 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() + let compressor = { + let compression_session = + vortex_btrblocks::CompressionSessionExt::fork_compression(session); + for id in default_string_scheme_ids() .into_iter() .filter(|&id| id != forced) - .chain([DeltaScheme::default().id()]), - ); - WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor) + .chain([DeltaScheme::default().id()]) + { + vortex_btrblocks::CompressionSessionExt::compression(&compression_session) + .unregister(id); + } + BtrBlocksCompressor::from_session(&compression_session) + }; + WriteStrategyBuilder::from_session(session) + .with_btrblocks_compressor(compressor) .build() } @@ -251,7 +261,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 +365,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::CompressionSessionExt; use vortex_btrblocks::SchemeExt; use super::*; @@ -365,8 +375,11 @@ 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 session = VortexSession::default(); + let mut actual = session + .registered_schemes() .iter() + .filter(|scheme| !scheme.id().to_string().starts_with("vortex.compressor.")) .filter(|scheme| scheme.matches(&canonical)) .map(|scheme| scheme.id()) .collect::>(); diff --git a/encodings/alp/Cargo.toml b/encodings/alp/Cargo.toml index 7da81b60d99..61d20c01058 100644 --- a/encodings/alp/Cargo.toml +++ b/encodings/alp/Cargo.toml @@ -23,6 +23,7 @@ num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-mask = { workspace = true } diff --git a/encodings/alp/src/lib.rs b/encodings/alp/src/lib.rs index 160ab568cd4..721cf62a55d 100644 --- a/encodings/alp/src/lib.rs +++ b/encodings/alp/src/lib.rs @@ -50,3 +50,6 @@ pub fn initialize(session: &VortexSession) { &compute::nan_count::ALPNanCountKernel, ); } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/encodings/alp/src/schemes/alp.rs similarity index 86% rename from vortex-btrblocks/src/schemes/float/alp.rs rename to encodings/alp/src/schemes/alp.rs index f9fc7066bf4..bd958e006bc 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/encodings/alp/src/schemes/alp.rs @@ -3,37 +3,42 @@ //! ALP (Adaptive Lossless floating-Point) encoding. -use vortex_alp::ALP; -use vortex_alp::ALPArrayExt; -use vortex_alp::ALPArraySlotsExt; -use vortex_alp::alp_encode; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::arrays::Constant; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; +use vortex_compressor::CascadingCompressor; +use vortex_compressor::compress_patches; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; -use crate::compress_patches; +use crate::ALP; +use crate::ALPArrayExt; +use crate::ALPArraySlotsExt; +use crate::alp_encode; /// ALP (Adaptive Lossless floating-Point) encoding. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct ALPScheme; impl Scheme for ALPScheme { + fn selection_priority(&self) -> u16 { + 10 + } + fn scheme_name(&self) -> &'static str { "vortex.float.alp" } @@ -43,7 +48,7 @@ impl Scheme for ALPScheme { } fn produced_encodings(&self) -> Vec { - let mut encodings = vec![ALP.id()]; + let mut encodings = vec![ALP.id(), Constant.id()]; if use_experimental_patches() { encodings.push(Patched.id()); } diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/encodings/alp/src/schemes/alprd.rs similarity index 84% rename from vortex-btrblocks/src/schemes/float/alprd.rs rename to encodings/alp/src/schemes/alprd.rs index 09c0ee85b0d..44a1adc68c4 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/encodings/alp/src/schemes/alprd.rs @@ -3,35 +3,40 @@ //! ALPRD (ALP with Real Double) encoding variant. -use vortex_alp::ALPRDArrayExt; -use vortex_alp::ALPRDArrayOwnedExt; -use vortex_alp::RDEncoder; -use vortex_alp::RDEncoderExt; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::arrays::Constant; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; +use vortex_compressor::CascadingCompressor; +use vortex_compressor::compress_patches; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; use vortex_error::vortex_panic; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::compress_patches; +use crate::ALPRDArrayExt; +use crate::ALPRDArrayOwnedExt; +use crate::RDEncoder; +use crate::RDEncoderExt; /// ALPRD (ALP with Real Double) encoding variant. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct ALPRDScheme; impl Scheme for ALPRDScheme { + fn selection_priority(&self) -> u16 { + 20 + } + fn scheme_name(&self) -> &'static str { "vortex.float.alprd" } @@ -41,7 +46,7 @@ impl Scheme for ALPRDScheme { } fn produced_encodings(&self) -> Vec { - vec![vortex_alp::ALPRD.id()] + vec![crate::ALPRD.id(), Constant.id()] } fn expected_compression_ratio( @@ -82,7 +87,7 @@ impl Scheme for ALPRDScheme { .map(|p| compress_patches(p, exec_ctx)) .transpose()?; - Ok(vortex_alp::ALPRD::try_new( + Ok(crate::ALPRD::try_new( dtype, parts.left_parts, parts.left_parts_dictionary, diff --git a/encodings/alp/src/schemes/mod.rs b/encodings/alp/src/schemes/mod.rs new file mode 100644 index 00000000000..56063daff0f --- /dev/null +++ b/encodings/alp/src/schemes/mod.rs @@ -0,0 +1,15 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod alp; +pub mod alprd; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&alp::ALPScheme); + session.register_scheme(&alprd::ALPRDScheme); +} diff --git a/encodings/datetime-parts/Cargo.toml b/encodings/datetime-parts/Cargo.toml index 19f2490573a..02f47b9230f 100644 --- a/encodings/datetime-parts/Cargo.toml +++ b/encodings/datetime-parts/Cargo.toml @@ -21,6 +21,7 @@ num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/datetime-parts/src/lib.rs b/encodings/datetime-parts/src/lib.rs index cdc65a44a71..94a112c771c 100644 --- a/encodings/datetime-parts/src/lib.rs +++ b/encodings/datetime-parts/src/lib.rs @@ -52,3 +52,6 @@ mod test { ); } } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/encodings/datetime-parts/src/schemes/mod.rs b/encodings/datetime-parts/src/schemes/mod.rs new file mode 100644 index 00000000000..e97b8019612 --- /dev/null +++ b/encodings/datetime-parts/src/schemes/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod temporal; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&temporal::TemporalScheme); +} diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/encodings/datetime-parts/src/schemes/temporal.rs similarity index 91% rename from vortex-btrblocks/src/schemes/temporal.rs rename to encodings/datetime-parts/src/schemes/temporal.rs index 79748b69450..5f70e4823bb 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/encodings/datetime-parts/src/schemes/temporal.rs @@ -17,18 +17,18 @@ use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::extension::Matcher; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::extension::datetime::TemporalMetadata; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::EstimateVerdict; -use vortex_datetime_parts::DateTimeParts; -use vortex_datetime_parts::TemporalParts; -use vortex_datetime_parts::split_temporal; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; +use crate::DateTimeParts; +use crate::TemporalParts; +use crate::split_temporal; /// Compression scheme for temporal timestamp arrays via datetime-part decomposition. /// @@ -38,6 +38,10 @@ use crate::SchemeExt; pub struct TemporalScheme; impl Scheme for TemporalScheme { + fn selection_priority(&self) -> u16 { + 10 + } + fn scheme_name(&self) -> &'static str { "vortex.ext.temporal" } diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index 6c3312be43f..6b9b8c6c42d 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -21,6 +21,7 @@ num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 2557555eac8..f68e9d8066f 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -33,3 +33,6 @@ pub fn initialize(session: &VortexSession) { &DecimalBytePartsIsConstantKernel, ); } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/encodings/decimal-byte-parts/src/schemes/decimal.rs similarity index 89% rename from vortex-btrblocks/src/schemes/decimal.rs rename to encodings/decimal-byte-parts/src/schemes/decimal.rs index f77a77d8c50..4599f58da56 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/encodings/decimal-byte-parts/src/schemes/decimal.rs @@ -12,17 +12,17 @@ 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::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::EstimateVerdict; -use vortex_decimal_byte_parts::DecimalByteParts; -use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; +use crate::DecimalByteParts; +use crate::decimal_byte_parts_v1_id; /// Compression scheme for decimal arrays via byte-part decomposition. /// @@ -32,6 +32,10 @@ use crate::SchemeExt; pub struct DecimalScheme; impl Scheme for DecimalScheme { + fn selection_priority(&self) -> u16 { + 10 + } + fn scheme_name(&self) -> &'static str { "vortex.decimal.byte_parts" } diff --git a/encodings/decimal-byte-parts/src/schemes/mod.rs b/encodings/decimal-byte-parts/src/schemes/mod.rs new file mode 100644 index 00000000000..49ca6f8bd88 --- /dev/null +++ b/encodings/decimal-byte-parts/src/schemes/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod decimal; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&decimal::DecimalScheme); +} diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index ac626a9e98c..f711fdd208f 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -28,6 +28,7 @@ prost = { workspace = true } rand = { workspace = true, optional = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 43d83c6fc7f..c5f48be360c 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -235,3 +235,6 @@ mod test { Ok(()) } } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/encodings/fastlanes/src/schemes/bitpacking.rs similarity index 88% rename from vortex-btrblocks/src/schemes/integer/bitpacking.rs rename to encodings/fastlanes/src/schemes/bitpacking.rs index 5ac7d0e4078..f9adc235761 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/encodings/fastlanes/src/schemes/bitpacking.rs @@ -9,29 +9,34 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::arrays::Constant; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::CascadingCompressor; +use vortex_compressor::compress_patches; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_fastlanes::BitPacked; -use vortex_fastlanes::bitpack_compress::bit_width_histogram; -use vortex_fastlanes::bitpack_compress::bitpack_encode; -use vortex_fastlanes::bitpack_compress::find_best_bit_width; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::compress_patches; +use crate::BitPacked; +use crate::bitpack_compress::bit_width_histogram; +use crate::bitpack_compress::bitpack_encode; +use crate::bitpack_compress::find_best_bit_width; /// BitPacking encoding for non-negative integers. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct BitPackingScheme; impl Scheme for BitPackingScheme { + fn selection_priority(&self) -> u16 { + 30 + } + fn scheme_name(&self) -> &'static str { "vortex.int.bitpacking" } @@ -41,7 +46,7 @@ impl Scheme for BitPackingScheme { } fn produced_encodings(&self) -> Vec { - let mut encodings = vec![BitPacked.id()]; + let mut encodings = vec![BitPacked.id(), Constant.id()]; if use_experimental_patches() { encodings.push(Patched.id()); } diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/encodings/fastlanes/src/schemes/delta.rs similarity index 94% rename from vortex-btrblocks/src/schemes/integer/delta.rs rename to encodings/fastlanes/src/schemes/delta.rs index 46b2f1e302e..264ef3c2786 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/encodings/fastlanes/src/schemes/delta.rs @@ -10,6 +10,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; +use vortex_compressor::CascadingCompressor; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; @@ -17,20 +18,19 @@ use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateScore; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; +use vortex_compressor::stats::GenerateStatsOptions; use vortex_error::VortexResult; -use vortex_fastlanes::Delta; -use vortex_fastlanes::FL_CHUNK_SIZE; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::GenerateStatsOptions; -use crate::Scheme; -use crate::SchemeExt; +use crate::Delta; +use crate::FL_CHUNK_SIZE; /// FastLanes Delta encoding for smooth / near-monotone integers. /// @@ -171,7 +171,7 @@ impl Scheme for DeltaScheme { // Measure the actual FastLanes transposed-delta span. This is the lane-stride // difference that gets bit-packed, not the lag-1 difference (which the transpose // makes optimistic), so it is what truly drives the compressed size. - let (_bases, deltas) = vortex_fastlanes::delta_compress(&primitive, exec_ctx)?; + let (_bases, deltas) = crate::delta_compress(&primitive, exec_ctx)?; let delta_stats = ArrayAndStats::new(deltas.into_array(), GenerateStatsOptions::default()); let span = delta_stats.integer_stats(exec_ctx).erased().max_minus_min(); @@ -201,7 +201,7 @@ impl Scheme for DeltaScheme { ) -> VortexResult { let primitive = data.array().clone().execute::(exec_ctx)?; let len = primitive.len(); - let (bases, deltas) = vortex_fastlanes::delta_compress(&primitive, exec_ctx)?; + let (bases, deltas) = crate::delta_compress(&primitive, exec_ctx)?; let compressed_bases = compressor.compress_child( &bases.into_array(), diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/encodings/fastlanes/src/schemes/float_rle.rs similarity index 87% rename from vortex-btrblocks/src/schemes/float/rle.rs rename to encodings/fastlanes/src/schemes/float_rle.rs index 71158b9dc3b..df3b87d37fb 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/encodings/fastlanes/src/schemes/float_rle.rs @@ -8,20 +8,20 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::VTable; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_fastlanes::RLE; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::schemes::integer::RUN_LENGTH_THRESHOLD; -use crate::schemes::integer::rle_compress; +use super::RUN_LENGTH_THRESHOLD; +use super::integer_rle::rle_compress; +use crate::RLE; use crate::schemes::rle_ancestor_exclusions; use crate::schemes::rle_descendant_exclusions; @@ -30,6 +30,10 @@ use crate::schemes::rle_descendant_exclusions; pub struct FloatRLEScheme; impl Scheme for FloatRLEScheme { + fn selection_priority(&self) -> u16 { + 50 + } + fn scheme_name(&self) -> &'static str { "vortex.float.rle" } diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/encodings/fastlanes/src/schemes/for_.rs similarity index 91% rename from vortex-btrblocks/src/schemes/integer/for_.rs rename to encodings/fastlanes/src/schemes/for_.rs index 476a0dec282..a0e073fa2e5 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/encodings/fastlanes/src/schemes/for_.rs @@ -10,6 +10,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; +use vortex_compressor::CascadingCompressor; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; @@ -17,25 +18,28 @@ use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_fastlanes::FoR; -use vortex_fastlanes::FoRArrayExt; -use vortex_fastlanes::FoRArraySlotsExt; -use super::BitPackingScheme; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; +use super::bitpacking::BitPackingScheme; +use crate::FoR; +use crate::FoRArrayExt; +use crate::FoRArraySlotsExt; /// Frame of Reference encoding. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct FoRScheme; impl Scheme for FoRScheme { + fn selection_priority(&self) -> u16 { + 10 + } + fn scheme_name(&self) -> &'static str { "vortex.int.for" } @@ -45,7 +49,9 @@ impl Scheme for FoRScheme { } fn produced_encodings(&self) -> Vec { - vec![FoR.id()] + let mut ids = BitPackingScheme.produced_encodings(); + ids.push(FoR.id()); + ids } /// Dict codes always start at 0, so FoR (which subtracts the min) is a no-op. diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/encodings/fastlanes/src/schemes/integer_rle.rs similarity index 92% rename from vortex-btrblocks/src/schemes/integer/rle.rs rename to encodings/fastlanes/src/schemes/integer_rle.rs index 86a40d0b36a..6c3c6aa2173 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/encodings/fastlanes/src/schemes/integer_rle.rs @@ -11,22 +11,22 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_fastlanes::RLE; -use vortex_fastlanes::RLEArrayExt; -use vortex_fastlanes::RLEArraySlotsExt; use super::RUN_LENGTH_THRESHOLD; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; +use crate::RLE; +use crate::RLEArrayExt; +use crate::RLEArraySlotsExt; use crate::schemes::rle_ancestor_exclusions; use crate::schemes::rle_descendant_exclusions; @@ -98,6 +98,10 @@ pub(crate) fn rle_compress( } impl Scheme for IntRLEScheme { + fn selection_priority(&self) -> u16 { + 80 + } + fn scheme_name(&self) -> &'static str { "vortex.int.rle" } diff --git a/encodings/fastlanes/src/schemes/mod.rs b/encodings/fastlanes/src/schemes/mod.rs new file mode 100644 index 00000000000..12bfd8dbda7 --- /dev/null +++ b/encodings/fastlanes/src/schemes/mod.rs @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +pub mod bitpacking; +pub mod delta; +pub mod float_rle; +pub mod for_; +pub mod integer_rle; + +const RUN_LENGTH_THRESHOLD: u32 = 4; + +use vortex_compressor::builtins::BinaryDictScheme; +use vortex_compressor::builtins::FloatDictScheme; +use vortex_compressor::builtins::IntDictScheme; +use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::ChildSelection; +use vortex_compressor::scheme::DescendantExclusion; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::scheme::SchemeId; +use vortex_compressor::session::CompressionSessionExt; + +/// Shared descendant exclusion rules for RLE schemes. +/// +/// RLE indices (child 1) and offsets (child 2) are monotonically increasing positions with all +/// unique values. Dict and Sparse are pointless on such data. Self-exclusion already prevents +/// RLE on RLE children. +fn rle_descendant_exclusions() -> Vec { + vec![ + DescendantExclusion { + excluded: IntDictScheme.id(), + children: ChildSelection::Many(&[1, 2]), + }, + // TODO(connor): This is wrong for some reason? + // DescendantExclusion { + // excluded: RunEndScheme.id(), + // children: ChildSelection::Many(&[1, 2]), + // }, + DescendantExclusion { + excluded: SchemeId::new("vortex.int.sparse"), + children: ChildSelection::Many(&[1, 2]), + }, + ] +} + +/// Shared ancestor exclusion rules for RLE schemes. +/// +/// Dict values (child 0) are all unique by definition, so RLE is pointless on them. +fn rle_ancestor_exclusions() -> Vec { + vec![ + AncestorExclusion { + ancestor: IntDictScheme.id(), + children: ChildSelection::One(0), + }, + AncestorExclusion { + ancestor: FloatDictScheme.id(), + children: ChildSelection::One(0), + }, + AncestorExclusion { + ancestor: StringDictScheme.id(), + children: ChildSelection::One(0), + }, + AncestorExclusion { + ancestor: BinaryDictScheme.id(), + children: ChildSelection::One(0), + }, + ] +} + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&for_::FoRScheme); + session.register_scheme(&bitpacking::BitPackingScheme); + session.register_scheme(&integer_rle::IntRLEScheme); + session.register_scheme(&float_rle::FloatRLEScheme); +} diff --git a/encodings/fsst/Cargo.toml b/encodings/fsst/Cargo.toml index 0603eee0519..982f68447f4 100644 --- a/encodings/fsst/Cargo.toml +++ b/encodings/fsst/Cargo.toml @@ -23,6 +23,7 @@ prost = { workspace = true } rand = { workspace = true, optional = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/fsst/src/lib.rs b/encodings/fsst/src/lib.rs index 70dcc705249..2c9ada95f7a 100644 --- a/encodings/fsst/src/lib.rs +++ b/encodings/fsst/src/lib.rs @@ -35,3 +35,6 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(FSST); kernel::initialize(session); } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/encodings/fsst/src/schemes/fsst.rs similarity index 87% rename from vortex-btrblocks/src/schemes/string/fsst.rs rename to encodings/fsst/src/schemes/fsst.rs index fd3fd28696a..21ce5f3a47a 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/encodings/fsst/src/schemes/fsst.rs @@ -16,32 +16,36 @@ use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_fsst::FSST; -use vortex_fsst::FSSTArrayExt; -use vortex_fsst::FSSTArraySlotsExt; -use vortex_fsst::FSSTSymbolTable; -use vortex_fsst::fsst_compress; -use vortex_fsst::fsst_train_compressor; - -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; + +use crate::FSST; +use crate::FSSTArrayExt; +use crate::FSSTArraySlotsExt; +use crate::FSSTSymbolTable; +use crate::fsst_compress; +use crate::fsst_train_compressor; /// 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 +/// the default session (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)] pub struct FSSTScheme; impl Scheme for FSSTScheme { + fn selection_priority(&self) -> u16 { + 20 + } + fn scheme_name(&self) -> &'static str { "vortex.string.fsst" } diff --git a/encodings/fsst/src/schemes/mod.rs b/encodings/fsst/src/schemes/mod.rs new file mode 100644 index 00000000000..6fc7ea9b60b --- /dev/null +++ b/encodings/fsst/src/schemes/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod fsst; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&fsst::FSSTScheme); +} diff --git a/encodings/onpair/Cargo.toml b/encodings/onpair/Cargo.toml index 4d782caa6d9..dba347078fe 100644 --- a/encodings/onpair/Cargo.toml +++ b/encodings/onpair/Cargo.toml @@ -22,6 +22,7 @@ onpair = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/onpair/src/lib.rs b/encodings/onpair/src/lib.rs index f9d8be5ff0d..740800fae38 100644 --- a/encodings/onpair/src/lib.rs +++ b/encodings/onpair/src/lib.rs @@ -38,3 +38,6 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(OnPair); kernel::initialize(session); } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/encodings/onpair/src/schemes/mod.rs b/encodings/onpair/src/schemes/mod.rs new file mode 100644 index 00000000000..56527e74acf --- /dev/null +++ b/encodings/onpair/src/schemes/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod onpair; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&onpair::OnPairScheme); +} diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/encodings/onpair/src/schemes/onpair.rs similarity index 90% rename from vortex-btrblocks/src/schemes/string/onpair.rs rename to encodings/onpair/src/schemes/onpair.rs index 06a7ff97f89..377e0fdbe2b 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/encodings/onpair/src/schemes/onpair.rs @@ -11,25 +11,25 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; use vortex_compressor::scheme::SchemeId; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_onpair::DEFAULT_CONFIG; -use vortex_onpair::OnPair; -use vortex_onpair::OnPairArrayExt; -use vortex_onpair::OnPairArraySlotsExt; -use vortex_onpair::onpair_compress; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; +use crate::DEFAULT_CONFIG; +use crate::OnPair; +use crate::OnPairArrayExt; +use crate::OnPairArraySlotsExt; +use crate::onpair_compress; /// OnPair short-string compression (dict-12). /// -/// A default string-fragmentation scheme (alongside [`super::FSSTScheme`]) — +/// A default string-fragmentation scheme (alongside FSST), /// targets large columns of short-to-medium strings with high lexical /// overlap, like URLs or log lines. Uses a learned dictionary of frequent /// adjacent substrings (built by the OnPair trainer at compress time) and @@ -40,6 +40,10 @@ use crate::SchemeExt; pub struct OnPairScheme; impl Scheme for OnPairScheme { + fn selection_priority(&self) -> u16 { + 30 + } + fn scheme_name(&self) -> &'static str { "vortex.string.onpair" } 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/encodings/pco/Cargo.toml b/encodings/pco/Cargo.toml index ded53693070..813ecec2d4b 100644 --- a/encodings/pco/Cargo.toml +++ b/encodings/pco/Cargo.toml @@ -21,6 +21,7 @@ pco = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/pco/src/lib.rs b/encodings/pco/src/lib.rs index 70d71abf7d2..606babe95a8 100644 --- a/encodings/pco/src/lib.rs +++ b/encodings/pco/src/lib.rs @@ -61,3 +61,6 @@ pub struct PcoMetadata { #[cfg(test)] mod tests; + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/encodings/pco/src/schemes/float.rs similarity index 85% rename from vortex-btrblocks/src/schemes/float/pco.rs rename to encodings/pco/src/schemes/float.rs index 416668c2fd0..b470ce981cf 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/encodings/pco/src/schemes/float.rs @@ -9,15 +9,14 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; - /// Pco (pcodec) compression for floats. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct PcoScheme; @@ -32,7 +31,7 @@ impl Scheme for PcoScheme { } fn produced_encodings(&self) -> Vec { - vec![vortex_pco::Pco.id()] + vec![crate::Pco.id()] } fn expected_compression_ratio( @@ -51,7 +50,7 @@ impl Scheme for PcoScheme { _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - Ok(vortex_pco::Pco::from_primitive( + Ok(crate::Pco::from_primitive( data.array_as_primitive(), pco::DEFAULT_COMPRESSION_LEVEL, 8192, diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/encodings/pco/src/schemes/integer.rs similarity index 87% rename from vortex-btrblocks/src/schemes/integer/pco.rs rename to encodings/pco/src/schemes/integer.rs index 675a112d44f..b142b541b1e 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/encodings/pco/src/schemes/integer.rs @@ -9,16 +9,15 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; - /// Pco (pcodec) compression for integers. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct PcoScheme; @@ -33,7 +32,7 @@ impl Scheme for PcoScheme { } fn produced_encodings(&self) -> Vec { - vec![vortex_pco::Pco.id()] + vec![crate::Pco.id()] } fn expected_compression_ratio( @@ -59,7 +58,7 @@ impl Scheme for PcoScheme { _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - Ok(vortex_pco::Pco::from_primitive( + Ok(crate::Pco::from_primitive( data.array_as_primitive(), pco::DEFAULT_COMPRESSION_LEVEL, 8192, diff --git a/encodings/pco/src/schemes/mod.rs b/encodings/pco/src/schemes/mod.rs new file mode 100644 index 00000000000..b7308a9627e --- /dev/null +++ b/encodings/pco/src/schemes/mod.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_array::session::ArraySessionExt; +use vortex_compressor::session::CompressionSessionExt; +pub mod float; +pub mod integer; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + session.arrays().register(crate::Pco); + session.register_scheme(&integer::PcoScheme); + session.register_scheme(&float::PcoScheme); +} diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 72a4581e5a6..a1399734c00 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -20,6 +20,7 @@ num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/runend/src/lib.rs b/encodings/runend/src/lib.rs index 98c661e6e36..2a2584c46b4 100644 --- a/encodings/runend/src/lib.rs +++ b/encodings/runend/src/lib.rs @@ -101,3 +101,6 @@ mod tests { ); } } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/encodings/runend/src/schemes/mod.rs b/encodings/runend/src/schemes/mod.rs new file mode 100644 index 00000000000..ecc5928b61c --- /dev/null +++ b/encodings/runend/src/schemes/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod runend; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&runend::RunEndScheme); +} diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/encodings/runend/src/schemes/runend.rs similarity index 90% rename from vortex-btrblocks/src/schemes/integer/runend.rs rename to encodings/runend/src/schemes/runend.rs index 6a97f7ec37d..1e03b16f964 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/encodings/runend/src/schemes/runend.rs @@ -10,6 +10,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; +use vortex_compressor::CascadingCompressor; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; @@ -17,20 +18,18 @@ use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::scheme::SchemeId; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_runend::RunEnd; -use vortex_runend::compress::runend_encode; -use super::IntRLEScheme; -use super::SparseScheme; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; +use crate::RunEnd; +use crate::compress::runend_encode; /// Threshold for the average run length in an array before we consider run-end encoding. const RUN_END_THRESHOLD: u32 = 4; @@ -40,6 +39,10 @@ const RUN_END_THRESHOLD: u32 = 4; pub struct RunEndScheme; impl Scheme for RunEndScheme { + fn selection_priority(&self) -> u16 { + 60 + } + fn scheme_name(&self) -> &'static str { "vortex.int.runend" } @@ -70,11 +73,11 @@ impl Scheme for RunEndScheme { children: ChildSelection::One(1), }, DescendantExclusion { - excluded: IntRLEScheme.id(), + excluded: SchemeId::new("vortex.int.rle"), children: ChildSelection::One(1), }, DescendantExclusion { - excluded: SparseScheme.id(), + excluded: SchemeId::new("vortex.int.sparse"), children: ChildSelection::One(1), }, ] diff --git a/encodings/sequence/Cargo.toml b/encodings/sequence/Cargo.toml index e307d07f075..3c35d4f11ed 100644 --- a/encodings/sequence/Cargo.toml +++ b/encodings/sequence/Cargo.toml @@ -19,6 +19,7 @@ prost = { workspace = true } smallvec = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/sequence/src/lib.rs b/encodings/sequence/src/lib.rs index 6c1a4acc5e8..a7a004578e2 100644 --- a/encodings/sequence/src/lib.rs +++ b/encodings/sequence/src/lib.rs @@ -48,3 +48,6 @@ pub fn initialize(session: &VortexSession) { // TODO(joe): hook up to the compressor // TODO(joe): support comparisons with other operators // TODO(joe): support list in expr pushdown + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/encodings/sequence/src/schemes/mod.rs b/encodings/sequence/src/schemes/mod.rs new file mode 100644 index 00000000000..a4a9d0c4056 --- /dev/null +++ b/encodings/sequence/src/schemes/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod sequence; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&sequence::SequenceScheme); +} diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/encodings/sequence/src/schemes/sequence.rs similarity index 94% rename from vortex-btrblocks/src/schemes/integer/sequence.rs rename to encodings/sequence/src/schemes/sequence.rs index edcefb99fc2..61d4ef5c32c 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/encodings/sequence/src/schemes/sequence.rs @@ -8,6 +8,7 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::VTable; +use vortex_compressor::CascadingCompressor; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; @@ -15,26 +16,29 @@ use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateScore; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_sequence::Sequence; -use vortex_sequence::sequence_encode; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; +use crate::Sequence; +use crate::sequence_encode; /// Sequence encoding for sequential patterns. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct SequenceScheme; impl Scheme for SequenceScheme { + fn selection_priority(&self) -> u16 { + 70 + } + fn scheme_name(&self) -> &'static str { "vortex.int.sequence" } diff --git a/encodings/sparse/Cargo.toml b/encodings/sparse/Cargo.toml index 8f5db518478..945405f302a 100644 --- a/encodings/sparse/Cargo.toml +++ b/encodings/sparse/Cargo.toml @@ -22,6 +22,7 @@ num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/sparse/src/lib.rs b/encodings/sparse/src/lib.rs index 79edc647c37..7f45cf9f231 100644 --- a/encodings/sparse/src/lib.rs +++ b/encodings/sparse/src/lib.rs @@ -1053,3 +1053,6 @@ mod test { assert_eq!(actual, expected); } } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/encodings/sparse/src/schemes/float.rs similarity index 90% rename from vortex-btrblocks/src/schemes/float/sparse.rs rename to encodings/sparse/src/schemes/float.rs index 3d9c25b18e4..124d216017d 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/encodings/sparse/src/schemes/float.rs @@ -11,20 +11,20 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_sparse::Sparse; -use vortex_sparse::SparseExt as _; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; -use crate::schemes::integer::SparseScheme as IntSparseScheme; +use super::integer::SparseScheme as IntSparseScheme; +use crate::Sparse; +use crate::SparseExt as _; /// Sparse encoding for null-dominated float arrays. /// @@ -33,6 +33,10 @@ use crate::schemes::integer::SparseScheme as IntSparseScheme; pub struct NullDominatedSparseScheme; impl Scheme for NullDominatedSparseScheme { + fn selection_priority(&self) -> u16 { + 40 + } + fn scheme_name(&self) -> &'static str { "vortex.float.sparse" } diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/encodings/sparse/src/schemes/integer.rs similarity index 92% rename from vortex-btrblocks/src/schemes/integer/sparse.rs rename to encodings/sparse/src/schemes/integer.rs index 429ff5c1a31..92ac6720d0d 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/encodings/sparse/src/schemes/integer.rs @@ -14,30 +14,33 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::scalar::Scalar; +use vortex_compressor::CascadingCompressor; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::scheme::SchemeId; +use vortex_compressor::stats::ArrayAndStats; +use vortex_compressor::stats::GenerateStatsOptions; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_sparse::Sparse; -use vortex_sparse::SparseExt as _; - -use super::IntRLEScheme; -use super::RunEndScheme; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::GenerateStatsOptions; -use crate::Scheme; -use crate::SchemeExt; + +use crate::Sparse; +use crate::SparseExt as _; /// Sparse encoding for single-value-dominated arrays. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct SparseScheme; impl Scheme for SparseScheme { + fn selection_priority(&self) -> u16 { + 40 + } + fn scheme_name(&self) -> &'static str { "vortex.int.sparse" } @@ -70,11 +73,11 @@ impl Scheme for SparseScheme { children: ChildSelection::One(1), }, DescendantExclusion { - excluded: RunEndScheme.id(), + excluded: SchemeId::new("vortex.int.runend"), children: ChildSelection::One(1), }, DescendantExclusion { - excluded: IntRLEScheme.id(), + excluded: SchemeId::new("vortex.int.rle"), children: ChildSelection::One(1), }, DescendantExclusion { diff --git a/encodings/sparse/src/schemes/mod.rs b/encodings/sparse/src/schemes/mod.rs new file mode 100644 index 00000000000..f192a242638 --- /dev/null +++ b/encodings/sparse/src/schemes/mod.rs @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod float; +pub mod integer; +pub mod string; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&integer::SparseScheme); + session.register_scheme(&float::NullDominatedSparseScheme); + session.register_scheme(&string::NullDominatedSparseScheme); +} diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/encodings/sparse/src/schemes/string.rs similarity index 90% rename from vortex-btrblocks/src/schemes/string/sparse.rs rename to encodings/sparse/src/schemes/string.rs index 8620c366f77..4febf1cd102 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/encodings/sparse/src/schemes/string.rs @@ -11,21 +11,21 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::CascadingCompressor; +use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_sparse::Sparse; -use vortex_sparse::SparseExt as _; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; -use crate::schemes::integer::IntDictScheme; -use crate::schemes::integer::SparseScheme as IntSparseScheme; +use super::integer::SparseScheme as IntSparseScheme; +use crate::Sparse; +use crate::SparseExt as _; /// Sparse encoding for null-dominated arrays. /// @@ -34,6 +34,10 @@ use crate::schemes::integer::SparseScheme as IntSparseScheme; pub struct NullDominatedSparseScheme; impl Scheme for NullDominatedSparseScheme { + fn selection_priority(&self) -> u16 { + 40 + } + fn scheme_name(&self) -> &'static str { "vortex.string.sparse" } diff --git a/encodings/zigzag/Cargo.toml b/encodings/zigzag/Cargo.toml index 83e6d324ed1..e5b7d26ea8d 100644 --- a/encodings/zigzag/Cargo.toml +++ b/encodings/zigzag/Cargo.toml @@ -16,6 +16,7 @@ version = { workspace = true } [dependencies] vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/encodings/zigzag/src/lib.rs b/encodings/zigzag/src/lib.rs index 01c3f51541b..88783b28547 100644 --- a/encodings/zigzag/src/lib.rs +++ b/encodings/zigzag/src/lib.rs @@ -18,3 +18,6 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(ZigZag); kernel::initialize(session); } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/encodings/zigzag/src/schemes/mod.rs b/encodings/zigzag/src/schemes/mod.rs new file mode 100644 index 00000000000..0d9069c9103 --- /dev/null +++ b/encodings/zigzag/src/schemes/mod.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod zigzag; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&zigzag::ZigZagScheme); +} diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/encodings/zigzag/src/schemes/zigzag.rs similarity index 89% rename from vortex-btrblocks/src/schemes/integer/zigzag.rs rename to encodings/zigzag/src/schemes/zigzag.rs index 0e4be01845a..f81d87e2d82 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/encodings/zigzag/src/schemes/zigzag.rs @@ -10,6 +10,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; +use vortex_compressor::CascadingCompressor; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; @@ -17,27 +18,29 @@ use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeExt; +use vortex_compressor::scheme::SchemeId; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use vortex_zigzag::ZigZag; -use vortex_zigzag::ZigZagArraySlotsExt; -use vortex_zigzag::zigzag_encode; - -use super::RunEndScheme; -use super::SparseScheme; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; + +use crate::ZigZag; +use crate::ZigZagArraySlotsExt; +use crate::zigzag_encode; /// ZigZag encoding for negative integers. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct ZigZagScheme; impl Scheme for ZigZagScheme { + fn selection_priority(&self) -> u16 { + 20 + } + fn scheme_name(&self) -> &'static str { "vortex.int.zigzag" } @@ -65,11 +68,11 @@ impl Scheme for ZigZagScheme { children: ChildSelection::All, }, DescendantExclusion { - excluded: RunEndScheme.id(), + excluded: SchemeId::new("vortex.int.runend"), children: ChildSelection::All, }, DescendantExclusion { - excluded: SparseScheme.id(), + excluded: SchemeId::new("vortex.int.sparse"), children: ChildSelection::All, }, ] diff --git a/encodings/zstd/Cargo.toml b/encodings/zstd/Cargo.toml index 6ab7b2ab31b..dbb17c8446d 100644 --- a/encodings/zstd/Cargo.toml +++ b/encodings/zstd/Cargo.toml @@ -25,6 +25,7 @@ num-traits = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-compressor = { workspace = true } vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } diff --git a/encodings/zstd/src/lib.rs b/encodings/zstd/src/lib.rs index 9280d99958c..f3146996f17 100644 --- a/encodings/zstd/src/lib.rs +++ b/encodings/zstd/src/lib.rs @@ -122,3 +122,6 @@ pub struct ZstdBuffersMetadata { #[prost(uint64, repeated, tag = "6")] pub child_lens: Vec, } + +/// Compression schemes and their session registration. +pub mod schemes; diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/encodings/zstd/src/schemes/binary.rs similarity index 83% rename from vortex-btrblocks/src/schemes/binary/zstd.rs rename to encodings/zstd/src/schemes/binary.rs index d652e344db2..9c6da830b4c 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/encodings/zstd/src/schemes/binary.rs @@ -9,15 +9,14 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; - /// Zstd compression without dictionaries for binary arrays. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct ZstdScheme; @@ -32,7 +31,7 @@ impl Scheme for ZstdScheme { } fn produced_encodings(&self) -> Vec { - vec![vortex_zstd::Zstd.id()] + vec![crate::Zstd.id()] } fn expected_compression_ratio( @@ -56,7 +55,7 @@ impl Scheme for ZstdScheme { .into_owned() .compact_buffers(exec_ctx)?; Ok( - vortex_zstd::Zstd::from_var_bin_view_without_dict(&compacted, 3, 8192, exec_ctx)? + crate::Zstd::from_var_bin_view_without_dict(&compacted, 3, 8192, exec_ctx)? .into_array(), ) } diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/encodings/zstd/src/schemes/binary_buffers.rs similarity index 82% rename from vortex-btrblocks/src/schemes/binary/zstd_buffers.rs rename to encodings/zstd/src/schemes/binary_buffers.rs index 3f06d65b061..0700ff72a48 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/encodings/zstd/src/schemes/binary_buffers.rs @@ -9,15 +9,14 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; - /// Zstd buffer-level compression preserving array layout for GPU decompression. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct ZstdBuffersScheme; @@ -32,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { } fn produced_encodings(&self) -> Vec { - vec![vortex_zstd::ZstdBuffers.id()] + vec![crate::ZstdBuffers.id()] } fn expected_compression_ratio( @@ -51,6 +50,6 @@ impl Scheme for ZstdBuffersScheme { _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - Ok(vortex_zstd::ZstdBuffers::compress(data.array(), 3, exec_ctx.session())?.into_array()) + Ok(crate::ZstdBuffers::compress(data.array(), 3, exec_ctx.session())?.into_array()) } } diff --git a/encodings/zstd/src/schemes/mod.rs b/encodings/zstd/src/schemes/mod.rs new file mode 100644 index 00000000000..89f6e412e19 --- /dev/null +++ b/encodings/zstd/src/schemes/mod.rs @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compression schemes owned by this encoding package. + +use vortex_compressor::session::CompressionSessionExt; +pub mod binary; +pub mod binary_buffers; +pub mod string; +pub mod string_buffers; + +/// Register the encoding plugins and their compression schemes. +pub fn initialize(session: &vortex_session::VortexSession) { + crate::initialize(session); + session.register_scheme(&string::ZstdScheme); + session.register_scheme(&binary::ZstdScheme); +} diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/encodings/zstd/src/schemes/string.rs similarity index 83% rename from vortex-btrblocks/src/schemes/string/zstd.rs rename to encodings/zstd/src/schemes/string.rs index 84e8860d626..99dd83c24b1 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/encodings/zstd/src/schemes/string.rs @@ -9,15 +9,14 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; - /// Zstd compression without dictionaries (nvCOMP compatible). #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct ZstdScheme; @@ -32,7 +31,7 @@ impl Scheme for ZstdScheme { } fn produced_encodings(&self) -> Vec { - vec![vortex_zstd::Zstd.id()] + vec![crate::Zstd.id()] } fn expected_compression_ratio( @@ -56,7 +55,7 @@ impl Scheme for ZstdScheme { .into_owned() .compact_buffers(exec_ctx)?; Ok( - vortex_zstd::Zstd::from_var_bin_view_without_dict(&compacted, 3, 8192, exec_ctx)? + crate::Zstd::from_var_bin_view_without_dict(&compacted, 3, 8192, exec_ctx)? .into_array(), ) } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/encodings/zstd/src/schemes/string_buffers.rs similarity index 82% rename from vortex-btrblocks/src/schemes/string/zstd_buffers.rs rename to encodings/zstd/src/schemes/string_buffers.rs index cf691c70fcb..9143c42888d 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/encodings/zstd/src/schemes/string_buffers.rs @@ -9,15 +9,14 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::CompressorContext; use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::Scheme; +use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; - /// Zstd buffer-level compression preserving array layout for GPU decompression. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct ZstdBuffersScheme; @@ -32,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { } fn produced_encodings(&self) -> Vec { - vec![vortex_zstd::ZstdBuffers.id()] + vec![crate::ZstdBuffers.id()] } fn expected_compression_ratio( @@ -51,6 +50,6 @@ impl Scheme for ZstdBuffersScheme { _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - Ok(vortex_zstd::ZstdBuffers::compress(data.array(), 3, exec_ctx.session())?.into_array()) + Ok(crate::ZstdBuffers::compress(data.array(), 3, exec_ctx.session())?.into_array()) } } diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index 6d9c8906fc9..2cf9c5fda85 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -18,7 +18,7 @@ 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::BtrBlocksCompressor; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_file::OpenOptionsSessionExt; @@ -65,8 +65,13 @@ 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_btrblocks_compressor({ + let compression_session = + vortex_btrblocks::CompressionSessionExt::fork_compression(&*SESSION); + vortex_btrblocks::initialize_compact(&compression_session); + BtrBlocksCompressor::from_session(&compression_session) + }) .build(), ), }; diff --git a/fuzz/src/array/mod.rs b/fuzz/src/array/mod.rs index e513c5daf81..8b12367598b 100644 --- a/fuzz/src/array/mod.rs +++ b/fuzz/src/array/mod.rs @@ -65,8 +65,6 @@ use vortex_array::search_sorted::SearchResult; use vortex_array::search_sorted::SearchSorted; use vortex_array::search_sorted::SearchSortedSide; use vortex_btrblocks::BtrBlocksCompressor; -#[cfg(feature = "zstd")] -use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_mask::Mask; @@ -249,7 +247,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 +559,17 @@ 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 => { + let compression_session = + vortex_btrblocks::CompressionSessionExt::fork_compression(&*SESSION); + vortex_btrblocks::initialize_compact(&compression_session); + BtrBlocksCompressor::from_session(&compression_session) + } + .compress(array, ctx) + .vortex_expect("Compact compress should succeed in fuzz test"), } } @@ -579,7 +580,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..f071a7240ea 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -37,7 +37,8 @@ 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::compressor::CompressionSessionExt; use vortex::dtype::DType; use vortex::dtype::FieldPath; use vortex::dtype::StructFields; @@ -66,7 +67,6 @@ use wkb::writer::write_geometry; use crate::CompactionStrategy; use crate::Format; use crate::SESSION; -use crate::retain_edition_encodings; use crate::utils::file::idempotent_async; /// Memory budget per concurrent conversion stream in GB. This is somewhat arbitary. @@ -246,12 +246,13 @@ 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_btrblocks_compressor({ + let compression_session = CompressionSessionExt::fork_compression(&*SESSION); + vortex::compressor::initialize_compact(&compression_session); + BtrBlocksCompressor::from_session(&compression_session) + }); } for name in binary_fields { builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); @@ -263,7 +264,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..276e2048875 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -28,7 +28,8 @@ use tpcds::TpcDsBenchmark; use tpch::benchmark::TpcHBenchmark; pub use utils::file::*; pub use utils::logging::*; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::BtrBlocksCompressor; +use vortex::compressor::CompressionSessionExt; use vortex::error::VortexExpect; use vortex::error::vortex_err; use vortex::file::VortexWriteOptions; @@ -70,8 +71,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 +253,13 @@ 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_btrblocks_compressor({ + let compression_session = + CompressionSessionExt::fork_compression(&*SESSION); + vortex::compressor::initialize_compact(&compression_session); + BtrBlocksCompressor::from_session(&compression_session) + }) .build(), ), CompactionStrategy::Default => options, @@ -266,21 +267,6 @@ 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) -} - /// Verify that local data has already been prepared for the requested benchmark formats. /// /// Engine-specific benchmark binaries call this before running queries. Data generation itself diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 24a03337768..f59f05bb5ed 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -26,6 +26,7 @@ vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } +vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-fsst = { workspace = true } @@ -33,6 +34,7 @@ vortex-onpair = { workspace = true } vortex-pco = { workspace = true, optional = true } vortex-runend = { workspace = true } vortex-sequence = { workspace = true } +vortex-session = { workspace = true } vortex-sparse = { workspace = true } vortex-utils = { workspace = true } vortex-zigzag = { workspace = true } @@ -49,9 +51,7 @@ tpchgen = { workspace = true } tpchgen-arrow = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-arrow = { workspace = true } -vortex-edition = { workspace = true } vortex-mask = { workspace = true } -vortex-session = { workspace = true } [features] pco = ["dep:pco", "dep:vortex-pco"] diff --git a/vortex-btrblocks/benches/compress.rs b/vortex-btrblocks/benches/compress.rs index b088c257489..bde8192b303 100644 --- a/vortex-btrblocks/benches/compress.rs +++ b/vortex-btrblocks/benches/compress.rs @@ -24,7 +24,11 @@ mod benchmarks { use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session + }); fn make_clickbench_window_name() -> ArrayRef { // A test that's meant to mirror the WindowName column from ClickBench. @@ -51,7 +55,7 @@ mod benchmarks { let array = make_clickbench_window_name() .execute::(&mut ctx) .unwrap(); - let compressor = BtrBlocksCompressor::default(); + let compressor = BtrBlocksCompressor::for_memory(&SESSION); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .input_counter(|(array, _)| ItemsCount::new(array.len())) diff --git a/vortex-btrblocks/benches/compress_listview.rs b/vortex-btrblocks/benches/compress_listview.rs index 881f0f0a0eb..ae92fd27535 100644 --- a/vortex-btrblocks/benches/compress_listview.rs +++ b/vortex-btrblocks/benches/compress_listview.rs @@ -31,7 +31,11 @@ mod benchmarks { const NUM_ROWS: usize = 8192; const SEED: u64 = 42; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session + }); const SHORT_STRINGS: &[&str] = &[ "alpha_one", @@ -183,7 +187,7 @@ mod benchmarks { fn compress_listview(bencher: Bencher, layout: OffsetLayout) { let array = build_nested_listview(NUM_ROWS, layout); let nbytes = array.nbytes(); - let compressor = BtrBlocksCompressor::default(); + let compressor = BtrBlocksCompressor::for_memory(&SESSION); bencher .with_inputs(|| (&array, SESSION.create_execution_ctx())) .input_counter(|_| ItemsCount::new(NUM_ROWS)) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs deleted file mode 100644 index 3bcda909227..00000000000 --- a/vortex-btrblocks/src/builder.rs +++ /dev/null @@ -1,321 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Builder for configuring `BtrBlocksCompressor` instances. - -use vortex_array::ArrayId; -use vortex_utils::aliases::hash_set::HashSet; - -use crate::BtrBlocksCompressor; -use crate::CascadingCompressor; -use crate::Scheme; -use crate::SchemeExt; -use crate::SchemeId; -use crate::schemes::binary; -use crate::schemes::decimal; -use crate::schemes::float; -use crate::schemes::integer; -use crate::schemes::string; -use crate::schemes::temporal; - -/// All available compression schemes. -/// -/// This list is order-sensitive: the builder preserves this order when constructing -/// the final scheme list, so that tie-breaking is deterministic. -pub const ALL_SCHEMES: &[&dyn Scheme] = &[ - //////////////////////////////////////////////////////////////////////////////////////////////// - // Integer schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - // NOTE: FoR must precede BitPacking to avoid unnecessary patches. - &integer::FoRScheme, - // NOTE: ZigZag should precede BitPacking because we don't want negative numbers. - &integer::ZigZagScheme, - &integer::BitPackingScheme, - &integer::SparseScheme, - &integer::IntDictScheme, - &integer::RunEndScheme, - &integer::SequenceScheme, - &integer::IntRLEScheme, - // Delta is omitted here: see [`DELTA_SCHEME`]. - //////////////////////////////////////////////////////////////////////////////////////////////// - // Float schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &float::ALPScheme, - &float::ALPRDScheme, - &float::FloatDictScheme, - &float::NullDominatedSparseScheme, - &float::FloatRLEScheme, - //////////////////////////////////////////////////////////////////////////////////////////////// - // String schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &string::StringDictScheme, - // Both string-fragmentation schemes are registered; the sample-based - // selector keeps whichever is smaller per column. - &string::FSSTScheme, - &string::OnPairScheme, - &string::NullDominatedSparseScheme, - //////////////////////////////////////////////////////////////////////////////////////////////// - // Binary schemes. - //////////////////////////////////////////////////////////////////////////////////////////////// - &binary::BinaryDictScheme, - &binary::VarBinScheme, - // Decimal schemes. - &decimal::DecimalScheme, - // Temporal schemes. - &temporal::TemporalScheme, -]; - -/// Delta, kept out of [`ALL_SCHEMES`] because it is slower to decompress than the schemes that -/// would otherwise win. Callers that want it opt in with -/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme). -/// -/// TODO(robert): Return it to [`ALL_SCHEMES`] once we have scheme filtering. -pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); - -/// Builder for creating configured [`BtrBlocksCompressor`] instances. -/// -/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated -/// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via -/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the -/// `zstd` feature is enabled. -/// -/// # Examples -/// -/// ```rust -/// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt}; -/// use vortex_btrblocks::schemes::integer::IntDictScheme; -/// -/// // Default compressor with all schemes in ALL_SCHEMES. -/// let compressor = BtrBlocksCompressorBuilder::default().build(); -/// -/// // Remove specific schemes. -/// let compressor = BtrBlocksCompressorBuilder::default() -/// .exclude_schemes([IntDictScheme.id()]) -/// .build(); -/// ``` -#[derive(Debug, Clone)] -pub struct BtrBlocksCompressorBuilder { - schemes: Vec<&'static dyn Scheme>, -} - -impl Default for BtrBlocksCompressorBuilder { - fn default() -> Self { - Self { - schemes: ALL_SCHEMES.to_vec(), - } - } -} - -impl BtrBlocksCompressorBuilder { - /// Creates a builder with no schemes registered. - /// - /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor. - pub fn empty() -> Self { - Self { - schemes: Vec::new(), - } - } - - /// Adds an external compression scheme not in [`ALL_SCHEMES`]. - /// - /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes - /// with the compressor. - /// - /// # Panics - /// - /// Panics if a scheme with the same [`SchemeId`] is already present. - pub fn with_new_scheme(mut self, scheme: &'static dyn Scheme) -> Self { - assert!( - !self.schemes.iter().any(|s| s.id() == scheme.id()), - "scheme {:?} is already present in the builder", - scheme.id(), - ); - - self.schemes.push(scheme); - self - } - - /// Adds compact encoding schemes (Zstd for strings and binary, Pco for numerics). - /// - /// This provides better compression ratios than the default, especially for floating-point - /// heavy datasets. Requires the `zstd` feature. When the `pco` feature is also enabled, - /// Pco schemes for integers and floats are included. - /// - /// # Panics - /// - /// Panics if any of the compact schemes are already present. - #[cfg(feature = "zstd")] - pub fn with_compact(self) -> Self { - let builder = self - .with_new_scheme(&string::ZstdScheme) - .with_new_scheme(&binary::ZstdScheme); - - #[cfg(feature = "pco")] - let builder = builder - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme); - - builder - } - - /// Excludes schemes without CUDA kernel support, keeps FSST for string compression, - /// and adds Zstd for binary compression. - /// - /// Both the array-level and the buffer-level Zstd schemes are added. Buffer-level - /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression, - /// but belongs to the opt-in `zstd` edition, so callers filter the two through - /// [`retain_allowed_encodings`](Self::retain_allowed_encodings). - /// - /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a - /// larger encoded representation than the default compressor. - pub fn only_cuda_compatible(self) -> Self { - // Keep FSST, which has a CUDA decoder and direct Arrow offset-based export. Other - // string fragmentation and dictionary schemes still require unsupported decode paths. - #[cfg_attr(not(any(feature = "pco", feature = "zstd")), allow(unused_mut))] - let mut excluded: Vec = vec![ - integer::SparseScheme.id(), - integer::IntRLEScheme.id(), - float::ALPRDScheme.id(), - float::FloatRLEScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - string::StringDictScheme.id(), - binary::BinaryDictScheme.id(), - ]; - // Delta now has a CUDA decode kernel, so arrays that reach the GPU already encoded with - // it — the Delta children OnPair emits, for instance — decode there. It stays excluded - // from this preset until GPU delta decode is benchmarked against the schemes it would - // displace, since the preset picks encodings rather than merely decoding them. - excluded.push(integer::DeltaScheme::default().id()); - #[cfg(feature = "pco")] - excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]); - let builder = self.exclude_schemes(excluded); - - #[cfg(feature = "zstd")] - let builder = builder - .with_new_scheme(&binary::ZstdScheme) - .with_new_scheme(&binary::ZstdBuffersScheme); - - builder - } - - /// Removes the specified compression schemes by their [`SchemeId`]. - pub fn exclude_schemes(mut self, ids: impl IntoIterator) -> Self { - let ids: HashSet<_> = ids.into_iter().collect(); - self.schemes.retain(|s| !ids.contains(&s.id())); - self - } - - /// Retains only schemes whose produced serialized IDs all belong to `allowed`. - /// - /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions - /// permit. - pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes - .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); - self - } - - /// Builds the configured [`BtrBlocksCompressor`]. - pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) - } -} - -#[cfg(test)] -mod tests { - use vortex_array::VTable; - use vortex_fastlanes::FoR; - - use super::*; - - #[test] - fn empty_starts_with_no_schemes() { - let builder = BtrBlocksCompressorBuilder::empty(); - assert!(builder.schemes.is_empty()); - } - - #[test] - fn default_includes_all_schemes() { - let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn retain_allowed_encodings_filters_schemes() { - let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), 1); - assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); - - let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); - assert!(none.schemes.is_empty()); - } - - #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn cuda_compatible_excludes_alprd() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - !builder - .schemes - .iter() - .any(|s| s.id() == float::ALPRDScheme.id()) - ); - } - - /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. - #[test] - fn cuda_compatible_excludes_every_sparse_scheme() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - for excluded in [ - integer::SparseScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - ] { - assert!( - !builder.schemes.iter().any(|s| s.id() == excluded), - "{excluded} should be excluded" - ); - } - } - - #[test] - fn cuda_compatible_uses_fsst_for_strings() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::FSSTScheme.id()) - ); - #[cfg(feature = "zstd")] - assert!( - !builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::ZstdScheme.id()) - ); - } - - #[test] - #[cfg(feature = "pco")] - fn cuda_compatible_excludes_pco() { - let builder = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme) - .only_cuda_compatible(); - for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { - assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); - } - } -} diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index d93be365550..2412122b032 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -5,32 +5,21 @@ use std::ops::Deref; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::session::ArraySessionExt; +use vortex_compressor::session::CompressionSessionExt; +use vortex_edition::ComponentKind; +use vortex_edition::EditionSessionExt; use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; -use crate::BtrBlocksCompressorBuilder; use crate::CascadingCompressor; -/// The BtrBlocks-style compressor with all built-in schemes pre-registered. -/// -/// This is a thin wrapper around [`CascadingCompressor`] that provides a default set of -/// compression schemes via [`BtrBlocksCompressorBuilder`]. -/// -/// # Examples -/// -/// ```rust -/// use vortex_btrblocks::{BtrBlocksCompressor, BtrBlocksCompressorBuilder, Scheme, SchemeExt}; -/// use vortex_btrblocks::schemes::integer::IntDictScheme; -/// -/// // Default compressor - all schemes allowed. -/// let compressor = BtrBlocksCompressor::default(); -/// -/// // Remove specific schemes using the builder. -/// let compressor = BtrBlocksCompressorBuilder::default() -/// .exclude_schemes([IntDictScheme.id()]) -/// .build(); -/// ``` +/// A fully resolved compressor over schemes registered on a session. +/// Construct with [`Self::from_session`] after registration and edition selection. #[derive(Clone)] pub struct BtrBlocksCompressor( /// The underlying cascading compressor. @@ -38,6 +27,43 @@ pub struct BtrBlocksCompressor( ); impl BtrBlocksCompressor { + /// Resolve registered schemes against the selected editions and registered serializers. + pub fn from_session(session: &VortexSession) -> Self { + Self::from_session_with_encodings(session, Self::permitted_encodings(session)) + } + + /// Intersect selected edition permissions with the registered serialized IDs. + pub fn permitted_encodings(session: &VortexSession) -> HashSet { + let registered = session + .arrays() + .registry() + .read(|r| r.keys().copied().collect::>()); + session + .enabled_component_ids(ComponentKind::Array) + .into_iter() + .filter(|id| registered.contains(id)) + .collect() + } + + /// Resolve against the exact IDs used by the writer's serialization context. + pub fn from_session_with_encodings( + session: &VortexSession, + ids: impl IntoIterator, + ) -> Self { + let allowed: HashSet<_> = ids.into_iter().collect(); + let schemes = session + .registered_schemes() + .into_iter() + .filter(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))) + .collect(); + Self(CascadingCompressor::new_with_encodings(schemes, allowed)) + } + + /// Resolve registered schemes for in-memory compression without a wire-format restriction. + pub fn for_memory(session: &VortexSession) -> Self { + Self(CascadingCompressor::new(session.registered_schemes())) + } + /// Compresses an array using BtrBlocks-inspired compression. pub fn compress(&self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.0.compress(array, ctx) @@ -52,12 +78,6 @@ impl Deref for BtrBlocksCompressor { } } -impl Default for BtrBlocksCompressor { - fn default() -> Self { - BtrBlocksCompressorBuilder::default().build() - } -} - #[cfg(test)] mod tests { use std::sync::LazyLock; @@ -65,9 +85,9 @@ mod tests { use rstest::rstest; #[cfg(feature = "zstd")] use vortex_array::ArrayId; - #[cfg(feature = "zstd")] - use vortex_array::ArrayPlugin; use vortex_array::IntoArray; + #[cfg(feature = "zstd")] + use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; @@ -75,6 +95,10 @@ mod tests { use vortex_array::arrays::List; use vortex_array::arrays::ListView; use vortex_array::arrays::ListViewArray; + #[cfg(feature = "zstd")] + use vortex_array::arrays::Primitive; + #[cfg(feature = "zstd")] + use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; @@ -85,13 +109,17 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; #[cfg(feature = "zstd")] - use vortex_utils::aliases::hash_set::HashSet; + use vortex_zstd::schemes::binary::ZstdScheme; + #[cfg(feature = "zstd")] + use vortex_zstd::schemes::binary_buffers::ZstdBuffersScheme; use crate::BtrBlocksCompressor; - #[cfg(feature = "zstd")] - use crate::BtrBlocksCompressorBuilder; - static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); #[rstest] #[case::zctl( @@ -120,7 +148,7 @@ mod tests { ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let array_ref = input.clone().into_array(); - let result = BtrBlocksCompressor::default() + let result = BtrBlocksCompressor::for_memory(&SESSION) .compress(&array_ref, &mut SESSION.create_execution_ctx())?; if expect_list { assert!(result.as_opt::().is_some()); @@ -135,7 +163,7 @@ mod tests { 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 btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -149,7 +177,7 @@ mod tests { 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 btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -166,7 +194,7 @@ mod tests { BitBuffer::from(vec![true; 100]), Validity::from(BitBuffer::from(vec![true; 100])), ); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -181,7 +209,7 @@ mod tests { 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 btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -198,7 +226,7 @@ mod tests { BitBuffer::from(vec![true, false, true, false, true]), Validity::NonNullable, ); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -213,7 +241,7 @@ mod tests { 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 btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -231,7 +259,7 @@ mod tests { .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 btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -257,7 +285,11 @@ mod tests { DType::Binary(Nullability::NonNullable), ); - let compressor = BtrBlocksCompressorBuilder::default().with_compact().build(); + let compressor = { + let compression_session = crate::CompressionSessionExt::fork_compression(&*SESSION); + crate::initialize_compact(&compression_session); + BtrBlocksCompressor::for_memory(&compression_session) + }; let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; @@ -290,12 +322,12 @@ mod tests { 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 session = vortex_array::array_session(); + vortex_zstd::initialize(&session); + crate::CompressionSessionExt::register_scheme(&session, &ZstdScheme); + crate::CompressionSessionExt::register_scheme(&session, &ZstdBuffersScheme); + let permitted = [allowed, VarBinView.id(), Primitive.id()]; + let compressor = BtrBlocksCompressor::from_session_with_encodings(&session, permitted); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; @@ -304,3 +336,87 @@ mod tests { Ok(()) } } + +#[cfg(test)] +mod permission_tests { + use vortex_array::IntoArray; + use vortex_array::VTable; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Constant; + use vortex_array::arrays::Primitive; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::session::ArraySessionExt; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_compressor::builtins::ConstantScheme; + use vortex_edition::EDITION_DECLARATIONS; + use vortex_edition::EDITION_FAMILIES; + use vortex_edition::EditionSessionExt; + use vortex_edition::declarations::core::CORE_2026_08_3; + use vortex_error::VortexResult; + use vortex_fastlanes::BitPacked; + use vortex_fastlanes::schemes::for_::FoRScheme; + use vortex_sequence::schemes::sequence::SequenceScheme; + + use crate::BtrBlocksCompressor; + use crate::CompressionSessionExt; + + #[test] + fn registered_encoding_does_not_enable_a_scheme() -> VortexResult<()> { + let session = vortex_array::array_session(); + let input = buffer![7i32; 128].into_array(); + let output = BtrBlocksCompressor::for_memory(&session) + .compress(&input, &mut session.create_execution_ctx())?; + assert!(output.is::()); + session.register_scheme(&ConstantScheme); + let output = BtrBlocksCompressor::for_memory(&session) + .compress(&input, &mut session.create_execution_ctx())?; + assert!(output.is::()); + Ok(()) + } + + #[test] + fn denied_constant_and_direct_bitpacking_are_not_produced() -> VortexResult<()> { + let session = vortex_array::array_session(); + crate::initialize(&session); + let allowed = session.arrays().registry().read(|r| { + r.keys() + .copied() + .filter(|id| *id != Constant.id() && *id != BitPacked.id()) + .collect::>() + }); + let compressor = BtrBlocksCompressor::from_session_with_encodings(&session, allowed); + let array = PrimitiveArray::new(buffer![42i32; 128], Validity::NonNullable).into_array(); + let result = compressor.compress(&array, &mut session.create_execution_ctx())?; + assert!(!result.is::()); + + let session = vortex_array::array_session(); + session.register_scheme(&FoRScheme); + let compressor = BtrBlocksCompressor::from_session_with_encodings( + &session, + [Primitive.id(), Constant.id(), vortex_fastlanes::FoR.id()], + ); + let array = buffer![100i32, 101, 102, 103].into_array(); + let result = compressor.compress(&array, &mut session.create_execution_ctx())?; + assert!(result.is::()); + Ok(()) + } + + #[test] + fn editions_do_not_enable_missing_serializers() -> VortexResult<()> { + let session = vortex_array::array_session(); + session.register_scheme(&SequenceScheme); + for family in EDITION_FAMILIES { + session.editions().declare_family(family)?; + } + for edition in EDITION_DECLARATIONS { + session.register_edition(edition)?; + } + session.enable_edition(CORE_2026_08_3)?; + let input = buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(); + let output = BtrBlocksCompressor::from_session(&session) + .compress(&input, &mut session.create_execution_ctx())?; + assert!(output.is::()); + Ok(()) + } +} diff --git a/vortex-btrblocks/src/initialize.rs b/vortex-btrblocks/src/initialize.rs new file mode 100644 index 00000000000..59a13c4e917 --- /dev/null +++ b/vortex-btrblocks/src/initialize.rs @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! First-party package initialization for the default compressor. + +use vortex_session::VortexSession; + +/// Register the default encoding packages and their compression schemes. +/// Repeated initialization is idempotent. +pub fn initialize(session: &VortexSession) { + vortex_compressor::builtins::initialize(session); + vortex_alp::schemes::initialize(session); + vortex_datetime_parts::schemes::initialize(session); + vortex_decimal_byte_parts::schemes::initialize(session); + vortex_fastlanes::schemes::initialize(session); + vortex_fsst::schemes::initialize(session); + vortex_onpair::schemes::initialize(session); + vortex_runend::schemes::initialize(session); + vortex_sequence::schemes::initialize(session); + vortex_sparse::schemes::initialize(session); + vortex_zigzag::schemes::initialize(session); +} + +/// Add the compact compression schemes supported by the enabled Cargo features. +pub fn initialize_compact(session: &VortexSession) { + #[cfg(feature = "pco")] + vortex_pco::schemes::initialize(session); + #[cfg(feature = "zstd")] + vortex_zstd::schemes::initialize(session); + let _ = session; +} diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 2e8ae484f90..e5ff2af5203 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -3,85 +3,32 @@ #![deny(missing_docs)] -//! Vortex's [BtrBlocks]-inspired adaptive compression framework. +//! First-party compression initialization and edition-aware compressor construction. //! -//! This crate provides a sophisticated multi-level compression system that adaptively selects -//! optimal compression schemes based on data characteristics. The compressor analyzes arrays -//! to determine the best encoding strategy, supporting cascaded compression with multiple -//! encoding layers for maximum efficiency. -//! -//! # Key Features -//! -//! - **Adaptive Compression**: Automatically selects the best compression scheme based on data -//! patterns. -//! - **Unified Scheme Trait**: A single [`Scheme`] trait covers all data types (integers, floats, -//! strings, etc.) with a [`SchemeId`] for identity. -//! - **Cascaded Encoding**: Multiple compression layers can be applied for optimal results. -//! - **Statistical Analysis**: Uses data sampling and statistics to predict compression ratios. -//! - **Recursive Structure Handling**: Compresses nested structures like structs and lists. -//! -//! # How It Works -//! -//! [`BtrBlocksCompressor::compress()`] takes an `&ArrayRef` plus a mutable execution context and -//! returns an `ArrayRef` that may use a different encoding. It first canonicalizes the input, then dispatches by type. -//! Primitives and strings go through `choose_and_compress`, which evaluates every enabled -//! [`Scheme`] and picks the one with the best compression ratio. Compound types like structs -//! and lists recurse into their fields and elements. -//! -//! 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`]. -//! -//! 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 -//! the same scheme from being applied twice in a chain. -//! -//! # Example -//! -//! ```rust -//! 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_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(); -//! 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 +//! Register encoding packages and their schemes, select editions, then construct +//! [`BtrBlocksCompressor::from_session`]. Registration opts a scheme into compression. +//! Construction removes configurations whose serialized outputs are not permitted. -mod builder; +mod initialize; +pub use initialize::initialize; +pub use initialize::initialize_compact; mod canonical_compressor; /// Compression scheme implementations. pub mod schemes; +mod session; #[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; +/// Optional delta scheme configuration. +pub static DELTA_SCHEME: schemes::integer::DeltaScheme = schemes::integer::DeltaScheme::new(1.25); pub use canonical_compressor::BtrBlocksCompressor; -pub use schemes::patches::compress_patches; +pub use session::CompressionSession; +pub use session::CompressionSessionExt; pub use vortex_compressor::CascadingCompressor; +pub use vortex_compressor::compress_patches; pub use vortex_compressor::scheme::CompressorContext; pub use vortex_compressor::scheme::MAX_CASCADE; pub use vortex_compressor::scheme::Scheme; diff --git a/vortex-btrblocks/src/schemes/binary/mod.rs b/vortex-btrblocks/src/schemes/binary/mod.rs index d83b6eef919..dbb5f075d63 100644 --- a/vortex-btrblocks/src/schemes/binary/mod.rs +++ b/vortex-btrblocks/src/schemes/binary/mod.rs @@ -4,15 +4,11 @@ //! Binary compression schemes. mod varbin; -#[cfg(feature = "zstd")] -mod zstd; -#[cfg(feature = "zstd")] -mod zstd_buffers; // Re-export builtin schemes from vortex-compressor. pub use varbin::VarBinScheme; pub use vortex_compressor::builtins::BinaryDictScheme; #[cfg(feature = "zstd")] -pub use zstd::ZstdScheme; +pub use vortex_zstd::schemes::binary::ZstdScheme; #[cfg(feature = "zstd")] -pub use zstd_buffers::ZstdBuffersScheme; +pub use vortex_zstd::schemes::binary_buffers::ZstdBuffersScheme; diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs index 849402a6493..cb1efeccc20 100644 --- a/vortex-btrblocks/src/schemes/binary/varbin.rs +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -1,100 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Offset-based storage for binary arrays. -//! -//! Canonical binary arrays are [`VarBinViewArray`](vortex_array::arrays::VarBinViewArray), which -//! spends a fixed 16 bytes per element on an opaque views buffer that no scheme can compress. -//! Re-encoding as [`VarBinArray`] replaces that buffer with an offsets child array, which the -//! cascading compressor can then compress with the ordinary integer schemes. For fixed-width -//! values the offsets are a constant-stride sequence and collapse to nothing. +//! Compatibility re-export of the shared compression implementation. -use vortex_array::ArrayId; -use vortex_array::ArrayRef; -use vortex_array::Canonical; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::VTable; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::VarBin; -use vortex_array::arrays::VarBinArray; -use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::arrays::varbin::VarBinArraySlotsExt; -use vortex_array::builders::VarBinBuilder; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; -use vortex_compressor::scheme::SchemeExt; -use vortex_error::VortexResult; - -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; - -/// Offset-based (rather than view-based) storage for binary arrays. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct VarBinScheme; - -impl Scheme for VarBinScheme { - fn scheme_name(&self) -> &'static str { - "vortex.binary.varbin" - } - - fn matches(&self, canonical: &Canonical) -> bool { - canonical.dtype().is_binary() - } - - fn produced_encodings(&self) -> Vec { - vec![VarBin.id()] - } - - fn num_children(&self) -> usize { - 1 - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) - } - - fn compress( - &self, - compressor: &CascadingCompressor, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - // `append_to_builder` resolves the views slice and data buffers once and appends - // borrowed slices into a single pre-sized allocation. Iterating the array per element - // instead would clone a buffer handle and allocate for every value. - let array = data.array(); - let mut builder = VarBinBuilder::::with_capacity_in( - array.dtype().clone(), - array.len(), - exec_ctx.allocator(), - ); - array.append_to_builder(&mut builder, exec_ctx)?; - let varbin = builder.finish_into_varbin(); - - let offsets = varbin - .offsets() - .clone() - .execute::(exec_ctx)? - .narrow(exec_ctx)? - .into_array(); - let compressed_offsets = - compressor.compress_child(&offsets, &compress_ctx, self.id(), 0, exec_ctx)?; - - Ok(VarBinArray::try_new( - compressed_offsets, - varbin.bytes().clone(), - varbin.dtype().clone(), - varbin.validity()?, - )? - .into_array()) - } -} +pub use vortex_compressor::builtins::VarBinScheme; diff --git a/vortex-btrblocks/src/schemes/float/mod.rs b/vortex-btrblocks/src/schemes/float/mod.rs index 1301184ac0c..cab89c1151d 100644 --- a/vortex-btrblocks/src/schemes/float/mod.rs +++ b/vortex-btrblocks/src/schemes/float/mod.rs @@ -3,23 +3,15 @@ //! Float compression schemes. -mod alp; -mod alprd; -mod rle; -mod sparse; - -#[cfg(feature = "pco")] -mod pco; - -pub use alp::ALPScheme; -pub use alprd::ALPRDScheme; -#[cfg(feature = "pco")] -pub use pco::PcoScheme; -pub use rle::FloatRLEScheme; -pub use sparse::NullDominatedSparseScheme; +pub use vortex_alp::schemes::alp::ALPScheme; +pub use vortex_alp::schemes::alprd::ALPRDScheme; // Re-export builtin schemes from vortex-compressor. pub use vortex_compressor::builtins::FloatDictScheme; pub use vortex_compressor::stats::FloatStats; +pub use vortex_fastlanes::schemes::float_rle::FloatRLEScheme; +#[cfg(feature = "pco")] +pub use vortex_pco::schemes::float::PcoScheme; +pub use vortex_sparse::schemes::float::NullDominatedSparseScheme; #[cfg(test)] mod scheme_selection_tests; diff --git a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs index c6f54a9dafe..d65e4fbe40c 100644 --- a/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/float/scheme_selection_tests.rs @@ -21,13 +21,17 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_constant_compressed() -> VortexResult<()> { let values: Vec = vec![42.5; 100]; let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -37,7 +41,7 @@ fn test_constant_compressed() -> VortexResult<()> { fn test_alp_compressed() -> VortexResult<()> { let values: Vec = (0..1000).map(|i| (i as f64) * 0.01).collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -50,7 +54,7 @@ fn test_dict_compressed() -> VortexResult<()> { .map(|i| distinct_values[i % distinct_values.len()]) .collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert!(compressed.children()[0].is::()); @@ -69,7 +73,7 @@ fn test_null_dominated_compressed() -> VortexResult<()> { } builder.append_nulls(95); let array = builder.finish_into_primitive(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; // Verify the compressed array preserves values. assert_eq!(compressed.len(), 100); diff --git a/vortex-btrblocks/src/schemes/float/tests.rs b/vortex-btrblocks/src/schemes/float/tests.rs index bb5301b807b..252c99cd0b9 100644 --- a/vortex-btrblocks/src/schemes/float/tests.rs +++ b/vortex-btrblocks/src/schemes/float/tests.rs @@ -22,11 +22,15 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::float::FloatRLEScheme; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_empty() -> VortexResult<()> { - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let array = PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable).into_array(); let result = btr.compress(&array, &mut SESSION.create_execution_ctx())?; @@ -42,7 +46,7 @@ fn test_compress() -> VortexResult<()> { } let array = values.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 1024); @@ -92,7 +96,7 @@ fn test_sparse_compression() -> VortexResult<()> { array.append_nulls(90); let array = array.finish_into_primitive().into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 96); diff --git a/vortex-btrblocks/src/schemes/integer/mod.rs b/vortex-btrblocks/src/schemes/integer/mod.rs index 5508ac4566c..423e3b5f42a 100644 --- a/vortex-btrblocks/src/schemes/integer/mod.rs +++ b/vortex-btrblocks/src/schemes/integer/mod.rs @@ -3,35 +3,19 @@ //! Integer compression schemes. -mod bitpacking; -mod delta; -mod for_; -mod rle; -mod runend; -mod sequence; -mod sparse; -mod zigzag; - -#[cfg(feature = "pco")] -mod pco; - -pub use bitpacking::BitPackingScheme; -pub use delta::DeltaScheme; -pub use for_::FoRScheme; -#[cfg(feature = "pco")] -pub use pco::PcoScheme; -pub use rle::IntRLEScheme; -pub(crate) use rle::rle_compress; -pub use runend::RunEndScheme; -pub use sequence::SequenceScheme; -pub use sparse::SparseScheme; // Re-export builtin schemes from vortex-compressor. pub use vortex_compressor::builtins::IntDictScheme; pub use vortex_compressor::stats::IntegerStats; -pub use zigzag::ZigZagScheme; - -/// Threshold for the average run length in an array before we consider run-length encoding. -pub(crate) const RUN_LENGTH_THRESHOLD: u32 = 4; +pub use vortex_fastlanes::schemes::bitpacking::BitPackingScheme; +pub use vortex_fastlanes::schemes::delta::DeltaScheme; +pub use vortex_fastlanes::schemes::for_::FoRScheme; +pub use vortex_fastlanes::schemes::integer_rle::IntRLEScheme; +#[cfg(feature = "pco")] +pub use vortex_pco::schemes::integer::PcoScheme; +pub use vortex_runend::schemes::runend::RunEndScheme; +pub use vortex_sequence::schemes::sequence::SequenceScheme; +pub use vortex_sparse::schemes::integer::SparseScheme; +pub use vortex_zigzag::schemes::zigzag::ZigZagScheme; #[cfg(test)] mod scheme_selection_tests; diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index b4726dab9b5..b6b59adaeca 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -28,15 +28,18 @@ use vortex_session::VortexSession; use vortex_sparse::Sparse; use crate::BtrBlocksCompressor; -use crate::BtrBlocksCompressorBuilder; use crate::DELTA_SCHEME; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[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::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -46,7 +49,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::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -56,7 +59,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::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); assert_eq!( @@ -85,7 +88,7 @@ fn test_sparse_compressed() -> VortexResult<()> { } } let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -109,7 +112,7 @@ fn test_dict_compressed() -> VortexResult<()> { } let array = PrimitiveArray::new(Buffer::copy_from(&codes), Validity::NonNullable); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -122,7 +125,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::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -132,7 +135,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::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) @@ -149,7 +152,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::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; eprintln!("{}", compressed.display_tree()); assert!(compressed.is::()); @@ -176,9 +179,11 @@ 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 = { + let compression_session = crate::CompressionSessionExt::fork_compression(&*SESSION); + crate::CompressionSessionExt::register_scheme(&compression_session, &DELTA_SCHEME); + BtrBlocksCompressor::for_memory(&compression_session) + }; let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -216,9 +221,11 @@ 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 = { + let compression_session = crate::CompressionSessionExt::fork_compression(&*SESSION); + crate::CompressionSessionExt::register_scheme(&compression_session, &DELTA_SCHEME); + BtrBlocksCompressor::for_memory(&compression_session) + }; let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -245,9 +252,11 @@ 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 = { + let compression_session = crate::CompressionSessionExt::fork_compression(&*SESSION); + crate::CompressionSessionExt::register_scheme(&compression_session, &DELTA_SCHEME); + BtrBlocksCompressor::for_memory(&compression_session) + }; 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..2a43c0e90b6 100644 --- a/vortex-btrblocks/src/schemes/integer/tests.rs +++ b/vortex-btrblocks/src/schemes/integer/tests.rs @@ -27,12 +27,16 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::schemes::integer::IntRLEScheme; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_empty() -> VortexResult<()> { // Make sure empty array compression does not fail. - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let array = PrimitiveArray::new(Buffer::::empty(), Validity::NonNullable); let result = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; @@ -60,7 +64,7 @@ fn test_dict_encodable() -> VortexResult<()> { } } - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress( &codes.freeze().into_array(), &mut SESSION.create_execution_ctx(), @@ -80,7 +84,7 @@ fn constant_mostly_nulls() -> VortexResult<()> { ); let validity = array.validity()?; - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -99,7 +103,7 @@ fn nullable_sequence() -> VortexResult<()> { let values = (0i32..20).step_by(7).collect_vec(); let array = PrimitiveArray::from_option_iter(values.clone().into_iter().map(Some)); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); @@ -141,7 +145,7 @@ fn compress_large_int() -> VortexResult<()> { .collect::() .into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); btr.compress(&prim, &mut SESSION.create_execution_ctx())?; Ok(()) diff --git a/vortex-btrblocks/src/schemes/mod.rs b/vortex-btrblocks/src/schemes/mod.rs index a0e9b042a66..ae46b51055f 100644 --- a/vortex-btrblocks/src/schemes/mod.rs +++ b/vortex-btrblocks/src/schemes/mod.rs @@ -8,65 +8,5 @@ pub mod float; pub mod integer; pub mod string; -pub mod decimal; -pub mod temporal; - -pub(crate) mod patches; - -use vortex_compressor::builtins::BinaryDictScheme; -use vortex_compressor::builtins::FloatDictScheme; -use vortex_compressor::builtins::IntDictScheme; -use vortex_compressor::builtins::StringDictScheme; -use vortex_compressor::scheme::AncestorExclusion; -use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::SchemeExt; - -use crate::schemes::integer::SparseScheme; - -/// Shared descendant exclusion rules for RLE schemes. -/// -/// RLE indices (child 1) and offsets (child 2) are monotonically increasing positions with all -/// unique values. Dict and Sparse are pointless on such data. Self-exclusion already prevents -/// RLE on RLE children. -fn rle_descendant_exclusions() -> Vec { - vec![ - DescendantExclusion { - excluded: IntDictScheme.id(), - children: ChildSelection::Many(&[1, 2]), - }, - // TODO(connor): This is wrong for some reason? - // DescendantExclusion { - // excluded: RunEndScheme.id(), - // children: ChildSelection::Many(&[1, 2]), - // }, - DescendantExclusion { - excluded: SparseScheme.id(), - children: ChildSelection::Many(&[1, 2]), - }, - ] -} - -/// Shared ancestor exclusion rules for RLE schemes. -/// -/// Dict values (child 0) are all unique by definition, so RLE is pointless on them. -fn rle_ancestor_exclusions() -> Vec { - vec![ - AncestorExclusion { - ancestor: IntDictScheme.id(), - children: ChildSelection::One(0), - }, - AncestorExclusion { - ancestor: FloatDictScheme.id(), - children: ChildSelection::One(0), - }, - AncestorExclusion { - ancestor: StringDictScheme.id(), - children: ChildSelection::One(0), - }, - AncestorExclusion { - ancestor: BinaryDictScheme.id(), - children: ChildSelection::One(0), - }, - ] -} +pub use vortex_datetime_parts::schemes::temporal; +pub use vortex_decimal_byte_parts::schemes::decimal; diff --git a/vortex-btrblocks/src/schemes/string/mod.rs b/vortex-btrblocks/src/schemes/string/mod.rs index 5d04de4777f..5fefadbd4ca 100644 --- a/vortex-btrblocks/src/schemes/string/mod.rs +++ b/vortex-btrblocks/src/schemes/string/mod.rs @@ -3,26 +3,16 @@ //! String compression schemes. -mod fsst; -mod sparse; - -#[cfg(feature = "zstd")] -mod zstd; -#[cfg(feature = "zstd")] -mod zstd_buffers; - -mod onpair; - -pub use fsst::FSSTScheme; -pub use onpair::OnPairScheme; -pub use sparse::NullDominatedSparseScheme; // Re-export builtin schemes from vortex-compressor. pub use vortex_compressor::builtins::StringDictScheme; pub use vortex_compressor::stats::StringStats; +pub use vortex_fsst::schemes::fsst::FSSTScheme; +pub use vortex_onpair::schemes::onpair::OnPairScheme; +pub use vortex_sparse::schemes::string::NullDominatedSparseScheme; #[cfg(feature = "zstd")] -pub use zstd::ZstdScheme; +pub use vortex_zstd::schemes::string::ZstdScheme; #[cfg(feature = "zstd")] -pub use zstd_buffers::ZstdBuffersScheme; +pub use vortex_zstd::schemes::string_buffers::ZstdBuffersScheme; #[cfg(test)] mod scheme_selection_tests; diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index aac0b4de4de..2fe7b1ec2f5 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -17,16 +17,22 @@ use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::SchemeExt; +use crate::schemes::string::OnPairScheme; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_constant_compressed() -> VortexResult<()> { let strings: Vec> = vec![Some("constant_value"); 100]; let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressed = - BtrBlocksCompressor::default().compress(&array_ref, &mut SESSION.create_execution_ctx())?; + let compressed = BtrBlocksCompressor::for_memory(&SESSION) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } @@ -40,21 +46,21 @@ 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::for_memory(&SESSION) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!(compressed.is::()); Ok(()) } #[test] -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(); +fn test_default_registration_includes_onpair() { + let ids: Vec<_> = crate::CompressionSessionExt::registered_schemes(&*SESSION) + .iter() + .map(|s| s.id()) + .collect(); assert!( ids.contains(&OnPairScheme.id()), - "OnPairScheme not registered in ALL_SCHEMES" + "OnPairScheme not registered on the default session" ); } @@ -71,8 +77,8 @@ fn test_default_btrblocks_compressor_selects_onpair() -> VortexResult<()> { } let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressed = - BtrBlocksCompressor::default().compress(&array_ref, &mut SESSION.create_execution_ctx())?; + let compressed = BtrBlocksCompressor::for_memory(&SESSION) + .compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), "expected OnPair, got {}", @@ -81,21 +87,21 @@ fn test_default_btrblocks_compressor_selects_onpair() -> VortexResult<()> { Ok(()) } -/// FSST is registered in the default scheme list, and an FSST-only builder -/// still produces an FSST array. +/// FSST is registered by default, and an FSST-only session still produces an FSST array. #[test] -fn test_fsst_in_default_scheme_list() -> VortexResult<()> { - use crate::BtrBlocksCompressorBuilder; +fn test_fsst_registration() -> VortexResult<()> { 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", + crate::CompressionSessionExt::registered_schemes(&*SESSION) + .iter() + .any(|s| s.id() == FSSTScheme.id()), + "FSSTScheme should be on the default session", ); - // An FSST-only builder still produces an FSST array for FSST-favourable + // An FSST-only session still produces an FSST array for FSST-favourable // input. let mut strings = Vec::with_capacity(1000); for i in 0..1000 { @@ -106,9 +112,11 @@ 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 = { + let compression_session = vortex_array::array_session(); + crate::CompressionSessionExt::register_scheme(&compression_session, &FSSTScheme); + BtrBlocksCompressor::for_memory(&compression_session) + }; let compressed = compressor.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), diff --git a/vortex-btrblocks/src/schemes/string/tests.rs b/vortex-btrblocks/src/schemes/string/tests.rs index 1928f0065a2..9eb8c964745 100644 --- a/vortex-btrblocks/src/schemes/string/tests.rs +++ b/vortex-btrblocks/src/schemes/string/tests.rs @@ -16,7 +16,11 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); #[test] fn test_strings() -> VortexResult<()> { @@ -30,7 +34,7 @@ fn test_strings() -> VortexResult<()> { let strings = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = strings.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 2048); @@ -57,7 +61,7 @@ fn test_sparse_nulls() -> VortexResult<()> { let strings = strings.finish_into_varbinview(); let array_ref = strings.into_array(); - let btr = BtrBlocksCompressor::default(); + let btr = BtrBlocksCompressor::for_memory(&SESSION); let compressed = btr.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert_eq!(compressed.len(), 100); diff --git a/vortex-btrblocks/src/session.rs b/vortex-btrblocks/src/session.rs new file mode 100644 index 00000000000..0a6e17973b5 --- /dev/null +++ b/vortex-btrblocks/src/session.rs @@ -0,0 +1,5 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compatibility exports for the compression registry. +pub use vortex_compressor::session::*; diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index e23e4ef0244..be08cec4703 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -53,7 +53,7 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; -use crate::BtrBlocksCompressorBuilder; +use crate::BtrBlocksCompressor; use crate::DELTA_SCHEME; /// A session with the default Vortex encodings registered. @@ -69,6 +69,7 @@ use crate::DELTA_SCHEME; /// this crate does not depend on them, so the compressor cannot emit them. fn trace_session() -> VortexSession { let session = VortexSession::empty().with::(); + crate::initialize(&session); vortex_fsst::initialize(&session); vortex_onpair::initialize(&session); @@ -127,10 +128,12 @@ fn lineitem() -> VortexResult { /// Delta is opt-in, and these traces cover the delta-encoded FSST offsets, so enable it here. fn compressed_lineitem() -> VortexResult { - BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build() - .compress(&lineitem()?, &mut execution_ctx()) + { + let compression_session = crate::CompressionSessionExt::fork_compression(&trace_session()); + crate::CompressionSessionExt::register_scheme(&compression_session, &DELTA_SCHEME); + BtrBlocksCompressor::for_memory(&compression_session) + } + .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..9359d7201c5 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -15,7 +15,7 @@ //! - `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. +//! [`vortex_btrblocks::initialize_compact`] — pins Zstd / Pco selection. //! //! Every corpus entry is longer than 1024 values so the sampling-based estimation path is //! exercised, and each entry is compressed twice per run to assert determinism directly. @@ -49,9 +49,8 @@ use vortex_array::dtype::Nullability; use vortex_array::extension::datetime::TimeUnit; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; -use vortex_btrblocks::BtrBlocksCompressorBuilder; +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; @@ -60,8 +59,14 @@ use vortex_edition::EditionSessionExt; use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexResult; use vortex_session::VortexSession; +#[cfg(all(feature = "zstd", feature = "pco"))] +use vortex_zstd::editions::ZSTD_2026_02; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session +}); /// Number of values in each numeric corpus entry: comfortably above the 1024-value sampling /// threshold so scheme selection runs on sampled estimates, as it does for real file chunks. @@ -393,15 +398,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(session: &VortexSession) -> VortexSession { + let session = vortex_btrblocks::CompressionSessionExt::fork_compression(session); + vortex_btrblocks::CompressionSessionExt::compression(&session) + .unregister(vortex_btrblocks::SchemeExt::id(&OnPairScheme)); + session } 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 +420,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::from_session(&without_onpair(&session)); golden_corpus_snapshots("regular", &compressor) } @@ -450,7 +431,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, @@ -463,10 +444,9 @@ fn golden_onpair() -> VortexResult<()> { 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(), - ); + session.enable_edition(ZSTD_2026_02)?; + let session = without_onpair(&session); + vortex_btrblocks::initialize_compact(&session); + let compressor = BtrBlocksCompressor::from_session(&session); golden_corpus_snapshots("compact", &compressor) } diff --git a/vortex-btrblocks/tests/onpair_roundtrip.rs b/vortex-btrblocks/tests/onpair_roundtrip.rs index 31734d6a60e..38b37805c3d 100644 --- a/vortex-btrblocks/tests/onpair_roundtrip.rs +++ b/vortex-btrblocks/tests/onpair_roundtrip.rs @@ -21,7 +21,11 @@ use vortex_array::dtype::Nullability; use vortex_btrblocks::BtrBlocksCompressor; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session +}); /// Helper: synthetic short-string corpus that the cascading compressor should /// route through OnPair. @@ -58,7 +62,7 @@ fn nonnullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::for_memory(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert a specific scheme — both OnPair and FSST are registered and @@ -101,7 +105,7 @@ fn nullable_roundtrip_via_default_compressor() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::for_memory(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); // Don't assert OnPair specifically here — the sample-based selector may @@ -137,7 +141,7 @@ fn large_unique_short_strings_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::for_memory(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); @@ -166,7 +170,7 @@ fn empty_and_short_string_roundtrip() { ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::for_memory(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); let decoded = compressed @@ -211,7 +215,7 @@ fn delta_dict_offsets_roundtrip() { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::for_memory(&SESSION) .compress(&array, &mut SESSION.create_execution_ctx()) .expect("compress"); let decoded = compressed diff --git a/vortex-btrblocks/tests/varbin_scheme.rs b/vortex-btrblocks/tests/varbin_scheme.rs index d47c8280af1..0b68a6bfdc3 100644 --- a/vortex-btrblocks/tests/varbin_scheme.rs +++ b/vortex-btrblocks/tests/varbin_scheme.rs @@ -14,14 +14,18 @@ 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::SchemeExt; use vortex_btrblocks::schemes::binary::VarBinScheme; use vortex_btrblocks::schemes::string::OnPairScheme; use vortex_error::VortexResult; use vortex_session::VortexSession; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_btrblocks::initialize(&session); + session +}); const N: usize = 100_000; @@ -70,10 +74,14 @@ 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::for_memory(&SESSION); + let without = { + let compression_session = + vortex_btrblocks::CompressionSessionExt::fork_compression(&*SESSION); + vortex_btrblocks::CompressionSessionExt::compression(&compression_session) + .unregister(VarBinScheme.id()); + BtrBlocksCompressor::for_memory(&compression_session) + }; println!( "{:<20}{:>12}{:>14}{:>14}{:>9}", @@ -116,9 +124,9 @@ 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 session = vortex_btrblocks::CompressionSessionExt::fork_compression(&*SESSION); + vortex_btrblocks::CompressionSessionExt::compression(&session).unregister(OnPairScheme.id()); + let compressor = BtrBlocksCompressor::for_memory(&session); 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..cd66305f774 100644 --- a/vortex-compressor/Cargo.toml +++ b/vortex-compressor/Cargo.toml @@ -24,6 +24,7 @@ vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } +vortex-session = { workspace = true } vortex-utils = { workspace = true } [dev-dependencies] diff --git a/vortex-compressor/src/builtins/constant.rs b/vortex-compressor/src/builtins/constant.rs new file mode 100644 index 00000000000..a1e70ce3688 --- /dev/null +++ b/vortex-compressor/src/builtins/constant.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Constant compression, selected through the same registry as other schemes. + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::Bool; +use vortex_array::arrays::Constant; +use vortex_array::arrays::Masked; +use vortex_error::VortexResult; + +use crate::CascadingCompressor; +use crate::compressor::constant::compress_constant; +use crate::compressor::constant::is_constant_for_compression; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::EstimateVerdict; +use crate::scheme::Scheme; +use crate::stats::ArrayAndStats; + +/// Compress entirely valid or entirely null leaves as constants. +#[derive(Debug)] +pub struct ConstantScheme; + +/// Compress partially valid constant leaves with a validity wrapper. +#[derive(Debug)] +pub struct MaskedConstantScheme; + +impl Scheme for ConstantScheme { + fn scheme_name(&self) -> &'static str { + "vortex.compressor.constant" + } + + fn selection_priority(&self) -> u16 { + 0 + } + + fn matches(&self, _array: &Canonical) -> bool { + true + } + + fn produced_encodings(&self) -> Vec { + vec![Constant.id()] + } + + fn expected_compression_ratio( + &self, + data: &ArrayAndStats, + ctx: CompressorContext, + exec: &mut ExecutionCtx, + ) -> CompressionEstimate { + constant_estimate(data, ctx, exec, false) + } + + fn compress( + &self, + _: &CascadingCompressor, + data: &ArrayAndStats, + _: CompressorContext, + exec: &mut ExecutionCtx, + ) -> VortexResult { + compress_constant(data.array(), exec) + } +} + +impl Scheme for MaskedConstantScheme { + fn scheme_name(&self) -> &'static str { + "vortex.compressor.masked_constant" + } + + fn selection_priority(&self) -> u16 { + 0 + } + + fn matches(&self, _array: &Canonical) -> bool { + true + } + + fn produced_encodings(&self) -> Vec { + vec![Constant.id(), Masked.id(), Bool.id()] + } + + fn expected_compression_ratio( + &self, + data: &ArrayAndStats, + ctx: CompressorContext, + exec: &mut ExecutionCtx, + ) -> CompressionEstimate { + constant_estimate(data, ctx, exec, true) + } + + fn compress( + &self, + _: &CascadingCompressor, + data: &ArrayAndStats, + _: CompressorContext, + exec: &mut ExecutionCtx, + ) -> VortexResult { + compress_constant(data.array(), exec) + } +} + +/// Keep constant samples from deciding whether a full input is constant. +fn constant_estimate( + data: &ArrayAndStats, + ctx: CompressorContext, + exec: &mut ExecutionCtx, + masked: bool, +) -> CompressionEstimate { + let result = (|| -> VortexResult { + let validity = data + .array() + .validity()? + .execute_mask(data.array_len(), exec)?; + if validity.all_false() { + return Ok(!masked); + } + if ctx.is_sample() || masked == validity.all_true() { + return Ok(false); + } + is_constant_for_compression(data, exec) + })(); + match result { + Ok(true) => CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse), + Ok(false) => CompressionEstimate::Verdict(EstimateVerdict::Skip), + Err(error) => CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( + move |_, _, _, _, _| Err(error), + ))), + } +} diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index c407d0251c6..a6b9a943e8f 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -40,6 +40,10 @@ use crate::stats::GenerateStatsOptions; pub struct BinaryDictScheme; impl Scheme for BinaryDictScheme { + fn selection_priority(&self) -> u16 { + 10 + } + fn scheme_name(&self) -> &'static str { "vortex.binary.dict" } diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index f962c3ff967..691dbc1760b 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -25,6 +25,7 @@ use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_utils::aliases::hash_map::HashMap; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; @@ -46,6 +47,10 @@ use crate::stats::GenerateStatsOptions; pub struct FloatDictScheme; impl Scheme for FloatDictScheme { + fn selection_priority(&self) -> u16 { + 30 + } + fn scheme_name(&self) -> &'static str { "vortex.float.dict" } @@ -228,7 +233,7 @@ macro_rules! impl_encode { #[expect(clippy::cast_possible_truncation)] fn encode(distinct: &[$typ], values: &[$typ]) -> Buffer<$ityp> { let mut codes = - vortex_utils::aliases::hash_map::HashMap::<$utyp, $ityp>::with_capacity( + HashMap::<$utyp, $ityp>::with_capacity( distinct.len(), ); for (code, &value) in distinct.iter().enumerate() { diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 27a17ef94ad..7d292f13f25 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -24,6 +24,7 @@ use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_utils::aliases::hash_map::HashMap; use crate::CascadingCompressor; use crate::scheme::CompressionEstimate; @@ -41,6 +42,10 @@ use crate::stats::IntegerStats; pub struct IntDictScheme; impl Scheme for IntDictScheme { + fn selection_priority(&self) -> u16 { + 50 + } + fn scheme_name(&self) -> &'static str { "vortex.int.dict" } @@ -229,7 +234,7 @@ macro_rules! impl_encode { #[expect(clippy::cast_possible_truncation)] fn encode(distinct: &[$typ], values: &[$typ]) -> Buffer<$ityp> { let mut codes = - vortex_utils::aliases::hash_map::HashMap::<$typ, $ityp>::with_capacity( + HashMap::<$typ, $ityp>::with_capacity( distinct.len(), ); for (code, &value) in distinct.iter().enumerate() { @@ -272,6 +277,7 @@ mod tests { use vortex_error::VortexResult; use super::dictionary_encode; + use crate::stats::GenerateStatsOptions; use crate::stats::IntegerStats; #[test] @@ -284,7 +290,7 @@ mod tests { let stats = IntegerStats::generate_opts( &array, - crate::stats::GenerateStatsOptions { + GenerateStatsOptions { count_distinct_values: true, }, &mut ctx, diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index f5cbcd54d89..40e5704e77b 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -40,6 +40,10 @@ use crate::stats::GenerateStatsOptions; pub struct StringDictScheme; impl Scheme for StringDictScheme { + fn selection_priority(&self) -> u16 { + 10 + } + fn scheme_name(&self) -> &'static str { "vortex.string.dict" } diff --git a/vortex-compressor/src/builtins/mod.rs b/vortex-compressor/src/builtins/mod.rs index 11059fa6ed4..96e02203c06 100644 --- a/vortex-compressor/src/builtins/mod.rs +++ b/vortex-compressor/src/builtins/mod.rs @@ -8,6 +8,7 @@ //! //! [`DictArray`]: vortex_array::arrays::DictArray +use crate::session::CompressionSessionExt; mod dict; pub use dict::BinaryDictScheme; @@ -16,3 +17,21 @@ pub use dict::IntDictScheme; pub use dict::StringDictScheme; pub use dict::float_dictionary_encode; pub use dict::integer_dictionary_encode; + +mod varbin; +pub use varbin::VarBinScheme; + +mod constant; +pub use constant::ConstantScheme; +pub use constant::MaskedConstantScheme; + +/// Register compression strategies for canonical array encodings. +pub fn initialize(session: &vortex_session::VortexSession) { + session.register_scheme(&ConstantScheme); + session.register_scheme(&MaskedConstantScheme); + session.register_scheme(&IntDictScheme); + session.register_scheme(&FloatDictScheme); + session.register_scheme(&StringDictScheme); + session.register_scheme(&BinaryDictScheme); + session.register_scheme(&VarBinScheme); +} diff --git a/vortex-compressor/src/builtins/varbin.rs b/vortex-compressor/src/builtins/varbin.rs new file mode 100644 index 00000000000..ecbee555e14 --- /dev/null +++ b/vortex-compressor/src/builtins/varbin.rs @@ -0,0 +1,100 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Offset-based storage for binary arrays. +//! +//! Canonical binary arrays are [`VarBinViewArray`](vortex_array::arrays::VarBinViewArray), which +//! spends a fixed 16 bytes per element on an opaque views buffer that no scheme can compress. +//! Re-encoding as [`VarBinArray`] replaces that buffer with an offsets child array, which the +//! cascading compressor can then compress with the ordinary integer schemes. For fixed-width +//! values the offsets are a constant-stride sequence and collapse to nothing. + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::VarBin; +use vortex_array::arrays::VarBinArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::arrays::varbin::VarBinArraySlotsExt; +use vortex_array::builders::VarBinBuilder; +use vortex_error::VortexResult; + +use crate::CascadingCompressor; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::DeferredEstimate; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::stats::ArrayAndStats; + +/// Offset-based (rather than view-based) storage for binary arrays. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct VarBinScheme; + +impl Scheme for VarBinScheme { + fn scheme_name(&self) -> &'static str { + "vortex.binary.varbin" + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_binary() + } + + fn produced_encodings(&self) -> Vec { + vec![VarBin.id()] + } + + fn num_children(&self) -> usize { + 1 + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Deferred(DeferredEstimate::Sample) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + // `append_to_builder` resolves the views slice and data buffers once and appends + // borrowed slices into a single pre-sized allocation. Iterating the array per element + // instead would clone a buffer handle and allocate for every value. + let array = data.array(); + let mut builder = VarBinBuilder::::with_capacity_in( + array.dtype().clone(), + array.len(), + exec_ctx.allocator(), + ); + array.append_to_builder(&mut builder, exec_ctx)?; + let varbin = builder.finish_into_varbin(); + + let offsets = varbin + .offsets() + .clone() + .execute::(exec_ctx)? + .narrow(exec_ctx)? + .into_array(); + let compressed_offsets = + compressor.compress_child(&offsets, &compress_ctx, self.id(), 0, exec_ctx)?; + + Ok(VarBinArray::try_new( + compressed_offsets, + varbin.bytes().clone(), + varbin.dtype().clone(), + varbin.validity()?, + )? + .into_array()) + } +} diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 86d45d2c0d9..588af9e803c 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -9,7 +9,6 @@ use vortex_array::CanonicalValidity; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Constant; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::Masked; @@ -27,11 +26,11 @@ use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::union::UnionArrayExt; use vortex_array::arrays::union::UnionArraySlotsExt; use vortex_array::arrays::variant::VariantArraySlotsExt; -use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use super::CascadingCompressor; -use super::constant; +use crate::builtins::ConstantScheme; +use crate::builtins::IntDictScheme; use crate::scheme::CompressorContext; use crate::scheme::Scheme; use crate::scheme::SchemeExt; @@ -52,6 +51,30 @@ impl CascadingCompressor { &self, array: &ArrayRef, exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + self.compress_with_context(array, CompressorContext::new(), exec_ctx) + } + + /// Compress dictionary codes with the existing dictionary cascade exclusions. + /// + /// # Errors + /// + /// Returns an error if canonicalization or compression fails. + pub fn compress_dictionary_codes( + &self, + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let ctx = CompressorContext::new().excluding_at_site(IntDictScheme.id()); + self.compress_with_context(array, ctx, exec_ctx) + } + + /// Canonicalize and compress an input while preserving its call-site context. + pub(super) fn compress_with_context( + &self, + array: &ArrayRef, + context: CompressorContext, + exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let before_nbytes = array.nbytes(); let span = trace::compress_span(array.len(), array.dtype(), before_nbytes); @@ -59,13 +82,25 @@ impl CascadingCompressor { let canonical = array.clone().execute::(exec_ctx)?.0; let compact = canonical.compact(exec_ctx)?; - let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; + let compressed = self.compress_canonical(compact, context, exec_ctx)?; trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); Ok(compressed) } + /// Reject a canonical fallback before estimating or compressing its descendants. + pub(super) fn validate_canonical_encoding(&self, array: &ArrayRef) -> VortexResult<()> { + if let Some(allowed) = &self.canonical_encodings { + vortex_error::vortex_ensure!( + allowed.contains(&array.encoding_id()), + "canonical encoding {} is not permitted for this compressor", + array.encoding_id() + ); + } + Ok(()) + } + /// Compresses a child array produced by a cascading scheme. /// /// If the cascade budget is exhausted, the canonical array is returned as-is. Otherwise, the @@ -84,6 +119,7 @@ impl CascadingCompressor { exec_ctx: &mut ExecutionCtx, ) -> VortexResult { if parent_ctx.finished_cascading() { + self.validate_canonical_encoding(child)?; trace::cascade_exhausted(parent_id, child_index); return Ok(child.clone()); } @@ -108,6 +144,9 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { + if !matches!(&array, Canonical::List(_)) { + self.validate_canonical_encoding(&array.clone().into_array())?; + } match array { Canonical::Null(null_array) => Ok(null_array.into_array()), Canonical::Bool(bool_array) => { @@ -122,7 +161,9 @@ impl CascadingCompressor { Canonical::Struct(struct_array) => { let fields = struct_array .iter_unmasked_fields() - .map(|field| self.compress(field, exec_ctx)) + .map(|field| { + self.compress_with_context(field, compress_ctx.for_structure(), exec_ctx) + }) .collect::, _>>()?; Ok(StructArray::try_new( @@ -134,10 +175,16 @@ impl CascadingCompressor { .into_array()) } Canonical::Union(union_array) => { - let type_ids = self.compress(union_array.type_ids(), exec_ctx)?; + let type_ids = self.compress_with_context( + union_array.type_ids(), + compress_ctx.for_structure(), + exec_ctx, + )?; let children = union_array .iter_children() - .map(|child| self.compress(child, exec_ctx)) + .map(|child| { + self.compress_with_context(child, compress_ctx.for_structure(), exec_ctx) + }) .collect::, _>>()?; Ok( @@ -146,7 +193,10 @@ impl CascadingCompressor { ) } Canonical::List(list_view_array) => { - if list_view_array.is_zero_copy_to_list() || list_view_array.elements().is_empty() { + if self.convert_lists + && (list_view_array.is_zero_copy_to_list() + || list_view_array.elements().is_empty()) + { let list_array = list_from_list_view(list_view_array, exec_ctx)?; self.compress_list_array(list_array, compress_ctx, exec_ctx) } else { @@ -155,7 +205,11 @@ impl CascadingCompressor { } Canonical::Map(map_array) => self.compress_map_array(map_array, compress_ctx, exec_ctx), Canonical::FixedSizeList(fsl_array) => { - let compressed_elems = self.compress(fsl_array.elements(), exec_ctx)?; + let compressed_elems = self.compress_with_context( + fsl_array.elements(), + compress_ctx.for_structure(), + exec_ctx, + )?; Ok(FixedSizeListArray::try_new( compressed_elems, @@ -172,7 +226,7 @@ impl CascadingCompressor { // Try scheme-based compression first. let scheme_compressed = self.choose_and_compress( Canonical::Extension(ext_array.clone()), - compress_ctx, + compress_ctx.clone(), exec_ctx, )?; @@ -189,7 +243,11 @@ impl CascadingCompressor { // Also compress the underlying storage array. Some extension schemes can beat the // extension storage but still lose to ordinary storage compression. - let compressed_storage = self.compress(ext_array.storage_array(), exec_ctx)?; + let compressed_storage = self.compress_with_context( + ext_array.storage_array(), + compress_ctx.for_structure(), + exec_ctx, + )?; let storage_compressed = ExtensionArray::new(ext_array.ext_dtype().clone(), compressed_storage) .into_array(); @@ -201,16 +259,19 @@ impl CascadingCompressor { } } Canonical::Variant(variant_array) => { - let core_storage = - self.compress_physical_slots(variant_array.core_storage(), exec_ctx)?; + let core_storage = self.compress_physical_slots( + variant_array.core_storage(), + &compress_ctx, + exec_ctx, + )?; let shredded = variant_array .shredded() .map(|arr| { // Avoid stack-overflow for variant shredded values if arr.is::() { - self.compress_physical_slots(arr, exec_ctx) + self.compress_physical_slots(arr, &compress_ctx, exec_ctx) } else { - self.compress(arr, exec_ctx) + self.compress_with_context(arr, compress_ctx.for_structure(), exec_ctx) } }) .transpose()?; @@ -229,8 +290,8 @@ impl CascadingCompressor { /// If a winner is found and its compressed output is actually smaller, that output is /// returned. Otherwise, the original array is returned unchanged. /// - /// Empty, all-null, and constant arrays are handled by the compressor itself before any - /// scheme evaluation (constant detection is skipped while compressing samples). + /// Empty arrays remain canonical. All-null arrays only evaluate the registered constant + /// scheme. Other constant arrays use ordinary scheme selection. /// /// [`matches`]: Scheme::matches /// [`stats_options`]: Scheme::stats_options @@ -240,7 +301,7 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let eligible_schemes: Vec<&'static dyn Scheme> = self + let mut eligible_schemes: Vec<&'static dyn Scheme> = self .schemes .iter() .copied() @@ -248,17 +309,15 @@ impl CascadingCompressor { .collect(); let array: ArrayRef = canonical.into(); + self.validate_canonical_encoding(&array)?; if array.is_empty() { return Ok(array); } if array.all_invalid(exec_ctx)? { - return Ok( - ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()).into_array(), - ); + eligible_schemes.retain(|s| s.id() == ConstantScheme.id()); } - let before_nbytes = array.nbytes(); let merged_opts = eligible_schemes @@ -270,27 +329,6 @@ impl CascadingCompressor { let data = ArrayAndStats::new(array, merged_opts); - // Constant detection is built into the compressor: a constant leaf always short-circuits - // scheme selection. Samples are exempt because a constant sample does not imply that the - // full array is constant. - if !compress_ctx.is_sample() && constant::is_constant_for_compression(&data, exec_ctx)? { - let _winner_span = - trace::winner_compress_span(constant::CONSTANT_SCHEME_ID, before_nbytes).entered(); - let compressed = constant::compress_constant(data.array(), exec_ctx)?; - - let after_nbytes = compressed.nbytes(); - let actual_ratio = - (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); - let accepted = after_nbytes < before_nbytes; - trace::record_winner_compress_result(after_nbytes, None, actual_ratio, accepted); - - return if accepted { - Ok(compressed) - } else { - Ok(data.into_array()) - }; - } - if eligible_schemes.is_empty() { return Ok(data.into_array()); } diff --git a/vortex-compressor/src/compressor/constant.rs b/vortex-compressor/src/compressor/constant.rs index 7e8edc73f2f..02c31c59825 100644 --- a/vortex-compressor/src/compressor/constant.rs +++ b/vortex-compressor/src/compressor/constant.rs @@ -1,14 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Built-in constant detection and encoding. -//! -//! Constant arrays are not compressed through a pluggable [`Scheme`]: the compressor always -//! detects constant leaf arrays itself, before evaluating any registered scheme. Detection is -//! skipped while compressing samples, since a constant sample does not imply that the full array -//! is constant. -//! -//! [`Scheme`]: crate::scheme::Scheme +//! Shared detection and encoding for the registered constant compression schemes. use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -20,14 +13,8 @@ use vortex_array::dtype::DType; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; -use crate::scheme::SchemeId; use crate::stats::ArrayAndStats; -/// Synthetic scheme ID reported in traces when the compressor's built-in constant encoding wins. -pub(crate) const CONSTANT_SCHEME_ID: SchemeId = SchemeId { - name: "vortex.compressor.constant", -}; - /// Returns `true` if all valid values of the canonical array are equal, meaning the array can be /// encoded by [`compress_constant`]. /// @@ -141,20 +128,21 @@ mod tests { use vortex_session::VortexSession; use crate::CascadingCompressor; + use crate::builtins::ConstantScheme; + use crate::builtins::MaskedConstantScheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); - /// Constant detection is built into the compressor, so it must work with no schemes at all. - fn empty_compressor() -> CascadingCompressor { - CascadingCompressor::new(Vec::new()) + fn constant_compressor() -> CascadingCompressor { + CascadingCompressor::new(vec![&ConstantScheme, &MaskedConstantScheme]) } #[test] - fn constant_int_compresses_without_schemes() -> VortexResult<()> { + fn constant_int_compresses_with_constant_scheme() -> VortexResult<()> { let array = PrimitiveArray::new(buffer![7i64; 100], Validity::NonNullable).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let compressed = empty_compressor().compress(&array, &mut ctx)?; + let compressed = constant_compressor().compress(&array, &mut ctx)?; assert!(compressed.is::()); Ok(()) } @@ -166,33 +154,33 @@ mod tests { let array = PrimitiveArray::new(buffer![7i64; 100], validity).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let compressed = empty_compressor().compress(&array, &mut ctx)?; + let compressed = constant_compressor().compress(&array, &mut ctx)?; assert!(compressed.is::()); Ok(()) } #[test] - fn constant_string_compresses_without_schemes() -> VortexResult<()> { + fn constant_string_compresses_with_constant_scheme() -> VortexResult<()> { let array = VarBinViewArray::from_iter_str(std::iter::repeat_n("hello", 100)).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let compressed = empty_compressor().compress(&array, &mut ctx)?; + let compressed = constant_compressor().compress(&array, &mut ctx)?; assert!(compressed.is::()); Ok(()) } #[test] - fn constant_bool_compresses_without_schemes() -> VortexResult<()> { + fn constant_bool_compresses_with_constant_scheme() -> VortexResult<()> { let array = BoolArray::from_iter(std::iter::repeat_n(true, 100)).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let compressed = empty_compressor().compress(&array, &mut ctx)?; + let compressed = constant_compressor().compress(&array, &mut ctx)?; assert!(compressed.is::()); Ok(()) } #[test] - fn constant_decimal_compresses_without_schemes() -> VortexResult<()> { + fn constant_decimal_compresses_with_constant_scheme() -> VortexResult<()> { let array = DecimalArray::new( buffer![123_456i128; 100], DecimalDType::new(20, 2), @@ -201,29 +189,29 @@ mod tests { .into_array(); let mut ctx = SESSION.create_execution_ctx(); - let compressed = empty_compressor().compress(&array, &mut ctx)?; + let compressed = constant_compressor().compress(&array, &mut ctx)?; assert!(compressed.is::()); Ok(()) } #[test] - fn constant_timestamp_compresses_without_schemes() -> VortexResult<()> { + fn constant_timestamp_compresses_with_constant_scheme() -> VortexResult<()> { let ts = PrimitiveArray::from_iter(std::iter::repeat_n(1_704_067_200_000i64, 100)); let array = TemporalArray::new_timestamp(ts.into_array(), TimeUnit::Milliseconds, None) .into_array(); let mut ctx = SESSION.create_execution_ctx(); - let compressed = empty_compressor().compress(&array, &mut ctx)?; + let compressed = constant_compressor().compress(&array, &mut ctx)?; assert!(compressed.is::()); Ok(()) } #[test] - fn non_constant_int_is_left_canonical_without_schemes() -> VortexResult<()> { + fn non_constant_int_is_left_canonical() -> VortexResult<()> { let array = PrimitiveArray::from_iter(0..100i64).into_array(); let mut ctx = SESSION.create_execution_ctx(); - let compressed = empty_compressor().compress(&array, &mut ctx)?; + let compressed = constant_compressor().compress(&array, &mut ctx)?; assert!(!compressed.is::()); assert_eq!(compressed.dtype(), array.dtype()); Ok(()) diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..7c444ac68d2 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -4,11 +4,15 @@ //! Cascading array compression implementation. mod cascade; -mod constant; +pub(crate) mod constant; mod sample; mod select; mod structural; +use vortex_array::VTable; +use vortex_array::arrays::List; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -42,17 +46,22 @@ pub(crate) const ROOT_SCHEME_ID: SchemeId = SchemeId { pub struct CascadingCompressor { /// The enabled compression schemes. schemes: Vec<&'static dyn Scheme>, + /// Canonical forms admitted at construction, before any compression work. + canonical_encodings: Option>, /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + /// Whether construction permits replacing a list view with List storage. + convert_lists: bool, } impl CascadingCompressor { /// Creates a new compressor with the given schemes. /// /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built automatically. - pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { + pub fn new(mut schemes: Vec<&'static dyn Scheme>) -> Self { + schemes.sort_by_key(|s| s.selection_priority()); // Root exclusion: exclude IntDict from list/listview offsets (monotonically // increasing data where dictionary encoding is wasteful). let root_exclusions = vec![DescendantExclusion { @@ -62,9 +71,23 @@ impl CascadingCompressor { Self { schemes, + canonical_encodings: None, root_exclusions, + convert_lists: true, } } + + /// Construct with an exact set of serializable output IDs. + /// Schemes must already have been resolved against these IDs. + pub fn new_with_encodings( + schemes: Vec<&'static dyn Scheme>, + encodings: HashSet, + ) -> Self { + let mut compressor = Self::new(schemes); + compressor.convert_lists = encodings.contains(&List.id()); + compressor.canonical_encodings = Some(encodings); + compressor + } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 3c73d2d4cdb..8bb144baf8d 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -59,7 +59,7 @@ impl CascadingCompressor { /// current best [`EstimateScore`] as an early-exit hint so the callback can return /// [`EstimateVerdict::Skip`] without doing expensive work when it cannot beat the threshold. /// - /// Ties are broken by registration order within each pass. + /// Ties follow scheme priority, then registration order, within each pass. /// /// [`expected_compression_ratio`]: Scheme::expected_compression_ratio pub(super) fn choose_best_scheme( @@ -140,6 +140,9 @@ impl CascadingCompressor { /// exclusion rules. pub(super) fn is_excluded(&self, candidate: &dyn Scheme, ctx: &CompressorContext) -> bool { let id = candidate.id(); + if ctx.excludes_at_site(id) { + return true; + } let history = ctx.cascade_history(); // Self-exclusion: no scheme appears twice in any chain. diff --git a/vortex-compressor/src/compressor/structural.rs b/vortex-compressor/src/compressor/structural.rs index e0d9462ede1..5a7216e12a6 100644 --- a/vortex-compressor/src/compressor/structural.rs +++ b/vortex-compressor/src/compressor/structural.rs @@ -40,9 +40,14 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { + self.validate_canonical_encoding(&list_array.clone().into_array())?; let list_array = list_array.reset_offsets(true, exec_ctx)?; - let compressed_elems = self.compress(list_array.elements(), exec_ctx)?; + let compressed_elems = self.compress_with_context( + list_array.elements(), + compress_ctx.for_structure(), + exec_ctx, + )?; // Record the root scheme with the offsets child index so root exclusion rules apply. let offset_ctx = @@ -72,7 +77,12 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - let compressed_elems = self.compress(list_view.elements(), exec_ctx)?; + self.validate_canonical_encoding(&list_view.clone().into_array())?; + let compressed_elems = self.compress_with_context( + list_view.elements(), + compress_ctx.for_structure(), + exec_ctx, + )?; let offset_ctx = compress_ctx .clone() @@ -130,6 +140,7 @@ impl CascadingCompressor { pub(super) fn compress_physical_slots( &self, array: &ArrayRef, + compress_ctx: &CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let slots = array @@ -137,7 +148,9 @@ impl CascadingCompressor { .iter() .map(|slot| { slot.as_ref() - .map(|child| self.compress(child, exec_ctx)) + .map(|child| { + self.compress_with_context(child, compress_ctx.for_structure(), exec_ctx) + }) .transpose() }) .collect::>()?; diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 3a2a6281047..0955d0575f5 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -32,8 +32,10 @@ use super::ROOT_SCHEME_ID; use super::sample::estimate_compression_ratio_with_sampling; use super::select::WinnerEstimate; use super::structural; +use crate::builtins::ConstantScheme; use crate::builtins::FloatDictScheme; use crate::builtins::IntDictScheme; +use crate::builtins::MaskedConstantScheme; use crate::builtins::StringDictScheme; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; @@ -48,7 +50,13 @@ use crate::stats::GenerateStatsOptions; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); fn compressor() -> CascadingCompressor { - CascadingCompressor::new(vec![&IntDictScheme, &FloatDictScheme, &StringDictScheme]) + CascadingCompressor::new(vec![ + &ConstantScheme, + &MaskedConstantScheme, + &IntDictScheme, + &FloatDictScheme, + &StringDictScheme, + ]) } fn estimate_test_data() -> ArrayAndStats { @@ -705,9 +713,8 @@ fn all_null_array_compresses_to_constant() -> VortexResult<()> { ) .into_array(); - // The compressor should produce a `ConstantArray` for an all-null array regardless of - // which schemes are registered. - let compressor = CascadingCompressor::new(vec![&IntDictScheme]); + // Constant production requires the constant scheme to be registered. + let compressor = CascadingCompressor::new(vec![&ConstantScheme, &IntDictScheme]); let mut exec_ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array, &mut exec_ctx)?; assert!(compressed.is::()); diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 55bb9b188f6..48b9fa645a0 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -19,9 +19,8 @@ //! # Example //! //! A [`CascadingCompressor`] can be created directly with a fixed scheme list. With no schemes it -//! still canonicalizes supported inputs, recursively handles nested structure, and encodes -//! constant leaves (constant detection is built into the compressor), but no other leaf -//! compression is selected. +//! still canonicalizes supported inputs and recursively handles nested structure. Leaf compression, +//! including constants, requires explicitly registered schemes. //! //! ```rust //! use vortex_array::{IntoArray, VortexSessionExecute, array_session}; @@ -70,3 +69,8 @@ mod compressor; pub use compressor::CascadingCompressor; mod trace; + +mod patches; +pub use patches::compress_patches; + +pub mod session; diff --git a/vortex-btrblocks/src/schemes/patches.rs b/vortex-compressor/src/patches.rs similarity index 91% rename from vortex-btrblocks/src/schemes/patches.rs rename to vortex-compressor/src/patches.rs index 69ca8450f12..27ec053a3a1 100644 --- a/vortex-btrblocks/src/schemes/patches.rs +++ b/vortex-compressor/src/patches.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Shared patch compression for encoding-owned schemes. + use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -12,6 +14,10 @@ use vortex_error::VortexError; use vortex_error::VortexResult; /// Compresses the given patches by downscaling integers and checking for constant values. +/// +/// # Errors +/// +/// Returns an error if patch arrays cannot be evaluated or rebuilt. pub fn compress_patches(patches: Patches, ctx: &mut ExecutionCtx) -> VortexResult { // Downscale the patch indices. let indices = patches diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 4eed7538daa..36bfde1f481 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -38,6 +38,8 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, + /// Site-level exclusions supplied by the enclosing file layout. + site_exclusions: Vec, } impl CompressorContext { @@ -50,6 +52,7 @@ impl CompressorContext { allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), + site_exclusions: Vec::new(), } } } @@ -62,6 +65,25 @@ impl Default for CompressorContext { } impl CompressorContext { + /// Reset structural ancestry while preserving the enclosing layout's exclusions. + pub(crate) fn for_structure(&self) -> Self { + Self { + site_exclusions: self.site_exclusions.clone(), + ..Self::new() + } + } + + /// Exclude a redundant transformation already handled by the enclosing layout. + pub(crate) fn excluding_at_site(mut self, scheme: SchemeId) -> Self { + self.site_exclusions.push(scheme); + self + } + + /// Whether the layout excludes this scheme at this compression site. + pub(crate) fn excludes_at_site(&self, scheme: SchemeId) -> bool { + self.site_exclusions.contains(&scheme) + } + /// Whether this context is for sample compression (ratio estimation). pub fn is_sample(&self) -> bool { self.is_sample diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 0ba1c90202a..4ff75f14732 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -35,18 +35,21 @@ use crate::stats::GenerateStatsOptions; /// Unique identifier for a compression scheme. /// -/// The only way to obtain a [`SchemeId`] is through [`SchemeExt::id()`], which is auto-implemented -/// for all [`Scheme`] types. There is no public constructor. -/// -/// The only exception to this is for the compressor's synthetic `ROOT_SCHEME_ID`. +/// Use [`SchemeExt::id()`] for a known implementation, or [`SchemeId::new`] for an exclusion +/// referring to a scheme in another package. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct SchemeId { - /// Only constructable within `vortex-compressor`. - /// - /// The only public way to obtain a [`SchemeId`] is through [`SchemeExt::id()`]. + /// The globally unique scheme name. pub(super) name: &'static str, } +impl SchemeId { + /// Name a scheme in an exclusion rule without depending on its implementation crate. + pub const fn new(name: &'static str) -> Self { + Self { name } + } +} + impl fmt::Display for SchemeId { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str(self.name) @@ -118,6 +121,11 @@ impl fmt::Display for SchemeId { /// [`descendant_exclusions`]: Scheme::descendant_exclusions /// [`ancestor_exclusions`]: Scheme::ancestor_exclusions pub trait Scheme: Debug + Send + Sync { + /// Tie-break priority. Lower values are evaluated first, then registration order. + fn selection_priority(&self) -> u16 { + 1000 + } + /// The globally unique name for this scheme (e.g. `"vortex.int.bitpacking"`). fn scheme_name(&self) -> &'static str; @@ -127,9 +135,9 @@ pub trait Scheme: Debug + Send + Sync { /// The serialized IDs this scheme itself may write into its compressed output. /// /// Every declared ID must be permitted for the scheme to be used. Cascaded children are - /// compressed by other schemes, which declare their own IDs, so only arrays constructed - /// directly by [`compress`](Scheme::compress) belong here. Canonical arrays the scheme - /// merely rearranges do not need to be declared. + /// compressed by other schemes, which declare their own IDs, so arrays constructed + /// by [`compress`](Scheme::compress), its direct helpers, and its wrappers belong here. + /// Canonical arrays the scheme merely rearranges do not need to be declared. /// /// For most encodings this is the in-memory encoding ID. An encoding with several wire /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. @@ -186,10 +194,9 @@ pub trait Scheme: Debug + Send + Sync { /// constancy). Implementations should check `ctx.is_sample` to make sure that they are /// returning the correct information. /// - /// The compressor guarantees that empty and all-null arrays are handled before this method is - /// called, so implementations may assume the array has at least one valid element. Outside of - /// sample compression, the compressor also encodes constant arrays itself before evaluating - /// schemes, so implementations only see constant arrays when `ctx.is_sample()` is `true`. + /// Empty arrays skip scheme evaluation. All-null arrays only use the registered constant + /// scheme, so other schemes may assume at least one valid element. Constant values can reach + /// any scheme when constant compression is not registered or permitted. fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-compressor/src/session.rs b/vortex-compressor/src/session.rs new file mode 100644 index 00000000000..7e4e971eb85 --- /dev/null +++ b/vortex-compressor/src/session.rs @@ -0,0 +1,137 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Schemes explicitly registered for compression in a session. + +use std::any::Any; +use std::sync::Arc; + +use parking_lot::RwLock; +use vortex_session::SessionExt; +use vortex_session::SessionGuard; +use vortex_session::SessionVar; + +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::scheme::SchemeId; + +/// An ordered registry. Clones share registrations and registration is atomic. +#[derive(Clone, Debug, Default)] +pub struct CompressionSession { + /// Shared registration order, protected across deduplication and insertion. + schemes: Arc>>, +} + +impl CompressionSession { + /// Create an empty registry. + pub fn empty() -> Self { + Self::default() + } + + /// Register a scheme. An already registered ID keeps its original instance and position. + pub fn register(&self, scheme: &'static dyn Scheme) { + let mut schemes = self.schemes.write(); + if !schemes.iter().any(|s| s.id() == scheme.id()) { + schemes.push(scheme); + } + } + + /// Remove a registered scheme from future compressor construction. + pub fn unregister(&self, id: SchemeId) { + self.schemes.write().retain(|s| s.id() != id); + } + + /// Snapshot the registered schemes in registration order. + pub fn schemes(&self) -> Vec<&'static dyn Scheme> { + self.schemes.read().clone() + } +} + +impl SessionVar for CompressionSession { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +/// Access to the session's compression registry. +pub trait CompressionSessionExt: SessionExt { + /// Return the registry, initially empty. + fn compression(&self) -> SessionGuard<'_, CompressionSession> { + self.get::() + } + + /// Register a scheme for future compression. + fn register_scheme(&self, scheme: &'static dyn Scheme) { + self.compression().register(scheme); + } + + /// Fork session configuration with an independent copy of the compression registry. + /// Other registered services remain shared. Use for per-operation scheme registration. + fn fork_compression(&self) -> vortex_session::VortexSession { + let session = self.session().fork(); + let registry = CompressionSession::empty(); + for scheme in self.registered_schemes() { + registry.register(scheme); + } + session.register(registry); + session + } + + /// Snapshot the schemes selected by registration. + fn registered_schemes(&self) -> Vec<&'static dyn Scheme> { + self.compression().schemes() + } +} +impl CompressionSessionExt for S {} + +#[cfg(test)] +mod tests { + use vortex_session::VortexSession; + + use super::CompressionSessionExt; + use crate::builtins::FloatDictScheme; + use crate::builtins::IntDictScheme; + use crate::builtins::StringDictScheme; + use crate::scheme::Scheme; + use crate::scheme::SchemeExt; + + #[test] + fn registration_is_explicit_and_atomic() { + let session = VortexSession::empty(); + assert!(session.registered_schemes().is_empty()); + let schemes: &[&'static dyn Scheme] = + &[&IntDictScheme, &FloatDictScheme, &StringDictScheme]; + std::thread::scope(|scope| { + for &scheme in schemes { + let session = session.clone(); + scope.spawn(move || { + for _ in 0..32 { + session.register_scheme(scheme); + } + }); + } + }); + let registered = session.registered_schemes(); + assert_eq!(registered.len(), schemes.len()); + assert!( + schemes + .iter() + .all(|expected| registered.iter().any(|actual| actual.id() == expected.id())) + ); + } + + #[test] + fn per_operation_registration_is_isolated() { + let session = VortexSession::empty(); + session.register_scheme(&IntDictScheme); + let fork = session.fork_compression(); + fork.register_scheme(&FloatDictScheme); + fork.compression().unregister(IntDictScheme.id()); + assert_eq!(session.registered_schemes()[0].id(), IntDictScheme.id()); + assert_eq!(fork.registered_schemes()[0].id(), FloatDictScheme.id()); + } +} diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 7cc7d0322da..b52a509edfd 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -29,16 +29,23 @@ use vortex::array::serde::SerializedArray; use vortex::array::stats::StatsSetRef; use vortex::buffer::BufferString; use vortex::buffer::ByteBuffer; -use vortex::compressor::BtrBlocksCompressorBuilder; +use vortex::compressor::BtrBlocksCompressor; +use vortex::compressor::CascadingCompressor; +use vortex::compressor::CompressionSessionExt; use vortex::dtype::DType; use vortex::dtype::FieldMask; -use vortex::editions::ComponentKind; +#[cfg(test)] use vortex::editions::Edition; +#[cfg(test)] use vortex::editions::EditionDeclaration; -use vortex::editions::EditionFamily; +#[cfg(test)] use vortex::editions::EditionId; -use vortex::editions::EditionMember; use vortex::editions::EditionSessionExt; +use vortex::editions::cuda::CUDA_2026_09_0 as CUDA_EDITION; +use vortex::editions::cuda::DECLARATION as CUDA_EDITION_DECLARATION; +use vortex::editions::cuda::FAMILY as CUDA_EDITION_FAMILY; +use vortex::encodings::zstd::schemes::binary::ZstdScheme; +use vortex::encodings::zstd::schemes::binary_buffers::ZstdBuffersScheme; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -51,6 +58,7 @@ use vortex::layout::LayoutEncodingRef; use vortex::layout::LayoutId; use vortex::layout::LayoutParts; use vortex::layout::LayoutReader; +use vortex::layout::LayoutReaderContext; use vortex::layout::LayoutReaderRef; use vortex::layout::LayoutRef; use vortex::layout::LayoutStrategy; @@ -198,7 +206,7 @@ impl VTable for CudaFlat { name: Arc, segment_source: Arc, session: &VortexSession, - _ctx: &vortex::layout::LayoutReaderContext, + _ctx: &LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(CudaFlatReader { layout: layout.clone(), @@ -549,26 +557,28 @@ fn extract_constant_buffers(chunk: &ArrayRef) -> Vec { /// Build a CUDA-flat writer using only CUDA-compatible, session-enabled array encodings. /// -/// Requires [`register_cuda_layout`]. Zero `block_rows` uses default sizing and dictionary policy; +/// Selects the complete CUDA edition on the shared session and registers the CUDA layout. +/// Zero `block_rows` uses default sizing and dictionary policy; /// nonzero sets row blocks without outer dictionaries or byte coalescing, retaining per-block /// dictionary compression. 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() + register_cuda_layout(session); + session + .set_enabled_editions([CUDA_EDITION]) + .vortex_expect("CUDA edition is registered"); + let compression_session = CompressionSessionExt::fork_compression(session); + CompressionSessionExt::register_scheme(&compression_session, &ZstdScheme); + CompressionSessionExt::register_scheme(&compression_session, &ZstdBuffersScheme); + let compressor = BtrBlocksCompressor::from_session(&compression_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_btrblocks_compressor(compressor).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(compressor) + .with_probe_compressor(BtrBlocksCompressor(CascadingCompressor::new(Vec::new()))) .with_row_block_size(block_rows) .with_data_block_target_bytes(None) .build() @@ -594,25 +604,10 @@ impl SessionVar for CudaLayoutRegistration { } } -const CUDA_EDITION_FAMILY: EditionFamily = EditionFamily { - name: "cuda", - origin: "vortex-cuda", - doc: "CUDA-readable layouts, enabled only when CUDA layout support is registered.", -}; -const CUDA_EDITION: EditionId = EditionId::new("cuda", 2026, 9, 0); -static CUDA_EDITION_DECLARATION: EditionDeclaration = EditionDeclaration { - edition: Edition { - id: CUDA_EDITION, - min_library_version: None, - }, - added: &[EditionMember::layout(&"vortex.cuda_flat")], -}; - /// Register [`CudaFlat`] and its draft `cuda` edition once per session. /// -/// Enables a newly registered edition only if no `cuda` edition is selected; otherwise preserves -/// writer policy, including on repeated calls. The draft has no cross-version compatibility -/// guarantee. Readers must also register the layout. +/// Registration preserves the selected output editions. CUDA writers explicitly select +/// the complete CUDA edition. This draft has no cross-version compatibility guarantee. /// /// Call alongside [`crate::initialize_cuda`]; registration itself needs no GPU. pub fn register_cuda_layout(session: &VortexSession) { @@ -633,16 +628,6 @@ pub fn register_cuda_layout(session: &VortexSession) { session .register_edition(&CUDA_EDITION_DECLARATION) .vortex_expect("CUDA edition declaration is valid"); - if !session - .enabled_editions() - .editions() - .iter() - .any(|edition| edition.family == CUDA_EDITION.family) - { - session - .enable_edition(CUDA_EDITION) - .vortex_expect("CUDA edition is registered"); - } }); } @@ -659,6 +644,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; @@ -766,7 +752,6 @@ mod tests { let session = VortexSession::default(); session.enable_edition(CORE_2025_05_0)?; let mut expected_editions = session.enabled_editions().editions(); - expected_editions.push(CUDA_EDITION); expected_editions.sort_unstable(); let expected_arrays = session.enabled_component_ids(ComponentKind::Array); let barrier = std::sync::Barrier::new(4); @@ -778,11 +763,7 @@ mod tests { scope.spawn(move || { barrier.wait(); register_cuda_layout(&session); - assert!( - session - .enabled_component_ids(ComponentKind::Layout) - .contains(&CudaFlat.id()) - ); + assert!(session.editions().find(&CUDA_EDITION).is_some()); }); } }); @@ -805,8 +786,6 @@ mod tests { let session = VortexSession::default(); if register_first { register_cuda_layout(&session); - } else { - session.editions().declare_family(&CUDA_EDITION_FAMILY)?; } session.register_edition(&EditionDeclaration { edition: Edition { @@ -835,8 +814,6 @@ mod tests { #[test] fn test_cuda_registration_preserves_disabled_pre_registered_edition() -> VortexResult<()> { let session = VortexSession::default(); - session.editions().declare_family(&CUDA_EDITION_FAMILY)?; - session.register_edition(&CUDA_EDITION_DECLARATION)?; let mut expected_editions = session.enabled_editions().editions(); expected_editions.sort_unstable(); let expected_layouts = session.enabled_component_ids(ComponentKind::Layout); diff --git a/vortex-edition/src/declarations/cuda.rs b/vortex-edition/src/declarations/cuda.rs new file mode 100644 index 00000000000..da6df228692 --- /dev/null +++ b/vortex-edition/src/declarations/cuda.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Draft output contract for files read through the CUDA reader. +//! +//! This edition includes the host-side metadata and canonical fallback representations. +//! It describes wire compatibility. It does not guarantee that every dtype or operation executes +//! on the GPU, or that a particular compression scheme is fast on the GPU. + +use crate::Edition; +use crate::EditionDeclaration; +use crate::EditionFamily; +use crate::EditionId; +use crate::EditionMember; + +/// The CUDA reader's independently versioned wire-format family. +pub static FAMILY: EditionFamily = EditionFamily { + name: "cuda", + origin: "vortex-cuda", + doc: "Serialized representations supported by the CUDA reader, including host-side metadata.", +}; + +/// The September 2026 draft CUDA output edition. +pub const CUDA_2026_09_0: EditionId = EditionId::new("cuda", 2026, 9, 0); + +/// Complete output permissions for the CUDA reader. +pub static DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: CUDA_2026_09_0, + min_library_version: None, + }, + added: &[ + EditionMember::array(&"fastlanes.bitpacked"), + EditionMember::array(&"fastlanes.for"), + EditionMember::array(&"vortex.alp"), + EditionMember::array(&"vortex.bool"), + EditionMember::array(&"vortex.bytebool"), + EditionMember::array(&"vortex.chunked"), + EditionMember::array(&"vortex.constant"), + EditionMember::array(&"vortex.datetimeparts"), + EditionMember::array(&"vortex.decimal"), + EditionMember::array(&"vortex.decimal_byte_parts"), + EditionMember::array(&"vortex.dict"), + EditionMember::array(&"vortex.ext"), + EditionMember::array(&"vortex.fsst"), + EditionMember::array(&"vortex.list"), + EditionMember::array(&"vortex.null"), + EditionMember::array(&"vortex.primitive"), + EditionMember::array(&"vortex.runend"), + EditionMember::array(&"vortex.struct"), + EditionMember::array(&"vortex.varbin"), + EditionMember::array(&"vortex.varbinview"), + EditionMember::array(&"vortex.zigzag"), + EditionMember::layout(&"vortex.chunked"), + EditionMember::layout(&"vortex.dict"), + EditionMember::layout(&"vortex.flat"), + EditionMember::layout(&"vortex.stats"), + EditionMember::layout(&"vortex.struct"), + EditionMember::dtype(&"vortex.date"), + EditionMember::dtype(&"vortex.time"), + EditionMember::dtype(&"vortex.timestamp"), + EditionMember::array(&"vortex.sequence"), + EditionMember::array(&"vortex.zstd"), + EditionMember::array(&"vortex.fixed_size_list"), + EditionMember::array(&"vortex.listview"), + EditionMember::array(&"vortex.masked"), + EditionMember::layout(&"vortex.zoned"), + EditionMember::aggregate(&"vortex.bounded_max"), + EditionMember::aggregate(&"vortex.bounded_min"), + EditionMember::aggregate(&"vortex.max"), + EditionMember::aggregate(&"vortex.min"), + EditionMember::aggregate(&"vortex.nan_count"), + EditionMember::aggregate(&"vortex.null_count"), + EditionMember::array(&"vortex.onpair"), + EditionMember::array(&"vortex.map"), + EditionMember::dtype(&"vortex.uuid"), + EditionMember::array(&"fastlanes.delta"), + EditionMember::array(&"vortex.zstd_buffers"), + EditionMember::layout(&"vortex.cuda_flat"), + ], +}; diff --git a/vortex-edition/src/declarations/mod.rs b/vortex-edition/src/declarations/mod.rs index 39132db9283..4bf639e08bb 100644 --- a/vortex-edition/src/declarations/mod.rs +++ b/vortex-edition/src/declarations/mod.rs @@ -12,13 +12,14 @@ //! the declarations and selecting which of them the default writer may emit. pub mod core; +pub mod cuda; pub mod preview; use crate::EditionDeclaration; use crate::EditionFamily; /// The first-party edition families. Every family must be declared before its editions. -pub static EDITION_FAMILIES: &[&EditionFamily] = &[&core::FAMILY, &preview::FAMILY]; +pub static EDITION_FAMILIES: &[&EditionFamily] = &[&core::FAMILY, &preview::FAMILY, &cuda::FAMILY]; /// The first-party Vortex edition declarations. pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ @@ -30,4 +31,5 @@ pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ &core::v2026_08_2::DECLARATION, &core::v2026_08_3::DECLARATION, &preview::v2026_08::DECLARATION, + &cuda::DECLARATION, ]; diff --git a/vortex-edition/src/session.rs b/vortex-edition/src/session.rs index 4b192e3d7d9..d60588494dd 100644 --- a/vortex-edition/src/session.rs +++ b/vortex-edition/src/session.rs @@ -353,6 +353,31 @@ pub trait EditionSessionExt: SessionExt { Ok(()) } + /// Replace the selected editions atomically, with at most one edition per family. + /// All editions must be registered. An invalid selection leaves the previous set intact. + fn set_enabled_editions( + &self, + editions: impl IntoIterator, + ) -> VortexResult<()> { + let enabled = self.enabled_editions(); + let mut selected = enabled.inner.snapshot().as_ref().clone(); + selected.clear(); + let registry = self.editions(); + for edition in editions { + if registry.find(&edition).is_none() { + vortex_bail!("cannot enable unregistered edition {edition}"); + } + if selected.insert(Id::from(edition.family), edition).is_some() { + vortex_bail!( + "only one edition per family may be selected, got {}", + edition.family + ); + } + } + enabled.inner.replace(selected); + Ok(()) + } + /// Resolve the ids of one [`ComponentKind`] across all enabled editions: what a writer may /// emit for that kind. /// diff --git a/vortex-edition/src/tests.rs b/vortex-edition/src/tests.rs index 63cf003d11f..07278d87c79 100644 --- a/vortex-edition/src/tests.rs +++ b/vortex-edition/src/tests.rs @@ -15,6 +15,8 @@ use crate::EditionMember; use crate::EditionSession; use crate::EditionSessionExt; use crate::EnabledEditions; +use crate::declarations::core::CORE_2026_08_3; +use crate::declarations::cuda::CUDA_2026_09_0; static TEST_FAMILY: EditionFamily = EditionFamily { name: "test", @@ -444,3 +446,38 @@ fn kinds_are_resolved_independently() -> VortexResult<()> { ); Ok(()) } + +#[test] +fn selecting_cuda_replaces_core_and_preserves_additive_composition() -> VortexResult<()> { + let session = VortexSession::empty(); + for family in crate::EDITION_FAMILIES { + session.editions().declare_family(family)?; + } + for declaration in crate::EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + let core = CORE_2026_08_3; + let cuda = CUDA_2026_09_0; + session.enable_edition(core)?; + session.set_enabled_editions([cuda])?; + assert_eq!(session.enabled_editions().editions(), vec![cuda]); + assert!( + !session + .enabled_component_ids(ComponentKind::Array) + .iter() + .any(|id| id.as_str() == "vortex.sparse") + ); + + let unknown = EditionId::new("missing", 2026, 9, 0); + assert!(session.set_enabled_editions([core, unknown]).is_err()); + assert_eq!(session.enabled_editions().editions(), vec![cuda]); + + session.enable_edition(core)?; + assert!( + session + .enabled_component_ids(ComponentKind::Array) + .iter() + .any(|id| id.as_str() == "vortex.sparse") + ); + Ok(()) +} 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..05397bd2c5b 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::BtrBlocksCompressor; 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_btrblocks_compressor(BtrBlocksCompressor::for_memory(&SESSION)) .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_btrblocks_compressor(BtrBlocksCompressor::for_memory(&SESSION)); 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..31f482ccdf2 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -198,6 +198,7 @@ pub fn register_default_encodings(session: &VortexSession) { #[cfg(test)] pub(crate) fn enable_all_registered_array_encodings(session: &VortexSession) { + vortex_btrblocks::initialize(session); use vortex_array::dtype::session::DTypeSessionExt; use vortex_edition::ComponentKind; use vortex_edition::Edition; diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index c8110fe88c3..0bb754f6471 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -7,9 +7,7 @@ use std::num::NonZeroUsize; use std::sync::Arc; use vortex_array::dtype::FieldPath; -use vortex_btrblocks::BtrBlocksCompressorBuilder; -use vortex_btrblocks::SchemeExt; -use vortex_btrblocks::schemes::integer::IntDictScheme; +use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexExpect; use vortex_layout::LayoutStrategy; use vortex_layout::layouts::buffered::BufferedStrategy; @@ -26,16 +24,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. + /// A resolved compressor, with dictionary-code context applied to data compression. /// `IntDictScheme` is automatically excluded from the data compressor to prevent recursive /// dictionary encoding. - BtrBlocks(BtrBlocksCompressorBuilder), + BtrBlocks(BtrBlocksCompressor), /// An opaque compressor used as-is for both data and stats compression. Opaque(Arc), } @@ -64,12 +63,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` and the + /// array encodings its enabled editions permit. It can be further configured, and then + /// finally built yielding the [`LayoutStrategy`]. + pub fn from_session(session: &VortexSession) -> Self { Self { - compressor: CompressorConfig::BtrBlocks(BtrBlocksCompressorBuilder::default()), + compressor: CompressorConfig::BtrBlocks(BtrBlocksCompressor::from_session(session)), row_block_size: 8192, data_block_target_bytes: Some(ONE_MEG), field_writers: HashMap::new(), @@ -132,12 +132,12 @@ impl WriteStrategyBuilder { self } - /// Override the default [`BtrBlocksCompressorBuilder`] used for compression. + /// Use a resolved compressor for data and statistics. /// - /// 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); + /// Data compression receives the outer dictionary exclusion. Statistics use the + /// same registered schemes without that exclusion. + pub fn with_btrblocks_compressor(mut self, compressor: BtrBlocksCompressor) -> Self { + self.compressor = CompressorConfig::BtrBlocks(compressor); self } @@ -177,12 +177,12 @@ 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::BtrBlocks(compressor) => Arc::new({ + let compressor = compressor.clone(); + move |array: &vortex_array::ArrayRef, ctx: &mut vortex_array::ExecutionCtx| { + compressor.compress_dictionary_codes(array, ctx) + } + }), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; let compressing = CompressingStrategy::new(buffered, data_compressor); @@ -206,7 +206,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::BtrBlocks(compressor) => Arc::new(compressor), 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..8f065fd670c 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -65,7 +65,7 @@ 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::SchemeExt; use vortex_btrblocks::schemes::string::StringDictScheme; use vortex_buffer::Buffer; @@ -101,6 +101,7 @@ use crate::VortexFile; use crate::WriteOptionsSessionExt; use crate::flatbuffers::footer as fb; use crate::footer::SegmentSpec; +use crate::strategy::WriteStrategyBuilder; static SESSION: LazyLock = LazyLock::new(|| { let session = array_session() .with::() @@ -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 = 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 = 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(WriteStrategyBuilder::from_session(&SESSION).build()) .write(&mut buf, strings.clone().to_array_stream()) .await?; assert!( @@ -2586,14 +2587,19 @@ 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 = { + let compression_session = + vortex_btrblocks::CompressionSessionExt::fork_compression(&*SESSION); + vortex_btrblocks::CompressionSessionExt::compression(&compression_session) + .unregister(StringDictScheme.id()); + BtrBlocksCompressor::from_session(&compression_session) + }; let mut buf = ByteBufferMut::empty(); let summary = SESSION .write_options() .with_strategy( - crate::strategy::WriteStrategyBuilder::default() - .with_btrblocks_builder(no_string_dict) + WriteStrategyBuilder::from_session(&SESSION) + .with_btrblocks_compressor(no_string_dict) .build(), ) .write(&mut buf, strings.to_array_stream()) @@ -2614,15 +2620,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 = { + let compression_session = + vortex_btrblocks::CompressionSessionExt::fork_compression(&*SESSION); + vortex_btrblocks::CompressionSessionExt::compression(&compression_session) + .unregister(StringDictScheme.id()); + BtrBlocksCompressor::from_session(&compression_session) + }; let mut buf = ByteBufferMut::empty(); let summary = SESSION .write_options() .with_strategy( - crate::strategy::WriteStrategyBuilder::default() + 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..9174e8a9313 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::BtrBlocksCompressor; use vortex_buffer::ByteBuffer; use vortex_edition::ComponentKind; use vortex_edition::EditionSessionExt; @@ -251,11 +251,11 @@ impl VortexWriteOptions { 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 => WriteStrategyBuilder::from_session(&self.session) + .with_btrblocks_compressor(BtrBlocksCompressor::from_session_with_encodings( + &self.session, + allowed_serialized_ids, + )) .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..3868c4aaaab 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -383,6 +383,7 @@ mod tests { use vortex_array::expr::lit; use vortex_array::expr::pack; use vortex_array::expr::root; + use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexExpect; @@ -411,10 +412,12 @@ mod tests { // FIXME(ngates): Deprecate the global `runtime::single::block_on` helper and require tests // to call `block_on` on an explicit runtime instance. fn session_with_handle(handle: Handle) -> VortexSession { - array_session() + let session = array_session() .with::() .with::() - .with_handle(handle) + .with_handle(handle); + vortex_btrblocks::initialize(&session); + session } async fn write_dict_layout( @@ -426,7 +429,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::for_memory(session)), ); let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); @@ -456,7 +459,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::for_memory(&session)), ); let array = VarBinArray::from_iter( @@ -559,7 +562,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::for_memory(&session)), ); let array = @@ -584,10 +587,7 @@ mod tests { let filter = eq( root(), - lit(vortex_array::scalar::Scalar::utf8( - filter_value, - Nullability::Nullable, - )), + lit(Scalar::utf8(filter_value, Nullability::Nullable)), ); let reader = layout .new_reader("".into(), segments, &session, &Default::default()) @@ -614,7 +614,7 @@ mod tests { FlatLayoutStrategy::default(), FlatLayoutStrategy::default(), DictLayoutOptions::default(), - Arc::new(BtrBlocksCompressor::default()), + Arc::new(BtrBlocksCompressor::for_memory(&session)), ); let array = VarBinArray::from_iter( diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..7334e7b76fd 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -117,8 +117,8 @@ impl TableStrategy { /// # use vortex_layout::layouts::table::TableStrategy; /// /// // A strategy for compressing data using the balanced BtrBlocks compressor. - /// let compress = - /// CompressingStrategy::new(FlatLayoutStrategy::default(), BtrBlocksCompressor::default()); + /// let compressor = BtrBlocksCompressor::from_session(&session); + /// let compress = CompressingStrategy::new(FlatLayoutStrategy::default(), compressor); /// /// // Our combined strategy uses no compression for validity buffers, BtrBlocks compression /// // for most columns, and stores a nested binary column uncompressed (flat) because it diff --git a/vortex-python/src/compress.rs b/vortex-python/src/compress.rs index 8688fbabf35..d3f9c2d7f74 100644 --- a/vortex-python/src/compress.rs +++ b/vortex-python/src/compress.rs @@ -56,7 +56,8 @@ pub fn compress(py: Python, array: PyArrayRef) -> PyVortexResult { let session = session(); let array = array.into_inner(); let compressed = py.detach(move || { - BtrBlocksCompressor::default().compress(&array, &mut session.create_execution_ctx()) + BtrBlocksCompressor::from_session(session) + .compress(&array, &mut session.create_execution_ctx()) })?; Ok(PyArrayRef::from(compressed)) } diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 7288fc5d1e3..6ab75c6ab64 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::BtrBlocksCompressor; +use vortex::compressor::CompressionSessionExt; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::file::WriteOptionsSessionExt; @@ -380,16 +379,12 @@ 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 compression_session = CompressionSessionExt::fork_compression(session); if self.use_compact_encodings { - compressor = compressor.with_compact(); + vortex::compressor::initialize_compact(&compression_session); } - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); + let strategy = WriteStrategyBuilder::from_session(session) + .with_btrblocks_compressor(BtrBlocksCompressor::from_session(&compression_session)); let strategy = strategy.build(); current_runtime().block_on(async move { match resolve_store(path, store.map(|x| x.into_inner()))? { diff --git a/vortex-session/src/arc_swap_map.rs b/vortex-session/src/arc_swap_map.rs index a0eb08d085d..54179009d70 100644 --- a/vortex-session/src/arc_swap_map.rs +++ b/vortex-session/src/arc_swap_map.rs @@ -60,6 +60,18 @@ impl Debug for ArcSwapMap { } impl ArcSwapMap { + /// Create an independent map initially containing this snapshot's values. + pub fn fork(&self) -> Self { + Self { + inner: Arc::new(ArcSwap::from(self.snapshot())), + } + } + + /// Atomically replace the map contents. + pub fn replace(&self, values: HashMap) { + self.inner.store(Arc::new(values)); + } + /// Return the currently published map snapshot. pub fn snapshot(&self) -> Arc> { self.inner.load_full() diff --git a/vortex-session/src/session.rs b/vortex-session/src/session.rs index 441a37031a5..43f742fbaf3 100644 --- a/vortex-session/src/session.rs +++ b/vortex-session/src/session.rs @@ -159,7 +159,8 @@ impl Debug for SessionMut<'_, V> { /// /// Cloning a session is cheap and shares the backing store: a variable registered through one /// clone (via [`VortexSession::with_some`] or one of the `with_*` helpers) is observed by all -/// clones. To build an *independent* session, start from [`VortexSession::empty`]. +/// clones. [`VortexSession::fork`] copies the variable map so selected variables can be replaced +/// independently. For independent services as well, start from [`VortexSession::empty`]. #[derive(Clone)] pub struct VortexSession(SharedSessionVars); @@ -176,6 +177,12 @@ impl VortexSession { Self(SharedSessionVars::default()) } + /// Snapshot the variable map into an independent session. + /// Existing variables remain shared until explicitly replaced in the fork. + pub fn fork(&self) -> Self { + Self(self.0.fork()) + } + /// Inserts `V::default()` if no variable of type `V` is present yet, copy-on-write. /// /// The default is constructed *before* the map is updated, so `V::default()` never runs under a 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..d0f505bb127 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::BtrBlocksCompressor; + use vortex::compressor::CompressionSessionExt; use vortex::editions::CORE_2026_08_3; use vortex::editions::EditionSessionExt; use vortex::file::WriteStrategyBuilder; @@ -44,15 +45,21 @@ mod tests { let regular_bytes = adapter::write_compressed_to_bytes_with_session( &session, array.clone(), - WriteStrategyBuilder::default().build(), + WriteStrategyBuilder::from_session(&session) + .with_btrblocks_compressor(BtrBlocksCompressor::for_memory(&session)) + .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_btrblocks_compressor({ + let compression_session = CompressionSessionExt::fork_compression(&session); + vortex::compressor::initialize_compact(&compression_session); + BtrBlocksCompressor::for_memory(&compression_session) + }) .build(), )?; let _compact = adapter::read_file(compact_bytes)?; diff --git a/vortex-test/compat-gen/src/fixtures/mod.rs b/vortex-test/compat-gen/src/fixtures/mod.rs index 291d9a5f5ff..ee567b8148c 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::BtrBlocksCompressor; +use vortex::compressor::CompressionSessionExt; use vortex::file::WriteStrategyBuilder; +use vortex::session::VortexSession; use vortex_array::ExecutionCtx; use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; @@ -137,13 +140,22 @@ 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_btrblocks_compressor({ + let compression_session = CompressionSessionExt::fork_compression(&session); + vortex::compressor::initialize_compact(&compression_session); + BtrBlocksCompressor::for_memory(&compression_session) + }) .build(); adapter::write_compressed(&path, array, strategy)?; } else { - let strategy = WriteStrategyBuilder::default().build(); + let strategy = WriteStrategyBuilder::from_session(&session) + .with_btrblocks_compressor(BtrBlocksCompressor::for_memory(&session)) + .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..b8fd72ecf3f 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::BtrBlocksCompressor; +use vortex::compressor::CompressionSessionExt; use vortex::error::VortexExpect; use vortex::error::vortex_err; use vortex::file::WriteOptionsSessionExt; @@ -98,16 +97,12 @@ 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 compression_session = CompressionSessionExt::fork_compression(session); if matches!(flags.strategy, Strategy::Compact) { - compressor = compressor.with_compact(); + vortex::compressor::initialize_compact(&compression_session); } - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); + let strategy = WriteStrategyBuilder::from_session(session) + .with_btrblocks_compressor(BtrBlocksCompressor::from_session(&compression_session)); let mut file = File::create(output_path).await?; session diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 49fc0a44ddf..4b75da7d916 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -83,7 +83,12 @@ tokio = [ "vortex-io/tokio", "vortex-layout/tokio", ] -zstd = ["dep:vortex-zstd", "vortex-file?/zstd"] +zstd = [ + "dep:vortex-zstd", + "vortex-file?/zstd", + "vortex-btrblocks/zstd", + "vortex-btrblocks/pco", +] 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..da9a9ba1ebf 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -48,7 +48,7 @@ //! # fn example() -> vortex::error::VortexResult<()> { //! let session = VortexSession::default(); //! let array = PrimitiveArray::new(buffer![42u64; 1024], Validity::NonNullable).into_array(); -//! let compressed = BtrBlocksCompressor::default() +//! let compressed = BtrBlocksCompressor::from_session(&session) //! .compress(&array, &mut session.create_execution_ctx())?; //! //! assert_eq!(compressed.dtype(), array.dtype()); @@ -114,6 +114,7 @@ pub use vortex_array::scalar_fn; use vortex_array::scalar_fn::session::ScalarFnSession; use vortex_array::session::ArraySession; use vortex_array::stats::session::StatsSession; +use vortex_btrblocks::CompressionSession; use vortex_io::session::RuntimeSession; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; @@ -144,9 +145,14 @@ pub mod buffer { /// Default adaptive compression APIs based on the maintained BtrBlocks-style compressor. pub mod compressor { pub use vortex_btrblocks::BtrBlocksCompressor; - pub use vortex_btrblocks::BtrBlocksCompressorBuilder; + pub use vortex_btrblocks::CascadingCompressor; + pub use vortex_btrblocks::CompressionSession; + pub use vortex_btrblocks::CompressionSessionExt; pub use vortex_btrblocks::Scheme; + pub use vortex_btrblocks::SchemeExt; pub use vortex_btrblocks::SchemeId; + pub use vortex_btrblocks::initialize; + pub use vortex_btrblocks::initialize_compact; } /// Vortex editions: versioned sets of serialized components. @@ -318,6 +324,7 @@ impl VortexSessionDefault for VortexSession { let session = VortexSession::empty() .with::() .with::() + .with::() .with::() .with::() .with::() @@ -325,6 +332,7 @@ impl VortexSessionDefault for VortexSession { .with::() .with::() .with::(); + vortex_btrblocks::initialize(&session); vortex_arrow::initialize(&session); vortex_parquet_variant::initialize(&session); editions::register_default_editions(&session); @@ -366,7 +374,7 @@ mod test { use vortex_array::expr::select; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; - use vortex_btrblocks::BtrBlocksCompressorBuilder; + use vortex_btrblocks::BtrBlocksCompressor; 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,13 @@ mod test { session .write_options() .with_strategy( - WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().with_compact()) + WriteStrategyBuilder::from_session(&session) + .with_btrblocks_compressor({ + let compression_session = + vortex_btrblocks::CompressionSessionExt::fork_compression(&session); + vortex_btrblocks::initialize_compact(&compression_session); + BtrBlocksCompressor::from_session(&compression_session) + }) .build(), ) .write( diff --git a/wasm-test/src/main.rs b/wasm-test/src/main.rs index 964d3a36c9a..f9c8da4630a 100644 --- a/wasm-test/src/main.rs +++ b/wasm-test/src/main.rs @@ -1,25 +1,27 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex::VortexSessionDefault; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::PrimitiveArray; use vortex::array::validity::Validity; -use vortex::buffer::buffer; +use vortex::buffer::Buffer; use vortex::compressor::BtrBlocksCompressor; use vortex::session::VortexSession; -use vortex::VortexSessionDefault; -//use wasm_bindgen::prelude::*; +// use wasm_bindgen::prelude::*; pub fn main() { // Extremely simple test of compression/decompression and a few compute functions. - let array = PrimitiveArray::new(buffer![1i32; 1024], Validity::AllValid).into_array(); + let array = + PrimitiveArray::new((0..1024i32).collect::>(), Validity::AllValid).into_array(); let session = VortexSession::default(); - let compressed = BtrBlocksCompressor::default() + let compressed = BtrBlocksCompressor::from_session(&session) .compress(&array, &mut session.create_execution_ctx()) .unwrap(); + assert_ne!(compressed.encoding_id(), array.encoding_id()); println!("Compressed size: {}", compressed.len()); println!("Tree view: {}", compressed.display_tree()); }