Skip to main content

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::aggregate_fn::session::AggregateFnSessionExt;
7use vortex_array::dtype::session::DTypeSessionExt;
8use vortex_array::scalar_fn::session::ScalarFnSessionExt;
9use vortex_array::stats::session::StatsSessionExt;
10use vortex_arrow::ArrowSessionExt;
11use vortex_session::VortexSession;
12
13use crate::aggregate_fn::GeometryAabb;
14use crate::extension::LineString;
15use crate::extension::MultiLineString;
16use crate::extension::MultiPoint;
17use crate::extension::MultiPolygon;
18use crate::extension::Point;
19use crate::extension::Polygon;
20use crate::extension::Rect;
21use crate::extension::WellKnownBinary;
22use crate::prune::GeoDistancePrune;
23use crate::prune::GeoIntersectsPrune;
24use crate::scalar_fn::contains::GeoContains;
25use crate::scalar_fn::distance::GeoDistance;
26use crate::scalar_fn::intersects::GeoIntersects;
27
28pub mod aggregate_fn;
29pub mod extension;
30pub mod prune;
31pub mod scalar_fn;
32#[cfg(test)]
33mod test_harness;
34#[cfg(test)]
35mod tests;
36
37/// Set up a session with support for geospatial extension types, encodings and layouts.
38pub fn initialize(session: &VortexSession) {
39    // Register the geospatial extension types.
40    session.dtypes().register(WellKnownBinary);
41    session.arrow().register_exporter(Arc::new(WellKnownBinary));
42    session.arrow().register_importer(Arc::new(WellKnownBinary));
43    session.dtypes().register(Point);
44    session.arrow().register_exporter(Arc::new(Point));
45    session.arrow().register_importer(Arc::new(Point));
46    session.dtypes().register(LineString);
47    session.arrow().register_exporter(Arc::new(LineString));
48    session.arrow().register_importer(Arc::new(LineString));
49    session.dtypes().register(MultiPoint);
50    session.arrow().register_exporter(Arc::new(MultiPoint));
51    session.arrow().register_importer(Arc::new(MultiPoint));
52    session.dtypes().register(Polygon);
53    session.arrow().register_exporter(Arc::new(Polygon));
54    session.arrow().register_importer(Arc::new(Polygon));
55    session.dtypes().register(MultiLineString);
56    session.arrow().register_exporter(Arc::new(MultiLineString));
57    session.arrow().register_importer(Arc::new(MultiLineString));
58    session.dtypes().register(MultiPolygon);
59    session.arrow().register_exporter(Arc::new(MultiPolygon));
60    session.arrow().register_importer(Arc::new(MultiPolygon));
61    session.dtypes().register(Rect);
62    session.arrow().register_exporter(Arc::new(Rect));
63    session.arrow().register_importer(Arc::new(Rect));
64
65    // Register the geometry scalar functions.
66    session.scalar_fns().register(GeoContains);
67    session.scalar_fns().register(GeoDistance);
68    session.scalar_fns().register(GeoIntersects);
69
70    // The axis-aligned bounding-box (AABB) aggregate; self-declares as a per-chunk zone stat for
71    // geometry columns.
72    session.aggregate_fns().register(GeometryAabb);
73
74    // Register the spatial pruning rules that use that AABB.
75    session.stats().register_rewrite(GeoDistancePrune);
76    session.stats().register_rewrite(GeoIntersectsPrune);
77}