Skip to content

Commit 0e0401c

Browse files
authored
Gate scalar fn array plugin for tensor exprs (#7447)
## Summary This should have been a part of #7437 ## Testing N/A Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent 89de477 commit 0e0401c

2 files changed

Lines changed: 22 additions & 6 deletions

File tree

vortex-tensor/public-api.lock

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -588,4 +588,6 @@ pub fn vortex_tensor::vector_search::build_similarity_search_tree<T: vortex_arra
588588

589589
pub fn vortex_tensor::vector_search::compress_turboquant(data: vortex_array::array::erased::ArrayRef, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
590590

591+
pub const vortex_tensor::SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str
592+
591593
pub fn vortex_tensor::initialize(session: &vortex_session::VortexSession)

vortex-tensor/src/lib.rs

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,25 +31,39 @@ pub mod vector_search;
3131

3232
mod utils;
3333

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.
39+
pub const SCALAR_FN_ARRAY_TENSOR_PLUGIN_ENV: &str = "VX_SCALAR_FN_ARRAY_TENSOR_PLUGIN";
40+
3441
/// Initialize the Vortex tensor library with a Vortex session.
3542
pub fn initialize(session: &VortexSession) {
3643
session.dtypes().register(Vector);
3744
session.dtypes().register(FixedShapeTensor);
3845

3946
let session_fns = session.scalar_fns();
40-
let session_arrays = session.arrays();
4147

4248
session_fns.register(CosineSimilarity);
4349
session_fns.register(InnerProduct);
4450
session_fns.register(L2Denorm);
4551
session_fns.register(L2Norm);
4652
session_fns.register(SorfTransform);
4753

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));
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+
}
5367
}
5468

5569
#[cfg(test)]

0 commit comments

Comments
 (0)