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

common_function/scalars/geo/
relation.rs
