feat(query): ST_CONTAINS (#15994)

kkk25641463 · web-flow · commit 0e55fa9bc10e · 2024-07-12T05:00:36.000Z
* ST_CONTAINS

Signed-off-by: Fan Yang &lt;yangfanlinux@gmail.com&gt;

* ST_CONTAINS

Signed-off-by: Fan Yang &lt;yangfanlinux@gmail.com&gt;

---------

Signed-off-by: Fan Yang &lt;yangfanlinux@gmail.com&gt;
diff --git a/scripts/benchmark/query/load/tpch10.sh b/scripts/benchmark/query/load/tpch10.sh
@@ -7,7 +7,7 @@ select version();
 SQL
 
 for t in customer lineitem nation orders partsupp part region supplier; do
-  echo "DROP TABLE IF EXISTS $t;" | bendsql
+	echo "DROP TABLE IF EXISTS $t;" | bendsql
 done
 
 cat <<SQL | bendsql
@@ -113,8 +113,8 @@ cat <<SQL | bendsql
 SQL
 
 for t in customer lineitem nation orders partsupp part region supplier; do
-  echo "loading into $t ..."
-  cat <<SQL | bendsql
+	echo "loading into $t ..."
+	cat <<SQL | bendsql
 COPY INTO $t FROM 's3://repo.databend.com/datasets/tpch10/${t}/' connection=(connection_name='repo') pattern ='${t}.*'
 file_format=(type='CSV' field_delimiter='|' record_delimiter='\\n' skip_header=0);
 ANALYZE TABLE "${t}";
diff --git a/scripts/benchmark/query/load/tpch100.sh b/scripts/benchmark/query/load/tpch100.sh
@@ -7,7 +7,7 @@ select version();
 SQL
 
 for t in customer lineitem nation orders partsupp part region supplier; do
-  echo "DROP TABLE IF EXISTS $t;" | bendsql
+	echo "DROP TABLE IF EXISTS $t;" | bendsql
 done
 
 cat <<SQL | bendsql
@@ -113,8 +113,8 @@ cat <<SQL | bendsql
 SQL
 
 for t in nation region; do
-  echo "loading into $t ..."
-  cat <<SQL | bendsql
+	echo "loading into $t ..."
+	cat <<SQL | bendsql
 COPY INTO $t FROM 's3://repo.databend.com/tpch100/${t}.tbl'
 credentials=(access_key_id ='$REPO_ACCESS_KEY_ID' secret_access_key ='$REPO_SECRET_ACCESS_KEY')
 file_format=(type='CSV' field_delimiter='|' record_delimiter='\\n' skip_header=1);
@@ -124,8 +124,8 @@ SQL
 done
 
 for t in customer lineitem orders partsupp part supplier; do
-  echo "loading into $t ..."
-  cat <<SQL | bendsql
+	echo "loading into $t ..."
+	cat <<SQL | bendsql
 COPY INTO $t FROM 's3://repo.databend.com/tpch100/${t}/' connection=(connection_name='repo') pattern ='${t}.tbl.*'
 file_format=(type='CSV' field_delimiter='|' record_delimiter='\\n' skip_header=1);
 ANALYZE TABLE "${t}";
diff --git a/src/query/functions/src/scalars/geometry.rs b/src/query/functions/src/scalars/geometry.rs
@@ -15,6 +15,7 @@
 use databend_common_exception::ErrorCode;
 use databend_common_expression::types::geometry::GeometryType;
 use databend_common_expression::types::BinaryType;
+use databend_common_expression::types::BooleanType;
 use databend_common_expression::types::Int32Type;
 use databend_common_expression::types::NullableType;
 use databend_common_expression::types::NumberType;
@@ -34,6 +35,7 @@ use databend_common_io::Extremum;
 use databend_common_io::GeometryDataType;
 use geo::dimensions::Dimensions;
 use geo::BoundingRect;
+use geo::Contains;
 use geo::EuclideanDistance;
 use geo::EuclideanLength;
 use geo::HasDimensions;
@@ -214,6 +216,45 @@ pub fn register(registry: &mut FunctionRegistry) {
         }),
     );
 
