common_function/scalars/geo/
relation.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_query::error::Result;
16use datafusion::arrow::datatypes::DataType;
17use datafusion_expr::{Signature, Volatility};
18use datatypes::scalars::ScalarVectorBuilder;
19use datatypes::vectors::{BooleanVectorBuilder, MutableVector, VectorRef};
20use derive_more::Display;
21use geo::algorithm::contains::Contains;
22use geo::algorithm::intersects::Intersects;
23use geo::algorithm::within::Within;
24
25use crate::function::{Function, FunctionContext};
26use crate::scalars::geo::helpers::{ensure_columns_len, ensure_columns_n};
27use crate::scalars::geo::wkt::parse_wkt;
28
29/// Test if spatial relationship: contains
30#[derive(Clone, Debug, Default, Display)]
31#[display("{}", self.name())]
32pub struct STContains;
33
34impl Function for STContains {
35    fn name(&self) -> &str {
36        "st_contains"
37    }
38
39    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
40        Ok(DataType::Boolean)
41    }
42
43    fn signature(&self) -> Signature {
44        Signature::string(2, Volatility::Stable)
45    }
46
47    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
48        ensure_columns_n!(columns, 2);
49
50        let wkt_this_vec = &columns[0];
51        let wkt_that_vec = &columns[1];
52
53        let size = wkt_this_vec.len();
54        let mut results = BooleanVectorBuilder::with_capacity(size);
55
56        for i in 0..size {
57            let wkt_this = wkt_this_vec.get(i).as_string();
58            let wkt_that = wkt_that_vec.get(i).as_string();
59
60            let result = match (wkt_this, wkt_that) {
61                (Some(wkt_this), Some(wkt_that)) => {
62                    let geom_this = parse_wkt(&wkt_this)?;
63                    let geom_that = parse_wkt(&wkt_that)?;
64
65                    Some(geom_this.contains(&geom_that))
66                }
67                _ => None,
68            };
69
70            results.push(result);
71        }
72
73        Ok(results.to_vector())
74    }
75}
76
77/// Test if spatial relationship: within
78#[derive(Clone, Debug, Default, Display)]
79#[display("{}", self.name())]
80pub struct STWithin;
81
82impl Function for STWithin {
83    fn name(&self) -> &str {
84        "st_within"
85    }
86
87    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
88        Ok(DataType::Boolean)
89    }
90
91    fn signature(&self) -> Signature {
92        Signature::string(2, Volatility::Stable)
93    }
94
95    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
96        ensure_columns_n!(columns, 2);
97
98        let wkt_this_vec = &columns[0];
99        let wkt_that_vec = &columns[1];
100
101        let size = wkt_this_vec.len();
102        let mut results = BooleanVectorBuilder::with_capacity(size);
103
104        for i in 0..size {
105            let wkt_this = wkt_this_vec.get(i).as_string();
106            let wkt_that = wkt_that_vec.get(i).as_string();
107
108            let result = match (wkt_this, wkt_that) {
109                (Some(wkt_this), Some(wkt_that)) => {
110                    let geom_this = parse_wkt(&wkt_this)?;
111                    let geom_that = parse_wkt(&wkt_that)?;
112
113                    Some(geom_this.is_within(&geom_that))
114                }
115                _ => None,
116            };
117
118            results.push(result);
119        }
120
121        Ok(results.to_vector())
122    }
123}
124
125/// Test if spatial relationship: within
126#[derive(Clone, Debug, Default, Display)]
127#[display("{}", self.name())]
128pub struct STIntersects;
129
130impl Function for STIntersects {
131    fn name(&self) -> &str {
132        "st_intersects"
133    }
134
135    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
136        Ok(DataType::Boolean)
137    }
138
139    fn signature(&self) -> Signature {
140        Signature::string(2, Volatility::Stable)
141    }
142
143    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
144        ensure_columns_n!(columns, 2);
145
146        let wkt_this_vec = &columns[0];
147        let wkt_that_vec = &columns[1];
148
149        let size = wkt_this_vec.len();
150        let mut results = BooleanVectorBuilder::with_capacity(size);
151
152        for i in 0..size {
153            let wkt_this = wkt_this_vec.get(i).as_string();
154            let wkt_that = wkt_that_vec.get(i).as_string();
155
156            let result = match (wkt_this, wkt_that) {
157                (Some(wkt_this), Some(wkt_that)) => {
158                    let geom_this = parse_wkt(&wkt_this)?;
159                    let geom_that = parse_wkt(&wkt_that)?;
160
161                    Some(geom_this.intersects(&geom_that))
162                }
163                _ => None,
164            };
165
166            results.push(result);
167        }
168
169        Ok(results.to_vector())
170    }
171}
common_function/scalars/geo/relation.rs

common_function/scalars/geo/
relation.rs
