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::envelope::GeoEnvelope;
27use crate::scalar_fn::intersects::GeoIntersects;
28
29pub mod aggregate_fn;
30pub mod extension;
31pub mod prune;
32pub mod scalar_fn;
33#[cfg(any(test, feature = "_test-harness"))]
34pub mod test_harness;
35#[cfg(test)]
36mod tests;
37
38/// Set up a session with support for geospatial extension types, encodings and layouts.
39pub fn initialize(session: &VortexSession) {
40    // Register the geospatial extension types.
41    session.dtypes().register(WellKnownBinary);
42    session.arrow().register_exporter(Arc::new(WellKnownBinary));
43    session.arrow().register_importer(Arc::new(WellKnownBinary));
44    session.dtypes().register(Point);
45    session.arrow().register_exporter(Arc::new(Point));
46    session.arrow().register_importer(Arc::new(Point));
47    session.dtypes().register(LineString);
48    session.arrow().register_exporter(Arc::new(LineString));
49    session.arrow().register_importer(Arc::new(LineString));
50    session.dtypes().register(MultiPoint);
51    session.arrow().register_exporter(Arc::new(MultiPoint));
52    session.arrow().register_importer(Arc::new(MultiPoint));
53    session.dtypes().register(Polygon);
54    session.arrow().register_exporter(Arc::new(Polygon));
55    session.arrow().register_importer(Arc::new(Polygon));
56    session.dtypes().register(MultiLineString);
57    session.arrow().register_exporter(Arc::new(MultiLineString));
58    session.arrow().register_importer(Arc::new(MultiLineString));
59    session.dtypes().register(MultiPolygon);
60    session.arrow().register_exporter(Arc::new(MultiPolygon));
61    session.arrow().register_importer(Arc::new(MultiPolygon));
62    session.dtypes().register(Rect);
63    session.arrow().register_exporter(Arc::new(Rect));
64    session.arrow().register_importer(Arc::new(Rect));
65
66    // Register the geometry scalar functions.
67    session.scalar_fns().register(GeoEnvelope);
68    session.scalar_fns().register(GeoContains);
69    session.scalar_fns().register(GeoDistance);
70    session.scalar_fns().register(GeoIntersects);
71
72    // The axis-aligned bounding-box (AABB) aggregate; self-declares as a per-chunk zone stat for
73    // geometry columns.
74    session.aggregate_fns().register(GeometryAabb);
75
76    // Register the spatial pruning rules that use that AABB.
77    session.stats().register_rewrite(GeoDistancePrune);
78    session.stats().register_rewrite(GeoIntersectsPrune);
79}
vortex_geo/lib.rs

vortex_geo/
lib.rs
