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 datafusion::arrow::datatypes::DataType;
19use datafusion_expr::{Signature, Volatility};
20use datatypes::scalars::ScalarVectorBuilder;
21use datatypes::vectors::{Float64VectorBuilder, MutableVector, VectorRef};
22use derive_more::Display;
23use geo::algorithm::line_measures::metric_spaces::Euclidean;
24use geo::{Area, Distance, Haversine};
25use geo_types::Geometry;
26use snafu::ResultExt;
27
28use crate::function::{Function, FunctionContext};
29use crate::scalars::geo::helpers::{ensure_columns_len, ensure_columns_n};
30use crate::scalars::geo::wkt::parse_wkt;
31
32/// Return WGS84(SRID: 4326) euclidean distance between two geometry object, in degree
33#[derive(Clone, Debug, Default, Display)]
34#[display("{}", self.name())]
35pub struct STDistance;
36
37impl Function for STDistance {
38    fn name(&self) -> &str {
39        "st_distance"
40    }
41
42    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
43        Ok(DataType::Float64)
44    }
45
46    fn signature(&self) -> Signature {
47        Signature::string(2, Volatility::Stable)
48    }
49
50    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
51        ensure_columns_n!(columns, 2);
52
53        let wkt_this_vec = &columns[0];
54        let wkt_that_vec = &columns[1];
55
56        let size = wkt_this_vec.len();
57        let mut results = Float64VectorBuilder::with_capacity(size);
58
59        for i in 0..size {
60            let wkt_this = wkt_this_vec.get(i).as_string();
61            let wkt_that = wkt_that_vec.get(i).as_string();
62
63            let result = match (wkt_this, wkt_that) {
64                (Some(wkt_this), Some(wkt_that)) => {
65                    let geom_this = parse_wkt(&wkt_this)?;
66                    let geom_that = parse_wkt(&wkt_that)?;
67
68                    Some(Euclidean::distance(&geom_this, &geom_that))
69                }
70                _ => None,
71            };
72
73            results.push(result);
74        }
75
76        Ok(results.to_vector())
77    }
78}
79
80/// Return great circle distance between two geometry object, in meters
81#[derive(Clone, Debug, Default, Display)]
82#[display("{}", self.name())]
83pub struct STDistanceSphere;
84
85impl Function for STDistanceSphere {
86    fn name(&self) -> &str {
87        "st_distance_sphere_m"
88    }
89
90    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
91        Ok(DataType::Float64)
92    }
93
94    fn signature(&self) -> Signature {
95        Signature::string(2, Volatility::Stable)
96    }
97
98    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
99        ensure_columns_n!(columns, 2);
100
101        let wkt_this_vec = &columns[0];
102        let wkt_that_vec = &columns[1];
103
104        let size = wkt_this_vec.len();
105        let mut results = Float64VectorBuilder::with_capacity(size);
106
107        for i in 0..size {
108            let wkt_this = wkt_this_vec.get(i).as_string();
109            let wkt_that = wkt_that_vec.get(i).as_string();
110
111            let result = match (wkt_this, wkt_that) {
112                (Some(wkt_this), Some(wkt_that)) => {
113                    let geom_this = parse_wkt(&wkt_this)?;
114                    let geom_that = parse_wkt(&wkt_that)?;
115
116                    match (geom_this, geom_that) {
117                        (Geometry::Point(this), Geometry::Point(that)) => {
118                            Some(Haversine::distance(this, that))
119                        }
120                        _ => {
121                            Err(BoxedError::new(PlainError::new(
122                                "Great circle distance between non-point objects are not supported for now.".to_string(),
123                                StatusCode::Unsupported,
124                            ))).context(error::ExecuteSnafu)?
125                        }
126                    }
127                }
128                _ => None,
129            };
130
131            results.push(result);
132        }
133
134        Ok(results.to_vector())
135    }
136}
137
138/// Return area of given geometry object
139#[derive(Clone, Debug, Default, Display)]
140#[display("{}", self.name())]
141pub struct STArea;
142
143impl Function for STArea {
144    fn name(&self) -> &str {
145        "st_area"
146    }
147
148    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
149        Ok(DataType::Float64)
150    }
151
152    fn signature(&self) -> Signature {
153        Signature::string(1, Volatility::Stable)
154    }
155
156    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
157        ensure_columns_n!(columns, 1);
158
159        let wkt_vec = &columns[0];
160
161        let size = wkt_vec.len();
162        let mut results = Float64VectorBuilder::with_capacity(size);
163
164        for i in 0..size {
165            let wkt = wkt_vec.get(i).as_string();
166
167            let result = if let Some(wkt) = wkt {
168                let geom = parse_wkt(&wkt)?;
169                Some(geom.unsigned_area())
170            } else {
171                None
172            };
173
174            results.push(result);
175        }
176
177        Ok(results.to_vector())
178    }
179}
common_function/scalars/geo/measure.rs

common_function/scalars/geo/
measure.rs
