[improvement](geo) Decode constant shapes once in st_contains and other spatial relation functions

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

@@ -504,7 +504,15 @@ struct StRelationFunction {
     static const size_t NUM_ARGS = 2;
     using Type = DataTypeUInt8;
 
-    static Status execute(Block& block, const ColumnNumbers& arguments, size_t result) {
+    static Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) {
+        if (scope == FunctionContext::FunctionStateScope::THREAD_LOCAL) {
+            context->set_function_state(scope, std::make_shared<StRelationState>());
+        }
+        return Status::OK();
+    }
+
+    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);
         auto left_col = ColumnView<TYPE_STRING>::create(block.get_by_position(arguments[0]).column);
@@ -517,26 +525,61 @@ struct StRelationFunction {
         auto null_map = ColumnUInt8::create(size, 0);
         auto& null_map_data = null_map->get_data();
 
-        for (int row = 0; row < size; ++row) {
-            auto lhs_value = left_col.value_at(row);
-            auto rhs_value = right_col.value_at(row);
+        auto* state = reinterpret_cast<StRelationState*>(
+                context->get_function_state(FunctionContext::FunctionStateScope::THREAD_LOCAL));
+        std::shared_ptr<GeoShape> const_shape1;
+        std::shared_ptr<GeoShape> const_shape2;
+        if (left_col.is_const) {
+            const_shape1 = decode_const(state, 0, left_col.value_at(0));
+        }
+        if (right_col.is_const) {
+            const_shape2 = decode_const(state, 1, right_col.value_at(0));
+        }
 
-            std::unique_ptr<GeoShape> shape1(
-                    GeoShape::from_encoded(lhs_value.data, lhs_value.size));
-            std::unique_ptr<GeoShape> shape2(
-                    GeoShape::from_encoded(rhs_value.data, rhs_value.size));
+        for (int row = 0; row < size; ++row) {
+            std::unique_ptr<GeoShape> row_shape1;
+            std::unique_ptr<GeoShape> row_shape2;
+            GeoShape* shape1 = const_shape1.get();
+            GeoShape* shape2 = const_shape2.get();
+            if (!left_col.is_const) {
+                auto lhs_value = left_col.value_at(row);
+                row_shape1 = GeoShape::from_encoded(lhs_value.data, lhs_value.size);
+                shape1 = row_shape1.get();
+            }
+            if (!right_col.is_const) {
+                auto rhs_value = right_col.value_at(row);
+                row_shape2 = GeoShape::from_encoded(rhs_value.data, rhs_value.size);
+                shape2 = row_shape2.get();
+            }
 
             if (!shape1 || !shape2) {
                 null_map_data[row] = 1;
                 continue;
             }
-            auto relation_value = Func::evaluate(shape1.get(), shape2.get());
+            auto relation_value = Func::evaluate(shape1, shape2);
             res->get_data()[row] = relation_value;
         }
         block.replace_by_position(result,
                                   ColumnNullable::create(std::move(res), std::move(null_map)));
         return Status::OK();
     }
+
+private:
+    // Decodes a constant argument, reusing the shape decoded for an earlier block when the value
+    // is unchanged. Returns nullptr if the value is not a valid shape.
+    static std::shared_ptr<GeoShape> decode_const(StRelationState* state, int arg,
+                                                  const StringRef& value) {
+        if (state == nullptr) {
+            return GeoShape::from_encoded(value.data, value.size);
+        }
+        auto& cached = state->args[arg];
+        if (!cached.decoded || StringRef(cached.encoded) != value) {
+            cached.encoded.assign(value.data, value.size);
+            cached.shape = GeoShape::from_encoded(cached.encoded.data(), cached.encoded.size());
+            cached.decoded = true;
+        }
+        return cached.shape;
+    }
 };
 
 struct StContainsFunc {

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

@@ -46,11 +46,18 @@ struct StConstructState {
     std::string encoded_buf;
 };
 
-struct StContainsState {
-    StContainsState() : is_null(false), shapes {nullptr, nullptr} {}
-    ~StContainsState() {}
-    bool is_null;
-    std::vector<std::shared_ptr<GeoShape>> shapes;
+// Thread-local state of st_contains, st_within and the other relation functions: the decoded
+// shape of each constant argument, keyed by its encoded value. Decoding a polygon validates it,
+// which builds an index over all its edges, so a constant shape is decoded once per thread
+// instead of once per row.
+struct StRelationState {
+    struct ConstShape {
+        bool decoded = false;
+        std::string encoded;
+        // nullptr if `encoded` is not a valid shape
+        std::shared_ptr<GeoShape> shape;
+    };
+    ConstShape args[2];
 };
 
 template <typename Impl>
@@ -67,9 +74,21 @@ class GeoFunction : public IFunction {
         return make_nullable(std::make_shared<ReturnType>());
     }
 
+    Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override {
+        if constexpr (requires { Impl::open(context, scope); }) {
+            return Impl::open(context, scope);
+        } else {
+            return Status::OK();
+        }
+    }
+
     Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                         uint32_t result, size_t input_rows_count) const override {
-        return Impl::execute(block, arguments, result);
+        if constexpr (requires { Impl::execute(context, block, arguments, result); }) {
+            return Impl::execute(context, block, arguments, result);
+        } else {
+            return Impl::execute(block, arguments, result);
+        }
     }
 };
 

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

