Skip to main content

parquet_geospatial/
types.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow_schema::{ArrowError, DataType, extension::ExtensionType};
19use serde::{Deserialize, Serialize};
20
21/// Hints at the likely Parquet geospatial logical type represented by a [`Metadata`].
22///
23/// Based on the `algorithm` field:
24/// - [`Hint::Geometry`]: WKB format with linear/planar edge interpolation
25/// - [`Hint::Geography`]: WKB format with explicit non-linear/non-planar edge interpolation
26///
27/// See the [Parquet Geospatial specification](https://github.com/apache/parquet-format/blob/master/Geospatial.md)
28/// for more details.
29#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
30pub enum Hint {
31    /// Geospatial features in WKB format with linear/planar edge interpolation
32    Geometry,
33    /// Geospatial features in WKB format with explicit non-linear/non-planar edge interpolation
34    Geography,
35}
36
37/// The edge interpolation algorithms used with `GEOMETRY` logical types.
38#[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(rename_all = "lowercase")]
40pub enum Edges {
41    /// Edges are interpolated as geodesics on a sphere.
42    #[default]
43    Spherical,
44    /// <https://en.wikipedia.org/wiki/Vincenty%27s_formulae>
45    Vincenty,
46    /// Thomas, Paul D. Spheroidal geodesics, reference systems, & local geometry. US Naval Oceanographic Office, 1970
47    Thomas,
48    /// Thomas, Paul D. Mathematical models for navigation systems. US Naval Oceanographic Office, 1965.
49    Andoyer,
50    /// Karney, Charles FF. "Algorithms for geodesics." Journal of Geodesy 87 (2013): 43-55
51    Karney,
52}
53
54/// The metadata associated with a [`WkbType`].
55#[derive(Clone, Debug, Default, Serialize, Deserialize)]
56pub struct Metadata {
57    /// The Coordinate Reference System (CRS) of the [`WkbType`], if present.
58    ///
59    /// This may be a raw string value (e.g., "EPSG:3857") or a JSON object (e.g., PROJJSON).
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub crs: Option<serde_json::Value>,
62    /// The edge interpolation algorithm of the [`WkbType`], if present.
63    #[serde(rename = "edges", skip_serializing_if = "Option::is_none")]
64    pub algorithm: Option<Edges>,
65}
66
67impl Metadata {
68    /// Constructs a new [`Metadata`] with the given CRS and algorithm.
69    ///
70    /// If a CRS is provided, and can be parsed as JSON, it will be stored as a JSON object instead
71    /// of its string representation.
72    pub fn new(crs: Option<&str>, algorithm: Option<Edges>) -> Self {
73        let crs = crs.map(|c| match serde_json::from_str(c) {
74            Ok(crs) => crs,
75            Err(_) => serde_json::Value::String(c.to_string()),
76        });
77
78        Self { crs, algorithm }
79    }
80
81    /// Returns a [`Hint`] to the likely underlying Logical Type that this [`Metadata`] represents.
82    pub fn type_hint(&self) -> Hint {
83        match &self.algorithm {
84            Some(_) => Hint::Geography,
85            None => Hint::Geometry,
86        }
87    }
88
89    /// Detect if the CRS is a common representation of lon/lat on the standard WGS84 ellipsoid.
90    ///
91    /// Returns `true` for OGC:CRS84, EPSG:4326, and PROJJSON representations thereof.
92    pub fn crs_is_lon_lat(&self) -> bool {
93        use serde_json::Value;
94
95        let Some(crs) = &self.crs else {
96            return false;
97        };
98
99        match crs {
100            Value::String(s) if s == "EPSG:4326" || s == "OGC:CRS84" => true,
101            Value::Object(_) => match (&crs["id"]["authority"], &crs["id"]["code"]) {
102                (Value::String(auth), Value::String(code)) if auth == "OGC" && code == "CRS84" => {
103                    true
104                }
105                (Value::String(auth), Value::String(code)) if auth == "EPSG" && code == "4326" => {
106                    true
107                }
108                (Value::String(auth), Value::Number(code))
109                    if auth == "EPSG" && code.as_i64() == Some(4326) =>
110                {
111                    true
112                }
113                _ => false,
114            },
115            _ => false,
116        }
117    }
118}
119
120/// Well-Known Binary (WKB) [`ExtensionType`] for geospatial data.
121///
122/// Represents the canonical Arrow Extension Type for storing
123/// [GeoArrow](https://github.com/geoarrow/geoarrow) data.
124#[derive(Debug, Default)]
125pub struct WkbType(Metadata);
126
127impl WkbType {
128    /// Constructs a new [`WkbType`] with the given [`Metadata`].
129    ///
130    /// If `None` is provided, default (empty) metadata is used.
131    pub fn new(metadata: Option<Metadata>) -> Self {
132        Self(metadata.unwrap_or_default())
133    }
134}
135
136type ArrowResult<T> = Result<T, ArrowError>;
137impl ExtensionType for WkbType {
138    const NAME: &'static str = "geoarrow.wkb";
139
140    type Metadata = Metadata;
141
142    fn metadata(&self) -> &Self::Metadata {
143        &self.0
144    }
145
146    fn serialize_metadata(&self) -> Option<String> {
147        serde_json::to_string(&self.0).ok()
148    }
149
150    fn deserialize_metadata(metadata: Option<&str>) -> ArrowResult<Self::Metadata> {
151        let Some(metadata) = metadata else {
152            return Ok(Self::Metadata::default());
153        };
154
155        serde_json::from_str(metadata).map_err(|e| ArrowError::JsonError(e.to_string()))
156    }
157
158    fn supports_data_type(&self, data_type: &arrow_schema::DataType) -> ArrowResult<()> {
159        match data_type {
160            DataType::Binary | DataType::LargeBinary | DataType::BinaryView => Ok(()),
161            dt => Err(ArrowError::InvalidArgumentError(format!(
162                "Geometry data type mismatch, expected one of Binary, LargeBinary, BinaryView. Found {dt}"
163            ))),
164        }
165    }
166
167    fn try_new(data_type: &arrow_schema::DataType, metadata: Self::Metadata) -> ArrowResult<Self> {
168        let wkb = Self(metadata);
169        wkb.supports_data_type(data_type)?;
170        Ok(wkb)
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177    use arrow_schema::Field;
178
179    /// Test metadata serialization and deserialization with empty/default metadata
180    #[test]
181    fn test_metadata_empty_roundtrip() -> ArrowResult<()> {
182        let metadata = Metadata::default();
183        let wkb = WkbType::new(Some(metadata));
184
185        let serialized = wkb.serialize_metadata().unwrap();
186        assert_eq!(serialized, "{}");
187
188        let deserialized = WkbType::deserialize_metadata(Some(&serialized))?;
189        assert!(deserialized.crs.is_none());
190        assert!(deserialized.algorithm.is_none());
191
192        Ok(())
193    }
194
195    /// Test metadata serialization with CRS as a simple string
196    #[test]
197    fn test_metadata_crs_string_roundtrip() -> ArrowResult<()> {
198        let metadata = Metadata::new(Some("srid:1234"), None);
199        let wkb = WkbType::new(Some(metadata));
200
201        let serialized = wkb.serialize_metadata().unwrap();
202        assert_eq!(serialized, r#"{"crs":"srid:1234"}"#);
203
204        let deserialized = WkbType::deserialize_metadata(Some(&serialized))?;
205        assert_eq!(
206            deserialized.crs.unwrap(),
207            serde_json::Value::String(String::from("srid:1234"))
208        );
209        assert!(deserialized.algorithm.is_none());
210
211        Ok(())
212    }
213
214    /// Test metadata serialization with CRS as a JSON object
215    #[test]
216    fn test_metadata_crs_json_object_roundtrip() -> ArrowResult<()> {
217        let crs_json = r#"{"type":"custom_json","properties":{"name":"EPSG:4326"}}"#;
218        let metadata = Metadata::new(Some(crs_json), None);
219        let wkb = WkbType::new(Some(metadata));
220
221        let serialized = wkb.serialize_metadata().unwrap();
222        // Validate by parsing the JSON and checking structure (field order is not guaranteed)
223        let parsed: serde_json::Value = serde_json::from_str(&serialized).unwrap();
224        assert_eq!(parsed["crs"]["type"], "custom_json");
225        assert_eq!(parsed["crs"]["properties"]["name"], "EPSG:4326");
226
227        let deserialized = WkbType::deserialize_metadata(Some(&serialized))?;
228
229        // Verify it's a JSON object with expected structure
230        let crs = deserialized.crs.unwrap();
231        assert!(crs.is_object());
232        assert_eq!(crs["type"], "custom_json");
233        assert_eq!(crs["properties"]["name"], "EPSG:4326");
234
235        Ok(())
236    }
237
238    /// Test metadata serialization with algorithm field
239    #[test]
240    fn test_metadata_algorithm_roundtrip() -> ArrowResult<()> {
241        let metadata = Metadata::new(None, Some(Edges::Spherical));
242        let wkb = WkbType::new(Some(metadata));
243
244        let serialized = wkb.serialize_metadata().unwrap();
245        assert_eq!(serialized, r#"{"edges":"spherical"}"#);
246
247        let deserialized = WkbType::deserialize_metadata(Some(&serialized))?;
248        assert!(deserialized.crs.is_none());
249        assert_eq!(deserialized.algorithm, Some(Edges::Spherical));
250
251        Ok(())
252    }
253
254    /// Test metadata serialization with both CRS and algorithm
255    #[test]
256    fn test_metadata_full_roundtrip() -> ArrowResult<()> {
257        let metadata = Metadata::new(Some("srid:1234"), Some(Edges::Spherical));
258        let wkb = WkbType::new(Some(metadata));
259
260        let serialized = wkb.serialize_metadata().unwrap();
261        assert_eq!(serialized, r#"{"crs":"srid:1234","edges":"spherical"}"#);
262
263        let deserialized = WkbType::deserialize_metadata(Some(&serialized))?;
264        assert_eq!(
265            deserialized.crs.unwrap(),
266            serde_json::Value::String("srid:1234".to_string())
267        );
268        assert_eq!(deserialized.algorithm, Some(Edges::Spherical));
269
270        Ok(())
271    }
272
273    /// Test deserialization of None metadata
274    #[test]
275    fn test_metadata_deserialize_none() -> ArrowResult<()> {
276        let deserialized = WkbType::deserialize_metadata(None)?;
277        assert!(deserialized.crs.is_none());
278        assert!(deserialized.algorithm.is_none());
279        Ok(())
280    }
281
282    /// Test deserialization of invalid JSON
283    #[test]
284    fn test_metadata_deserialize_invalid_json() {
285        let result = WkbType::deserialize_metadata(Some("not valid json {"));
286        assert!(matches!(result, Err(ArrowError::JsonError(_))));
287    }
288
289    /// Test metadata that results in a Geometry type hint
290    #[test]
291    fn test_type_hint_geometry() {
292        let metadata = Metadata::new(None, None);
293        assert!(matches!(metadata.type_hint(), Hint::Geometry));
294    }
295
296    /// Test metadata that results in a Geography type hint
297    #[test]
298    fn test_type_hint_edges_is_geography() {
299        let algorithms = vec![
300            Edges::Spherical,
301            Edges::Vincenty,
302            Edges::Thomas,
303            Edges::Andoyer,
304            Edges::Karney,
305        ];
306        for algo in algorithms {
307            let metadata = Metadata::new(None, Some(algo));
308            assert!(matches!(metadata.type_hint(), Hint::Geography));
309        }
310    }
311
312    /// Test extension type integration using a Field
313    #[test]
314    fn test_extension_type_with_field() -> ArrowResult<()> {
315        let metadata = Metadata::new(Some("srid:1234"), None);
316        let wkb_type = WkbType::new(Some(metadata));
317
318        let mut field = Field::new("geometry", DataType::Binary, false);
319        field.try_with_extension_type(wkb_type)?;
320
321        // Verify we can extract the extension type back
322        let extracted = field.try_extension_type::<WkbType>()?;
323        assert_eq!(
324            extracted.metadata().crs.as_ref().unwrap(),
325            &serde_json::Value::String(String::from("srid:1234"))
326        );
327
328        Ok(())
329    }
330
331    /// Test extension type DataType support
332    #[test]
333    fn test_extension_type_support() -> ArrowResult<()> {
334        let wkb = WkbType::default();
335        // supported types
336        wkb.supports_data_type(&DataType::Binary)?;
337        wkb.supports_data_type(&DataType::LargeBinary)?;
338        wkb.supports_data_type(&DataType::BinaryView)?;
339
340        // reject unsupported types with an error
341        let result = wkb.supports_data_type(&DataType::Utf8);
342        assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
343
344        Ok(())
345    }
346
347    /// Test crs_is_lon_lat() detection for common lon/lat representations
348    #[test]
349    fn test_crs_is_lon_lat() -> ArrowResult<()> {
350        // EPSG:4326 as string should be detected
351        let metadata = Metadata::new(Some("EPSG:4326"), None);
352        assert!(metadata.crs_is_lon_lat());
353
354        // OGC:CRS84 as string should be detected
355        let metadata = Metadata::new(Some("OGC:CRS84"), None);
356        assert!(metadata.crs_is_lon_lat());
357
358        // A JSON object that reasonably looks like PROJJSON for EPSG:4326 should be detected
359        // detect "4326" as a string
360        let crs_json = r#"{"id":{"authority":"EPSG","code":"4326"}}"#;
361        let metadata = Metadata::new(Some(crs_json), None);
362        assert!(metadata.crs_is_lon_lat());
363
364        // detect 4326 as a number
365        let crs_json = r#"{"id":{"authority":"EPSG","code":4326}}"#;
366        let metadata = Metadata::new(Some(crs_json), None);
367        assert!(metadata.crs_is_lon_lat());
368
369        // A JSON object that reasonably looks like PROJJSON for OGC:CRS84 should be detected
370        let crs_json = r#"{"id":{"authority":"OGC","code":"CRS84"}}"#;
371        let metadata = Metadata::new(Some(crs_json), None);
372        assert!(metadata.crs_is_lon_lat());
373
374        // Other CRS values should NOT be detected as lon/lat
375        let metadata = Metadata::new(Some("srid:1234"), None);
376        assert!(!metadata.crs_is_lon_lat());
377
378        let metadata = Metadata::new(Some("EPSG:3857"), None);
379        assert!(!metadata.crs_is_lon_lat());
380
381        // None CRS should NOT be detected as lon/lat
382        let metadata = Metadata::new(None, None);
383        assert!(!metadata.crs_is_lon_lat());
384
385        Ok(())
386    }
387}