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 point;
6mod polygon;
7mod wkb;
8
9use std::fmt::Display;
10use std::sync::Arc;
11
12use geo_types::Geometry;
13use geoarrow::datatypes::Crs;
14use geoarrow::datatypes::Metadata;
15pub use point::*;
16pub use polygon::*;
17use vortex_array::ArrayRef;
18use vortex_array::ExecutionCtx;
19use vortex_array::IntoArray;
20use vortex_array::arrays::ConstantArray;
21use vortex_array::arrays::ExtensionArray;
22use vortex_array::arrays::extension::ExtensionArrayExt;
23use vortex_array::scalar::Scalar;
24use vortex_error::VortexResult;
25use vortex_error::vortex_bail;
26use vortex_error::vortex_err;
27pub use wkb::*;
28
29/// Decode a native geometry column to `geo_types`. A non-geometry operand is an error.
30pub(crate) fn geometries(
31    array: &ArrayRef,
32    ctx: &mut ExecutionCtx,
33) -> VortexResult<Vec<Geometry<f64>>> {
34    let Some(ext) = array.dtype().as_extension_opt() else {
35        vortex_bail!(
36            "geo: operand is not a geometry extension type, was {}",
37            array.dtype()
38        );
39    };
40    let storage = array
41        .clone()
42        .execute::<ExtensionArray>(ctx)?
43        .storage_array()
44        .clone();
45    if ext.is::<Point>() {
46        point_geometries(&storage, ctx)
47    } else if ext.is::<Polygon>() {
48        polygon_geometries(&storage, ctx)
49    } else {
50        vortex_bail!("geo: unsupported geometry extension {}", array.dtype())
51    }
52}
53
54/// Decode a constant operand scalar to one geo geometry, a constant of any
55/// supported geometry type is decoded exactly like a column.
56pub(crate) fn single_geometry(
57    scalar: &Scalar,
58    ctx: &mut ExecutionCtx,
59) -> VortexResult<Geometry<f64>> {
60    let array = ConstantArray::new(scalar.clone(), 1).into_array();
61    geometries(&array, ctx)?
62        .pop()
63        .ok_or_else(|| vortex_err!("geo: constant operand decoded to no geometry"))
64}
65
66/// Extension metadata that is common to all the geospatial extension types.
67///
68/// Currently, this is just the coordinate reference system (CRS).
69/// We may wish to add a second field for edges interpretation in the future similar to
70/// the GeoArrow standard.
71#[derive(Clone, PartialEq, Eq, Hash, prost::Message)]
72pub struct GeoMetadata {
73    #[prost(optional, string, tag = "1")]
74    pub crs: Option<String>,
75}
76
77impl Display for GeoMetadata {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        match self.crs.as_ref() {
80            Some(crs) => write!(f, "Geometry(crs={crs})"),
81            None => write!(f, "Geometry(unreferenced)"),
82        }
83    }
84}
85
86/// The GeoArrow [`Metadata`] equivalent of `geo_metadata`.
87pub(crate) fn geoarrow_metadata(geo_metadata: &GeoMetadata) -> Arc<Metadata> {
88    Arc::new(Metadata::new(
89        geo_metadata
90            .crs
91            .as_ref()
92            .map(|crs| Crs::from_unknown_crs_type(crs.to_string()))
93            .unwrap_or_default(),
94        None,
95    ))
96}
97
98/// Recover [`GeoMetadata`] from GeoArrow metadata.
99pub(crate) fn geo_metadata_from_arrow(metadata: &Metadata) -> GeoMetadata {
100    let crs = metadata.crs().crs_value().map(|value| {
101        // `Crs::from_unknown_crs_type` stores the user's string verbatim as a JSON string
102        // value, so prefer the raw string when available to round-trip cleanly. For other
103        // CRS encodings (PROJJSON object, etc.), fall back to the JSON-encoded form.
104        value
105            .as_str()
106            .map(str::to_string)
107            .unwrap_or_else(|| value.to_string())
108    });
109    GeoMetadata { crs }
110}
111
112#[cfg(test)]
113mod tests {
114    use prost::Message;
115
116    use crate::extension::GeoMetadata;
117
118    #[test]
119    fn test_metadata() {
120        let meta = GeoMetadata {
121            crs: Some("EPSG:4326".to_string()),
122        };
123
124        assert_eq!(meta.to_string(), "Geometry(crs=EPSG:4326)");
125        // round trip
126        let bytes = meta.encode_to_vec();
127        let decoded = GeoMetadata::decode(bytes.as_slice()).unwrap();
128        assert_eq!(decoded, meta);
129    }
130}
vortex_geo/extension/mod.rs

vortex_geo/extension/
mod.rs
