TurboQuant encoding for Vectors

_typos.toml

@@ -1,5 +1,5 @@
 [default]
-extend-ignore-identifiers-re = ["ffor", "FFOR", "FoR", "typ", "ratatui"]
+extend-ignore-identifiers-re = ["ffor", "FFOR", "FoR", "typ", "ratatui", "wht", "WHT"]
 # We support a few common special comments to tell the checker to ignore sections of code
 extend-ignore-re = [
     "(#|//)\\s*spellchecker:ignore-next-line\\n.*", # Ignore the next line

fuzz/fuzz_targets/file_io.rs

@@ -60,10 +60,9 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus {
     let write_options = match compressor_strategy {
         CompressorStrategy::Default => SESSION.write_options(),
         CompressorStrategy::Compact => {
-            let strategy = WriteStrategyBuilder::default()
-                .with_compact_encodings()
-                .build();
-            SESSION.write_options().with_strategy(strategy)
+            let mut strategy = WriteStrategyBuilder::default();
+            strategy.with_compact_encodings();
+            SESSION.write_options().with_strategy(strategy.build())
         }
     };
 

fuzz/src/array/mod.rs

@@ -541,11 +541,14 @@ pub fn compress_array(array: &ArrayRef, strategy: CompressorStrategy) -> ArrayRe
         CompressorStrategy::Default => BtrBlocksCompressor::default()
             .compress(array)
             .vortex_expect("BtrBlocksCompressor compress should succeed in fuzz test"),
-        CompressorStrategy::Compact => BtrBlocksCompressorBuilder::default()
-            .with_compact()
-            .build()
-            .compress(array)
-            .vortex_expect("Compact compress should succeed in fuzz test"),
+        CompressorStrategy::Compact => {
+            let mut builder = BtrBlocksCompressorBuilder::default();
+            builder.add_compact();
+            builder
+                .build()
+                .compress(array)
+                .vortex_expect("Compact compress should succeed in fuzz test")
+        }
     }
 }
 

vortex-bench/src/lib.rs

