From 3f6f3420ea1d6f1b29f8d5743e457c2245005696 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 20 Apr 2026 14:03:10 -0400 Subject: [PATCH] move tensor types under `types/` submodule Signed-off-by: Connor Tsui --- vortex-tensor/src/encodings/turboquant/compress.rs | 4 ++-- vortex-tensor/src/encodings/turboquant/mod.rs | 4 ++-- vortex-tensor/src/encodings/turboquant/tests/mod.rs | 2 +- .../src/encodings/turboquant/tests/structural.rs | 2 +- vortex-tensor/src/lib.rs | 10 ++++++---- vortex-tensor/src/matcher.rs | 8 ++++---- vortex-tensor/src/scalar_fns/cosine_similarity.rs | 2 +- vortex-tensor/src/scalar_fns/inner_product.rs | 4 ++-- vortex-tensor/src/scalar_fns/l2_denorm.rs | 2 +- vortex-tensor/src/scalar_fns/l2_norm.rs | 2 +- vortex-tensor/src/scalar_fns/sorf_transform/tests.rs | 7 +++++-- vortex-tensor/src/scalar_fns/sorf_transform/vtable.rs | 4 ++-- vortex-tensor/src/{ => types}/fixed_shape/matcher.rs | 6 +++--- vortex-tensor/src/{ => types}/fixed_shape/metadata.rs | 0 vortex-tensor/src/{ => types}/fixed_shape/mod.rs | 0 vortex-tensor/src/{ => types}/fixed_shape/proto.rs | 2 +- vortex-tensor/src/{ => types}/fixed_shape/vtable.rs | 10 +++++----- vortex-tensor/src/types/mod.rs | 7 +++++++ vortex-tensor/src/{ => types}/vector/matcher.rs | 6 +++--- vortex-tensor/src/{ => types}/vector/mod.rs | 0 vortex-tensor/src/{ => types}/vector/vtable.rs | 4 ++-- vortex-tensor/src/utils.rs | 6 +++--- vortex-tensor/src/vector_search.rs | 2 +- 23 files changed, 53 insertions(+), 41 deletions(-) rename vortex-tensor/src/{ => types}/fixed_shape/matcher.rs (96%) rename vortex-tensor/src/{ => types}/fixed_shape/metadata.rs (100%) rename vortex-tensor/src/{ => types}/fixed_shape/mod.rs (100%) rename vortex-tensor/src/{ => types}/fixed_shape/proto.rs (98%) rename vortex-tensor/src/{ => types}/fixed_shape/vtable.rs (94%) create mode 100644 vortex-tensor/src/types/mod.rs rename vortex-tensor/src/{ => types}/vector/matcher.rs (96%) rename vortex-tensor/src/{ => types}/vector/mod.rs (100%) rename vortex-tensor/src/{ => types}/vector/vtable.rs (98%) diff --git a/vortex-tensor/src/encodings/turboquant/compress.rs b/vortex-tensor/src/encodings/turboquant/compress.rs index 40e24807091..58c5666525a 100644 --- a/vortex-tensor/src/encodings/turboquant/compress.rs +++ b/vortex-tensor/src/encodings/turboquant/compress.rs @@ -39,9 +39,9 @@ use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm; use crate::scalar_fns::sorf_transform::SorfMatrix; use crate::scalar_fns::sorf_transform::SorfOptions; use crate::scalar_fns::sorf_transform::SorfTransform; +use crate::types::vector::AnyVector; +use crate::types::vector::Vector; use crate::utils::cast_to_f32; -use crate::vector::AnyVector; -use crate::vector::Vector; /// Configuration for TurboQuant encoding. #[derive(Clone, Debug)] diff --git a/vortex-tensor/src/encodings/turboquant/mod.rs b/vortex-tensor/src/encodings/turboquant/mod.rs index 51480955594..463e8bc7b25 100644 --- a/vortex-tensor/src/encodings/turboquant/mod.rs +++ b/vortex-tensor/src/encodings/turboquant/mod.rs @@ -154,8 +154,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; -use crate::vector::AnyVector; -use crate::vector::VectorMatcherMetadata; +use crate::types::vector::AnyVector; +use crate::types::vector::VectorMatcherMetadata; /// Validates that `dtype` is a [`Vector`](crate::vector::Vector) extension type with /// dimension >= [`MIN_DIMENSION`]. diff --git a/vortex-tensor/src/encodings/turboquant/tests/mod.rs b/vortex-tensor/src/encodings/turboquant/tests/mod.rs index 5b7e325b9b7..7dfe53c444a 100644 --- a/vortex-tensor/src/encodings/turboquant/tests/mod.rs +++ b/vortex-tensor/src/encodings/turboquant/tests/mod.rs @@ -32,7 +32,7 @@ use vortex_error::VortexResult; use crate::encodings::turboquant::TurboQuantConfig; use crate::encodings::turboquant::turboquant_encode; use crate::tests::SESSION; -use crate::vector::Vector; +use crate::types::vector::Vector; /// Create a FixedSizeListArray of random f32 vectors with the given validity. fn make_fsl_with_validity( diff --git a/vortex-tensor/src/encodings/turboquant/tests/structural.rs b/vortex-tensor/src/encodings/turboquant/tests/structural.rs index 7cd2cdfcf66..193d340285a 100644 --- a/vortex-tensor/src/encodings/turboquant/tests/structural.rs +++ b/vortex-tensor/src/encodings/turboquant/tests/structural.rs @@ -236,7 +236,7 @@ fn sorf_transform_roundtrip_isolation() -> VortexResult<()> { use crate::scalar_fns::sorf_transform::SorfMatrix; use crate::scalar_fns::sorf_transform::SorfOptions; use crate::scalar_fns::sorf_transform::SorfTransform; - use crate::vector::Vector; + use crate::types::vector::Vector; let dim = 128usize; let seed = 99u64; diff --git a/vortex-tensor/src/lib.rs b/vortex-tensor/src/lib.rs index 68269407473..c748bdd9f43 100644 --- a/vortex-tensor/src/lib.rs +++ b/vortex-tensor/src/lib.rs @@ -11,19 +11,21 @@ use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; -use crate::fixed_shape::FixedShapeTensor; use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::scalar_fns::inner_product::InnerProduct; use crate::scalar_fns::l2_denorm::L2Denorm; use crate::scalar_fns::l2_norm::L2Norm; use crate::scalar_fns::sorf_transform::SorfTransform; -use crate::vector::Vector; +use crate::types::fixed_shape::FixedShapeTensor; +use crate::types::vector::Vector; pub mod matcher; pub mod scalar_fns; -pub mod fixed_shape; -pub mod vector; +mod types; + +pub use types::fixed_shape; +pub use types::vector; pub mod encodings; diff --git a/vortex-tensor/src/matcher.rs b/vortex-tensor/src/matcher.rs index 9ff58d17c4f..4566dcb3a38 100644 --- a/vortex-tensor/src/matcher.rs +++ b/vortex-tensor/src/matcher.rs @@ -7,10 +7,10 @@ use vortex_array::dtype::PType; use vortex_array::dtype::extension::ExtDTypeRef; use vortex_array::dtype::extension::Matcher; -use crate::fixed_shape::AnyFixedShapeTensor; -use crate::fixed_shape::FixedShapeTensorMatcherMetadata; -use crate::vector::AnyVector; -use crate::vector::VectorMatcherMetadata; +use crate::types::fixed_shape::AnyFixedShapeTensor; +use crate::types::fixed_shape::FixedShapeTensorMatcherMetadata; +use crate::types::vector::AnyVector; +use crate::types::vector::VectorMatcherMetadata; /// Matcher for any tensor-like extension type. /// diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 85d16236c8c..9f1cd39fa1a 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -327,12 +327,12 @@ mod tests { use crate::scalar_fns::cosine_similarity::CosineSimilarity; use crate::scalar_fns::l2_denorm::L2Denorm; use crate::tests::SESSION; + use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::constant_tensor_array; use crate::utils::test_helpers::l2_denorm_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; - use crate::vector::Vector; /// Evaluates cosine similarity between two tensor arrays and returns the result as `Vec`. fn eval_cosine_similarity(lhs: ArrayRef, rhs: ArrayRef, len: usize) -> VortexResult> { diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index 22b28f9e4d5..f001c74498e 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -55,11 +55,11 @@ use crate::matcher::AnyTensor; use crate::scalar_fns::l2_denorm::DenormOrientation; use crate::scalar_fns::sorf_transform::SorfMatrix; use crate::scalar_fns::sorf_transform::SorfTransform; +use crate::types::vector::Vector; use crate::utils::extract_constant_flat_row; use crate::utils::extract_flat_elements; use crate::utils::extract_l2_denorm_children; use crate::utils::validate_binary_tensor_float_inputs; -use crate::vector::Vector; /// Inner product (dot product) between two columns. /// @@ -906,10 +906,10 @@ mod tests { use crate::scalar_fns::sorf_transform::SorfOptions; use crate::scalar_fns::sorf_transform::SorfTransform; use crate::tests::SESSION; + use crate::types::vector::Vector; use crate::utils::extract_flat_elements; use crate::utils::test_helpers::literal_vector_array; use crate::utils::test_helpers::vector_array; - use crate::vector::Vector; /// Build a `Vector` whose storage is `FSL(DictArray(codes: u8, values: /// f32))`. This mirrors the shape that TurboQuant produces as the SorfTransform child. diff --git a/vortex-tensor/src/scalar_fns/l2_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs index 875d7f57e2f..8474216156a 100644 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ b/vortex-tensor/src/scalar_fns/l2_denorm.rs @@ -767,11 +767,11 @@ mod tests { use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm; use crate::scalar_fns::l2_denorm::validate_l2_normalized_rows_against_norms; use crate::tests::SESSION; + use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::constant_tensor_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; - use crate::vector::Vector; /// Evaluates L2 denorm on a tensor/vector array and returns the executed array. fn eval_l2_denorm(normalized: ArrayRef, norms: ArrayRef, len: usize) -> VortexResult { diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index 47ba68a35c4..d9341b282db 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -286,11 +286,11 @@ mod tests { use crate::scalar_fns::l2_norm::L2Norm; use crate::tests::SESSION; + use crate::types::vector::Vector; use crate::utils::test_helpers::assert_close; use crate::utils::test_helpers::literal_vector_array; use crate::utils::test_helpers::tensor_array; use crate::utils::test_helpers::vector_array; - use crate::vector::Vector; /// Evaluates L2 norm on a tensor/vector array and returns the result as `Vec`. fn eval_l2_norm(input: ArrayRef, len: usize) -> VortexResult> { diff --git a/vortex-tensor/src/scalar_fns/sorf_transform/tests.rs b/vortex-tensor/src/scalar_fns/sorf_transform/tests.rs index 64308c39562..4bc93871594 100644 --- a/vortex-tensor/src/scalar_fns/sorf_transform/tests.rs +++ b/vortex-tensor/src/scalar_fns/sorf_transform/tests.rs @@ -36,7 +36,7 @@ use crate::encodings::turboquant::centroids::compute_centroid_boundaries; use crate::encodings::turboquant::centroids::find_nearest_centroid; use crate::encodings::turboquant::centroids::get_centroids; use crate::tests::SESSION; -use crate::vector::Vector; +use crate::types::vector::Vector; /// Build a unit-normalized input vector array and forward-transform + quantize it, returning /// `(input_f32, Vector(FSL(Dict(codes, centroids))), padded_dim)`. @@ -300,7 +300,10 @@ fn return_dtype_is_vector_extension() -> VortexResult<()> { let ext = return_dtype .as_extension_opt() .expect("return dtype should be an extension type"); - assert!(ext.metadata_opt::().is_some()); + assert!( + ext.metadata_opt::() + .is_some() + ); // Inner FSL should have the original (unpadded) dimension. let DType::FixedSizeList(_, inner_dim, _) = ext.storage_dtype() else { diff --git a/vortex-tensor/src/scalar_fns/sorf_transform/vtable.rs b/vortex-tensor/src/scalar_fns/sorf_transform/vtable.rs index b54158a753f..81e17f665fa 100644 --- a/vortex-tensor/src/scalar_fns/sorf_transform/vtable.rs +++ b/vortex-tensor/src/scalar_fns/sorf_transform/vtable.rs @@ -47,8 +47,8 @@ use super::SorfOptions; use super::SorfTransform; use super::rotation::SorfMatrix; use super::validate_sorf_options; -use crate::vector::AnyVector; -use crate::vector::Vector; +use crate::types::vector::AnyVector; +use crate::types::vector::Vector; impl ScalarFnVTable for SorfTransform { type Options = SorfOptions; diff --git a/vortex-tensor/src/fixed_shape/matcher.rs b/vortex-tensor/src/types/fixed_shape/matcher.rs similarity index 96% rename from vortex-tensor/src/fixed_shape/matcher.rs rename to vortex-tensor/src/types/fixed_shape/matcher.rs index f703ccba978..32c0dd55136 100644 --- a/vortex-tensor/src/fixed_shape/matcher.rs +++ b/vortex-tensor/src/types/fixed_shape/matcher.rs @@ -8,8 +8,8 @@ use vortex_array::dtype::extension::Matcher; use vortex_error::VortexExpect; use vortex_error::vortex_panic; -use crate::fixed_shape::FixedShapeTensor; -use crate::fixed_shape::FixedShapeTensorMetadata; +use crate::types::fixed_shape::FixedShapeTensor; +use crate::types::fixed_shape::FixedShapeTensorMetadata; pub struct AnyFixedShapeTensor; @@ -98,7 +98,7 @@ mod tests { use vortex_error::VortexResult; use super::*; - use crate::vector::Vector; + use crate::types::vector::Vector; fn tensor_storage_dtype(element_ptype: PType, list_size: u32) -> DType { DType::FixedSizeList( diff --git a/vortex-tensor/src/fixed_shape/metadata.rs b/vortex-tensor/src/types/fixed_shape/metadata.rs similarity index 100% rename from vortex-tensor/src/fixed_shape/metadata.rs rename to vortex-tensor/src/types/fixed_shape/metadata.rs diff --git a/vortex-tensor/src/fixed_shape/mod.rs b/vortex-tensor/src/types/fixed_shape/mod.rs similarity index 100% rename from vortex-tensor/src/fixed_shape/mod.rs rename to vortex-tensor/src/types/fixed_shape/mod.rs diff --git a/vortex-tensor/src/fixed_shape/proto.rs b/vortex-tensor/src/types/fixed_shape/proto.rs similarity index 98% rename from vortex-tensor/src/fixed_shape/proto.rs rename to vortex-tensor/src/types/fixed_shape/proto.rs index 9d89c56bcec..89b3db4289d 100644 --- a/vortex-tensor/src/fixed_shape/proto.rs +++ b/vortex-tensor/src/types/fixed_shape/proto.rs @@ -8,7 +8,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; -use crate::fixed_shape::FixedShapeTensorMetadata; +use crate::types::fixed_shape::FixedShapeTensorMetadata; /// Protobuf representation of [`FixedShapeTensorMetadata`]. /// diff --git a/vortex-tensor/src/fixed_shape/vtable.rs b/vortex-tensor/src/types/fixed_shape/vtable.rs similarity index 94% rename from vortex-tensor/src/fixed_shape/vtable.rs rename to vortex-tensor/src/types/fixed_shape/vtable.rs index 52e32ef0cf6..7d69b9a7c44 100644 --- a/vortex-tensor/src/fixed_shape/vtable.rs +++ b/vortex-tensor/src/types/fixed_shape/vtable.rs @@ -11,9 +11,9 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use crate::fixed_shape::FixedShapeTensor; -use crate::fixed_shape::FixedShapeTensorMetadata; -use crate::fixed_shape::proto; +use crate::types::fixed_shape::FixedShapeTensor; +use crate::types::fixed_shape::FixedShapeTensorMetadata; +use crate::types::fixed_shape::proto; impl ExtVTable for FixedShapeTensor { type Metadata = FixedShapeTensorMetadata; @@ -80,8 +80,8 @@ mod tests { use vortex_array::dtype::extension::ExtVTable; use vortex_error::VortexResult; - use crate::fixed_shape::FixedShapeTensor; - use crate::fixed_shape::FixedShapeTensorMetadata; + use crate::types::fixed_shape::FixedShapeTensor; + use crate::types::fixed_shape::FixedShapeTensorMetadata; /// Serializes and deserializes the given metadata through protobuf, asserting equality. fn assert_roundtrip(metadata: &FixedShapeTensorMetadata) -> VortexResult<()> { diff --git a/vortex-tensor/src/types/mod.rs b/vortex-tensor/src/types/mod.rs new file mode 100644 index 00000000000..3ecd2826743 --- /dev/null +++ b/vortex-tensor/src/types/mod.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Internal homes for tensor extension types. + +pub mod fixed_shape; +pub mod vector; diff --git a/vortex-tensor/src/vector/matcher.rs b/vortex-tensor/src/types/vector/matcher.rs similarity index 96% rename from vortex-tensor/src/vector/matcher.rs rename to vortex-tensor/src/types/vector/matcher.rs index e4838b77426..7ac75f097db 100644 --- a/vortex-tensor/src/vector/matcher.rs +++ b/vortex-tensor/src/types/vector/matcher.rs @@ -10,7 +10,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; -use crate::vector::Vector; +use crate::types::vector::Vector; pub struct AnyVector; @@ -101,8 +101,8 @@ mod tests { use vortex_error::VortexResult; use super::*; - use crate::fixed_shape::FixedShapeTensor; - use crate::fixed_shape::FixedShapeTensorMetadata; + use crate::types::fixed_shape::FixedShapeTensor; + use crate::types::fixed_shape::FixedShapeTensorMetadata; fn vector_storage_dtype(element_ptype: PType, dimensions: u32) -> DType { DType::FixedSizeList( diff --git a/vortex-tensor/src/vector/mod.rs b/vortex-tensor/src/types/vector/mod.rs similarity index 100% rename from vortex-tensor/src/vector/mod.rs rename to vortex-tensor/src/types/vector/mod.rs diff --git a/vortex-tensor/src/vector/vtable.rs b/vortex-tensor/src/types/vector/vtable.rs similarity index 98% rename from vortex-tensor/src/vector/vtable.rs rename to vortex-tensor/src/types/vector/vtable.rs index 19fe57fed41..d870e7a8e0d 100644 --- a/vortex-tensor/src/vector/vtable.rs +++ b/vortex-tensor/src/types/vector/vtable.rs @@ -11,7 +11,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use crate::vector::Vector; +use crate::types::vector::Vector; impl ExtVTable for Vector { type Metadata = EmptyMetadata; @@ -70,7 +70,7 @@ mod tests { use vortex_array::extension::EmptyMetadata; use vortex_error::VortexResult; - use crate::vector::Vector; + use crate::types::vector::Vector; /// Constructs a `FixedSizeList` storage dtype with the given float [`PType`], list size, and /// [`Nullability`]. diff --git a/vortex-tensor/src/utils.rs b/vortex-tensor/src/utils.rs index 3434058f082..92ead274bcb 100644 --- a/vortex-tensor/src/utils.rs +++ b/vortex-tensor/src/utils.rs @@ -242,10 +242,10 @@ pub mod test_helpers { use vortex_buffer::Buffer; use vortex_error::VortexResult; - use crate::fixed_shape::FixedShapeTensor; - use crate::fixed_shape::FixedShapeTensorMetadata; use crate::scalar_fns::l2_denorm::L2Denorm; - use crate::vector::Vector; + use crate::types::fixed_shape::FixedShapeTensor; + use crate::types::fixed_shape::FixedShapeTensorMetadata; + use crate::types::vector::Vector; /// Builds a `FixedSizeList` storage array from flat `elements`. The row count is /// inferred from `elements.len() / list_size`. diff --git a/vortex-tensor/src/vector_search.rs b/vortex-tensor/src/vector_search.rs index c5551c99f2d..2a036accdfe 100644 --- a/vortex-tensor/src/vector_search.rs +++ b/vortex-tensor/src/vector_search.rs @@ -48,7 +48,7 @@ use vortex_array::scalar_fn::fns::operators::Operator; use vortex_error::VortexResult; use crate::scalar_fns::cosine_similarity::CosineSimilarity; -use crate::vector::Vector; +use crate::types::vector::Vector; /// Build the lazy similarity-search expression tree for a prepared database array and a /// single query vector.