use a single trait for extension types

a10y · a10y · commit a2139a59c431 · 2025-05-07T17:25:25.000-04:00
diff --git a/extensions/geo/src/array.rs b/extensions/geo/src/array.rs
@@ -51,7 +51,7 @@ impl PointArray {
             vortex_bail!("points must be Struct typed, was {}", points.dtype())
         };
 
-        validate_coord_schema(&schema, metadata.dimension)?;
+        validate_coord_schema(schema, metadata.dimension)?;
         let point_type =
             OwnedGeometryType::Point(metadata.clone()).into_ext_dtype(points.dtype().nullability());
         let inner = ExtensionArray::new(Arc::new(point_type), points);
@@ -153,7 +153,7 @@ mod tests {
         );
         assert_eq!(ext.id(), &*POINT_ID);
     }
-    
+
     #[test]
     fn test_polygon() {
         // Create two lists of points: one defining an interior and exterior.
diff --git a/extensions/geo/src/arrow.rs b/extensions/geo/src/arrow.rs
@@ -19,7 +19,7 @@ use vortex_array::arrow::ArrowArray;
 use vortex_array::arrow::compute::{ToArrowArgs, ToArrowKernelRef};
 use vortex_array::compute::{InvocationArgs, Kernel, Output, cast};
 use vortex_array::{Array, ToCanonical, register_kernel};
-use vortex_dtype::arrow::{ArrowMetadata, ArrowMetadataRef, ArrowToDType, ArrowToDTypeRef};
+use vortex_dtype::arrow::{TypeConversion, TypeConversionRef};
 use vortex_dtype::{DType, ExtDType, PType, register_extension_type};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 
@@ -68,7 +68,7 @@ fn to_arrow_point(
     point_array: &ExtensionArray,
     metadata: &GeoMetadata,
 ) -> VortexResult<arrow_array::ArrayRef> {
-    let coords = coordinate_buffer(point_array.storage(), &metadata)?;
+    let coords = coordinate_buffer(point_array.storage(), metadata)?;
     let nulls = point_array.validity_mask()?.to_null_buffer();
     let crs = match metadata.crs {
         None => Crs::default(),
@@ -88,7 +88,7 @@ fn to_arrow_linestring(
     metadata: &GeoMetadata,
 ) -> VortexResult<arrow_array::ArrayRef> {
     let list = line_string.to_list()?;
-    let coords = coordinate_buffer(list.elements().as_ref(), &metadata)?;
+    let coords = coordinate_buffer(list.elements().as_ref(), metadata)?;
     let offsets = cast(&list.offsets().to_primitive()?, &PType::I32.into())?
         .to_primitive()?
         .into_buffer::<i32>()
@@ -114,26 +114,24 @@ fn to_arrow_linestring(
 pub struct GeoArrowConversion;
 
 // Conversion between Vortex geospatial extension types and GeoArrow extension types.
-impl ArrowToDType for GeoArrowConversion {
-    fn can_convert(&self, field: &Field) -> bool {
-        // The field needs to support all of our types
-        if let Some(ext_type) = field.extension_type_name() {
-            if ext_type == PointType::NAME
-                || ext_type == LineStringType::NAME
-                || ext_type == PolygonType::NAME
-                || ext_type == WkbType::NAME
-            {
-                return true;
-            }
-        }
-        false
-    }
-
-    fn to_vortex(&self, field: &Field) -> VortexResult<DType> {
+impl TypeConversion for GeoArrowConversion {
+    fn to_vortex(&self, field: &Field) -> VortexResult<Option<DType>> {
+        // Validate that the field is one of the supported geospatial
+        // extension types.
         let Some(ext_type) = field.extension_type_name() else {
-            vortex_bail!("field does not have an extension type")
+            return Ok(None);
         };
 
+        println!("ext_type = {}", ext_type);
+
+        if ext_type != PointType::NAME
+            && ext_type != LineStringType::NAME
+            && ext_type != PolygonType::NAME
+            && ext_type != WkbType::NAME
+        {
+            return Ok(None);
+        }
+
         macro_rules! dim_from_fields {
             ($fields:expr) => {{
                 match $fields.len() {
@@ -232,13 +230,15 @@ impl ArrowToDType for GeoArrowConversion {
             _ => vortex_bail!("extension type {} not supported", ext_type),
         };
 
-        Ok(DType::Extension(Arc::new(ext_dtype)))
+        Ok(Some(DType::Extension(Arc::new(ext_dtype))))
     }
-}
 
-impl ArrowMetadata for GeoArrowConversion {
     #[allow(clippy::disallowed_types)]
-    fn arrow_metadata(&self, vortex_extension_type: &ExtDType) -> Option<HashMap<String, String>> {
+    fn to_arrow(
+        &self,
+        vortex_extension_type: &ExtDType,
+        field: &mut Field,
+    ) -> VortexResult<Option<()>> {
         if let Ok(geometry) = GeometryType::try_from(vortex_extension_type) {
             let mut extension_metadata = HashMap::new();
             let ext_type_name = match geometry {
@@ -249,7 +249,7 @@ impl ArrowMetadata for GeoArrowConversion {
             };
             extension_metadata.insert(EXTENSION_TYPE_NAME_KEY.to_string(), ext_type_name);
 
-            match geometry {
+            let extension_metadata = match geometry {
                 GeometryType::Point(meta)
                 | GeometryType::LineString(meta)
                 | GeometryType::Polygon(meta)
@@ -264,19 +264,23 @@ impl ArrowMetadata for GeoArrowConversion {
                                 .vortex_expect("failed to serialize geoarrow metadata");
                         extension_metadata
                             .insert(EXTENSION_TYPE_METADATA_KEY.to_string(), metadata_json);
+                        extension_metadata
+                    } else {
+                        extension_metadata
                     }
                 }
             };
 
-            Some(extension_metadata)
+            field.set_metadata(extension_metadata);
+
+            Ok(Some(()))
         } else {
-            None
+            Ok(None)
         }
     }
 }
 
-register_extension_type!(ArrowToDTypeRef(ArcRef::new_ref(&GeoArrowConversion)));
-register_extension_type!(ArrowMetadataRef(ArcRef::new_ref(&GeoArrowConversion)));
+register_extension_type!(TypeConversionRef::new(ArcRef::new_ref(&GeoArrowConversion)));
 
 /// Unpack the geoarrow CoordBuffer. Errors if the dimensions specified in the metadata do not
 /// match the actual encoding.
@@ -400,7 +404,7 @@ mod tests {
             dimension: crate::Dimension::XY,
             crs: None,
         })
-        .into_ext_dtype(false.into());
+        .into_ext_dtype(true.into());
         assert_eq!(dtype, DType::Extension(Arc::new(owned_type)));
 
         // round trip back to Arrow type.
@@ -424,7 +428,4 @@ mod tests {
         let exported = imported.into_arrow_preferred().unwrap();
         assert_eq!(exported.data_type(), struct_array.data_type());
     }
-
-    #[test]
-    fn test_roundtrip() {}
 }
diff --git a/vortex-array/src/arrow/convert.rs b/vortex-array/src/arrow/convert.rs
@@ -18,7 +18,7 @@ use arrow_buffer::buffer::{NullBuffer, OffsetBuffer};
 use arrow_buffer::{ArrowNativeType, BooleanBuffer, Buffer as ArrowBuffer, ScalarBuffer};
 use arrow_schema::{DataType, TimeUnit as ArrowTimeUnit};
 use vortex_buffer::{Alignment, Buffer, ByteBuffer};
-use vortex_dtype::arrow::ArrowToDTypeRef;
+use vortex_dtype::arrow::geo_field_to_dtype;
 use vortex_dtype::datetime::TimeUnit;
 use vortex_dtype::{DType, DecimalDType, FieldNames, NativePType, PType};
 use vortex_error::{VortexExpect as _, vortex_panic};
@@ -243,7 +243,7 @@ impl FromArrowArray<&ArrowStructArray> for ArrayRef {
         for (field, column) in value.fields().iter().zip(value.columns()) {
             if field.extension_type_name().is_some() {
                 // Check if we have a conversion function. Wrap using the ExtensionArray.
-                if let Some(Ok(DType::Extension(ext_dtype))) = ArrowToDTypeRef::convert(field) {
+                if let Ok(Some(DType::Extension(ext_dtype))) = geo_field_to_dtype(field) {
                     let storage = Self::from_arrow(column.clone(), field.is_nullable());
                     field_names.push(field.name().to_string().into());
                     field_arrays.push(ExtensionArray::new(ext_dtype, storage).into_array());
diff --git a/vortex-dtype/src/arrow.rs b/vortex-dtype/src/arrow.rs
@@ -10,8 +10,6 @@
 //! For this reason, it's recommended to do as much computation as possible within Vortex, and then
 //! materialize an Arrow ArrayRef at the very end of the processing chain.
 
-#[allow(clippy::disallowed_types)]
-use std::collections::HashMap;
 use std::sync::Arc;
 
 use arcref::ArcRef;
@@ -137,12 +135,13 @@ impl FromArrowType<&Field> for DType {
             Some(ext_type) => {
                 // Check the registry for any Vortex extension types that represent the Arrow
                 // extension type named here.
-                for converter in inventory::iter::<ArrowToDTypeRef> {
-                    if converter.0.can_convert(field) {
-                        return converter
-                            .0
-                            .to_vortex(field)
-                            .vortex_expect("arrow extension type to vortex dtype");
+                for converter in inventory::iter::<TypeConversionRef> {
+                    if let Some(converted) = converter
+                        .0
+                        .to_vortex(field)
+                        .vortex_expect("Conversion from GeoArrow type to GeoVortex")
+                    {
+                        return converted;
                     }
                 }
 
@@ -210,9 +209,8 @@ impl DType {
                     // Optionally: attach any Arrow extension type metadata, if an ArrowMetadata
                     // kernel is defined for the extension type.
                     if let DType::Extension(ext_type) = field_dt {
-                        for kernel in inventory::iter::<ArrowMetadataRef>() {
-                            if let Some(datum) = kernel.0.arrow_metadata(ext_type.as_ref()) {
-                                field = field.with_metadata(datum);
+                        for kernel in inventory::iter::<TypeConversionRef> {
+                            if kernel.0.to_arrow(ext_type.as_ref(), &mut field)?.is_some() {
                                 break;
                             }
                         }
@@ -242,60 +240,60 @@ impl DType {
     }
 }
 
-/// Type-erased pointer to a [`ArrowToDType`] implementation.
-pub struct ArrowToDTypeRef(pub ArcRef<dyn ArrowToDType>);
-inventory::collect!(ArrowToDTypeRef);
-
-impl ArrowToDTypeRef {
-    /// Attempt to convert a field to a DType. If the there is no registered converter that
-    /// can handle the field type, `None` is returned.
-    ///
-    /// If a converter is resolved, it is used to convert the Field and the result is returned in
-    /// a `Some`.
-    pub fn convert(field: impl AsRef<Field>) -> Option<VortexResult<DType>> {
-        for converter in inventory::iter::<ArrowToDTypeRef> {
-            if converter.0.can_convert(field.as_ref()) {
-                return Some(converter.0.to_vortex(field.as_ref()));
-            }
+/// Attempt to convert a field to a DType. If the there is no registered converter that
+/// can handle the field type, `None` is returned.
+///
+/// If a converter is resolved, it is used to convert the Field and the result is returned in
+/// a `Some`.
+pub fn geo_field_to_dtype(field: impl AsRef<Field>) -> VortexResult<Option<DType>> {
+    for converter in inventory::iter::<TypeConversionRef> {
+        if let Some(converted) = converter.0.to_vortex(field.as_ref())? {
+            return Ok(Some(converted));
         }
-
-        None
     }
-}
 
-/// Get Arrow extension type metadata for a type.
-pub trait ArrowMetadata: 'static + Send + Sync {
-    /// Optionally get the Arrow metadata for this type. None indicates it is not an Arrow
-    /// extension type, Some indicates that it is.
-    #[allow(clippy::disallowed_types)]
-    fn arrow_metadata(&self, vortex_extension_type: &ExtDType) -> Option<HashMap<String, String>>;
+    Ok(None)
 }
 
-/// Type-erased pointer to a thing that can implement the DType conversion instead.
-pub struct ArrowMetadataRef(pub ArcRef<dyn ArrowMetadata>);
-inventory::collect!(ArrowMetadataRef);
-
-/// Conversion for extension types.
+/// Conversions between Arrow [`Field`] type and Vortex logical `DType`.
 ///
-/// If we have custom conversions we can register them via a plugin. This is something that is
-/// determined elsewhere however.
-pub trait ArrowToDType: Send + Sync {
-    /// If this returns `true`, the implementor is able to convert the given Vortex [`DType`] to
-    /// an Arrow [`Field`]. The caller can then call the `to_vortex` method with the argument.
-    fn can_convert(&self, data_type: &Field) -> bool;
-
+/// This crate provides infra for registering and discovering extension types.
+/// Once you implement this trait, you need to register it using [`register_extension_type`]
+/// to have it get discovered.
+///
+/// See also: [`register_extension_type`]
+pub trait TypeConversion: 'static + Send + Sync {
     /// Convert the given Arrow [`Field`] to a Vortex [`DType`].
-    fn to_vortex(&self, field: &Field) -> VortexResult<DType>;
+    fn to_vortex(&self, _field: &Field) -> VortexResult<Option<DType>> {
+        Ok(None)
+    }
+
+    /// Convert a Vortex [`ExtDType`] to an Arrow schema [`Field`].
+    ///
+    /// The returned field will have the type, nullability, and metadata
+    /// that should be propagated forward, but its name may be changed.
+    fn to_arrow(&self, _dtype: &ExtDType, _field_builder: &mut Field) -> VortexResult<Option<()>> {
+        Ok(None)
+    }
 }
 
-/// Register an extension type globally. This should be a type that implements
-/// the [`ArrowToDType`] trait.
+/// Conversion token
+pub struct TypeConversionRef(ArcRef<dyn TypeConversion>);
+inventory::collect!(TypeConversionRef);
+
+impl TypeConversionRef {
+    /// Create a new `TypeConversionRef` from a pointer to the implementation.
+    pub const fn new(conversion: ArcRef<dyn TypeConversion>) -> Self {
+        Self(conversion)
+    }
+}
+
+/// Register an extension type for lookup.
 #[macro_export]
 macro_rules! register_extension_type {
     ($extension:expr) => {
-        $crate::inventory::submit! {
-            $extension
-        }
+        const _: $crate::arrow::TypeConversionRef = $extension;
+        $crate::inventory::submit! { $extension }
     };
 }
 