+    registry.register_combine_nullable_2_arg::<GeometryType, GeometryType, BooleanType, _, _>(
+        "st_contains",
+        |_, _, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_2_arg::<GeometryType, GeometryType, NullableType<BooleanType>>(
+            |l_geometry, r_geometry, builder, ctx| {
+                let l_ewkb = Ewkb(l_geometry);
+                let r_ewkb = Ewkb(r_geometry);
+                let l_geos: Geometry = l_ewkb.to_geos().unwrap();
+                let r_geos: Geometry = r_ewkb.to_geos().unwrap();
+                let l_srid = l_geos.srid();
+                let r_srid = r_geos.srid();
+                if l_srid != r_srid {
+                    builder.push_null();
+                    ctx.set_error(
+                        builder.len(),
+                        ErrorCode::GeometryError("Srid does not match!").to_string(),
+                    );
+                } else {
+                    let l_geo: geo::Geometry = l_geos.to_geo().unwrap();
+                    let r_geo: geo::Geometry = r_geos.to_geo().unwrap();
+                    if matches!(l_geo, geo::Geometry::GeometryCollection(_))
+                        || matches!(r_geo, geo::Geometry::GeometryCollection(_))
+                    {
+                        builder.push_null();
+                        ctx.set_error(
+                            builder.len(),
+                            ErrorCode::GeometryError(
+                                "A GEOMETRY object that is a GeometryCollection",
+                            )
+                            .to_string(),
+                        );
+                    } else {
+                        builder.push(l_geo.contains(&r_geo));
+                    }
+                }
+            },
+        ),
+    );
+
     registry.register_combine_nullable_2_arg::<GeometryType, GeometryType, NumberType<F64>, _, _>(
         "st_distance",
         |_, _, _| FunctionDomain::MayThrow,
diff --git a/src/query/functions/tests/it/scalars/geometry.rs b/src/query/functions/tests/it/scalars/geometry.rs
@@ -31,6 +31,7 @@ fn test_geometry() {
     test_st_aswkb(file);
     test_st_asewkt(file);
     test_st_aswkt(file);
+    test_st_contains(file);
     test_st_endpoint(file);
     test_st_dimension(file);
     test_st_distance(file);
@@ -180,6 +181,29 @@ fn test_st_asgeojson(file: &mut impl Write) {
     );
 }
 
+fn test_st_contains(file: &mut impl Write) {
+    run_ast(
+        file,
+        "ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'))",
+        &[],
+    );
+    run_ast(
+        file,
+        "ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('POLYGON((-1 0, 0 1, 1 0, -1 0))'))",
+        &[],
+    );
+    run_ast(
+        file,
+        "ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('LINESTRING(-1 1, 0 2, 1 1))'))",
+        &[],
+    );
+    run_ast(
+        file,
+        "ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('LINESTRING(-2 0, 0 0, 0 1))'))",
+        &[],
+    );
+}
+
 fn test_st_endpoint(file: &mut impl Write) {
     run_ast(
         file,
diff --git a/src/query/functions/tests/it/scalars/testdata/function_list.txt b/src/query/functions/tests/it/scalars/testdata/function_list.txt
@@ -3552,6 +3552,8 @@ Functions overloads:
 1 st_aswkb(Geometry NULL) :: Binary NULL
 0 st_aswkt(Geometry) :: String
 1 st_aswkt(Geometry NULL) :: String NULL
+0 st_contains(Geometry, Geometry) :: Boolean NULL
+1 st_contains(Geometry NULL, Geometry NULL) :: Boolean NULL
 0 st_dimension(Geometry) :: Int32 NULL
 1 st_dimension(Geometry NULL) :: Int32 NULL
 0 st_distance(Geometry, Geometry) :: Float64 NULL
diff --git a/src/query/functions/tests/it/scalars/testdata/geometry.txt b/src/query/functions/tests/it/scalars/testdata/geometry.txt
@@ -79,6 +79,42 @@ output domain  : {"LINESTRING(0.75 0.75,-10 20)"..="LINESTRING(0.75 0.75,-10 20)
 output         : 'LINESTRING(0.75 0.75,-10 20)'
 
 
+ast            : ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'))
+raw expr       : ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'))
+checked expr   : st_contains<Geometry, Geometry>(to_geometry<String>("POLYGON((-2 0, 0 2, 2 0, -2 0))"), to_geometry<String>("POLYGON((-2 0, 0 2, 2 0, -2 0))"))
+optimized expr : true
+output type    : Boolean NULL
+output domain  : {TRUE}
+output         : true
+
+
+ast            : ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('POLYGON((-1 0, 0 1, 1 0, -1 0))'))
+raw expr       : ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('POLYGON((-1 0, 0 1, 1 0, -1 0))'))
+checked expr   : st_contains<Geometry, Geometry>(to_geometry<String>("POLYGON((-2 0, 0 2, 2 0, -2 0))"), to_geometry<String>("POLYGON((-1 0, 0 1, 1 0, -1 0))"))
+optimized expr : true
+output type    : Boolean NULL
+output domain  : {TRUE}
+output         : true
+
+
+ast            : ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('LINESTRING(-1 1, 0 2, 1 1))'))
+raw expr       : ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('LINESTRING(-1 1, 0 2, 1 1))'))
+checked expr   : st_contains<Geometry, Geometry>(to_geometry<String>("POLYGON((-2 0, 0 2, 2 0, -2 0))"), to_geometry<String>("LINESTRING(-1 1, 0 2, 1 1))"))
+optimized expr : false
+output type    : Boolean NULL
+output domain  : {FALSE}
+output         : false
+
+
+ast            : ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('LINESTRING(-2 0, 0 0, 0 1))'))
+raw expr       : ST_CONTAINS(TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))'), TO_GEOMETRY('LINESTRING(-2 0, 0 0, 0 1))'))
+checked expr   : st_contains<Geometry, Geometry>(to_geometry<String>("POLYGON((-2 0, 0 2, 2 0, -2 0))"), to_geometry<String>("LINESTRING(-2 0, 0 0, 0 1))"))
+optimized expr : true
+output type    : Boolean NULL
+output domain  : {TRUE}
+output         : true
+
+
 ast            : st_endpoint(to_geometry('LINESTRING(1 1, 2 2, 3 3, 4 4)'))
 raw expr       : st_endpoint(to_geometry('LINESTRING(1 1, 2 2, 3 3, 4 4)'))
 checked expr   : st_endpoint<Geometry>(to_geometry<String>("LINESTRING(1 1, 2 2, 3 3, 4 4)"))
