Skip to content

Commit

Permalink
Inspection functions all support VARCHAR input (#110)
Browse files Browse the repository at this point in the history 
* inspection functions support VARCHAR

* h3_get_icosahedron_faces supports VARCHAR
  • Loading branch information
@isaacbrodsky
isaacbrodsky authored Jun 21, 2024
1 parent c394c1c commit 3abe7bf
Show file tree
Hide file tree
Showing 3 changed files with 147 additions and 10 deletions.
10 changes: 5 additions & 5 deletions README.md
View file
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
View file
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
View file
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

0 comments on commit 3abe7bf

Please sign in to comment.