Skip to content

Commit 3d7dfed

Browse files
connortsui20connortsui20
committed
rebase to new vtable world
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent af3cb41 commit 3d7dfed

10 files changed

Lines changed: 112 additions & 221 deletions

File tree

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

Lines changed: 4 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ use vortex_array::ArrayRef;
77
use vortex_array::dtype::DType;
88
use vortex_array::dtype::Nullability;
99
use vortex_array::dtype::PType;
10-
use vortex_array::stats::ArrayStats;
1110
use vortex_error::VortexExpect;
1211
use vortex_error::VortexResult;
1312
use vortex_error::vortex_ensure;
@@ -29,12 +28,6 @@ use crate::utils::extension_list_size;
2928
/// A degenerate TurboQuant array has zero rows and `bit_width == 0`, with all slots empty.
3029
#[derive(Clone, Debug)]
3130
pub struct TurboQuantData {
32-
/// The [`Vector`](crate::vector::Vector) extension dtype that this array encodes.
33-
///
34-
/// The storage dtype within the extension determines the element type (f16, f32, or f64) and
35-
/// the list size (dimension).
36-
pub(crate) dtype: DType,
37-
3831
/// Child arrays stored as slots. See [`Slot`] for positions:
3932
///
4033
/// - [`Codes`](Slot::Codes): Non-nullable `FixedSizeListArray<u8>` with
@@ -66,9 +59,6 @@ pub struct TurboQuantData {
6659
///
6760
/// This is 0 for degenerate empty arrays.
6861
pub(crate) bit_width: u8,
69-
70-
/// The stats for this array.
71-
pub(crate) stats_set: ArrayStats,
7262
}
7363

7464
impl TurboQuantData {
@@ -83,13 +73,13 @@ impl TurboQuantData {
8373
/// Returns an error if the provided components do not satisfy the invariants documented
8474
/// in [`new_unchecked`](Self::new_unchecked).
8575
pub fn try_new(
86-
dtype: DType,
76+
dtype: &DType,
8777
codes: ArrayRef,
8878
norms: ArrayRef,
8979
centroids: ArrayRef,
9080
rotation_signs: ArrayRef,
9181
) -> VortexResult<Self> {
92-
Self::validate(&dtype, &codes, &norms, &centroids, &rotation_signs)?;
82+
Self::validate(dtype, &codes, &norms, &centroids, &rotation_signs)?;
9383

9484
// SAFETY: we validate that the inputs are valid above.
9585
Ok(unsafe { Self::new_unchecked(dtype, codes, norms, centroids, rotation_signs) })
@@ -115,14 +105,14 @@ impl TurboQuantData {
115105
///
116106
/// Violating these invariants may produce incorrect results during decompression.
117107
pub unsafe fn new_unchecked(
118-
dtype: DType,
108+
dtype: &DType,
119109
codes: ArrayRef,
120110
norms: ArrayRef,
121111
centroids: ArrayRef,
122112
rotation_signs: ArrayRef,
123113
) -> Self {
124114
#[cfg(debug_assertions)]
125-
Self::validate(&dtype, &codes, &norms, &centroids, &rotation_signs)
115+
Self::validate(dtype, &codes, &norms, &centroids, &rotation_signs)
126116
.vortex_expect("[Debug Assertion]: Invalid TurboQuantData parameters");
127117

128118
let dimension = dtype
@@ -147,11 +137,9 @@ impl TurboQuantData {
147137
slots[Slot::RotationSigns as usize] = Some(rotation_signs);
148138

149139
Self {
150-
dtype,
151140
slots,
152141
dimension,
153142
bit_width,
154-
stats_set: Default::default(),
155143
}
156144
}
157145

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

Lines changed: 0 additions & 12 deletions
This file was deleted.

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
//! and rotation signs.
66
77
pub(crate) mod data;
8-
pub(crate) mod metadata;
98
pub(crate) mod slots;
109

1110
pub(crate) mod centroids;

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,12 @@ 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::TurboQuantData;
23+
use crate::encodings::turboquant::TurboQuant;
2424
use crate::encodings::turboquant::array::centroids::compute_boundaries;
2525
use crate::encodings::turboquant::array::centroids::find_nearest_centroid;
2626
use crate::encodings::turboquant::array::centroids::get_centroids;
2727
use crate::encodings::turboquant::array::rotation::RotationMatrix;
28+
use crate::encodings::turboquant::vtable::TurboQuantArray;
2829
use crate::scalar_fns::ApproxOptions;
2930
use crate::scalar_fns::l2_norm::L2Norm;
3031

@@ -168,7 +169,7 @@ fn build_turboquant(
168169
fsl: &FixedSizeListArray,
169170
core: QuantizationResult,
170171
ext_dtype: DType,
171-
) -> VortexResult<TurboQuantData> {
172+
) -> VortexResult<TurboQuantArray> {
172173
let num_rows = fsl.len();
173174
let padded_dim = core.padded_dim;
174175
let codes_elements =
@@ -190,7 +191,7 @@ fn build_turboquant(
190191

191192
let rotation_signs = bitpack_rotation_signs(&core.rotation)?;
192193

193-
TurboQuantData::try_new(
194+
TurboQuant::try_new_array(
194195
ext_dtype,
195196
codes,
196197
core.norms_array,
@@ -241,7 +242,7 @@ pub fn turboquant_encode(
241242

242243
let empty_centroids = PrimitiveArray::empty::<f32>(Nullability::NonNullable);
243244
let empty_signs = PrimitiveArray::empty::<u8>(Nullability::NonNullable);
244-
return Ok(TurboQuantData::try_new(
245+
return Ok(TurboQuant::try_new_array(
245246
ext_dtype,
246247
empty_codes.into_array(),
247248
empty_norms,

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ use vortex_array::arrays::slice::SliceReduce;
1010
use vortex_error::VortexResult;
1111

1212
use crate::encodings::turboquant::TurboQuant;
13-
use crate::encodings::turboquant::TurboQuantData;
1413

1514
impl SliceReduce for TurboQuant {
1615
fn slice(
@@ -20,14 +19,15 @@ impl SliceReduce for TurboQuant {
2019
let sliced_codes = array.codes().slice(range.clone())?;
2120
let sliced_norms = array.norms().slice(range)?;
2221

23-
let result = TurboQuantData::try_new(
24-
array.dtype.clone(),
25-
sliced_codes,
26-
sliced_norms,
27-
array.centroids().clone(),
28-
array.rotation_signs().clone(),
29-
)?;
30-
31-
Ok(Some(result.into_array()))
22+
Ok(Some(
23+
TurboQuant::try_new_array(
24+
array.dtype().clone(),
25+
sliced_codes,
26+
sliced_norms,
27+
array.centroids().clone(),
28+
array.rotation_signs().clone(),
29+
)?
30+
.into_array(),
31+
))
3232
}
3333
}

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ use vortex_array::arrays::dict::TakeExecute;
99
use vortex_error::VortexResult;
1010

1111
use crate::encodings::turboquant::TurboQuant;
12-
use crate::encodings::turboquant::TurboQuantData;
1312

1413
impl TakeExecute for TurboQuant {
1514
fn take(
@@ -21,14 +20,15 @@ impl TakeExecute for TurboQuant {
2120
let taken_codes = array.codes().take(indices.clone())?;
2221
let taken_norms = array.norms().take(indices.clone())?;
2322

24-
let result = TurboQuantData::try_new(
25-
array.dtype.clone(),
26-
taken_codes,
27-
taken_norms,
28-
array.centroids().clone(),
29-
array.rotation_signs().clone(),
30-
)?;
31-
32-
Ok(Some(result.into_array()))
23+
Ok(Some(
24+
TurboQuant::try_new_array(
25+
array.dtype().clone(),
26+
taken_codes,
27+
taken_norms,
28+
array.centroids().clone(),
29+
array.rotation_signs().clone(),
30+
)?
31+
.into_array(),
32+
))
3333
}
3434
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ pub fn execute_decompress(
3636
let dim = array.dimension() as usize;
3737
let padded_dim = array.padded_dim() as usize;
3838
let num_rows = array.norms().len();
39-
let ext_dtype = array.dtype.as_extension().clone();
39+
let ext_dtype = array.dtype().as_extension().clone();
4040
let element_ptype = extension_element_ptype(&ext_dtype)?;
4141

4242
if num_rows == 0 {

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

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,6 @@ pub fn initialize(session: &mut VortexSession) {
100100

101101
mod array;
102102
pub use array::data::TurboQuantData;
103-
pub use array::metadata::TurboQuantMetadata;
104103
pub use array::scheme::TurboQuantScheme;
105104

106105
pub(crate) mod compute;

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

Lines changed: 0 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ use rand::rngs::StdRng;
88
use rand_distr::Distribution;
99
use rand_distr::Normal;
1010
use rstest::rstest;
11-
use vortex_array::ArrayRef;
1211
use vortex_array::IntoArray;
1312
use vortex_array::VortexSessionExecute;
1413
use vortex_array::arrays::ExtensionArray;
@@ -451,82 +450,6 @@ fn stored_rotation_signs_produce_correct_decode() -> VortexResult<()> {
451450
Ok(())
452451
}
453452

454-
// -----------------------------------------------------------------------
455-
// Serde roundtrip
456-
// -----------------------------------------------------------------------
457-
458-
#[test]
459-
fn serde_roundtrip() -> VortexResult<()> {
460-
use vortex_array::vtable::VTable;
461-
462-
let fsl = make_fsl(10, 128, 42);
463-
let ext = make_vector_ext(&fsl);
464-
let config = TurboQuantConfig {
465-
bit_width: 3,
466-
seed: Some(123),
467-
};
468-
let mut ctx = SESSION.create_execution_ctx();
469-
let encoded = turboquant_encode(&ext, &config, &mut ctx)?;
470-
let encoded = encoded.as_opt::<TurboQuant>().unwrap();
471-
472-
// Serialize metadata.
473-
let metadata = <TurboQuant as VTable>::metadata(encoded)?;
474-
let serialized =
475-
<TurboQuant as VTable>::serialize(metadata)?.expect("metadata should serialize");
476-
477-
// Collect children.
478-
let nchildren = <TurboQuant as VTable>::nchildren(encoded);
479-
assert_eq!(nchildren, 4);
480-
let children: Vec<ArrayRef> = (0..nchildren)
481-
.map(|i| <TurboQuant as VTable>::child(encoded, i))
482-
.collect();
483-
484-
// Deserialize and rebuild.
485-
let deserialized = <TurboQuant as VTable>::deserialize(
486-
&serialized,
487-
encoded.dtype(),
488-
encoded.len(),
489-
&[],
490-
&SESSION,
491-
)?;
492-
493-
// Verify metadata fields survived roundtrip.
494-
assert_eq!(deserialized.bit_width, encoded.bit_width());
495-
496-
// Verify the rebuilt array decodes identically.
497-
let mut ctx = SESSION.create_execution_ctx();
498-
let decoded_original = encoded
499-
.array()
500-
.clone()
501-
.execute::<ExtensionArray>(&mut ctx)?;
502-
let original_fsl = decoded_original
503-
.storage_array()
504-
.to_canonical()?
505-
.into_fixed_size_list();
506-
let original_elements = original_fsl.elements().to_canonical()?.into_primitive();
507-
508-
// Rebuild from children (simulating deserialization).
509-
let rebuilt = crate::encodings::turboquant::TurboQuantData::try_new(
510-
encoded.dtype().clone(),
511-
children[0].clone(),
512-
children[1].clone(),
513-
children[2].clone(),
514-
children[3].clone(),
515-
)?;
516-
let decoded_rebuilt = rebuilt.into_array().execute::<ExtensionArray>(&mut ctx)?;
517-
let rebuilt_fsl = decoded_rebuilt
518-
.storage_array()
519-
.to_canonical()?
520-
.into_fixed_size_list();
521-
let rebuilt_elements = rebuilt_fsl.elements().to_canonical()?.into_primitive();
522-
523-
assert_eq!(
524-
original_elements.as_slice::<f32>(),
525-
rebuilt_elements.as_slice::<f32>()
526-
);
527-
Ok(())
528-
}
529-
530453
// -----------------------------------------------------------------------
531454
// Compute pushdown tests
532455
// -----------------------------------------------------------------------

0 commit comments

Comments
 (0)