[Enhancement] Support some spatial functions (#48695)

Support for ST_Intersects, ST_Disjoint, ST_Touches sql functions.

Author: koi2000

be/src/geo/geo_types.cpp

@@ -60,6 +60,13 @@ class GeoShape {
     virtual std::string as_wkt() const = 0;
 
     virtual bool contains(const GeoShape* rhs) const { return false; }
+
+    virtual bool disjoint(const GeoShape* rhs) const { return false; }
+
+    virtual bool intersects(const GeoShape* rhs) const { return false; }
+
+    virtual bool touches(const GeoShape* rhs) const { return false; }
+
     virtual std::string to_string() const { return ""; }
     static std::string as_binary(GeoShape* rhs);
 
@@ -82,6 +89,10 @@ class GeoPoint : public GeoShape {
 
     GeoCoordinateList to_coords() const;
 
+    bool intersects(const GeoShape* rhs) const override;
+    bool disjoint(const GeoShape* rhs) const override;
+    bool touches(const GeoShape* rhs) const override;
+
     GeoShapeType type() const override { return GEO_SHAPE_POINT; }
 
     const S2Point* point() const { return _point.get(); }
@@ -119,6 +130,10 @@ class GeoLine : public GeoShape {
 
     GeoCoordinateList to_coords() const;
 
+    bool intersects(const GeoShape* rhs) const override;
+    bool disjoint(const GeoShape* rhs) const override;
+    bool touches(const GeoShape* rhs) const override;
+
     GeoShapeType type() const override { return GEO_SHAPE_LINE_STRING; }
     const S2Polyline* polyline() const { return _polyline.get(); }
 
@@ -148,7 +163,14 @@ class GeoPolygon : public GeoShape {
     GeoShapeType type() const override { return GEO_SHAPE_POLYGON; }
     const S2Polygon* polygon() const { return _polygon.get(); }
 
+    bool intersects(const GeoShape* rhs) const override;
+    bool disjoint(const GeoShape* rhs) const override;
+    bool touches(const GeoShape* rhs) const override;
     bool contains(const GeoShape* rhs) const override;
+
+    bool polygon_touch_point(const S2Polygon* polygon, const S2Point* point) const;
+    bool polygon_touch_polygon(const S2Polygon* polygon1, const S2Polygon* polygon2) const;
+
     std::string as_wkt() const override;
 
     int numLoops() const;
@@ -174,6 +196,11 @@ class GeoCircle : public GeoShape {
 
     GeoShapeType type() const override { return GEO_SHAPE_CIRCLE; }
 
+    const S2Cap* circle() const { return _cap.get(); }
+
+    bool intersects(const GeoShape* rhs) const override;
+    bool disjoint(const GeoShape* rhs) const override;
+    bool touches(const GeoShape* rhs) const override;
     bool contains(const GeoShape* rhs) const override;
     std::string as_wkt() const override;
 

be/src/geo/geo_types.h

@@ -626,11 +626,13 @@ struct StCircle {
     }
 };
 
-struct StContains {
+template <typename Func>
+struct StRelationFunction {
     static constexpr auto NEED_CONTEXT = true;
-    static constexpr auto NAME = "st_contains";
+    static constexpr auto NAME = Func::NAME;
     static const size_t NUM_ARGS = 2;
     using Type = DataTypeUInt8;
+
     static Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                           size_t result) {
         DCHECK_EQ(arguments.size(), 2);
@@ -642,8 +644,7 @@ struct StContains {
 
         const auto size = std::max(left_column->size(), right_column->size());
 
-        auto res = ColumnUInt8::create();
-        res->reserve(size);
+        auto res = ColumnUInt8::create(size, 0);
         auto null_map = ColumnUInt8::create(size, 0);
         auto& null_map_data = null_map->get_data();
 
@@ -660,55 +661,50 @@ struct StContains {
     }
 
     static void loop_do(StringRef& lhs_value, StringRef& rhs_value,
-                        std::vector<std::shared_ptr<GeoShape>>& shapes, int& i,
+                        std::vector<std::unique_ptr<GeoShape>>& shapes,
                         ColumnUInt8::MutablePtr& res, NullMap& null_map, int row) {
         StringRef* strs[2] = {&lhs_value, &rhs_value};
-        for (i = 0; i < 2; ++i) {
-            shapes[i] =
-                    std::shared_ptr<GeoShape>(GeoShape::from_encoded(strs[i]->data, strs[i]->size));
-            if (shapes[i] == nullptr) {
+        for (int i = 0; i < 2; ++i) {
+            std::unique_ptr<GeoShape> shape(GeoShape::from_encoded(strs[i]->data, strs[i]->size));
+            shapes[i] = std::move(shape);
+            if (!shapes[i]) {
                 null_map[row] = 1;
-                res->insert_default();
                 break;
             }
         }
-
-        if (i == 2) {
-            auto contains_value = shapes[0]->contains(shapes[1].get());
-            res->insert_data(const_cast<const char*>((char*)&contains_value), 0);
+        if (shapes[0] && shapes[1]) {
+            auto relation_value = Func::evaluate(shapes[0].get(), shapes[1].get());
+            res->get_data()[row] = relation_value;
         }
     }
 
     static void const_vector(const ColumnPtr& left_column, const ColumnPtr& right_column,
                              ColumnUInt8::MutablePtr& res, NullMap& null_map, const size_t size) {
-        int i;
         auto lhs_value = left_column->get_data_at(0);
-        std::vector<std::shared_ptr<GeoShape>> shapes = {nullptr, nullptr};
+        std::vector<std::unique_ptr<GeoShape>> shapes(2);
         for (int row = 0; row < size; ++row) {
             auto rhs_value = right_column->get_data_at(row);
-            loop_do(lhs_value, rhs_value, shapes, i, res, null_map, row);
+            loop_do(lhs_value, rhs_value, shapes, res, null_map, row);
         }
     }
 
     static void vector_const(const ColumnPtr& left_column, const ColumnPtr& right_column,
                              ColumnUInt8::MutablePtr& res, NullMap& null_map, const size_t size) {
-        int i;
         auto rhs_value = right_column->get_data_at(0);
-        std::vector<std::shared_ptr<GeoShape>> shapes = {nullptr, nullptr};
+        std::vector<std::unique_ptr<GeoShape>> shapes(2);
         for (int row = 0; row < size; ++row) {
             auto lhs_value = left_column->get_data_at(row);
-            loop_do(lhs_value, rhs_value, shapes, i, res, null_map, row);
+            loop_do(lhs_value, rhs_value, shapes, res, null_map, row);
         }
     }
 
     static void vector_vector(const ColumnPtr& left_column, const ColumnPtr& right_column,
                               ColumnUInt8::MutablePtr& res, NullMap& null_map, const size_t size) {
-        int i;
-        std::vector<std::shared_ptr<GeoShape>> shapes = {nullptr, nullptr};
+        std::vector<std::unique_ptr<GeoShape>> shapes(2);
         for (int row = 0; row < size; ++row) {
             auto lhs_value = left_column->get_data_at(row);
             auto rhs_value = right_column->get_data_at(row);
-            loop_do(lhs_value, rhs_value, shapes, i, res, null_map, row);
+            loop_do(lhs_value, rhs_value, shapes, res, null_map, row);
         }
     }
 
@@ -719,7 +715,27 @@ struct StContains {
     static Status close(FunctionContext* context, FunctionContext::FunctionStateScope scope) {
         return Status::OK();
     }
-}; // namespace doris::vectorized
+};
+
+struct StContainsFunc {
+    static constexpr auto NAME = "st_contains";
+    static bool evaluate(GeoShape* shape1, GeoShape* shape2) { return shape1->contains(shape2); }
+};
+
+struct StIntersectsFunc {
+    static constexpr auto NAME = "st_intersects";
+    static bool evaluate(GeoShape* shape1, GeoShape* shape2) { return shape1->intersects(shape2); }
+};
+
+struct StDisjointFunc {
+    static constexpr auto NAME = "st_disjoint";
+    static bool evaluate(GeoShape* shape1, GeoShape* shape2) { return shape1->disjoint(shape2); }
+};
+
+struct StTouchesFunc {
+    static constexpr auto NAME = "st_touches";
+    static bool evaluate(GeoShape* shape1, GeoShape* shape2) { return shape1->touches(shape2); }
+};
 
 struct StGeometryFromText {
     static constexpr auto NAME = "st_geometryfromtext";
@@ -914,7 +930,10 @@ void register_function_geo(SimpleFunctionFactory& factory) {
     factory.register_function<GeoFunction<StAngleSphere>>();
     factory.register_function<GeoFunction<StAngle>>();
     factory.register_function<GeoFunction<StAzimuth>>();
-    factory.register_function<GeoFunction<StContains>>();
+    factory.register_function<GeoFunction<StRelationFunction<StContainsFunc>>>();
+    factory.register_function<GeoFunction<StRelationFunction<StIntersectsFunc>>>();
+    factory.register_function<GeoFunction<StRelationFunction<StDisjointFunc>>>();
+    factory.register_function<GeoFunction<StRelationFunction<StTouchesFunc>>>();
     factory.register_function<GeoFunction<StCircle>>();
     factory.register_function<GeoFunction<StGeoFromText<StGeometryFromText>>>();
     factory.register_function<GeoFunction<StGeoFromText<StGeomFromText>>>();