Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(query): implement ST_MAKEPOLYGON #15194

Merged
merged 1 commit into from
Apr 10, 2024
Merged

feat(query): implement ST_MAKEPOLYGON #15194

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/common/io/src/geometry.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ pub fn parse_to_subtype(buf: &[u8]) -> Result<GeometryDataType> {
Ok(GeometryDataType::EWKB)
} else {
Err(ErrorCode::GeometryError(
"Invalid geometry type format1".to_string(),
"Invalid geometry type format".to_string(),
))
}
}
Expand All @@ -113,7 +113,7 @@ pub fn parse_to_subtype(buf: &[u8]) -> Result<GeometryDataType> {
Ok(GeometryDataType::EWKB)
} else {
Err(ErrorCode::GeometryError(
"Invalid geometry type format1".to_string(),
"Invalid geometry type format".to_string(),
))
}
}
Expand Down
128 changes: 110 additions & 18 deletions src/query/functions/src/scalars/geometry.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,11 @@ use databend_common_expression::vectorize_with_builder_2_arg;
use databend_common_expression::FunctionDomain;
use databend_common_expression::FunctionRegistry;
use databend_common_io::parse_to_ewkb;
use databend_common_io::parse_to_subtype;
use databend_common_io::GeometryDataType;
use geo::MultiPoint;
use geo::Point;
use geo_types::Polygon;
use geohash::decode_bbox;
use geos::geo_types;
use geos::geo_types::Coord;
Expand All @@ -34,8 +37,10 @@ use geos::Geom;
use geos::Geometry;
use geozero::geojson::GeoJson;
use geozero::wkb::Ewkb;
use geozero::wkt::Wkt;
use geozero::CoordDimensions;
use geozero::GeozeroGeometry;
use geozero::ToGeo;
use geozero::ToGeos;
use geozero::ToJson;
use geozero::ToWkb;
Expand All @@ -46,20 +51,19 @@ use geozero::ToWkt;
pub fn register(registry: &mut FunctionRegistry) {
// aliases
registry.register_aliases("st_makegeompoint", &["st_geom_point"]);
registry.register_aliases("st_makepolygon", &["st_polygon"]);
registry.register_aliases("st_makeline", &["st_make_line"]);
registry.register_aliases("st_geometryfromwkb", &[
"st_geomfromwkb",
"st_geometryfromewkb",
"st_geomfromewkb",
"to_geometry",
]);
registry.register_aliases("st_geometryfromwkt", &[
"st_geomfromwkt",
"st_geometryfromewkt",
"st_geomfromewkt",
"st_geometryfromtext",
"st_geomfromtext",
"to_geometry",
b41sh marked this conversation as resolved.
Show resolved Hide resolved
]);

// functions
Expand All @@ -76,6 +80,7 @@ pub fn register(registry: &mut FunctionRegistry) {
if geohash.len() > 12 {
ctx.set_error(builder.len(), "");
builder.put_str("Currently the precision only implement within 12 digits!");
builder.commit_row();
return;
}

Expand All @@ -87,6 +92,7 @@ pub fn register(registry: &mut FunctionRegistry) {
ErrorCode::GeometryError(e.to_string()).to_string(),
);
builder.put_str("");
builder.commit_row();
return;
}
};
Expand Down Expand Up @@ -114,6 +120,7 @@ pub fn register(registry: &mut FunctionRegistry) {
if geohash.len() > 12 {
ctx.set_error(builder.len(), "");
builder.put_str("Currently the precision only implement within 12 digits!");
builder.commit_row();
return;
}

Expand All @@ -125,6 +132,7 @@ pub fn register(registry: &mut FunctionRegistry) {
ErrorCode::GeometryError(e.to_string()).to_string(),
);
builder.put_str("");
builder.commit_row();
return;
}
};
Expand Down Expand Up @@ -163,6 +171,76 @@ pub fn register(registry: &mut FunctionRegistry) {
})
);