@@ -229,11 +229,11 @@ pub enum CompactionStrategy {
 impl CompactionStrategy {
     pub fn apply_options(&self, options: VortexWriteOptions) -> VortexWriteOptions {
         match self {
-            CompactionStrategy::Compact => options.with_strategy(
-                WriteStrategyBuilder::default()
-                    .with_compact_encodings()
-                    .build(),
-            ),
+            CompactionStrategy::Compact => {
+                let mut strategy = WriteStrategyBuilder::default();
+                strategy.with_compact_encodings();
+                options.with_strategy(strategy.build())
+            }
             CompactionStrategy::Default => options,
         }
     }

vortex-btrblocks/public-api.lock

@@ -612,13 +612,15 @@ pub struct vortex_btrblocks::BtrBlocksCompressorBuilder
 
 impl vortex_btrblocks::BtrBlocksCompressorBuilder
 
-pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::build(self) -> vortex_btrblocks::BtrBlocksCompressor
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::add_compact(&mut self) -> &mut Self
+
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::add_cuda_compatible(&mut self) -> &mut Self
 
-pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::exclude(self, ids: impl core::iter::traits::collect::IntoIterator<Item = vortex_compressor::scheme::SchemeId>) -> Self
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::add_scheme(&mut self, scheme: &'static dyn vortex_compressor::scheme::Scheme) -> &mut Self
 
-pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::with_compact(self) -> Self
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::build(self) -> vortex_btrblocks::BtrBlocksCompressor
 
-pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::with_new_scheme(self, scheme: &'static dyn vortex_compressor::scheme::Scheme) -> Self
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::remove(&mut self, ids: impl core::iter::traits::collect::IntoIterator<Item = vortex_compressor::scheme::SchemeId>) -> &mut Self
 
 impl core::clone::Clone for vortex_btrblocks::BtrBlocksCompressorBuilder
 

vortex-btrblocks/src/builder.rs

@@ -67,8 +67,8 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
 ///
 /// By default, all schemes in [`ALL_SCHEMES`] are enabled. 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`](BtrBlocksCompressorBuilder::with_compact).
+/// [`add_scheme`](BtrBlocksCompressorBuilder::add_scheme) or
+/// [`add_compact`](BtrBlocksCompressorBuilder::add_compact).
 ///
 /// # Examples
 ///
@@ -79,10 +79,10 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
 /// // Default compressor with all schemes in ALL_SCHEMES.
 /// let compressor = BtrBlocksCompressorBuilder::default().build();
 ///
-/// // Exclude specific schemes.
-/// let compressor = BtrBlocksCompressorBuilder::default()
-///     .exclude([IntDictScheme.id()])
-///     .build();
+/// // Remove specific schemes.
+/// let mut builder = BtrBlocksCompressorBuilder::default();
+/// builder.remove([IntDictScheme.id()]);
+/// let compressor = builder.build();
 /// ```
 #[derive(Debug, Clone)]
 pub struct BtrBlocksCompressorBuilder {
@@ -100,19 +100,18 @@ impl Default for BtrBlocksCompressorBuilder {
 impl BtrBlocksCompressorBuilder {
     /// 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.
+    /// 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 {
+    pub fn add_scheme(&mut self, scheme: &'static dyn Scheme) -> &mut 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
     }
@@ -127,20 +126,40 @@ impl BtrBlocksCompressorBuilder {
     ///
     /// Panics if any of the compact schemes are already present.
     #[cfg(feature = "zstd")]
-    pub fn with_compact(self) -> Self {
-        // This should be fast since we don't have that many schemes.
-        let builder = self.with_new_scheme(&string::ZstdScheme);
-
+    pub fn add_compact(&mut self) -> &mut Self {
+        self.add_scheme(&string::ZstdScheme);
         #[cfg(feature = "pco")]
-        let builder = builder
-            .with_new_scheme(&integer::PcoScheme)
-            .with_new_scheme(&float::PcoScheme);
+        {
+            self.add_scheme(&integer::PcoScheme);
+            self.add_scheme(&float::PcoScheme);
+        }
+        self
+    }
 
-        builder
+    /// Excludes schemes without CUDA kernel support and adds Zstd for string compression.
+    ///
+    /// With the `unstable_encodings` feature, buffer-level Zstd compression is used which
+    /// preserves the array buffer layout for zero-conversion GPU decompression. Without it,
+    /// interleaved Zstd compression is used.
+    #[cfg(feature = "zstd")]
+    pub fn add_cuda_compatible(&mut self) -> &mut Self {
+        self.remove([
+            integer::SparseScheme.id(),
+            rle::RLE_INTEGER_SCHEME.id(),
+            rle::RLE_FLOAT_SCHEME.id(),
+            float::NullDominatedSparseScheme.id(),
+            string::StringDictScheme.id(),
+            string::FSSTScheme.id(),
+        ]);
+        #[cfg(feature = "unstable_encodings")]
+        self.add_scheme(&string::ZstdBuffersScheme);
+        #[cfg(not(feature = "unstable_encodings"))]
+        self.add_scheme(&string::ZstdScheme);
+        self
     }
 
-    /// Excludes the specified compression schemes by their [`SchemeId`].
-    pub fn exclude(mut self, ids: impl IntoIterator<Item = SchemeId>) -> Self {
+    /// Removes the specified compression schemes by their [`SchemeId`].
+    pub fn remove(&mut self, ids: impl IntoIterator<Item = SchemeId>) -> &mut Self {
         let ids: HashSet<_> = ids.into_iter().collect();
         self.schemes.retain(|s| !ids.contains(&s.id()));
         self

vortex-btrblocks/src/canonical_compressor.rs

@@ -25,10 +25,10 @@ use crate::CascadingCompressor;
 /// // Default compressor - all schemes allowed.
 /// let compressor = BtrBlocksCompressor::default();
 ///
-/// // Exclude specific schemes using the builder.
-/// let compressor = BtrBlocksCompressorBuilder::default()
-///     .exclude([IntDictScheme.id()])
-///     .build();
+/// // Remove specific schemes using the builder.
+/// let mut builder = BtrBlocksCompressorBuilder::default();
+/// builder.remove([IntDictScheme.id()]);
+/// let compressor = builder.build();
 /// ```
 #[derive(Clone)]
 pub struct BtrBlocksCompressor(

vortex-btrblocks/src/lib.rs

@@ -46,10 +46,10 @@
 //! // Default compressor with all schemes enabled.
 //! let compressor = BtrBlocksCompressor::default();
 //!
-//! // Configure with builder to exclude specific schemes.
-//! let compressor = BtrBlocksCompressorBuilder::default()
-//!     .exclude([IntDictScheme.id()])
-//!     .build();
+//! // Remove specific schemes using the builder.
+//! let mut builder = BtrBlocksCompressorBuilder::default();
+//! builder.remove([IntDictScheme.id()]);
+//! let compressor = builder.build();
 //! ```
 //!
 //! [BtrBlocks]: https://www.cs.cit.tum.de/fileadmin/w00cfj/dis/papers/btrblocks.pdf

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

@@ -91,8 +91,9 @@ async fn main() -> VortexResult<()> {
 /// Build the write strategy used for CUDA-compatible file output.
 #[cuda_available]
 fn cuda_write_strategy() -> Arc<dyn vortex::layout::LayoutStrategy> {
-    WriteStrategyBuilder::default()
-        .with_cuda_compatible_encodings()
+    let mut strategy = WriteStrategyBuilder::default();
+    strategy.with_cuda_compatible_encodings();
+    strategy
         .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default()))
         .build()
 }

vortex-file/public-api.lock

@@ -342,15 +342,17 @@ pub struct vortex_file::WriteStrategyBuilder
 
 impl vortex_file::WriteStrategyBuilder
 
+pub fn vortex_file::WriteStrategyBuilder::btrblocks_mut(&mut self) -> core::option::Option<&mut vortex_btrblocks::builder::BtrBlocksCompressorBuilder>
+
 pub fn vortex_file::WriteStrategyBuilder::build(self) -> alloc::sync::Arc<dyn vortex_layout::strategy::LayoutStrategy>
 
 pub fn vortex_file::WriteStrategyBuilder::with_allow_encodings(self, allow_encodings: vortex_array::session::ArrayRegistry) -> Self
 
-pub fn vortex_file::WriteStrategyBuilder::with_compact_encodings(self) -> Self
+pub fn vortex_file::WriteStrategyBuilder::with_compact_encodings(&mut self) -> &mut Self
 
 pub fn vortex_file::WriteStrategyBuilder::with_compressor<C: vortex_layout::layouts::compressed::CompressorPlugin>(self, compressor: C) -> Self
 
-pub fn vortex_file::WriteStrategyBuilder::with_cuda_compatible_encodings(self) -> Self
+pub fn vortex_file::WriteStrategyBuilder::with_cuda_compatible_encodings(&mut self) -> &mut Self
 
 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