@@ -37,6 +37,7 @@
 
 #include "core/assert_cast.h"
 // IWYU pragma: no_include <bits/std_abs.h>
+#include <algorithm>
 #include <cmath>
 #include <iomanip>
 #include <limits>
@@ -788,6 +789,7 @@ const S2Point* GeoLine::getPoint(int i) const {
 }
 
 GeoParseStatus GeoPolygon::from_coords(const GeoCoordinateListList& list) {
+    _loop_lat_lngs.clear();
     return to_s2polygon(list, &_polygon);
 }
 
@@ -800,6 +802,7 @@ void GeoPolygon::encode(std::string* buf) {
 bool GeoPolygon::decode(const void* data, size_t size) {
     Decoder decoder(data, size);
     _polygon.reset(new S2Polygon());
+    _loop_lat_lngs.clear();
     return _polygon->Decode(&decoder) && _polygon->IsValid();
 }
 
@@ -1043,17 +1046,7 @@ bool GeoPolygon::contains(const GeoShape* rhs) const {
         }
 
         // Point on the edge of polygon doesn't count as "Contians"
-        for (int i = 0; i < _polygon->num_loops(); ++i) {
-            const S2Loop* loop = _polygon->loop(i);
-            for (int j = 0; j < loop->num_vertices(); ++j) {
-                const S2Point& p1 = loop->vertex(j);
-                const S2Point& p2 = loop->vertex((j + 1) % loop->num_vertices());
-                if (compute_distance_to_line(*point->point(), p1, p2) < TOLERANCE) {
-                    return false;
-                }
-            }
-        }
-        return true;
+        return !point_on_edge(*point->point());
     }
     case GeoShapeType::GEO_SHAPE_LINE_STRING: {
         const GeoLine* line = (const GeoLine*)rhs;
@@ -1097,6 +1090,46 @@ bool GeoPolygon::contains(const GeoShape* rhs) const {
     }
 }
 
+// Same check as compute_distance_to_line(point, p1, p2) < TOLERANCE over every edge, with the
+// vertices converted to lat/lng once per polygon instead of once per edge and call.
+bool GeoPolygon::point_on_edge(const S2Point& point) const {
+    if (_loop_lat_lngs.empty()) {
+        _loop_lat_lngs.resize(_polygon->num_loops());
+        for (int i = 0; i < _polygon->num_loops(); ++i) {
+            const S2Loop* loop = _polygon->loop(i);
+            auto& lat_lngs = _loop_lat_lngs[i];
+            lat_lngs.reserve(loop->num_vertices());
+            for (int j = 0; j < loop->num_vertices(); ++j) {
+                S2LatLng vertex(loop->vertex(j));
+                lat_lngs.emplace_back(vertex.lat().degrees(), vertex.lng().degrees(), 0.0);
+            }
+        }
+    }
+
+    S2LatLng point_lat_lng(point);
+    const S2Point query(point_lat_lng.lat().degrees(), point_lat_lng.lng().degrees(), 0.0);
+    // The point can only be within TOLERANCE of an edge if it lies in the edge's bounding box
+    // grown by TOLERANCE. The margin is doubled to stay clear of rounding in project_distance.
+    constexpr double margin = 2 * TOLERANCE;
+    for (const auto& lat_lngs : _loop_lat_lngs) {
+        const size_t num_vertices = lat_lngs.size();
+        for (size_t j = 0; j < num_vertices; ++j) {
+            const S2Point& p1 = lat_lngs[j];
+            const S2Point& p2 = lat_lngs[j + 1 == num_vertices ? 0 : j + 1];
+            if (query.x() < std::min(p1.x(), p2.x()) - margin ||
+                query.x() > std::max(p1.x(), p2.x()) + margin ||
+                query.y() < std::min(p1.y(), p2.y()) - margin ||
+                query.y() > std::max(p1.y(), p2.y()) + margin) {
+                continue;
+            }
+            if (project_distance(query, p1, p2) < TOLERANCE) {
+                return true;
+            }
+        }
+    }
+    return false;
+}
+
 std::double_t GeoPolygon::getArea() const {
     return _polygon->GetArea();
 }

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

