Skip to content

Commit

Permalink
feat: ensure column number and length
Browse files Browse the repository at this point in the history
  • Loading branch information
@sunng87
sunng87 committed Sep 30, 2024
1 parent acf9f5b commit 33aa437
Showing 1 changed file with 61 additions and 153 deletions.
214 changes: 61 additions & 153 deletions src/common/function/src/scalars/geo/h3.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,50 @@ const POSITION_TYPES: Lazy<Vec<ConcreteDataType>> = Lazy::new(|| {
]
});

macro_rules! ensure_columns_len {
($columns:ident) => {
ensure!(
$columns.windows(2).all(|c| c[0].len() == c[1].len()),
InvalidFuncArgsSnafu {
err_msg: "The length of input columns are in different size"
}
)
};
($column_a:ident, $column_b:ident, $($column_n:ident),*) => {
ensure!(
{
let mut result = $column_a.len() == $column_b.len();
$(
result = result && ($column_a.len() == $column_n.len());
)*
result
}
InvalidFuncArgsSnafu {
err_msg: "The length of input columns are in different size"
}
)
};
}

macro_rules! ensure_columns_n {
($columns:ident, $n:literal) => {
ensure!(
$columns.len() == $n,
InvalidFuncArgsSnafu {
err_msg: format!(
"The length of arguments is not correct, expect {}, provided : {}",
stringify!($n),
$columns.len()
),
}
);

if $n > 1 {
ensure_columns_len!($columns);
}
};
}

/// Function that returns [h3] encoding cellid for a given geospatial coordinate.
///
/// [h3]: https://h3geo.org/
Expand Down Expand Up @@ -118,15 +162,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];
Expand Down Expand Up @@ -198,15 +234,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];
Expand Down Expand Up @@ -262,15 +290,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();
Expand Down Expand Up @@ -308,15 +328,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();
Expand Down Expand Up @@ -368,15 +380,7 @@ impl Function for H3CellCenterLatLng {
}

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 @@ -465,15 +469,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();
Expand Down Expand Up @@ -509,15 +505,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();
Expand Down Expand Up @@ -553,15 +541,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];
Expand Down Expand Up @@ -601,15 +581,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];
Expand Down Expand Up @@ -649,15 +621,7 @@ impl Function for H3CellToChildren {
}

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];
Expand Down Expand Up @@ -706,15 +670,7 @@ impl Function for H3CellToChildrenSize {
}

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];
Expand Down Expand Up @@ -751,15 +707,7 @@ impl Function for H3CellToChildPos {
}

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];
Expand Down Expand Up @@ -808,15 +756,7 @@ impl Function for H3ChildPosToCell {
}

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 pos_vec = &columns[0];
let cell_vec = &columns[1];
Expand Down Expand Up @@ -857,15 +797,7 @@ impl Function for H3GridDisk {
}

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 k_vec = &columns[1];
Expand Down Expand Up @@ -918,15 +850,7 @@ impl Function for H3GridDiskDistances {
}

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 k_vec = &columns[1];
Expand Down Expand Up @@ -976,15 +900,7 @@ impl Function for H3GridDistance {
}

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_this_vec = &columns[0];
let cell_that_vec = &columns[1];
Expand Down Expand Up @@ -1040,15 +956,7 @@ impl Function for H3GridPathCells {
}

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_this_vec = &columns[0];
let cell_that_vec = &columns[1];
Expand Down

0 comments on commit 33aa437

Please sign in to comment.