Skip to content

Commit 4a5b7d7

Browse files
connortsui20connortsui20
authored
Support serializion tensor scalar fns (#7437)
## Summary Tracking issue: #7297 Adds serialization support to the tensor scalar fns. This is gated by the `VORTEX_TENSOR_SCALAR_FN_ARRAY_PLUGINS` environment variable. ## Testing Basic testing. ## Unresolved Questions - I should probably write this to a file and roundtrip that, right? - It is also not clear if it makes sense to have a `SorfTransformMetadata` when it is basically identical to `SorfOptions` but with prost. --------- Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent c4ad4e2 commit 4a5b7d7

9 files changed

Lines changed: 624 additions & 42 deletions

File tree

‎vortex-tensor/public-api.lock‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,12 @@ impl core::clone::Clone for vortex_tensor::scalar_fns::cosine_similarity::Cosine
248248

249249
pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::clone(&self) -> vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity
250250

251+
impl vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable for vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity
252+
253+
pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::deserialize(&self, _dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], children: &dyn vortex_array::serde::ArrayChildren, session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts<Self>>
254+
255+
pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::serialize(&self, view: &vortex_array::arrays::scalar_fn::vtable::ScalarFnArrayView<'_, Self>, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
256+
251257
impl vortex_array::scalar_fn::vtable::ScalarFnVTable for vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity
252258

253259
pub type vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::Options = vortex_array::scalar_fn::vtable::EmptyOptions
@@ -284,6 +290,12 @@ impl core::clone::Clone for vortex_tensor::scalar_fns::inner_product::InnerProdu
284290

285291
pub fn vortex_tensor::scalar_fns::inner_product::InnerProduct::clone(&self) -> vortex_tensor::scalar_fns::inner_product::InnerProduct
286292

293+
impl vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable for vortex_tensor::scalar_fns::inner_product::InnerProduct
294+
295+
pub fn vortex_tensor::scalar_fns::inner_product::InnerProduct::deserialize(&self, _dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], children: &dyn vortex_array::serde::ArrayChildren, session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts<Self>>
296+
297+
pub fn vortex_tensor::scalar_fns::inner_product::InnerProduct::serialize(&self, view: &vortex_array::arrays::scalar_fn::vtable::ScalarFnArrayView<'_, Self>, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
298+
287299
impl vortex_array::scalar_fn::vtable::ScalarFnVTable for vortex_tensor::scalar_fns::inner_product::InnerProduct
288300

289301
pub type vortex_tensor::scalar_fns::inner_product::InnerProduct::Options = vortex_array::scalar_fn::vtable::EmptyOptions
@@ -322,6 +334,12 @@ impl core::clone::Clone for vortex_tensor::scalar_fns::l2_denorm::L2Denorm
322334

323335
pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::clone(&self) -> vortex_tensor::scalar_fns::l2_denorm::L2Denorm
324336

337+
impl vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable for vortex_tensor::scalar_fns::l2_denorm::L2Denorm
338+
339+
pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::deserialize(&self, _dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], children: &dyn vortex_array::serde::ArrayChildren, session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts<Self>>
340+
341+
pub fn vortex_tensor::scalar_fns::l2_denorm::L2Denorm::serialize(&self, view: &vortex_array::arrays::scalar_fn::vtable::ScalarFnArrayView<'_, Self>, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
342+
325343
impl vortex_array::scalar_fn::vtable::ScalarFnVTable for vortex_tensor::scalar_fns::l2_denorm::L2Denorm
326344

327345
pub type vortex_tensor::scalar_fns::l2_denorm::L2Denorm::Options = vortex_array::scalar_fn::vtable::EmptyOptions
@@ -362,6 +380,12 @@ impl core::clone::Clone for vortex_tensor::scalar_fns::l2_norm::L2Norm
362380

363381
pub fn vortex_tensor::scalar_fns::l2_norm::L2Norm::clone(&self) -> vortex_tensor::scalar_fns::l2_norm::L2Norm
364382

383+
impl vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable for vortex_tensor::scalar_fns::l2_norm::L2Norm
384+
385+
pub fn vortex_tensor::scalar_fns::l2_norm::L2Norm::deserialize(&self, _dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], children: &dyn vortex_array::serde::ArrayChildren, session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts<Self>>
386+
387+
pub fn vortex_tensor::scalar_fns::l2_norm::L2Norm::serialize(&self, view: &vortex_array::arrays::scalar_fn::vtable::ScalarFnArrayView<'_, Self>, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
388+
365389
impl vortex_array::scalar_fn::vtable::ScalarFnVTable for vortex_tensor::scalar_fns::l2_norm::L2Norm
366390

367391
pub type vortex_tensor::scalar_fns::l2_norm::L2Norm::Options = vortex_array::scalar_fn::vtable::EmptyOptions
@@ -444,6 +468,12 @@ impl core::clone::Clone for vortex_tensor::scalar_fns::sorf_transform::SorfTrans
444468

445469
pub fn vortex_tensor::scalar_fns::sorf_transform::SorfTransform::clone(&self) -> vortex_tensor::scalar_fns::sorf_transform::SorfTransform
446470

471+
impl vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable for vortex_tensor::scalar_fns::sorf_transform::SorfTransform
472+
473+
pub fn vortex_tensor::scalar_fns::sorf_transform::SorfTransform::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts<Self>>
474+
475+
pub fn vortex_tensor::scalar_fns::sorf_transform::SorfTransform::serialize(&self, view: &vortex_array::arrays::scalar_fn::vtable::ScalarFnArrayView<'_, Self>, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
476+
447477
impl vortex_array::scalar_fn::vtable::ScalarFnVTable for vortex_tensor::scalar_fns::sorf_transform::SorfTransform
448478

449479
pub type vortex_tensor::scalar_fns::sorf_transform::SorfTransform::Options = vortex_tensor::scalar_fns::sorf_transform::SorfOptions

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

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,6 @@ mod nullable;
88
mod roundtrip;
99
mod structural;
1010

11-
use std::sync::LazyLock;
12-
1311
use rand::SeedableRng;
1412
use rand::rngs::StdRng;
1513
use rand_distr::Distribution;
@@ -29,22 +27,18 @@ use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
2927
use vortex_array::arrays::scalar_fn::ScalarFnArrayExt;
3028
use vortex_array::dtype::extension::ExtDType;
3129
use vortex_array::extension::EmptyMetadata;
32-
use vortex_array::session::ArraySession;
3330
use vortex_array::validity::Validity;
3431
use vortex_buffer::BufferMut;
3532
use vortex_error::VortexExpect;
3633
use vortex_error::VortexResult;
37-
use vortex_session::VortexSession;
3834

3935
use crate::encodings::turboquant::TurboQuantConfig;
4036
use crate::encodings::turboquant::turboquant_encode_unchecked;
4137
use crate::scalar_fns::l2_denorm::L2Denorm;
4238
use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm;
39+
use crate::tests::SESSION;
4340
use crate::vector::Vector;
4441

45-
static SESSION: LazyLock<VortexSession> =
46-
LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
47-
4842
/// Create a FixedSizeListArray of random f32 vectors with the given validity.
4943
fn make_fsl_with_validity(
5044
num_rows: usize,

‎vortex-tensor/src/lib.rs‎

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@
55
//! including unit vectors, spherical coordinates, and similarity measures such as cosine
66
//! similarity.
77
8+
use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin;
89
use vortex_array::dtype::session::DTypeSessionExt;
910
use vortex_array::scalar_fn::session::ScalarFnSessionExt;
11+
use vortex_array::session::ArraySessionExt;
1012
use vortex_session::VortexSession;
1113

1214
use crate::fixed_shape::FixedShapeTensor;
@@ -34,9 +36,32 @@ pub fn initialize(session: &VortexSession) {
3436
session.dtypes().register(Vector);
3537
session.dtypes().register(FixedShapeTensor);
3638

37-
session.scalar_fns().register(CosineSimilarity);
38-
session.scalar_fns().register(InnerProduct);
39-
session.scalar_fns().register(L2Denorm);
40-
session.scalar_fns().register(L2Norm);
41-
session.scalar_fns().register(SorfTransform);
39+
let session_fns = session.scalar_fns();
40+
let session_arrays = session.arrays();
41+
42+
session_fns.register(CosineSimilarity);
43+
session_fns.register(InnerProduct);
44+
session_fns.register(L2Denorm);
45+
session_fns.register(L2Norm);
46+
session_fns.register(SorfTransform);
47+
48+
session_arrays.register(ScalarFnArrayPlugin::new(CosineSimilarity));
49+
session_arrays.register(ScalarFnArrayPlugin::new(InnerProduct));
50+
session_arrays.register(ScalarFnArrayPlugin::new(L2Denorm));
51+
session_arrays.register(ScalarFnArrayPlugin::new(L2Norm));
52+
session_arrays.register(ScalarFnArrayPlugin::new(SorfTransform));
53+
}
54+
55+
#[cfg(test)]
56+
mod tests {
57+
use std::sync::LazyLock;
58+
59+
use vortex_array::session::ArraySession;
60+
use vortex_session::VortexSession;
61+
62+
pub static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
63+
let session = VortexSession::empty().with::<ArraySession>();
64+
crate::initialize(&session);
65+
session
66+
});
4267
}

‎vortex-tensor/src/scalar_fns/cosine_similarity.rs‎

Lines changed: 79 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@ use vortex_array::IntoArray;
1212
use vortex_array::arrays::PrimitiveArray;
1313
use vortex_array::arrays::ScalarFnArray;
1414
use vortex_array::arrays::scalar_fn::ExactScalarFn;
15+
use vortex_array::arrays::scalar_fn::ScalarFnArrayView;
16+
use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayParts;
17+
use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable;
1518
use vortex_array::builtins::ArrayBuiltins;
1619
use vortex_array::dtype::DType;
1720
use vortex_array::dtype::Nullability;
@@ -25,11 +28,14 @@ use vortex_array::scalar_fn::ExecutionArgs;
2528
use vortex_array::scalar_fn::ScalarFn;
2629
use vortex_array::scalar_fn::ScalarFnId;
2730
use vortex_array::scalar_fn::ScalarFnVTable;
31+
use vortex_array::serde::ArrayChildren;
2832
use vortex_array::validity::Validity;
2933
use vortex_buffer::Buffer;
3034
use vortex_error::VortexResult;
3135
use vortex_error::vortex_ensure;
36+
use vortex_session::VortexSession;
3237

38+
use crate::scalar_fns::inner_product::BinaryTensorOpMetadata;
3339
use crate::scalar_fns::inner_product::InnerProduct;
3440
use crate::scalar_fns::l2_denorm::L2Denorm;
3541
use crate::scalar_fns::l2_denorm::try_build_constant_l2_denorm;
@@ -221,6 +227,37 @@ impl ScalarFnVTable for CosineSimilarity {
221227
}
222228
}
223229

230+
impl ScalarFnArrayVTable for CosineSimilarity {
231+
fn serialize(
232+
&self,
233+
view: &ScalarFnArrayView<Self>,
234+
_session: &VortexSession,
235+
) -> VortexResult<Option<Vec<u8>>> {
236+
Ok(Some(BinaryTensorOpMetadata::encode_from_view(view)?))
237+
}
238+
239+
fn deserialize(
240+
&self,
241+
_dtype: &DType,
242+
len: usize,
243+
metadata: &[u8],
244+
children: &dyn ArrayChildren,
245+
session: &VortexSession,
246+
) -> VortexResult<ScalarFnArrayParts<Self>> {
247+
let reconstructed = BinaryTensorOpMetadata::decode_children(
248+
metadata,
249+
len,
250+
children,
251+
session,
252+
"CosineSimilarity",
253+
)?;
254+
Ok(ScalarFnArrayParts {
255+
options: EmptyOptions,
256+
children: reconstructed,
257+
})
258+
}
259+
}
260+
224261
impl CosineSimilarity {
225262
/// Both sides are `L2Denorm`: treat the normalized children as authoritative, so
226263
/// `cosine_similarity = dot(n_l, n_r)`.
@@ -292,31 +329,28 @@ impl CosineSimilarity {
292329

293330
#[cfg(test)]
294331
mod tests {
295-
use std::sync::LazyLock;
296332

297333
use rstest::rstest;
334+
use vortex_array::ArrayPlugin;
298335
use vortex_array::ArrayRef;
299336
use vortex_array::IntoArray;
300337
use vortex_array::VortexSessionExecute;
301338
use vortex_array::arrays::MaskedArray;
302339
use vortex_array::arrays::PrimitiveArray;
303340
use vortex_array::arrays::ScalarFnArray;
304-
use vortex_array::session::ArraySession;
341+
use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin;
305342
use vortex_array::validity::Validity;
306343
use vortex_error::VortexResult;
307-
use vortex_session::VortexSession;
308344

309345
use crate::scalar_fns::cosine_similarity::CosineSimilarity;
310346
use crate::scalar_fns::l2_denorm::L2Denorm;
347+
use crate::tests::SESSION;
311348
use crate::utils::test_helpers::assert_close;
312349
use crate::utils::test_helpers::constant_tensor_array;
313350
use crate::utils::test_helpers::constant_vector_array;
314351
use crate::utils::test_helpers::tensor_array;
315352
use crate::utils::test_helpers::vector_array;
316353

317-
static SESSION: LazyLock<VortexSession> =
318-
LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
319-
320354
/// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec<f64>`.
321355
fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef, len: usize) -> VortexResult<Vec<f64>> {
322356
let scalar_fn = CosineSimilarity::new().erased();
@@ -693,4 +727,43 @@ mod tests {
693727
);
694728
Ok(())
695729
}
730+
731+
#[rstest]
732+
#[case::vector(
733+
vector_array(3, &[1.0, 0.0, 0.0, 3.0, 4.0, 0.0]).unwrap(),
734+
vector_array(3, &[0.0, 1.0, 0.0, 3.0, 4.0, 0.0]).unwrap(),
735+
2,
736+
)]
737+
#[case::fixed_shape_tensor(
738+
tensor_array(&[2], &[1.0, 0.0, 3.0, 4.0]).unwrap(),
739+
tensor_array(&[2], &[0.0, 1.0, 3.0, 4.0]).unwrap(),
740+
2,
741+
)]
742+
fn serde_round_trip(
743+
#[case] lhs: ArrayRef,
744+
#[case] rhs: ArrayRef,
745+
#[case] len: usize,
746+
) -> VortexResult<()> {
747+
let original = CosineSimilarity::try_new_array(lhs.clone(), rhs.clone(), len)?.into_array();
748+
749+
let plugin = ScalarFnArrayPlugin::new(CosineSimilarity);
750+
let metadata = plugin
751+
.serialize(&original, &SESSION)?
752+
.expect("CosineSimilarity serialize must produce metadata");
753+
754+
let children = vec![lhs, rhs];
755+
let recovered = plugin.deserialize(
756+
original.dtype(),
757+
original.len(),
758+
&metadata,
759+
&[],
760+
&children,
761+
&SESSION,
762+
)?;
763+
764+
assert_eq!(recovered.dtype(), original.dtype());
765+
assert_eq!(recovered.len(), original.len());
766+
assert_eq!(recovered.encoding_id(), original.encoding_id());
767+
Ok(())
768+
}
696769
}

0 commit comments

Comments
 (0)