Vortex Fixed-Shape Tensor (#6812)

## Summary

Adds an experimental fixed-shape tensor extension type in a new
`vortex-tensor` crate.

See https://vortex-data.github.io/rfcs/rfc/0024.html for info about the
design of this tensor type.

Additionally adds a `CosineSimilarity` expression that takes 2 tensor
arrays and computes the cosine similarity of tensors in the arrays
(resulting in a `PrimitiveArray`).

## Testing

Adds some very basic tests for cosine similarity and tensor metadata
operations.

## Future Work

I think this was a good way to see if our `ExtVTable` is not completely
wrong, but at the same time this tells us nothing about what we might
want to add for extension arrays on the `ExtVTable` because we as long
as the storage `DType` is correct, any storage array is valid.

The more interesting expressions have not been implemented here. Those
would include:

- Cast
- Index / Slice (and lazily get back another tensor array, potentially
non-contiguous)
- Maybe others?

Additional work includes exporting to Arrow, NumPy, and PyTorch. Arrow
will require a cheap translation from logical to physical shape, but
other than that those conversions should be easy.

---------

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

Cargo.lock

@@ -11,6 +11,7 @@ members = [
     "vortex-io",
     "vortex-proto",
     "vortex-array",
+    "vortex-tensor",
     "vortex-btrblocks",
     "vortex-layout",
     "vortex-scan",
@@ -271,6 +272,7 @@ vortex-scan = { version = "0.1.0", path = "./vortex-scan", default-features = fa
 vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-features = false }
 vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false }
 vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false }
+vortex-tensor = { version = "0.1.0", path = "./vortex-tensor", default-features = false }
 vortex-utils = { version = "0.1.0", path = "./vortex-utils", default-features = false }
 vortex-zigzag = { version = "0.1.0", path = "./encodings/zigzag", default-features = false }
 vortex-zstd = { version = "0.1.0", path = "./encodings/zstd", default-features = false }

Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "vortex-tensor"
+authors = { workspace = true }
+categories = { workspace = true }
+description = "Vortex tensor extension type"
+edition = { workspace = true }
+homepage = { workspace = true }
+include = { workspace = true }
+keywords = { workspace = true }
+license = { workspace = true }
+readme = { workspace = true }
+repository = { workspace = true }
+rust-version = { workspace = true }
+version = { workspace = true }
+
+[lints]
+workspace = true
+
+[dependencies]
+vortex = { workspace = true }
+
+itertools = { workspace = true }
+num-traits = { workspace = true }
+prost = { workspace = true }
+
+[dev-dependencies]
+rstest = { workspace = true }

vortex-tensor/Cargo.toml

