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
8#![cfg_attr(
9    test,
10    allow(clippy::unwrap_used, clippy::expect_used, clippy::unwrap_in_result)
11)]
12
13use std::sync::Arc;
14
15use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayPlugin;
16use vortex_array::arrow::ArrowSessionExt;
17use vortex_array::dtype::session::DTypeSessionExt;
18use vortex_array::scalar_fn::session::ScalarFnSessionExt;
19use vortex_array::session::ArraySessionExt;
20use vortex_session::VortexSession;
21
22use crate::scalar_fns::cosine_similarity::CosineSimilarity;
23use crate::scalar_fns::inner_product::InnerProduct;
24use crate::scalar_fns::l2_denorm::L2Denorm;
25use crate::scalar_fns::l2_norm::L2Norm;
26use crate::scalar_fns::sorf_transform::SorfTransform;
27use crate::types::fixed_shape_tensor::FixedShapeTensor;
28use crate::types::vector::Vector;
29
30pub mod matcher;
31pub mod scalar_fns;
32
33mod types;
34
35pub use types::fixed_shape_tensor;
36pub use types::vector;
37
38pub mod encodings;
39
40pub mod vector_search;
41
42mod utils;
43
44/// Environment variable that gates registration of the tensor scalar-fn array plugins (the array
45/// encodings that let [`CosineSimilarity`], [`InnerProduct`], [`L2Denorm`], [`L2Norm`], and
46/// [`SorfTransform`] persist in a Vortex file). When unset, only the scalar functions themselves
47/// are registered; readers of files containing serialized tensor scalar-fn arrays will fail to
48/// deserialize. Opt-in by setting the variable to any non-empty value.
49pub const SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str = "VX_SCALAR_FN_ARRAY_TENSOR_PLUGIN";
50
51/// Initialize the Vortex tensor library with a Vortex session.
52pub fn initialize(session: &VortexSession) {
53    session.dtypes().register(Vector);
54    session.dtypes().register(FixedShapeTensor);
55
56    let arrow_session = session.arrow();
57    arrow_session.register_exporter(Arc::new(Vector));
58    arrow_session.register_importer(Arc::new(Vector));
59
60    let session_fns = session.scalar_fns();
61
62    session_fns.register(CosineSimilarity);
63    session_fns.register(InnerProduct);
64    session_fns.register(L2Denorm);
65    session_fns.register(L2Norm);
66    session_fns.register(SorfTransform);
67
68    // Registering the scalar-fn array plugins lets the tensor scalar fns be serialized as array
69    // encodings inside Vortex files. Gate this on an env var so applications that do not intend
70    // to persist these encodings do not pay the registry cost or widen their stable-encoding
71    // surface unintentionally.
72    if std::env::var_os(SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV).is_some_and(|v| !v.is_empty()) {
73        let session_arrays = session.arrays();
74
75        session_arrays.register(ScalarFnArrayPlugin::new(CosineSimilarity));
76        session_arrays.register(ScalarFnArrayPlugin::new(InnerProduct));
77        session_arrays.register(ScalarFnArrayPlugin::new(L2Denorm));
78        session_arrays.register(ScalarFnArrayPlugin::new(L2Norm));
79        session_arrays.register(ScalarFnArrayPlugin::new(SorfTransform));
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use std::sync::LazyLock;
86
87    use vortex_array::session::ArraySession;
88    use vortex_session::VortexSession;
89
90    pub static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
91        let session = VortexSession::empty().with::<ArraySession>();
92        crate::initialize(&session);
93        session
94    });
95}
vortex_tensor/lib.rs

vortex_tensor/
lib.rs
