Skip to content

Commit 0ec1d1c

Browse files
committed
add CompressorConfig
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 5d0d627 commit 0ec1d1c

File tree

16 files changed

+150
-180
lines changed

16 files changed

+150
-180
lines changed

fuzz/fuzz_targets/file_io.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,9 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus {
6060
let write_options = match compressor_strategy {
6161
CompressorStrategy::Default => SESSION.write_options(),
6262
CompressorStrategy::Compact => {
63-
let strategy = WriteStrategyBuilder::default()
64-
.with_compact_encodings()
65-
.build();
66-
SESSION.write_options().with_strategy(strategy)
63+
let mut strategy = WriteStrategyBuilder::default();
64+
strategy.with_compact_encodings();
65+
SESSION.write_options().with_strategy(strategy.build())
6766
}
6867
};
6968

fuzz/src/array/mod.rs

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -541,11 +541,14 @@ pub fn compress_array(array: &ArrayRef, strategy: CompressorStrategy) -> ArrayRe
541541
CompressorStrategy::Default => BtrBlocksCompressor::default()
542542
.compress(array)
543543
.vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test"),
544-
CompressorStrategy::Compact => BtrBlocksCompressorBuilder::default()
545-
.with_compact()
546-
.build()
547-
.compress(array)
548-
.vortex_expect("Compact compress should succeed in fuzz test"),
544+
CompressorStrategy::Compact => {
545+
let mut builder = BtrBlocksCompressorBuilder::default();
546+
builder.add_compact();
547+
builder
548+
.build()
549+
.compress(array)
550+
.vortex_expect("Compact compress should succeed in fuzz test")
551+
}
549552
}
550553
}
551554

