common_function/scalars/geo/
s2.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::{InvalidFuncArgsSnafu, Result};
16use common_query::prelude::{Signature, TypeSignature};
17use datafusion::logical_expr::Volatility;
18use datatypes::prelude::ConcreteDataType;
19use datatypes::scalars::ScalarVectorBuilder;
20use datatypes::value::Value;
21use datatypes::vectors::{MutableVector, StringVectorBuilder, UInt64VectorBuilder, VectorRef};
22use derive_more::Display;
23use once_cell::sync::Lazy;
24use s2::cellid::{CellID, MAX_LEVEL};
25use s2::latlng::LatLng;
26use snafu::ensure;
27
28use crate::function::{Function, FunctionContext};
29use crate::scalars::geo::helpers::{ensure_and_coerce, ensure_columns_len, ensure_columns_n};
30
31static CELL_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
32    vec![
33        ConcreteDataType::int64_datatype(),
34        ConcreteDataType::uint64_datatype(),
35    ]
36});
37
38static COORDINATE_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
39    vec![
40        ConcreteDataType::float32_datatype(),
41        ConcreteDataType::float64_datatype(),
42    ]
43});
44
45static LEVEL_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
46    vec![
47        ConcreteDataType::int8_datatype(),
48        ConcreteDataType::int16_datatype(),
49        ConcreteDataType::int32_datatype(),
50        ConcreteDataType::int64_datatype(),
51        ConcreteDataType::uint8_datatype(),
52        ConcreteDataType::uint16_datatype(),
53        ConcreteDataType::uint32_datatype(),
54        ConcreteDataType::uint64_datatype(),
55    ]
56});
57
58/// Function that returns [s2] encoding cellid for a given geospatial coordinate.
59///
60/// [s2]: http://s2geometry.io
61#[derive(Clone, Debug, Default, Display)]
62#[display("{}", self.name())]
63pub struct S2LatLngToCell;
64
65impl Function for S2LatLngToCell {
66    fn name(&self) -> &str {
67        "s2_latlng_to_cell"
68    }
69
70    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
71        Ok(ConcreteDataType::uint64_datatype())
72    }
73
74    fn signature(&self) -> Signature {
75        let mut signatures = Vec::with_capacity(COORDINATE_TYPES.len());
76        for coord_type in COORDINATE_TYPES.as_slice() {
77            signatures.push(TypeSignature::Exact(vec![
78                // latitude
79                coord_type.clone(),
80                // longitude
81                coord_type.clone(),
82            ]));
83        }
84        Signature::one_of(signatures, Volatility::Stable)
85    }
86
87    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
88        ensure_columns_n!(columns, 2);
89
90        let lat_vec = &columns[0];
91        let lon_vec = &columns[1];
92
93        let size = lat_vec.len();
94        let mut results = UInt64VectorBuilder::with_capacity(size);
95
96        for i in 0..size {
97            let lat = lat_vec.get(i).as_f64_lossy();
98            let lon = lon_vec.get(i).as_f64_lossy();
99
100            let result = match (lat, lon) {
101                (Some(lat), Some(lon)) => {
102                    let coord = LatLng::from_degrees(lat, lon);
103                    ensure!(
104                        coord.is_valid(),
105                        InvalidFuncArgsSnafu {
106                            err_msg: "The input coordinates are invalid",
107                        }
108                    );
109                    let cellid = CellID::from(coord);
110                    let encoded: u64 = cellid.0;
111                    Some(encoded)
112                }
113                _ => None,
114            };
115
116            results.push(result);
117        }
118
119        Ok(results.to_vector())
120    }
121}
122
123/// Return the level of current s2 cell
124#[derive(Clone, Debug, Default, Display)]
125#[display("{}", self.name())]
126pub struct S2CellLevel;
127
128impl Function for S2CellLevel {
129    fn name(&self) -> &str {
130        "s2_cell_level"
131    }
132
133    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
134        Ok(ConcreteDataType::uint64_datatype())
135    }
136
137    fn signature(&self) -> Signature {
138        signature_of_cell()
139    }
140
141    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
142        ensure_columns_n!(columns, 1);
143
144        let cell_vec = &columns[0];
145        let size = cell_vec.len();
146        let mut results = UInt64VectorBuilder::with_capacity(size);
147
148        for i in 0..size {
149            let cell = cell_from_value(cell_vec.get(i));
150            let res = cell.map(|cell| cell.level());
151
152            results.push(res);
153        }
154
155        Ok(results.to_vector())
156    }
157}
158
159/// Return the string presentation of the cell
160#[derive(Clone, Debug, Default, Display)]
161#[display("{}", self.name())]
162pub struct S2CellToToken;
163
164impl Function for S2CellToToken {
165    fn name(&self) -> &str {
166        "s2_cell_to_token"
167    }
168
169    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
170        Ok(ConcreteDataType::string_datatype())
171    }
172
173    fn signature(&self) -> Signature {
174        signature_of_cell()
175    }
176
177    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
178        ensure_columns_n!(columns, 1);
179
180        let cell_vec = &columns[0];
181        let size = cell_vec.len();
182        let mut results = StringVectorBuilder::with_capacity(size);
183
184        for i in 0..size {
185            let cell = cell_from_value(cell_vec.get(i));
186            let res = cell.map(|cell| cell.to_token());
187
188            results.push(res.as_deref());
189        }
190
191        Ok(results.to_vector())
192    }
193}
194
195/// Return parent at given level of current s2 cell
196#[derive(Clone, Debug, Default, Display)]
197#[display("{}", self.name())]
198pub struct S2CellParent;
199
200impl Function for S2CellParent {
201    fn name(&self) -> &str {
202        "s2_cell_parent"
203    }
204
205    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
206        Ok(ConcreteDataType::uint64_datatype())
207    }
208
209    fn signature(&self) -> Signature {
210        signature_of_cell_and_level()
211    }
212
213    fn eval(&self, _func_ctx: &FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
214        ensure_columns_n!(columns, 2);
215
216        let cell_vec = &columns[0];
217        let level_vec = &columns[1];
218        let size = cell_vec.len();
219        let mut results = UInt64VectorBuilder::with_capacity(size);
220
221        for i in 0..size {
222            let cell = cell_from_value(cell_vec.get(i));
223            let level = value_to_level(level_vec.get(i))?;
224            let result = cell.map(|cell| cell.parent(level).0);
225
226            results.push(result);
227        }
228
229        Ok(results.to_vector())
230    }
231}
232
233fn signature_of_cell() -> Signature {
234    let mut signatures = Vec::with_capacity(CELL_TYPES.len());
235    for cell_type in CELL_TYPES.as_slice() {
236        signatures.push(TypeSignature::Exact(vec![cell_type.clone()]));
237    }
238
239    Signature::one_of(signatures, Volatility::Stable)
240}
241
242fn signature_of_cell_and_level() -> Signature {
243    let mut signatures = Vec::with_capacity(CELL_TYPES.len() * LEVEL_TYPES.len());
244    for cell_type in CELL_TYPES.as_slice() {
245        for level_type in LEVEL_TYPES.as_slice() {
246            signatures.push(TypeSignature::Exact(vec![
247                cell_type.clone(),
248                level_type.clone(),
249            ]));
250        }
251    }
252    Signature::one_of(signatures, Volatility::Stable)
253}
254
255fn cell_from_value(v: Value) -> Option<CellID> {
256    match v {
257        Value::Int64(v) => Some(CellID(v as u64)),
258        Value::UInt64(v) => Some(CellID(v)),
259        _ => None,
260    }
261}
262
263fn value_to_level(v: Value) -> Result<u64> {
264    match v {
265        Value::Int8(v) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i8, v as u64),
266        Value::Int16(v) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i16, v as u64),
267        Value::Int32(v) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i32, v as u64),
268        Value::Int64(v) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i64, v as u64),
269        Value::UInt8(v) => ensure_and_coerce!(v <= MAX_LEVEL as u8, v as u64),
270        Value::UInt16(v) => ensure_and_coerce!(v <= MAX_LEVEL as u16, v as u64),
271        Value::UInt32(v) => ensure_and_coerce!(v <= MAX_LEVEL as u32, v as u64),
272        Value::UInt64(v) => ensure_and_coerce!(v <= MAX_LEVEL, v),
273        _ => unreachable!(),
274    }
275}
common_function/scalars/geo/s2.rs

common_function/scalars/geo/
s2.rs