@@ -218,7 +218,13 @@ class GeoPolygon : public GeoShape {
     bool decode(const void* data, size_t size) override;
 
 private:
+    bool point_on_edge(const S2Point& point) const;
+
     std::unique_ptr<S2Polygon> _polygon;
+    // (lat, lng, 0) in degrees of every loop's vertices, for the point-on-edge check in
+    // contains(). Built on first use, so a polygon that is checked against many points converts
+    // its vertices only once.
+    mutable std::vector<std::vector<S2Point>> _loop_lat_lngs;
 };
 
 class GeoMultiPolygon : public GeoShape {

be/test/exprs/function/function_geo_test.cpp

@@ -25,13 +25,20 @@
 #include <vector>
 
 #include "common/status.h"
+#include "core/assert_cast.h"
+#include "core/block/block.h"
+#include "core/column/column_const.h"
+#include "core/column/column_nullable.h"
+#include "core/column/column_string.h"
+#include "core/column/column_vector.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"
 #include "core/types.h"
 #include "exprs/function/function_test_util.h"
 #include "exprs/function/geo/geo_common.h"
 #include "exprs/function/geo/geo_types.h"
+#include "exprs/function/simple_function_factory.h"
 #include "gtest/gtest_pred_impl.h"
 #include "testutil/any_type.h"
 
@@ -317,6 +324,104 @@ TEST(VGeoFunctionsTest, function_geo_st_contains) {
     }
 }
 
