feat: turboquant encoding for vectors

Cargo.toml

@@ -51,6 +51,7 @@ members = [
     "encodings/zstd",
     "encodings/bytebool",
     "encodings/parquet-variant",
+    "encodings/turboquant",
     # Benchmarks
     "benchmarks/lance-bench",
     "benchmarks/compress-bench",
@@ -285,6 +286,7 @@ vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-feat
 vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false }
 vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false }
 vortex-tensor = { version = "0.1.0", path = "./vortex-tensor", default-features = false }
+vortex-turboquant = { version = "0.1.0", path = "./encodings/turboquant", default-features = false }
 vortex-utils = { version = "0.1.0", path = "./vortex-utils", default-features = false }
 vortex-zigzag = { version = "0.1.0", path = "./encodings/zigzag", default-features = false }
 vortex-zstd = { version = "0.1.0", path = "./encodings/zstd", default-features = false }

encodings/turboquant/Cargo.toml

@@ -0,0 +1,33 @@
+[package]
+name = "vortex-turboquant"
+authors = { workspace = true }
+categories = { workspace = true }
+description = "Vortex TurboQuant vector quantization encoding"
+edition = { workspace = true }
+homepage = { workspace = true }
+include = { workspace = true }
+keywords = { workspace = true }
+license = { workspace = true }
+readme = { workspace = true }
+repository = { workspace = true }
+rust-version = { workspace = true }
+version = { workspace = true }
+
+[lints]
+workspace = true
+
+[dependencies]
+half = { workspace = true }
+prost = { workspace = true }
+rand = { workspace = true }
+vortex-array = { workspace = true }
+vortex-buffer = { workspace = true }
+vortex-error = { workspace = true }
+vortex-fastlanes = { workspace = true }
+vortex-session = { workspace = true }
+vortex-utils = { workspace = true }
+
+[dev-dependencies]
+rand_distr = { workspace = true }
+rstest = { workspace = true }
+vortex-array = { workspace = true, features = ["_test-harness"] }

encodings/turboquant/public-api.lock

