vortex_geo/extension/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4pub(crate) mod coordinate;
5mod linestring;
6mod multilinestring;
7mod multipoint;
8mod multipolygon;
9mod point;
10mod polygon;
11mod wkb;
12
13use std::fmt::Display;
14use std::sync::Arc;
15
16use ::wkb::reader::GeometryType;
17use arrow_array::BinaryArray;
18use geo_types::Geometry;
19use geoarrow::array::GenericWkbArray;
20use geoarrow::array::GeoArrowArray;
21use geoarrow::datatypes::CoordType;
22use geoarrow::datatypes::Crs;
23use geoarrow::datatypes::Dimension;
24use geoarrow::datatypes::GeoArrowType;
25use geoarrow::datatypes::LineStringType;
26use geoarrow::datatypes::Metadata;
27use geoarrow::datatypes::MultiLineStringType;
28use geoarrow::datatypes::MultiPointType;
29use geoarrow::datatypes::MultiPolygonType;
30use geoarrow::datatypes::PointType;
31use geoarrow::datatypes::PolygonType;
32use geoarrow::datatypes::WkbType;
33use geoarrow_cast::cast::cast;
34pub use linestring::*;
35pub use multilinestring::*;
36pub use multipoint::*;
37pub use multipolygon::*;
38pub use point::*;
39pub use polygon::*;
40use vortex_array::ArrayRef;
41use vortex_array::ExecutionCtx;
42use vortex_array::IntoArray;
43use vortex_array::arrays::ConstantArray;
44use vortex_array::arrays::ExtensionArray;
45use vortex_array::arrays::extension::ExtensionArrayExt;
46use vortex_array::arrow::FromArrowArray;
47use vortex_array::dtype::extension::ExtDType;
48use vortex_array::dtype::extension::ExtVTable;
49use vortex_array::scalar::Scalar;
50use vortex_error::VortexResult;
51use vortex_error::vortex_bail;
52use vortex_error::vortex_err;
53pub use wkb::*;
54
55/// Decode a native geometry column to `geo_types`. A non-geometry operand is an error.
56pub(crate) fn geometries(
57    array: &ArrayRef,
58    ctx: &mut ExecutionCtx,
59) -> VortexResult<Vec<Geometry<f64>>> {
60    let Some(ext) = array.dtype().as_extension_opt() else {
61        vortex_bail!(
62            "geo: operand is not a geometry extension type, was {}",
63            array.dtype()
64        );
65    };
66    let storage = array
67        .clone()
68        .execute::<ExtensionArray>(ctx)?
69        .storage_array()
70        .clone();
71    if ext.is::<Point>() {
72        point_geometries(&storage, ctx)
73    } else if ext.is::<LineString>() {
74        linestring_geometries(&storage, ctx)
75    } else if ext.is::<MultiPoint>() {
76        multipoint_geometries(&storage, ctx)
77    } else if ext.is::<Polygon>() {
78        polygon_geometries(&storage, ctx)
79    } else if ext.is::<MultiLineString>() {
80        multilinestring_geometries(&storage, ctx)
81    } else if ext.is::<MultiPolygon>() {
82        multipolygon_geometries(&storage, ctx)
83    } else {
84        vortex_bail!("geo: unsupported geometry extension {}", array.dtype())
85    }
86}
87
88/// Decode a constant operand scalar to one geo geometry, a constant of any
89/// supported geometry type is decoded exactly like a column.
90pub(crate) fn single_geometry(
91    scalar: &Scalar,
92    ctx: &mut ExecutionCtx,
93) -> VortexResult<Geometry<f64>> {
94    let array = ConstantArray::new(scalar.clone(), 1).into_array();
95    geometries(&array, ctx)?
96        .pop()
97        .ok_or_else(|| vortex_err!("geo: constant operand decoded to no geometry"))
98}
99
100/// Decode a WKB geometry literal (DuckDB's wire form for `GEOMETRY` constants) to its native
101/// `Point`/`Polygon`/`MultiPolygon` scalar. `None` for unsupported types. Plan-time, one value only.
102pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult<Option<Scalar>> {
103    let metadata = geoarrow_metadata(&GeoMetadata::default());
104    let binary = BinaryArray::from(vec![Some(bytes)]);
105    let wkb = GenericWkbArray::<i32>::try_from((
106        &binary as &dyn arrow_array::Array,
107        WkbType::new(Arc::clone(&metadata)),
108    ))
109    .map_err(|e| vortex_err!("failed to read WKB literal: {e}"))?;
110
111    // Cast the WKB value to `target`, import its native storage as a Vortex array.
112    let to_storage = |target: &GeoArrowType| -> VortexResult<ArrayRef> {
113        let native =
114            cast(&wkb, target).map_err(|e| vortex_err!("failed to cast WKB literal: {e}"))?;
115        ArrayRef::from_arrow(native.to_array_ref().as_ref(), false)
116    };
117
118    let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() {
119        GeometryType::Point => {
120            let target = GeoArrowType::Point(
121                PointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated),
122            );
123            geo_ext_scalar(Point, to_storage(&target)?)?
124        }
125        GeometryType::LineString => {
126            let target = GeoArrowType::LineString(
127                LineStringType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated),
128            );
129            geo_ext_scalar(LineString, to_storage(&target)?)?
130        }
131        GeometryType::Polygon => {
132            let target = GeoArrowType::Polygon(
133                PolygonType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated),
134            );
135            geo_ext_scalar(Polygon, to_storage(&target)?)?
136        }
137        GeometryType::MultiPoint => {
138            let target = GeoArrowType::MultiPoint(
139                MultiPointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated),
140            );
141            geo_ext_scalar(MultiPoint, to_storage(&target)?)?
142        }
143        GeometryType::MultiLineString => {
144            let target = GeoArrowType::MultiLineString(
145                MultiLineStringType::new(Dimension::XY, metadata)
146                    .with_coord_type(CoordType::Separated),
147            );
148            geo_ext_scalar(MultiLineString, to_storage(&target)?)?
149        }
150        GeometryType::MultiPolygon => {
151            let target = GeoArrowType::MultiPolygon(
152                MultiPolygonType::new(Dimension::XY, metadata)
153                    .with_coord_type(CoordType::Separated),
154            );
155            geo_ext_scalar(MultiPolygon, to_storage(&target)?)?
156        }
157        _ => return Ok(None),
158    };
159    Ok(Some(scalar))
160}
161
162/// Wrap cast-from-WKB `storage` in its `vtable` extension type and pull out the single scalar.
163// `scalar_at` is deprecated for `execute_scalar`, but there is no execution context at plan time.
164#[allow(deprecated)]
165fn geo_ext_scalar<V: ExtVTable<Metadata = GeoMetadata>>(
166    vtable: V,
167    storage: ArrayRef,
168) -> VortexResult<Scalar> {
169    let ext = ExtDType::try_with_vtable(vtable, GeoMetadata::default(), storage.dtype().clone())?
170        .erased();
171    ExtensionArray::try_new(ext, storage)?
172        .into_array()
173        .scalar_at(0)
174}
175
176/// Extension metadata that is common to all the geospatial extension types.
177///
178/// Currently, this is just the coordinate reference system (CRS).
179/// We may wish to add a second field for edges interpretation in the future similar to
180/// the GeoArrow standard.
181#[derive(Clone, PartialEq, Eq, Hash, prost::Message)]
182pub struct GeoMetadata {
183    #[prost(optional, string, tag = "1")]
184    pub crs: Option<String>,
185}
186
187impl Display for GeoMetadata {
188    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189        match self.crs.as_ref() {
190            Some(crs) => write!(f, "Geometry(crs={crs})"),
191            None => write!(f, "Geometry(unreferenced)"),
192        }
193    }
194}
195
196/// The GeoArrow [`Metadata`] equivalent of `geo_metadata`.
197pub(crate) fn geoarrow_metadata(geo_metadata: &GeoMetadata) -> Arc<Metadata> {
198    Arc::new(Metadata::new(
199        geo_metadata
200            .crs
201            .as_ref()
202            .map(|crs| Crs::from_unknown_crs_type(crs.to_string()))
203            .unwrap_or_default(),
204        None,
205    ))
206}
207
208/// Serialize a native geometry array to WKB (a `WkbView` array) via geoarrow's cast.
209/// Shared by the `to_wkb` methods on the geometry extension types.
210pub(crate) fn geoarrow_to_wkb(geo_array: &dyn GeoArrowArray) -> VortexResult<ArrayRef> {
211    let wkb_type = GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&GeoMetadata::default())));
212    let wkb = cast(geo_array, &wkb_type)
213        .map_err(|e| vortex_err!("failed to cast geometry to WKB: {e}"))?;
214    ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false)
215}
216
217/// Recover [`GeoMetadata`] from GeoArrow metadata.
218pub(crate) fn geo_metadata_from_arrow(metadata: &Metadata) -> GeoMetadata {
219    let crs = metadata.crs().crs_value().map(|value| {
220        // `Crs::from_unknown_crs_type` stores the user's string verbatim as a JSON string
221        // value, so prefer the raw string when available to round-trip cleanly. For other
222        // CRS encodings (PROJJSON object, etc.), fall back to the JSON-encoded form.
223        value
224            .as_str()
225            .map(str::to_string)
226            .unwrap_or_else(|| value.to_string())
227    });
228    GeoMetadata { crs }
229}
230
231#[cfg(test)]
232mod tests {
233    use prost::Message;
234    use vortex_array::dtype::DType;
235    use vortex_error::VortexResult;
236    use vortex_error::vortex_err;
237
238    use super::LineString;
239    use super::MultiLineString;
240    use super::MultiPoint;
241    use super::Point;
242    use super::Polygon;
243    use super::native_geometry_scalar_from_wkb;
244    use crate::extension::GeoMetadata;
245
246    #[test]
247    fn test_metadata() {
248        let meta = GeoMetadata {
249            crs: Some("EPSG:4326".to_string()),
250        };
251
252        assert_eq!(meta.to_string(), "Geometry(crs=EPSG:4326)");
253        // round trip
254        let bytes = meta.encode_to_vec();
255        let decoded = GeoMetadata::decode(bytes.as_slice()).unwrap();
256        assert_eq!(decoded, meta);
257    }
258
259    /// A little-endian WKB `POINT` literal decodes to the native `Point` extension scalar.
260    #[test]
261    fn decodes_wkb_point_to_native() -> VortexResult<()> {
262        let mut wkb = vec![1u8]; // little-endian byte order
263        wkb.extend_from_slice(&1u32.to_le_bytes()); // geometry type: point
264        wkb.extend_from_slice(&1.0f64.to_le_bytes()); // x
265        wkb.extend_from_slice(&2.0f64.to_le_bytes()); // y
266
267        let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a point scalar");
268        let DType::Extension(ext) = scalar.dtype() else {
269            panic!("expected an extension dtype, got {}", scalar.dtype());
270        };
271        assert!(ext.is::<Point>());
272        Ok(())
273    }
274
275    /// A little-endian WKB `POLYGON` literal decodes to the native `Polygon` extension scalar.
276    #[test]
277    fn decodes_wkb_polygon_to_native() -> VortexResult<()> {
278        let ring = [(0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (0.0, 0.0)];
279        let mut wkb = vec![1u8]; // little-endian byte order
280        wkb.extend_from_slice(&3u32.to_le_bytes()); // geometry type: polygon
281        wkb.extend_from_slice(&1u32.to_le_bytes()); // one ring
282        let ring_len = u32::try_from(ring.len()).map_err(|e| vortex_err!("{e}"))?;
283        wkb.extend_from_slice(&ring_len.to_le_bytes());
284        for (x, y) in ring {
285            wkb.extend_from_slice(&f64::to_le_bytes(x));
286            wkb.extend_from_slice(&f64::to_le_bytes(y));
287        }
288
289        let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a polygon scalar");
290        let DType::Extension(ext) = scalar.dtype() else {
291            panic!("expected an extension dtype, got {}", scalar.dtype());
292        };
293        assert!(ext.is::<Polygon>());
294        Ok(())
295    }
296
297    /// A little-endian WKB `LINESTRING` literal decodes to the native `LineString` extension scalar.
298    #[test]
299    fn decodes_wkb_linestring_to_native() -> VortexResult<()> {
300        let points = [(0.0, 0.0), (1.0, 1.0)];
301        let mut wkb = vec![1u8]; // little-endian byte order
302        wkb.extend_from_slice(&2u32.to_le_bytes()); // geometry type: linestring
303        let len = u32::try_from(points.len()).map_err(|e| vortex_err!("{e}"))?;
304        wkb.extend_from_slice(&len.to_le_bytes());
305        for (x, y) in points {
306            wkb.extend_from_slice(&f64::to_le_bytes(x));
307            wkb.extend_from_slice(&f64::to_le_bytes(y));
308        }
309
310        let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a linestring scalar");
311        let DType::Extension(ext) = scalar.dtype() else {
312            panic!("expected an extension dtype, got {}", scalar.dtype());
313        };
314        assert!(ext.is::<LineString>());
315        Ok(())
316    }
317
318    /// A little-endian WKB `MULTIPOINT` literal decodes to the native `MultiPoint` extension scalar.
319    #[test]
320    fn decodes_wkb_multipoint_to_native() -> VortexResult<()> {
321        let points = [(0.0, 0.0), (1.0, 1.0)];
322        let mut wkb = vec![1u8]; // little-endian byte order
323        wkb.extend_from_slice(&4u32.to_le_bytes()); // geometry type: multipoint
324        let len = u32::try_from(points.len()).map_err(|e| vortex_err!("{e}"))?;
325        wkb.extend_from_slice(&len.to_le_bytes());
326        for (x, y) in points {
327            // each member is a full WKB point
328            wkb.push(1u8);
329            wkb.extend_from_slice(&1u32.to_le_bytes());
330            wkb.extend_from_slice(&f64::to_le_bytes(x));
331            wkb.extend_from_slice(&f64::to_le_bytes(y));
332        }
333
334        let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a multipoint scalar");
335        let DType::Extension(ext) = scalar.dtype() else {
336            panic!("expected an extension dtype, got {}", scalar.dtype());
337        };
338        assert!(ext.is::<MultiPoint>());
339        Ok(())
340    }
341
342    /// A little-endian WKB `MULTILINESTRING` literal decodes to the native `MultiLineString` scalar.
343    #[test]
344    fn decodes_wkb_multilinestring_to_native() -> VortexResult<()> {
345        let lines = [[(0.0, 0.0), (1.0, 1.0)], [(2.0, 2.0), (3.0, 3.0)]];
346        let mut wkb = vec![1u8]; // little-endian byte order
347        wkb.extend_from_slice(&5u32.to_le_bytes()); // geometry type: multilinestring
348        let num_lines = u32::try_from(lines.len()).map_err(|e| vortex_err!("{e}"))?;
349        wkb.extend_from_slice(&num_lines.to_le_bytes());
350        for line in lines {
351            // each member is a full WKB linestring
352            wkb.push(1u8);
353            wkb.extend_from_slice(&2u32.to_le_bytes());
354            let len = u32::try_from(line.len()).map_err(|e| vortex_err!("{e}"))?;
355            wkb.extend_from_slice(&len.to_le_bytes());
356            for (x, y) in line {
357                wkb.extend_from_slice(&f64::to_le_bytes(x));
358                wkb.extend_from_slice(&f64::to_le_bytes(y));
359            }
360        }
361
362        let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a multilinestring scalar");
363        let DType::Extension(ext) = scalar.dtype() else {
364            panic!("expected an extension dtype, got {}", scalar.dtype());
365        };
366        assert!(ext.is::<MultiLineString>());
367        Ok(())
368    }
369}
vortex_geo/extension/mod.rs

vortex_geo/extension/
mod.rs
