vortex_tensor/vector/
vtable.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex::dtype::DType;
5use vortex::dtype::extension::ExtDType;
6use vortex::dtype::extension::ExtId;
7use vortex::dtype::extension::ExtVTable;
8use vortex::error::VortexResult;
9use vortex::error::vortex_bail;
10use vortex::error::vortex_ensure;
11use vortex::extension::EmptyMetadata;
12use vortex::scalar::ScalarValue;
13
14use crate::vector::Vector;
15
16impl ExtVTable for Vector {
17    type Metadata = EmptyMetadata;
18
19    // TODO(connor): This is just a placeholder for now.
20    type NativeValue<'a> = &'a ScalarValue;
21
22    fn id(&self) -> ExtId {
23        ExtId::new_ref("vortex.tensor.vector")
24    }
25
26    fn serialize_metadata(&self, _metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
27        Ok(Vec::new())
28    }
29
30    fn deserialize_metadata(&self, _metadata: &[u8]) -> VortexResult<Self::Metadata> {
31        Ok(EmptyMetadata)
32    }
33
34    fn validate_dtype(ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
35        let storage_dtype = ext_dtype.storage_dtype();
36        let DType::FixedSizeList(element_dtype, _list_size, _nullability) = storage_dtype else {
37            vortex_bail!("Vector storage dtype must be a FixedSizeList, got {storage_dtype}");
38        };
39
40        vortex_ensure!(
41            element_dtype.is_float(),
42            "Vector element dtype must be a float, got {element_dtype}"
43        );
44        vortex_ensure!(
45            !element_dtype.is_nullable(),
46            "Vector element dtype must be non-nullable"
47        );
48
49        Ok(())
50    }
51
52    fn unpack_native<'a>(
53        _ext_dtype: &'a ExtDType<Self>,
54        storage_value: &'a ScalarValue,
55    ) -> VortexResult<Self::NativeValue<'a>> {
56        Ok(storage_value)
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use std::sync::Arc;
63
64    use rstest::rstest;
65    use vortex::dtype::DType;
66    use vortex::dtype::Nullability;
67    use vortex::dtype::PType;
68    use vortex::dtype::extension::ExtDType;
69    use vortex::dtype::extension::ExtVTable;
70    use vortex::error::VortexResult;
71    use vortex::extension::EmptyMetadata;
72
73    use crate::vector::Vector;
74
75    /// Constructs a `FixedSizeList` storage dtype with the given float [`PType`], list size, and
76    /// [`Nullability`].
77    fn vector_storage_dtype(ptype: PType, size: u32, nullability: Nullability) -> DType {
78        DType::FixedSizeList(
79            Arc::new(DType::Primitive(ptype, Nullability::NonNullable)),
80            size,
81            nullability,
82        )
83    }
84
85    #[rstest]
86    #[case::f16(PType::F16)]
87    #[case::f32(PType::F32)]
88    #[case::f64(PType::F64)]
89    fn validate_accepts_float_types(#[case] ptype: PType) -> VortexResult<()> {
90        let storage = vector_storage_dtype(ptype, 128, Nullability::NonNullable);
91        ExtDType::<Vector>::try_new(EmptyMetadata, storage)?;
92        Ok(())
93    }
94
95    #[rstest]
96    #[case::nullable(Nullability::Nullable)]
97    #[case::non_nullable(Nullability::NonNullable)]
98    fn validate_accepts_any_outer_nullability(
99        #[case] nullability: Nullability,
100    ) -> VortexResult<()> {
101        let storage = vector_storage_dtype(PType::F32, 128, nullability);
102        ExtDType::<Vector>::try_new(EmptyMetadata, storage)?;
103        Ok(())
104    }
105
106    #[test]
107    fn validate_rejects_non_fsl() {
108        let storage = DType::Primitive(PType::F32, Nullability::NonNullable);
109        assert!(ExtDType::<Vector>::try_new(EmptyMetadata, storage).is_err());
110    }
111
112    #[test]
113    fn validate_rejects_integer_elements() {
114        let storage = DType::FixedSizeList(
115            Arc::new(DType::Primitive(PType::U32, Nullability::NonNullable)),
116            128,
117            Nullability::NonNullable,
118        );
119        assert!(ExtDType::<Vector>::try_new(EmptyMetadata, storage).is_err());
120    }
121
122    #[test]
123    fn validate_rejects_nullable_elements() {
124        let storage = DType::FixedSizeList(
125            Arc::new(DType::Primitive(PType::F32, Nullability::Nullable)),
126            128,
127            Nullability::NonNullable,
128        );
129        assert!(ExtDType::<Vector>::try_new(EmptyMetadata, storage).is_err());
130    }
131
132    #[test]
133    fn roundtrip_metadata() -> VortexResult<()> {
134        let vtable = Vector;
135        let bytes = vtable.serialize_metadata(&EmptyMetadata)?;
136        assert!(bytes.is_empty());
137        let deserialized = vtable.deserialize_metadata(&bytes)?;
138        assert_eq!(deserialized, EmptyMetadata);
139        Ok(())
140    }
141}
vortex_tensor/vector/vtable.rs

vortex_tensor/vector/
vtable.rs
