common_function/scalars/geo/
h3.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::str::FromStr;
16use std::sync::{Arc, LazyLock};
17
18use common_error::ext::{BoxedError, PlainError};
19use common_error::status_code::StatusCode;
20use common_query::error;
21use datafusion::arrow::array::{
22    Array, AsArray, BooleanBuilder, Float64Builder, Int32Builder, ListBuilder, StringViewArray,
23    StringViewBuilder, UInt8Builder, UInt64Builder,
24};
25use datafusion::arrow::compute;
26use datafusion::arrow::datatypes::{Float64Type, Int64Type, UInt8Type, UInt64Type};
27use datafusion::logical_expr::ColumnarValue;
28use datafusion_common::{DataFusionError, ScalarValue};
29use datafusion_expr::type_coercion::aggregates::INTEGERS;
30use datafusion_expr::{ScalarFunctionArgs, Signature, TypeSignature, Volatility};
31use datatypes::arrow::datatypes::{DataType, Field};
32use derive_more::Display;
33use h3o::{CellIndex, LatLng, Resolution};
34use snafu::prelude::*;
35
36use crate::function::{Function, extract_args};
37use crate::scalars::geo::helpers;
38
39static CELL_TYPES: LazyLock<Vec<DataType>> =
40    LazyLock::new(|| vec![DataType::Int64, DataType::UInt64, DataType::Utf8]);
41
42static COORDINATE_TYPES: LazyLock<Vec<DataType>> =
43    LazyLock::new(|| vec![DataType::Float32, DataType::Float64]);
44
45static RESOLUTION_TYPES: &[DataType] = INTEGERS;
46
47static DISTANCE_TYPES: &[DataType] = INTEGERS;
48
49static POSITION_TYPES: &[DataType] = INTEGERS;
50
51/// Function that returns [h3] encoding cellid for a given geospatial coordinate.
52///
53/// [h3]: https://h3geo.org/
54#[derive(Clone, Debug, Display)]
55#[display("{}", self.name())]
56pub(crate) struct H3LatLngToCell {
57    signature: Signature,
58}
59
60impl Default for H3LatLngToCell {
61    fn default() -> Self {
62        let mut signatures = Vec::new();
63        for coord_type in COORDINATE_TYPES.as_slice() {
64            for resolution_type in RESOLUTION_TYPES {
65                signatures.push(TypeSignature::Exact(vec![
66                    // latitude
67                    coord_type.clone(),
68                    // longitude
69                    coord_type.clone(),
70                    // resolution
71                    resolution_type.clone(),
72                ]));
73            }
74        }
75        Self {
76            signature: Signature::one_of(signatures, Volatility::Stable),
77        }
78    }
79}
80
81impl Function for H3LatLngToCell {
82    fn name(&self) -> &str {
83        "h3_latlng_to_cell"
84    }
85
86    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
87        Ok(DataType::UInt64)
88    }
89
90    fn signature(&self) -> &Signature {
91        &self.signature
92    }
93
94    fn invoke_with_args(
95        &self,
96        args: ScalarFunctionArgs,
97    ) -> datafusion_common::Result<ColumnarValue> {
98        let [lat_vec, lon_vec, resolution_vec] = extract_args(self.name(), &args)?;
99
100        let lat_vec = helpers::cast::<Float64Type>(&lat_vec)?;
101        let lat_vec = lat_vec.as_primitive::<Float64Type>();
102        let lon_vec = helpers::cast::<Float64Type>(&lon_vec)?;
103        let lon_vec = lon_vec.as_primitive::<Float64Type>();
104        let resolutions = helpers::cast::<UInt8Type>(&resolution_vec)?;
105        let resolution_vec = resolutions.as_primitive::<UInt8Type>();
106
107        let size = lat_vec.len();
108        let mut builder = UInt64Builder::with_capacity(size);
109
110        for i in 0..size {
111            let lat = lat_vec.is_valid(i).then(|| lat_vec.value(i));
112            let lon = lon_vec.is_valid(i).then(|| lon_vec.value(i));
113            let r = resolution_vec
114                .is_valid(i)
115                .then(|| value_to_resolution(resolution_vec.value(i)))
116                .transpose()?;
117
118            let result = match (lat, lon, r) {
119                (Some(lat), Some(lon), Some(r)) => {
120                    let coord = LatLng::new(lat, lon)
121                        .map_err(|e| {
122                            BoxedError::new(PlainError::new(
123                                format!("H3 error: {}", e),
124                                StatusCode::EngineExecuteQuery,
125                            ))
126                        })
127                        .context(error::ExecuteSnafu)?;
128                    let encoded: u64 = coord.to_cell(r).into();
129                    Some(encoded)
130                }
131                _ => None,
132            };
133
134            builder.append_option(result);
135        }
136
137        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
138    }
139}
140
141/// Function that returns [h3] encoding cellid in string form for a given
142/// geospatial coordinate.
143///
144/// [h3]: https://h3geo.org/
145#[derive(Clone, Debug, Display)]
146#[display("{}", self.name())]
147pub(crate) struct H3LatLngToCellString {
148    signature: Signature,
149}
150
151impl Default for H3LatLngToCellString {
152    fn default() -> Self {
153        let mut signatures = Vec::new();
154        for coord_type in COORDINATE_TYPES.as_slice() {
155            for resolution_type in RESOLUTION_TYPES {
156                signatures.push(TypeSignature::Exact(vec![
157                    // latitude
158                    coord_type.clone(),
159                    // longitude
160                    coord_type.clone(),
161                    // resolution
162                    resolution_type.clone(),
163                ]));
164            }
165        }
166        Self {
167            signature: Signature::one_of(signatures, Volatility::Stable),
168        }
169    }
170}
171
172impl Function for H3LatLngToCellString {
173    fn name(&self) -> &str {
174        "h3_latlng_to_cell_string"
175    }
176
177    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
178        Ok(DataType::Utf8View)
179    }
180
181    fn signature(&self) -> &Signature {
182        &self.signature
183    }
184
185    fn invoke_with_args(
186        &self,
187        args: ScalarFunctionArgs,
188    ) -> datafusion_common::Result<ColumnarValue> {
189        let [lat_vec, lon_vec, resolution_vec] = extract_args(self.name(), &args)?;
190
191        let lat_vec = helpers::cast::<Float64Type>(&lat_vec)?;
192        let lat_vec = lat_vec.as_primitive::<Float64Type>();
193        let lon_vec = helpers::cast::<Float64Type>(&lon_vec)?;
194        let lon_vec = lon_vec.as_primitive::<Float64Type>();
195        let resolutions = helpers::cast::<UInt8Type>(&resolution_vec)?;
196        let resolution_vec = resolutions.as_primitive::<UInt8Type>();
197
198        let size = lat_vec.len();
199        let mut builder = StringViewBuilder::with_capacity(size);
200
201        for i in 0..size {
202            let lat = lat_vec.is_valid(i).then(|| lat_vec.value(i));
203            let lon = lon_vec.is_valid(i).then(|| lon_vec.value(i));
204            let r = resolution_vec
205                .is_valid(i)
206                .then(|| value_to_resolution(resolution_vec.value(i)))
207                .transpose()?;
208
209            let result = match (lat, lon, r) {
210                (Some(lat), Some(lon), Some(r)) => {
211                    let coord = LatLng::new(lat, lon)
212                        .map_err(|e| {
213                            BoxedError::new(PlainError::new(
214                                format!("H3 error: {}", e),
215                                StatusCode::EngineExecuteQuery,
216                            ))
217                        })
218                        .context(error::ExecuteSnafu)?;
219                    let encoded = coord.to_cell(r).to_string();
220                    Some(encoded)
221                }
222                _ => None,
223            };
224
225            builder.append_option(result);
226        }
227
228        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
229    }
230}
231
232/// Function that converts cell id to its string form
233#[derive(Clone, Debug, Display)]
234#[display("{}", self.name())]
235pub(crate) struct H3CellToString {
236    signature: Signature,
237}
238
239impl Default for H3CellToString {
240    fn default() -> Self {
241        Self {
242            signature: signature_of_cell(),
243        }
244    }
245}
246
247impl Function for H3CellToString {
248    fn name(&self) -> &str {
249        "h3_cell_to_string"
250    }
251
252    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
253        Ok(DataType::Utf8View)
254    }
255
256    fn signature(&self) -> &Signature {
257        &self.signature
258    }
259
260    fn invoke_with_args(
261        &self,
262        args: ScalarFunctionArgs,
263    ) -> datafusion_common::Result<ColumnarValue> {
264        let [cell_vec] = extract_args(self.name(), &args)?;
265
266        let size = cell_vec.len();
267        let mut builder = StringViewBuilder::with_capacity(size);
268
269        for i in 0..size {
270            let v = ScalarValue::try_from_array(&cell_vec, i)
271                .and_then(cell_from_value)?
272                .map(|x| x.to_string());
273            builder.append_option(v);
274        }
275
276        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
277    }
278}
279
280/// Function that converts cell string id to uint64 number
281#[derive(Clone, Debug, Display)]
282#[display("{}", self.name())]
283pub(crate) struct H3StringToCell {
284    signature: Signature,
285}
286
287impl Default for H3StringToCell {
288    fn default() -> Self {
289        Self {
290            signature: Signature::string(1, Volatility::Stable),
291        }
292    }
293}
294
295impl Function for H3StringToCell {
296    fn name(&self) -> &str {
297        "h3_string_to_cell"
298    }
299
300    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
301        Ok(DataType::UInt64)
302    }
303
304    fn signature(&self) -> &Signature {
305        &self.signature
306    }
307
308    fn invoke_with_args(
309        &self,
310        args: ScalarFunctionArgs,
311    ) -> datafusion_common::Result<ColumnarValue> {
312        let [string_vec] = extract_args(self.name(), &args)?;
313        let string_vec = compute::cast(string_vec.as_ref(), &DataType::Utf8View)?;
314        let string_vec = datafusion_common::downcast_value!(string_vec, StringViewArray);
315
316        let size = string_vec.len();
317        let mut builder = UInt64Builder::with_capacity(size);
318
319        for i in 0..size {
320            let cell_id = string_vec
321                .is_valid(i)
322                .then(|| {
323                    CellIndex::from_str(string_vec.value(i))
324                        .map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))
325                        .map(Into::into)
326                })
327                .transpose()?;
328
329            builder.append_option(cell_id);
330        }
331
332        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
333    }
334}
335
336/// Function that returns centroid latitude and longitude of given cell id
337#[derive(Clone, Debug, Display)]
338#[display("{}", self.name())]
339pub(crate) struct H3CellCenterLatLng {
340    signature: Signature,
341}
342
343impl Default for H3CellCenterLatLng {
344    fn default() -> Self {
345        Self {
346            signature: signature_of_cell(),
347        }
348    }
349}
350
351impl Function for H3CellCenterLatLng {
352    fn name(&self) -> &str {
353        "h3_cell_center_latlng"
354    }
355
356    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
357        Ok(DataType::List(Arc::new(Field::new(
358            "x",
359            DataType::Float64,
360            false,
361        ))))
362    }
363
364    fn signature(&self) -> &Signature {
365        &self.signature
366    }
367
368    fn invoke_with_args(
369        &self,
370        args: ScalarFunctionArgs,
371    ) -> datafusion_common::Result<ColumnarValue> {
372        let [cell_vec] = extract_args(self.name(), &args)?;
373
374        let size = cell_vec.len();
375        let mut builder = ListBuilder::new(Float64Builder::new());
376
377        for i in 0..size {
378            let cell = ScalarValue::try_from_array(&cell_vec, i).and_then(cell_from_value)?;
379            let latlng = cell.map(LatLng::from);
380
381            if let Some(latlng) = latlng {
382                builder.values().append_value(latlng.lat());
383                builder.values().append_value(latlng.lng());
384                builder.append(true);
385            } else {
386                builder.append_null();
387            }
388        }
389
390        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
391    }
392}
393
394/// Function that returns resolution of given cell id
395#[derive(Clone, Debug, Display)]
396#[display("{}", self.name())]
397pub(crate) struct H3CellResolution {
398    signature: Signature,
399}
400
401impl Default for H3CellResolution {
402    fn default() -> Self {
403        Self {
404            signature: signature_of_cell(),
405        }
406    }
407}
408
409impl Function for H3CellResolution {
410    fn name(&self) -> &str {
411        "h3_cell_resolution"
412    }
413
414    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
415        Ok(DataType::UInt8)
416    }
417
418    fn signature(&self) -> &Signature {
419        &self.signature
420    }
421
422    fn invoke_with_args(
423        &self,
424        args: ScalarFunctionArgs,
425    ) -> datafusion_common::Result<ColumnarValue> {
426        let [cell_vec] = extract_args(self.name(), &args)?;
427
428        let size = cell_vec.len();
429        let mut builder = UInt8Builder::with_capacity(cell_vec.len());
430
431        for i in 0..size {
432            let cell = ScalarValue::try_from_array(&cell_vec, i).and_then(cell_from_value)?;
433            let res = cell.map(|cell| cell.resolution().into());
434            builder.append_option(res);
435        }
436
437        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
438    }
439}
440
441/// Function that returns base cell of given cell id
442#[derive(Clone, Debug, Display)]
443#[display("{}", self.name())]
444pub(crate) struct H3CellBase {
445    signature: Signature,
446}
447
448impl Default for H3CellBase {
449    fn default() -> Self {
450        Self {
451            signature: signature_of_cell(),
452        }
453    }
454}
455
456impl Function for H3CellBase {
457    fn name(&self) -> &str {
458        "h3_cell_base"
459    }
460
461    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
462        Ok(DataType::UInt8)
463    }
464
465    fn signature(&self) -> &Signature {
466        &self.signature
467    }
468
469    fn invoke_with_args(
470        &self,
471        args: ScalarFunctionArgs,
472    ) -> datafusion_common::Result<ColumnarValue> {
473        let [cell_vec] = extract_args(self.name(), &args)?;
474
475        let size = cell_vec.len();
476        let mut builder = UInt8Builder::with_capacity(size);
477
478        for i in 0..size {
479            let cell = ScalarValue::try_from_array(&cell_vec, i).and_then(cell_from_value)?;
480            let res = cell.map(|cell| cell.base_cell().into());
481
482            builder.append_option(res);
483        }
484
485        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
486    }
487}
488
489/// Function that check if given cell id is a pentagon
490#[derive(Clone, Debug, Display)]
491#[display("{}", self.name())]
492pub(crate) struct H3CellIsPentagon {
493    signature: Signature,
494}
495
496impl Default for H3CellIsPentagon {
497    fn default() -> Self {
498        Self {
499            signature: signature_of_cell(),
500        }
501    }
502}
503
504impl Function for H3CellIsPentagon {
505    fn name(&self) -> &str {
506        "h3_cell_is_pentagon"
507    }
508
509    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
510        Ok(DataType::Boolean)
511    }
512
513    fn signature(&self) -> &Signature {
514        &self.signature
515    }
516
517    fn invoke_with_args(
518        &self,
519        args: ScalarFunctionArgs,
520    ) -> datafusion_common::Result<ColumnarValue> {
521        let [cell_vec] = extract_args(self.name(), &args)?;
522
523        let size = cell_vec.len();
524        let mut builder = BooleanBuilder::with_capacity(size);
525
526        for i in 0..size {
527            let cell = ScalarValue::try_from_array(&cell_vec, i).and_then(cell_from_value)?;
528            let res = cell.map(|cell| cell.is_pentagon());
529
530            builder.append_option(res);
531        }
532
533        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
534    }
535}
536
537/// Function that returns center child cell of given cell id
538#[derive(Clone, Debug, Display)]
539#[display("{}", self.name())]
540pub(crate) struct H3CellCenterChild {
541    signature: Signature,
542}
543
544impl Default for H3CellCenterChild {
545    fn default() -> Self {
546        Self {
547            signature: signature_of_cell_and_resolution(),
548        }
549    }
550}
551
552impl Function for H3CellCenterChild {
553    fn name(&self) -> &str {
554        "h3_cell_center_child"
555    }
556
557    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
558        Ok(DataType::UInt64)
559    }
560
561    fn signature(&self) -> &Signature {
562        &self.signature
563    }
564
565    fn invoke_with_args(
566        &self,
567        args: ScalarFunctionArgs,
568    ) -> datafusion_common::Result<ColumnarValue> {
569        calculate_cell_child_property(self.name(), args, |cell, resolution| {
570            cell.center_child(resolution).map(Into::into)
571        })
572    }
573}
574
575/// Function that returns parent cell of given cell id and resolution
576#[derive(Clone, Debug, Display)]
577#[display("{}", self.name())]
578pub(crate) struct H3CellParent {
579    signature: Signature,
580}
581
582impl Default for H3CellParent {
583    fn default() -> Self {
584        Self {
585            signature: signature_of_cell_and_resolution(),
586        }
587    }
588}
589
590impl Function for H3CellParent {
591    fn name(&self) -> &str {
592        "h3_cell_parent"
593    }
594
595    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
596        Ok(DataType::UInt64)
597    }
598
599    fn signature(&self) -> &Signature {
600        &self.signature
601    }
602
603    fn invoke_with_args(
604        &self,
605        args: ScalarFunctionArgs,
606    ) -> datafusion_common::Result<ColumnarValue> {
607        calculate_cell_child_property(self.name(), args, |cell, resolution| {
608            cell.parent(resolution).map(Into::into)
609        })
610    }
611}
612
613/// Function that returns children cell list
614#[derive(Clone, Debug, Display)]
615#[display("{}", self.name())]
616pub(crate) struct H3CellToChildren {
617    signature: Signature,
618}
619
620impl Default for H3CellToChildren {
621    fn default() -> Self {
622        Self {
623            signature: signature_of_cell_and_resolution(),
624        }
625    }
626}
627
628impl Function for H3CellToChildren {
629    fn name(&self) -> &str {
630        "h3_cell_to_children"
631    }
632
633    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
634        Ok(DataType::List(Arc::new(Field::new(
635            "item",
636            DataType::UInt64,
637            true,
638        ))))
639    }
640
641    fn signature(&self) -> &Signature {
642        &self.signature
643    }
644
645    fn invoke_with_args(
646        &self,
647        args: ScalarFunctionArgs,
648    ) -> datafusion_common::Result<ColumnarValue> {
649        let [cell_vec, res_vec] = extract_args(self.name(), &args)?;
650        let resolutions = helpers::cast::<UInt8Type>(&res_vec)?;
651        let resolutions = resolutions.as_primitive::<UInt8Type>();
652
653        let size = cell_vec.len();
654        let mut builder = ListBuilder::new(UInt64Builder::new());
655
656        for i in 0..size {
657            let cell = ScalarValue::try_from_array(&cell_vec, i).and_then(cell_from_value)?;
658            let resolution = resolutions
659                .is_valid(i)
660                .then(|| value_to_resolution(resolutions.value(i)))
661                .transpose()?;
662
663            match (cell, resolution) {
664                (Some(c), Some(r)) => {
665                    for x in c.children(r) {
666                        builder.values().append_value(u64::from(x));
667                    }
668                    builder.append(true);
669                }
670                _ => builder.append_null(),
671            }
672        }
673
674        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
675    }
676}
677
678/// Function that returns children cell count
679#[derive(Clone, Debug, Display)]
680#[display("{}", self.name())]
681pub(crate) struct H3CellToChildrenSize {
682    signature: Signature,
683}
684
685impl Default for H3CellToChildrenSize {
686    fn default() -> Self {
687        Self {
688            signature: signature_of_cell_and_resolution(),
689        }
690    }
691}
692
693impl Function for H3CellToChildrenSize {
694    fn name(&self) -> &str {
695        "h3_cell_to_children_size"
696    }
697
698    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
699        Ok(DataType::UInt64)
700    }
701
702    fn signature(&self) -> &Signature {
703        &self.signature
704    }
705
706    fn invoke_with_args(
707        &self,
708        args: ScalarFunctionArgs,
709    ) -> datafusion_common::Result<ColumnarValue> {
710        calculate_cell_child_property(self.name(), args, |cell, resolution| {
711            Some(cell.children_count(resolution))
712        })
713    }
714}
715
716/// Function that returns the cell position if its parent at given resolution
717#[derive(Clone, Debug, Display)]
718#[display("{}", self.name())]
719pub(crate) struct H3CellToChildPos {
720    signature: Signature,
721}
722
723impl Default for H3CellToChildPos {
724    fn default() -> Self {
725        Self {
726            signature: signature_of_cell_and_resolution(),
727        }
728    }
729}
730
731impl Function for H3CellToChildPos {
732    fn name(&self) -> &str {
733        "h3_cell_to_child_pos"
734    }
735
736    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
737        Ok(DataType::UInt64)
738    }
739
740    fn signature(&self) -> &Signature {
741        &self.signature
742    }
743
744    fn invoke_with_args(
745        &self,
746        args: ScalarFunctionArgs,
747    ) -> datafusion_common::Result<ColumnarValue> {
748        calculate_cell_child_property(self.name(), args, |cell, resolution| {
749            cell.child_position(resolution)
750        })
751    }
752}
753
754fn calculate_cell_child_property<F>(
755    name: &str,
756    args: ScalarFunctionArgs,
757    calculator: F,
758) -> datafusion_common::Result<ColumnarValue>
759where
760    F: Fn(CellIndex, Resolution) -> Option<u64>,
761{
762    let [cells, resolutions] = extract_args(name, &args)?;
763    let resolutions = helpers::cast::<UInt8Type>(&resolutions)?;
764    let resolutions = resolutions.as_primitive::<UInt8Type>();
765
766    let mut builder = UInt64Builder::with_capacity(cells.len());
767    for i in 0..cells.len() {
768        let cell = ScalarValue::try_from_array(&cells, i).and_then(cell_from_value)?;
769        let resolution = resolutions
770            .is_valid(i)
771            .then(|| value_to_resolution(resolutions.value(i)))
772            .transpose()?;
773        let v = match (cell, resolution) {
774            (Some(c), Some(r)) => calculator(c, r),
775            _ => None,
776        };
777        builder.append_option(v);
778    }
779
780    Ok(ColumnarValue::Array(Arc::new(builder.finish())))
781}
782
783/// Function that returns the cell at given position of the parent at given resolution
784#[derive(Clone, Debug, Display)]
785#[display("{}", self.name())]
786pub(crate) struct H3ChildPosToCell {
787    signature: Signature,
788}
789
790impl Default for H3ChildPosToCell {
791    fn default() -> Self {
792        let mut signatures =
793            Vec::with_capacity(POSITION_TYPES.len() * CELL_TYPES.len() * RESOLUTION_TYPES.len());
794        for position_type in POSITION_TYPES {
795            for cell_type in CELL_TYPES.as_slice() {
796                for resolution_type in RESOLUTION_TYPES {
797                    signatures.push(TypeSignature::Exact(vec![
798                        position_type.clone(),
799                        cell_type.clone(),
800                        resolution_type.clone(),
801                    ]));
802                }
803            }
804        }
805        Self {
806            signature: Signature::one_of(signatures, Volatility::Stable),
807        }
808    }
809}
810
811impl Function for H3ChildPosToCell {
812    fn name(&self) -> &str {
813        "h3_child_pos_to_cell"
814    }
815
816    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
817        Ok(DataType::UInt64)
818    }
819
820    fn signature(&self) -> &Signature {
821        &self.signature
822    }
823
824    fn invoke_with_args(
825        &self,
826        args: ScalarFunctionArgs,
827    ) -> datafusion_common::Result<ColumnarValue> {
828        let [pos_vec, cell_vec, res_vec] = extract_args(self.name(), &args)?;
829        let resolutions = helpers::cast::<UInt8Type>(&res_vec)?;
830        let resolutions = resolutions.as_primitive::<UInt8Type>();
831
832        let size = cell_vec.len();
833        let mut builder = UInt64Builder::with_capacity(size);
834
835        for i in 0..size {
836            let cell = ScalarValue::try_from_array(&cell_vec, i).and_then(cell_from_value)?;
837            let pos = ScalarValue::try_from_array(&pos_vec, i).and_then(value_to_position)?;
838            let resolution = resolutions
839                .is_valid(i)
840                .then(|| value_to_resolution(resolutions.value(i)))
841                .transpose()?;
842            let result = match (cell, resolution) {
843                (Some(c), Some(r)) => c.child_at(pos, r).map(u64::from),
844                _ => None,
845            };
846            builder.append_option(result);
847        }
848
849        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
850    }
851}
852
853/// Function that returns cells with k distances of given cell
854#[derive(Clone, Debug, Display)]
855#[display("{}", self.name())]
856pub(crate) struct H3GridDisk {
857    signature: Signature,
858}
859
860impl Default for H3GridDisk {
861    fn default() -> Self {
862        Self {
863            signature: signature_of_cell_and_distance(),
864        }
865    }
866}
867
868impl Function for H3GridDisk {
869    fn name(&self) -> &str {
870        "h3_grid_disk"
871    }
872
873    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
874        Ok(DataType::List(Arc::new(Field::new(
875            "item",
876            DataType::UInt64,
877            true,
878        ))))
879    }
880
881    fn signature(&self) -> &Signature {
882        &self.signature
883    }
884
885    fn invoke_with_args(
886        &self,
887        args: ScalarFunctionArgs,
888    ) -> datafusion_common::Result<ColumnarValue> {
889        let [cell_vec, k_vec] = extract_args(self.name(), &args)?;
890
891        let size = cell_vec.len();
892        let mut builder = ListBuilder::new(UInt64Builder::new());
893
894        for i in 0..size {
895            let cell = ScalarValue::try_from_array(&cell_vec, i).and_then(cell_from_value)?;
896            let k = ScalarValue::try_from_array(&k_vec, i).and_then(value_to_distance)?;
897
898            if let Some(cell) = cell {
899                for x in cell.grid_disk::<Vec<_>>(k) {
900                    builder.values().append_value(u64::from(x));
901                }
902                builder.append(true);
903            } else {
904                builder.append_null();
905            }
906        }
907
908        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
909    }
910}
911
912/// Function that returns all cells within k distances of given cell
913#[derive(Clone, Debug, Display)]
914#[display("{}", self.name())]
915pub(crate) struct H3GridDiskDistances {
916    signature: Signature,
917}
918
919impl Default for H3GridDiskDistances {
920    fn default() -> Self {
921        Self {
922            signature: signature_of_cell_and_distance(),
923        }
924    }
925}
926
927impl Function for H3GridDiskDistances {
928    fn name(&self) -> &str {
929        "h3_grid_disk_distances"
930    }
931
932    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
933        Ok(DataType::List(Arc::new(Field::new(
934            "item",
935            DataType::UInt64,
936            true,
937        ))))
938    }
939
940    fn signature(&self) -> &Signature {
941        &self.signature
942    }
943
944    fn invoke_with_args(
945        &self,
946        args: ScalarFunctionArgs,
947    ) -> datafusion_common::Result<ColumnarValue> {
948        let [cell_vec, k_vec] = extract_args(self.name(), &args)?;
949
950        let size = cell_vec.len();
951        let mut builder = ListBuilder::new(UInt64Builder::new());
952
953        for i in 0..size {
954            let cell = ScalarValue::try_from_array(&cell_vec, i).and_then(cell_from_value)?;
955            let k = ScalarValue::try_from_array(&k_vec, i).and_then(value_to_distance)?;
956
957            if let Some(cell) = cell {
958                for (x, _) in cell.grid_disk_distances::<Vec<_>>(k) {
959                    builder.values().append_value(u64::from(x));
960                }
961                builder.append(true);
962            } else {
963                builder.append_null();
964            }
965        }
966
967        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
968    }
969}
970
971/// Function that returns distance between two cells
972#[derive(Clone, Debug, Display)]
973#[display("{}", self.name())]
974pub(crate) struct H3GridDistance {
975    signature: Signature,
976}
977
978impl Default for H3GridDistance {
979    fn default() -> Self {
980        Self {
981            signature: signature_of_double_cells(),
982        }
983    }
984}
985
986impl Function for H3GridDistance {
987    fn name(&self) -> &str {
988        "h3_grid_distance"
989    }
990    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
991        Ok(DataType::Int32)
992    }
993
994    fn signature(&self) -> &Signature {
995        &self.signature
996    }
997
998    fn invoke_with_args(
999        &self,
1000        args: ScalarFunctionArgs,
1001    ) -> datafusion_common::Result<ColumnarValue> {
1002        let [cell_this_vec, cell_that_vec] = extract_args(self.name(), &args)?;
1003
1004        let size = cell_this_vec.len();
1005        let mut builder = Int32Builder::with_capacity(size);
1006
1007        for i in 0..size {
1008            let cell_this =
1009                ScalarValue::try_from_array(&cell_this_vec, i).and_then(cell_from_value)?;
1010            let cell_that =
1011                ScalarValue::try_from_array(&cell_that_vec, i).and_then(cell_from_value)?;
1012            let result = match (cell_this, cell_that) {
1013                (Some(cell_this), Some(cell_that)) => {
1014                    let dist = cell_this
1015                        .grid_distance(cell_that)
1016                        .map_err(|e| {
1017                            BoxedError::new(PlainError::new(
1018                                format!("H3 error: {}", e),
1019                                StatusCode::EngineExecuteQuery,
1020                            ))
1021                        })
1022                        .context(error::ExecuteSnafu)?;
1023                    Some(dist)
1024                }
1025                _ => None,
1026            };
1027
1028            builder.append_option(result);
1029        }
1030
1031        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
1032    }
1033}
1034
1035/// Function that returns path cells between two cells
1036#[derive(Clone, Debug, Display)]
1037#[display("{}", self.name())]
1038pub(crate) struct H3GridPathCells {
1039    signature: Signature,
1040}
1041
1042impl Default for H3GridPathCells {
1043    fn default() -> Self {
1044        Self {
1045            signature: signature_of_double_cells(),
1046        }
1047    }
1048}
1049
1050impl Function for H3GridPathCells {
1051    fn name(&self) -> &str {
1052        "h3_grid_path_cells"
1053    }
1054
1055    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
1056        Ok(DataType::List(Arc::new(Field::new(
1057            "item",
1058            DataType::UInt64,
1059            true,
1060        ))))
1061    }
1062
1063    fn signature(&self) -> &Signature {
1064        &self.signature
1065    }
1066
1067    fn invoke_with_args(
1068        &self,
1069        args: ScalarFunctionArgs,
1070    ) -> datafusion_common::Result<ColumnarValue> {
1071        let [cell_this_vec, cell_that_vec] = extract_args(self.name(), &args)?;
1072
1073        let size = cell_this_vec.len();
1074        let mut builder = ListBuilder::new(UInt64Builder::new());
1075
1076        for i in 0..size {
1077            let cell_this =
1078                ScalarValue::try_from_array(&cell_this_vec, i).and_then(cell_from_value)?;
1079            let cell_that =
1080                ScalarValue::try_from_array(&cell_that_vec, i).and_then(cell_from_value)?;
1081            match (cell_this, cell_that) {
1082                (Some(cell_this), Some(cell_that)) => {
1083                    let cells = cell_this
1084                        .grid_path_cells(cell_that)
1085                        .map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))?;
1086                    for cell in cells {
1087                        let cell = cell
1088                            .map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))?;
1089                        builder.values().append_value(u64::from(cell));
1090                    }
1091                    builder.append(true);
1092                }
1093                _ => builder.append_null(),
1094            };
1095        }
1096
1097        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
1098    }
1099}
1100
1101/// Tests if cells contains given cells
1102#[derive(Clone, Debug, Display)]
1103#[display("{}", self.name())]
1104pub(crate) struct H3CellContains {
1105    signature: Signature,
1106}
1107
1108impl Default for H3CellContains {
1109    fn default() -> Self {
1110        let multi_cell_types = vec![
1111            DataType::new_list(DataType::Int64, true),
1112            DataType::new_list(DataType::UInt64, true),
1113            DataType::new_list(DataType::Utf8, true),
1114            DataType::Utf8,
1115        ];
1116
1117        let mut signatures = Vec::with_capacity(multi_cell_types.len() * CELL_TYPES.len());
1118        for multi_cell_type in &multi_cell_types {
1119            for cell_type in CELL_TYPES.as_slice() {
1120                signatures.push(TypeSignature::Exact(vec![
1121                    multi_cell_type.clone(),
1122                    cell_type.clone(),
1123                ]));
1124            }
1125        }
1126        Self {
1127            signature: Signature::one_of(signatures, Volatility::Stable),
1128        }
1129    }
1130}
1131
1132impl Function for H3CellContains {
1133    fn name(&self) -> &str {
1134        "h3_cells_contains"
1135    }
1136
1137    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
1138        Ok(DataType::Boolean)
1139    }
1140
1141    fn signature(&self) -> &Signature {
1142        &self.signature
1143    }
1144
1145    fn invoke_with_args(
1146        &self,
1147        args: ScalarFunctionArgs,
1148    ) -> datafusion_common::Result<ColumnarValue> {
1149        let [cells_vec, cell_this_vec] = extract_args(self.name(), &args)?;
1150
1151        let size = cell_this_vec.len();
1152        let mut builder = BooleanBuilder::with_capacity(size);
1153
1154        for i in 0..size {
1155            let cells = ScalarValue::try_from_array(&cells_vec, i).and_then(cells_from_value)?;
1156            let cell_this =
1157                ScalarValue::try_from_array(&cell_this_vec, i).and_then(cell_from_value)?;
1158            let mut result = None;
1159            if let (cells, Some(cell_this)) = (cells, cell_this) {
1160                result = Some(false);
1161
1162                for cell_that in cells.iter() {
1163                    // get cell resolution, and find cell_this's parent at
1164                    //  this solution, test if cell_that equals the parent
1165                    let resolution = cell_that.resolution();
1166                    if let Some(cell_this_parent) = cell_this.parent(resolution)
1167                        && cell_this_parent == *cell_that
1168                    {
1169                        result = Some(true);
1170                        break;
1171                    }
1172                }
1173            }
1174
1175            builder.append_option(result);
1176        }
1177
1178        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
1179    }
1180}
1181
1182/// Get WGS84 great circle distance of two cell centroid
1183#[derive(Clone, Debug, Display)]
1184#[display("{}", self.name())]
1185pub(crate) struct H3CellDistanceSphereKm {
1186    signature: Signature,
1187}
1188
1189impl Default for H3CellDistanceSphereKm {
1190    fn default() -> Self {
1191        Self {
1192            signature: signature_of_double_cells(),
1193        }
1194    }
1195}
1196
1197impl Function for H3CellDistanceSphereKm {
1198    fn name(&self) -> &str {
1199        "h3_distance_sphere_km"
1200    }
1201    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
1202        Ok(DataType::Float64)
1203    }
1204
1205    fn signature(&self) -> &Signature {
1206        &self.signature
1207    }
1208
1209    fn invoke_with_args(
1210        &self,
1211        args: ScalarFunctionArgs,
1212    ) -> datafusion_common::Result<ColumnarValue> {
1213        let [cell_this_vec, cell_that_vec] = extract_args(self.name(), &args)?;
1214
1215        let size = cell_this_vec.len();
1216        let mut builder = Float64Builder::with_capacity(size);
1217
1218        for i in 0..size {
1219            let cell_this =
1220                ScalarValue::try_from_array(&cell_this_vec, i).and_then(cell_from_value)?;
1221            let cell_that =
1222                ScalarValue::try_from_array(&cell_that_vec, i).and_then(cell_from_value)?;
1223            let result = match (cell_this, cell_that) {
1224                (Some(cell_this), Some(cell_that)) => {
1225                    let centroid_this = LatLng::from(cell_this);
1226                    let centroid_that = LatLng::from(cell_that);
1227
1228                    Some(centroid_this.distance_km(centroid_that))
1229                }
1230                _ => None,
1231            };
1232
1233            builder.append_option(result);
1234        }
1235
1236        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
1237    }
1238}
1239
1240/// Get Euclidean distance of two cell centroid
1241#[derive(Clone, Debug, Display)]
1242#[display("{}", self.name())]
1243pub(crate) struct H3CellDistanceEuclideanDegree {
1244    signature: Signature,
1245}
1246
1247impl Default for H3CellDistanceEuclideanDegree {
1248    fn default() -> Self {
1249        Self {
1250            signature: signature_of_double_cells(),
1251        }
1252    }
1253}
1254
1255impl H3CellDistanceEuclideanDegree {
1256    fn distance(centroid_this: LatLng, centroid_that: LatLng) -> f64 {
1257        ((centroid_this.lat() - centroid_that.lat()).powi(2)
1258            + (centroid_this.lng() - centroid_that.lng()).powi(2))
1259        .sqrt()
1260    }
1261}
1262
1263impl Function for H3CellDistanceEuclideanDegree {
1264    fn name(&self) -> &str {
1265        "h3_distance_degree"
1266    }
1267    fn return_type(&self, _: &[DataType]) -> datafusion_common::Result<DataType> {
1268        Ok(DataType::Float64)
1269    }
1270
1271    fn signature(&self) -> &Signature {
1272        &self.signature
1273    }
1274
1275    fn invoke_with_args(
1276        &self,
1277        args: ScalarFunctionArgs,
1278    ) -> datafusion_common::Result<ColumnarValue> {
1279        let [cell_this_vec, cell_that_vec] = extract_args(self.name(), &args)?;
1280
1281        let size = cell_this_vec.len();
1282        let mut builder = Float64Builder::with_capacity(size);
1283
1284        for i in 0..size {
1285            let cell_this =
1286                ScalarValue::try_from_array(&cell_this_vec, i).and_then(cell_from_value)?;
1287            let cell_that =
1288                ScalarValue::try_from_array(&cell_that_vec, i).and_then(cell_from_value)?;
1289            let result = match (cell_this, cell_that) {
1290                (Some(cell_this), Some(cell_that)) => {
1291                    let centroid_this = LatLng::from(cell_this);
1292                    let centroid_that = LatLng::from(cell_that);
1293
1294                    let dist = Self::distance(centroid_this, centroid_that);
1295                    Some(dist)
1296                }
1297                _ => None,
1298            };
1299
1300            builder.append_option(result);
1301        }
1302
1303        Ok(ColumnarValue::Array(Arc::new(builder.finish())))
1304    }
1305}
1306
1307fn value_to_resolution(r: u8) -> datafusion_common::Result<Resolution> {
1308    Resolution::try_from(r).map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))
1309}
1310
1311macro_rules! ensure_then_coerce {
1312    ($compare:expr, $coerce:expr) => {{
1313        if !$compare {
1314            return Err(datafusion_common::DataFusionError::Execution(
1315                "Argument was outside of acceptable range".to_string(),
1316            ));
1317        }
1318        Ok($coerce)
1319    }};
1320}
1321
1322fn value_to_position(v: ScalarValue) -> datafusion_common::Result<u64> {
1323    match v {
1324        ScalarValue::Int8(Some(v)) => ensure_then_coerce!(v >= 0, v as u64),
1325        ScalarValue::Int16(Some(v)) => ensure_then_coerce!(v >= 0, v as u64),
1326        ScalarValue::Int32(Some(v)) => ensure_then_coerce!(v >= 0, v as u64),
1327        ScalarValue::Int64(Some(v)) => ensure_then_coerce!(v >= 0, v as u64),
1328        ScalarValue::UInt8(Some(v)) => Ok(v as u64),
1329        ScalarValue::UInt16(Some(v)) => Ok(v as u64),
1330        ScalarValue::UInt32(Some(v)) => Ok(v as u64),
1331        ScalarValue::UInt64(Some(v)) => Ok(v),
1332        _ => unreachable!(),
1333    }
1334}
1335
1336fn value_to_distance(v: ScalarValue) -> datafusion_common::Result<u32> {
1337    match v {
1338        ScalarValue::Int8(Some(v)) => ensure_then_coerce!(v >= 0, v as u32),
1339        ScalarValue::Int16(Some(v)) => ensure_then_coerce!(v >= 0, v as u32),
1340        ScalarValue::Int32(Some(v)) => ensure_then_coerce!(v >= 0, v as u32),
1341        ScalarValue::Int64(Some(v)) => ensure_then_coerce!(v >= 0, v as u32),
1342        ScalarValue::UInt8(Some(v)) => Ok(v as u32),
1343        ScalarValue::UInt16(Some(v)) => Ok(v as u32),
1344        ScalarValue::UInt32(Some(v)) => Ok(v),
1345        ScalarValue::UInt64(Some(v)) => Ok(v as u32),
1346        _ => unreachable!(),
1347    }
1348}
1349
1350fn signature_of_cell() -> Signature {
1351    let mut signatures = Vec::with_capacity(CELL_TYPES.len());
1352    for cell_type in CELL_TYPES.as_slice() {
1353        signatures.push(TypeSignature::Exact(vec![cell_type.clone()]));
1354    }
1355
1356    Signature::one_of(signatures, Volatility::Stable)
1357}
1358
1359fn signature_of_double_cells() -> Signature {
1360    let mut signatures = Vec::with_capacity(CELL_TYPES.len() * CELL_TYPES.len());
1361    for cell_type in CELL_TYPES.as_slice() {
1362        for cell_type2 in CELL_TYPES.as_slice() {
1363            signatures.push(TypeSignature::Exact(vec![
1364                cell_type.clone(),
1365                cell_type2.clone(),
1366            ]));
1367        }
1368    }
1369
1370    Signature::one_of(signatures, Volatility::Stable)
1371}
1372
1373fn signature_of_cell_and_resolution() -> Signature {
1374    let mut signatures = Vec::with_capacity(CELL_TYPES.len() * RESOLUTION_TYPES.len());
1375    for cell_type in CELL_TYPES.as_slice() {
1376        for resolution_type in RESOLUTION_TYPES {
1377            signatures.push(TypeSignature::Exact(vec![
1378                cell_type.clone(),
1379                resolution_type.clone(),
1380            ]));
1381        }
1382    }
1383    Signature::one_of(signatures, Volatility::Stable)
1384}
1385
1386fn signature_of_cell_and_distance() -> Signature {
1387    let mut signatures = Vec::with_capacity(CELL_TYPES.len() * DISTANCE_TYPES.len());
1388    for cell_type in CELL_TYPES.as_slice() {
1389        for distance_type in DISTANCE_TYPES {
1390            signatures.push(TypeSignature::Exact(vec![
1391                cell_type.clone(),
1392                distance_type.clone(),
1393            ]));
1394        }
1395    }
1396    Signature::one_of(signatures, Volatility::Stable)
1397}
1398
1399fn cell_from_value(v: ScalarValue) -> datafusion_common::Result<Option<CellIndex>> {
1400    match v {
1401        ScalarValue::Int64(Some(v)) => Some(CellIndex::try_from(v as u64)),
1402        ScalarValue::UInt64(Some(v)) => Some(CellIndex::try_from(v)),
1403        ScalarValue::Utf8(Some(s)) => Some(CellIndex::from_str(&s)),
1404        _ => None,
1405    }
1406    .transpose()
1407    .map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))
1408}
1409
1410/// extract cell array from all possible types including:
1411/// - int64 list
1412/// - uint64 list
1413/// - string list
1414/// - comma-separated string
1415fn cells_from_value(v: ScalarValue) -> datafusion_common::Result<Vec<CellIndex>> {
1416    match v {
1417        ScalarValue::List(list) => match list.value_type() {
1418            DataType::Int64 => list
1419                .values()
1420                .as_primitive::<Int64Type>()
1421                .iter()
1422                .map(|v| {
1423                    if let Some(v) = v {
1424                        CellIndex::try_from(v as u64)
1425                            .map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))
1426                    } else {
1427                        Err(DataFusionError::Execution(
1428                            "Invalid data type in array".to_string(),
1429                        ))
1430                    }
1431                })
1432                .collect::<datafusion_common::Result<Vec<CellIndex>>>(),
1433            DataType::UInt64 => list
1434                .values()
1435                .as_primitive::<UInt64Type>()
1436                .iter()
1437                .map(|v| {
1438                    if let Some(v) = v {
1439                        CellIndex::try_from(v)
1440                            .map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))
1441                    } else {
1442                        Err(DataFusionError::Execution(
1443                            "Invalid data type in array".to_string(),
1444                        ))
1445                    }
1446                })
1447                .collect::<datafusion_common::Result<Vec<CellIndex>>>(),
1448            DataType::Utf8 => list
1449                .values()
1450                .as_string::<i32>()
1451                .iter()
1452                .map(|v| {
1453                    if let Some(v) = v {
1454                        CellIndex::from_str(v)
1455                            .map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))
1456                    } else {
1457                        Err(DataFusionError::Execution(
1458                            "Invalid data type in array".to_string(),
1459                        ))
1460                    }
1461                })
1462                .collect::<datafusion_common::Result<Vec<CellIndex>>>(),
1463            _ => Ok(vec![]),
1464        },
1465        ScalarValue::Utf8(Some(csv)) => {
1466            let str_seq = csv.split(',');
1467            str_seq
1468                .map(|v| {
1469                    CellIndex::from_str(v.trim())
1470                        .map_err(|e| DataFusionError::Execution(format!("H3 error: {}", e)))
1471                })
1472                .collect::<datafusion_common::Result<Vec<CellIndex>>>()
1473        }
1474        _ => Ok(vec![]),
1475    }
1476}
1477
1478#[cfg(test)]
1479mod tests {
1480    use super::*;
1481
1482    #[test]
1483    fn test_h3_euclidean_distance() {
1484        let point_this = LatLng::new(42.3521, -72.1235).expect("incorrect lat lng");
1485        let point_that = LatLng::new(42.45, -72.1260).expect("incorrect lat lng");
1486
1487        let dist = H3CellDistanceEuclideanDegree::distance(point_this, point_that);
1488        assert_eq!(dist, 0.09793191512474639);
1489    }
1490}