vortex_geo/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_array::arrow::ArrowSessionExt;
7use vortex_array::dtype::session::DTypeSessionExt;
8use vortex_array::scalar_fn::session::ScalarFnSessionExt;
9use vortex_session::VortexSession;
10
11use crate::extension::Point;
12use crate::extension::WellKnownBinary;
13use crate::scalar_fn::distance::GeoDistance;
14
15pub mod extension;
16pub mod scalar_fn;
17/// Set up a session with support for geospatial extension types, encodings and layouts.
18pub fn initialize(session: &VortexSession) {
19    // Register the geospatial extension types.
20    session.dtypes().register(WellKnownBinary);
21    session.arrow().register_exporter(Arc::new(WellKnownBinary));
22    session.arrow().register_importer(Arc::new(WellKnownBinary));
23    session.dtypes().register(Point);
24
25    // Register the geometry scalar functions.
26    session.scalar_fns().register(GeoDistance);
27}
28
29#[cfg(test)]
30mod tests {
31    use std::sync::Arc;
32    use std::sync::LazyLock;
33
34    use arrow_array::Array as _;
35    use arrow_array::ArrayRef as ArrowArrayRef;
36    use arrow_array::BinaryArray;
37    use arrow_array::BinaryViewArray;
38    use arrow_array::LargeBinaryArray;
39    use arrow_array::cast::AsArray;
40    use arrow_schema::DataType;
41    use arrow_schema::Field;
42    use arrow_schema::extension::ExtensionType as _;
43    use geo_traits::to_geo::ToGeoGeometry;
44    use geo_types::Coord;
45    use geo_types::Geometry;
46    use geo_types::LineString;
47    use geo_types::Polygon;
48    use geoarrow::datatypes::Crs;
49    use geoarrow::datatypes::Metadata;
50    use geoarrow::datatypes::WkbType;
51    use vortex_array::IntoArray;
52    use vortex_array::VortexSessionExecute;
53    use vortex_array::arrays::ExtensionArray;
54    use vortex_array::arrays::varbin::builder::VarBinBuilder;
55    use vortex_array::arrow::ArrowSessionExt;
56    use vortex_array::dtype::DType;
57    use vortex_array::dtype::Nullability;
58    use vortex_array::dtype::extension::ExtDType;
59    use vortex_array::dtype::extension::ExtVTable;
60    use vortex_array::session::ArraySession;
61    use vortex_error::VortexResult;
62    use vortex_error::vortex_err;
63    use vortex_session::VortexSession;
64    use wkb::writer::WriteOptions;
65
66    use crate::extension::GeoMetadata;
67    use crate::extension::WellKnownBinary;
68
69    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
70        let session = VortexSession::empty().with::<ArraySession>();
71        crate::initialize(&session);
72        session
73    });
74
75    #[test]
76    fn test_array() -> VortexResult<()> {
77        let mut execution_ctx = SESSION.create_execution_ctx();
78
79        let mut buf = Vec::new();
80
81        let mut builder = VarBinBuilder::<i32>::with_capacity(3);
82
83        let polygon = Geometry::Polygon(Polygon::new(
84            LineString::new(vec![
85                Coord::zero(),
86                Coord { x: 100.0, y: 0.0 },
87                Coord { x: 100.0, y: 100.0 },
88                Coord { x: 0.0, y: 100.0 },
89                Coord::zero(),
90            ]),
91            vec![],
92        ));
93
94        // We should always prefer to write little-endian, which is the default option.
95        wkb::writer::write_geometry(&mut buf, &polygon, &WriteOptions::default())
96            .map_err(|e| vortex_err!("writing WKB failed: {e}"))?;
97
98        // Push same polygon 3 times
99        builder.append_value(&buf);
100        builder.append_value(&buf);
101        builder.append_value(&buf);
102
103        let dtype = ExtDType::<WellKnownBinary>::try_new(
104            GeoMetadata {
105                crs: Some("EPSG:4326".to_string()),
106            },
107            DType::Binary(Nullability::NonNullable),
108        )?;
109
110        let array = builder.finish(DType::Binary(Nullability::NonNullable));
111        let array = ExtensionArray::new(dtype.clone().erased(), array.into_array()).into_array();
112
113        for idx in 0..3 {
114            let geom = array.execute_scalar(idx, &mut execution_ctx)?;
115            let wkb = WellKnownBinary::unpack_native(&dtype, geom.value().unwrap())?;
116
117            assert_eq!(wkb.to_geometry(), polygon);
118        }
119
120        Ok(())
121    }
122
123    fn wkb_extension_array() -> VortexResult<(Vec<u8>, vortex_array::ArrayRef)> {
124        let polygon = Geometry::Polygon(Polygon::new(
125            LineString::new(vec![
126                Coord::zero(),
127                Coord { x: 100.0, y: 0.0 },
128                Coord { x: 100.0, y: 100.0 },
129                Coord { x: 0.0, y: 100.0 },
130                Coord::zero(),
131            ]),
132            vec![],
133        ));
134        let mut buf = Vec::new();
135        wkb::writer::write_geometry(&mut buf, &polygon, &WriteOptions::default())
136            .map_err(|e| vortex_err!("writing WKB failed: {e}"))?;
137
138        let mut builder = VarBinBuilder::<i32>::with_capacity(3);
139        builder.append_value(&buf);
140        builder.append_value(&buf);
141        builder.append_value(&buf);
142
143        let dtype = ExtDType::<WellKnownBinary>::try_new(
144            GeoMetadata {
145                crs: Some("EPSG:4326".to_string()),
146            },
147            DType::Binary(Nullability::NonNullable),
148        )?;
149        let storage = builder.finish(DType::Binary(Nullability::NonNullable));
150        let array = ExtensionArray::new(dtype.erased(), storage.into_array()).into_array();
151        Ok((buf, array))
152    }
153
154    #[test]
155    fn export_arrow_field_carries_wkb_extension() -> VortexResult<()> {
156        let (_, array) = wkb_extension_array()?;
157        let field = SESSION.arrow().to_arrow_field("geom", array.dtype())?;
158        assert_eq!(field.extension_type_name(), Some(WkbType::NAME));
159        // Vortex's canonical mapping of `DType::Binary` is Arrow's `BinaryView`.
160        assert_eq!(field.data_type(), &DataType::BinaryView);
161        Ok(())
162    }
163
164    #[test]
165    fn execute_arrow_exports_wkb_to_binary() -> VortexResult<()> {
166        let mut ctx = SESSION.create_execution_ctx();
167        let (wkb_bytes, array) = wkb_extension_array()?;
168
169        let field = Field::new("geom", DataType::Binary, false)
170            .with_extension_type(WkbType::new(Default::default()));
171        let exported = SESSION
172            .arrow()
173            .execute_arrow(array, Some(&field), &mut ctx)?;
174
175        assert_eq!(exported.data_type(), &DataType::Binary);
176        let binary = exported.as_binary::<i32>();
177        assert_eq!(binary.len(), 3);
178        for idx in 0..3 {
179            assert_eq!(binary.value(idx), wkb_bytes.as_slice());
180        }
181        Ok(())
182    }
183
184    #[test]
185    fn execute_arrow_exports_wkb_to_large_binary() -> VortexResult<()> {
186        let mut ctx = SESSION.create_execution_ctx();
187        let (wkb_bytes, array) = wkb_extension_array()?;
188
189        let field = Field::new("geom", DataType::LargeBinary, false)
190            .with_extension_type(WkbType::new(Default::default()));
191        let exported = SESSION
192            .arrow()
193            .execute_arrow(array, Some(&field), &mut ctx)?;
194
195        assert_eq!(exported.data_type(), &DataType::LargeBinary);
196        let binary = exported.as_binary::<i64>();
197        assert_eq!(binary.len(), 3);
198        for idx in 0..3 {
199            assert_eq!(binary.value(idx), wkb_bytes.as_slice());
200        }
201        Ok(())
202    }
203
204    #[test]
205    fn execute_arrow_exports_wkb_to_binary_view() -> VortexResult<()> {
206        let mut ctx = SESSION.create_execution_ctx();
207        let (wkb_bytes, array) = wkb_extension_array()?;
208
209        let field = Field::new("geom", DataType::BinaryView, false)
210            .with_extension_type(WkbType::new(Default::default()));
211        let exported = SESSION
212            .arrow()
213            .execute_arrow(array, Some(&field), &mut ctx)?;
214
215        assert_eq!(exported.data_type(), &DataType::BinaryView);
216        let binary = exported.as_binary_view();
217        assert_eq!(binary.len(), 3);
218        for idx in 0..3 {
219            assert_eq!(binary.value(idx), wkb_bytes.as_slice());
220        }
221        Ok(())
222    }
223
224    fn wkb_field(name: &str, data_type: DataType, nullable: bool, crs: Option<&str>) -> Field {
225        let crs = crs
226            .map(|crs| Crs::from_unknown_crs_type(crs.to_string()))
227            .unwrap_or_default();
228        let wkb_type = WkbType::new(Arc::new(Metadata::new(crs, None)));
229        Field::new(name, data_type, nullable).with_extension_type(wkb_type)
230    }
231
232    fn assert_imported_wkb_dtype(dtype: &DType, expected_crs: Option<&str>, nullable: bool) {
233        let DType::Extension(ext) = dtype else {
234            panic!("expected Extension dtype, got {dtype}");
235        };
236        assert!(ext.is::<WellKnownBinary>());
237        assert_eq!(ext.storage_dtype(), &DType::Binary(nullable.into()));
238        let geo = ext.metadata::<WellKnownBinary>();
239        assert_eq!(geo.crs.as_deref(), expected_crs);
240    }
241
242    #[test]
243    fn import_arrow_field_recovers_wkb_extension() -> VortexResult<()> {
244        let field = wkb_field("geom", DataType::Binary, false, Some("EPSG:4326"));
245        let dtype = SESSION.arrow().from_arrow_field(&field)?;
246        assert_imported_wkb_dtype(&dtype, Some("EPSG:4326"), false);
247        Ok(())
248    }
249
250    #[test]
251    fn import_arrow_field_without_crs() -> VortexResult<()> {
252        let field = wkb_field("geom", DataType::BinaryView, true, None);
253        let dtype = SESSION.arrow().from_arrow_field(&field)?;
254        assert_imported_wkb_dtype(&dtype, None, true);
255        Ok(())
256    }
257
258    #[test]
259    fn import_arrow_array_from_binary() -> VortexResult<()> {
260        let mut ctx = SESSION.create_execution_ctx();
261        let (wkb_bytes, _) = wkb_extension_array()?;
262        let arrow: ArrowArrayRef = Arc::new(BinaryArray::from_iter_values([
263            wkb_bytes.as_slice(),
264            wkb_bytes.as_slice(),
265            wkb_bytes.as_slice(),
266        ]));
267        let field = wkb_field("geom", DataType::Binary, false, Some("EPSG:4326"));
268
269        let imported = SESSION.arrow().from_arrow_array(arrow, &field)?;
270        assert_imported_wkb_dtype(imported.dtype(), Some("EPSG:4326"), false);
271
272        for idx in 0..3 {
273            let geom = imported.execute_scalar(idx, &mut ctx)?;
274            assert_eq!(geom.value().unwrap().as_binary().as_slice(), wkb_bytes);
275        }
276        Ok(())
277    }
278
279    #[test]
280    fn import_arrow_array_from_large_binary() -> VortexResult<()> {
281        let (wkb_bytes, _) = wkb_extension_array()?;
282        let arrow: ArrowArrayRef = Arc::new(LargeBinaryArray::from_iter_values([
283            wkb_bytes.as_slice(),
284            wkb_bytes.as_slice(),
285        ]));
286        let field = wkb_field("geom", DataType::LargeBinary, false, Some("EPSG:4326"));
287
288        let imported = SESSION.arrow().from_arrow_array(arrow, &field)?;
289        assert_imported_wkb_dtype(imported.dtype(), Some("EPSG:4326"), false);
290        assert_eq!(imported.len(), 2);
291        Ok(())
292    }
293
294    #[test]
295    fn import_arrow_array_from_binary_view() -> VortexResult<()> {
296        let (wkb_bytes, _) = wkb_extension_array()?;
297        let arrow: ArrowArrayRef = Arc::new(BinaryViewArray::from_iter_values([
298            wkb_bytes.as_slice(),
299            wkb_bytes.as_slice(),
300            wkb_bytes.as_slice(),
301            wkb_bytes.as_slice(),
302        ]));
303        let field = wkb_field("geom", DataType::BinaryView, false, None);
304
305        let imported = SESSION.arrow().from_arrow_array(arrow, &field)?;
306        assert_imported_wkb_dtype(imported.dtype(), None, false);
307        assert_eq!(imported.len(), 4);
308        Ok(())
309    }
310
311    #[test]
312    fn import_arrow_array_roundtrips_through_export() -> VortexResult<()> {
313        let mut ctx = SESSION.create_execution_ctx();
314        let (wkb_bytes, original) = wkb_extension_array()?;
315
316        let field = Field::new("geom", DataType::BinaryView, false)
317            .with_extension_type(WkbType::new(Default::default()));
318        let exported = SESSION
319            .arrow()
320            .execute_arrow(original, Some(&field), &mut ctx)?;
321        let arrow_field = wkb_field("geom", DataType::BinaryView, false, Some("EPSG:4326"));
322        let reimported = SESSION.arrow().from_arrow_array(exported, &arrow_field)?;
323
324        assert_imported_wkb_dtype(reimported.dtype(), Some("EPSG:4326"), false);
325        for idx in 0..3 {
326            let geom = reimported.execute_scalar(idx, &mut ctx)?;
327            assert_eq!(geom.value().unwrap().as_binary().as_slice(), wkb_bytes);
328        }
329        Ok(())
330    }
331}
vortex_geo/lib.rs

vortex_geo/
lib.rs
