Skip to main content

nodedb_lite/query/spatial_ops/
reads.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Spatial read operations: Scan with OGC predicate refinement.
3
4use 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
16/// Parameters for `SpatialOp::Scan`.
17pub 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
30/// Execute a spatial scan: R-tree candidate generation → prefilter → OGC
31/// refinement → attribute filters → RLS → projection + limit.
32pub 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    // Expand bbox for DWithin; use exact bbox for all other predicates.
50    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    // Decode attribute filters (empty bytes → no filters).
58    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    // Decode RLS filters.
67    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    // Candidate entry IDs from R-tree range search.
76    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    // Resolve each entry_id to doc_id while holding the lock.
84    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        // Parse surrogate from doc_id string.
100        // Non-numeric doc_id — surrogate was not assigned via SpatialOp::Insert.
101        // Fall back to treating the doc_id as surrogate 0 for prefilter purposes.
102        let surrogate_val: u32 = doc_id.parse::<u32>().unwrap_or_default();
103        let surrogate = Surrogate(surrogate_val);
104
105        // Apply cross-engine prefilter bitmap.
106        if let Some(ref pf) = prefilter
107            && !pf.contains(surrogate)
108        {
109            continue;
110        }
111
112        // Fetch the document from the CRDT store for geometry refinement and
113        // attribute filtering.
114        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        // OGC predicate refinement: extract candidate geometry from the document.
121        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        // Build a Value::Object from the document fields for filter evaluation.
143        let doc_value = Value::Object(doc.fields.clone());
144
145        // Apply attribute filters.
146        if !ScanFilter::all_match_value(&attr_filters, &doc_value)? {
147            continue;
148        }
149
150        // Apply RLS filters.
151        if !ScanFilter::all_match_value(&rls, &doc_value)? {
152            continue;
153        }
154
155        // Build result row applying projection.
156        let row = if projection.is_empty() {
157            // Return id + full document.
158            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        // Verify that DWithin expands the bbox correctly via expand_meters.
207        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        // Point at (10.0, 20.0).
213        mgr.index_document("col", "geom", "1", &Geometry::point(10.0, 20.0));
214        // Point far away at (50.0, 50.0).
215        mgr.index_document("col", "geom", "2", &Geometry::point(50.0, 50.0));
216
217        // Query geometry at (10.0, 20.0), large distance should include nearby point.
218        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        // Only the nearby point should be within 100 km.
222        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        // Query bbox covers only the first point.
236        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        // Resolve doc_id via entry_to_doc inverse map.
241        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}