Skip to content

Commit 6111f27

Browse files
connortsui20connortsui20
committed
restructure modules for turboquant
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 8f41d20 commit 6111f27

17 files changed

Lines changed: 159 additions & 174 deletions

File tree

‎vortex-btrblocks/src/builder.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -150,8 +150,8 @@ impl BtrBlocksCompressorBuilder {
150150
/// [`Vector`]: vortex_tensor::vector::Vector
151151
#[cfg(feature = "unstable_encodings")]
152152
pub fn with_turboquant(self) -> Self {
153-
use vortex_tensor::encodings::turboquant::scheme::TURBOQUANT_SCHEME;
154-
self.with_new_scheme(&TURBOQUANT_SCHEME)
153+
use vortex_tensor::encodings::turboquant::TurboQuantScheme;
154+
self.with_new_scheme(&TurboQuantScheme)
155155
}
156156

157157
/// Excludes schemes without CUDA kernel support and adds Zstd for string compression.

vortex-tensor/src/encodings/turboquant/centroids.rs renamed to vortex-tensor/src/encodings/turboquant/array/centroids.rs

File renamed without changes.

vortex-tensor/src/encodings/turboquant/array.rs renamed to vortex-tensor/src/encodings/turboquant/array/data.rs

Lines changed: 3 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -1,76 +1,19 @@
11
// SPDX-License-Identifier: Apache-2.0
22
// SPDX-FileCopyrightText: Copyright the Vortex contributors
33

4-
//! TurboQuant array definition: stores quantized coordinate codes, norms,
5-
//! centroids (codebook), and rotation signs.
6-
7-
use vortex_array::ArrayId;
84
use vortex_array::ArrayRef;
95
use vortex_array::dtype::DType;
106
use vortex_array::dtype::Nullability;
117
use vortex_array::dtype::PType;
128
use vortex_array::stats::ArrayStats;
13-
use vortex_array::vtable;
149
use vortex_error::VortexExpect;
1510
use vortex_error::VortexResult;
1611
use vortex_error::vortex_ensure;
1712

13+
use crate::encodings::turboquant::array::slots::Slot;
14+
use crate::encodings::turboquant::vtable::TurboQuant;
1815
use crate::utils::extension_element_ptype;
1916
use crate::utils::extension_list_size;
20-
use crate::vector::Vector;
21-
22-
/// Encoding marker type for TurboQuant.
23-
#[derive(Clone, Debug)]
24-
pub struct TurboQuant;
25-
26-
impl TurboQuant {
27-
pub const ID: ArrayId = ArrayId::new_ref("vortex.turboquant");
28-
}
29-
30-
vtable!(TurboQuant, TurboQuant, TurboQuantData);
31-
32-
/// Serialized metadata for TurboQuant encoding: a single byte holding the `bit_width` (0-8).
33-
///
34-
/// All other fields (dimension, element type) are derived from the dtype and children.
35-
/// A `bit_width` of 0 indicates a degenerate empty array.
36-
#[derive(Clone, Debug)]
37-
pub struct TurboQuantMetadata {
38-
/// MSE bits per coordinate (0 for degenerate empty arrays, 1-8 otherwise).
39-
pub bit_width: u8,
40-
}
41-
42-
/// Slot positions for TurboQuantArray children.
43-
#[repr(usize)]
44-
#[derive(Clone, Copy, Debug)]
45-
pub(crate) enum Slot {
46-
Codes = 0,
47-
Norms = 1,
48-
Centroids = 2,
49-
RotationSigns = 3,
50-
}
51-
52-
impl Slot {
53-
pub(crate) const COUNT: usize = 4;
54-
55-
pub(crate) fn name(self) -> &'static str {
56-
match self {
57-
Self::Codes => "codes",
58-
Self::Norms => "norms",
59-
Self::Centroids => "centroids",
60-
Self::RotationSigns => "rotation_signs",
61-
}
62-
}
63-
64-
pub(crate) fn from_index(idx: usize) -> Self {
65-
match idx {
66-
0 => Self::Codes,
67-
1 => Self::Norms,
68-
2 => Self::Centroids,
69-
3 => Self::RotationSigns,
70-
_ => vortex_error::vortex_panic!("invalid slot index {idx}"),
71-
}
72-
}
73-
}
7417

