Skip to content

Commit

Permalink
feat(query): Support geometry relation functions (#16927)
Browse files Browse the repository at this point in the history 
* feat(query): Support geometry relation functions

* fix

* fix
  • Loading branch information
@b41sh
b41sh authored Nov 28, 2024
1 parent 5b4e61a commit fb89467
Show file tree
Hide file tree
Showing 5 changed files with 404 additions and 46 deletions.
206 changes: 160 additions & 46 deletions src/query/functions/src/scalars/geometry.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ use databend_common_expression::vectorize_with_builder_1_arg;
use databend_common_expression::vectorize_with_builder_2_arg;
use databend_common_expression::vectorize_with_builder_3_arg;
use databend_common_expression::vectorize_with_builder_4_arg;
use databend_common_expression::EvalContext;
use databend_common_expression::FunctionDomain;
use databend_common_expression::FunctionRegistry;
use databend_common_io::ewkb_to_geo;
Expand All @@ -53,6 +54,7 @@ use geo::EuclideanLength;
use geo::Geometry;
use geo::HasDimensions;
use geo::HaversineDistance;
use geo::Intersects;
use geo::Line;
use geo::LineString;
use geo::MultiLineString;
Expand All @@ -63,6 +65,7 @@ use geo::Rect;
use geo::ToDegrees;
use geo::ToRadians;
use geo::Triangle;
use geo::Within;
use geohash::decode_bbox;
use geohash::encode;
use geozero::geojson::GeoJson;
Expand Down Expand Up @@ -255,29 +258,144 @@ pub fn register(registry: &mut FunctionRegistry) {
ewkb_to_geo(&mut Ewkb(r_ewkb)),
) {
(Ok((l_geo, l_srid)), Ok((r_geo, r_srid))) => {
if !l_srid.eq(&r_srid) {
ctx.set_error(
builder.len(),
format!(
"Incompatible SRID: {} and {}",
l_srid.unwrap_or_default(),
r_srid.unwrap_or_default()
),
);
if !check_incompatible_srid(l_srid, r_srid, builder.len(), ctx) {
builder.push(false);
return;
}
if matches!(l_geo, Geometry::GeometryCollection(_))
|| matches!(r_geo, Geometry::GeometryCollection(_))
{
ctx.set_error(
builder.len(),
"A GEOMETRY object that is a GeometryCollection".to_string(),
);
let is_contains = l_geo.contains(&r_geo);
builder.push(is_contains);
}
(Err(e), _) | (_, Err(e)) => {
ctx.set_error(builder.len(), e.to_string());
builder.push(false);
}
}
},
),
);

registry.register_passthrough_nullable_2_arg::<GeometryType, GeometryType, BooleanType, _, _>(
"st_intersects",
|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<GeometryType, GeometryType, BooleanType>(
|l_ewkb, r_ewkb, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(false);
return;
}
}

match (
ewkb_to_geo(&mut Ewkb(l_ewkb)),
ewkb_to_geo(&mut Ewkb(r_ewkb)),
) {
(Ok((l_geo, l_srid)), Ok((r_geo, r_srid))) => {
if !check_incompatible_srid(l_srid, r_srid, builder.len(), ctx) {
builder.push(false);
} else {
builder.push(l_geo.contains(&r_geo));
return;
}
let is_intersects = l_geo.intersects(&r_geo);
builder.push(is_intersects);
}
(Err(e), _) | (_, Err(e)) => {
ctx.set_error(builder.len(), e.to_string());
builder.push(false);
}
}
},
),
);

registry.register_passthrough_nullable_2_arg::<GeometryType, GeometryType, BooleanType, _, _>(
"st_disjoint",
|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<GeometryType, GeometryType, BooleanType>(
|l_ewkb, r_ewkb, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(false);
return;
}
}

match (
ewkb_to_geo(&mut Ewkb(l_ewkb)),
ewkb_to_geo(&mut Ewkb(r_ewkb)),
) {
(Ok((l_geo, l_srid)), Ok((r_geo, r_srid))) => {
if !check_incompatible_srid(l_srid, r_srid, builder.len(), ctx) {
builder.push(false);
return;
}
let is_disjoint = !l_geo.intersects(&r_geo);
builder.push(is_disjoint);
}
(Err(e), _) | (_, Err(e)) => {
ctx.set_error(builder.len(), e.to_string());
builder.push(false);
}
}
},
),
);

registry.register_passthrough_nullable_2_arg::<GeometryType, GeometryType, BooleanType, _, _>(
"st_within",
|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<GeometryType, GeometryType, BooleanType>(
|l_ewkb, r_ewkb, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(false);
return;
}
}

match (
ewkb_to_geo(&mut Ewkb(l_ewkb)),
ewkb_to_geo(&mut Ewkb(r_ewkb)),
) {
(Ok((l_geo, l_srid)), Ok((r_geo, r_srid))) => {
if !check_incompatible_srid(l_srid, r_srid, builder.len(), ctx) {
builder.push(false);
return;
}
let is_within = l_geo.is_within(&r_geo);
builder.push(is_within);
}
(Err(e), _) | (_, Err(e)) => {
ctx.set_error(builder.len(), e.to_string());
builder.push(false);
}
}
},
),
);

