Skip to content

Commit

Permalink
[feat][test](nereids)(vec) support the ST_Intersects, ST_Disjoint, ST…
Browse files Browse the repository at this point in the history 
…_Touches function and test (apache#48203)
  • Loading branch information
@koi2000
koi2000 committed Mar 5, 2025
1 parent 54e7e94 commit 78ad92b
Show file tree
Hide file tree
Showing 14 changed files with 2,246 additions and 1 deletion.
446 changes: 446 additions & 0 deletions be/src/geo/geo_types.cpp
View file

Large diffs are not rendered by default.

23 changes: 23 additions & 0 deletions be/src/geo/geo_types.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand All @@ -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(); }
Expand Down Expand Up @@ -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(); }

Expand Down Expand Up @@ -148,6 +163,9 @@ 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;
std::string as_wkt() const override;

Expand All @@ -174,6 +192,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;

Expand Down
288 changes: 288 additions & 0 deletions be/src/vec/functions/functions_geo.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -721,6 +721,291 @@ struct StContains {
}
}; // namespace doris::vectorized

struct StIntersects {
static constexpr auto NEED_CONTEXT = true;
static constexpr auto NAME = "st_intersects";
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);
auto return_type = block.get_data_type(result);
const auto& [left_column, left_const] =
unpack_if_const(block.get_by_position(arguments[0]).column);
const auto& [right_column, right_const] =
unpack_if_const(block.get_by_position(arguments[1]).column);

const auto size = std::max(left_column->size(), right_column->size());

auto res = ColumnUInt8::create();
res->reserve(size);
auto null_map = ColumnUInt8::create(size, 0);
auto& null_map_data = null_map->get_data();

if (left_const) {
const_vector(left_column, right_column, res, null_map_data, size);
} else if (right_const) {
vector_const(left_column, right_column, res, null_map_data, size);
} else {
vector_vector(left_column, right_column, res, null_map_data, size);
}
block.replace_by_position(result,
ColumnNullable::create(std::move(res), std::move(null_map)));
return Status::OK();
}

static void loop_do(StringRef& lhs_value, StringRef& rhs_value,
std::vector<std::shared_ptr<GeoShape>>& shapes, int& i,
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) {
null_map[row] = 1;
res->insert_default();
break;
}
}

if (i == 2) {
auto contains_value = shapes[0]->intersects(shapes[1].get());
res->insert_data(const_cast<const char*>((char*)&contains_value), 0);
}
}

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};
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);
}
}

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};
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);
}
}

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};
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);
}
}

static Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) {
return Status::OK();
}

static Status close(FunctionContext* context, FunctionContext::FunctionStateScope scope) {
return Status::OK();
}
}; // namespace doris::vectorized

struct StDisjoint {
static constexpr auto NEED_CONTEXT = true;
static constexpr auto NAME = "st_disjoint";
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);
auto return_type = block.get_data_type(result);
const auto& [left_column, left_const] =
unpack_if_const(block.get_by_position(arguments[0]).column);
const auto& [right_column, right_const] =
unpack_if_const(block.get_by_position(arguments[1]).column);

const auto size = std::max(left_column->size(), right_column->size());

auto res = ColumnUInt8::create();
res->reserve(size);
auto null_map = ColumnUInt8::create(size, 0);
auto& null_map_data = null_map->get_data();

if (left_const) {
const_vector(left_column, right_column, res, null_map_data, size);
} else if (right_const) {
vector_const(left_column, right_column, res, null_map_data, size);
} else {
vector_vector(left_column, right_column, res, null_map_data, size);
}
block.replace_by_position(result,
ColumnNullable::create(std::move(res), std::move(null_map)));
return Status::OK();
}

static void loop_do(StringRef& lhs_value, StringRef& rhs_value,
std::vector<std::shared_ptr<GeoShape>>& shapes, int& i,
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) {
null_map[row] = 1;
res->insert_default();
break;
}
}

if (i == 2) {
auto contains_value = shapes[0]->disjoint(shapes[1].get());
res->insert_data(const_cast<const char*>((char*)&contains_value), 0);
}
}

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};
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);
}
}

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};
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);
}
}

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};
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);
}
}

static Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) {
return Status::OK();
}

static Status close(FunctionContext* context, FunctionContext::FunctionStateScope scope) {
return Status::OK();
}
}; // namespace doris::vectorized

struct StTouches {
static constexpr auto NEED_CONTEXT = true;
static constexpr auto NAME = "st_touches";
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);
auto return_type = block.get_data_type(result);
const auto& [left_column, left_const] =
unpack_if_const(block.get_by_position(arguments[0]).column);
const auto& [right_column, right_const] =
unpack_if_const(block.get_by_position(arguments[1]).column);

const auto size = std::max(left_column->size(), right_column->size());

auto res = ColumnUInt8::create();
res->reserve(size);
auto null_map = ColumnUInt8::create(size, 0);
auto& null_map_data = null_map->get_data();

if (left_const) {
const_vector(left_column, right_column, res, null_map_data, size);
} else if (right_const) {
vector_const(left_column, right_column, res, null_map_data, size);
} else {
vector_vector(left_column, right_column, res, null_map_data, size);
}
block.replace_by_position(result,
ColumnNullable::create(std::move(res), std::move(null_map)));
return Status::OK();
}

static void loop_do(StringRef& lhs_value, StringRef& rhs_value,
std::vector<std::shared_ptr<GeoShape>>& shapes, int& i,
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) {
null_map[row] = 1;
res->insert_default();
break;
}
}

if (i == 2) {
auto contains_value = shapes[0]->touches(shapes[1].get());
res->insert_data(const_cast<const char*>((char*)&contains_value), 0);
}
}

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};
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);
}
}

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};
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);
}
}

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};
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);
}
}

static Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) {
return Status::OK();
}

static Status close(FunctionContext* context, FunctionContext::FunctionStateScope scope) {
return Status::OK();
}
}; // namespace doris::vectorized

struct StGeometryFromText {
static constexpr auto NAME = "st_geometryfromtext";
static constexpr GeoShapeType shape_type = GEO_SHAPE_ANY;
Expand Down Expand Up @@ -915,6 +1200,9 @@ void register_function_geo(SimpleFunctionFactory& factory) {
factory.register_function<GeoFunction<StAngle>>();
factory.register_function<GeoFunction<StAzimuth>>();
factory.register_function<GeoFunction<StContains>>();
factory.register_function<GeoFunction<StIntersects>>();
factory.register_function<GeoFunction<StDisjoint>>();
factory.register_function<GeoFunction<StTouches>>();
factory.register_function<GeoFunction<StCircle>>();
factory.register_function<GeoFunction<StGeoFromText<StGeometryFromText>>>();
factory.register_function<GeoFunction<StGeoFromText<StGeomFromText>>>();
Expand Down
Loading

0 comments on commit 78ad92b

Please sign in to comment.