h3_get_icosahedron_faces supports VARCHAR

isaacbrodsky · Jun 21, 2024 · ad531d8 · ad531d8
1 parent cb77ed2
commit ad531d8
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -57,7 +57,7 @@ one to use. The unsigned and signed APIs are identical. Many functions also supp
 | `h3_is_valid_cell` | [v](#fv) | True if this is a valid cell ID
 | `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` | [i](#fi) | List of icosahedron face IDs the cell is on
+| `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

diff --git a/src/h3_inspection.cpp b/src/h3_inspection.cpp
@@ -194,6 +194,51 @@ 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");
   funcs.AddFunction(ScalarFunction({LogicalType::UBIGINT}, LogicalType::INTEGER,
@@ -266,13 +311,15 @@ CreateScalarFunctionInfo H3Functions::GetIsPentagonFunction() {
 
 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);
 }
 

diff --git a/test/sql/h3/h3_functions_inspection.test b/test/sql/h3/h3_functions_inspection.test
@@ -70,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)
 ----
@@ -79,3 +84,8 @@ query I
 select h3_get_icosahedron_faces(9223372036854775807::bigint)
 ----
 NULL
+
+query I
+select h3_get_icosahedron_faces('7fffffffffffffff')
+----
+NULL