Skip to content

Commit

Permalink
feat(query): ST_CONTAINS (#15994)
Browse files Browse the repository at this point in the history 
* ST_CONTAINS

Signed-off-by: Fan Yang <[email protected]>

* ST_CONTAINS

Signed-off-by: Fan Yang <[email protected]>

---------

Signed-off-by: Fan Yang <[email protected]>
  • Loading branch information
@kkk25641463
kkk25641463 authored Jul 12, 2024
1 parent 2338e00 commit 0e55fa9
Show file tree
Hide file tree
Showing 9 changed files with 173 additions and 57 deletions.
6 changes: 3 additions & 3 deletions scripts/benchmark/query/load/tpch10.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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}";
Expand Down
10 changes: 5 additions & 5 deletions scripts/benchmark/query/load/tpch100.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand All @@ -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}";
Expand Down
41 changes: 41 additions & 0 deletions src/query/functions/src/scalars/geometry.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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,
Expand Down
24 changes: 24 additions & 0 deletions src/query/functions/tests/it/scalars/geometry.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions src/query/functions/tests/it/scalars/testdata/function_list.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
36 changes: 36 additions & 0 deletions src/query/functions/tests/it/scalars/testdata/geometry.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)"))
Expand Down
6 changes: 3 additions & 3 deletions tests/cloud_control_server/simple_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
13 changes: 13 additions & 0 deletions tests/sqllogictests/suites/query/functions/02_0060_function_geometry.test
View file
Original file line number Diff line number Diff line change
@@ -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);
----
Expand Down
92 changes: 46 additions & 46 deletions tests/udf/udf_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down

0 comments on commit 0e55fa9

Please sign in to comment.