[feature](function) Add ST_NumGeometries, ST_NumPoints, ST_Geometries functions

Author: yixuanyxzhu

be/src/exprs/function/geo/functions_geo.cpp

@@ -28,8 +28,10 @@
 #include "core/block/block.h"
 #include "core/block/column_with_type_and_name.h"
 #include "core/column/column.h"
+#include "core/column/column_array.h"
 #include "core/column/column_execute_util.h"
 #include "core/column/column_nullable.h"
+#include "core/data_type/data_type_array.h"
 #include "core/data_type/data_type_nullable.h"
 #include "core/data_type/data_type_number.h"
 #include "core/data_type/data_type_string.h"
@@ -917,6 +919,167 @@ struct StDistance {
     }
 };
 
+struct StNumGeometries {
+    static constexpr auto NAME = "st_numgeometries";
+    static const size_t NUM_ARGS = 1;
+    using Type = DataTypeInt64;
+
+    static Status execute(Block& block, const ColumnNumbers& arguments, size_t result) {
+        DCHECK_EQ(arguments.size(), 1);
+
+        auto col = block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
+        const auto size = col->size();
+
+        auto res = ColumnInt64::create();
+        res->reserve(size);
+
+        auto null_map = ColumnUInt8::create(size, 0);
+        auto& null_map_data = null_map->get_data();
+
+        for (int row = 0; row < size; ++row) {
+            auto value = col->get_data_at(row);
+            auto shape = GeoShape::from_encoded(value.data, value.size);
+            if (!shape) {
+                null_map_data[row] = 1;
+                res->insert_default();
+                continue;
+            }
+
+            res->insert_value(shape->num_geometries());
+        }
+
+        block.replace_by_position(result,
+                                  ColumnNullable::create(std::move(res), std::move(null_map)));
+        return Status::OK();
+    }
+};
+
+struct StGeometries {
+    static constexpr auto NAME = "st_geometries";
+    static const size_t NUM_ARGS = 1;
+
+    static Status execute(Block& block, const ColumnNumbers& arguments, size_t result) {
+        DCHECK_EQ(arguments.size(), 1);
+
+        auto col = block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
+        const auto size = col->size();
+
+        auto nested_data = ColumnString::create();
+        auto offsets_col = ColumnArray::ColumnOffsets::create();
+        auto& offsets = offsets_col->get_data();
+        offsets.reserve(size);
+
+        auto null_map = ColumnUInt8::create(size, 0);
+        auto& null_map_data = null_map->get_data();
+
+        size_t current_offset = 0;
+        std::string buf;
+
+        for (size_t row = 0; row < size; ++row) {
+            auto shape_value = col->get_data_at(row);
+            auto shape = GeoShape::from_encoded(shape_value.data, shape_value.size);
+
+            if (!shape) {
+                null_map_data[row] = 1;
+                offsets.push_back(current_offset);
+                continue;
+            }
+
+            if (shape->type() == GEO_SHAPE_MULTI_POLYGON) {
+                auto* multi_polygon = static_cast<GeoMultiPolygon*>(shape.get());
+                const auto& polygons = multi_polygon->polygons();
+
+                if (polygons.empty()) {
+                    null_map_data[row] = 1;
+                    offsets.push_back(current_offset);
+                    continue;
+                }
+
+                for (const auto& polygon : polygons) {
+                    DCHECK(polygon != nullptr);
+                    buf.clear();
+                    polygon->encode_to(&buf);
+                    nested_data->insert_data(buf.data(), buf.size());
+                    ++current_offset;
+                }
+            } else {
+                nested_data->insert_data(shape_value.data, shape_value.size);
+                ++current_offset;
+            }
+
+            offsets.push_back(current_offset);
+        }
+
+        auto nested_null_map = ColumnUInt8::create(nested_data->size(), 0);
+        auto nested_nullable =
+                ColumnNullable::create(std::move(nested_data), std::move(nested_null_map));
+        auto array_col = ColumnArray::create(std::move(nested_nullable), std::move(offsets_col));
+
+        block.replace_by_position(
+                result, ColumnNullable::create(std::move(array_col), std::move(null_map)));
+
+        return Status::OK();
+    }
+};
+
+class FunctionStGeometries final : public IFunction {
+public:
+    static constexpr auto name = StGeometries::NAME;
+
+    static FunctionPtr create() { return std::make_shared<FunctionStGeometries>(); }
+
+    String get_name() const override { return name; }
+
+    size_t get_number_of_arguments() const override { return StGeometries::NUM_ARGS; }
+
+    bool is_variadic() const override { return false; }
+
+    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
+        return make_nullable(
+                std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeString>())));
+    }
+
+    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
+                        uint32_t result, size_t input_rows_count) const override {
+        return StGeometries::execute(block, arguments, result);
+    }
+};
+
+struct StNumPoints {
+    static constexpr auto NAME = "st_numpoints";
+    static const size_t NUM_ARGS = 1;
+    using Type = DataTypeInt64;
+
+    static Status execute(Block& block, const ColumnNumbers& arguments, size_t result) {
+        DCHECK_EQ(arguments.size(), 1);
+
+        auto col = block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
+        const auto size = col->size();
+
+        auto res = ColumnInt64::create();
+        res->reserve(size);
+
+        auto null_map = ColumnUInt8::create(size, 0);
+        auto& null_map_data = null_map->get_data();
+
+        for (int row = 0; row < size; ++row) {
+            auto value = col->get_data_at(row);
+            auto shape = GeoShape::from_encoded(value.data, value.size);
+            if (!shape) {
+                null_map_data[row] = 1;
+                res->insert_default();
+                continue;
+            }
+
+            res->insert_value(shape->num_points());
+        }
+
+        block.replace_by_position(result,
+                                  ColumnNullable::create(std::move(res), std::move(null_map)));
+        return Status::OK();
+    }
+};
+
 void register_function_geo(SimpleFunctionFactory& factory) {
     factory.register_function<GeoFunction<StPoint>>();
     factory.register_function<GeoFunction<StAsText<StAsWktName>>>();
@@ -947,6 +1110,9 @@ void register_function_geo(SimpleFunctionFactory& factory) {
     factory.register_function<GeoFunction<StLength>>();
     factory.register_function<GeoFunction<StGeometryType>>();
     factory.register_function<GeoFunction<StDistance>>();
+    factory.register_function<GeoFunction<StNumGeometries>>();
+    factory.register_function<GeoFunction<StNumPoints>>();
+    factory.register_function<FunctionStGeometries>();
 }
 
 } // namespace doris