@@ -0,0 +1,185 @@
+pub mod vortex_turboquant
+
+pub struct vortex_turboquant::QjlCorrection
+
+impl vortex_turboquant::QjlCorrection
+
+pub fn vortex_turboquant::QjlCorrection::residual_norms(&self) -> &vortex_array::array::ArrayRef
+
+pub fn vortex_turboquant::QjlCorrection::rotation_signs(&self) -> &vortex_array::array::ArrayRef
+
+pub fn vortex_turboquant::QjlCorrection::signs(&self) -> &vortex_array::array::ArrayRef
+
+impl core::clone::Clone for vortex_turboquant::QjlCorrection
+
+pub fn vortex_turboquant::QjlCorrection::clone(&self) -> vortex_turboquant::QjlCorrection
+
+impl core::fmt::Debug for vortex_turboquant::QjlCorrection
+
+pub fn vortex_turboquant::QjlCorrection::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub struct vortex_turboquant::TurboQuant
+
+impl vortex_turboquant::TurboQuant
+
+pub const vortex_turboquant::TurboQuant::ID: vortex_array::vtable::dyn_::ArrayId
+
+impl core::clone::Clone for vortex_turboquant::TurboQuant
+
+pub fn vortex_turboquant::TurboQuant::clone(&self) -> vortex_turboquant::TurboQuant
+
+impl core::fmt::Debug for vortex_turboquant::TurboQuant
+
+pub fn vortex_turboquant::TurboQuant::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl vortex_array::arrays::dict::take::TakeExecute for vortex_turboquant::TurboQuant
+
+pub fn vortex_turboquant::TurboQuant::take(array: &vortex_turboquant::TurboQuantArray, indices: &vortex_array::array::ArrayRef, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
+impl vortex_array::arrays::slice::SliceReduce for vortex_turboquant::TurboQuant
+
+pub fn vortex_turboquant::TurboQuant::slice(array: &vortex_turboquant::TurboQuantArray, range: core::ops::range::Range<usize>) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
+impl vortex_array::vtable::VTable for vortex_turboquant::TurboQuant
+
+pub type vortex_turboquant::TurboQuant::Array = vortex_turboquant::TurboQuantArray
+
+pub type vortex_turboquant::TurboQuant::Metadata = vortex_array::metadata::ProstMetadata<vortex_turboquant::array::TurboQuantMetadata>
+
+pub type vortex_turboquant::TurboQuant::OperationsVTable = vortex_turboquant::TurboQuant
+
+pub type vortex_turboquant::TurboQuant::ValidityVTable = vortex_array::vtable::validity::ValidityVTableFromChild
+
+pub fn vortex_turboquant::TurboQuant::array_eq(array: &vortex_turboquant::TurboQuantArray, other: &vortex_turboquant::TurboQuantArray, precision: vortex_array::hash::Precision) -> bool
+
+pub fn vortex_turboquant::TurboQuant::array_hash<H: core::hash::Hasher>(array: &vortex_turboquant::TurboQuantArray, state: &mut H, precision: vortex_array::hash::Precision)
+
+pub fn vortex_turboquant::TurboQuant::buffer(_array: &vortex_turboquant::TurboQuantArray, idx: usize) -> vortex_array::buffer::BufferHandle
+
+pub fn vortex_turboquant::TurboQuant::buffer_name(_array: &vortex_turboquant::TurboQuantArray, _idx: usize) -> core::option::Option<alloc::string::String>
+
+pub fn vortex_turboquant::TurboQuant::build(dtype: &vortex_array::dtype::DType, len: usize, metadata: &Self::Metadata, _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren) -> vortex_error::VortexResult<vortex_turboquant::TurboQuantArray>
+
+pub fn vortex_turboquant::TurboQuant::child(array: &vortex_turboquant::TurboQuantArray, idx: usize) -> vortex_array::array::ArrayRef
+
+pub fn vortex_turboquant::TurboQuant::child_name(_array: &vortex_turboquant::TurboQuantArray, idx: usize) -> alloc::string::String
+
+pub fn vortex_turboquant::TurboQuant::deserialize(bytes: &[u8], _dtype: &vortex_array::dtype::DType, _len: usize, _buffers: &[vortex_array::buffer::BufferHandle], _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<Self::Metadata>
+
+pub fn vortex_turboquant::TurboQuant::dtype(array: &vortex_turboquant::TurboQuantArray) -> &vortex_array::dtype::DType
+
+pub fn vortex_turboquant::TurboQuant::execute(array: alloc::sync::Arc<vortex_array::vtable::typed::Array<Self>>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
+
+pub fn vortex_turboquant::TurboQuant::execute_parent(array: &vortex_array::vtable::typed::Array<Self>, parent: &vortex_array::array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
+pub fn vortex_turboquant::TurboQuant::id(&self) -> vortex_array::vtable::dyn_::ArrayId
+
+pub fn vortex_turboquant::TurboQuant::len(array: &vortex_turboquant::TurboQuantArray) -> usize
+
+pub fn vortex_turboquant::TurboQuant::metadata(array: &vortex_turboquant::TurboQuantArray) -> vortex_error::VortexResult<Self::Metadata>
+
+pub fn vortex_turboquant::TurboQuant::nbuffers(_array: &vortex_turboquant::TurboQuantArray) -> usize
+
+pub fn vortex_turboquant::TurboQuant::nchildren(array: &vortex_turboquant::TurboQuantArray) -> usize
+
+pub fn vortex_turboquant::TurboQuant::reduce_parent(array: &vortex_array::vtable::typed::Array<Self>, parent: &vortex_array::array::ArrayRef, child_idx: usize) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
+pub fn vortex_turboquant::TurboQuant::serialize(metadata: Self::Metadata) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
+
+pub fn vortex_turboquant::TurboQuant::stats(array: &vortex_turboquant::TurboQuantArray) -> vortex_array::stats::array::StatsSetRef<'_>
+
+pub fn vortex_turboquant::TurboQuant::vtable(_array: &Self::Array) -> &Self
+
+pub fn vortex_turboquant::TurboQuant::with_children(array: &mut Self::Array, children: alloc::vec::Vec<vortex_array::array::ArrayRef>) -> vortex_error::VortexResult<()>
+
+impl vortex_array::vtable::operations::OperationsVTable<vortex_turboquant::TurboQuant> for vortex_turboquant::TurboQuant
+
+pub fn vortex_turboquant::TurboQuant::scalar_at(array: &vortex_turboquant::TurboQuantArray, index: usize, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
+
+impl vortex_array::vtable::validity::ValidityChild<vortex_turboquant::TurboQuant> for vortex_turboquant::TurboQuant
+
+pub fn vortex_turboquant::TurboQuant::validity_child(array: &vortex_turboquant::TurboQuantArray) -> &vortex_array::array::ArrayRef
+
+pub struct vortex_turboquant::TurboQuantArray
+
+impl vortex_turboquant::TurboQuantArray
+
+pub fn vortex_turboquant::TurboQuantArray::bit_width(&self) -> u8
+
+pub fn vortex_turboquant::TurboQuantArray::centroids(&self) -> &vortex_array::array::ArrayRef
+
+pub fn vortex_turboquant::TurboQuantArray::codes(&self) -> &vortex_array::array::ArrayRef
+
+pub fn vortex_turboquant::TurboQuantArray::dimension(&self) -> u32
+
+pub fn vortex_turboquant::TurboQuantArray::has_qjl(&self) -> bool
+
+pub fn vortex_turboquant::TurboQuantArray::norms(&self) -> &vortex_array::array::ArrayRef
+
+pub fn vortex_turboquant::TurboQuantArray::padded_dim(&self) -> u32
+
+pub fn vortex_turboquant::TurboQuantArray::qjl(&self) -> core::option::Option<&vortex_turboquant::QjlCorrection>
+
+pub fn vortex_turboquant::TurboQuantArray::rotation_signs(&self) -> &vortex_array::array::ArrayRef
+
+pub fn vortex_turboquant::TurboQuantArray::try_new_mse(dtype: vortex_array::dtype::DType, codes: vortex_array::array::ArrayRef, norms: vortex_array::array::ArrayRef, centroids: vortex_array::array::ArrayRef, rotation_signs: vortex_array::array::ArrayRef, dimension: u32, bit_width: u8) -> vortex_error::VortexResult<Self>
+
+pub fn vortex_turboquant::TurboQuantArray::try_new_qjl(dtype: vortex_array::dtype::DType, codes: vortex_array::array::ArrayRef, norms: vortex_array::array::ArrayRef, centroids: vortex_array::array::ArrayRef, rotation_signs: vortex_array::array::ArrayRef, qjl: vortex_turboquant::QjlCorrection, dimension: u32, bit_width: u8) -> vortex_error::VortexResult<Self>
+
+impl vortex_turboquant::TurboQuantArray
+
+pub fn vortex_turboquant::TurboQuantArray::to_array(&self) -> vortex_array::array::ArrayRef
+
+impl core::clone::Clone for vortex_turboquant::TurboQuantArray
+
+pub fn vortex_turboquant::TurboQuantArray::clone(&self) -> vortex_turboquant::TurboQuantArray
+
+impl core::convert::AsRef<dyn vortex_array::array::DynArray> for vortex_turboquant::TurboQuantArray
+
+pub fn vortex_turboquant::TurboQuantArray::as_ref(&self) -> &dyn vortex_array::array::DynArray
+
+impl core::convert::From<vortex_turboquant::TurboQuantArray> for vortex_array::array::ArrayRef
+
+pub fn vortex_array::array::ArrayRef::from(value: vortex_turboquant::TurboQuantArray) -> vortex_array::array::ArrayRef
+
+impl core::fmt::Debug for vortex_turboquant::TurboQuantArray
+
+pub fn vortex_turboquant::TurboQuantArray::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::ops::deref::Deref for vortex_turboquant::TurboQuantArray
+
+pub type vortex_turboquant::TurboQuantArray::Target = dyn vortex_array::array::DynArray
+
+pub fn vortex_turboquant::TurboQuantArray::deref(&self) -> &Self::Target
+
+impl vortex_array::array::IntoArray for vortex_turboquant::TurboQuantArray
+
+pub fn vortex_turboquant::TurboQuantArray::into_array(self) -> vortex_array::array::ArrayRef
+
+pub struct vortex_turboquant::TurboQuantConfig
+
+pub vortex_turboquant::TurboQuantConfig::bit_width: u8
+
+pub vortex_turboquant::TurboQuantConfig::seed: core::option::Option<u64>
+
+impl core::clone::Clone for vortex_turboquant::TurboQuantConfig
+
+pub fn vortex_turboquant::TurboQuantConfig::clone(&self) -> vortex_turboquant::TurboQuantConfig
+
+impl core::default::Default for vortex_turboquant::TurboQuantConfig
+
+pub fn vortex_turboquant::TurboQuantConfig::default() -> Self
+
+impl core::fmt::Debug for vortex_turboquant::TurboQuantConfig
+
+pub fn vortex_turboquant::TurboQuantConfig::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub const vortex_turboquant::FIXED_SHAPE_TENSOR_EXT_ID: &str
+
+pub const vortex_turboquant::VECTOR_EXT_ID: &str
+
+pub fn vortex_turboquant::initialize(session: &mut vortex_session::VortexSession)
+
+pub fn vortex_turboquant::turboquant_encode_mse(fsl: &vortex_array::arrays::fixed_size_list::array::FixedSizeListArray, config: &vortex_turboquant::TurboQuantConfig) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
+
+pub fn vortex_turboquant::turboquant_encode_qjl(fsl: &vortex_array::arrays::fixed_size_list::array::FixedSizeListArray, config: &vortex_turboquant::TurboQuantConfig) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>