vrl/stdlib/
haversine.rs

1use std::collections::BTreeMap;
2
3use crate::compiler::prelude::*;
4use crate::value;
5
6use super::util::round_to_precision;
7
8const EARTH_RADIUS_IN_METERS: f64 = 6_371_008.8;
9const EARTH_RADIUS_IN_KILOMETERS: f64 = EARTH_RADIUS_IN_METERS / 1000.0;
10const EARTH_RADIUS_IN_MILES: f64 = EARTH_RADIUS_IN_KILOMETERS * 0.621_371_2;
11
12fn haversine_distance(
13    latitude1: Value,
14    longitude1: Value,
15    latitude2: Value,
16    longitude2: Value,
17    measurement_unit: &MeasurementUnit,
18) -> Resolved {
19    let latitude1 = latitude1.try_float()?.to_radians();
20    let longitude1 = longitude1.try_float()?.to_radians();
21    let latitude2 = latitude2.try_float()?.to_radians();
22    let longitude2 = longitude2.try_float()?.to_radians();
23
24    let mut result = ObjectMap::new();
25
26    // Distance calculation
27    let dlon = longitude2 - longitude1;
28    let dlat = latitude2 - latitude1;
29    let a =
30        (dlat / 2.0).sin().powi(2) + latitude1.cos() * latitude2.cos() * (dlon / 2.0).sin().powi(2);
31    let distance = 2.0 * a.sqrt().asin();
32
33    result.insert(
34        "distance".into(),
35        match measurement_unit {
36            MeasurementUnit::Kilometers => Value::from_f64_or_zero(round_to_precision(
37                distance * EARTH_RADIUS_IN_KILOMETERS,
38                7,
39                f64::round,
40            )),
41            MeasurementUnit::Miles => Value::from_f64_or_zero(round_to_precision(
42                distance * EARTH_RADIUS_IN_MILES,
43                7,
44                f64::round,
45            )),
46        },
47    );
48
49    // Bearing calculation
50    let y = dlon.sin() * latitude2.cos();
51    let x = latitude1.cos() * latitude2.sin() - latitude1.sin() * latitude2.cos() * dlon.cos();
52    let bearing = (y.atan2(x).to_degrees() + 360.0) % 360.0;
53
54    result.insert(
55        "bearing".into(),
56        Value::from_f64_or_zero(round_to_precision(bearing, 3, f64::round)),
57    );
58
59    Ok(result.into())
60}
61
62fn measurement_systems() -> Vec<Value> {
63    vec![value!("kilometers"), value!("miles")]
64}
65
66#[derive(Clone, Debug)]
67enum MeasurementUnit {
68    Kilometers,
69    Miles,
70}
71
72#[derive(Clone, Copy, Debug)]
73pub struct Haversine;
74
75impl Function for Haversine {
76    fn identifier(&self) -> &'static str {
77        "haversine"
78    }
79
80    fn usage(&self) -> &'static str {
81        "Calculates [haversine](https://en.wikipedia.org/wiki/Haversine_formula) distance and bearing between two points."
82    }
83
84    fn parameters(&self) -> &'static [Parameter] {
85        &[
86            Parameter {
87                keyword: "latitude1",
88                kind: kind::FLOAT,
89                required: true,
90            },
91            Parameter {
92                keyword: "longitude1",
93                kind: kind::FLOAT,
94                required: true,
95            },
96            Parameter {
97                keyword: "latitude2",
98                kind: kind::FLOAT,
99                required: true,
100            },
101            Parameter {
102                keyword: "longitude2",
103                kind: kind::FLOAT,
104                required: true,
105            },
106            Parameter {
107                keyword: "measurement_unit",
108                kind: kind::BYTES,
109                required: false,
110            },
111        ]
112    }
113
114    fn compile(
115        &self,
116        state: &state::TypeState,
117        _ctx: &mut FunctionCompileContext,
118        arguments: ArgumentList,
119    ) -> Compiled {
120        let latitude1 = arguments.required("latitude1");
121        let longitude1 = arguments.required("longitude1");
122        let latitude2 = arguments.required("latitude2");
123        let longitude2 = arguments.required("longitude2");
124
125        let measurement_unit = match arguments
126            .optional_enum("measurement_unit", &measurement_systems(), state)?
127            .unwrap_or_else(|| value!("kilometers"))
128            .try_bytes()
129            .ok()
130            .as_deref()
131        {
132            Some(b"kilometers") => MeasurementUnit::Kilometers,
133            Some(b"miles") => MeasurementUnit::Miles,
134            _ => return Err(Box::new(ExpressionError::from("invalid measurement unit"))),
135        };
136
137        Ok(HaversineFn {
138            latitude1,
139            longitude1,
140            latitude2,
141            longitude2,
142            measurement_unit,
143        }
144        .as_expr())
145    }
146
147    fn examples(&self) -> &'static [Example] {
148        &[
149            example! {
150                title: "Haversine in kilometers",
151                source: "haversine(0.0, 0.0, 10.0, 10.0)",
152                result: Ok(indoc!(
153                    r#"{
154                        "distance": 1568.5227233,
155                        "bearing": 44.561
156                    }"#
157                )),
158            },
159            example! {
160                title: "Haversine in miles",
161                source: r#"haversine(0.0, 0.0, 10.0, 10.0, measurement_unit: "miles")"#,
162                result: Ok(indoc!(
163                    r#"{
164                        "distance": 974.6348468,
165                        "bearing": 44.561
166                    }"#
167                )),
168            },
169        ]
170    }
171}
172
173#[derive(Clone, Debug)]
174struct HaversineFn {
175    latitude1: Box<dyn Expression>,
176    longitude1: Box<dyn Expression>,
177    latitude2: Box<dyn Expression>,
178    longitude2: Box<dyn Expression>,
179    measurement_unit: MeasurementUnit,
180}
181
182impl FunctionExpression for HaversineFn {
183    fn resolve(&self, ctx: &mut Context) -> Resolved {
184        let latitude1 = self.latitude1.resolve(ctx)?;
185        let longitude1 = self.longitude1.resolve(ctx)?;
186        let latitude2 = self.latitude2.resolve(ctx)?;
187        let longitude2 = self.longitude2.resolve(ctx)?;
188
189        haversine_distance(
190            latitude1,
191            longitude1,
192            latitude2,
193            longitude2,
194            &self.measurement_unit,
195        )
196    }
197
198    fn type_def(&self, _state: &state::TypeState) -> TypeDef {
199        TypeDef::object(inner_kind()).infallible()
200    }
201}
202
203fn inner_kind() -> BTreeMap<Field, Kind> {
204    BTreeMap::from([
205        (Field::from("distance"), Kind::float()),
206        (Field::from("bearing"), Kind::float()),
207    ])
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213    use crate::value;
214
215    test_function![
216        haversine => Haversine;
217
218        basic_kilometers {
219            args: func_args![latitude1: value!(0.0), longitude1: value!(0.0), latitude2: value!(10.0), longitude2: value!(10.0)],
220            want: Ok(value!({ "distance": 1_568.522_723_3, "bearing": 44.561 })),
221            tdef: TypeDef::object(inner_kind()).infallible(),
222        }
223
224        basic_miles {
225            args: func_args![latitude1: value!(0.0), longitude1: value!(0.0), latitude2: value!(10.0), longitude2: value!(10.0), measurement_unit: value!("miles")],
226            want: Ok(value!({ "distance": 974.634_846_8, "bearing": 44.561 })),
227            tdef: TypeDef::object(inner_kind()).infallible(),
228        }
229    ];
230}