registry.register_passthrough_nullable_1_arg::<GeometryType, GeometryType, _, _>(
"st_makepolygon",
|_, _| FunctionDomain::MayThrow,
vectorize_with_builder_1_arg::<GeometryType, GeometryType>(|binary, builder, ctx| {
if let Some(validity) = &ctx.validity {
if !validity.get_bit(builder.len()) {
builder.commit_row();
return;
}
}

let subtype = match parse_to_subtype(binary) {
Ok(subtype) => subtype,
Err(e) => {
ctx.set_error(
builder.len(),
ErrorCode::GeometryError(e.to_string()).to_string(),
);
builder.put_str("");
kkk25641463 marked this conversation as resolved.
Show resolved Hide resolved
builder.commit_row();
return;
}
};

let geometry = match subtype {
GeometryDataType::EWKT => Wkt(binary).to_geo(),
GeometryDataType::EWKB => Ewkb(binary).to_geo(),
GeometryDataType::GEOJSON => GeoJson(std::str::from_utf8(binary).unwrap()).to_geo(),
};

let line_string = match geometry {
Ok(geo) => geo.try_into(),
Err(e) => {
ctx.set_error(
builder.len(),
ErrorCode::GeometryError(e.to_string()).to_string(),
);
builder.put_str("");
kkk25641463 marked this conversation as resolved.
Show resolved Hide resolved
builder.commit_row();
return;
}
};

let wkt = line_string
.map_err(|e: geo_types::Error| ErrorCode::GeometryError(e.to_string()))
.and_then(|line_string: LineString| {
let points = line_string.into_points();
if points.len() < 4 {
Err(ErrorCode::GeometryError(
"Input lines must have at least 4 points!",
))
} else if points.last() != points.first() {
Err(ErrorCode::GeometryError(
"The first and last elements are not equal.",
))
} else {
let polygon = Polygon::new(LineString::from(points), vec![]);
geo_types::Geometry::from(polygon)
.to_json()
.map_err(|e| ErrorCode::GeometryError(e.to_string()))
}
});
match wkt {
Ok(data) => builder.put_slice(data.as_bytes()),
Err(e) => ctx.set_error(builder.len(), e.to_string()),
}
builder.commit_row();
}),
);

registry.register_passthrough_nullable_2_arg::<GeometryType, GeometryType, GeometryType, _, _>(
"st_makeline",
|_, _, _| FunctionDomain::Full,
Expand All @@ -180,21 +258,24 @@ pub fn register(registry: &mut FunctionRegistry) {
match binary_to_geos(params)
{
Ok(geos) => {
match get_shared_srid(&geos){
Ok(s) => {
srid = s;
geos
},
Err(e) => {
ctx.set_error(builder.len(), ErrorCode::GeometryError(e).to_string());
builder.put_str("");
return;
}
match get_shared_srid(&geos){
Ok(s) => {
srid = s;
geos
},
Err(e) => {
ctx.set_error(builder.len(), ErrorCode::GeometryError(e).to_string());
builder.put_str("");
builder.commit_row();
return;
}
}
},
Err(e) => {
ctx.set_error(builder.len(), ErrorCode::GeometryError(e.to_string()).to_string());
return builder.put_str("");
builder.put_str("");
builder.commit_row();
return;
}
};

Expand All @@ -207,7 +288,9 @@ pub fn register(registry: &mut FunctionRegistry) {
Ok(point) => point,
Err(e) => {
ctx.set_error(builder.len(), ErrorCode::GeometryError(e.to_string()).to_string());
return builder.put_str("");
builder.put_str("");
builder.commit_row();
return;
}
};
coords.push(point.into());
Expand All @@ -217,7 +300,9 @@ pub fn register(registry: &mut FunctionRegistry) {
Ok(line) => line,
Err(e) => {
ctx.set_error(builder.len(), e.to_string());
return builder.put_str("");
builder.put_str("");
builder.commit_row();
return;
}
};
coords.append(&mut line.into_inner());
Expand All @@ -227,7 +312,9 @@ pub fn register(registry: &mut FunctionRegistry) {
Ok(multi_point) => multi_point,
Err(e) => {
ctx.set_error(builder.len(), e.to_string());
return builder.put_str("");
builder.put_str("");
builder.commit_row();
return;
}
};
for point in multi_point.into_iter() {
Expand All @@ -239,7 +326,9 @@ pub fn register(registry: &mut FunctionRegistry) {
builder.len(),
ErrorCode::GeometryError("Geometry expression must be a Point, MultiPoint, or LineString.").to_string(),
);
return builder.put_str("");
builder.put_str("");
builder.commit_row();
return;
}
}
}
Expand Down Expand Up @@ -271,7 +360,9 @@ pub fn register(registry: &mut FunctionRegistry) {
builder.len(),
ErrorCode::GeometryError(e.to_string()).to_string(),
);
return builder.put_str("");
builder.put_str("");
builder.commit_row();
return;
}
};
match Ewkb(&b_ewkb).to_geos() {
Expand Down Expand Up @@ -437,6 +528,7 @@ pub fn register(registry: &mut FunctionRegistry) {
);
ctx.set_error(builder.len(), ErrorCode::GeometryError(error).to_string());
builder.put_str("");
builder.commit_row();
return;
}
},
Expand Down
33 changes: 27 additions & 6 deletions src/query/functions/tests/it/scalars/geometry.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ fn test_geometry() {
test_st_geompointfromgeohash(file);
test_st_makeline(file);
test_st_makepoint(file);
test_st_makepolygon(file);
test_to_string(file);
test_st_geometryfromwkb(file);
test_st_geometryfromwkt(file);
Expand All @@ -50,22 +51,22 @@ fn test_st_makeline(file: &mut impl Write) {
run_ast(
file,
"st_makeline(
to_geometry('SRID=4326;POINT(1.0 2.0)'),
to_geometry('SRID=4326;POINT(3.5 4.5)'))",
st_geometryfromwkt('SRID=4326;POINT(1.0 2.0)'),
st_geometryfromwkt('SRID=4326;POINT(3.5 4.5)'))",
&[],
);
run_ast(
file,
"st_makeline(
to_geometry('SRID=3857;POINT(1.0 2.0)'),
to_geometry('SRID=3857;LINESTRING(1.0 2.0, 10.1 5.5)'))",
st_geometryfromwkt('SRID=3857;POINT(1.0 2.0)'),
st_geometryfromwkt('SRID=3857;LINESTRING(1.0 2.0, 10.1 5.5)'))",
&[],
);
run_ast(
file,
"st_makeline(
to_geometry('LINESTRING(1.0 2.0, 10.1 5.5)'),
to_geometry('MULTIPOINT(3.5 4.5, 6.1 7.9)'))",
st_geometryfromwkt('LINESTRING(1.0 2.0, 10.1 5.5)'),
st_geometryfromwkt('MULTIPOINT(3.5 4.5, 6.1 7.9)'))",
&[],
);
}
Expand All @@ -79,6 +80,26 @@ fn test_st_makepoint(file: &mut impl Write) {
]);
}