7518
/// TurboQuant array data.
7619
///
@@ -207,22 +150,8 @@ impl TurboQuantData {
207150
centroids: &ArrayRef,
208151
rotation_signs: &ArrayRef,
209152
) -> VortexResult<()> {
210-
// Dtype must be a Vector extension type.
211-
let ext = dtype
212-
.as_extension_opt()
213-
.filter(|e| e.is::<Vector>())
214-
.ok_or_else(|| {
215-
vortex_error::vortex_err!(
216-
"TurboQuant dtype must be a Vector extension type, got {dtype}"
217-
)
218-
})?;
219-
220-
// Dimension is derived from the storage dtype's list size and must be >= 3.
153+
let ext = TurboQuant::validate_dtype(dtype)?;
221154
let dimension = extension_list_size(ext)?;
222-
vortex_ensure!(
223-
dimension >= 3,
224-
"TurboQuant requires dimension >= 3, got {dimension}"
225-
);
226155

227156
let num_rows = norms.len();
228157

‎vortex-tensor/src/encodings/turboquant/array/metadata.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
/// Serialized metadata for TurboQuant encoding: a single byte holding the `bit_width` (0-8).
5+
///
6+
/// All other fields (dimension, element type) are derived from the dtype and children.
7+
/// A `bit_width` of 0 indicates a degenerate empty array.
8+
#[derive(Clone, Debug)]
9+
pub struct TurboQuantMetadata {
10+
/// MSE bits per coordinate (0 for degenerate empty arrays, 1-8 otherwise).
11+
pub bit_width: u8,
12+
}

‎vortex-tensor/src/encodings/turboquant/array/mod.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! TurboQuant array definition: stores quantized coordinate codes, norms, centroids (codebook),
5+
//! and rotation signs.
6+
7+
pub(crate) mod data;
8+
pub(crate) mod metadata;
9+
pub(crate) mod slots;
10+
11+
pub(crate) mod scheme;
12+
13+
pub(crate) mod centroids;
14+
pub(crate) mod rotation;

vortex-tensor/src/encodings/turboquant/rotation.rs renamed to vortex-tensor/src/encodings/turboquant/array/rotation.rs

File renamed without changes.

vortex-tensor/src/encodings/turboquant/scheme.rs renamed to vortex-tensor/src/encodings/turboquant/array/scheme.rs

Lines changed: 21 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -5,42 +5,37 @@
55
66
use vortex_array::ArrayRef;
77
use vortex_array::Canonical;
8-
use vortex_array::dtype::DType;
9-
use vortex_array::dtype::Nullability;
10-
use vortex_array::dtype::PType;
118
use vortex_compressor::CascadingCompressor;
129
use vortex_compressor::ctx::CompressorContext;
1310
use vortex_compressor::scheme::Scheme;
1411
use vortex_compressor::stats::ArrayAndStats;
1512
use vortex_error::VortexResult;
16-
use vortex_error::vortex_bail;
17-
use vortex_error::vortex_ensure;
1813

19-
use super::FIXED_SHAPE_TENSOR_EXT_ID;
20-
use super::TurboQuantConfig;
21-
use super::VECTOR_EXT_ID;
22-
use super::turboquant_encode;
14+
use crate::encodings::turboquant::TurboQuant;
15+
use crate::encodings::turboquant::TurboQuantConfig;
16+
use crate::encodings::turboquant::turboquant_encode;
17+
use crate::utils::extension_element_ptype;
18+
use crate::utils::extension_list_size;
2319

24-
/// TurboQuant compression scheme for tensor extension types.
20+
/// TurboQuant compression scheme for [`Vector`] extension types.
2521
///
26-
/// Applies lossy vector quantization to `Vector` and `FixedShapeTensor` extension
27-
/// arrays using the TurboQuant algorithm with MSE-optimal encoding.
22+
/// Applies lossy vector quantization to [`Vector`] extension arrays using the TurboQuant
23+
/// algorithm with MSE-optimal encoding.
2824
///
2925
/// Register this scheme with the compressor builder via `with_scheme`:
3026
/// ```ignore
3127
/// use vortex_btrblocks::BtrBlocksCompressorBuilder;
32-
/// use vortex_tensor::encodings::turboquant::scheme::TURBOQUANT_SCHEME;
28+
/// use vortex_tensor::encodings::turboquant::TurboQuantScheme;
3329
///
3430
/// let compressor = BtrBlocksCompressorBuilder::default()
35-
/// .with_scheme(&TURBOQUANT_SCHEME)
31+
/// .with_scheme(&TurboQuantScheme)
3632
/// .build();
3733
/// ```
34+
///
35+
/// [`Vector`]: crate::vector::Vector
3836
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
3937
pub struct TurboQuantScheme;
4038

41-
/// Static instance for registration with `BtrBlocksCompressorBuilder::with_scheme`.
42-
pub static TURBOQUANT_SCHEME: TurboQuantScheme = TurboQuantScheme;
43-
4439
impl Scheme for TurboQuantScheme {
4540
fn scheme_name(&self) -> &'static str {
4641
"vortex.tensor.turboquant"
@@ -51,7 +46,7 @@ impl Scheme for TurboQuantScheme {
5146
return false;
5247
};
5348

54-
get_tensor_element_ptype_and_length(ext.dtype()).is_ok()
49+
TurboQuant::validate_dtype(ext.dtype()).is_ok()
5550
}
5651

5752
fn expected_compression_ratio(
@@ -62,10 +57,14 @@ impl Scheme for TurboQuantScheme {
6257
) -> VortexResult<f64> {
6358
let dtype = data.array().dtype();
6459
let len = data.array().len();
65-
let (element_ptype, dimensions) = get_tensor_element_ptype_and_length(dtype)?;
60+
61+
let ext = TurboQuant::validate_dtype(dtype)?;
62+
let element_ptype = extension_element_ptype(ext)?;
63+
let dimension = extension_list_size(ext)?;
64+
6665
Ok(estimate_compression_ratio(
6766
element_ptype.bit_width(),
68-
dimensions,
67+
dimension,
6968
len,
7069
))
7170
}
@@ -76,8 +75,8 @@ impl Scheme for TurboQuantScheme {
7675
data: &mut ArrayAndStats,
7776
_ctx: CompressorContext,
7877
) -> VortexResult<ArrayRef> {
79-
let array = data.array().clone();
80-
let ext_array = array.to_canonical()?.into_extension();
78+
// TODO(connor): Fix this once we ensure that the data array is always canonical.
79+
let ext_array = data.array().to_canonical()?.into_extension();
8180

8281
let config = TurboQuantConfig::default();
8382
turboquant_encode(&ext_array, &config, &mut compressor.execution_ctx())
@@ -104,40 +103,6 @@ fn estimate_compression_ratio(bits_per_element: usize, dimensions: u32, num_vect
104103
uncompressed_size_bits as f64 / compressed_size_bits as f64
105104
}
106105

107-
fn get_tensor_element_ptype_and_length(dtype: &DType) -> VortexResult<(PType, u32)> {
108-
let ext_id = dtype.as_extension().id();
109-
let is_tensor = dtype.is_extension()
110-
&& (ext_id.as_ref() == VECTOR_EXT_ID || ext_id.as_ref() == FIXED_SHAPE_TENSOR_EXT_ID);
111-
vortex_ensure!(is_tensor, "expected tensor extension dtype, got {}", dtype);
112-
113-
let storage_dtype = dtype.as_extension().storage_dtype();
114-
let (element_dtype, fsl_len) = match storage_dtype {
115-
DType::FixedSizeList(element_dtype, list_size, _) => (element_dtype, list_size),
116-
_ => vortex_bail!(
117-
"expected FixedSizeList storage dtype, got {}",
118-
storage_dtype
119-
),
120-
};
121-
122-
// TurboQuant requires dimension >= 3: the marginal coordinate distribution
123-
// (1 - x^2)^((d-3)/2) has a singularity at d=2 (arcsine distribution) that
124-
// causes NaN in the Max-Lloyd centroid computation.
125-
vortex_ensure!(
126-
*fsl_len >= 3,
127-
"TurboQuant requires dimension >= 3, got {}",
128-
fsl_len
129-
);
130-
131-
if let &DType::Primitive(ptype, Nullability::NonNullable) = element_dtype.as_ref() {
132-
Ok((ptype, *fsl_len))
133-
} else {
134-
vortex_bail!(
135-
"expected non-nullable primitive element type, got {}",
136-
element_dtype
137-
);
138-
}
139-
}
140-
141106
#[cfg(test)]
142107
mod tests {
143108
use rstest::rstest;

‎vortex-tensor/src/encodings/turboquant/array/slots.rs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
/// Slot positions for TurboQuantArray children.
5+
#[repr(usize)]
6+
#[derive(Clone, Copy, Debug)]
7+
pub(crate) enum Slot {
8+
Codes = 0,
9+
Norms = 1,
10+
Centroids = 2,
11+
RotationSigns = 3,
12+
}
13+
14+
impl Slot {
15+
pub(crate) const COUNT: usize = 4;
16+
17+
pub(crate) fn name(self) -> &'static str {
18+
match self {
19+
Self::Codes => "codes",
20+
Self::Norms => "norms",
21+
Self::Centroids => "centroids",
22+
Self::RotationSigns => "rotation_signs",
23+
}
24+
}
25+
26+
pub(crate) fn from_index(idx: usize) -> Self {
27+
match idx {
28+
0 => Self::Codes,
29+
1 => Self::Norms,
30+
2 => Self::Centroids,
31+
3 => Self::RotationSigns,
32+
_ => vortex_error::vortex_panic!("invalid slot index {idx}"),
33+
}
34+
}
35+
}

‎vortex-tensor/src/encodings/turboquant/compress.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,11 @@ use vortex_error::vortex_bail;
2020
use vortex_error::vortex_ensure;
2121
use vortex_fastlanes::bitpack_compress::bitpack_encode;
2222

23-
use crate::encodings::turboquant::array::TurboQuantData;
24-
use crate::encodings::turboquant::centroids::compute_boundaries;
25-
use crate::encodings::turboquant::centroids::find_nearest_centroid;
26-
use crate::encodings::turboquant::centroids::get_centroids;
27-
use crate::encodings::turboquant::rotation::RotationMatrix;
23+
use crate::encodings::turboquant::TurboQuantData;
24+
use crate::encodings::turboquant::array::centroids::compute_boundaries;
25+
use crate::encodings::turboquant::array::centroids::find_nearest_centroid;
26+
use crate::encodings::turboquant::array::centroids::get_centroids;
27+
use crate::encodings::turboquant::array::rotation::RotationMatrix;
2828
use crate::scalar_fns::ApproxOptions;
2929
use crate::scalar_fns::l2_norm::L2Norm;
3030

‎vortex-tensor/src/encodings/turboquant/compute/ops.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ use vortex_array::vtable::OperationsVTable;
1010
use vortex_error::VortexResult;
1111
use vortex_error::vortex_bail;
1212

13-
use crate::encodings::turboquant::array::TurboQuant;
13+
use crate::encodings::turboquant::TurboQuant;
1414

1515
impl OperationsVTable<TurboQuant> for TurboQuant {
1616
fn scalar_at(

0 commit comments

Comments
 (0)