nodedb_spatial/rtree/
bulk_load.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Sort-Tile-Recursive (STR) bulk loading for R-tree.
4
5use nodedb_mem::ScopedMemory;
6
7use super::node::{ChildRef, INTERNAL_CAPACITY, LEAF_CAPACITY, Node, NodeKind, RTreeEntry};
8use super::tree::RTree;
9
10impl RTree {
11    /// Bulk load entries using Sort-Tile-Recursive packing, tracked against `memory`.
12    ///
13    /// More efficient than repeated single inserts for large datasets.
14    /// Produces better packing (less overlap between nodes).
15    pub fn bulk_load(entries: Vec<RTreeEntry>, memory: ScopedMemory) -> Self {
16        if entries.is_empty() {
17            return Self::new(memory);
18        }
19        let len = entries.len();
20
21        // Reserve budget for nodes vec (approximately len / LEAF_CAPACITY nodes,
22        // each holding LEAF_CAPACITY RTreeEntry slots). Best-effort: budget
23        // pressure is a backpressure signal, not a hard gate on bulk load.
24        let node_count = len.div_ceil(LEAF_CAPACITY) * 2; // leaves + internals
25        let bytes = node_count
26            * (std::mem::size_of::<Node>() + LEAF_CAPACITY * std::mem::size_of::<RTreeEntry>());
27        let _guard = memory.reserve(bytes).ok();
28
29        let mut tree = Self {
30            nodes: Vec::new(),
31            root: 0,
32            len,
33            memory,
34        };
35        tree.root = str_pack(&mut tree.nodes, entries);
36        tree
37    }
38}
39
40/// Recursively pack entries into leaf nodes, then group into internal nodes.
41fn str_pack(nodes: &mut Vec<Node>, mut entries: Vec<RTreeEntry>) -> usize {
42    if entries.len() <= LEAF_CAPACITY {
43        let mut node = Node::new_leaf();
44        if let NodeKind::Leaf {
45            entries: ref mut leaf,
46        } = node.kind
47        {
48            *leaf = entries;
49        }
50        node.recompute_bbox();
51        let idx = nodes.len();
52        nodes.push(node);
53        return idx;
54    }
55
56    let num_leaves = entries.len().div_ceil(LEAF_CAPACITY);
57    let num_slices = (num_leaves as f64).sqrt().ceil() as usize;
58    let slice_size = entries.len().div_ceil(num_slices);
59
60    // Sort by longitude (X).
61    entries.sort_by(|a, b| {
62        let ca = (a.bbox.min_lng + a.bbox.max_lng) / 2.0;
63        let cb = (b.bbox.min_lng + b.bbox.max_lng) / 2.0;
64        ca.partial_cmp(&cb).unwrap_or(std::cmp::Ordering::Equal)
65    });
66
67    let mut child_nodes: Vec<usize> = Vec::new();
68
69    for slice in entries.chunks(slice_size) {
70        let mut slice_vec: Vec<RTreeEntry> = slice.to_vec();
71        // Sort each slice by latitude (Y).
72        slice_vec.sort_by(|a, b| {
73            let ca = (a.bbox.min_lat + a.bbox.max_lat) / 2.0;
74            let cb = (b.bbox.min_lat + b.bbox.max_lat) / 2.0;
75            ca.partial_cmp(&cb).unwrap_or(std::cmp::Ordering::Equal)
76        });
77
78        for chunk in slice_vec.chunks(LEAF_CAPACITY) {
79            let child_idx = str_pack(nodes, chunk.to_vec());
80            child_nodes.push(child_idx);
81        }
82    }
83
84    pack_internal(nodes, child_nodes)
85}
86
87/// Recursively group child node indices into internal nodes.
88fn pack_internal(nodes: &mut Vec<Node>, child_indices: Vec<usize>) -> usize {
89    if child_indices.len() <= INTERNAL_CAPACITY {
90        let level = nodes[child_indices[0]].level + 1;
91        let mut node = Node::new_internal(level);
92        if let NodeKind::Internal { children } = &mut node.kind {
93            for idx in child_indices {
94                children.push(ChildRef {
95                    bbox: nodes[idx].bbox,
96                    node_idx: idx,
97                });
98            }
99        }
100        node.recompute_bbox();
101        let idx = nodes.len();
102        nodes.push(node);
103        return idx;
104    }
105
106    let mut new_children = Vec::new();
107    for chunk in child_indices.chunks(INTERNAL_CAPACITY) {
108        let level = nodes[chunk[0]].level + 1;
109        let mut node = Node::new_internal(level);
110        if let NodeKind::Internal { children } = &mut node.kind {
111            for &idx in chunk {
112                children.push(ChildRef {
113                    bbox: nodes[idx].bbox,
114                    node_idx: idx,
115                });
116            }
117        }
118        node.recompute_bbox();
119        let idx = nodes.len();
120        nodes.push(node);
121        new_children.push(idx);
122    }
123
124    pack_internal(nodes, new_children)
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use crate::test_support::test_memory;
131    use nodedb_types::BoundingBox;
132
133    fn make_entry(id: u64, lng: f64, lat: f64) -> RTreeEntry {
134        RTreeEntry {
135            id,
136            bbox: BoundingBox::from_point(lng, lat),
137        }
138    }
139
140    #[test]
141    fn bulk_load_500() {
142        let entries: Vec<RTreeEntry> = (0..500)
143            .map(|i| make_entry(i, (i as f64) * 0.1 - 25.0, (i as f64) * 0.06 - 15.0))
144            .collect();
145        let tree = RTree::bulk_load(entries, test_memory());
146        assert_eq!(tree.len(), 500);
147
148        let all = tree.search(&BoundingBox::new(-180.0, -90.0, 180.0, 90.0));
149        assert_eq!(all.len(), 500);
150
151        let subset = tree.search(&BoundingBox::new(-5.0, -5.0, 5.0, 5.0));
152        assert!(!subset.is_empty());
153    }
154
155    #[test]
156    fn bulk_load_small() {
157        let entries = vec![make_entry(1, 0.0, 0.0), make_entry(2, 1.0, 1.0)];
158        let tree = RTree::bulk_load(entries, test_memory());
159        assert_eq!(tree.len(), 2);
160    }
161
162    #[test]
163    fn bulk_load_empty() {
164        let tree = RTree::bulk_load(Vec::new(), test_memory());
165        assert!(tree.is_empty());
166    }
167}
nodedb_spatial/rtree/bulk_load.rs

nodedb_spatial/rtree/
bulk_load.rs
