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