vortex-bench/src/lib.rs

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -229,11 +229,11 @@ pub enum CompactionStrategy {
229229
impl CompactionStrategy {
230230
pub fn apply_options(&self, options: VortexWriteOptions) -> VortexWriteOptions {
231231
match self {
232-
CompactionStrategy::Compact => options.with_strategy(
233-
WriteStrategyBuilder::default()
234-
.with_compact_encodings()
235-
.build(),
236-
),
232+
CompactionStrategy::Compact => {
233+
let mut strategy = WriteStrategyBuilder::default();
234+
strategy.with_compact_encodings();
235+
options.with_strategy(strategy.build())
236+
}
237237
CompactionStrategy::Default => options,
238238
}
239239
}

vortex-btrblocks/public-api.lock

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -612,13 +612,15 @@ pub struct vortex_btrblocks::BtrBlocksCompressorBuilder
612612

613613
impl vortex_btrblocks::BtrBlocksCompressorBuilder
614614

615-
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::build(self) -> vortex_btrblocks::BtrBlocksCompressor
615+
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::add_compact(&mut self) -> &mut Self
616+
617+
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::add_cuda_compatible(&mut self) -> &mut Self
616618

617-
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::exclude(self, ids: impl core::iter::traits::collect::IntoIterator<Item = vortex_compressor::scheme::SchemeId>) -> Self
619+
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::add_scheme(&mut self, scheme: &'static dyn vortex_compressor::scheme::Scheme) -> &mut Self
618620

619-
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::with_compact(self) -> Self
621+
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::build(self) -> vortex_btrblocks::BtrBlocksCompressor
620622

621-
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::with_new_scheme(self, scheme: &'static dyn vortex_compressor::scheme::Scheme) -> Self
623+
pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::remove(&mut self, ids: impl core::iter::traits::collect::IntoIterator<Item = vortex_compressor::scheme::SchemeId>) -> &mut Self
622624

623625
impl core::clone::Clone for vortex_btrblocks::BtrBlocksCompressorBuilder
624626

vortex-btrblocks/src/builder.rs

Lines changed: 39 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -67,8 +67,8 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
6767
///
6868
/// By default, all schemes in [`ALL_SCHEMES`] are enabled. Feature-gated schemes (Pco, Zstd)
6969
/// are not in `ALL_SCHEMES` and must be added explicitly via
70-
/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or
71-
/// [`with_compact`](BtrBlocksCompressorBuilder::with_compact).
70+
/// [`add_scheme`](BtrBlocksCompressorBuilder::add_scheme) or
71+
/// [`add_compact`](BtrBlocksCompressorBuilder::add_compact).
7272
///
7373
/// # Examples
7474
///
@@ -79,10 +79,10 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
7979
/// // Default compressor with all schemes in ALL_SCHEMES.
8080
/// let compressor = BtrBlocksCompressorBuilder::default().build();
8181
///
82-
/// // Exclude specific schemes.
83-
/// let compressor = BtrBlocksCompressorBuilder::default()
84-
/// .exclude([IntDictScheme.id()])
85-
/// .build();
82+
/// // Remove specific schemes.
83+
/// let mut builder = BtrBlocksCompressorBuilder::default();
84+
/// builder.remove([IntDictScheme.id()]);
85+
/// let compressor = builder.build();
8686
/// ```
8787
#[derive(Debug, Clone)]
8888
pub struct BtrBlocksCompressorBuilder {
@@ -100,19 +100,18 @@ impl Default for BtrBlocksCompressorBuilder {
100100
impl BtrBlocksCompressorBuilder {
101101
/// Adds an external compression scheme not in [`ALL_SCHEMES`].
102102
///
103-
/// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes with
104-
/// the compressor.
103+
/// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes
104+
/// with the compressor.
105105
///
106106
/// # Panics
107107
///
108108
/// Panics if a scheme with the same [`SchemeId`] is already present.
109-
pub fn with_new_scheme(mut self, scheme: &'static dyn Scheme) -> Self {
109+
pub fn add_scheme(&mut self, scheme: &'static dyn Scheme) -> &mut Self {
110110
assert!(
111111
!self.schemes.iter().any(|s| s.id() == scheme.id()),
112112
"scheme {:?} is already present in the builder",
113113
scheme.id(),
114114
);
115-
116115
self.schemes.push(scheme);
117116
self
118117
}
@@ -127,20 +126,40 @@ impl BtrBlocksCompressorBuilder {
127126
///
128127
/// Panics if any of the compact schemes are already present.
129128
#[cfg(feature = "zstd")]
130-
pub fn with_compact(self) -> Self {
131-
// This should be fast since we don't have that many schemes.
132-
let builder = self.with_new_scheme(&string::ZstdScheme);
133-
129+
pub fn add_compact(&mut self) -> &mut Self {
130+
self.add_scheme(&string::ZstdScheme);
134131
#[cfg(feature = "pco")]
135-
let builder = builder
136-
.with_new_scheme(&integer::PcoScheme)
137-
.with_new_scheme(&float::PcoScheme);
132+
{
133+
self.add_scheme(&integer::PcoScheme);
134+
self.add_scheme(&float::PcoScheme);
135+
}
136+
self
137+
}
138138

139-
builder
139+
/// Excludes schemes without CUDA kernel support and adds Zstd for string compression.
140+
///
141+
/// With the `unstable_encodings` feature, buffer-level Zstd compression is used which
142+
/// preserves the array buffer layout for zero-conversion GPU decompression. Without it,
143+
/// interleaved Zstd compression is used.
144+
#[cfg(feature = "zstd")]
145+
pub fn add_cuda_compatible(&mut self) -> &mut Self {
146+
self.remove([
147+
integer::SparseScheme.id(),
148+
rle::RLE_INTEGER_SCHEME.id(),
149+
rle::RLE_FLOAT_SCHEME.id(),
150+
float::NullDominatedSparseScheme.id(),
151+
string::StringDictScheme.id(),
152+
string::FSSTScheme.id(),
153+
]);
154+
#[cfg(feature = "unstable_encodings")]
155+
self.add_scheme(&string::ZstdBuffersScheme);
156+
#[cfg(not(feature = "unstable_encodings"))]
157+
self.add_scheme(&string::ZstdScheme);
158+
self
140159
}
141160

142-
/// Excludes the specified compression schemes by their [`SchemeId`].
143-
pub fn exclude(mut self, ids: impl IntoIterator<Item = SchemeId>) -> Self {
161+
/// Removes the specified compression schemes by their [`SchemeId`].
162+
pub fn remove(&mut self, ids: impl IntoIterator<Item = SchemeId>) -> &mut Self {
144163
let ids: HashSet<_> = ids.into_iter().collect();
145164
self.schemes.retain(|s| !ids.contains(&s.id()));
146165
self

vortex-btrblocks/src/canonical_compressor.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,10 @@ use crate::CascadingCompressor;
2525
/// // Default compressor - all schemes allowed.
2626
/// let compressor = BtrBlocksCompressor::default();
2727
///
28-
/// // Exclude specific schemes using the builder.
29-
/// let compressor = BtrBlocksCompressorBuilder::default()
30-
/// .exclude([IntDictScheme.id()])
31-
/// .build();
28+
/// // Remove specific schemes using the builder.
29+
/// let mut builder = BtrBlocksCompressorBuilder::default();
30+
/// builder.remove([IntDictScheme.id()]);
31+
/// let compressor = builder.build();
3232
/// ```
3333
#[derive(Clone)]
3434
pub struct BtrBlocksCompressor(

vortex-btrblocks/src/lib.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,10 @@
4646
//! // Default compressor with all schemes enabled.
4747
//! let compressor = BtrBlocksCompressor::default();
4848
//!
49-
//! // Configure with builder to exclude specific schemes.
50-
//! let compressor = BtrBlocksCompressorBuilder::default()
51-
//! .exclude([IntDictScheme.id()])
52-
//! .build();
49+
//! // Remove specific schemes using the builder.
50+
//! let mut builder = BtrBlocksCompressorBuilder::default();
51+
//! builder.remove([IntDictScheme.id()]);
52+
//! let compressor = builder.build();
5353
//! ```
5454
//!
5555
//! [BtrBlocks]: https://www.cs.cit.tum.de/fileadmin/w00cfj/dis/papers/btrblocks.pdf

vortex-cuda/gpu-scan-cli/src/main.rs

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,8 +91,9 @@ async fn main() -> VortexResult<()> {
9191
/// Build the write strategy used for CUDA-compatible file output.
9292
#[cuda_available]
9393
fn cuda_write_strategy() -> Arc<dyn vortex::layout::LayoutStrategy> {
94-
WriteStrategyBuilder::default()
95-
.with_cuda_compatible_encodings()
94+
let mut strategy = WriteStrategyBuilder::default();
95+
strategy.with_cuda_compatible_encodings();
96+
strategy
9697
.with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default()))
9798
.build()
9899
}

vortex-file/public-api.lock

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -342,15 +342,17 @@ pub struct vortex_file::WriteStrategyBuilder
342342

343343
impl vortex_file::WriteStrategyBuilder
344344

345+
pub fn vortex_file::WriteStrategyBuilder::btrblocks_mut(&mut self) -> core::option::Option<&mut vortex_btrblocks::builder::BtrBlocksCompressorBuilder>
346+
345347
pub fn vortex_file::WriteStrategyBuilder::build(self) -> alloc::sync::Arc<dyn vortex_layout::strategy::LayoutStrategy>
346348

347349
pub fn vortex_file::WriteStrategyBuilder::with_allow_encodings(self, allow_encodings: vortex_array::session::ArrayRegistry) -> Self
348350

349-
pub fn vortex_file::WriteStrategyBuilder::with_compact_encodings(self) -> Self
351+
pub fn vortex_file::WriteStrategyBuilder::with_compact_encodings(&mut self) -> &mut Self
350352

351353
pub fn vortex_file::WriteStrategyBuilder::with_compressor<C: vortex_layout::layouts::compressed::CompressorPlugin>(self, compressor: C) -> Self
352354

353-
pub fn vortex_file::WriteStrategyBuilder::with_cuda_compatible_encodings(self) -> Self
355+
pub fn vortex_file::WriteStrategyBuilder::with_cuda_compatible_encodings(&mut self) -> &mut Self
354356

355357
pub fn vortex_file::WriteStrategyBuilder::with_field_writer(self, field: impl core::convert::Into<vortex_array::dtype::field::FieldPath>, writer: alloc::sync::Arc<dyn vortex_layout::strategy::LayoutStrategy>) -> Self
356358

0 commit comments

Comments
 (0)