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::{Arc, LazyLock};
16
17use common_query::error::InvalidFuncArgsSnafu;
18use datafusion_common::ScalarValue;
19use datafusion_common::arrow::array::{Array, AsArray, StringViewBuilder, UInt64Builder};
20use datafusion_common::arrow::datatypes::{DataType, Float64Type};
21use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature, TypeSignature, Volatility};
22use derive_more::Display;
23use s2::cellid::{CellID, MAX_LEVEL};
24use s2::latlng::LatLng;
25use snafu::ensure;
26
27use crate::function::{Function, extract_args};
28use crate::helper::INTEGER_TYPES;
29use crate::scalars::geo::helpers;
30use crate::scalars::geo::helpers::ensure_and_coerce;
31
32static CELL_TYPES: LazyLock<Vec<DataType>> =
33    LazyLock::new(|| vec![DataType::Int64, DataType::UInt64]);
34
35static COORDINATE_TYPES: LazyLock<Vec<DataType>> =
36    LazyLock::new(|| vec![DataType::Float32, DataType::Float64]);
37
38static LEVEL_TYPES: &[DataType] = INTEGER_TYPES;
39
40/// Function that returns [s2] encoding cellid for a given geospatial coordinate.
41///
42/// [s2]: http://s2geometry.io
43#[derive(Clone, Debug, Display)]
44#[display("{}", self.name())]
45pub(crate) struct S2LatLngToCell {
46    signature: Signature,
47}
48
49impl Default for S2LatLngToCell {
50    fn default() -> Self {
51        let mut signatures = Vec::with_capacity(COORDINATE_TYPES.len());
52        for coord_type in COORDINATE_TYPES.as_slice() {
53            signatures.push(TypeSignature::Exact(vec![
54                // latitude
55                coord_type.clone(),
56                // longitude
57                coord_type.clone(),
58            ]));
59        }
60        Self {
61            signature: Signature::one_of(signatures, Volatility::Stable),
62        }
63    }
64}
65
66impl Function for S2LatLngToCell {
67    fn name(&self) -> &str {
68        "s2_latlng_to_cell"
69    }
70
71    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
72        Ok(DataType::UInt64)
73    }
74
75    fn signature(&self) -> &Signature {
76        &self.signature
77    }
78
79    fn invoke_with_args(
80        &self,
81        args: ScalarFunctionArgs,
82    ) -> datafusion_common::Result<ColumnarValue> {
83        let [arg0, arg1] = extract_args(self.name(), &args)?;
84
85        let arg0 = helpers::cast::<Float64Type>(&arg0)?;
86        let lat_vec = arg0.as_primitive::<Float64Type>();
87        let arg1 = helpers::cast::<Float64Type>(&arg1)?;
88        let lon_vec = arg1.as_primitive::<Float64Type>();
89
90        let size = lat_vec.len();
91        let mut builder = UInt64Builder::with_capacity(size);
92
93        for i in 0..size {
94            let lat = lat_vec.is_valid(i).then(|| lat_vec.value(i));
95            let lon = lon_vec.is_valid(i).then(|| lon_vec.value(i));
96
97            let result = match (lat, lon) {
98                (Some(lat), Some(lon)) => {
99                    let coord = LatLng::from_degrees(lat, lon);
100                    ensure!(
101                        coord.is_valid(),
102                        InvalidFuncArgsSnafu {
103                            err_msg: "The input coordinates are invalid",
104                        }
105                    );
106                    let cellid = CellID::from(coord);
107                    let encoded: u64 = cellid.0;
108                    Some(encoded)
109                }
110                _ => None,
111            };
112
113            builder.append_option(result);
114        }
115
116        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
117    }
118}
119
120/// Return the level of current s2 cell
121#[derive(Clone, Debug, Display)]
122#[display("{}", self.name())]
123pub(crate) struct S2CellLevel {
124    signature: Signature,
125}
126
127impl Default for S2CellLevel {
128    fn default() -> Self {
129        Self {
130            signature: signature_of_cell(),
131        }
132    }
133}
134
135impl Function for S2CellLevel {
136    fn name(&self) -> &str {
137        "s2_cell_level"
138    }
139
140    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
141        Ok(DataType::UInt64)
142    }
143
144    fn signature(&self) -> &Signature {
145        &self.signature
146    }
147
148    fn invoke_with_args(
149        &self,
150        args: ScalarFunctionArgs,
151    ) -> datafusion_common::Result<ColumnarValue> {
152        let [cell_vec] = extract_args(self.name(), &args)?;
153
154        let size = cell_vec.len();
155        let mut builder = UInt64Builder::with_capacity(size);
156
157        for i in 0..size {
158            let v = ScalarValue::try_from_array(&cell_vec, i)?;
159            let v = cell_from_value(v).map(|x| x.level());
160
161            builder.append_option(v);
162        }
163
164        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
165    }
166}
167
168/// Return the string presentation of the cell
169#[derive(Clone, Debug, Display)]
170#[display("{}", self.name())]
171pub(crate) struct S2CellToToken {
172    signature: Signature,
173}
174
175impl Default for S2CellToToken {
176    fn default() -> Self {
177        Self {
178            signature: signature_of_cell(),
179        }
180    }
181}
182
183impl Function for S2CellToToken {
184    fn name(&self) -> &str {
185        "s2_cell_to_token"
186    }
187
188    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
189        Ok(DataType::Utf8View)
190    }
191
192    fn signature(&self) -> &Signature {
193        &self.signature
194    }
195
196    fn invoke_with_args(
197        &self,
198        args: ScalarFunctionArgs,
199    ) -> datafusion_common::Result<ColumnarValue> {
200        let [cell_vec] = extract_args(self.name(), &args)?;
201
202        let size = cell_vec.len();
203        let mut builder = StringViewBuilder::with_capacity(size);
204
205        for i in 0..size {
206            let v = ScalarValue::try_from_array(&cell_vec, i)?;
207            let v = cell_from_value(v).map(|x| x.to_token());
208
209            builder.append_option(v.as_deref());
210        }
211
212        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
213    }
214}
215
216/// Return parent at given level of current s2 cell
217#[derive(Clone, Debug, Display)]
218#[display("{}", self.name())]
219pub(crate) struct S2CellParent {
220    signature: Signature,
221}
222
223impl Default for S2CellParent {
224    fn default() -> Self {
225        Self {
226            signature: signature_of_cell_and_level(),
227        }
228    }
229}
230
231impl Function for S2CellParent {
232    fn name(&self) -> &str {
233        "s2_cell_parent"
234    }
235
236    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
237        Ok(DataType::UInt64)
238    }
239
240    fn signature(&self) -> &Signature {
241        &self.signature
242    }
243
244    fn invoke_with_args(
245        &self,
246        args: ScalarFunctionArgs,
247    ) -> datafusion_common::Result<ColumnarValue> {
248        let [cell_vec, levels] = extract_args(self.name(), &args)?;
249
250        let size = cell_vec.len();
251        let mut builder = UInt64Builder::with_capacity(size);
252
253        for i in 0..size {
254            let cell = ScalarValue::try_from_array(&cell_vec, i).map(cell_from_value)?;
255            let level = ScalarValue::try_from_array(&levels, i).and_then(value_to_level)?;
256            let result = if let (Some(cell), Some(level)) = (cell, level) {
257                Some(cell.parent(level).0)
258            } else {
259                None
260            };
261
262            builder.append_option(result);
263        }
264
265        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
266    }
267}
268
269fn signature_of_cell() -> Signature {
270    let mut signatures = Vec::with_capacity(CELL_TYPES.len());
271    for cell_type in CELL_TYPES.as_slice() {
272        signatures.push(TypeSignature::Exact(vec![cell_type.clone()]));
273    }
274
275    Signature::one_of(signatures, Volatility::Stable)
276}
277
278fn signature_of_cell_and_level() -> Signature {
279    let mut signatures = Vec::with_capacity(CELL_TYPES.len() * LEVEL_TYPES.len());
280    for cell_type in CELL_TYPES.as_slice() {
281        for level_type in LEVEL_TYPES {
282            signatures.push(TypeSignature::Exact(vec![
283                cell_type.clone(),
284                level_type.clone(),
285            ]));
286        }
287    }
288    Signature::one_of(signatures, Volatility::Stable)
289}
290
291fn cell_from_value(v: ScalarValue) -> Option<CellID> {
292    match v {
293        ScalarValue::Int64(v) => v.map(|x| CellID(x as u64)),
294        ScalarValue::UInt64(v) => v.map(CellID),
295        _ => None,
296    }
297}
298
299fn value_to_level(v: ScalarValue) -> datafusion_common::Result<Option<u64>> {
300    match v {
301        ScalarValue::Int8(Some(v)) => ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i8, v as u64),
302        ScalarValue::Int16(Some(v)) => {
303            ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i16, v as u64)
304        }
305        ScalarValue::Int32(Some(v)) => {
306            ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i32, v as u64)
307        }
308        ScalarValue::Int64(Some(v)) => {
309            ensure_and_coerce!(v >= 0 && v <= MAX_LEVEL as i64, v as u64)
310        }
311        ScalarValue::UInt8(Some(v)) => ensure_and_coerce!(v <= MAX_LEVEL as u8, v as u64),
312        ScalarValue::UInt16(Some(v)) => ensure_and_coerce!(v <= MAX_LEVEL as u16, v as u64),
313        ScalarValue::UInt32(Some(v)) => ensure_and_coerce!(v <= MAX_LEVEL as u32, v as u64),
314        ScalarValue::UInt64(Some(v)) => ensure_and_coerce!(v <= MAX_LEVEL, v),
315        _ => Ok(None),
316    }
317}
common_function/scalars/geo/s2.rs

common_function/scalars/geo/
s2.rs
