Skip to content

Commit

Permalink
feat: add some s2 geo functions (#4823)
Browse files Browse the repository at this point in the history
* feat: add first batch of s2 functions

* refactor: update reusable code from main

* test: add sqlness tests for s2

* feat: add tostring function for s2

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

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>

* Apply suggestions from code review

* one more change

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
Co-authored-by: Ruihang Xia <waynestxia@gmail.com>
Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>
  • Loading branch information
3 people authored Oct 15, 2024
1 parent b5233e5 commit bb8b54b
Show file tree
Hide file tree
Showing 9 changed files with 383 additions and 40 deletions.
52 changes: 48 additions & 4 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion src/common/function/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ workspace = true

[features]
default = ["geo"]
geo = ["geohash", "h3o"]
geo = ["geohash", "h3o", "s2"]

[dependencies]
api.workspace = true
Expand All @@ -35,6 +35,7 @@ num = "0.4"
num-traits = "0.2"
once_cell.workspace = true
paste = "1.0"
s2 = { version = "0.0.12", optional = true }
serde.workspace = true
serde_json.workspace = true
session.workspace = true
Expand Down
13 changes: 9 additions & 4 deletions src/common/function/src/scalars/geo.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,7 @@ pub(crate) mod encoding;
mod geohash;
mod h3;
mod helpers;

use geohash::{GeohashFunction, GeohashNeighboursFunction};
mod s2;

use crate::function_registry::FunctionRegistry;

Expand All @@ -27,8 +26,8 @@ pub(crate) struct GeoFunctions;
impl GeoFunctions {
pub fn register(registry: &FunctionRegistry) {
// geohash
registry.register(Arc::new(GeohashFunction));
registry.register(Arc::new(GeohashNeighboursFunction));
registry.register(Arc::new(geohash::GeohashFunction));
registry.register(Arc::new(geohash::GeohashNeighboursFunction));

// h3 index
registry.register(Arc::new(h3::H3LatLngToCell));
Expand All @@ -55,5 +54,11 @@ impl GeoFunctions {
registry.register(Arc::new(h3::H3GridDiskDistances));
registry.register(Arc::new(h3::H3GridDistance));
registry.register(Arc::new(h3::H3GridPathCells));

// s2
registry.register(Arc::new(s2::S2LatLngToCell));
registry.register(Arc::new(s2::S2CellLevel));
registry.register(Arc::new(s2::S2CellToToken));
registry.register(Arc::new(s2::S2CellParent));
}
}
2 changes: 1 addition & 1 deletion src/common/function/src/scalars/geo/encoding.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use std::sync::Arc;
use common_error::ext::{BoxedError, PlainError};
use common_error::status_code::StatusCode;
use common_macro::{as_aggr_func_creator, AggrFuncTypeStore};
use common_query::error::{self, InvalidFuncArgsSnafu, InvalidInputStateSnafu, Result};
use common_query::error::{self, InvalidInputStateSnafu, Result};
use common_query::logical_plan::accumulator::AggrFuncTypeStore;
use common_query::logical_plan::{Accumulator, AggregateFunctionCreator};
use common_query::prelude::AccumulatorCreatorFunction;
Expand Down
28 changes: 4 additions & 24 deletions src/common/function/src/scalars/geo/h3.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use std::str::FromStr;

use common_error::ext::{BoxedError, PlainError};
use common_error::status_code::StatusCode;
use common_query::error::{self, InvalidFuncArgsSnafu, Result};
use common_query::error::{self, Result};
use common_query::prelude::{Signature, TypeSignature};
use datafusion::logical_expr::Volatility;
use datatypes::prelude::ConcreteDataType;
Expand All @@ -29,9 +29,9 @@ use datatypes::vectors::{
use derive_more::Display;
use h3o::{CellIndex, LatLng, Resolution};
use once_cell::sync::Lazy;
use snafu::{ensure, ResultExt};
use snafu::ResultExt;

use super::helpers::{ensure_columns_len, ensure_columns_n};
use super::helpers::{ensure_and_coerce, ensure_columns_len, ensure_columns_n};
use crate::function::{Function, FunctionContext};

static CELL_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
Expand Down Expand Up @@ -382,15 +382,7 @@ impl Function for H3CellResolution {
}

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();
Expand Down Expand Up @@ -982,18 +974,6 @@ 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),
Expand Down
26 changes: 20 additions & 6 deletions src/common/function/src/scalars/geo/helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,23 +14,23 @@

macro_rules! ensure_columns_len {
($columns:ident) => {
ensure!(
snafu::ensure!(
$columns.windows(2).all(|c| c[0].len() == c[1].len()),
InvalidFuncArgsSnafu {
common_query::error::InvalidFuncArgsSnafu {
err_msg: "The length of input columns are in different size"
}
)
};
($column_a:ident, $column_b:ident, $($column_n:ident),*) => {
ensure!(
snafu::ensure!(
{
let mut result = $column_a.len() == $column_b.len();
$(
result = result && ($column_a.len() == $column_n.len());
)*
result
}
InvalidFuncArgsSnafu {
common_query::error::InvalidFuncArgsSnafu {
err_msg: "The length of input columns are in different size"
}
)
Expand All @@ -41,9 +41,9 @@ pub(super) use ensure_columns_len;

macro_rules! ensure_columns_n {
($columns:ident, $n:literal) => {
ensure!(
snafu::ensure!(
$columns.len() == $n,
InvalidFuncArgsSnafu {
common_query::error::InvalidFuncArgsSnafu {
err_msg: format!(
"The length of arguments is not correct, expect {}, provided : {}",
stringify!($n),
Expand All @@ -59,3 +59,17 @@ macro_rules! ensure_columns_n {
}

pub(super) use ensure_columns_n;

macro_rules! ensure_and_coerce {
($compare:expr, $coerce:expr) => {{
snafu::ensure!(
$compare,
common_query::error::InvalidFuncArgsSnafu {
err_msg: "Argument was outside of acceptable range "
}
);
Ok($coerce)
}};
}

pub(super) use ensure_and_coerce;
Loading

0 comments on commit bb8b54b

Please sign in to comment.