feat: Implement mapping from GeoArrow extension types to Polars extension types

Author: kylebarron

Cargo.lock

@@ -19,6 +19,7 @@ arrow-data = { version = "56", features = ["ffi"] }
 arrow-schema = "56"
 geoarrow-array = "0.6"
 geoarrow-schema = "0.6"
+geopolars-arrow = { path = "geopolars/geopolars-arrow" }
 polars = { git = "https://github.com/pola-rs/polars", rev = "0c6dca54ae4a33b3f31191748ad6a5ffd5452822", features = [
     "dtype-extension",
 ] }

Cargo.toml

@@ -18,9 +18,10 @@ arrow-data = { workspace = true }
 arrow-schema = { workspace = true }
 geoarrow-array = { workspace = true }
 geoarrow-schema = { workspace = true }
+geopolars-arrow = { workspace = true }
+polars = { workspace = true }
 polars-arrow = { workspace = true }
 polars-core = { workspace = true }
-polars = { workspace = true }
 
 [dev-dependencies]
 geoarrow-array = { workspace = true, features = ["test-data"] }

geopolars/geopolars-extension/Cargo.toml

@@ -1,6 +1,16 @@
-use polars::prelude::DataType;
+use std::collections::HashMap;
+
+use geoarrow_schema::GeoArrowType;
+use geopolars_arrow::to_arrow::polars_field_to_arrow;
 use polars::prelude::extension::{ExtensionTypeFactory, ExtensionTypeImpl};
+use polars::prelude::{CompatLevel, DataType};
+
+use crate::geoarrow::{
+    BoxType, GeometryCollectionType, GeometryType, LineStringType, MultiLineStringType,
+    MultiPointType, MultiPolygonType, PointType, PolygonType, WkbType, WktType,
+};
 
+/// A factory for creating GeoArrow extension types in Polars.
 pub struct GeoArrowExtensionTypeFactory;
 
 impl ExtensionTypeFactory for GeoArrowExtensionTypeFactory {
@@ -10,6 +20,55 @@ impl ExtensionTypeFactory for GeoArrowExtensionTypeFactory {
         storage: &DataType,
         metadata: Option<&str>,
     ) -> Box<dyn ExtensionTypeImpl> {
-        todo!()
+        let arrow_data_type = polars_storage_type_to_arrow_data_type(storage);
+
+        // Create Arrow field with extension metadata
+        let mut arrow_metadata = HashMap::<String, String>::with_capacity(2);
+        arrow_metadata.insert(
+            arrow_schema::extension::EXTENSION_TYPE_NAME_KEY.to_string(),
+            name.to_string(),
+        );
+        if let Some(metadata) = metadata {
+            arrow_metadata.insert(
+                arrow_schema::extension::EXTENSION_TYPE_METADATA_KEY.to_string(),
+                metadata.to_string(),
+            );
+        }
+
+        let arrow_field =
+            arrow_schema::Field::new("", arrow_data_type, true).with_metadata(arrow_metadata);
+
+        // TODO: are we assured that the name matches the type here?
+        // What do we do if the storage type isn't compatible with the extension type?
+        let geoarrow_type = GeoArrowType::from_extension_field(&arrow_field)
+            .expect("Creation of GeoArrow extension type");
+
+        match geoarrow_type {
+            GeoArrowType::Point(t) => Box::new(PointType::new(t)),
+            GeoArrowType::LineString(t) => Box::new(LineStringType::new(t)),
+            GeoArrowType::Polygon(t) => Box::new(PolygonType::new(t)),
+            GeoArrowType::MultiPoint(t) => Box::new(MultiPointType::new(t)),
+            GeoArrowType::MultiLineString(t) => Box::new(MultiLineStringType::new(t)),
+            GeoArrowType::MultiPolygon(t) => Box::new(MultiPolygonType::new(t)),
+            GeoArrowType::GeometryCollection(t) => Box::new(GeometryCollectionType::new(t)),
+            GeoArrowType::Geometry(t) => Box::new(GeometryType::new(t)),
+            GeoArrowType::Rect(t) => Box::new(BoxType::new(t)),
+            GeoArrowType::Wkb(t) => Box::new(WkbType::new(t)),
+            GeoArrowType::LargeWkb(t) => Box::new(WkbType::new(t)),
+            GeoArrowType::WkbView(t) => Box::new(WkbType::new(t)),
+            GeoArrowType::Wkt(t) => Box::new(WktType::new(t)),
+            GeoArrowType::LargeWkt(t) => Box::new(WktType::new(t)),
+            GeoArrowType::WktView(t) => Box::new(WktType::new(t)),
+        }
     }
 }
+
+fn polars_storage_type_to_arrow_data_type(storage: &DataType) -> arrow_schema::DataType {
+    let polars_field = polars_arrow::datatypes::Field::new(
+        "".into(),
+        storage.to_arrow(CompatLevel::newest()),
+        true,
+    );
+    let arrow_field_from_ffi = polars_field_to_arrow(&polars_field);
+    arrow_field_from_ffi.data_type().clone()
+}

geopolars/geopolars-extension/src/factory.rs

@@ -1,3 +1,8 @@
+//! Definitions of GeoArrow extension types for the Polars extension registry.
+//!
+//! Each of the types contained in this module is a thin wrapper around the corresponding type in
+//! [`geoarrow_schema`] to implement Polars' [`ExtensionTypeImpl`].
+
 use std::any::Any;
 use std::borrow::Cow;
 use std::hash::{BuildHasher, Hash};

geopolars/geopolars-extension/src/geoarrow.rs

@@ -1,6 +1,39 @@
-mod factory;
-mod geoarrow;
+pub mod factory;
+pub mod geoarrow;
 
+use std::sync::Arc;
+
+use arrow_schema::extension::ExtensionType;
+use polars::error::PolarsResult;
 use polars_core::datatypes::extension::register_extension_type;
 
-pub fn register_all_extensions() {}
+use crate::factory::GeoArrowExtensionTypeFactory;
+
+/// Register all GeoArrow extension types into the Polars extension registry.
+pub fn register_all_extensions() -> PolarsResult<()> {
+    let factory = Arc::new(GeoArrowExtensionTypeFactory);
+
+    register_extension_type(geoarrow_schema::PointType::NAME, Some(factory.clone()))?;
+    register_extension_type(geoarrow_schema::LineStringType::NAME, Some(factory.clone()))?;
+    register_extension_type(geoarrow_schema::PolygonType::NAME, Some(factory.clone()))?;
+    register_extension_type(geoarrow_schema::MultiPointType::NAME, Some(factory.clone()))?;
+    register_extension_type(
+        geoarrow_schema::MultiLineStringType::NAME,
+        Some(factory.clone()),
+    )?;
+    register_extension_type(
+        geoarrow_schema::MultiPolygonType::NAME,
+        Some(factory.clone()),
+    )?;
+    register_extension_type(
+        geoarrow_schema::GeometryCollectionType::NAME,
+        Some(factory.clone()),
+    )?;
+    register_extension_type(geoarrow_schema::GeometryType::NAME, Some(factory.clone()))?;
+    register_extension_type(geoarrow_schema::BoxType::NAME, Some(factory.clone()))?;
+
+    register_extension_type(geoarrow_schema::WkbType::NAME, Some(factory.clone()))?;
+    register_extension_type(geoarrow_schema::WktType::NAME, Some(factory))?;
+
+    Ok(())
+}