+namespace {
+std::string encode_wkt(const std::string& wkt) {
+    GeoParseStatus status;
+    std::unique_ptr<GeoShape> shape(GeoShape::from_wkt(wkt.data(), wkt.size(), status));
+    EXPECT_TRUE(status == GeoParseStatus::GEO_PARSE_OK);
+    std::string buf;
+    shape->encode_to(&buf);
+    return buf;
+}
+
+std::string encode_point(double x, double y) {
+    GeoPoint point;
+    EXPECT_TRUE(point.from_coord(x, y) == GeoParseStatus::GEO_PARSE_OK);
+    std::string buf;
+    point.encode_to(&buf);
+    return buf;
+}
+} // namespace
+
+// A constant polygon is decoded once and then checked against every row of the block.
+TEST(VGeoFunctionsTest, function_geo_st_contains_const_polygon) {
+    std::string polygon = encode_wkt("POLYGON ((0 0, 10 0, 10 10, 0 10, 0 0))");
+    DataSet data_set = {{{polygon, encode_point(5, 5)}, (uint8_t)1},
+                        {{polygon, encode_point(50, 50)}, (uint8_t)0},
+                        {{polygon, encode_point(10, 5)}, (uint8_t)0},
+                        {{polygon, encode_point(10 - 5e-7, 5)}, (uint8_t)0},
+                        {{polygon, encode_point(10 - 2e-6, 5)}, (uint8_t)1},
+                        {{polygon, Null()}, Null()},
+                        {{polygon, encode_point(1, 9)}, (uint8_t)1}};
+    InputTypeSet input_types = {Consted {PrimitiveType::TYPE_VARCHAR}, PrimitiveType::TYPE_VARCHAR};
+    static_cast<void>(check_function<DataTypeUInt8, true>("st_contains", input_types, data_set));
+
+    // The polygon is the second argument of st_within.
+    DataSet within_data_set = {{{encode_point(5, 5), polygon}, (uint8_t)1},
+                               {{encode_point(50, 50), polygon}, (uint8_t)0},
+                               {{encode_point(1, 9), polygon}, (uint8_t)1}};
+    InputTypeSet within_input_types = {PrimitiveType::TYPE_VARCHAR,
+                                       Consted {PrimitiveType::TYPE_VARCHAR}};
+    static_cast<void>(
+            check_function<DataTypeUInt8, true>("st_within", within_input_types, within_data_set));
+
+    // An invalid constant makes every row NULL.
+    DataSet invalid_data_set = {{{std::string("invalid"), encode_point(5, 5)}, Null()},
+                                {{std::string("invalid"), encode_point(50, 50)}, Null()}};
+    static_cast<void>(
+            check_function<DataTypeUInt8, true>("st_contains", input_types, invalid_data_set));
+}
+
+// The decoded constant is kept in the thread-local state across blocks. A block whose constant
+// differs from the previous one must not reuse the shape decoded for it.
+TEST(VGeoFunctionsTest, function_geo_st_contains_const_polygon_across_blocks) {
+    std::string small_polygon = encode_wkt("POLYGON ((0 0, 10 0, 10 10, 0 10, 0 0))");
+    std::string large_polygon = encode_wkt("POLYGON ((0 0, 30 0, 30 30, 0 30, 0 0))");
+    std::string inside_both = encode_point(5, 5);
+    std::string inside_large = encode_point(20, 20);
+
+    DataTypePtr string_type = std::make_shared<DataTypeString>();
+    DataTypePtr return_type = make_nullable(std::make_shared<DataTypeUInt8>());
+    auto make_block = [&](const std::string& polygon) {
+        auto polygon_column = ColumnString::create();
+        polygon_column->insert_data(polygon.data(), polygon.size());
+        auto point_column = ColumnString::create();
+        point_column->insert_data(inside_both.data(), inside_both.size());
+        point_column->insert_data(inside_large.data(), inside_large.size());
+        Block block;
+        block.insert({ColumnConst::create(std::move(polygon_column), 2), string_type, "polygon"});
+        block.insert({std::move(point_column), string_type, "point"});
+        return block;
+    };
+
+    Block sample_block = make_block(small_polygon);
+    auto func = SimpleFunctionFactory::instance().get_function(
+            "st_contains", sample_block.get_columns_with_type_and_name(), return_type);
+    ASSERT_NE(nullptr, func);
+    FunctionUtils fn_utils(return_type, {string_type, string_type}, false);
+    auto* fn_ctx = fn_utils.get_fn_ctx();
+    ASSERT_TRUE(func->open(fn_ctx, FunctionContext::FunctionStateScope::FRAGMENT_LOCAL).ok());
+    ASSERT_TRUE(func->open(fn_ctx, FunctionContext::FunctionStateScope::THREAD_LOCAL).ok());
+
+    auto contains = [&](const std::string& polygon) {
+        Block block = make_block(polygon);
+        block.insert({nullptr, return_type, "result"});
+        EXPECT_TRUE(func->execute(fn_ctx, block, {0, 1}, 2, 2).ok());
+        const auto& result = assert_cast<const ColumnNullable&>(*block.get_by_position(2).column);
+        EXPECT_FALSE(result.is_null_at(0));
+        EXPECT_FALSE(result.is_null_at(1));
+        const auto& data = assert_cast<const ColumnUInt8&>(result.get_nested_column()).get_data();
+        return std::vector<uint8_t> {data[0], data[1]};
+    };
+    EXPECT_EQ((std::vector<uint8_t> {1, 0}), contains(small_polygon));
+    EXPECT_EQ((std::vector<uint8_t> {1, 0}), contains(small_polygon));
+    EXPECT_EQ((std::vector<uint8_t> {1, 1}), contains(large_polygon));
+    EXPECT_EQ((std::vector<uint8_t> {1, 0}), contains(small_polygon));
+
+    static_cast<void>(func->close(fn_ctx, FunctionContext::FunctionStateScope::THREAD_LOCAL));
+    static_cast<void>(func->close(fn_ctx, FunctionContext::FunctionStateScope::FRAGMENT_LOCAL));
+}
+
 TEST(VGeoFunctionsTest, function_geo_st_circle) {
     std::string func_name = "st_circle";
     {