jpx_core/extensions/
geo.rs

1//! Geographic/geospatial functions.
2
3use std::collections::HashSet;
4
5use geoutils::Location;
6use serde_json::Value;
7
8use crate::functions::{Function, number_value};
9use crate::interpreter::SearchResult;
10use crate::registry::register_if_enabled;
11use crate::{Context, Runtime, arg, defn};
12
13// =============================================================================
14// geo_distance(lat1, lon1, lat2, lon2) -> number (meters)
15// =============================================================================
16
17defn!(
18    GeoDistanceFn,
19    vec![arg!(number), arg!(number), arg!(number), arg!(number)],
20    None
21);
22
23impl Function for GeoDistanceFn {
24    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
25        self.signature.validate(args, ctx)?;
26        let lat1 = args[0].as_f64().unwrap();
27        let lon1 = args[1].as_f64().unwrap();
28        let lat2 = args[2].as_f64().unwrap();
29        let lon2 = args[3].as_f64().unwrap();
30
31        let loc1 = Location::new(lat1, lon1);
32        let loc2 = Location::new(lat2, lon2);
33
34        let distance = loc1.haversine_distance_to(&loc2);
35        Ok(number_value(distance.meters()))
36    }
37}
38
39// =============================================================================
40// geo_distance_km(lat1, lon1, lat2, lon2) -> number (kilometers)
41// =============================================================================
42
43defn!(
44    GeoDistanceKmFn,
45    vec![arg!(number), arg!(number), arg!(number), arg!(number)],
46    None
47);
48
49impl Function for GeoDistanceKmFn {
50    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
51        self.signature.validate(args, ctx)?;
52        let lat1 = args[0].as_f64().unwrap();
53        let lon1 = args[1].as_f64().unwrap();
54        let lat2 = args[2].as_f64().unwrap();
55        let lon2 = args[3].as_f64().unwrap();
56
57        let loc1 = Location::new(lat1, lon1);
58        let loc2 = Location::new(lat2, lon2);
59
60        let distance = loc1.haversine_distance_to(&loc2);
61        Ok(number_value(distance.meters() / 1000.0))
62    }
63}
64
65// =============================================================================
66// geo_distance_miles(lat1, lon1, lat2, lon2) -> number (miles)
67// =============================================================================
68
69defn!(
70    GeoDistanceMilesFn,
71    vec![arg!(number), arg!(number), arg!(number), arg!(number)],
72    None
73);
74
75impl Function for GeoDistanceMilesFn {
76    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
77        self.signature.validate(args, ctx)?;
78        let lat1 = args[0].as_f64().unwrap();
79        let lon1 = args[1].as_f64().unwrap();
80        let lat2 = args[2].as_f64().unwrap();
81        let lon2 = args[3].as_f64().unwrap();
82
83        let loc1 = Location::new(lat1, lon1);
84        let loc2 = Location::new(lat2, lon2);
85
86        // 1 meter = 0.000621371 miles
87        const METERS_TO_MILES: f64 = 0.000621371;
88
89        let distance = loc1.haversine_distance_to(&loc2);
90        Ok(number_value(distance.meters() * METERS_TO_MILES))
91    }
92}
93
94// =============================================================================
95// geo_bearing(lat1, lon1, lat2, lon2) -> number (degrees 0-360)
96// =============================================================================
97
98defn!(
99    GeoBearingFn,
100    vec![arg!(number), arg!(number), arg!(number), arg!(number)],
101    None
102);
103
104impl Function for GeoBearingFn {
105    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
106        self.signature.validate(args, ctx)?;
107        let lat1 = args[0].as_f64().unwrap();
108        let lon1 = args[1].as_f64().unwrap();
109        let lat2 = args[2].as_f64().unwrap();
110        let lon2 = args[3].as_f64().unwrap();
111
112        // Calculate initial bearing using the forward azimuth formula
113        let lat1_rad = lat1.to_radians();
114        let lat2_rad = lat2.to_radians();
115        let delta_lon = (lon2 - lon1).to_radians();
116
117        let x = delta_lon.sin() * lat2_rad.cos();
118        let y = lat1_rad.cos() * lat2_rad.sin() - lat1_rad.sin() * lat2_rad.cos() * delta_lon.cos();
119
120        let bearing_rad = x.atan2(y);
121        let mut bearing = bearing_rad.to_degrees();
122
123        // Normalize to 0-360
124        if bearing < 0.0 {
125            bearing += 360.0;
126        }
127
128        Ok(number_value(bearing))
129    }
130}
131
132/// Register geo functions filtered by the enabled set.
133pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
134    register_if_enabled(
135        runtime,
136        "geo_distance",
137        enabled,
138        Box::new(GeoDistanceFn::new()),
139    );
140    register_if_enabled(
141        runtime,
142        "geo_distance_km",
143        enabled,
144        Box::new(GeoDistanceKmFn::new()),
145    );
146    register_if_enabled(
147        runtime,
148        "geo_distance_miles",
149        enabled,
150        Box::new(GeoDistanceMilesFn::new()),
151    );
152    register_if_enabled(
153        runtime,
154        "geo_bearing",
155        enabled,
156        Box::new(GeoBearingFn::new()),
157    );
158}
159
160#[cfg(test)]
161mod tests {
162    use crate::Runtime;
163    use serde_json::json;
164
165    fn setup_runtime() -> Runtime {
166        Runtime::builder()
167            .with_standard()
168            .with_all_extensions()
169            .build()
170    }
171
172    #[test]
173    fn test_geo_distance() {
174        let runtime = setup_runtime();
175        // NYC to LA: approximately 3940 km
176        let data = json!({"nyc": [40.7128, -74.0060], "la": [34.0522, -118.2437]});
177        let expr = runtime
178            .compile("geo_distance(nyc[0], nyc[1], la[0], la[1])")
179            .unwrap();
180        let result = expr.search(&data).unwrap();
181        let meters = result.as_f64().unwrap();
182        // Should be approximately 3940000 meters
183        assert!(meters > 3900000.0 && meters < 4000000.0);
184    }
185
186    #[test]
187    fn test_geo_distance_km() {
188        let runtime = setup_runtime();
189        let data = json!({"nyc": [40.7128, -74.0060], "la": [34.0522, -118.2437]});
190        let expr = runtime
191            .compile("geo_distance_km(nyc[0], nyc[1], la[0], la[1])")
192            .unwrap();
193        let result = expr.search(&data).unwrap();
194        let km = result.as_f64().unwrap();
195        // Should be approximately 3940 km
196        assert!(km > 3900.0 && km < 4000.0);
197    }
198
199    #[test]
200    fn test_geo_distance_miles() {
201        let runtime = setup_runtime();
202        let data = json!({"nyc": [40.7128, -74.0060], "la": [34.0522, -118.2437]});
203        let expr = runtime
204            .compile("geo_distance_miles(nyc[0], nyc[1], la[0], la[1])")
205            .unwrap();
206        let result = expr.search(&data).unwrap();
207        let miles = result.as_f64().unwrap();
208        // Should be approximately 2450 miles
209        assert!(miles > 2400.0 && miles < 2500.0);
210    }
211
212    #[test]
213    fn test_geo_bearing() {
214        let runtime = setup_runtime();
215        // NYC to LA should be roughly west (270 degrees)
216        let data = json!({"nyc": [40.7128, -74.0060], "la": [34.0522, -118.2437]});
217        let expr = runtime
218            .compile("geo_bearing(nyc[0], nyc[1], la[0], la[1])")
219            .unwrap();
220        let result = expr.search(&data).unwrap();
221        let bearing = result.as_f64().unwrap();
222        // Should be roughly 273 degrees (west-southwest)
223        assert!(bearing > 260.0 && bearing < 290.0);
224    }
225
226    #[test]
227    fn test_geo_distance_same_point() {
228        let runtime = setup_runtime();
229        let data = json!([40.7128, -74.0060]);
230        let expr = runtime
231            .compile("geo_distance(@[0], @[1], @[0], @[1])")
232            .unwrap();
233        let result = expr.search(&data).unwrap();
234        let meters = result.as_f64().unwrap();
235        assert!(meters < 1.0); // Should be essentially 0
236    }
237}
jpx_core/extensions/geo.rs

jpx_core/extensions/
geo.rs
