1use 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#[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 coord_type.clone(),
112 coord_type.clone(),
114 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#[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 coord_type.clone(),
184 coord_type.clone(),
186 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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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 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#[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#[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
1267fn 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}