rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-21 23:40:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add more h3 functions (#4770)

Browse Source 
* feat: add more h3 grid functions

* feat: add more traversal functions

* refactor: update some function definitions

* style: format

* refactor: avoid creating slice in nested loop

* feat: ensure column number and length

* refactor: fix lint warnings

* refactor: merge main

* Apply suggestions from code review

Co-authored-by: LFC <990479+MichaelScofield@users.noreply.github.com>

* Update src/common/function/src/scalars/geo/h3.rs

Co-authored-by: LFC <990479+MichaelScofield@users.noreply.github.com>

* style: format

---------

Co-authored-by: LFC <990479+MichaelScofield@users.noreply.github.com>
This commit is contained in:
Ning Sun
2024-10-11 10:57:54 -07:00
committed by GitHub
parent 4bb1f4f184
commit aaa9b32908
4 changed files with 595 additions and 243 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/common/function/src/scalars/geo.rs
View File

@@ -29,18 +29,31 @@ impl GeoFunctions {
// geohash
registry.register(Arc::new(GeohashFunction));
registry.register(Arc::new(GeohashNeighboursFunction));
// h3 family
// h3 index
registry.register(Arc::new(h3::H3LatLngToCell));
registry.register(Arc::new(h3::H3LatLngToCellString));
// h3 index inspection
registry.register(Arc::new(h3::H3CellBase));
registry.register(Arc::new(h3::H3CellCenterChild));
registry.register(Arc::new(h3::H3CellCenterLat));
registry.register(Arc::new(h3::H3CellCenterLng));
registry.register(Arc::new(h3::H3CellIsPentagon));
registry.register(Arc::new(h3::H3CellParent));
registry.register(Arc::new(h3::H3CellResolution));
registry.register(Arc::new(h3::H3CellToString));
registry.register(Arc::new(h3::H3IsNeighbour));
registry.register(Arc::new(h3::H3StringToCell));
registry.register(Arc::new(h3::H3CellToString));
registry.register(Arc::new(h3::H3CellCenterLatLng));
registry.register(Arc::new(h3::H3CellResolution));
// h3 hierarchical grid
registry.register(Arc::new(h3::H3CellCenterChild));
registry.register(Arc::new(h3::H3CellParent));
registry.register(Arc::new(h3::H3CellToChildren));
registry.register(Arc::new(h3::H3CellToChildrenSize));
registry.register(Arc::new(h3::H3CellToChildPos));
registry.register(Arc::new(h3::H3ChildPosToCell));
// h3 grid traversal
registry.register(Arc::new(h3::H3GridDisk));
registry.register(Arc::new(h3::H3GridDiskDistances));
registry.register(Arc::new(h3::H3GridDistance));
registry.register(Arc::new(h3::H3GridPathCells));
}
}

728
src/common/function/src/scalars/geo/h3.rs
View File

