Skip to main content

nodedb_query/
geo_functions.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Geospatial SQL function evaluation.
4//!
5//! All ST_* predicates, geometry operations, and geo_* utility functions.
6//! Called from `functions::eval_function` for any geo-prefixed or ST_-prefixed name.
7
8use crate::value_ops::{to_value_number, value_to_f64};
9use nodedb_types::Value;
10
11/// Try to evaluate a geo/spatial function. Returns `Some(result)` if the
12/// function name matched, `None` if unrecognized (caller falls through).
13pub fn eval_geo_function(name: &str, args: &[Value]) -> Option<Value> {
14    let result = match name {
15        "geo_distance" | "haversine_distance" => {
16            let lng1 = num_arg(args, 0).unwrap_or(0.0);
17            let lat1 = num_arg(args, 1).unwrap_or(0.0);
18            let lng2 = num_arg(args, 2).unwrap_or(0.0);
19            let lat2 = num_arg(args, 3).unwrap_or(0.0);
20            to_value_number(nodedb_types::geometry::haversine_distance(
21                lng1, lat1, lng2, lat2,
22            ))
23        }
24        "geo_bearing" | "haversine_bearing" => {
25            let lng1 = num_arg(args, 0).unwrap_or(0.0);
26            let lat1 = num_arg(args, 1).unwrap_or(0.0);
27            let lng2 = num_arg(args, 2).unwrap_or(0.0);
28            let lat2 = num_arg(args, 3).unwrap_or(0.0);
29            to_value_number(nodedb_types::geometry::haversine_bearing(
30                lng1, lat1, lng2, lat2,
31            ))
32        }
33        "geo_point" | "st_point" => {
34            let lng = num_arg(args, 0).unwrap_or(0.0);
35            let lat = num_arg(args, 1).unwrap_or(0.0);
36            Value::Geometry(nodedb_types::geometry::Geometry::point(lng, lat))
37        }
38        "geo_geohash" => {
39            let lng = num_arg(args, 0).unwrap_or(0.0);
40            let lat = num_arg(args, 1).unwrap_or(0.0);
41            let precision = num_arg(args, 2).unwrap_or(6.0) as u8;
42            Value::String(nodedb_spatial::geohash_encode(lng, lat, precision))
43        }
44        "geo_geohash_decode" => {
45            let hash = str_arg(args, 0).unwrap_or_default();
46            match nodedb_spatial::geohash_decode(&hash) {
47                Some(bb) => {
48                    let mut map = std::collections::HashMap::new();
49                    map.insert("min_lng".to_string(), Value::Float(bb.min_lng));
50                    map.insert("min_lat".to_string(), Value::Float(bb.min_lat));
51                    map.insert("max_lng".to_string(), Value::Float(bb.max_lng));
52                    map.insert("max_lat".to_string(), Value::Float(bb.max_lat));
53                    Value::Object(map)
54                }
55                None => Value::Null,
56            }
57        }
58        "geo_geohash_neighbors" => {
59            let hash = str_arg(args, 0).unwrap_or_default();
60            let neighbors = nodedb_spatial::geohash_neighbors(&hash);
61            let arr: Vec<Value> = neighbors
62                .into_iter()
63                .map(|(dir, h)| {
64                    let mut map = std::collections::HashMap::new();
65                    map.insert("direction".to_string(), Value::String(format!("{dir:?}")));
66                    map.insert("hash".to_string(), Value::String(h));
67                    Value::Object(map)
68                })
69                .collect();
70            Value::Array(arr)
71        }
72
73        // ── Spatial predicates (ST_*) ──
74        "st_contains" => geo_predicate_2(args, nodedb_spatial::st_contains),
75        "st_intersects" => geo_predicate_2(args, nodedb_spatial::st_intersects),
76        "st_within" => geo_predicate_2(args, nodedb_spatial::st_within),
77        "st_disjoint" => geo_predicate_2(args, nodedb_spatial::st_disjoint),
78        "st_dwithin" => {
79            let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else {
80                return Some(Value::Null);
81            };
82            let dist = num_arg(args, 2).unwrap_or(0.0);
83            Value::Bool(nodedb_spatial::st_dwithin(&a, &b, dist))
84        }
85        "st_distance" => {
86            let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else {
87                return Some(Value::Null);
88            };
89            to_value_number(nodedb_spatial::st_distance(&a, &b))
90        }
91        "st_buffer" => {
92            let Some(geom) = geom_arg(args, 0) else {
93                return Some(Value::Null);
94            };
95            let dist = num_arg(args, 1).unwrap_or(0.0);
96            let segs = num_arg(args, 2).unwrap_or(32.0) as usize;
97            Value::Geometry(nodedb_spatial::st_buffer(&geom, dist, segs))
98        }
99        "st_envelope" => {
100            let Some(geom) = geom_arg(args, 0) else {
101                return Some(Value::Null);
102            };
103            Value::Geometry(nodedb_spatial::st_envelope(&geom))
104        }
105        "st_union" => {
106            let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else {
107                return Some(Value::Null);
108            };
109            Value::Geometry(nodedb_spatial::st_union(&a, &b))
110        }
111
112        // ── Extended geo functions ──
113        "geo_length" => {
114            let Some(geom) = geom_arg(args, 0) else {
115                return Some(Value::Null);
116            };
117            to_value_number(geo_linestring_length(&geom))
118        }
119        "geo_perimeter" => {
120            let Some(geom) = geom_arg(args, 0) else {
121                return Some(Value::Null);
122            };
123            to_value_number(geo_polygon_perimeter(&geom))
124        }
125        "geo_line" => {
126            let coords: Vec<[f64; 2]> = args
127                .iter()
128                .filter_map(|v| {
129                    let geom = v.as_geometry()?;
130                    if let nodedb_types::geometry::Geometry::Point { coordinates } = geom {
131                        Some(*coordinates)
132                    } else {
133                        None
134                    }
135                })
136                .collect();
137            if coords.len() < 2 {
138                Value::Null
139            } else {
140                Value::Geometry(nodedb_types::geometry::Geometry::line_string(coords))
141            }
142        }
143        "geo_polygon" => {
144            // Args are arrays of coordinate pairs. Convert from Value::Array.
145            let rings: Vec<Vec<[f64; 2]>> = args
146                .iter()
147                .filter_map(|v| {
148                    let arr = v.as_array()?;
149                    let coords: Vec<[f64; 2]> = arr
150                        .iter()
151                        .filter_map(|pt| {
152                            let inner = pt.as_array()?;
153                            if inner.len() >= 2 {
154                                Some([inner[0].as_f64()?, inner[1].as_f64()?])
155                            } else {
156                                None
157                            }
158                        })
159                        .collect();
160                    if coords.is_empty() {
161                        None
162                    } else {
163                        Some(coords)
164                    }
165                })
166                .collect();
167            if rings.is_empty() {
168                Value::Null
169            } else {
170                Value::Geometry(nodedb_types::geometry::Geometry::polygon(rings))
171            }
172        }
173        "geo_circle" => {
174            let lng = num_arg(args, 0).unwrap_or(0.0);
175            let lat = num_arg(args, 1).unwrap_or(0.0);
176            let radius = num_arg(args, 2).unwrap_or(0.0);
177            let segs = num_arg(args, 3).unwrap_or(32.0) as usize;
178            let circle = nodedb_spatial::st_buffer(
179                &nodedb_types::geometry::Geometry::point(lng, lat),
180                radius,
181                segs,
182            );
183            Value::Geometry(circle)
184        }
185        "geo_bbox" => {
186            let min_lng = num_arg(args, 0).unwrap_or(0.0);
187            let min_lat = num_arg(args, 1).unwrap_or(0.0);
188            let max_lng = num_arg(args, 2).unwrap_or(0.0);
189            let max_lat = num_arg(args, 3).unwrap_or(0.0);
190            Value::Geometry(nodedb_types::geometry::Geometry::polygon(vec![vec![
191                [min_lng, min_lat],
192                [max_lng, min_lat],
193                [max_lng, max_lat],
194                [min_lng, max_lat],
195                [min_lng, min_lat],
196            ]]))
197        }
198        "geo_as_geojson" => {
199            let Some(geom) = geom_arg(args, 0) else {
200                return Some(Value::Null);
201            };
202            match sonic_rs::to_string(&geom) {
203                Ok(s) => Value::String(s),
204                Err(_) => Value::Null,
205            }
206        }
207        "geo_from_geojson" => {
208            let s = str_arg(args, 0).unwrap_or_default();
209            match sonic_rs::from_str::<nodedb_types::geometry::Geometry>(&s) {
210                Ok(g) => Value::Geometry(g),
211                Err(_) => Value::Null,
212            }
213        }
214        "geo_as_wkt" => {
215            let Some(geom) = geom_arg(args, 0) else {
216                return Some(Value::Null);
217            };
218            Value::String(nodedb_spatial::geometry_to_wkt(&geom))
219        }
220        "geo_from_wkt" => {
221            let s = str_arg(args, 0).unwrap_or_default();
222            match nodedb_spatial::geometry_from_wkt(&s) {
223                Some(g) => Value::Geometry(g),
224                None => Value::Null,
225            }
226        }
227        "geo_x" => {
228            let Some(geom) = geom_arg(args, 0) else {
229                return Some(Value::Null);
230            };
231            if let nodedb_types::geometry::Geometry::Point { coordinates } = geom {
232                to_value_number(coordinates[0])
233            } else {
234                Value::Null
235            }
236        }
237        "geo_y" => {
238            let Some(geom) = geom_arg(args, 0) else {
239                return Some(Value::Null);
240            };
241            if let nodedb_types::geometry::Geometry::Point { coordinates } = geom {
242                to_value_number(coordinates[1])
243            } else {
244                Value::Null
245            }
246        }
247        "geo_num_points" => {
248            let Some(geom) = geom_arg(args, 0) else {
249                return Some(Value::Null);
250            };
251            Value::Integer(count_points(&geom) as i64)
252        }
253        "geo_type" => {
254            let Some(geom) = geom_arg(args, 0) else {
255                return Some(Value::Null);
256            };
257            Value::String(geom.geometry_type().to_string())
258        }
259        "geo_is_valid" => {
260            let Some(geom) = geom_arg(args, 0) else {
261                return Some(Value::Null);
262            };
263            Value::Bool(nodedb_spatial::is_valid(&geom))
264        }
265
266        // ── H3 hexagonal index ──
267        "geo_h3" => {
268            let lng = num_arg(args, 0).unwrap_or(0.0);
269            let lat = num_arg(args, 1).unwrap_or(0.0);
270            let resolution = num_arg(args, 2).unwrap_or(7.0) as u8;
271            match nodedb_spatial::h3::h3_encode_string(lng, lat, resolution) {
272                Some(hex) => Value::String(hex),
273                None => Value::Null,
274            }
275        }
276        "geo_h3_to_boundary" => {
277            let h3_str = str_arg(args, 0).unwrap_or_default();
278            let h3_idx = u64::from_str_radix(&h3_str, 16).unwrap_or(0);
279            if !nodedb_spatial::h3::h3_is_valid(h3_idx) {
280                return Some(Value::Null);
281            }
282            match nodedb_spatial::h3::h3_to_boundary(h3_idx) {
283                Some(geom) => Value::Geometry(geom),
284                None => Value::Null,
285            }
286        }
287        "geo_h3_resolution" => {
288            let h3_str = str_arg(args, 0).unwrap_or_default();
289            let h3_idx = u64::from_str_radix(&h3_str, 16).unwrap_or(0);
290            if !nodedb_spatial::h3::h3_is_valid(h3_idx) {
291                return Some(Value::Null);
292            }
293            match nodedb_spatial::h3::h3_resolution(h3_idx) {
294                Some(r) => Value::Integer(r as i64),
295                None => Value::Null,
296            }
297        }
298        // ── SQL-aliased geohash/H3 encode/decode ──
299        "st_geohash" => {
300            let lng = num_arg(args, 0).unwrap_or(0.0);
301            let lat = num_arg(args, 1).unwrap_or(0.0);
302            let precision = num_arg(args, 2).unwrap_or(6.0) as u8;
303            Value::String(nodedb_spatial::geohash_encode(lng, lat, precision))
304        }
305        "st_geohashdecode" => {
306            let hash = str_arg(args, 0).unwrap_or_default();
307            match nodedb_spatial::geohash_decode(&hash) {
308                Some(bb) => {
309                    let mut map = std::collections::HashMap::new();
310                    map.insert("min_lng".to_string(), Value::Float(bb.min_lng));
311                    map.insert("min_lat".to_string(), Value::Float(bb.min_lat));
312                    map.insert("max_lng".to_string(), Value::Float(bb.max_lng));
313                    map.insert("max_lat".to_string(), Value::Float(bb.max_lat));
314                    Value::Object(map)
315                }
316                None => Value::Null,
317            }
318        }
319        "h3_latlngtocell" => {
320            let lat = num_arg(args, 0).unwrap_or(0.0);
321            let lng = num_arg(args, 1).unwrap_or(0.0);
322            let resolution = num_arg(args, 2).unwrap_or(7.0) as u8;
323            match nodedb_spatial::h3::h3_encode_string(lng, lat, resolution) {
324                Some(hex) => Value::String(hex),
325                None => Value::Null,
326            }
327        }
328        "h3_celltolatlng" => {
329            let h3_str = str_arg(args, 0).unwrap_or_default();
330            let h3_idx = u64::from_str_radix(&h3_str, 16).unwrap_or(0);
331            if !nodedb_spatial::h3::h3_is_valid(h3_idx) {
332                return Some(Value::Null);
333            }
334            match nodedb_spatial::h3::h3_to_center(h3_idx) {
335                Some((lng, lat)) => {
336                    let mut map = std::collections::HashMap::new();
337                    map.insert("lat".to_string(), Value::Float(lat));
338                    map.insert("lng".to_string(), Value::Float(lng));
339                    Value::Object(map)
340                }
341                None => Value::Null,
342            }
343        }
344        "st_intersection" => {
345            let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else {
346                return Some(Value::Null);
347            };
348            Value::Geometry(nodedb_spatial::st_intersection(&a, &b))
349        }
350        _ => return None,
351    };
352    Some(result)
353}
354
355// ── Helpers ──
356
357fn str_arg(args: &[Value], idx: usize) -> Option<String> {
358    args.get(idx)?.as_str().map(|s| s.to_string())
359}
360
361fn num_arg(args: &[Value], idx: usize) -> Option<f64> {
362    args.get(idx).and_then(|v| value_to_f64(v, true))
363}
364
365/// Extract a geometry argument. Supports `Value::Geometry` directly,
366/// or falls back to JSON deserialization for `Value::Object` (GeoJSON).
367fn geom_arg(args: &[Value], idx: usize) -> Option<nodedb_types::geometry::Geometry> {
368    match args.get(idx)? {
369        Value::Geometry(g) => Some(g.clone()),
370        // GeoJSON string: produced when a Geometry constant was round-tripped
371        // through the const-fold SqlValue path (e.g. ST_Distance(ST_Point(...), ST_Point(...))).
372        Value::String(s) => sonic_rs::from_str::<nodedb_types::geometry::Geometry>(s).ok(),
373        other => {
374            // Fallback: convert Value → JSON → Geometry for GeoJSON objects.
375            let json = serde_json::Value::from(other.clone());
376            serde_json::from_value(json).ok()
377        }
378    }
379}
380
381fn geo_predicate_2(
382    args: &[Value],
383    f: fn(&nodedb_types::geometry::Geometry, &nodedb_types::geometry::Geometry) -> bool,
384) -> Value {
385    let (Some(a), Some(b)) = (geom_arg(args, 0), geom_arg(args, 1)) else {
386        return Value::Null;
387    };
388    Value::Bool(f(&a, &b))
389}
390
391fn geo_linestring_length(geom: &nodedb_types::geometry::Geometry) -> f64 {
392    let coords = match geom {
393        nodedb_types::geometry::Geometry::LineString { coordinates } => coordinates,
394        _ => return 0.0,
395    };
396    let mut total = 0.0;
397    for i in 0..coords.len().saturating_sub(1) {
398        total += nodedb_types::geometry::haversine_distance(
399            coords[i][0],
400            coords[i][1],
401            coords[i + 1][0],
402            coords[i + 1][1],
403        );
404    }
405    total
406}
407
408fn geo_polygon_perimeter(geom: &nodedb_types::geometry::Geometry) -> f64 {
409    let rings = match geom {
410        nodedb_types::geometry::Geometry::Polygon { coordinates } => coordinates,
411        _ => return 0.0,
412    };
413    let Some(exterior) = rings.first() else {
414        return 0.0;
415    };
416    let mut total = 0.0;
417    for i in 0..exterior.len().saturating_sub(1) {
418        total += nodedb_types::geometry::haversine_distance(
419            exterior[i][0],
420            exterior[i][1],
421            exterior[i + 1][0],
422            exterior[i + 1][1],
423        );
424    }
425    total
426}
427
428fn count_points(geom: &nodedb_types::geometry::Geometry) -> usize {
429    use nodedb_types::geometry::Geometry;
430    match geom {
431        Geometry::Point { .. } => 1,
432        Geometry::LineString { coordinates } => coordinates.len(),
433        Geometry::Polygon { coordinates } => coordinates.iter().map(|r| r.len()).sum(),
434        Geometry::MultiPoint { coordinates } => coordinates.len(),
435        Geometry::MultiLineString { coordinates } => coordinates.iter().map(|ls| ls.len()).sum(),
436        Geometry::MultiPolygon { coordinates } => coordinates
437            .iter()
438            .flat_map(|poly| poly.iter())
439            .map(|ring| ring.len())
440            .sum(),
441        Geometry::GeometryCollection { geometries } => geometries.iter().map(count_points).sum(),
442        _ => 0,
443    }
444}