Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Inspection functions all support VARCHAR input #110

Merged
merged 2 commits into from
Jun 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 +50,14 @@ one to use. The unsigned and signed APIs are identical. Many functions also supp
| `h3_cell_to_lng` | [v](#fv) | Convert cell ID to longitude
| `h3_cell_to_latlng` | [v](#fv) | Convert cell ID to latitude/longitude
| `h3_cell_to_boundary_wkt` | [v](#fv) | Convert cell ID to cell boundary
| `h3_get_resolution` | [i](#fi) | Get resolution number of cell ID
| `h3_get_base_cell_number` | [i](#fi) | Get base cell number of cell ID
| `h3_get_resolution` | [v](#fv) | Get resolution number of cell ID
| `h3_get_base_cell_number` | [v](#fv) | Get base cell number of cell ID
| `h3_string_to_h3` | [u](#fu) | Convert VARCHAR cell ID to UBIGINT
| `h3_h3_to_string` | [i](#fi) | Convert BIGINT or UBIGINT cell ID to VARCHAR
| `h3_is_valid_cell` | [v](#fv) | True if this is a valid cell ID
| `h3_is_res_class_iii` | [i](#fi) | True if the cell's resolution is class III
| `h3_is_pentagon` | [i](#fi) | True if the cell is a pentagon
| `h3_get_icosahedron_faces` | [i](#fi) | List of icosahedron face IDs the cell is on
| `h3_is_res_class_iii` | [v](#fv) | True if the cell's resolution is class III
| `h3_is_pentagon` | [v](#fv) | True if the cell is a pentagon
| `h3_get_icosahedron_faces` | [v](#fv) | List of icosahedron face IDs the cell is on
| `h3_cell_to_parent` | [i](#fi) | Get coarser cell for a cell
| `h3_cell_to_children` | [i](#fi) | Get finer cells for a cell
| `h3_cell_to_center_child` | [i](#fi) | Get the center finer cell for a cell
Expand Down
132 changes: 127 additions & 5 deletions src/h3_inspection.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,24 @@ static void GetResolutionFunction(DataChunk &args, ExpressionState &state,
[&](H3Index cell) { return getResolution(cell); });
}

static void GetResolutionVarcharFunction(DataChunk &args,
ExpressionState &state,
Vector &result) {
auto &inputs = args.data[0];
UnaryExecutor::ExecuteWithNulls<string_t, int>(
inputs, result, args.size(),
[&](string_t cellAddress, ValidityMask &mask, idx_t idx) {
H3Index cell;
H3Error err0 = stringToH3(cellAddress.GetString().c_str(), &cell);
if (err0) {
mask.SetInvalid(idx);
return 0;
} else {
return getResolution(cell);
}
});
}

static void GetBaseCellNumberFunction(DataChunk &args, ExpressionState &state,
Vector &result) {
auto &inputs = args.data[0];
Expand All @@ -19,6 +37,24 @@ static void GetBaseCellNumberFunction(DataChunk &args, ExpressionState &state,
[&](H3Index cell) { return getBaseCellNumber(cell); });
}

static void GetBaseCellNumberVarcharFunction(DataChunk &args,
ExpressionState &state,
Vector &result) {
auto &inputs = args.data[0];
UnaryExecutor::ExecuteWithNulls<string_t, int>(
inputs, result, args.size(),
[&](string_t cellAddress, ValidityMask &mask, idx_t idx) {
H3Index cell;
H3Error err0 = stringToH3(cellAddress.GetString().c_str(), &cell);
if (err0) {
mask.SetInvalid(idx);
return 0;
} else {
return getBaseCellNumber(cell);
}
});
}

static void StringToH3Function(DataChunk &args, ExpressionState &state,
Vector &result) {
auto &inputs = args.data[0];
Expand Down Expand Up @@ -81,6 +117,24 @@ static void IsResClassIIIFunction(DataChunk &args, ExpressionState &state,
[&](uint64_t cell) { return bool(isResClassIII(cell)); });
}

static void IsResClassIIIVarcharFunction(DataChunk &args,
ExpressionState &state,
Vector &result) {
auto &inputs = args.data[0];
UnaryExecutor::ExecuteWithNulls<string_t, bool>(
inputs, result, args.size(),
[&](string_t cellAddress, ValidityMask &mask, idx_t idx) {
H3Index cell;
H3Error err0 = stringToH3(cellAddress.GetString().c_str(), &cell);
if (err0) {
mask.SetInvalid(idx);
return false;
} else {
return bool(isResClassIII(cell));
}
});
}

static void IsPentagonFunction(DataChunk &args, ExpressionState &state,
Vector &result) {
auto &inputs = args.data[0];
Expand All @@ -89,6 +143,23 @@ static void IsPentagonFunction(DataChunk &args, ExpressionState &state,
[&](uint64_t cell) { return bool(isPentagon(cell)); });
}

static void IsPentagonVarcharFunction(DataChunk &args, ExpressionState &state,
Vector &result) {
auto &inputs = args.data[0];
UnaryExecutor::ExecuteWithNulls<string_t, bool>(
inputs, result, args.size(),
[&](string_t cellAddress, ValidityMask &mask, idx_t idx) {
H3Index cell;
H3Error err0 = stringToH3(cellAddress.GetString().c_str(), &cell);
if (err0) {
mask.SetInvalid(idx);
return false;
} else {
return bool(isPentagon(cell));
}
});
}

static void GetIcosahedronFacesFunction(DataChunk &args, ExpressionState &state,
Vector &result) {
auto result_data = FlatVector::GetData<list_entry_t>(result);
Expand Down Expand Up @@ -123,23 +194,70 @@ static void GetIcosahedronFacesFunction(DataChunk &args, ExpressionState &state,
result.Verify(args.size());
}

static void GetIcosahedronFacesVarcharFunction(DataChunk &args,
ExpressionState &state,
Vector &result) {
UnifiedVectorFormat vdata;
args.data[0].ToUnifiedFormat(args.size(), vdata);

auto ldata = UnifiedVectorFormat::GetData<string_t>(vdata);

result.SetVectorType(VectorType::FLAT_VECTOR);
auto result_data = FlatVector::GetData<list_entry_t>(result);
for (idx_t i = 0; i < args.size(); i++) {
result_data[i].offset = ListVector::GetListSize(result);

int faceCount;
int64_t actual = 0;
string_t cellAddress = ldata[i];
H3Index cell;
H3Error err0 = stringToH3(cellAddress.GetString().c_str(), &cell);
if (err0) {
result.SetValue(i, Value(LogicalType::SQLNULL));
} else {
H3Error err1 = maxFaceCount(cell, &faceCount);
if (err1) {
result.SetValue(i, Value(LogicalType::SQLNULL));
} else {
std::vector<int> out(faceCount);
H3Error err2 = getIcosahedronFaces(cell, out.data());
if (err2) {
result.SetValue(i, Value(LogicalType::SQLNULL));
} else {
for (auto val : out) {
if (val != -1) {
ListVector::PushBack(result, Value::INTEGER(val));
actual++;
}
}
}
}
}

result_data[i].length = actual;
}
result.Verify(args.size());
}

CreateScalarFunctionInfo H3Functions::GetGetResolutionFunction() {
ScalarFunctionSet funcs("h3_get_resolution");
// TODO: VARCHAR variant of this function
funcs.AddFunction(ScalarFunction({LogicalType::UBIGINT}, LogicalType::INTEGER,
GetResolutionFunction));
funcs.AddFunction(ScalarFunction({LogicalType::BIGINT}, LogicalType::INTEGER,
GetResolutionFunction));
funcs.AddFunction(ScalarFunction({LogicalType::VARCHAR}, LogicalType::INTEGER,
GetResolutionVarcharFunction));
return CreateScalarFunctionInfo(funcs);
}

CreateScalarFunctionInfo H3Functions::GetGetBaseCellNumberFunction() {
ScalarFunctionSet funcs("h3_get_base_cell_number");
// TODO: VARCHAR variant of this function
funcs.AddFunction(ScalarFunction({LogicalType::UBIGINT}, LogicalType::INTEGER,
GetBaseCellNumberFunction));
funcs.AddFunction(ScalarFunction({LogicalType::BIGINT}, LogicalType::INTEGER,
GetBaseCellNumberFunction));
funcs.AddFunction(ScalarFunction({LogicalType::VARCHAR}, LogicalType::INTEGER,
GetBaseCellNumberFunction));
return CreateScalarFunctionInfo(funcs);
}

Expand Down Expand Up @@ -171,33 +289,37 @@ CreateScalarFunctionInfo H3Functions::GetIsValidCellFunctions() {

CreateScalarFunctionInfo H3Functions::GetIsResClassIIIFunction() {
ScalarFunctionSet funcs("h3_is_res_class_iii");
// TODO: VARCHAR variant of this function
funcs.AddFunction(ScalarFunction({LogicalType::UBIGINT}, LogicalType::BOOLEAN,
IsResClassIIIFunction));
funcs.AddFunction(ScalarFunction({LogicalType::BIGINT}, LogicalType::BOOLEAN,
IsResClassIIIFunction));
funcs.AddFunction(ScalarFunction({LogicalType::VARCHAR}, LogicalType::BOOLEAN,
IsResClassIIIFunction));
return CreateScalarFunctionInfo(funcs);
}

CreateScalarFunctionInfo H3Functions::GetIsPentagonFunction() {
ScalarFunctionSet funcs("h3_is_pentagon");
// TODO: VARCHAR variant of this function
funcs.AddFunction(ScalarFunction({LogicalType::UBIGINT}, LogicalType::BOOLEAN,
IsPentagonFunction));
funcs.AddFunction(ScalarFunction({LogicalType::BIGINT}, LogicalType::BOOLEAN,
IsPentagonFunction));
funcs.AddFunction(ScalarFunction({LogicalType::VARCHAR}, LogicalType::BOOLEAN,
IsPentagonFunction));
return CreateScalarFunctionInfo(funcs);
}

CreateScalarFunctionInfo H3Functions::GetGetIcosahedronFacesFunction() {
ScalarFunctionSet funcs("h3_get_icosahedron_faces");
// TODO: VARCHAR variant of this function
funcs.AddFunction(ScalarFunction({LogicalType::UBIGINT},
LogicalType::LIST(LogicalType::INTEGER),
GetIcosahedronFacesFunction));
funcs.AddFunction(ScalarFunction({LogicalType::BIGINT},
LogicalType::LIST(LogicalType::INTEGER),
GetIcosahedronFacesFunction));
funcs.AddFunction(ScalarFunction({LogicalType::VARCHAR},
LogicalType::LIST(LogicalType::INTEGER),
GetIcosahedronFacesVarcharFunction));
return CreateScalarFunctionInfo(funcs);
}

Expand Down
15 changes: 15 additions & 0 deletions test/sql/h3/h3_functions_inspection.test
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,11 @@ SELECT h3_get_resolution(cast(586265647244115967 as bigint));
----
2

query I
SELECT h3_get_resolution('822d57fffffffff');
----
2

query I
SELECT h3_is_valid_cell(cast(586265647244115967 as ubigint));
----
Expand Down Expand Up @@ -65,6 +70,11 @@ select h3_get_icosahedron_faces(599686042433355775::bigint), h3_get_icosahedron_
----
[7] [1, 6, 11, 7, 2]

query II
select h3_get_icosahedron_faces('85283473fffffff'), h3_get_icosahedron_faces('801dfffffffffff')
----
[7] [1, 6, 11, 7, 2]

query I
select h3_get_icosahedron_faces(18446744073709551615::ubigint)
----
Expand All @@ -74,3 +84,8 @@ query I
select h3_get_icosahedron_faces(9223372036854775807::bigint)
----
NULL

query I
select h3_get_icosahedron_faces('7fffffffffffffff')
----
NULL
Loading