fn test_st_makepolygon(file: &mut impl Write) {
run_ast(
file,
"st_makepolygon(st_geometryfromwkt('LINESTRING(0.0 0.0, 1.0 0.0, 1.0 2.0, 0.0 2.0, 0.0 0.0)'))",
&[],
);
run_ast(
file,
"st_makepolygon(st_geometryfromwkb(unhex('01020000000500000000000000000000000000000000000000000000000000f03f0000000000000000000000000000f03f00000000000000400000000000000000000000000000004000000000000000000000000000000000')))",
&[],
);
run_ast(file, "st_makepolygon(a)", &[(
"a",
StringType::from_data(vec![
"LINESTRING(0.0 0.0, 1.0 0.0, 1.0 2.0, 0.0 2.0, 0.0 0.0)",
"LINESTRING(10.1 5.2, 15.2 7.3, 20.2 8.3, 10.9 7.7, 10.1 5.2)",
]),
)]);
}

fn test_to_string(file: &mut impl Write) {
run_ast(file, "to_string(st_makegeompoint(7.0, -8.0))", &[]);
run_ast(file, "to_string(st_makegeompoint(a, b))", &[
Expand Down
4 changes: 3 additions & 1 deletion src/query/functions/tests/it/scalars/testdata/function_list.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ st_geomfromtext -> st_geometryfromwkt
st_geomfromwkb -> st_geometryfromwkb
st_geomfromwkt -> st_geometryfromwkt
st_make_line -> st_makeline
st_polygon -> st_makepolygon
str_to_date -> to_date
str_to_timestamp -> to_timestamp
str_to_year -> to_year
Expand All @@ -52,7 +53,6 @@ substring -> substr
substring_utf8 -> substr
subtract -> minus
to_datetime -> to_timestamp
to_geometry -> st_geometryfromwkt
to_text -> to_string
to_varchar -> to_string
try_ipv4_num_to_string -> try_inet_ntoa
Expand Down Expand Up @@ -3518,6 +3518,8 @@ Functions overloads:
1 st_makegeompoint(Float64 NULL, Float64 NULL) :: Geometry NULL
0 st_makeline(Geometry, Geometry) :: Geometry
1 st_makeline(Geometry NULL, Geometry NULL) :: Geometry NULL
0 st_makepolygon(Geometry) :: Geometry
1 st_makepolygon(Geometry NULL) :: Geometry NULL
0 strcmp(String, String) :: Int8
1 strcmp(String NULL, String NULL) :: Int8 NULL
0 string_to_h3(String) :: UInt64
Expand Down
Loading
Loading