nodedb_spatial/
spatial_join.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Spatial join strategy: R-tree probe join.
4//!
5//! When two collections are joined on a spatial predicate
6//! (`ST_Intersects(a.geom, b.geom)`), the naive approach is nested-loop
7//! O(N*M). With an R-tree index on one side, we get O(N * log M):
8//!
9//! 1. Build R-tree on the smaller collection (or use existing index)
10//! 2. For each geometry in the larger collection, R-tree range search
11//! 3. For each candidate pair, apply exact predicate
12//!
13//! This module provides the join logic; the planner decides which side
14//! to index based on collection cardinality.
15
16use nodedb_mem::ScopedMemory;
17use nodedb_types::geometry::Geometry;
18use nodedb_types::{BoundingBox, geometry_bbox};
19
20use crate::predicates;
21use crate::rtree::{RTree, RTreeEntry};
22
23/// Result of a spatial join between two collections.
24pub struct SpatialJoinResult {
25    /// Matched pairs: (left_entry_id, right_entry_id).
26    pub pairs: Vec<(u64, u64)>,
27    /// Number of R-tree probes performed.
28    pub probes: usize,
29    /// Number of exact predicate evaluations (after R-tree filter).
30    pub exact_evals: usize,
31}
32
33/// Execute a spatial join using R-tree probe.
34///
35/// `indexed_side`: R-tree built on one collection.
36/// `probe_side`: entries from the other collection to probe against.
37/// `get_geometry`: callback to retrieve the full geometry for an entry ID
38///   (needed for exact predicate evaluation after R-tree bbox filter).
39/// `predicate`: which spatial predicate to apply (intersects, contains, etc.).
40pub fn spatial_join(
41    indexed_side: &RTree,
42    probe_entries: &[(u64, BoundingBox)],
43    get_indexed_geom: &dyn Fn(u64) -> Option<Geometry>,
44    get_probe_geom: &dyn Fn(u64) -> Option<Geometry>,
45    predicate: SpatialJoinPredicate,
46) -> SpatialJoinResult {
47    let mut pairs = Vec::new();
48    let mut probes = 0;
49    let mut exact_evals = 0;
50
51    for &(probe_id, ref probe_bbox) in probe_entries {
52        // R-tree range search: find indexed entries whose bbox intersects probe bbox.
53        let candidates = indexed_side.search(probe_bbox);
54        probes += 1;
55
56        for candidate in &candidates {
57            // Exact predicate evaluation.
58            let Some(indexed_geom) = get_indexed_geom(candidate.id) else {
59                continue;
60            };
61            let Some(probe_geom) = get_probe_geom(probe_id) else {
62                continue;
63            };
64            exact_evals += 1;
65
66            let matches = match predicate {
67                SpatialJoinPredicate::Intersects => {
68                    predicates::st_intersects(&probe_geom, &indexed_geom)
69                }
70                SpatialJoinPredicate::Contains => {
71                    predicates::st_contains(&probe_geom, &indexed_geom)
72                }
73                SpatialJoinPredicate::Within => predicates::st_within(&probe_geom, &indexed_geom),
74                SpatialJoinPredicate::DWithin(dist) => {
75                    predicates::st_dwithin(&probe_geom, &indexed_geom, dist)
76                }
77            };
78
79            if matches {
80                pairs.push((probe_id, candidate.id));
81            }
82        }
83    }
84
85    SpatialJoinResult {
86        pairs,
87        probes,
88        exact_evals,
89    }
90}
91
92/// Build an R-tree from a list of (entry_id, geometry) pairs for join.
93pub fn build_join_index(entries: &[(u64, Geometry)], memory: ScopedMemory) -> RTree {
94    let rtree_entries: Vec<RTreeEntry> = entries
95        .iter()
96        .map(|(id, geom)| RTreeEntry {
97            id: *id,
98            bbox: geometry_bbox(geom),
99        })
100        .collect();
101    RTree::bulk_load(rtree_entries, memory)
102}
103
104/// Which spatial predicate to use for the join.
105#[derive(Debug, Clone, Copy)]
106#[non_exhaustive]
107pub enum SpatialJoinPredicate {
108    Intersects,
109    Contains,
110    Within,
111    DWithin(f64),
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use crate::test_support::test_memory;
118
119    #[test]
120    fn join_overlapping_squares() {
121        // Left side: 5 squares at positions (0,0), (5,5), (10,10), (15,15), (20,20).
122        let left: Vec<(u64, Geometry)> = (0..5)
123            .map(|i| {
124                let base = (i * 5) as f64;
125                (
126                    i as u64,
127                    Geometry::polygon(vec![vec![
128                        [base, base],
129                        [base + 4.0, base],
130                        [base + 4.0, base + 4.0],
131                        [base, base + 4.0],
132                        [base, base],
133                    ]]),
134                )
135            })
136            .collect();
137
138        // Right side: 5 squares offset by 2.
139        let right: Vec<(u64, Geometry)> = (0..5)
140            .map(|i| {
141                let base = (i * 5 + 2) as f64;
142                (
143                    100 + i as u64,
144                    Geometry::polygon(vec![vec![
145                        [base, base],
146                        [base + 4.0, base],
147                        [base + 4.0, base + 4.0],
148                        [base, base + 4.0],
149                        [base, base],
150                    ]]),
151                )
152            })
153            .collect();
154
155        // Build index on left side.
156        let index = build_join_index(&left, test_memory());
157
158        // Probe with right side.
159        let probe_entries: Vec<(u64, BoundingBox)> = right
160            .iter()
161            .map(|(id, geom)| (*id, geometry_bbox(geom)))
162            .collect();
163
164        let left_map: std::collections::HashMap<u64, Geometry> = left.into_iter().collect();
165        let right_map: std::collections::HashMap<u64, Geometry> = right.into_iter().collect();
166
167        let result = spatial_join(
168            &index,
169            &probe_entries,
170            &|id| left_map.get(&id).cloned(),
171            &|id| right_map.get(&id).cloned(),
172            SpatialJoinPredicate::Intersects,
173        );
174
175        // Adjacent overlapping squares should produce matches.
176        assert!(!result.pairs.is_empty(), "expected some join matches");
177        assert!(result.probes == 5); // One probe per right entry.
178    }
179
180    #[test]
181    fn join_no_overlap() {
182        let left = vec![(
183            1u64,
184            Geometry::polygon(vec![vec![
185                [0.0, 0.0],
186                [1.0, 0.0],
187                [1.0, 1.0],
188                [0.0, 1.0],
189                [0.0, 0.0],
190            ]]),
191        )];
192        let right = vec![(
193            100u64,
194            Geometry::polygon(vec![vec![
195                [50.0, 50.0],
196                [51.0, 50.0],
197                [51.0, 51.0],
198                [50.0, 51.0],
199                [50.0, 50.0],
200            ]]),
201        )];
202
203        let index = build_join_index(&left, test_memory());
204        let probes: Vec<(u64, BoundingBox)> = right
205            .iter()
206            .map(|(id, g)| (*id, geometry_bbox(g)))
207            .collect();
208
209        let left_map: std::collections::HashMap<u64, Geometry> = left.into_iter().collect();
210        let right_map: std::collections::HashMap<u64, Geometry> = right.into_iter().collect();
211
212        let result = spatial_join(
213            &index,
214            &probes,
215            &|id| left_map.get(&id).cloned(),
216            &|id| right_map.get(&id).cloned(),
217            SpatialJoinPredicate::Intersects,
218        );
219
220        assert!(result.pairs.is_empty());
221    }
222
223    #[test]
224    fn join_with_dwithin() {
225        let left = vec![(1u64, Geometry::point(0.0, 0.0))];
226        let right = vec![(100u64, Geometry::point(0.001, 0.0))]; // ~111m away
227
228        let index = build_join_index(&left, test_memory());
229        let probes: Vec<(u64, BoundingBox)> = right
230            .iter()
231            .map(|(id, g)| {
232                // Expand bbox by distance for R-tree search.
233                (*id, geometry_bbox(g).expand_meters(500.0))
234            })
235            .collect();
236
237        let left_map: std::collections::HashMap<u64, Geometry> = left.into_iter().collect();
238        let right_map: std::collections::HashMap<u64, Geometry> = right.into_iter().collect();
239
240        let result = spatial_join(
241            &index,
242            &probes,
243            &|id| left_map.get(&id).cloned(),
244            &|id| right_map.get(&id).cloned(),
245            SpatialJoinPredicate::DWithin(500.0), // 500m
246        );
247
248        assert_eq!(result.pairs.len(), 1);
249    }
250}
nodedb_spatial/spatial_join.rs

nodedb_spatial/
spatial_join.rs