diff --git a/tests/cloud_control_server/simple_server.py b/tests/cloud_control_server/simple_server.py
@@ -44,9 +44,9 @@ def load_data_from_json():
                 notification_history_data = json.load(f)
                 notification_history = notification_pb2.NotificationHistory()
                 json_format.ParseDict(notification_history_data, notification_history)
-                NOTIFICATION_HISTORY_DB[notification_history.name] = (
-                    notification_history
-                )
+                NOTIFICATION_HISTORY_DB[
+                    notification_history.name
+                ] = notification_history
 
 
 def create_task_request_to_task(id, create_task_request):
diff --git a/tests/sqllogictests/suites/query/functions/02_0060_function_geometry.test b/tests/sqllogictests/suites/query/functions/02_0060_function_geometry.test
@@ -1,6 +1,19 @@
 statement ok
 DROP TABLE IF EXISTS t1
 
+query T
+SELECT ST_CONTAINS(poly, poly_inside),
+      ST_CONTAINS(poly, poly),
+      ST_CONTAINS(poly, line_on_boundary),
+      ST_CONTAINS(poly, line_inside)
+FROM (SELECT
+    TO_GEOMETRY('POLYGON((-2 0, 0 2, 2 0, -2 0))') AS poly,
+    TO_GEOMETRY('POLYGON((-1 0, 0 1, 1 0, -1 0))') AS poly_inside,
+    TO_GEOMETRY('LINESTRING(-1 1, 0 2, 1 1)') AS line_on_boundary,
+    TO_GEOMETRY('LINESTRING(-2 0, 0 0, 0 1)') AS line_inside);
+----
+1 1 0 1
+
 query T
 SELECT haversine(40.7127, -74.0059, 34.0500, -118.2500);
 ----
diff --git a/tests/udf/udf_server.py b/tests/udf/udf_server.py
@@ -162,7 +162,7 @@ def json_concat(list: List[Any]) -> Any:
     result_type="TUPLE(VARIANT NULL, VARIANT NULL)",
 )
 def tuple_access(
-        tup: Tuple[List[Any], int, str], idx1: int, idx2: int
+    tup: Tuple[List[Any], int, str], idx1: int, idx2: int
 ) -> Tuple[Any, Any]:
     v1 = None if idx1 == 0 or idx1 > len(tup) else tup[idx1 - 1]
     v2 = None if idx2 == 0 or idx2 > len(tup) else tup[idx2 - 1]
@@ -193,21 +193,21 @@ def tuple_access(
     result_type=f"TUPLE({','.join(f'{t} NULL' for t in ALL_SCALAR_TYPES)})",
 )
 def return_all(
-        bool,
-        i8,
-        i16,
-        i32,
-        i64,
-        u8,
-        u16,
-        u32,
-        u64,
-        f32,
-        f64,
-        date,
-        timestamp,
-        varchar,
-        json,
+    bool,
+    i8,
+    i16,
+    i32,
+    i64,
+    u8,
+    u16,
+    u32,
+    u64,
+    f32,
+    f64,
+    date,
+    timestamp,
+    varchar,
+    json,
 ):
     return (
         bool,
@@ -233,21 +233,21 @@ def return_all(
     result_type=f"TUPLE({','.join(f'ARRAY({t})' for t in ALL_SCALAR_TYPES)})",
 )
 def return_all_arrays(
-        bool,
-        i8,
-        i16,
-        i32,
-        i64,
-        u8,
-        u16,
-        u32,
-        u64,
-        f32,
-        f64,
-        date,
-        timestamp,
-        varchar,
-        json,
+    bool,
+    i8,
+    i16,
+    i32,
+    i64,
+    u8,
+    u16,
+    u32,
+    u64,
+    f32,
+    f64,
+    date,
+    timestamp,
+    varchar,
+    json,
 ):
     return (
         bool,
@@ -273,21 +273,21 @@ def return_all_arrays(
     result_type=f"TUPLE({','.join(f'{t}' for t in ALL_SCALAR_TYPES)})",
 )
 def return_all_non_nullable(
-        bool,
-        i8,
-        i16,
-        i32,
-        i64,
-        u8,
-        u16,
-        u32,
-        u64,
-        f32,
-        f64,
-        date,
-        timestamp,
-        varchar,
-        json,
+    bool,
+    i8,
+    i16,
+    i32,
+    i64,
+    u8,
+    u16,
+    u32,
+    u64,
+    f32,
+    f64,
+    date,
+    timestamp,
+    varchar,
+    json,
 ):
     return (
         bool,
