nodedb_lite/query/spatial_ops/
reads.rs1use nodedb_physical::physical_plan::SpatialPredicate;
5use nodedb_query::scan_filter::ScanFilter;
6use nodedb_spatial::predicates::{st_contains, st_dwithin, st_intersects, st_within};
7use nodedb_types::geometry::Geometry;
8use nodedb_types::result::QueryResult;
9use nodedb_types::value::Value;
10use nodedb_types::{BoundingBox, Surrogate, SurrogateBitmap, geometry_bbox};
11
12use crate::error::LiteError;
13use crate::query::engine::LiteQueryEngine;
14use crate::storage::engine::StorageEngine;
15
16pub struct ScanParams {
18 pub collection: String,
19 pub field: String,
20 pub predicate: SpatialPredicate,
21 pub query_geometry: Geometry,
22 pub distance_meters: f64,
23 pub attribute_filters: Vec<u8>,
24 pub limit: usize,
25 pub projection: Vec<String>,
26 pub rls_filters: Vec<u8>,
27 pub prefilter: Option<SurrogateBitmap>,
28}
29
30pub fn spatial_scan<S: StorageEngine>(
33 engine: &LiteQueryEngine<S>,
34 params: ScanParams,
35) -> Result<QueryResult, LiteError> {
36 let ScanParams {
37 collection,
38 field,
39 predicate,
40 query_geometry,
41 distance_meters,
42 attribute_filters,
43 limit,
44 projection,
45 rls_filters,
46 prefilter,
47 } = params;
48
49 let bbox: BoundingBox = match predicate {
51 SpatialPredicate::DWithin => geometry_bbox(&query_geometry).expand_meters(distance_meters),
52 SpatialPredicate::Contains | SpatialPredicate::Intersects | SpatialPredicate::Within => {
53 geometry_bbox(&query_geometry)
54 }
55 };
56
57 let attr_filters: Vec<ScanFilter> = if attribute_filters.is_empty() {
59 Vec::new()
60 } else {
61 zerompk::from_msgpack(&attribute_filters).map_err(|e| LiteError::Serialization {
62 detail: format!("decode attribute_filters: {e}"),
63 })?
64 };
65
66 let rls: Vec<ScanFilter> = if rls_filters.is_empty() {
68 Vec::new()
69 } else {
70 zerompk::from_msgpack(&rls_filters).map_err(|e| LiteError::Serialization {
71 detail: format!("decode rls_filters: {e}"),
72 })?
73 };
74
75 let spatial_guard = engine.spatial.lock().map_err(|_| LiteError::LockPoisoned)?;
77 let candidate_entries: Vec<u64> = spatial_guard
78 .search(&collection, &field, &bbox)
79 .into_iter()
80 .map(|e| e.id)
81 .collect();
82
83 let candidates: Vec<(String, u64)> = candidate_entries
85 .into_iter()
86 .filter_map(|eid| {
87 spatial_guard
88 .doc_id_for_entry(eid)
89 .map(|doc_id| (doc_id.to_string(), eid))
90 })
91 .collect();
92 drop(spatial_guard);
93
94 let crdt_guard = engine.crdt.lock().map_err(|_| LiteError::LockPoisoned)?;
95
96 let mut rows: Vec<Vec<Value>> = Vec::new();
97
98 for (doc_id, _eid) in candidates {
99 let surrogate_val: u32 = doc_id.parse::<u32>().unwrap_or_default();
103 let surrogate = Surrogate(surrogate_val);
104
105 if let Some(ref pf) = prefilter
107 && !pf.contains(surrogate)
108 {
109 continue;
110 }
111
112 let loro_val = match crdt_guard.read(&collection, &doc_id) {
115 Some(v) => v,
116 None => continue,
117 };
118 let doc = crate::nodedb::convert::loro_value_to_document(&doc_id, &loro_val);
119
120 let candidate_geom: Geometry = match doc.fields.get(&field) {
122 Some(Value::String(s)) => match sonic_rs::from_str::<Geometry>(s) {
123 Ok(g) => g,
124 Err(_) => continue,
125 },
126 Some(Value::Geometry(g)) => g.clone(),
127 _ => continue,
128 };
129
130 let passes = match predicate {
131 SpatialPredicate::DWithin => {
132 st_dwithin(&candidate_geom, &query_geometry, distance_meters)
133 }
134 SpatialPredicate::Contains => st_contains(&query_geometry, &candidate_geom),
135 SpatialPredicate::Intersects => st_intersects(&candidate_geom, &query_geometry),
136 SpatialPredicate::Within => st_within(&candidate_geom, &query_geometry),
137 };
138 if !passes {
139 continue;
140 }
141
142 let doc_value = Value::Object(doc.fields.clone());
144
145 if !ScanFilter::all_match_value(&attr_filters, &doc_value)? {
147 continue;
148 }
149
150 if !ScanFilter::all_match_value(&rls, &doc_value)? {
152 continue;
153 }
154
155 let row = if projection.is_empty() {
157 let msgpack =
159 zerompk::to_msgpack_vec(&doc_value).map_err(|e| LiteError::Serialization {
160 detail: format!("serialize spatial doc: {e}"),
161 })?;
162 vec![Value::String(doc_id), Value::Bytes(msgpack)]
163 } else {
164 projection
165 .iter()
166 .map(|col| {
167 if col == "id" || col == "_id" {
168 Value::String(doc_id.clone())
169 } else {
170 doc.fields.get(col).cloned().unwrap_or(Value::Null)
171 }
172 })
173 .collect()
174 };
175 rows.push(row);
176
177 if rows.len() >= limit && limit > 0 {
178 break;
179 }
180 }
181 drop(crdt_guard);
182
183 let columns = if projection.is_empty() {
184 vec!["id".to_string(), "data".to_string()]
185 } else {
186 projection
187 };
188
189 Ok(QueryResult {
190 columns,
191 rows,
192 rows_affected: 0,
193 command: None,
194 })
195}
196
197#[cfg(test)]
198mod tests {
199 use nodedb_types::BoundingBox;
200 use nodedb_types::geometry::Geometry;
201
202 use crate::engine::spatial::SpatialIndexManager;
203
204 #[test]
205 fn dwithin_bbox_expansion() {
206 let governor = crate::query::engine::test_governor();
208 let mut mgr = SpatialIndexManager::new(crate::query::engine::test_scoped_memory(
209 &governor,
210 nodedb_mem::EngineId::Spatial,
211 ));
212 mgr.index_document("col", "geom", "1", &Geometry::point(10.0, 20.0));
214 mgr.index_document("col", "geom", "2", &Geometry::point(50.0, 50.0));
216
217 let query_geom = Geometry::point(10.0, 20.0);
219 let bbox = nodedb_types::geometry_bbox(&query_geom).expand_meters(100_000.0);
220 let hits = mgr.search("col", "geom", &bbox);
221 assert!(!hits.is_empty());
223 }
224
225 #[test]
226 fn intersects_bbox() {
227 let governor = crate::query::engine::test_governor();
228 let mut mgr = SpatialIndexManager::new(crate::query::engine::test_scoped_memory(
229 &governor,
230 nodedb_mem::EngineId::Spatial,
231 ));
232 mgr.index_document("col", "geom", "3", &Geometry::point(5.0, 5.0));
233 mgr.index_document("col", "geom", "4", &Geometry::point(90.0, 80.0));
234
235 let bbox = BoundingBox::new(0.0, 0.0, 10.0, 10.0);
237 let hits = mgr.search("col", "geom", &bbox);
238 assert_eq!(hits.len(), 1);
239
240 let entry_id = hits[0].id;
242 let doc_id = mgr.doc_id_for_entry(entry_id).unwrap();
243 assert_eq!(doc_id, "3");
244 }
245}