DataFusion extension types working proof of concept!

Cargo.toml

@@ -9,9 +9,7 @@ members = [
     "rust/geoarrow-test",
     "rust/geoarrow",
     "rust/pyo3-geoarrow",
-    # Comment out until datafusion 47 release so that the workspace can upgrade
-    # to arrow 55
-    # "rust/geodatafusion",
+    "rust/geodatafusion",
 ]
 exclude = ["js"]
 resolver = "2"
@@ -40,8 +38,8 @@ arrow-schema = "55"
 async-stream = "0.3"
 async-trait = "0.1"
 bytes = "1.10.0"
-# 47 rc https://lists.apache.org/thread/zrq9x9gf51r8b6m9qokf2q75kh251rm6
-datafusion = { git = "https://github.com/apache/datafusion", rev = "e4433049b04ca2c1e2031eb05d1a0990210f11d6" }
+# https://github.com/apache/datafusion/pull/15646
+datafusion = { git = "https://github.com/apache/datafusion", rev = "2d801940c3cb0cec3209aa890688590ded791865" }
 flatgeobuf = { version = "4.6", default-features = false }
 futures = "0.3"
 geo = "0.30.0"

rust/geoarrow-array/src/array/geometry.rs

@@ -375,8 +375,6 @@ impl GeometryArray {
     pub fn into_coord_type(self, coord_type: CoordType) -> Self {
         Self {
             data_type: self.data_type.with_coord_type(coord_type),
-            type_ids: self.type_ids,
-            offsets: self.offsets,
             points: self.points.map(|arr| arr.into_coord_type(coord_type)),
             line_strings: self.line_strings.map(|arr| arr.into_coord_type(coord_type)),
             polygons: self.polygons.map(|arr| arr.into_coord_type(coord_type)),
@@ -386,6 +384,15 @@ impl GeometryArray {
                 .map(|arr| arr.into_coord_type(coord_type)),
             mpolygons: self.mpolygons.map(|arr| arr.into_coord_type(coord_type)),
             gcs: self.gcs.map(|arr| arr.into_coord_type(coord_type)),
+            ..self
+        }
+    }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
         }
     }
 
@@ -486,6 +493,10 @@ impl GeoArrowArray for GeometryArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for GeometryArray {

rust/geoarrow-array/src/array/geometrycollection.rs

@@ -100,8 +100,15 @@ impl GeometryCollectionArray {
         Self {
             data_type: self.data_type.with_coord_type(coord_type),
             array: self.array.into_coord_type(coord_type),
-            geom_offsets: self.geom_offsets,
-            nulls: self.nulls,
+            ..self
+        }
+    }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
         }
     }
 }
@@ -149,6 +156,10 @@ impl GeoArrowArray for GeometryCollectionArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for GeometryCollectionArray {

rust/geoarrow-array/src/array/linestring.rs

@@ -154,8 +154,15 @@ impl LineStringArray {
         Self {
             data_type: self.data_type.with_coord_type(coord_type),
             coords: self.coords.into_coord_type(coord_type),
-            geom_offsets: self.geom_offsets,
-            nulls: self.nulls,
+            ..self
+        }
+    }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
         }
     }
 }
@@ -202,6 +209,10 @@ impl GeoArrowArray for LineStringArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for LineStringArray {

rust/geoarrow-array/src/array/mixed.rs

@@ -6,7 +6,7 @@ use arrow_array::{Array, ArrayRef, UnionArray};
 use arrow_buffer::ScalarBuffer;
 use arrow_schema::{DataType, UnionMode};
 use geoarrow_schema::{
-    CoordType, Dimension, GeometryCollectionType, LineStringType, Metadata, MultiLineStringType,
+    CoordType, Dimension, GeometryCollectionType, LineStringType, MultiLineStringType,
     MultiPointType, MultiPolygonType, PointType, PolygonType,
 };
 
@@ -65,8 +65,6 @@ pub struct MixedGeometryArray {
     pub(crate) coord_type: CoordType,
     pub(crate) dim: Dimension,
 
-    pub(crate) metadata: Arc<Metadata>,
-
     /// Invariant: every item in `type_ids` is `> 0 && < fields.len()` if `type_ids` are not provided. If `type_ids` exist in the NativeType, then every item in `type_ids` is `> 0 && `
     pub(crate) type_ids: ScalarBuffer<i8>,
 
@@ -107,7 +105,6 @@ impl MixedGeometryArray {
         multi_points: Option<MultiPointArray>,
         multi_line_strings: Option<MultiLineStringArray>,
         multi_polygons: Option<MultiPolygonArray>,
-        metadata: Arc<Metadata>,
     ) -> Self {
         let mut coord_types = HashSet::new();
         if let Some(points) = &points {
@@ -191,7 +188,6 @@ impl MixedGeometryArray {
                 ))
                 .finish(),
             ),
-            metadata,
             slice_offset: 0,
         }
     }
@@ -424,25 +420,20 @@ impl MixedGeometryArray {
             multi_points: self.multi_points.clone(),
             multi_line_strings: self.multi_line_strings.clone(),
             multi_polygons: self.multi_polygons.clone(),
-            metadata: self.metadata.clone(),
             slice_offset: self.slice_offset + offset,
         }
     }
 
     pub fn into_coord_type(self, coord_type: CoordType) -> Self {
         Self {
             coord_type,
-            dim: self.dim,
-            metadata: self.metadata,
-            type_ids: self.type_ids,
-            offsets: self.offsets,
             points: self.points.into_coord_type(coord_type),
             line_strings: self.line_strings.into_coord_type(coord_type),
             polygons: self.polygons.into_coord_type(coord_type),
             multi_points: self.multi_points.into_coord_type(coord_type),
             multi_line_strings: self.multi_line_strings.into_coord_type(coord_type),
             multi_polygons: self.multi_polygons.into_coord_type(coord_type),
-            slice_offset: self.slice_offset,
+            ..self
         }
     }
 
