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

common_function/scalars/geo/
s2.rs