be/src/exprs/function/geo/geo_types.cpp

@@ -1929,6 +1929,27 @@ double GeoMultiPolygon::Distance(const GeoShape* rhs) const {
     return (min_distance == std::numeric_limits<double>::max()) ? -1.0 : min_distance;
 }
 
+int GeoPolygon::num_points() const {
+    int total = 0;
+    for (int i = 0; i < numLoops(); ++i) {
+        auto* loop = getLoop(i);
+        DCHECK(loop != nullptr);
+        total += loop->num_vertices() + 1;
+    }
+    return total;
+}
+
+int GeoMultiPolygon::num_points() const {
+    int total = 0;
+    for (const auto& polygon : _polygons) {
+        DCHECK(polygon != nullptr);
+        int point_count = polygon->num_points();
+        DCHECK_GE(point_count, 0);
+        total += point_count;
+    }
+    return total;
+}
+
 double GeoCircle::Distance(const GeoShape* rhs) const {
     // Both rhs and self are guaranteed to be valid by StDistance (functions_geo.cpp)
     double circle_radius = S2Earth::ToMeters(_cap->radius());

be/src/exprs/function/geo/geo_types.h

@@ -81,6 +81,9 @@ class GeoShape {
 
     static bool ComputeArea(GeoShape* rhs, double* angle, std::string square_unit);
 
+    virtual int num_geometries() const { return 1; }
+    virtual int num_points() const { return -1; }
+
 protected:
     virtual void encode(std::string* buf) = 0;
     virtual bool decode(const void* data, size_t size) = 0;
@@ -125,6 +128,9 @@ class GeoPoint : public GeoShape {
     double x() const;
     double y() const;
 
+    int num_geometries() const override { return 1; }
+    int num_points() const override { return 1; }
+
 protected:
     void encode(std::string* buf) override;
     bool decode(const void* data, size_t size) override;
@@ -161,6 +167,9 @@ class GeoLine : public GeoShape {
     int numPoint() const;
     const S2Point* getPoint(int i) const;
 
+    int num_geometries() const override { return 1; }
+    int num_points() const override { return numPoint(); }
+
 protected:
     void encode(std::string* buf) override;
     bool decode(const void* data, size_t size) override;
@@ -199,6 +208,9 @@ class GeoPolygon : public GeoShape {
     double Distance(const GeoShape* rhs) const override;
     S2Loop* getLoop(int i) const;
 
+    int num_geometries() const override { return 1; }
+    int num_points() const override;
+
 protected:
     void encode(std::string* buf) override;
     bool decode(const void* data, size_t size) override;
@@ -232,6 +244,9 @@ class GeoMultiPolygon : public GeoShape {
     double Length() const override;
     double Distance(const GeoShape* rhs) const override;
 
+    int num_geometries() const override { return static_cast<int>(_polygons.size()); }
+    int num_points() const override;
+
 protected:
     void encode(std::string* buf) override;
     bool decode(const void* data, size_t size) override;

fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java

@@ -480,6 +480,7 @@
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StDistance;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StDistanceSphere;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeomFromWKB;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeometries;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeometryFromWKB;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeometryType;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeometryfromtext;
@@ -488,6 +489,8 @@
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StLength;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StLinefromtext;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StLinestringfromtext;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.StNumGeometries;
+import org.apache.doris.nereids.trees.expressions.functions.scalar.StNumPoints;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StPoint;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StPolyfromtext;
 import org.apache.doris.nereids.trees.expressions.functions.scalar.StPolygon;
@@ -1058,6 +1061,9 @@ public class BuiltinScalarFunctions implements FunctionHelper {
             scalar(StTouches.class, "st_touches"),
             scalar(StLength.class, "st_length"),
             scalar(StGeometryType.class, "st_geometrytype"),
+            scalar(StNumGeometries.class, "st_numgeometries"),
+            scalar(StGeometries.class, "st_geometries"),
+            scalar(StNumPoints.class, "st_numpoints"),
             scalar(StDistance.class, "st_distance"),
             scalar(StDistanceSphere.class, "st_distance_sphere"),
             scalar(StAngleSphere.class, "st_angle_sphere"),

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/StGeometries.java

@@ -0,0 +1,69 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package org.apache.doris.nereids.trees.expressions.functions.scalar;
+
+import org.apache.doris.catalog.FunctionSignature;
+import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable;
+import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
+import org.apache.doris.nereids.trees.expressions.functions.PropagateNullLiteral;
+import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression;
+import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
+import org.apache.doris.nereids.types.ArrayType;
+import org.apache.doris.nereids.types.VarcharType;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
+
+import java.util.List;
+
+/**
+ * ScalarFunction 'st_geometries'.
+ */
+public class StGeometries extends ScalarFunction
+        implements UnaryExpression, ExplicitlyCastableSignature, AlwaysNullable, PropagateNullLiteral {
+
+    public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
+            FunctionSignature.ret(ArrayType.of(VarcharType.SYSTEM_DEFAULT, true))
+                    .args(VarcharType.SYSTEM_DEFAULT)
+    );
+
+    public StGeometries(Expression arg0) {
+        super("st_geometries", arg0);
+    }
+
+    private StGeometries(ScalarFunctionParams functionParams) {
+        super(functionParams);
+    }
+
+    @Override
+    public StGeometries withChildren(List<Expression> children) {
+        Preconditions.checkArgument(children.size() == 1);
+        return new StGeometries(getFunctionParams(children));
+    }
+
+    @Override
+    public List<FunctionSignature> getSignatures() {
+        return SIGNATURES;
+    }
+
+    @Override
+    public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
+        return visitor.visitStGeometries(this, context);
+    }
+}