@@ -0,0 +1,121 @@
+pub mod vortex_tensor
+
+pub mod vortex_tensor::scalar_fns
+
+pub mod vortex_tensor::scalar_fns::cosine_similarity
+
+pub struct vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity
+
+impl core::clone::Clone for vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::clone(&self) -> vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity
+
+impl vortex_array::scalar_fn::vtable::ScalarFnVTable for vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity
+
+pub type vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::Options = vortex_array::scalar_fn::vtable::EmptyOptions
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::arity(&self, _options: &Self::Options) -> vortex_array::scalar_fn::vtable::Arity
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::child_name(&self, _options: &Self::Options, child_idx: usize) -> vortex_array::scalar_fn::vtable::ChildName
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::execute(&self, _options: &Self::Options, args: &dyn vortex_array::scalar_fn::vtable::ExecutionArgs, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::fmt_sql(&self, _options: &Self::Options, expr: &vortex_array::expr::expression::Expression, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::id(&self) -> vortex_array::scalar_fn::ScalarFnId
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::is_fallible(&self, _options: &Self::Options) -> bool
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::is_null_sensitive(&self, _options: &Self::Options) -> bool
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::return_dtype(&self, _options: &Self::Options, arg_dtypes: &[vortex_array::dtype::DType]) -> vortex_error::VortexResult<vortex_array::dtype::DType>
+
+pub fn vortex_tensor::scalar_fns::cosine_similarity::CosineSimilarity::validity(&self, _options: &Self::Options, expression: &vortex_array::expr::expression::Expression) -> vortex_error::VortexResult<core::option::Option<vortex_array::expr::expression::Expression>>
+
+pub struct vortex_tensor::FixedShapeTensor
+
+impl core::clone::Clone for vortex_tensor::FixedShapeTensor
+
+pub fn vortex_tensor::FixedShapeTensor::clone(&self) -> vortex_tensor::FixedShapeTensor
+
+impl core::cmp::Eq for vortex_tensor::FixedShapeTensor
+
+impl core::cmp::PartialEq for vortex_tensor::FixedShapeTensor
+
+pub fn vortex_tensor::FixedShapeTensor::eq(&self, other: &vortex_tensor::FixedShapeTensor) -> bool
+
+impl core::default::Default for vortex_tensor::FixedShapeTensor
+
+pub fn vortex_tensor::FixedShapeTensor::default() -> vortex_tensor::FixedShapeTensor
+
+impl core::fmt::Debug for vortex_tensor::FixedShapeTensor
+
+pub fn vortex_tensor::FixedShapeTensor::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::hash::Hash for vortex_tensor::FixedShapeTensor
+
+pub fn vortex_tensor::FixedShapeTensor::hash<__H: core::hash::Hasher>(&self, state: &mut __H)
+
+impl core::marker::StructuralPartialEq for vortex_tensor::FixedShapeTensor
+
+impl vortex_array::dtype::extension::vtable::ExtVTable for vortex_tensor::FixedShapeTensor
+
+pub type vortex_tensor::FixedShapeTensor::Metadata = vortex_tensor::FixedShapeTensorMetadata
+
+pub type vortex_tensor::FixedShapeTensor::NativeValue<'a> = &'a vortex_array::scalar::scalar_value::ScalarValue
+
+pub fn vortex_tensor::FixedShapeTensor::deserialize_metadata(&self, metadata: &[u8]) -> vortex_error::VortexResult<Self::Metadata>
+
+pub fn vortex_tensor::FixedShapeTensor::id(&self) -> vortex_array::dtype::extension::ExtId
+
+pub fn vortex_tensor::FixedShapeTensor::serialize_metadata(&self, metadata: &Self::Metadata) -> vortex_error::VortexResult<alloc::vec::Vec<u8>>
+
+pub fn vortex_tensor::FixedShapeTensor::unpack_native<'a>(&self, _metadata: &'a Self::Metadata, _storage_dtype: &'a vortex_array::dtype::DType, storage_value: &'a vortex_array::scalar::scalar_value::ScalarValue) -> vortex_error::VortexResult<Self::NativeValue>
+
+pub fn vortex_tensor::FixedShapeTensor::validate_dtype(&self, metadata: &Self::Metadata, storage_dtype: &vortex_array::dtype::DType) -> vortex_error::VortexResult<()>
+
+pub struct vortex_tensor::FixedShapeTensorMetadata
+
+impl vortex_tensor::FixedShapeTensorMetadata
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::dim_names(&self) -> core::option::Option<&[alloc::string::String]>
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::logical_shape(&self) -> &[usize]
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::ndim(&self) -> usize
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::new(shape: alloc::vec::Vec<usize>) -> Self
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::permutation(&self) -> core::option::Option<&[usize]>
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::physical_shape(&self) -> impl core::iter::traits::iterator::Iterator<Item = usize> + '_
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::strides(&self) -> impl core::iter::traits::iterator::Iterator<Item = usize> + '_
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::with_dim_names(self, names: alloc::vec::Vec<alloc::string::String>) -> vortex_error::VortexResult<Self>
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::with_permutation(self, permutation: alloc::vec::Vec<usize>) -> vortex_error::VortexResult<Self>
+
+impl core::clone::Clone for vortex_tensor::FixedShapeTensorMetadata
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::clone(&self) -> vortex_tensor::FixedShapeTensorMetadata
+
+impl core::cmp::Eq for vortex_tensor::FixedShapeTensorMetadata
+
+impl core::cmp::PartialEq for vortex_tensor::FixedShapeTensorMetadata
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::eq(&self, other: &vortex_tensor::FixedShapeTensorMetadata) -> bool
+
+impl core::fmt::Debug for vortex_tensor::FixedShapeTensorMetadata
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::fmt::Display for vortex_tensor::FixedShapeTensorMetadata
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
+impl core::hash::Hash for vortex_tensor::FixedShapeTensorMetadata
+
+pub fn vortex_tensor::FixedShapeTensorMetadata::hash<__H: core::hash::Hasher>(&self, state: &mut __H)
+
+impl core::marker::StructuralPartialEq for vortex_tensor::FixedShapeTensorMetadata

vortex-tensor/public-api.lock

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Tensor extension type.
+
+mod metadata;
+pub use metadata::FixedShapeTensorMetadata;
+
+mod proto;
+mod vtable;
+
+pub mod scalar_fns;
+
+/// The VTable for the Tensor extension type.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
+pub struct FixedShapeTensor;