TurboQuant encoding for Vectors (#7167)

Lossy quantization for vector data (e.g., embeddings) based on TurboQuant
(https://arxiv.org/abs/2504.19874). Supports both MSE-optimal and
inner-product-optimal (Prod with QJL correction) variants at 1-8 bits per
coordinate.

Key components:
- Single TurboQuant array encoding with optional QJL correction fields,
  storing quantized codes, norms, centroids, and rotation signs as children.
- Structured Random Hadamard Transform (SRHT) for O(d log d) rotation,
  fully self-contained with no external linear algebra library.
- Max-Lloyd centroid computation on Beta(d/2, d/2) distribution.
- Approximate cosine similarity and dot product compute directly on
  quantized arrays without full decompression.
- Pluggable TurboQuantScheme for BtrBlocks, exposed via
  WriteStrategyBuilder::with_vector_quantization().
- Benchmarks covering common embedding dimensions (128, 768, 1024, 1536).

Also refactors CompressingStrategy to a single constructor, and adds
vortex_tensor::initialize() for session registration of tensor types,
encodings, and scalar functions.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Will Manning <will@willmanning.io>
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: 2 people

Cargo.lock

@@ -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

_typos.toml

@@ -618,6 +618,8 @@ pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::exclude(self, ids: impl cor
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::include(self, ids: impl core::iter::traits::collect::IntoIterator<Item = vortex_compressor::scheme::SchemeId>) -> Self
 
+pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::with_scheme(self, scheme: &'static dyn vortex_compressor::scheme::Scheme) -> Self
+
 impl core::clone::Clone for vortex_btrblocks::BtrBlocksCompressorBuilder
 
 pub fn vortex_btrblocks::BtrBlocksCompressorBuilder::clone(&self) -> vortex_btrblocks::BtrBlocksCompressorBuilder

vortex-btrblocks/public-api.lock

@@ -144,6 +144,15 @@ impl BtrBlocksCompressorBuilder {
         self
     }
 
+    /// 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.
+    pub fn with_scheme(mut self, scheme: &'static dyn Scheme) -> Self {
+        self.schemes.insert(scheme);
+        self
+    }
+
     /// Excludes the specified compression schemes by their [`SchemeId`].
     pub fn exclude(mut self, ids: impl IntoIterator<Item = SchemeId>) -> Self {
         let ids: HashSet<_> = ids.into_iter().collect();

vortex-btrblocks/src/builder.rs

@@ -28,9 +28,11 @@ use vortex_array::arrays::VarBinView;
 use vortex_array::dtype::FieldPath;
 use vortex_array::session::ArrayRegistry;
 use vortex_array::session::ArraySession;
+use vortex_btrblocks::BtrBlocksCompressorBuilder;
 use vortex_bytebool::ByteBool;
 use vortex_datetime_parts::DateTimeParts;
 use vortex_decimal_byte_parts::DecimalByteParts;
+use vortex_error::vortex_panic;
 use vortex_fastlanes::BitPacked;
 use vortex_fastlanes::Delta;
 use vortex_fastlanes::FoR;
@@ -53,13 +55,14 @@ use vortex_pco::Pco;
 use vortex_runend::RunEnd;
 use vortex_sequence::Sequence;
 use vortex_sparse::Sparse;
+#[cfg(feature = "unstable_encodings")]
+use vortex_tensor::encodings::turboquant::TurboQuant;
 use vortex_utils::aliases::hash_map::HashMap;
 use vortex_zigzag::ZigZag;
 
 #[rustfmt::skip]
 #[cfg(feature = "zstd")]
 use vortex_btrblocks::{
-    BtrBlocksCompressorBuilder,
     SchemeExt,
     schemes::float,
     schemes::integer,
@@ -111,6 +114,8 @@ pub static ALLOWED_ENCODINGS: LazyLock<ArrayRegistry> = LazyLock::new(|| {
     session.register(RunEnd);
     session.register(Sequence);
     session.register(Sparse);
+    #[cfg(feature = "unstable_encodings")]
+    session.register(TurboQuant);
     session.register(ZigZag);
 
     #[cfg(feature = "zstd")]
@@ -127,23 +132,26 @@ pub static ALLOWED_ENCODINGS: LazyLock<ArrayRegistry> = LazyLock::new(|| {
 /// repartitioning and compressing them to strike a balance between size on-disk,
 /// bulk decoding performance, and IOPS required to perform an indexed read.
 pub struct WriteStrategyBuilder {
-    compressor: Option<Arc<dyn CompressorPlugin>>,
     row_block_size: usize,
     field_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
     allow_encodings: Option<ArrayRegistry>,
     flat_strategy: Option<Arc<dyn LayoutStrategy>>,
+    // builder and compressor are mutually exclusive
+    builder: Option<BtrBlocksCompressorBuilder>,
+    compressor: Option<Arc<dyn CompressorPlugin>>,
 }
 
 impl Default for WriteStrategyBuilder {
     /// Create a new empty builder. It can be further configured,
     /// and then finally built yielding the [`LayoutStrategy`].
     fn default() -> Self {
         Self {
-            compressor: None,
             row_block_size: 8192,
             field_writers: HashMap::new(),
             allow_encodings: Some(ALLOWED_ENCODINGS.clone()),
             flat_strategy: None,
+            builder: None,
+            compressor: None,
         }
     }
 }
@@ -154,6 +162,9 @@ impl WriteStrategyBuilder {
     /// If not provided, this will use a BtrBlocks-style cascading compressor that tries to balance
     /// total size with decoding performance.
     pub fn with_compressor<C: CompressorPlugin>(mut self, compressor: C) -> Self {
+        if self.builder.is_some() {
+            vortex_panic!("Cannot configure both a custom compressor and custom builder schemes");
+        }
         self.compressor = Some(Arc::new(compressor));
         self
     }
@@ -198,7 +209,12 @@ impl WriteStrategyBuilder {
     /// GPU decompression. Without it, strings use interleaved Zstd compression.
     #[cfg(feature = "zstd")]
     pub fn with_cuda_compatible_encodings(mut self) -> Self {
-        let mut builder = BtrBlocksCompressorBuilder::default().exclude([
+        if self.compressor.is_some() {
+            vortex_panic!(
+                "Cannot configure both a custom compressor and CUDA compatible encodings"
+            );
+        }
+        let b = self.builder.take().unwrap_or_default().exclude([
             integer::SparseScheme.id(),
             integer::RLE_INTEGER_SCHEME.id(),
             float::RLE_FLOAT_SCHEME.id(),
@@ -209,14 +225,13 @@ impl WriteStrategyBuilder {
 
         #[cfg(feature = "unstable_encodings")]
         {
-            builder = builder.include([string::ZstdBuffersScheme.id()]);
+            self.builder = Some(b.include([string::ZstdBuffersScheme.id()]));
         }
         #[cfg(not(feature = "unstable_encodings"))]
         {
-            builder = builder.include([string::ZstdScheme.id()]);
+            self.builder = Some(b.include([string::ZstdScheme.id()]));
         }
 
-        self.compressor = Some(Arc::new(builder.build()));
         self
     }
 
@@ -227,21 +242,47 @@ impl WriteStrategyBuilder {
     /// especially for floating-point heavy datasets.
     #[cfg(feature = "zstd")]
     pub fn with_compact_encodings(mut self) -> Self {
-        let btrblocks = BtrBlocksCompressorBuilder::default()
-            .include([
-                string::ZstdScheme.id(),
-                integer::PcoScheme.id(),
-                float::PcoScheme.id(),
-            ])
-            .build();
-
-        self.compressor = Some(Arc::new(btrblocks));
+        if self.compressor.is_some() {
+            vortex_panic!("Cannot configure both a custom compressor and compact encodings");
+        }
+        self.builder = Some(self.builder.take().unwrap_or_default().include([
+            string::ZstdScheme.id(),
+            integer::PcoScheme.id(),
+            float::PcoScheme.id(),
+        ]));
+        self
+    }
+
+    /// Enable TurboQuant lossy vector quantization for tensor columns.
+    ///
+    /// When enabled, `Vector` and `FixedShapeTensor` extension arrays are
+    /// compressed using the TurboQuant algorithm with QJL correction for
+    /// unbiased inner product estimation.
+    ///
+    /// This augments any existing compressor configuration rather than
+    /// replacing it. If no compressor has been set, the default BtrBlocks
+    /// compressor is used with TurboQuant added.
+    #[cfg(feature = "unstable_encodings")]
+    pub fn with_vector_quantization(mut self) -> Self {
+        if self.compressor.is_some() {
+            vortex_panic!("Cannot configure both a custom compressor and vector quantization");
+        }
+        use vortex_tensor::encodings::turboquant::scheme::TURBOQUANT_SCHEME;
+        self.builder = Some(
+            self.builder
+                .take()
+                .unwrap_or_default()
+                .with_scheme(&TURBOQUANT_SCHEME),
+        );
         self
     }
 
     /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides
     /// applied.
     pub fn build(self) -> Arc<dyn LayoutStrategy> {
+        use vortex_btrblocks::SchemeExt as _;
+        use vortex_btrblocks::schemes::integer::IntDictScheme;
+
         let flat: Arc<dyn LayoutStrategy> = if let Some(flat) = self.flat_strategy {
             flat
         } else if let Some(allow_encodings) = self.allow_encodings {
@@ -254,12 +295,24 @@ impl WriteStrategyBuilder {
         let chunked = ChunkedLayoutStrategy::new(flat.clone());
         // 6. buffer chunks so they end up with closer segment ids physically
         let buffered = BufferedStrategy::new(chunked, 2 * ONE_MEG); // 2MB
+
         // 5. compress each chunk
-        let compressing = if let Some(ref compressor) = self.compressor {
-            CompressingStrategy::new_opaque(buffered, compressor.clone())
-        } else {
-            CompressingStrategy::new_btrblocks(buffered, true)
-        };
+        let data_compressor: Arc<dyn CompressorPlugin> =
+            if let Some(ref compressor) = self.compressor {
+                assert!(
+                    self.builder.is_none(),
+                    "Cannot configure both a custom compressor and custom builder schemes"
+                );
+                compressor.clone()
+            } else {
+                Arc::new(
+                    self.builder
+                        .unwrap_or_default()
+                        .exclude([IntDictScheme.id()])
+                        .build(),
+                )
+            };
+        let compressing = CompressingStrategy::new(buffered, data_compressor.clone());
 
         // 4. prior to compression, coalesce up to a minimum size
         let coalescing = RepartitionStrategy::new(
@@ -279,11 +332,12 @@ impl WriteStrategyBuilder {
         );
 
         // 2.1. | 3.1. compress stats tables and dict values.
-        let compress_then_flat = if let Some(ref compressor) = self.compressor {
-            CompressingStrategy::new_opaque(flat, compressor.clone())
+        let stats_compressor = if let Some(compressor) = self.compressor {
+            compressor.clone()
         } else {
-            CompressingStrategy::new_btrblocks(flat, false)
+            Arc::new(BtrBlocksCompressorBuilder::default().build())
         };
+        let compress_then_flat = CompressingStrategy::new(flat, stats_compressor);
 
         // 3. apply dict encoding or fallback
         let dict = DictStrategy::new(

vortex-file/src/strategy.rs

@@ -20,6 +20,7 @@ use vortex_array::field_path;
 use vortex_array::scalar_fn::session::ScalarFnSession;
 use vortex_array::session::ArraySession;
 use vortex_array::validity::Validity;
+use vortex_btrblocks::BtrBlocksCompressor;
 use vortex_buffer::ByteBuffer;
 use vortex_file::OpenOptionsSessionExt;
 use vortex_file::WriteOptionsSessionExt;
@@ -67,9 +68,9 @@ async fn test_file_roundtrip() {
 
     // Create a writer which by default uses the BtrBlocks compressor for a.compressed, but leaves
     // the b and the a.raw columns uncompressed.
-    let default_strategy = Arc::new(CompressingStrategy::new_btrblocks(
+    let default_strategy = Arc::new(CompressingStrategy::new(
         FlatLayoutStrategy::default(),
-        false,
+        BtrBlocksCompressor::default(),
     ));
 
     let writer = Arc::new(

vortex-file/tests/test_write_table.rs

@@ -168,9 +168,7 @@ pub struct vortex_layout::layouts::compressed::CompressingStrategy
 
 impl vortex_layout::layouts::compressed::CompressingStrategy
 
-pub fn vortex_layout::layouts::compressed::CompressingStrategy::new_btrblocks<S: vortex_layout::LayoutStrategy>(child: S, exclude_int_dict_encoding: bool) -> Self
-
-pub fn vortex_layout::layouts::compressed::CompressingStrategy::new_opaque<S: vortex_layout::LayoutStrategy, C: vortex_layout::layouts::compressed::CompressorPlugin>(child: S, compressor: C) -> Self
+pub fn vortex_layout::layouts::compressed::CompressingStrategy::new<S: vortex_layout::LayoutStrategy, C: vortex_layout::layouts::compressed::CompressorPlugin>(child: S, compressor: C) -> Self
 
 pub fn vortex_layout::layouts::compressed::CompressingStrategy::with_concurrency(self, concurrency: usize) -> Self
 

vortex-layout/public-api.lock

@@ -9,9 +9,6 @@ use vortex_array::ArrayContext;
 use vortex_array::ArrayRef;
 use vortex_array::expr::stats::Stat;
 use vortex_btrblocks::BtrBlocksCompressor;
-use vortex_btrblocks::BtrBlocksCompressorBuilder;
-use vortex_btrblocks::SchemeExt;
-use vortex_btrblocks::schemes::integer::IntDictScheme;
 use vortex_error::VortexResult;
 use vortex_io::runtime::Handle;
 
@@ -60,32 +57,11 @@ pub struct CompressingStrategy {
 }
 
 impl CompressingStrategy {
-    /// Create a new writer that uses the BtrBlocks-style cascading compressor to compress chunks.
-    ///
-    /// This provides a good balance between decoding speed and small file size.
-    ///
-    /// Set `exclude_int_dict_encoding` to true to prevent dictionary encoding of integer arrays,
-    /// which is useful when compressing dictionary codes to avoid recursive dictionary encoding.
-    pub fn new_btrblocks<S: LayoutStrategy>(child: S, exclude_int_dict_encoding: bool) -> Self {
-        let compressor = if exclude_int_dict_encoding {
-            BtrBlocksCompressorBuilder::default()
-                .exclude([IntDictScheme.id()])
-                .build()
-        } else {
-            BtrBlocksCompressor::default()
-        };
-        Self::new(child, Arc::new(compressor))
-    }
-
-    /// Create a new compressor from a plugin interface.
-    pub fn new_opaque<S: LayoutStrategy, C: CompressorPlugin>(child: S, compressor: C) -> Self {
-        Self::new(child, Arc::new(compressor))
-    }
-
-    fn new<S: LayoutStrategy>(child: S, compressor: Arc<dyn CompressorPlugin>) -> Self {
+    /// Create a new compressing strategy that wraps a child strategy with a compressor plugin.
+    pub fn new<S: LayoutStrategy, C: CompressorPlugin>(child: S, compressor: C) -> Self {
         Self {
             child: Arc::new(child),
-            compressor,
+            compressor: Arc::new(compressor),
             concurrency: std::thread::available_parallelism()
                 .map(|v| v.get())
                 .unwrap_or(1),

vortex-layout/src/layouts/compressed.rs

@@ -86,12 +86,13 @@ impl TableStrategy {
     /// ```ignore
     /// # use std::sync::Arc;
     /// # use vortex_array::dtype::{field_path, Field, FieldPath};
+    /// # use vortex_btrblocks::BtrBlocksCompressor;
     /// # use vortex_layout::layouts::compressed::CompressingStrategy;
     /// # use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
     /// # use vortex_layout::layouts::table::TableStrategy;
     ///
     /// // A strategy for compressing data using the balanced BtrBlocks compressor.
-    /// let compress = CompressingStrategy::new_btrblocks(FlatLayoutStrategy::default(), true);
+    /// let compress = CompressingStrategy::new(FlatLayoutStrategy::default(), BtrBlocksCompressor::default());
     ///
     /// // Our combined strategy uses no compression for validity buffers, BtrBlocks compression
     /// // for most columns, and stores a nested binary column uncompressed (flat) because it

vortex-layout/src/layouts/table.rs

@@ -291,7 +291,7 @@ impl PyVortexWriteOptions {
     /// ```python
     /// >>> vx.io.VortexWriteOptions.compact().write(sprl, "tiny.vortex")
     /// >>> os.path.getsize('tiny.vortex')
-    /// 55120
+    /// 55460
     /// ```
     ///
     /// Random numbers are not (usually) composed of random bytes!