registry.register_passthrough_nullable_2_arg::<GeometryType, GeometryType, BooleanType, _, _>(
"st_equals",
|_, _, _| FunctionDomain::MayThrow,
vectorize_with_builder_2_arg::<GeometryType, GeometryType, BooleanType>(
|l_ewkb, r_ewkb, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.push(false);
return;
}
}

match (
ewkb_to_geo(&mut Ewkb(l_ewkb)),
ewkb_to_geo(&mut Ewkb(r_ewkb)),
) {
(Ok((l_geo, l_srid)), Ok((r_geo, r_srid))) => {
if !check_incompatible_srid(l_srid, r_srid, builder.len(), ctx) {
builder.push(false);
return;
}
let is_equal = l_geo.is_within(&r_geo) && r_geo.is_within(&l_geo);
builder.push(is_equal);
}
(Err(e), _) | (_, Err(e)) => {
ctx.set_error(builder.len(), e.to_string());
Expand Down Expand Up @@ -306,32 +424,11 @@ pub fn register(registry: &mut FunctionRegistry) {
ewkb_to_geo(&mut Ewkb(r_ewkb)),
) {
(Ok((l_geo, l_srid)), Ok((r_geo, r_srid))) => {
if !l_srid.eq(&r_srid) {
ctx.set_error(
builder.len(),
format!(
"Incompatible SRID: {} and {}",
l_srid.unwrap_or_default(),
r_srid.unwrap_or_default()
),
);
if !check_incompatible_srid(l_srid, r_srid, builder.len(), ctx) {
builder.push(F64::from(0_f64));
return;
}

let distance = match l_geo {
Geometry::Point(l) => l.euclidean_distance(&r_geo),
Geometry::Line(l) => l.euclidean_distance(&r_geo),
Geometry::LineString(l) => l.euclidean_distance(&r_geo),
Geometry::Polygon(l) => l.euclidean_distance(&r_geo),
Geometry::MultiPoint(l) => l.euclidean_distance(&r_geo),
Geometry::MultiLineString(l) => l.euclidean_distance(&r_geo),
Geometry::MultiPolygon(l) => l.euclidean_distance(&r_geo),
Geometry::GeometryCollection(l) => l.euclidean_distance(&r_geo),
Geometry::Rect(l) => l.euclidean_distance(&r_geo),
Geometry::Triangle(l) => l.euclidean_distance(&r_geo),
};

let distance = l_geo.euclidean_distance(&r_geo);
let distance =
(distance * 1_000_000_000_f64).round() / 1_000_000_000_f64;
builder.push(distance.into());
Expand Down Expand Up @@ -656,11 +753,7 @@ pub fn register(registry: &mut FunctionRegistry) {

match (ewkb_to_geo(&mut Ewkb(l_ewkb)), ewkb_to_geo(&mut Ewkb(r_ewkb))) {
(Ok((l_geo, l_srid)), Ok((r_geo, r_srid))) => {
if !l_srid.eq(&r_srid) {
ctx.set_error(
builder.len(),
format!("Incompatible SRID: {} and {}", l_srid.unwrap_or_default(), r_srid.unwrap_or_default())
);
if !check_incompatible_srid(l_srid, r_srid, builder.len(), ctx) {
builder.commit_row();
return;
}
Expand Down Expand Up @@ -1887,3 +1980,24 @@ fn round_geometry_coordinates(geom: Geometry<f64>) -> Geometry<f64> {
)),
}
}

fn check_incompatible_srid(
l_srid: Option<i32>,
r_srid: Option<i32>,
len: usize,
ctx: &mut EvalContext,
) -> bool {
if !l_srid.eq(&r_srid) {
ctx.set_error(
len,
format!(
"Incompatible SRID: {} and {}",
l_srid.unwrap_or_default(),
r_srid.unwrap_or_default()
),
);
false
} else {
true
}
}
81 changes: 81 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 @@ -60,6 +60,10 @@ fn test_geometry() {
test_st_ymax(file);
test_st_ymin(file);
test_st_transform(file);
test_st_intersects(file);
test_st_disjoint(file);
test_st_within(file);
test_st_equals(file);
}

fn test_haversine(file: &mut impl Write) {
Expand Down Expand Up @@ -679,3 +683,80 @@ fn test_st_transform(file: &mut impl Write) {
("c", Int32Type::from_data(vec![28992])),
]);
}

fn test_st_intersects(file: &mut impl Write) {
run_ast(
file,
"ST_INTERSECTS(TO_GEOMETRY('POINT(0 0)'), TO_GEOMETRY('LINESTRING(2 0, 0 2)'))",
&[],
);
run_ast(
file,
"ST_INTERSECTS(TO_GEOMETRY('POINT(0 0)'), TO_GEOMETRY('LINESTRING(0 0, 0 2)'))",
&[],
);
run_ast(
file,
"ST_INTERSECTS(TO_GEOMETRY('POLYGON((0 0, 0 2, 2 2, 2 0, 0 0))'), TO_GEOMETRY('POLYGON((1 1, 3 1, 3 3, 1 3, 1 1))'))",
&[],
);
}

fn test_st_disjoint(file: &mut impl Write) {
run_ast(
file,
"ST_DISJOINT(TO_GEOMETRY('POINT(0 0)'), TO_GEOMETRY('LINESTRING(2 0, 0 2)'))",
&[],
);
run_ast(
file,
"ST_DISJOINT(TO_GEOMETRY('POINT(0 0)'), TO_GEOMETRY('LINESTRING(0 0, 0 2)'))",
&[],
);
run_ast(
file,
"ST_DISJOINT(TO_GEOMETRY('POLYGON((0 0, 0 2, 2 2, 2 0, 0 0))'), TO_GEOMETRY('POLYGON((1 1, 3 1, 3 3, 1 3, 1 1))'))",
&[],
);
}

fn test_st_within(file: &mut impl Write) {
run_ast(
file,
"ST_WITHIN(TO_GEOMETRY('POINT(1 2)'), TO_GEOMETRY('LINESTRING(0 0, 2 4)'))",
&[],
);
run_ast(
file,
"ST_WITHIN(TO_GEOMETRY('POINT(10 20)'), TO_GEOMETRY('POLYGON((0 0, 0 40, 40 40, 40 0, 0 0))'))",
&[],
);
run_ast(
file,
"ST_WITHIN(TO_GEOMETRY('POLYGON((0 0, 0 40, 40 40, 40 0, 0 0))'), TO_GEOMETRY('POINT(10 20)'))",
&[],
);
}

fn test_st_equals(file: &mut impl Write) {
run_ast(
file,
"ST_EQUALS(TO_GEOMETRY('LINESTRING(0 0, 10 10)'), TO_GEOMETRY('LINESTRING(0 0, 5 5, 10 10)'))",
&[],
);
run_ast(
file,
"ST_EQUALS(TO_GEOMETRY('LINESTRING(0 0, 5 5, 10 10)'), TO_GEOMETRY('LINESTRING(0 0, 10 10)'))",
&[],
);
run_ast(
file,
"ST_EQUALS(TO_GEOMETRY('LINESTRING(10 10, 1 2)'), TO_GEOMETRY('LINESTRING(0 0, 5 5, 10 10)'))",
&[],
);
run_ast(
file,
"ST_EQUALS(TO_GEOMETRY('POINT(10 10)'), TO_GEOMETRY('LINESTRING(10 10, 10 10)'))",
&[],
);
}
8 changes: 8 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 @@ -3639,10 +3639,14 @@ Functions overloads:
1 st_contains(Geometry NULL, Geometry NULL) :: Boolean NULL
0 st_dimension(Geometry) :: Int32 NULL
1 st_dimension(Geometry NULL) :: Int32 NULL
0 st_disjoint(Geometry, Geometry) :: Boolean
1 st_disjoint(Geometry NULL, Geometry NULL) :: Boolean NULL
0 st_distance(Geometry, Geometry) :: Float64
1 st_distance(Geometry NULL, Geometry NULL) :: Float64 NULL
0 st_endpoint(Geometry) :: Geometry NULL
1 st_endpoint(Geometry NULL) :: Geometry NULL
0 st_equals(Geometry, Geometry) :: Boolean
1 st_equals(Geometry NULL, Geometry NULL) :: Boolean NULL
0 st_geographyfromewkt(String) :: Geography
1 st_geographyfromewkt(String NULL) :: Geography NULL
0 st_geohash(Geometry) :: String
Expand All @@ -3665,6 +3669,8 @@ Functions overloads:
1 st_geomfromgeohash(String NULL) :: Geometry NULL
0 st_geompointfromgeohash(String) :: Geometry
1 st_geompointfromgeohash(String NULL) :: Geometry NULL
0 st_intersects(Geometry, Geometry) :: Boolean
1 st_intersects(Geometry NULL, Geometry NULL) :: Boolean NULL
0 st_length(Geometry) :: Float64
1 st_length(Geometry NULL) :: Float64 NULL
0 st_makegeompoint(Float64, Float64) :: Geometry
Expand All @@ -3689,6 +3695,8 @@ Functions overloads:
1 st_transform(Geometry NULL, Int32 NULL) :: Geometry NULL
2 st_transform(Geometry, Int32, Int32) :: Geometry
3 st_transform(Geometry NULL, Int32 NULL, Int32 NULL) :: Geometry NULL
0 st_within(Geometry, Geometry) :: Boolean
1 st_within(Geometry NULL, Geometry NULL) :: Boolean NULL
0 st_x(Geometry) :: Float64
1 st_x(Geometry NULL) :: Float64 NULL
0 st_xmax(Geometry) :: Float64 NULL
Expand Down
Loading

0 comments on commit fb89467

Please sign in to comment.