common_function/scalars/geo/
measure.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use common_error::ext::{BoxedError, PlainError};
16use common_error::status_code::StatusCode;
17use common_query::error::{self, Result};
18use common_query::prelude::{Signature, TypeSignature};
19use datafusion::logical_expr::Volatility;
20use datatypes::prelude::ConcreteDataType;
21use datatypes::scalars::ScalarVectorBuilder;
22use datatypes::vectors::{Float64VectorBuilder, MutableVector, VectorRef};
23use derive_more::Display;
24use geo::algorithm::line_measures::metric_spaces::Euclidean;
25use geo::{Area, Distance, Haversine};
26use geo_types::Geometry;
27use snafu::ResultExt;
28
29use crate::function::{Function, FunctionContext};
30use crate::scalars::geo::helpers::{ensure_columns_len, ensure_columns_n};
31use crate::scalars::geo::wkt::parse_wkt;
32
33/// Return WGS84(SRID: 4326) euclidean distance between two geometry object, in degree
34#[derive(Clone, Debug, Default, Display)]
35#[display("{}", self.name())]
36pub struct STDistance;
37
38impl Function for STDistance {
39    fn name(&self) -> &str {
40        "st_distance"
41    }
42
43    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
44        Ok(ConcreteDataType::float64_datatype())
45    }
46
47    fn signature(&self) -> Signature {
48        Signature::new(
49            TypeSignature::Exact(vec![
50                ConcreteDataType::string_datatype(),
51                ConcreteDataType::string_datatype(),
52            ]),
53            Volatility::Stable,
54        )
55    }
56
57    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
58        ensure_columns_n!(columns, 2);
59
60        let wkt_this_vec = &columns[0];
61        let wkt_that_vec = &columns[1];
62
63        let size = wkt_this_vec.len();
64        let mut results = Float64VectorBuilder::with_capacity(size);
65
66        for i in 0..size {
67            let wkt_this = wkt_this_vec.get(i).as_string();
68            let wkt_that = wkt_that_vec.get(i).as_string();
69
70            let result = match (wkt_this, wkt_that) {
71                (Some(wkt_this), Some(wkt_that)) => {
72                    let geom_this = parse_wkt(&wkt_this)?;
73                    let geom_that = parse_wkt(&wkt_that)?;
74
75                    Some(Euclidean::distance(&geom_this, &geom_that))
76                }
77                _ => None,
78            };
79
80            results.push(result);
81        }
82
83        Ok(results.to_vector())
84    }
85}
86
87/// Return great circle distance between two geometry object, in meters
88#[derive(Clone, Debug, Default, Display)]
89#[display("{}", self.name())]
90pub struct STDistanceSphere;
91
92impl Function for STDistanceSphere {
93    fn name(&self) -> &str {
94        "st_distance_sphere_m"
95    }
96
97    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
98        Ok(ConcreteDataType::float64_datatype())
99    }
100
101    fn signature(&self) -> Signature {
102        Signature::new(
103            TypeSignature::Exact(vec![
104                ConcreteDataType::string_datatype(),
105                ConcreteDataType::string_datatype(),
106            ]),
107            Volatility::Stable,
108        )
109    }
110
111    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
112        ensure_columns_n!(columns, 2);
113
114        let wkt_this_vec = &columns[0];
115        let wkt_that_vec = &columns[1];
116
117        let size = wkt_this_vec.len();
118        let mut results = Float64VectorBuilder::with_capacity(size);
119
120        for i in 0..size {
121            let wkt_this = wkt_this_vec.get(i).as_string();
122            let wkt_that = wkt_that_vec.get(i).as_string();
123
124            let result = match (wkt_this, wkt_that) {
125                (Some(wkt_this), Some(wkt_that)) => {
126                    let geom_this = parse_wkt(&wkt_this)?;
127                    let geom_that = parse_wkt(&wkt_that)?;
128
129                    match (geom_this, geom_that) {
130                        (Geometry::Point(this), Geometry::Point(that)) => {
131                            Some(Haversine::distance(this, that))
132                        }
133                        _ => {
134                            Err(BoxedError::new(PlainError::new(
135                                "Great circle distance between non-point objects are not supported for now.".to_string(),
136                                StatusCode::Unsupported,
137                            ))).context(error::ExecuteSnafu)?
138                        }
139                    }
140                }
141                _ => None,
142            };
143
144            results.push(result);
145        }
146
147        Ok(results.to_vector())
148    }
149}
150
151/// Return area of given geometry object
152#[derive(Clone, Debug, Default, Display)]
153#[display("{}", self.name())]
154pub struct STArea;
155
156impl Function for STArea {
157    fn name(&self) -> &str {
158        "st_area"
159    }
160
161    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
162        Ok(ConcreteDataType::float64_datatype())
163    }
164
165    fn signature(&self) -> Signature {
166        Signature::new(
167            TypeSignature::Exact(vec![ConcreteDataType::string_datatype()]),
168            Volatility::Stable,
169        )
170    }
171
172    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
173        ensure_columns_n!(columns, 1);
174
175        let wkt_vec = &columns[0];
176
177        let size = wkt_vec.len();
178        let mut results = Float64VectorBuilder::with_capacity(size);
179
180        for i in 0..size {
181            let wkt = wkt_vec.get(i).as_string();
182
183            let result = if let Some(wkt) = wkt {
184                let geom = parse_wkt(&wkt)?;
185                Some(geom.unsigned_area())
186            } else {
187                None
188            };
189
190            results.push(result);
191        }
192
193        Ok(results.to_vector())
194    }
195}
common_function/scalars/geo/measure.rs

common_function/scalars/geo/
measure.rs
