Skip to content

Commit e870bd2

Browse files
author
yixuanyxzhu
committed
[feature](function) Add ST_NumGeometries, ST_NumPoints, ST_Geometries functions
1 parent da73f0b commit e870bd2

10 files changed

Lines changed: 1236 additions & 0 deletions

File tree

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

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,10 @@
2828
#include "core/block/block.h"
2929
#include "core/block/column_with_type_and_name.h"
3030
#include "core/column/column.h"
31+
#include "core/column/column_array.h"
3132
#include "core/column/column_execute_util.h"
3233
#include "core/column/column_nullable.h"
34+
#include "core/data_type/data_type_array.h"
3335
#include "core/data_type/data_type_nullable.h"
3436
#include "core/data_type/data_type_number.h"
3537
#include "core/data_type/data_type_string.h"
@@ -917,6 +919,167 @@ struct StDistance {
917919
}
918920
};
919921

922+
struct StNumGeometries {
923+
static constexpr auto NAME = "st_numgeometries";
924+
static const size_t NUM_ARGS = 1;
925+
using Type = DataTypeInt64;
926+
927+
static Status execute(Block& block, const ColumnNumbers& arguments, size_t result) {
928+
DCHECK_EQ(arguments.size(), 1);
929+
930+
auto col = block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
931+
const auto size = col->size();
932+
933+
auto res = ColumnInt64::create();
934+
res->reserve(size);
935+
936+
auto null_map = ColumnUInt8::create(size, 0);
937+
auto& null_map_data = null_map->get_data();
938+
939+
for (int row = 0; row < size; ++row) {
940+
auto value = col->get_data_at(row);
941+
auto shape = GeoShape::from_encoded(value.data, value.size);
942+
if (!shape) {
943+
null_map_data[row] = 1;
944+
res->insert_default();
945+
continue;
946+
}
947+
948+
res->insert_value(shape->num_geometries());
949+
}
950+
951+
block.replace_by_position(result,
952+
ColumnNullable::create(std::move(res), std::move(null_map)));
953+
return Status::OK();
954+
}
955+
};
956+
957+
struct StGeometries {
958+
static constexpr auto NAME = "st_geometries";
959+
static const size_t NUM_ARGS = 1;
960+
961+
static Status execute(Block& block, const ColumnNumbers& arguments, size_t result) {
962+
DCHECK_EQ(arguments.size(), 1);
963+
964+
auto col = block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
965+
const auto size = col->size();
966+
967+
auto nested_data = ColumnString::create();
968+
auto offsets_col = ColumnArray::ColumnOffsets::create();
969+
auto& offsets = offsets_col->get_data();
970+
offsets.reserve(size);
971+
972+
auto null_map = ColumnUInt8::create(size, 0);
973+
auto& null_map_data = null_map->get_data();
974+
975+
size_t current_offset = 0;
976+
std::string buf;
977+
978+
for (size_t row = 0; row < size; ++row) {
979+
auto shape_value = col->get_data_at(row);
980+
auto shape = GeoShape::from_encoded(shape_value.data, shape_value.size);
981+
982+
if (!shape) {
983+
null_map_data[row] = 1;
984+
offsets.push_back(current_offset);
985+
continue;
986+
}
987+
988+
if (shape->type() == GEO_SHAPE_MULTI_POLYGON) {
989+
auto* multi_polygon = static_cast<GeoMultiPolygon*>(shape.get());
990+
const auto& polygons = multi_polygon->polygons();
991+
992+
if (polygons.empty()) {
993+
null_map_data[row] = 1;
994+
offsets.push_back(current_offset);
995+
continue;
996+
}
997+
998+
for (const auto& polygon : polygons) {
999+
DCHECK(polygon != nullptr);
1000+
buf.clear();
1001+
polygon->encode_to(&buf);
1002+
nested_data->insert_data(buf.data(), buf.size());
1003+
++current_offset;
1004+
}
1005+
} else {
1006+
nested_data->insert_data(shape_value.data, shape_value.size);
1007+
++current_offset;
1008+
}
1009+
1010+
offsets.push_back(current_offset);
1011+
}
1012+
1013+
auto nested_null_map = ColumnUInt8::create(nested_data->size(), 0);
1014+
auto nested_nullable =
1015+
ColumnNullable::create(std::move(nested_data), std::move(nested_null_map));
1016+
auto array_col = ColumnArray::create(std::move(nested_nullable), std::move(offsets_col));
1017+
1018+
block.replace_by_position(
1019+
result, ColumnNullable::create(std::move(array_col), std::move(null_map)));
1020+
1021+
return Status::OK();
1022+
}
1023+
};
1024+
1025+
class FunctionStGeometries final : public IFunction {
1026+
public:
1027+
static constexpr auto name = StGeometries::NAME;
1028+
1029+
static FunctionPtr create() { return std::make_shared<FunctionStGeometries>(); }
1030+
1031+
String get_name() const override { return name; }
1032+
1033+
size_t get_number_of_arguments() const override { return StGeometries::NUM_ARGS; }
1034+
1035+
bool is_variadic() const override { return false; }
1036+
1037+
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
1038+
return make_nullable(
1039+
std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeString>())));
1040+
}
1041+
1042+
Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
1043+
uint32_t result, size_t input_rows_count) const override {
1044+
return StGeometries::execute(block, arguments, result);
1045+
}
1046+
};
1047+
1048+
struct StNumPoints {
1049+
static constexpr auto NAME = "st_numpoints";
1050+
static const size_t NUM_ARGS = 1;
1051+
using Type = DataTypeInt64;
1052+
1053+
static Status execute(Block& block, const ColumnNumbers& arguments, size_t result) {
1054+
DCHECK_EQ(arguments.size(), 1);
1055+
1056+
auto col = block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
1057+
const auto size = col->size();
1058+
1059+
auto res = ColumnInt64::create();
1060+
res->reserve(size);
1061+
1062+
auto null_map = ColumnUInt8::create(size, 0);
1063+
auto& null_map_data = null_map->get_data();
1064+
1065+
for (int row = 0; row < size; ++row) {
1066+
auto value = col->get_data_at(row);
1067+
auto shape = GeoShape::from_encoded(value.data, value.size);
1068+
if (!shape) {
1069+
null_map_data[row] = 1;
1070+
res->insert_default();
1071+
continue;
1072+
}
1073+
1074+
res->insert_value(shape->num_points());
1075+
}
1076+
1077+
block.replace_by_position(result,
1078+
ColumnNullable::create(std::move(res), std::move(null_map)));
1079+
return Status::OK();
1080+
}
1081+
};
1082+
9201083
void register_function_geo(SimpleFunctionFactory& factory) {
9211084
factory.register_function<GeoFunction<StPoint>>();
9221085
factory.register_function<GeoFunction<StAsText<StAsWktName>>>();
@@ -947,6 +1110,9 @@ void register_function_geo(SimpleFunctionFactory& factory) {
9471110
factory.register_function<GeoFunction<StLength>>();
9481111
factory.register_function<GeoFunction<StGeometryType>>();
9491112
factory.register_function<GeoFunction<StDistance>>();
1113+
factory.register_function<GeoFunction<StNumGeometries>>();
1114+
factory.register_function<GeoFunction<StNumPoints>>();
1115+
factory.register_function<FunctionStGeometries>();
9501116
}
9511117

9521118
} // namespace doris

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1929,6 +1929,27 @@ double GeoMultiPolygon::Distance(const GeoShape* rhs) const {
19291929
return (min_distance == std::numeric_limits<double>::max()) ? -1.0 : min_distance;
19301930
}
19311931