@@ -20,18 +20,71 @@ use common_query::error::{self, InvalidFuncArgsSnafu, Result};
use common_query::prelude::{Signature, TypeSignature};
use datafusion::logical_expr::Volatility;
use datatypes::prelude::ConcreteDataType;
use datatypes::scalars::ScalarVectorBuilder;
use datatypes::value::Value;
use datatypes::scalars::{Scalar, ScalarVectorBuilder};
use datatypes::value::{ListValue, Value};
use datatypes::vectors::{
BooleanVectorBuilder, Float64VectorBuilder, MutableVector, StringVectorBuilder,
UInt64VectorBuilder, UInt8VectorBuilder, VectorRef,
BooleanVectorBuilder, Int32VectorBuilder, ListVectorBuilder, MutableVector,
StringVectorBuilder, UInt64VectorBuilder, UInt8VectorBuilder, VectorRef,
};
use derive_more::Display;
use h3o::{CellIndex, LatLng, Resolution};
use once_cell::sync::Lazy;
use snafu::{ensure, ResultExt};
use super::helpers::{ensure_columns_len, ensure_columns_n};
use crate::function::{Function, FunctionContext};
static CELL_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
vec![
ConcreteDataType::int64_datatype(),
ConcreteDataType::uint64_datatype(),
]
});
static COORDINATE_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
vec![
ConcreteDataType::float32_datatype(),
ConcreteDataType::float64_datatype(),
]
});
static RESOLUTION_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
vec![
ConcreteDataType::int8_datatype(),
ConcreteDataType::int16_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int64_datatype(),
ConcreteDataType::uint8_datatype(),
ConcreteDataType::uint16_datatype(),
ConcreteDataType::uint32_datatype(),
ConcreteDataType::uint64_datatype(),
]
});
static DISTANCE_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
vec![
ConcreteDataType::int8_datatype(),
ConcreteDataType::int16_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int64_datatype(),
ConcreteDataType::uint8_datatype(),
ConcreteDataType::uint16_datatype(),
ConcreteDataType::uint32_datatype(),
ConcreteDataType::uint64_datatype(),
]
});
static POSITION_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
vec![
ConcreteDataType::int8_datatype(),
ConcreteDataType::int16_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int64_datatype(),
ConcreteDataType::uint8_datatype(),
ConcreteDataType::uint16_datatype(),
ConcreteDataType::uint32_datatype(),
ConcreteDataType::uint64_datatype(),
]
});
/// Function that returns [h3] encoding cellid for a given geospatial coordinate.
///
/// [h3]: https://h3geo.org/
@@ -50,20 +103,8 @@ impl Function for H3LatLngToCell {
fn signature(&self) -> Signature {
let mut signatures = Vec::new();
for coord_type in &[
ConcreteDataType::float32_datatype(),
ConcreteDataType::float64_datatype(),
] {
for resolution_type in &[
ConcreteDataType::int8_datatype(),
ConcreteDataType::int16_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int64_datatype(),
ConcreteDataType::uint8_datatype(),
ConcreteDataType::uint16_datatype(),
ConcreteDataType::uint32_datatype(),
ConcreteDataType::uint64_datatype(),
] {
for coord_type in COORDINATE_TYPES.as_slice() {
for resolution_type in RESOLUTION_TYPES.as_slice() {
signatures.push(TypeSignature::Exact(vec![
// latitude
coord_type.clone(),
@@ -78,15 +119,7 @@ impl Function for H3LatLngToCell {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 3,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 3, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 3);
let lat_vec = &columns[0];
let lon_vec = &columns[1];
@@ -142,20 +175,8 @@ impl Function for H3LatLngToCellString {
fn signature(&self) -> Signature {
let mut signatures = Vec::new();
for coord_type in &[
ConcreteDataType::float32_datatype(),
ConcreteDataType::float64_datatype(),
] {
for resolution_type in &[
ConcreteDataType::int8_datatype(),
ConcreteDataType::int16_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int64_datatype(),
ConcreteDataType::uint8_datatype(),
ConcreteDataType::uint16_datatype(),
ConcreteDataType::uint32_datatype(),
ConcreteDataType::uint64_datatype(),
] {
for coord_type in COORDINATE_TYPES.as_slice() {
for resolution_type in RESOLUTION_TYPES.as_slice() {
signatures.push(TypeSignature::Exact(vec![
// latitude
coord_type.clone(),
@@ -170,15 +191,7 @@ impl Function for H3LatLngToCellString {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 3,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 3, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 3);
let lat_vec = &columns[0];
let lon_vec = &columns[1];
@@ -234,15 +247,7 @@ impl Function for H3CellToString {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 1,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 1, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 1);
let cell_vec = &columns[0];
let size = cell_vec.len();
@@ -280,15 +285,7 @@ impl Function for H3StringToCell {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 1,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 1, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 1);
let string_vec = &columns[0];
let size = string_vec.len();
@@ -319,18 +316,20 @@ impl Function for H3StringToCell {
}
}
/// Function that returns centroid latitude of given cell id
/// Function that returns centroid latitude and longitude of given cell id
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3CellCenterLat;
pub struct H3CellCenterLatLng;
impl Function for H3CellCenterLat {
impl Function for H3CellCenterLatLng {
fn name(&self) -> &str {
"h3_cell_center_lat"
"h3_cell_center_latlng"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::float64_datatype())
Ok(ConcreteDataType::list_datatype(
ConcreteDataType::float64_datatype(),
))
}
fn signature(&self) -> Signature {
@@ -338,69 +337,26 @@ impl Function for H3CellCenterLat {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 1,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 1, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 1);
let cell_vec = &columns[0];
let size = cell_vec.len();
let mut results = Float64VectorBuilder::with_capacity(size);
let mut results =
ListVectorBuilder::with_type_capacity(ConcreteDataType::float64_datatype(), size);
for i in 0..size {
let cell = cell_from_value(cell_vec.get(i))?;
let lat = cell.map(|cell| LatLng::from(cell).lat());
let latlng = cell.map(LatLng::from);
results.push(lat);
}
Ok(results.to_vector())
}
}
/// Function that returns centroid longitude of given cell id
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3CellCenterLng;
impl Function for H3CellCenterLng {
fn name(&self) -> &str {
"h3_cell_center_lng"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::float64_datatype())
}
fn signature(&self) -> Signature {
signature_of_cell()
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 1,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 1, provided : {}",
columns.len()
),
if let Some(latlng) = latlng {
let result = ListValue::new(
vec![latlng.lat().into(), latlng.lng().into()],
ConcreteDataType::float64_datatype(),
);
results.push(Some(result.as_scalar_ref()));
} else {
results.push(None);
}
);
let cell_vec = &columns[0];
let size = cell_vec.len();
let mut results = Float64VectorBuilder::with_capacity(size);
for i in 0..size {
let cell = cell_from_value(cell_vec.get(i))?;
let lat = cell.map(|cell| LatLng::from(cell).lng());
results.push(lat);
}
Ok(results.to_vector())
@@ -470,15 +426,7 @@ impl Function for H3CellBase {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 1,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 1, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 1);
let cell_vec = &columns[0];
let size = cell_vec.len();
@@ -514,15 +462,7 @@ impl Function for H3CellIsPentagon {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 1,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 1, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 1);
let cell_vec = &columns[0];
let size = cell_vec.len();
@@ -558,15 +498,7 @@ impl Function for H3CellCenterChild {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 2,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 2, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 2);
let cell_vec = &columns[0];
let res_vec = &columns[1];
@@ -606,15 +538,7 @@ impl Function for H3CellParent {
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 2,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 2, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 2);
let cell_vec = &columns[0];
let res_vec = &columns[1];
@@ -633,48 +557,323 @@ impl Function for H3CellParent {
}
}
/// Function that checks if two cells are neighbour
/// Function that returns children cell list
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3IsNeighbour;
pub struct H3CellToChildren;
impl Function for H3IsNeighbour {
impl Function for H3CellToChildren {
fn name(&self) -> &str {
"h3_is_neighbour"
"h3_cell_to_children"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::boolean_datatype())
Ok(ConcreteDataType::list_datatype(
ConcreteDataType::uint64_datatype(),
))
}
fn signature(&self) -> Signature {
signature_of_double_cell()
signature_of_cell_and_resolution()
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 2,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of the args is not correct, expect 2, provided : {}",
columns.len()
),
}
);
ensure_columns_n!(columns, 2);
let cell_vec = &columns[0];
let cell2_vec = &columns[1];
let res_vec = &columns[1];
let size = cell_vec.len();
let mut results = BooleanVectorBuilder::with_capacity(size);
let mut results =
ListVectorBuilder::with_type_capacity(ConcreteDataType::uint64_datatype(), size);
for i in 0..size {
let cell = cell_from_value(cell_vec.get(i))?;
let res = value_to_resolution(res_vec.get(i))?;
let result = cell.map(|cell| {
let children: Vec<Value> = cell
.children(res)
.map(|child| Value::from(u64::from(child)))
.collect();
ListValue::new(children, ConcreteDataType::uint64_datatype())
});
if let Some(list_value) = result {
results.push(Some(list_value.as_scalar_ref()));
} else {
results.push(None);
}
}
Ok(results.to_vector())
}
}
/// Function that returns children cell count
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3CellToChildrenSize;
impl Function for H3CellToChildrenSize {
fn name(&self) -> &str {
"h3_cell_to_children_size"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::uint64_datatype())
}
fn signature(&self) -> Signature {
signature_of_cell_and_resolution()
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure_columns_n!(columns, 2);
let cell_vec = &columns[0];
let res_vec = &columns[1];
let size = cell_vec.len();
let mut results = UInt64VectorBuilder::with_capacity(size);
for i in 0..size {
let cell = cell_from_value(cell_vec.get(i))?;
let res = value_to_resolution(res_vec.get(i))?;
let result = cell.map(|cell| cell.children_count(res));
results.push(result);
}
Ok(results.to_vector())
}
}
/// Function that returns the cell position if its parent at given resolution
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3CellToChildPos;
impl Function for H3CellToChildPos {
fn name(&self) -> &str {
"h3_cell_to_child_pos"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::uint64_datatype())
}
fn signature(&self) -> Signature {
signature_of_cell_and_resolution()
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure_columns_n!(columns, 2);
let cell_vec = &columns[0];
let res_vec = &columns[1];
let size = cell_vec.len();
let mut results = UInt64VectorBuilder::with_capacity(size);
for i in 0..size {
let cell = cell_from_value(cell_vec.get(i))?;
let res = value_to_resolution(res_vec.get(i))?;
let result = cell.and_then(|cell| cell.child_position(res));
results.push(result);
}
Ok(results.to_vector())
}
}
/// Function that returns the cell at given position of the parent at given resolution
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3ChildPosToCell;
impl Function for H3ChildPosToCell {
fn name(&self) -> &str {
"h3_child_pos_to_cell"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::uint64_datatype())
}
fn signature(&self) -> Signature {
let mut signatures =
Vec::with_capacity(POSITION_TYPES.len() * CELL_TYPES.len() * RESOLUTION_TYPES.len());
for position_type in POSITION_TYPES.as_slice() {
for cell_type in CELL_TYPES.as_slice() {
for resolution_type in RESOLUTION_TYPES.as_slice() {
signatures.push(TypeSignature::Exact(vec![
position_type.clone(),
cell_type.clone(),
resolution_type.clone(),
]));
}
}
}
Signature::one_of(signatures, Volatility::Stable)
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure_columns_n!(columns, 3);
let pos_vec = &columns[0];
let cell_vec = &columns[1];
let res_vec = &columns[2];
let size = cell_vec.len();
let mut results = UInt64VectorBuilder::with_capacity(size);
for i in 0..size {
let cell = cell_from_value(cell_vec.get(i))?;
let pos = value_to_position(pos_vec.get(i))?;
let res = value_to_resolution(res_vec.get(i))?;
let result = cell.and_then(|cell| cell.child_at(pos, res).map(u64::from));
results.push(result);
}
Ok(results.to_vector())
}
}
/// Function that returns cells with k distances of given cell
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3GridDisk;
impl Function for H3GridDisk {
fn name(&self) -> &str {
"h3_grid_disk"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::list_datatype(
ConcreteDataType::uint64_datatype(),
))
}
fn signature(&self) -> Signature {
signature_of_cell_and_distance()
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure_columns_n!(columns, 2);
let cell_vec = &columns[0];
let k_vec = &columns[1];
let size = cell_vec.len();
let mut results =
ListVectorBuilder::with_type_capacity(ConcreteDataType::uint64_datatype(), size);
for i in 0..size {
let cell = cell_from_value(cell_vec.get(i))?;
let k = value_to_distance(k_vec.get(i))?;
let result = cell.map(|cell| {
let children: Vec<Value> = cell
.grid_disk::<Vec<_>>(k)
.into_iter()
.map(|child| Value::from(u64::from(child)))
.collect();
ListValue::new(children, ConcreteDataType::uint64_datatype())
});
if let Some(list_value) = result {
results.push(Some(list_value.as_scalar_ref()));
} else {
results.push(None);
}
}
Ok(results.to_vector())
}
}
/// Function that returns all cells within k distances of given cell
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3GridDiskDistances;
impl Function for H3GridDiskDistances {
fn name(&self) -> &str {
"h3_grid_disk_distances"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::list_datatype(
ConcreteDataType::uint64_datatype(),
))
}
fn signature(&self) -> Signature {
signature_of_cell_and_distance()
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure_columns_n!(columns, 2);
let cell_vec = &columns[0];
let k_vec = &columns[1];
let size = cell_vec.len();
let mut results =
ListVectorBuilder::with_type_capacity(ConcreteDataType::uint64_datatype(), size);
for i in 0..size {
let cell = cell_from_value(cell_vec.get(i))?;
let k = value_to_distance(k_vec.get(i))?;
let result = cell.map(|cell| {
let children: Vec<Value> = cell
.grid_disk_distances::<Vec<_>>(k)
.into_iter()
.map(|(child, _distance)| Value::from(u64::from(child)))
.collect();
ListValue::new(children, ConcreteDataType::uint64_datatype())
});
if let Some(list_value) = result {
results.push(Some(list_value.as_scalar_ref()));
} else {
results.push(None);
}
}
Ok(results.to_vector())
}
}
/// Function that returns distance between two cells
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3GridDistance;
impl Function for H3GridDistance {
fn name(&self) -> &str {
"h3_grid_distance"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::int32_datatype())
}
fn signature(&self) -> Signature {
signature_of_double_cells()
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure_columns_n!(columns, 2);
let cell_this_vec = &columns[0];
let cell_that_vec = &columns[1];
let size = cell_this_vec.len();
let mut results = Int32VectorBuilder::with_capacity(size);
for i in 0..size {
let result = match (
cell_from_value(cell_vec.get(i))?,
cell_from_value(cell2_vec.get(i))?,
cell_from_value(cell_this_vec.get(i))?,
cell_from_value(cell_that_vec.get(i))?,
) {
(Some(cell_this), Some(cell_that)) => {
let is_neighbour = cell_this
.is_neighbor_with(cell_that)
let dist = cell_this
.grid_distance(cell_that)
.map_err(|e| {
BoxedError::new(PlainError::new(
format!("H3 error: {}", e),
@@ -682,7 +881,7 @@ impl Function for H3IsNeighbour {
))
})
.context(error::ExecuteSnafu)?;
Some(is_neighbour)
Some(dist)
}
_ => None,
};
@@ -694,6 +893,73 @@ impl Function for H3IsNeighbour {
}
}
/// Function that returns path cells between two cells
#[derive(Clone, Debug, Default, Display)]
#[display("{}", self.name())]
pub struct H3GridPathCells;
impl Function for H3GridPathCells {
fn name(&self) -> &str {
"h3_grid_path_cells"
}
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
Ok(ConcreteDataType::list_datatype(
ConcreteDataType::uint64_datatype(),
))
}
fn signature(&self) -> Signature {
signature_of_double_cells()
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure_columns_n!(columns, 2);
let cell_this_vec = &columns[0];
let cell_that_vec = &columns[1];
let size = cell_this_vec.len();
let mut results =
ListVectorBuilder::with_type_capacity(ConcreteDataType::uint64_datatype(), size);
for i in 0..size {
let result = match (
cell_from_value(cell_this_vec.get(i))?,
cell_from_value(cell_that_vec.get(i))?,
) {
(Some(cell_this), Some(cell_that)) => {
let cells = cell_this
.grid_path_cells(cell_that)
.and_then(|x| x.collect::<std::result::Result<Vec<CellIndex>, _>>())
.map_err(|e| {
BoxedError::new(PlainError::new(
format!("H3 error: {}", e),
StatusCode::EngineExecuteQuery,
))
})
.context(error::ExecuteSnafu)?;
Some(ListValue::new(
cells
.into_iter()
.map(|c| Value::from(u64::from(c)))
.collect(),
ConcreteDataType::uint64_datatype(),
))
}
_ => None,
};
if let Some(list_value) = result {
results.push(Some(list_value.as_scalar_ref()));
} else {
results.push(None);
}
}
Ok(results.to_vector())
}
}
fn value_to_resolution(v: Value) -> Result<Resolution> {
let r = match v {
Value::Int8(v) => v as u8,
@@ -716,26 +982,59 @@ fn value_to_resolution(v: Value) -> Result<Resolution> {
.context(error::ExecuteSnafu)
}
macro_rules! ensure_and_coerce {
($compare:expr, $coerce:expr) => {{
ensure!(
$compare,
InvalidFuncArgsSnafu {
err_msg: "Argument was outside of acceptable range "
}
);
Ok($coerce)
}};
}
fn value_to_position(v: Value) -> Result<u64> {
match v {
Value::Int8(v) => ensure_and_coerce!(v >= 0, v as u64),
Value::Int16(v) => ensure_and_coerce!(v >= 0, v as u64),
Value::Int32(v) => ensure_and_coerce!(v >= 0, v as u64),
Value::Int64(v) => ensure_and_coerce!(v >= 0, v as u64),
Value::UInt8(v) => Ok(v as u64),
Value::UInt16(v) => Ok(v as u64),
Value::UInt32(v) => Ok(v as u64),
Value::UInt64(v) => Ok(v),
_ => unreachable!(),
}
}
fn value_to_distance(v: Value) -> Result<u32> {
match v {
Value::Int8(v) => ensure_and_coerce!(v >= 0, v as u32),
Value::Int16(v) => ensure_and_coerce!(v >= 0, v as u32),
Value::Int32(v) => ensure_and_coerce!(v >= 0, v as u32),
Value::Int64(v) => ensure_and_coerce!(v >= 0, v as u32),
Value::UInt8(v) => Ok(v as u32),
Value::UInt16(v) => Ok(v as u32),
Value::UInt32(v) => Ok(v),
Value::UInt64(v) => Ok(v as u32),
_ => unreachable!(),
}
}
fn signature_of_cell() -> Signature {
let mut signatures = Vec::new();
for cell_type in &[
ConcreteDataType::uint64_datatype(),
ConcreteDataType::int64_datatype(),
] {
let mut signatures = Vec::with_capacity(CELL_TYPES.len());
for cell_type in CELL_TYPES.as_slice() {
signatures.push(TypeSignature::Exact(vec![cell_type.clone()]));
}
Signature::one_of(signatures, Volatility::Stable)
}
fn signature_of_double_cell() -> Signature {
let mut signatures = Vec::new();
let cell_types = &[
ConcreteDataType::uint64_datatype(),
ConcreteDataType::int64_datatype(),
];
for cell_type in cell_types {
for cell_type2 in cell_types {
fn signature_of_double_cells() -> Signature {
let mut signatures = Vec::with_capacity(CELL_TYPES.len() * CELL_TYPES.len());
for cell_type in CELL_TYPES.as_slice() {
for cell_type2 in CELL_TYPES.as_slice() {
signatures.push(TypeSignature::Exact(vec![
cell_type.clone(),
cell_type2.clone(),
@@ -747,21 +1046,9 @@ fn signature_of_double_cell() -> Signature {
}
fn signature_of_cell_and_resolution() -> Signature {
let mut signatures = Vec::new();
for cell_type in &[
ConcreteDataType::uint64_datatype(),
ConcreteDataType::int64_datatype(),
] {
for resolution_type in &[
ConcreteDataType::int8_datatype(),
ConcreteDataType::int16_datatype(),
ConcreteDataType::int32_datatype(),
ConcreteDataType::int64_datatype(),
ConcreteDataType::uint8_datatype(),
ConcreteDataType::uint16_datatype(),
ConcreteDataType::uint32_datatype(),
ConcreteDataType::uint64_datatype(),
] {
let mut signatures = Vec::with_capacity(CELL_TYPES.len() * RESOLUTION_TYPES.len());
for cell_type in CELL_TYPES.as_slice() {
for resolution_type in RESOLUTION_TYPES.as_slice() {
signatures.push(TypeSignature::Exact(vec![
cell_type.clone(),
resolution_type.clone(),
@@ -771,6 +1058,19 @@ fn signature_of_cell_and_resolution() -> Signature {
Signature::one_of(signatures, Volatility::Stable)
}
fn signature_of_cell_and_distance() -> Signature {
let mut signatures = Vec::with_capacity(CELL_TYPES.len() * DISTANCE_TYPES.len());
for cell_type in CELL_TYPES.as_slice() {
for distance_type in DISTANCE_TYPES.as_slice() {
signatures.push(TypeSignature::Exact(vec![
cell_type.clone(),
distance_type.clone(),
]));
}
}
Signature::one_of(signatures, Volatility::Stable)
}
fn cell_from_value(v: Value) -> Result<Option<CellIndex>> {
let cell = match v {
Value::Int64(v) => Some(

59
tests/cases/standalone/common/function/geo.result
View File

File diff suppressed because one or more lines are too long

22
tests/cases/standalone/common/function/geo.sql
View File

@@ -26,18 +26,34 @@ SELECT h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64), h3_latlng_to_cell_stri
SELECT h3_cell_to_string(h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64)) AS cell_str, h3_string_to_cell(h3_latlng_to_cell_string(37.76938, -122.3889, 8::UInt64)) AS cell_index;
SELECT h3_cell_center_lat(h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64)) AS cell_center_lat, h3_cell_center_lng(h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64)) AS cell_center_lng;
SELECT h3_cell_center_latlng(h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64)) AS cell_center;
SELECT
h3_cell_resolution(cell) AS resolution,
h3_cell_base(cell) AS base,
h3_cell_is_pentagon(cell) AS pentagon,
h3_cell_parent(cell, 6::UInt64) AS parent,
h3_cell_to_children(cell, 10::UInt64) AS children,
h3_cell_to_children_size(cell, 10) AS children_count,
h3_cell_to_child_pos(cell, 6) AS child_pos,
h3_child_pos_to_cell(25, cell, 11) AS child
FROM (SELECT h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64) AS cell);
SELECT h3_is_neighbour(cell1, cell2)
FROM (SELECT h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64) AS cell1, h3_latlng_to_cell(36.76938, -122.3889, 8::UInt64) AS cell2);
SELECT
h3_grid_disk(cell, 0) AS current_cell,
h3_grid_disk(cell, 3) AS grids,
h3_grid_disk_distances(cell, 3) AS all_grids,
FROM (SELECT h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64) AS cell);
SELECT
h3_grid_distance(cell1, cell2) AS distance,
h3_grid_path_cells(cell1, cell2) AS path_cells,
FROM
(
SELECT
h3_latlng_to_cell(37.76938, -122.3889, 8::UInt64) AS cell1,
h3_latlng_to_cell(39.634, -104.999, 8::UInt64) AS cell2
);
SELECT geohash(37.76938, -122.3889, 9);