vector type first draft

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

Author: connortsui20

vortex-tensor/src/lib.rs

@@ -6,5 +6,6 @@
 //! similarity.
 
 pub mod fixed_shape;
+pub mod vector;
 
 pub mod scalar_fns;

vortex-tensor/src/vector/mod.rs

@@ -0,0 +1,10 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Vector extension type for fixed-length float vectors (e.g., embeddings).
+
+/// The VTable for the vector extension type.
+#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
+pub struct Vector;
+
+mod vtable;

vortex-tensor/src/vector/vtable.rs

@@ -0,0 +1,153 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex::dtype::DType;
+use vortex::dtype::extension::ExtDType;
+use vortex::dtype::extension::ExtId;
+use vortex::dtype::extension::ExtVTable;
+use vortex::error::VortexResult;
+use vortex::error::vortex_bail;
+use vortex::error::vortex_ensure;
+use vortex::extension::EmptyMetadata;
+use vortex::scalar::ScalarValue;
+
+use crate::vector::Vector;
+
+impl ExtVTable for Vector {
+    type Metadata = EmptyMetadata;
+
+    // TODO(connor): This is just a placeholder for now.
+    type NativeValue<'a> = &'a ScalarValue;
+
+    fn id(&self) -> ExtId {
+        ExtId::new_ref("vortex.vector")
+    }
+
+    fn serialize_metadata(&self, _metadata: &Self::Metadata) -> VortexResult<Vec<u8>> {
+        Ok(Vec::new())
+    }
+
+    fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult<Self::Metadata> {
+        vortex_ensure!(
+            metadata.is_empty(),
+            "Vector metadata must be empty, got {} bytes",
+            metadata.len()
+        );
+        Ok(EmptyMetadata)
+    }
+
+    fn validate_dtype(&self, ext_dtype: &ExtDType<Self>) -> VortexResult<()> {
+        let storage_dtype = ext_dtype.storage_dtype();
+        let DType::FixedSizeList(element_dtype, _list_size, _nullability) = storage_dtype else {
+            vortex_bail!("Vector storage dtype must be a FixedSizeList, got {storage_dtype}");
+        };
+
+        vortex_ensure!(
+            element_dtype.is_float(),
+            "Vector element dtype must be a float, got {element_dtype}"
+        );
+        vortex_ensure!(
+            !element_dtype.is_nullable(),
+            "Vector element dtype must be non-nullable"
+        );
+
+        Ok(())
+    }
+
+    fn unpack_native<'a>(
+        &self,
+        _ext_dtype: &'a ExtDType<Self>,
+        storage_value: &'a ScalarValue,
+    ) -> VortexResult<Self::NativeValue<'a>> {
+        Ok(storage_value)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use rstest::rstest;
+    use vortex::dtype::DType;
+    use vortex::dtype::Nullability;
+    use vortex::dtype::PType;
+    use vortex::dtype::extension::ExtDType;
+    use vortex::dtype::extension::ExtVTable;
+    use vortex::error::VortexResult;
+    use vortex::extension::EmptyMetadata;
+
+    use crate::vector::Vector;
+
+    /// Constructs a `FixedSizeList` storage dtype with the given float [`PType`], list size, and
+    /// [`Nullability`].
+    fn vector_storage_dtype(ptype: PType, size: u32, nullability: Nullability) -> DType {
+        DType::FixedSizeList(
+            Arc::new(DType::Primitive(ptype, Nullability::NonNullable)),
+            size,
+            nullability,
+        )
+    }
+
+    #[rstest]
+    #[case::f16(PType::F16)]
+    #[case::f32(PType::F32)]
+    #[case::f64(PType::F64)]
+    fn validate_accepts_float_types(#[case] ptype: PType) -> VortexResult<()> {
+        let storage = vector_storage_dtype(ptype, 128, Nullability::NonNullable);
+        ExtDType::<Vector>::try_new(EmptyMetadata, storage)?;
+        Ok(())
+    }
+
+    #[rstest]
+    #[case::nullable(Nullability::Nullable)]
+    #[case::non_nullable(Nullability::NonNullable)]
+    fn validate_accepts_any_outer_nullability(
+        #[case] nullability: Nullability,
+    ) -> VortexResult<()> {
+        let storage = vector_storage_dtype(PType::F32, 128, nullability);
+        ExtDType::<Vector>::try_new(EmptyMetadata, storage)?;
+        Ok(())
+    }
+
+    #[test]
+    fn validate_rejects_non_fsl() {
+        let storage = DType::Primitive(PType::F32, Nullability::NonNullable);
+        assert!(ExtDType::<Vector>::try_new(EmptyMetadata, storage).is_err());
+    }
+
+    #[test]
+    fn validate_rejects_integer_elements() {
+        let storage = DType::FixedSizeList(
+            Arc::new(DType::Primitive(PType::U32, Nullability::NonNullable)),
+            128,
+            Nullability::NonNullable,
+        );
+        assert!(ExtDType::<Vector>::try_new(EmptyMetadata, storage).is_err());
+    }
+
+    #[test]
+    fn validate_rejects_nullable_elements() {
+        let storage = DType::FixedSizeList(
+            Arc::new(DType::Primitive(PType::F32, Nullability::Nullable)),
+            128,
+            Nullability::NonNullable,
+        );
+        assert!(ExtDType::<Vector>::try_new(EmptyMetadata, storage).is_err());
+    }
+
+    #[test]
+    fn roundtrip_metadata() -> VortexResult<()> {
+        let vtable = Vector;
+        let bytes = vtable.serialize_metadata(&EmptyMetadata)?;
+        assert!(bytes.is_empty());
+        let deserialized = vtable.deserialize_metadata(&bytes)?;
+        assert_eq!(deserialized, EmptyMetadata);
+        Ok(())
+    }
+
+    #[test]
+    fn deserialize_rejects_non_empty_bytes() {
+        let vtable = Vector;
+        assert!(vtable.deserialize_metadata(&[0x01]).is_err());
+    }
+}