1use arrow_schema::{ArrowError, DataType, extension::ExtensionType};
19use serde::{Deserialize, Serialize};
20
21#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
30pub enum Hint {
31 Geometry,
33 Geography,
35}
36
37#[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Serialize, Deserialize)]
39#[serde(rename_all = "lowercase")]
40pub enum Edges {
41 #[default]
43 Spherical,
44 Vincenty,
46 Thomas,
48 Andoyer,
50 Karney,
52}
53
54#[derive(Clone, Debug, Default, Serialize, Deserialize)]
56pub struct Metadata {
57 #[serde(skip_serializing_if = "Option::is_none")]
61 pub crs: Option<serde_json::Value>,
62 #[serde(rename = "edges", skip_serializing_if = "Option::is_none")]
64 pub algorithm: Option<Edges>,
65}
66
67impl Metadata {
68 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 pub fn type_hint(&self) -> Hint {
83 match &self.algorithm {
84 Some(_) => Hint::Geography,
85 None => Hint::Geometry,
86 }
87 }
88
89 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#[derive(Debug, Default)]
125pub struct WkbType(Metadata);
126
127impl WkbType {
128 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]
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]
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]
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 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 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]
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]
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]
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]
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]
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]
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]
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 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]
333 fn test_extension_type_support() -> ArrowResult<()> {
334 let wkb = WkbType::default();
335 wkb.supports_data_type(&DataType::Binary)?;
337 wkb.supports_data_type(&DataType::LargeBinary)?;
338 wkb.supports_data_type(&DataType::BinaryView)?;
339
340 let result = wkb.supports_data_type(&DataType::Utf8);
342 assert!(matches!(result, Err(ArrowError::InvalidArgumentError(_))));
343
344 Ok(())
345 }
346
347 #[test]
349 fn test_crs_is_lon_lat() -> ArrowResult<()> {
350 let metadata = Metadata::new(Some("EPSG:4326"), None);
352 assert!(metadata.crs_is_lon_lat());
353
354 let metadata = Metadata::new(Some("OGC:CRS84"), None);
356 assert!(metadata.crs_is_lon_lat());
357
358 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 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 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 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 let metadata = Metadata::new(None, None);
383 assert!(!metadata.crs_is_lon_lat());
384
385 Ok(())
386 }
387}