Geo spatial: 3. geography schema, data, index and optimization (#3043)

* Geo spatial: 3. geography data and index
vesoft-inc · Oct 20, 2021 · 912c5fd · 912c5fd
1 parent 324b839
commit 912c5fd
Show file tree

Hide file tree

Showing 75 changed files with 2,289 additions and 432 deletions.
diff --git a/src/clients/meta/MetaClient.cpp b/src/clients/meta/MetaClient.cpp
@@ -346,6 +346,8 @@ bool MetaClient::loadSchemas(GraphSpaceID spaceId,
     bool hasDef = col.default_value_ref().has_value();
     auto& colType = col.get_type();
     size_t len = colType.type_length_ref().has_value() ? *colType.get_type_length() : 0;
+    cpp2::GeoShape geoShape =
+        colType.geo_shape_ref().has_value() ? *colType.get_geo_shape() : cpp2::GeoShape::ANY;
     bool nullable = col.nullable_ref().has_value() ? *col.get_nullable() : false;
     Expression* defaultValueExpr = nullptr;
     if (hasDef) {
@@ -359,8 +361,12 @@ bool MetaClient::loadSchemas(GraphSpaceID spaceId,
       }
     }
 
-    schema->addField(
-        col.get_name(), colType.get_type(), len, nullable, hasDef ? defaultValueExpr : nullptr);
+    schema->addField(col.get_name(),
+                     colType.get_type(),
+                     len,
+                     nullable,
+                     hasDef ? defaultValueExpr : nullptr,
+                     geoShape);
   };
 
   for (auto& tagIt : tagItemVec) {

diff --git a/src/codec/RowReaderV2.cpp b/src/codec/RowReaderV2.cpp
@@ -176,8 +176,22 @@ Value RowReaderV2::getValueByIndex(const int64_t index) const noexcept {
       return dt;
     }
     case meta::cpp2::PropertyType::GEOGRAPHY: {
-      // TODO(jie)
-      return Geography();
+      int32_t strOffset;
+      int32_t strLen;
+      memcpy(reinterpret_cast<void*>(&strOffset), &data_[offset], sizeof(int32_t));
+      memcpy(reinterpret_cast<void*>(&strLen), &data_[offset + sizeof(int32_t)], sizeof(int32_t));
+      if (static_cast<size_t>(strOffset) == data_.size() && strLen == 0) {
+        return Value::kEmpty;  // Is it ok to return Value::kEmpty?
+      }
+      CHECK_LT(strOffset, data_.size());
+      auto wkb = std::string(&data_[strOffset], strLen);
+      // Parse a geography from the wkb, normalize it and then verify its validity.
+      auto geogRet = Geography::fromWKB(wkb, true, true);
+      if (!geogRet.ok()) {
+        LOG(ERROR) << "Geography::fromWKB failed: " << geogRet.status();
+        return Value::kNullBadData;  // Is it ok to return Value::kNullBadData?
+      }
+      return std::move(geogRet).value();
     }
     case meta::cpp2::PropertyType::UNKNOWN:
       break;

diff --git a/src/codec/RowWriterV2.cpp b/src/codec/RowWriterV2.cpp
@@ -127,7 +127,7 @@ RowWriterV2::RowWriterV2(RowReader& reader) : RowWriterV2(reader.getSchema()) {
         set(i, v.moveDateTime());
         break;
       case Value::Type::GEOGRAPHY:
-        // TODO(jie)
+        set(i, v.moveGeography());
         break;
       default:
         LOG(FATAL) << "Invalid data: " << v << ", type: " << v.typeName();
@@ -207,8 +207,7 @@ WriteResult RowWriterV2::setValue(ssize_t index, const Value& val) noexcept {
     case Value::Type::DATETIME:
       return write(index, val.getDateTime());
     case Value::Type::GEOGRAPHY:
-      // TODO(jie)
-      return WriteResult::TYPE_MISMATCH;
+      return write(index, val.getGeography());
     default:
       return WriteResult::TYPE_MISMATCH;
   }
@@ -762,6 +761,17 @@ WriteResult RowWriterV2::write(ssize_t index, const DateTime& v) noexcept {
   return WriteResult::SUCCEEDED;
 }
 
+WriteResult RowWriterV2::write(ssize_t index, const Geography& v) noexcept {
+  auto field = schema_->field(index);
+  auto geoShape = field->geoShape();
+  if (geoShape != meta::cpp2::GeoShape::ANY &&
+      folly::to<uint32_t>(geoShape) != folly::to<uint32_t>(v.shape())) {
+    return WriteResult::TYPE_MISMATCH;
+  }
+  std::string wkb = v.asWKB();
+  return write(index, folly::StringPiece(wkb));
+}
+
 WriteResult RowWriterV2::checkUnsetFields() noexcept {
   DefaultValueContext expCtx;
   for (size_t i = 0; i < schema_->getNumFields(); i++) {
@@ -802,7 +812,7 @@ WriteResult RowWriterV2::checkUnsetFields() noexcept {
             r = write(i, defVal.getDateTime());
             break;
           case Value::Type::GEOGRAPHY:
-            // TODO(jie)
+            r = write(i, defVal.getGeography());
             break;
           default:
             LOG(FATAL) << "Unsupported default value type: " << defVal.typeName()

diff --git a/src/codec/RowWriterV2.h b/src/codec/RowWriterV2.h
@@ -63,6 +63,7 @@ enum class WriteResult {
         TIMESTAMP       (8 bytes)
         DATE            (4 bytes)
         DATETIME        (15 bytes)
+        GEOGRAPHY       (8 bytes) *
 
   All except STRING typed properties are stored in-place. The STRING property
   stored the offset of the string content in the first 4 bytes and the length
@@ -188,6 +189,8 @@ class RowWriterV2 {
   WriteResult write(ssize_t index, const Date& v) noexcept;
   WriteResult write(ssize_t index, const Time& v) noexcept;
   WriteResult write(ssize_t index, const DateTime& v) noexcept;
+
+  WriteResult write(ssize_t index, const Geography& v) noexcept;
 };
 
 }  // namespace nebula

diff --git a/src/codec/test/ResultSchemaProvider.cpp b/src/codec/test/ResultSchemaProvider.cpp
@@ -23,14 +23,16 @@ ResultSchemaProvider::ResultSchemaField::ResultSchemaField(std::string name,
                                                            bool nullable,
                                                            int32_t offset,
                                                            size_t nullFlagPos,
-                                                           Expression* defaultValue)
+                                                           Expression* defaultValue,
+                                                           meta::cpp2::GeoShape geoShape)
     : name_(std::move(name)),
       type_(type),
       size_(size),
       nullable_(nullable),
       offset_(offset),
       nullFlagPos_(nullFlagPos),
-      defaultValue_(defaultValue) {}
+      defaultValue_(defaultValue),
+      geoShape_(geoShape) {}
 
 const char* ResultSchemaProvider::ResultSchemaField::name() const { return name_.c_str(); }
 
@@ -50,6 +52,8 @@ size_t ResultSchemaProvider::ResultSchemaField::offset() const { return offset_;
 
 size_t ResultSchemaProvider::ResultSchemaField::nullFlagPos() const { return nullFlagPos_; }
 
+meta::cpp2::GeoShape ResultSchemaProvider::ResultSchemaField::geoShape() const { return geoShape_; }
+
 /***********************************
  *
  * ResultSchemaProvider

diff --git a/src/codec/test/ResultSchemaProvider.h b/src/codec/test/ResultSchemaProvider.h
@@ -22,7 +22,8 @@ class ResultSchemaProvider : public meta::SchemaProviderIf {
                                bool nullable,
                                int32_t offset,
                                size_t nullFlagPos,
-                               Expression* defaultValue = nullptr);
+                               Expression* defaultValue = nullptr,
+                               meta::cpp2::GeoShape = meta::cpp2::GeoShape::ANY);
 
     const char* name() const override;
     meta::cpp2::PropertyType type() const override;
@@ -32,6 +33,7 @@ class ResultSchemaProvider : public meta::SchemaProviderIf {
     size_t size() const override;
     size_t offset() const override;
     size_t nullFlagPos() const override;
+    meta::cpp2::GeoShape geoShape() const override;
 
    private:
     std::string name_;
@@ -41,6 +43,7 @@ class ResultSchemaProvider : public meta::SchemaProviderIf {
     int32_t offset_;
     size_t nullFlagPos_;
     Expression* defaultValue_;
+    meta::cpp2::GeoShape geoShape_;
   };
 
  public:

diff --git a/src/codec/test/SchemaWriter.cpp b/src/codec/test/SchemaWriter.cpp
@@ -17,7 +17,8 @@ SchemaWriter& SchemaWriter::appendCol(folly::StringPiece name,
                                       PropertyType type,
                                       int32_t fixedStrLen,
                                       bool nullable,
-                                      Expression* defaultValue) noexcept {
+                                      Expression* defaultValue,
+                                      meta::cpp2::GeoShape geoShape) noexcept {
   using folly::hash::SpookyHashV2;
   uint64_t hash = SpookyHashV2::Hash64(name.data(), name.size(), 0);
   DCHECK(nameIndex_.find(hash) == nameIndex_.end());
@@ -87,7 +88,8 @@ SchemaWriter& SchemaWriter::appendCol(folly::StringPiece name,
     nullFlagPos = numNullableFields_++;
   }
 
-  columns_.emplace_back(name.toString(), type, size, nullable, offset, nullFlagPos, defaultValue);
+  columns_.emplace_back(
+      name.toString(), type, size, nullable, offset, nullFlagPos, defaultValue, geoShape);
   nameIndex_.emplace(std::make_pair(hash, columns_.size() - 1));
 
   return *this;

diff --git a/src/codec/test/SchemaWriter.h b/src/codec/test/SchemaWriter.h
@@ -21,7 +21,8 @@ class SchemaWriter : public ResultSchemaProvider {
                           meta::cpp2::PropertyType type,
                           int32_t fixedStrLen = 0,
                           bool nullable = false,
-                          Expression* defaultValue = nullptr) noexcept;
+                          Expression* defaultValue = nullptr,
+                          meta::cpp2::GeoShape geoShape = meta::cpp2::GeoShape::ANY) noexcept;
 
  private:
 };

diff --git a/src/common/datatypes/Geography.cpp b/src/common/datatypes/Geography.cpp
@@ -19,20 +19,25 @@
 
 namespace nebula {
 
+constexpr double kMaxLongitude = 180.0;
+constexpr double kMaxLatitude = 90.0;
+
 void Coordinate::normalize() {
   // Reduce the x(longitude) to the range [-180, 180] degrees
   x = std::remainder(x, 360.0);
 
   // Reduce the y(latitude) to the range [-90, 90] degrees
   double tmp = remainder(y, 360.0);
-  if (tmp > 90.0) {
+  if (tmp > kMaxLatitude) {
     y = 180.0 - tmp;
-  } else if (tmp < -90.0) {
+  } else if (tmp < -kMaxLatitude) {
     y = -180.0 - tmp;
   }
 }
 
-bool Coordinate::isValid() const { return std::abs(x) <= 180.0 && std::abs(y) <= 90.0; }
+bool Coordinate::isValid() const {
+  return std::abs(x) <= kMaxLongitude && std::abs(y) <= kMaxLatitude;
+}
 
 void Point::normalize() {}
 
@@ -46,6 +51,7 @@ bool LineString::isValid() const {
     return false;
   }
   auto s2Region = geo::GeoUtils::s2RegionFromGeography(*this);
+  CHECK_NOTNULL(s2Region);
   return static_cast<S2Polyline*>(s2Region.get())->IsValid();
 }
 
@@ -67,6 +73,7 @@ bool Polygon::isValid() const {
     }
   }
   auto s2Region = geo::GeoUtils::s2RegionFromGeography(*this);
+  CHECK_NOTNULL(s2Region);
   return static_cast<S2Polygon*>(s2Region.get())->IsValid();
 }
 
@@ -91,6 +98,27 @@ StatusOr<Geography> Geography::fromWKT(const std::string& wkt,
   return geog;
 }
 
+StatusOr<Geography> Geography::fromWKB(const std::string& wkb,
+                                       bool needNormalize,
+                                       bool verifyValidity) {
+  auto geogRet = geo::WKBReader().read(wkb);
+  if (!geogRet.ok()) {
+    return geogRet;
+  }
+  auto geog = std::move(geogRet).value();
+  if (needNormalize) {
+    geog.normalize();
+  }
+  if (verifyValidity) {
+    if (!geog.isValid()) {
+      return Status::Error("Failed to parse an valid Geography instance from the wkb `%s'",
+                           wkb.c_str());
+    }
+  }
+
+  return geog;
+}
+
 GeoShape Geography::shape() const {
   switch (geo_.index()) {
     case 0:
@@ -184,6 +212,32 @@ bool Geography::isValid() const {
   }
 }
 
+Point Geography::centroid() const {
+  switch (shape()) {
+    case GeoShape::POINT: {
+      return this->point();
+    }
+    case GeoShape::LINESTRING: {
+      auto s2Region = geo::GeoUtils::s2RegionFromGeography(*this);
+      CHECK_NOTNULL(s2Region);
+      S2Point s2Point = static_cast<S2Polyline*>(s2Region.get())->GetCentroid();
+      return Point(geo::GeoUtils::coordinateFromS2Point(s2Point));
+    }
+    case GeoShape::POLYGON: {
+      auto s2Region = geo::GeoUtils::s2RegionFromGeography(*this);
+      CHECK_NOTNULL(s2Region);
+      S2Point s2Point = static_cast<S2Polygon*>(s2Region.get())->GetCentroid();
+      return Point(geo::GeoUtils::coordinateFromS2Point(s2Point));
+    }
+    case GeoShape::UNKNOWN:
+    default: {
+      LOG(ERROR)
+          << "Geography shapes other than Point/LineString/Polygon are not currently supported";
+      return Point();
+    }
+  }
+}
+
 std::string Geography::asWKT() const { return geo::WKTWriter().write(*this); }
 
 std::string Geography::asWKB() const { return geo::WKBWriter().write(*this); }

diff --git a/src/common/datatypes/Geography.h b/src/common/datatypes/Geography.h
@@ -58,8 +58,9 @@ struct Coordinate {
   }
   void __clear() { clear(); }
 
-  // TODO(jie) compare double correctly
-  bool operator==(const Coordinate& rhs) const { return x == rhs.x && y == rhs.y; }
+  bool operator==(const Coordinate& rhs) const {
+    return std::abs(x - rhs.x) < kEpsilon && std::abs(y - rhs.y) < kEpsilon;
+  }
   bool operator!=(const Coordinate& rhs) const { return !(*this == rhs); }
   bool operator<(const Coordinate& rhs) const {
     if (x != rhs.x) {
@@ -135,6 +136,10 @@ struct Geography {
                                      bool needNormalize = false,
                                      bool verifyValidity = false);
 
+  static StatusOr<Geography> fromWKB(const std::string& wkb,
+                                     bool needNormalize = false,
+                                     bool verifyValidity = false);
+
   Geography() {}
   Geography(const Point& v) : geo_(v) {}             // NOLINT
   Geography(Point&& v) : geo_(std::move(v)) {}       // NOLINT
@@ -156,6 +161,8 @@ struct Geography {
   void normalize();
   bool isValid() const;
 
+  Point centroid() const;
+
   std::string asWKT() const;
 
   std::string asWKB() const;

diff --git a/src/common/datatypes/test/GeographyTest.cpp b/src/common/datatypes/test/GeographyTest.cpp
@@ -21,14 +21,14 @@ TEST(Geography, shape) {
     EXPECT_EQ(GeoShape::POINT, g.shape());
   }
   {
-    std::string wkt = "LINESTRING(28.4 79.20,134.25 -28.34)";
+    std::string wkt = "LINESTRING(28.4 79.20, 134.25 -28.34)";
     auto gRet = Geography::fromWKT(wkt);
     ASSERT_TRUE(gRet.ok());
     auto g = gRet.value();
     EXPECT_EQ(GeoShape::LINESTRING, g.shape());
   }
   {
-    std::string wkt = "POLYGON((1 2,3 4,5 6,1 2))";
+    std::string wkt = "POLYGON((1 2, 3 4, 5 6, 1 2))";
     auto gRet = Geography::fromWKT(wkt);
     ASSERT_TRUE(gRet.ok());
     auto g = gRet.value();
@@ -46,15 +46,15 @@ TEST(Geography, asWKT) {
     EXPECT_EQ(wkt, got);
   }
   {
-    std::string wkt = "LINESTRING(28.4 79.2,134.25 -28.34)";
+    std::string wkt = "LINESTRING(28.4 79.2, 134.25 -28.34)";
     auto gRet = Geography::fromWKT(wkt);
     ASSERT_TRUE(gRet.ok());
     auto g = gRet.value();
     std::string got = g.asWKT();
     EXPECT_EQ(wkt, got);
   }
   {
-    std::string wkt = "POLYGON((1 2,3 4,5 6,1 2))";
+    std::string wkt = "POLYGON((1 2, 3 4, 5 6, 1 2))";
     auto gRet = Geography::fromWKT(wkt);
     ASSERT_TRUE(gRet.ok());
     auto g = gRet.value();