add s2 util test

vesoft-inc · Oct 10, 2021 · 6072b41 · 6072b41
1 parent ee127f5
commit 6072b41
Show file tree

Hide file tree

Showing 5 changed files with 123 additions and 7 deletions.
diff --git a/src/common/function/FunctionManager.cpp b/src/common/function/FunctionManager.cpp
@@ -2529,7 +2529,7 @@ FunctionManager::FunctionManager() {
         }
       }
       // TODO(jie) Should return uint64_t Value
-      uint64_t cellId = S2CellIdFromPoint(args[0].get().getGeography(), level);
+      uint64_t cellId = s2CellIdFromPoint(args[0].get().getGeography(), level);
       const char *tmp = reinterpret_cast<const char *>(&cellId);
       int64_t cellId2 = *reinterpret_cast<const int64_t *>(tmp);
       return cellId2;
@@ -2587,7 +2587,7 @@ FunctionManager::FunctionManager() {
           return Value::kNullBadData;
         }
       }
-      std::vector<uint64_t> cellIds = S2CoveringCellIds(
+      std::vector<uint64_t> cellIds = s2CoveringCellIds(
           args[0].get().getGeography(), minLevel, maxLevel, maxCells, bufferInMeters);
       // TODO(jie) Should return uint64_t List
       std::vector<Value> vals;

diff --git a/src/common/geo/function/S2Util.cpp b/src/common/geo/function/S2Util.cpp
@@ -18,7 +18,7 @@
 
 namespace nebula {
 
-uint64_t S2CellIdFromPoint(const Geography& a, int level) {
+uint64_t s2CellIdFromPoint(const Geography& a, int level) {
   auto aRegion = a.asS2();
   if (!aRegion) {
     return -1;
@@ -46,7 +46,7 @@ uint64_t S2CellIdFromPoint(const Geography& a, int level) {
   return false;
 }
 
-std::vector<uint64_t> S2CoveringCellIds(
+std::vector<uint64_t> s2CoveringCellIds(
     const Geography& a, int minLevel, int maxLevel, int maxCells, double bufferInMeters) {
   auto aRegion = a.asS2();
   if (!aRegion) {

diff --git a/src/common/geo/function/S2Util.h b/src/common/geo/function/S2Util.h
@@ -12,10 +12,13 @@
 
 namespace nebula {
 
-uint64_t S2CellIdFromPoint(const Geography& a, int level);
+uint64_t s2CellIdFromPoint(const Geography& a, int level = 30);
 
-std::vector<uint64_t> S2CoveringCellIds(
-    const Geography& a, int minLevel, int maxLevel, int maxCells, double bufferInMeters);
+std::vector<uint64_t> s2CoveringCellIds(const Geography& a,
+                                        int minLevel = 0,
+                                        int maxLevel = 30,
+                                        int maxCells = 8,
+                                        double bufferInMeters = 0.0);
 
 std::vector<uint64_t> coveringCellIds(const S2Region& r, const S2RegionCoverer::Options& opts);
 

diff --git a/src/common/geo/function/test/CMakeLists.txt b/src/common/geo/function/test/CMakeLists.txt
@@ -92,3 +92,21 @@ nebula_add_test(
         gtest
         gtest_main
 )
+
+nebula_add_test(
+    NAME
+        s2util_test
+    SOURCES
+        S2UtilTest.cpp
+    OBJECTS
+        $<TARGET_OBJECTS:function_manager_obj>
+        $<TARGET_OBJECTS:wkt_wkb_io_obj>
+        $<TARGET_OBJECTS:base_obj>
+        $<TARGET_OBJECTS:datatypes_obj>
+        $<TARGET_OBJECTS:time_utils_obj>
+        $<TARGET_OBJECTS:time_obj>
+        $<TARGET_OBJECTS:fs_obj>
+    LIBRARIES
+        gtest
+        gtest_main
+)
diff --git a/src/common/geo/function/test/S2UtilTest.cpp b/src/common/geo/function/test/S2UtilTest.cpp
@@ -3,3 +3,98 @@
  * This source code is licensed under Apache 2.0 License,
  * attached with Common Clause Condition 1.0, found in the LICENSES directory.
  */
+
+#include <gtest/gtest.h>
+
+#include "common/base/Base.h"
+#include "common/geo/function/S2Util.h"
+
+namespace nebula {
+
+TEST(s2CellIdFromPoint, point) {
+  {
+    auto point = Geography::fromWKT("POINT(1.0 1.0)").value();
+    uint64_t i = s2CellIdFromPoint(point);
+    EXPECT_EQ(1153277837650709461, i);
+  }
+  {
+    auto point = Geography::fromWKT("POINT(179.0 89.9)").value();
+    uint64_t i = s2CellIdFromPoint(point);
+    EXPECT_EQ(6533220958669205147, i);
+  }
+  {
+    auto point = Geography::fromWKT("POINT(-45.4 28.7652)").value();
+    uint64_t i = s2CellIdFromPoint(point);
+    int64_t expectInt64 = -8444974090143026723;
+    const char* c = reinterpret_cast<const char*>(&expectInt64);
+    uint64_t expect = *reinterpret_cast<const uint64_t*>(c);
+    EXPECT_EQ(expect, i);
+  }
+  {
+    auto point = Geography::fromWKT("POINT(0.0 0.0)").value();
+    uint64_t i = s2CellIdFromPoint(point);
+    EXPECT_EQ(1152921504606846977, i);
+  }
+}
+
+// s2CellIdFromPoint() just supports point
+TEST(s2CellIdFromPoint, lineString) {
+  {
+    auto line = Geography::fromWKT("LINESTRING(1.0 1.0, 2.0 2.0)").value();
+    uint64_t i = s2CellIdFromPoint(line);
+    EXPECT_EQ(-1, i);
+  }
+}
+
+// s2CellIdFromPoint() just supports point
+TEST(s2CellIdFromPoint, polygon) {
+  {
+    auto polygon = Geography::fromWKT("POLYGON((1.0 1.0, 2.0 2.0, 0.0 2.0, 1.0 1.0))").value();
+    uint64_t i = s2CellIdFromPoint(polygon);
+    EXPECT_EQ(-1, i);
+  }
+}
+
+TEST(coveringCellIds, point) {
+  {
+    auto point = Geography::fromWKT("POINT(1.0 1.0)").value();
+    std::vector<uint64_t> cellIds = s2CoveringCellIds(point);
+    EXPECT_EQ(std::vector<uint64_t>{1153277837650709461}, cellIds);
+  }
+  {
+    auto point = Geography::fromWKT("POINT(179.0 89.9)").value();
+    std::vector<uint64_t> cellIds = s2CoveringCellIds(point);
+    EXPECT_EQ(std::vector<uint64_t>{6533220958669205147}, cellIds);
+  }
+  {
+    auto point = Geography::fromWKT("POINT(-45.4 28.7652)").value();
+    std::vector<uint64_t> cellIds = s2CoveringCellIds(point);
+    int64_t expectInt64 = -8444974090143026723;
+    const char* c = reinterpret_cast<const char*>(&expectInt64);
+    uint64_t expect = *reinterpret_cast<const uint64_t*>(c);
+    EXPECT_EQ(std::vector<uint64_t>{expect}, cellIds);
+  }
+  {
+    auto point = Geography::fromWKT("POINT(0.0 0.0)").value();
+    std::vector<uint64_t> cellIds = s2CoveringCellIds(point);
+    EXPECT_EQ(std::vector<uint64_t>{1152921504606846977}, cellIds);
+  }
+}
+
+TEST(coveringCellIds, linestring) {
+  {
+    auto line = Geography::fromWKT("LINESTRING(1.0 1.0, 2.0 2.0)").value();
+    std::vector<uint64_t> cellIds = s2CoveringCellIds(line);
+    EXPECT_EQ(std::vector<uint64_t>{1153277837650709461}, cellIds);
+  }
+}
+
+}  // namespace nebula
+
+int main(int argc, char** argv) {
+  testing::InitGoogleTest(&argc, argv);
+  folly::init(&argc, &argv, true);
+  google::SetStderrLogging(google::INFO);
+
+  return RUN_ALL_TESTS();
+}