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

@@ -35,6 +35,7 @@ vortex-pco = { workspace = true, optional = true }
 vortex-runend = { workspace = true }
 vortex-sequence = { workspace = true }
 vortex-sparse = { workspace = true }
+vortex-tensor = { workspace = true, optional = true }
 vortex-utils = { workspace = true }
 vortex-zigzag = { workspace = true }
 vortex-zstd = { workspace = true, optional = true }
@@ -47,7 +48,7 @@ vortex-array = { workspace = true, features = ["_test-harness"] }
 
 [features]
 # This feature enabled unstable encodings for which we don't guarantee stability.
-unstable_encodings = ["vortex-zstd?/unstable_encodings"]
+unstable_encodings = ["dep:vortex-tensor", "vortex-zstd?/unstable_encodings"]
 pco = ["dep:pco", "dep:vortex-pco"]
 zstd = ["dep:vortex-zstd"]
 

vortex-btrblocks/Cargo.toml

@@ -120,7 +120,7 @@ impl BtrBlocksCompressorBuilder {
     /// Adds compact encoding schemes (Zstd for strings, 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,
+    /// heavy datasets. Requires the `zstd` feature. When the `pco` rfeature is also enabled,
     /// Pco schemes for integers and floats are included.
     ///
     /// # Panics
@@ -138,6 +138,23 @@ impl BtrBlocksCompressorBuilder {
         builder
     }
 
+    /// Adds the TurboQuant lossy vector quantization scheme.
+    ///
+    /// When enabled, [`Vector`] extension arrays are compressed using the TurboQuant algorithm with
+    /// QJL correction for unbiased inner product estimation.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the TurboQuant scheme is already present.
+    ///
+    /// [`Vector`]: vortex_tensor::vector::Vector
+    /// [`FixedShapeTensor`]: vortex_tensor::fixed_shape::FixedShapeTensor
+    #[cfg(feature = "unstable_encodings")]
+    pub fn with_turboquant(self) -> Self {
+        use vortex_tensor::encodings::turboquant::scheme::TURBOQUANT_SCHEME;
+        self.with_new_scheme(&TURBOQUANT_SCHEME)
+    }
+
     /// Excludes schemes without CUDA kernel support and adds Zstd for string compression.
     ///
     /// With the `unstable_encodings` feature, buffer-level Zstd compression is used which

vortex-btrblocks/src/builder.rs

@@ -56,6 +56,8 @@ 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;
 #[cfg(feature = "zstd")]
@@ -104,6 +106,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")]

vortex-file/src/strategy.rs

@@ -86,13 +86,14 @@ 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;
     ///
-    /// # use vortex_btrblocks::BtrBlocksCompressor;
     /// // A strategy for compressing data using the balanced BtrBlocks compressor.
-    /// let compress = CompressingStrategy::new(FlatLayoutStrategy::default(), BtrBlocksCompressor::default());
+    /// 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

@@ -19,13 +19,18 @@ workspace = true
 [dependencies]
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
+vortex-compressor = { workspace = true }
 vortex-error = { workspace = true }
+vortex-fastlanes = { workspace = true }
 vortex-session = { workspace = true }
+vortex-utils = { workspace = true }
 
+half = { workspace = true }
 itertools = { workspace = true }
 num-traits = { workspace = true }
 prost = { workspace = true }
+rand = { workspace = true }
 
 [dev-dependencies]
+rand_distr = { workspace = true }
 rstest = { workspace = true }
-vortex-buffer = { workspace = true }