vortex_tensor/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Types and functionality for working with tensors, vectors, and related mathematical constructs
5//! including unit vectors, spherical coordinates, and similarity measures such as cosine
6//! similarity.
7
8use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin;
9use vortex_array::dtype::session::DTypeSessionExt;
10use vortex_array::scalar_fn::session::ScalarFnSessionExt;
11use vortex_array::session::ArraySessionExt;
12use vortex_session::VortexSession;
13
14use crate::fixed_shape::FixedShapeTensor;
15use crate::scalar_fns::cosine_similarity::CosineSimilarity;
16use crate::scalar_fns::inner_product::InnerProduct;
17use crate::scalar_fns::l2_denorm::L2Denorm;
18use crate::scalar_fns::l2_norm::L2Norm;
19use crate::scalar_fns::sorf_transform::SorfTransform;
20use crate::vector::Vector;
21
22pub mod matcher;
23pub mod scalar_fns;
24
25pub mod fixed_shape;
26pub mod vector;
27
28pub mod encodings;
29
30pub mod vector_search;
31
32mod utils;
33
34/// Environment variable that gates registration of the tensor scalar-fn array plugins (the array
35/// encodings that let [`CosineSimilarity`], [`InnerProduct`], [`L2Denorm`], [`L2Norm`], and
36/// [`SorfTransform`] persist in a Vortex file). When unset, only the scalar functions themselves
37/// are registered; readers of files containing serialized tensor scalar-fn arrays will fail to
38/// deserialize. Opt-in by setting the variable to any non-empty value.
39pub const SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str = "VX_SCALAR_FN_ARRAY_TENSOR_PLUGIN";
40
41/// Initialize the Vortex tensor library with a Vortex session.
42pub fn initialize(session: &VortexSession) {
43    session.dtypes().register(Vector);
44    session.dtypes().register(FixedShapeTensor);
45
46    let session_fns = session.scalar_fns();
47
48    session_fns.register(CosineSimilarity);
49    session_fns.register(InnerProduct);
50    session_fns.register(L2Denorm);
51    session_fns.register(L2Norm);
52    session_fns.register(SorfTransform);
53
54    // Registering the scalar-fn array plugins lets the tensor scalar fns be serialized as array
55    // encodings inside Vortex files. Gate this on an env var so applications that do not intend
56    // to persist these encodings do not pay the registry cost or widen their stable-encoding
57    // surface unintentionally.
58    if std::env::var_os(SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV).is_some_and(|v| !v.is_empty()) {
59        let session_arrays = session.arrays();
60
61        session_arrays.register(ScalarFnArrayPlugin::new(CosineSimilarity));
62        session_arrays.register(ScalarFnArrayPlugin::new(InnerProduct));
63        session_arrays.register(ScalarFnArrayPlugin::new(L2Denorm));
64        session_arrays.register(ScalarFnArrayPlugin::new(L2Norm));
65        session_arrays.register(ScalarFnArrayPlugin::new(SorfTransform));
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use std::sync::LazyLock;
72
73    use vortex_array::session::ArraySession;
74    use vortex_session::VortexSession;
75
76    pub static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
77        let session = VortexSession::empty().with::<ArraySession>();
78        crate::initialize(&session);
79        session
80    });
81}
vortex_tensor/lib.rs

vortex_tensor/
lib.rs