1932+
int GeoPolygon::num_points() const {
1933+
int total = 0;
1934+
for (int i = 0; i < numLoops(); ++i) {
1935+
auto* loop = getLoop(i);
1936+
DCHECK(loop != nullptr);
1937+
total += loop->num_vertices() + 1;
1938+
}
1939+
return total;
1940+
}
1941+
1942+
int GeoMultiPolygon::num_points() const {
1943+
int total = 0;
1944+
for (const auto& polygon : _polygons) {
1945+
DCHECK(polygon != nullptr);
1946+
int point_count = polygon->num_points();
1947+
DCHECK_GE(point_count, 0);
1948+
total += point_count;
1949+
}
1950+
return total;
1951+
}
1952+
19321953
double GeoCircle::Distance(const GeoShape* rhs) const {
19331954
// Both rhs and self are guaranteed to be valid by StDistance (functions_geo.cpp)
19341955
double circle_radius = S2Earth::ToMeters(_cap->radius());

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,9 @@ class GeoShape {
8181

8282
static bool ComputeArea(GeoShape* rhs, double* angle, std::string square_unit);
8383

84+
virtual int num_geometries() const { return 1; }
85+
virtual int num_points() const { return -1; }
86+
8487
protected:
8588
virtual void encode(std::string* buf) = 0;
8689
virtual bool decode(const void* data, size_t size) = 0;
@@ -125,6 +128,9 @@ class GeoPoint : public GeoShape {
125128
double x() const;
126129
double y() const;
127130

131+
int num_geometries() const override { return 1; }
132+
int num_points() const override { return 1; }
133+
128134
protected:
129135
void encode(std::string* buf) override;
130136
bool decode(const void* data, size_t size) override;
@@ -161,6 +167,9 @@ class GeoLine : public GeoShape {
161167
int numPoint() const;
162168
const S2Point* getPoint(int i) const;
163169

170+
int num_geometries() const override { return 1; }
171+
int num_points() const override { return numPoint(); }
172+
164173
protected:
165174
void encode(std::string* buf) override;
166175
bool decode(const void* data, size_t size) override;
@@ -199,6 +208,9 @@ class GeoPolygon : public GeoShape {
199208
double Distance(const GeoShape* rhs) const override;
200209
S2Loop* getLoop(int i) const;
201210

211+
int num_geometries() const override { return 1; }
212+
int num_points() const override;
213+
202214
protected:
203215
void encode(std::string* buf) override;
204216
bool decode(const void* data, size_t size) override;
@@ -232,6 +244,9 @@ class GeoMultiPolygon : public GeoShape {
232244
double Length() const override;
233245
double Distance(const GeoShape* rhs) const override;
234246

247+
int num_geometries() const override { return static_cast<int>(_polygons.size()); }
248+
int num_points() const override;
249+
235250
protected:
236251
void encode(std::string* buf) override;
237252
bool decode(const void* data, size_t size) override;

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -479,6 +479,7 @@
479479
import org.apache.doris.nereids.trees.expressions.functions.scalar.StDistance;
480480
import org.apache.doris.nereids.trees.expressions.functions.scalar.StDistanceSphere;
481481
import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeomFromWKB;
482+
import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeometries;
482483
import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeometryFromWKB;
483484
import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeometryType;
484485
import org.apache.doris.nereids.trees.expressions.functions.scalar.StGeometryfromtext;
@@ -487,6 +488,8 @@
487488
import org.apache.doris.nereids.trees.expressions.functions.scalar.StLength;
488489
import org.apache.doris.nereids.trees.expressions.functions.scalar.StLinefromtext;
489490
import org.apache.doris.nereids.trees.expressions.functions.scalar.StLinestringfromtext;
491+
import org.apache.doris.nereids.trees.expressions.functions.scalar.StNumGeometries;
492+
import org.apache.doris.nereids.trees.expressions.functions.scalar.StNumPoints;
490493
import org.apache.doris.nereids.trees.expressions.functions.scalar.StPoint;
491494
import org.apache.doris.nereids.trees.expressions.functions.scalar.StPolyfromtext;
492495
import org.apache.doris.nereids.trees.expressions.functions.scalar.StPolygon;
@@ -1056,6 +1059,9 @@ public class BuiltinScalarFunctions implements FunctionHelper {
10561059
scalar(StTouches.class, "st_touches"),
10571060
scalar(StLength.class, "st_length"),
10581061
scalar(StGeometryType.class, "st_geometrytype"),
1062+
scalar(StNumGeometries.class, "st_numgeometries"),
1063+
scalar(StGeometries.class, "st_geometries"),
1064+
scalar(StNumPoints.class, "st_numpoints"),
10591065
scalar(StDistance.class, "st_distance"),
10601066
scalar(StDistanceSphere.class, "st_distance_sphere"),
10611067
scalar(StAngleSphere.class, "st_angle_sphere"),

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
package org.apache.doris.nereids.trees.expressions.functions.scalar;
19+
20+
import org.apache.doris.catalog.FunctionSignature;
21+
import org.apache.doris.nereids.trees.expressions.Expression;
22+
import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable;
23+
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
24+
import org.apache.doris.nereids.trees.expressions.functions.PropagateNullLiteral;
25+
import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression;
26+
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
27+
import org.apache.doris.nereids.types.ArrayType;
28+
import org.apache.doris.nereids.types.VarcharType;
29+
30+
import com.google.common.base.Preconditions;
31+
import com.google.common.collect.ImmutableList;
32+
33+
import java.util.List;
34+
35+
/**
36+
* ScalarFunction 'st_geometries'.
37+
*/
38+
public class StGeometries extends ScalarFunction
39+
implements UnaryExpression, ExplicitlyCastableSignature, AlwaysNullable, PropagateNullLiteral {
40+
41+
public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
42+
FunctionSignature.ret(ArrayType.of(VarcharType.SYSTEM_DEFAULT, true))
43+
.args(VarcharType.SYSTEM_DEFAULT)
44+
);
45+
46+
public StGeometries(Expression arg0) {
47+
super("st_geometries", arg0);
48+
}
49+
50+
private StGeometries(ScalarFunctionParams functionParams) {
51+
super(functionParams);
52+
}
53+
54+
@Override
55+
public StGeometries withChildren(List<Expression> children) {
56+
Preconditions.checkArgument(children.size() == 1);
57+
return new StGeometries(getFunctionParams(children));
58+
}
59+
60+
@Override
61+
public List<FunctionSignature> getSignatures() {
62+
return SIGNATURES;
63+
}
64+
65+
@Override
66+
public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
67+
return visitor.visitStGeometries(this, context);
68+
}
69+
}

0 commit comments

Comments
 (0)