TurboQuant encoding for Vectors (#7269)

Continuation of #7167,
authored by @lwwmanning

## Summary

Lossy quantization for vector data (e.g., embeddings) based on
TurboQuant (https://arxiv.org/abs/2504.19874). Implements the MSE-only
variant (Stage 1 of RFC 0033) at 1-8 bits per coordinate (0 for empty
arrays), defaulting to 8-bit near-lossless compression (but still a
small amount because we use SRHT instead of random orthogonal rotation
matrix, something about not satisfying the Haar assumption?).

Key components:
- TurboQuant array encoding with 4 slots: quantized codes, norms,
centroids, and rotation signs. Note that we should probably abstract the
codes and centroids as a dictionary encoded thing so we don't have to
duplicate pushdown rules, and we might want to make a matrix
multiplication `ScalarFn` expression for the rotations.
- Structured Random Hadamard Transform (SRHT) for `O(d log d)` rotation,
fully self-contained with no external linear algebra library. This is
what claude came up with, but we can see in testing that while this is
practical and more efficient, we lose some of the assumptions that a
Haar-random orthogonal matrix gives us. I think this is something we can
play around with because it's abstracted into a discrete step of the
algorithm.
- Max-Lloyd centroid computation on the Beta distribution for the given
dimension.
- Approximate cosine similarity and dot product computed directly on
quantized arrays without full decompression.
- Pluggable `TurboQuantScheme` for the cascading compressor.
- Minimum dimension of 128 (`TurboQuant::MIN_DIMENSION`) for SRHT
quality guarantees.
- Default 8-bit encoding (MSE ~4e-5, exact 4x compression on `f32`).
- Adds `vortex_tensor::initialize()` for session registration of tensor
types, encodings, and scalar functions.

## API Changes

- Adds `TurboQuant` encoding in `vortex-tensor` with
`turboquant_encode()` and `TurboQuantConfig`, and new types
`TurboQuantData` and `TurboQuantArray`.
- Adds `TurboQuantScheme` for compressor integration.
- Adds `TurboQuant::MIN_DIMENSION` (128) constant.
- Adds `float_from_f32<T: Float + FromPrimitive>` shared helper for
infallible f32-to-float conversion.

## Testing (claude-generated)

- Roundtrip tests across bit widths (1-8) and dimensions (128, 256, 768,
1024).
- MSE quality bounds verified against the theoretical bound from Theorem
1.
- Edge cases: empty arrays, single-row arrays, all-zero vectors,
dimension rejection below 128.
- Float type coverage: f16, f32, f64 input encoding and roundtrip.
- Nullable vector support: validity propagation through encode, decode,
slice, take, L2 norm readthrough.
- Quantized-domain cosine similarity and dot product accuracy tests.
- Serde roundtrip (serialize/deserialize).
- Compute pushdown tests: slice, take, scalar_at.
- Compression ratio estimates for typical embedding dimensions.
- Centroid correctness: count, sorted, symmetric, within bounds,
caching, boundary rejection.
- SRHT rotation: determinism, roundtrip, norm preservation, sign
export/import roundtrip.

---------

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

Author: 3 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

@@ -138,6 +138,22 @@ impl BtrBlocksCompressorBuilder {
         builder
     }
 
+    /// Adds the TurboQuant lossy vector quantization scheme.
+    ///
+    /// When enabled, [`Vector`] extension arrays are compressed using the TurboQuant algorithm
+    /// with MSE-optimal scalar quantization.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the TurboQuant scheme is already present.
+    ///
+    /// [`Vector`]: vortex_tensor::vector::Vector
+    #[cfg(feature = "unstable_encodings")]
+    pub fn with_turboquant(self) -> Self {
+        use vortex_tensor::encodings::turboquant::TurboQuantScheme;
+        self.with_new_scheme(&TurboQuantScheme)
+    }
+
     /// 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

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

vortex-tensor/Cargo.toml

@@ -2,6 +2,166 @@ pub mod vortex_tensor
 
 pub mod vortex_tensor::encodings
 
+pub mod vortex_tensor::encodings::turboquant
+
+pub struct vortex_tensor::encodings::turboquant::TurboQuant
+
+impl vortex_tensor::encodings::turboquant::TurboQuant
+
+pub const vortex_tensor::encodings::turboquant::TurboQuant::ID: vortex_array::array::ArrayId
+
+pub const vortex_tensor::encodings::turboquant::TurboQuant::MIN_DIMENSION: u32
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::try_new_array(dtype: vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<vortex_tensor::encodings::turboquant::TurboQuantArray>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::validate_dtype(dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<&vortex_array::dtype::extension::erased::ExtDTypeRef>
+
+impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuant
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::clone(&self) -> vortex_tensor::encodings::turboquant::TurboQuant
+
+impl core::fmt::Debug for vortex_tensor::encodings::turboquant::TurboQuant
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::array::vtable::VTable for vortex_tensor::encodings::turboquant::TurboQuant
+
+pub type vortex_tensor::encodings::turboquant::TurboQuant::ArrayData = vortex_tensor::encodings::turboquant::TurboQuantData
+
+pub type vortex_tensor::encodings::turboquant::TurboQuant::OperationsVTable = vortex_tensor::encodings::turboquant::TurboQuant
+
+pub type vortex_tensor::encodings::turboquant::TurboQuant::ValidityVTable = vortex_array::array::vtable::validity::ValidityVTableFromChild
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::array_eq(array: &vortex_tensor::encodings::turboquant::TurboQuantData, other: &vortex_tensor::encodings::turboquant::TurboQuantData, precision: vortex_array::hash::Precision) -> bool
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::array_hash<H: core::hash::Hasher>(array: &vortex_tensor::encodings::turboquant::TurboQuantData, state: &mut H, precision: vortex_array::hash::Precision)
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::buffer(_array: vortex_array::array::view::ArrayView<'_, Self>, idx: usize) -> vortex_array::buffer::BufferHandle
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::buffer_name(_array: vortex_array::array::view::ArrayView<'_, Self>, _idx: usize) -> core::option::Option<alloc::string::String>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::ArrayData>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::execute(array: vortex_array::array::typed::Array<Self>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::execute_parent(array: vortex_array::array::view::ArrayView<'_, Self>, parent: &vortex_array::array::erased::ArrayRef, child_idx: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::id(&self) -> vortex_array::array::ArrayId
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::nbuffers(_array: vortex_array::array::view::ArrayView<'_, Self>) -> usize
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::reduce_parent(array: vortex_array::array::view::ArrayView<'_, Self>, parent: &vortex_array::array::erased::ArrayRef, child_idx: usize) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::serialize(array: vortex_array::array::view::ArrayView<'_, Self>) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::slot_name(_array: vortex_array::array::view::ArrayView<'_, Self>, idx: usize) -> alloc::string::String
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::slots(array: vortex_array::array::view::ArrayView<'_, Self>) -> &[core::option::Option<vortex_array::array::erased::ArrayRef>]
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::validate(&self, data: &Self::ArrayData, dtype: &vortex_array::dtype::DType, len: usize) -> vortex_error::VortexResult<()>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::with_slots(array: &mut vortex_tensor::encodings::turboquant::TurboQuantData, slots: alloc::vec::Vec<core::option::Option<vortex_array::array::erased::ArrayRef>>) -> vortex_error::VortexResult<()>
+
+impl vortex_array::array::vtable::operations::OperationsVTable<vortex_tensor::encodings::turboquant::TurboQuant> for vortex_tensor::encodings::turboquant::TurboQuant
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::scalar_at(array: vortex_array::array::view::ArrayView<'_, vortex_tensor::encodings::turboquant::TurboQuant>, index: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+impl vortex_array::array::vtable::validity::ValidityChild<vortex_tensor::encodings::turboquant::TurboQuant> for vortex_tensor::encodings::turboquant::TurboQuant
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::validity_child(array: &vortex_tensor::encodings::turboquant::TurboQuantData) -> &vortex_array::array::erased::ArrayRef
+
+impl vortex_array::arrays::dict::take::TakeExecute for vortex_tensor::encodings::turboquant::TurboQuant
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::take(array: vortex_array::array::view::ArrayView<'_, vortex_tensor::encodings::turboquant::TurboQuant>, indices: &vortex_array::array::erased::ArrayRef, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
+
+impl vortex_array::arrays::slice::SliceReduce for vortex_tensor::encodings::turboquant::TurboQuant
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuant::slice(array: vortex_array::array::view::ArrayView<'_, vortex_tensor::encodings::turboquant::TurboQuant>, range: core::ops::range::Range<usize>) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
+
+pub struct vortex_tensor::encodings::turboquant::TurboQuantConfig
+
+pub vortex_tensor::encodings::turboquant::TurboQuantConfig::bit_width: u8
+
+pub vortex_tensor::encodings::turboquant::TurboQuantConfig::seed: core::option::Option<u64>
+
+impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuantConfig
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantConfig::clone(&self) -> vortex_tensor::encodings::turboquant::TurboQuantConfig
+
+impl core::default::Default for vortex_tensor::encodings::turboquant::TurboQuantConfig
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantConfig::default() -> Self
+
+impl core::fmt::Debug for vortex_tensor::encodings::turboquant::TurboQuantConfig
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantConfig::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub struct vortex_tensor::encodings::turboquant::TurboQuantData
+
+impl vortex_tensor::encodings::turboquant::TurboQuantData
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::bit_width(&self) -> u8
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::centroids(&self) -> &vortex_array::array::erased::ArrayRef
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::codes(&self) -> &vortex_array::array::erased::ArrayRef
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::dimension(&self) -> u32
+
+pub unsafe fn vortex_tensor::encodings::turboquant::TurboQuantData::new_unchecked(dtype: &vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> Self
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::norms(&self) -> &vortex_array::array::erased::ArrayRef
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::padded_dim(&self) -> u32
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::rotation_signs(&self) -> &vortex_array::array::erased::ArrayRef
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::try_new(dtype: &vortex_array::dtype::DType, codes: vortex_array::array::erased::ArrayRef, norms: vortex_array::array::erased::ArrayRef, centroids: vortex_array::array::erased::ArrayRef, rotation_signs: vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<Self>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::validate(dtype: &vortex_array::dtype::DType, codes: &vortex_array::array::erased::ArrayRef, norms: &vortex_array::array::erased::ArrayRef, centroids: &vortex_array::array::erased::ArrayRef, rotation_signs: &vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<()>
+
+impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuantData
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::clone(&self) -> vortex_tensor::encodings::turboquant::TurboQuantData
+
+impl core::fmt::Debug for vortex_tensor::encodings::turboquant::TurboQuantData
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantData::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub struct vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+impl core::clone::Clone for vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::clone(&self) -> vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+impl core::cmp::Eq for vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+impl core::cmp::PartialEq for vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::eq(&self, other: &vortex_tensor::encodings::turboquant::TurboQuantScheme) -> bool
+
+impl core::fmt::Debug for vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::marker::Copy for vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+impl core::marker::StructuralPartialEq for vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+impl vortex_compressor::scheme::Scheme for vortex_tensor::encodings::turboquant::TurboQuantScheme
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::compress(&self, compressor: &vortex_compressor::compressor::CascadingCompressor, data: &mut vortex_compressor::stats::cache::ArrayAndStats, _ctx: vortex_compressor::ctx::CompressorContext) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::expected_compression_ratio(&self, _compressor: &vortex_compressor::compressor::CascadingCompressor, data: &mut vortex_compressor::stats::cache::ArrayAndStats, _ctx: vortex_compressor::ctx::CompressorContext) -> vortex_error::VortexResult<f64>
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::matches(&self, canonical: &vortex_array::canonical::Canonical) -> bool
+
+pub fn vortex_tensor::encodings::turboquant::TurboQuantScheme::scheme_name(&self) -> &'static str
+
+pub fn vortex_tensor::encodings::turboquant::turboquant_encode(ext: &vortex_array::arrays::extension::vtable::ExtensionArray, config: &vortex_tensor::encodings::turboquant::TurboQuantConfig, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+
+pub type vortex_tensor::encodings::turboquant::TurboQuantArray = vortex_array::array::typed::Array<vortex_tensor::encodings::turboquant::TurboQuant>
+
 pub mod vortex_tensor::fixed_shape
 
 pub struct vortex_tensor::fixed_shape::FixedShapeTensor
@@ -180,7 +340,7 @@ pub fn vortex_tensor::scalar_fns::inner_product::InnerProduct::arity(&self, _opt
 
 pub fn vortex_tensor::scalar_fns::inner_product::InnerProduct::child_name(&self, _options: &Self::Options, child_idx: usize) -> vortex_array::scalar_fn::vtable::ChildName
 
-pub fn vortex_tensor::scalar_fns::inner_product::InnerProduct::execute(&self, _options: &Self::Options, args: &dyn vortex_array::scalar_fn::vtable::ExecutionArgs, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+pub fn vortex_tensor::scalar_fns::inner_product::InnerProduct::execute(&self, options: &Self::Options, args: &dyn vortex_array::scalar_fn::vtable::ExecutionArgs, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
 
 pub fn vortex_tensor::scalar_fns::inner_product::InnerProduct::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::expression::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
 
@@ -236,6 +396,12 @@ pub vortex_tensor::scalar_fns::ApproxOptions::Approximate
 
 pub vortex_tensor::scalar_fns::ApproxOptions::Exact
 
+impl vortex_tensor::scalar_fns::ApproxOptions
+
+pub fn vortex_tensor::scalar_fns::ApproxOptions::is_approx(&self) -> bool
+
+pub fn vortex_tensor::scalar_fns::ApproxOptions::is_exact(&self) -> bool
+
 impl core::clone::Clone for vortex_tensor::scalar_fns::ApproxOptions
 
 pub fn vortex_tensor::scalar_fns::ApproxOptions::clone(&self) -> vortex_tensor::scalar_fns::ApproxOptions

vortex-tensor/public-api.lock

@@ -7,5 +7,4 @@
 // pub mod norm; // Unit-normalized vectors.
 // pub mod spherical; // Spherical transform on unit-normalized vectors.
 
-// TODO(will):
-// pub mod turboquant;
+pub mod turboquant;