add ST_Centroid

vesoft-inc · Oct 14, 2021 · e30c3f5 · e30c3f5
1 parent f752557
commit e30c3f5
Show file tree

Hide file tree

Showing 7 changed files with 159 additions and 15 deletions.
diff --git a/src/common/datatypes/Geography.cpp b/src/common/datatypes/Geography.cpp
@@ -51,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();
 }
 
@@ -72,6 +73,7 @@ bool Polygon::isValid() const {
     }
   }
   auto s2Region = geo::GeoUtils::s2RegionFromGeography(*this);
+  CHECK_NOTNULL(s2Region);
   return static_cast<S2Polygon*>(s2Region.get())->IsValid();
 }
 
@@ -210,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) {
@@ -160,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/function/FunctionManager.cpp b/src/common/function/FunctionManager.cpp
@@ -338,6 +338,11 @@ std::unordered_map<std::string, std::vector<TypeSignature>> FunctionManager::typ
      {
          TypeSignature({Value::Type::GEOGRAPHY}, Value::Type::STRING),
      }},
+    // geo transformations
+    {"st_centroid",
+     {
+         TypeSignature({Value::Type::GEOGRAPHY}, Value::Type::GEOGRAPHY),
+     }},
     // geo accessors
     {"st_isvalid",
      {
@@ -2414,6 +2419,19 @@ FunctionManager::FunctionManager() {
       return g.asWKBHex();
     };
   }
+  // geo transformations
+  {
+    auto &attr = functions_["st_centroid"];
+    attr.minArity_ = 1;
+    attr.maxArity_ = 1;
+    attr.isPure_ = true;
+    attr.body_ = [](const auto &args) -> Value {
+      if (!args[0].get().isGeography()) {
+        return Value::kNullBadType;
+      }
+      return Geography(args[0].get().getGeography().centroid());
+    };
+  }
   // geo accessors
   {
     auto &attr = functions_["st_isvalid"];

diff --git a/src/common/geo/GeoIndex.cpp b/src/common/geo/GeoIndex.cpp
@@ -30,7 +30,7 @@ namespace geo {
 
 nebula::storage::cpp2::IndexColumnHint ScanRange::toIndexColumnHint() {
   nebula::storage::cpp2::IndexColumnHint hint;
-  // set_column_name should be called later
+  // column_name should be set by the caller
   if (isRangeScan) {
     hint.set_scan_type(nebula::storage::cpp2::ScanType::RANGE);
     hint.set_begin_value(

diff --git a/src/common/geo/GeoUtils.h b/src/common/geo/GeoUtils.h
@@ -6,6 +6,7 @@
 
 #pragma once
 
+#include <s2/s2latlng.h>
 #include <s2/s2loop.h>
 #include <s2/s2polygon.h>
 
@@ -26,28 +27,23 @@ class GeoUtils final {
       }
       case GeoShape::LINESTRING: {
         const auto& lineString = geog.lineString();
-        auto coordList = lineString.coordList;
+        const auto& coordList = lineString.coordList;
         auto s2Points = s2PointsFromCoordinateList(coordList);
-        auto s2Polyline = std::make_unique<S2Polyline>(s2Points, S2Debug::DISABLE);
-        return s2Polyline;
+        return std::make_unique<S2Polyline>(s2Points, S2Debug::DISABLE);
       }
       case GeoShape::POLYGON: {
         const auto& polygon = geog.polygon();
         uint32_t numCoordList = polygon.numCoordList();
         std::vector<std::unique_ptr<S2Loop>> s2Loops;
         s2Loops.reserve(numCoordList);
-        for (size_t i = 0; i < numCoordList; ++i) {
-          auto coordList = polygon.coordListList[i];
-          if (!coordList.empty()) {
-            coordList.pop_back();  // Remove redundant last coordinate
-          }
-          auto s2Points = s2PointsFromCoordinateList(coordList);
+        for (const auto& coordList : polygon.coordListList) {
+          auto s2Points = s2PointsFromCoordinateList(
+              coordList, true);  // S2 doesn't need the redundant last point
           auto s2Loop = std::make_unique<S2Loop>(std::move(s2Points), S2Debug::DISABLE);
           s2Loop->Normalize();  // All loops must be oriented CCW(counterclockwise) for S2
           s2Loops.emplace_back(std::move(s2Loop));
         }
-        auto s2Polygon = std::make_unique<S2Polygon>(std::move(s2Loops), S2Debug::DISABLE);
-        return s2Polygon;
+        return std::make_unique<S2Polygon>(std::move(s2Loops), S2Debug::DISABLE);
       }
       default:
         LOG(FATAL)
@@ -62,9 +58,22 @@ class GeoUtils final {
     return latlng.ToPoint();
   }
 
-  static std::vector<S2Point> s2PointsFromCoordinateList(const std::vector<Coordinate>& coordList) {
+  static Coordinate coordinateFromS2Point(const S2Point& s2Point) {
+    S2LatLng s2Latlng(s2Point);
+    return Coordinate(s2Latlng.lat().degrees(), s2Latlng.lng().degrees());
+  }
+
+  static std::vector<S2Point> s2PointsFromCoordinateList(const std::vector<Coordinate>& coordList,
+                                                         bool excludeTheLast = false) {
     std::vector<S2Point> s2Points;
     uint32_t numCoords = coordList.size();
+    if (excludeTheLast) {
+      numCoords -= 1;
+    }
+    if (numCoords == 0) {
+      return {};
+    }
+
     s2Points.reserve(numCoords);
     for (size_t i = 0; i < numCoords; ++i) {
       auto coord = coordList[i];

diff --git a/src/graph/validator/LookupValidator.cpp b/src/graph/validator/LookupValidator.cpp
@@ -273,6 +273,15 @@ StatusOr<Expression*> LookupValidator::checkFilter(Expression* expr) {
     auto relExpr = static_cast<RelationalExpression*>(expr);
     NG_RETURN_IF_ERROR(checkRelExpr(relExpr));
     return rewriteRelExpr(relExpr);
+  } else if (expr->kind() == Expression::Kind::kFunctionCall) {
+    auto funcExpr = static_cast<FunctionCallExpression*>(expr);
+    std::unordered_set<std::string> geoIndexAcceleratedPredicates{
+        "st_intersects", "st_covers", "st_coveredby", "st_dwithin"};
+    if (geoIndexAcceleratedPredicates.find(funcExpr->name()) !=
+        geoIndexAcceleratedPredicates.end()) {
+      NG_RETURN_IF_ERROR(checkGeoFuncExpr(funcExpr));
+      return rewriteGeoFuncExpr(funcExpr);
+    }
   }
   switch (expr->kind()) {
     case ExprKind::kLogicalOr: {
@@ -302,6 +311,25 @@ Status LookupValidator::checkRelExpr(RelationalExpression* expr) {
   return Status::SemanticError("Expression %s not supported yet", expr->toString().c_str());
 }
 
+Status LookupValidator::checkGeoFuncExpr(const FunctionCallExpression* expr) {
+  if (expr->args()->numArgs() < 2) {
+    return Status::SemanticError("Expression %s has not enough arguments",
+                                 expr->toString().c_str());
+  }
+
+  auto* first = expr->args()->args()[0];
+  auto* second = expr->args()->args()[1];
+
+  if (first->kind() == ExprKind::kLabelAttribute && second->kind() == ExprKind::kLabelAttribute) {
+    return Status::SemanticError("Expression %s not supported yet", expr->toString().c_str());
+  }
+  if (first->kind() == ExprKind::kLabelAttribute || second->kind() == ExprKind::kLabelAttribute) {
+    return Status::OK();
+  }
+
+  return Status::SemanticError("Expression %s not supported yet", expr->toString().c_str());
+}
+
 StatusOr<Expression*> LookupValidator::rewriteRelExpr(RelationalExpression* expr) {
   // swap LHS and RHS of relExpr if LabelAttributeExpr in on the right,
   // so that LabelAttributeExpr is always on the left
@@ -336,6 +364,40 @@ StatusOr<Expression*> LookupValidator::rewriteRelExpr(RelationalExpression* expr
   return expr;
 }
 
+StatusOr<FunctionCallExpression*> LookupValidator::rewriteGeoFuncExpr(
+    FunctionCallExpression* expr) {
+  // swap LHS and RHS of relExpr if LabelAttributeExpr in on the right,
+  // so that LabelAttributeExpr is always on the left
+  auto* second = expr->args()->args()[1];
+  if (second->kind() == ExprKind::kLabelAttribute) {
+    expr = reverseGeoFunc(expr);
+  }
+
+  auto* first = expr->args()->args()[0];
+  auto* la = static_cast<LabelAttributeExpression*>(first);
+  if (la->left()->name() != sentence()->from()) {
+    return Status::SemanticError("Schema name error: %s", la->left()->name().c_str());
+  }
+
+  // fold constant expression
+  auto foldRes = ExpressionUtils::foldConstantExpr(expr);
+  NG_RETURN_IF_ERROR(foldRes);
+  expr = static_cast<FunctionCallExpression*>(foldRes.value());
+  DCHECK_EQ(expr->args()->args()[0]->kind(), ExprKind::kLabelAttribute);
+
+  // Check schema and value type
+  std::string prop = la->right()->value().getStr();
+  auto c = checkConstExpr(expr->args()->args()[1], prop, expr->kind());
+  NG_RETURN_IF_ERROR(c);
+  expr->args()->setArg(1, std::move(c).value());
+
+  // rewrite PropertyExpression
+  auto propExpr = lookupCtx_->isEdge ? ExpressionUtils::rewriteLabelAttr2EdgeProp(la)
+                                     : ExpressionUtils::rewriteLabelAttr2TagProp(la);
+  expr->args()->setArg(0, propExpr);
+  return expr;
+}
+
 StatusOr<Expression*> LookupValidator::checkConstExpr(Expression* expr,
                                                       const std::string& prop,
                                                       const ExprKind kind) {
@@ -437,6 +499,27 @@ Expression* LookupValidator::reverseRelKind(RelationalExpression* expr) {
   return RelationalExpression::makeKind(pool, reversedKind, right->clone(), left->clone());
 }
 
+FunctionCallExpression* LookupValidator::reverseGeoFunc(const FunctionCallExpression* expr) {
+  const auto& name = expr->name();
+  auto newName = name;
+
+  if (name == "st_covers") {
+    newName = "st_coveredby";
+  } else if (name == "st_coveredby") {
+    newName = "st_covers";
+  } else if (name == "st_intersects") {
+  } else if (name == "st_dwithin") {
+  }
+
+  auto* newArgList = ArgumentList::make(qctx_->objPool(), expr->args()->numArgs());
+  for (auto& arg : expr->args()->args()) {
+    newArgList->addArgument(arg->clone());
+  }
+  newArgList->setArg(0, expr->args()->args()[1]);
+  newArgList->setArg(1, expr->args()->args()[0]);
+  return FunctionCallExpression::make(qctx_->objPool(), newName, newArgList);
+}
+
 Status LookupValidator::getSchemaProvider(shared_ptr<const NebulaSchemaProvider>* provider) const {
   auto from = sentence()->from();
   auto schemaMgr = qctx_->schemaMng();

diff --git a/src/graph/validator/LookupValidator.h b/src/graph/validator/LookupValidator.h
@@ -39,12 +39,15 @@ class LookupValidator final : public Validator {
 
   StatusOr<Expression*> checkFilter(Expression* expr);
   Status checkRelExpr(RelationalExpression* expr);
+  Status checkGeoFuncExpr(const FunctionCallExpression* expr);
   StatusOr<std::string> checkTSExpr(Expression* expr);
   StatusOr<Expression*> checkConstExpr(Expression* expr,
                                        const std::string& prop,
                                        const Expression::Kind kind);
   StatusOr<Expression*> rewriteRelExpr(RelationalExpression* expr);
+  StatusOr<FunctionCallExpression*> rewriteGeoFuncExpr(FunctionCallExpression* expr);
   Expression* reverseRelKind(RelationalExpression* expr);
+  FunctionCallExpression* reverseGeoFunc(const FunctionCallExpression* expr);
 
   const LookupSentence* sentence() const;
   int32_t schemaId() const;