@@ -660,9 +651,6 @@ impl TryFrom<(&UnionArray, Dimension, CoordType)> for MixedGeometryArray {
             multi_points,
             multi_line_strings,
             multi_polygons,
-            // Mixed array is only used inside of GeometryCollectionArray, and this array does not
-            // hold its own metadata
-            Default::default(),
         ))
     }
 }
@@ -684,7 +672,6 @@ impl TryFrom<(&dyn Array, Dimension, CoordType)> for MixedGeometryArray {
 impl PartialEq for MixedGeometryArray {
     fn eq(&self, other: &Self) -> bool {
         self.dim == other.dim
-            && self.metadata == other.metadata
             && self.type_ids == other.type_ids
             && self.offsets == other.offsets
             && self.points == other.points

rust/geoarrow-array/src/array/multilinestring.rs

@@ -181,9 +181,15 @@ impl MultiLineStringArray {
         Self {
             data_type: self.data_type.with_coord_type(coord_type),
             coords: self.coords.into_coord_type(coord_type),
-            geom_offsets: self.geom_offsets,
-            ring_offsets: self.ring_offsets,
-            nulls: self.nulls,
+            ..self
+        }
+    }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
         }
     }
 }
@@ -231,6 +237,10 @@ impl GeoArrowArray for MultiLineStringArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for MultiLineStringArray {

rust/geoarrow-array/src/array/multipoint.rs

@@ -158,8 +158,15 @@ impl MultiPointArray {
         Self {
             data_type: self.data_type.with_coord_type(coord_type),
             coords: self.coords.into_coord_type(coord_type),
-            geom_offsets: self.geom_offsets,
-            nulls: self.nulls,
+            ..self
+        }
+    }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
         }
     }
 }
@@ -207,6 +214,10 @@ impl GeoArrowArray for MultiPointArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for MultiPointArray {

rust/geoarrow-array/src/array/multipolygon.rs

@@ -234,10 +234,15 @@ impl MultiPolygonArray {
         Self {
             data_type: self.data_type.with_coord_type(coord_type),
             coords: self.coords.into_coord_type(coord_type),
-            geom_offsets: self.geom_offsets,
-            polygon_offsets: self.polygon_offsets,
-            ring_offsets: self.ring_offsets,
-            nulls: self.nulls,
+            ..self
+        }
+    }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
         }
     }
 }
@@ -285,6 +290,10 @@ impl GeoArrowArray for MultiPolygonArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for MultiPolygonArray {

rust/geoarrow-array/src/array/point.rs

@@ -119,7 +119,15 @@ impl PointArray {
         Self {
             data_type: self.data_type.with_coord_type(coord_type),
             coords: self.coords.into_coord_type(coord_type),
-            nulls: self.nulls,
+            ..self
+        }
+    }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
         }
     }
 }
@@ -167,6 +175,10 @@ impl GeoArrowArray for PointArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for PointArray {

rust/geoarrow-array/src/array/polygon.rs

@@ -185,9 +185,15 @@ impl PolygonArray {
         Self {
             data_type: self.data_type.with_coord_type(coord_type),
             coords: self.coords.into_coord_type(coord_type),
-            geom_offsets: self.geom_offsets,
-            ring_offsets: self.ring_offsets,
-            nulls: self.nulls,
+            ..self
+        }
+    }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
         }
     }
 }
@@ -235,6 +241,10 @@ impl GeoArrowArray for PolygonArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for PolygonArray {

rust/geoarrow-array/src/array/rect.rs

@@ -86,6 +86,14 @@ impl RectArray {
             nulls: self.nulls.as_ref().map(|v| v.slice(offset, length)),
         }
     }
+
+    /// Change the [`Metadata`] of this array.
+    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
+        Self {
+            data_type: self.data_type.with_metadata(metadata),
+            ..self
+        }
+    }
 }
 
 impl GeoArrowArray for RectArray {
@@ -131,6 +139,10 @@ impl GeoArrowArray for RectArray {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(self.with_metadata(metadata))
+    }
 }
 
 impl<'a> ArrayAccessor<'a> for RectArray {

rust/geoarrow-array/src/array/wkb.rs

@@ -131,6 +131,10 @@ impl<O: OffsetSizeTrait> GeoArrowArray for WkbArray<O> {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(Self::with_metadata(&self, metadata))
+    }
 }
 
 impl<'a, O: OffsetSizeTrait> ArrayAccessor<'a> for WkbArray<O> {

rust/geoarrow-array/src/array/wkt.rs

@@ -122,6 +122,10 @@ impl<O: OffsetSizeTrait> GeoArrowArray for WktArray<O> {
     fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
         Arc::new(self.slice(offset, length))
     }
+
+    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
+        Arc::new(Self::with_metadata(&self, metadata))
+    }
 }
 
 impl<'a, O: OffsetSizeTrait> ArrayAccessor<'a> for WktArray<O> {

rust/geoarrow-array/src/builder/geometrycollection.rs

@@ -45,7 +45,6 @@ impl<'a> GeometryCollectionBuilder {
                 typ.dimension(),
                 capacity.mixed_capacity,
                 typ.coord_type(),
-                typ.metadata().clone(),
             ),
             geom_offsets: OffsetsBuilder::with_capacity(capacity.geom_capacity),
             validity: NullBufferBuilder::new(capacity.geom_capacity),