1use 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#[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 coord_type.clone(),
68 coord_type.clone(),
70 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#[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 coord_type.clone(),
159 coord_type.clone(),
161 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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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 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#[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#[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
1410fn 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}