Skip to content

Commit b8996f5

Browse files
kkk25641463kkk25641463
authored
feat(query): implement st_asewkb/st_aswkb/st_asewkt/st_aswkt/st_geohash (#15325)
* impl st_asewkb/st_aswkb/st_asewkt/st_aswkt/st_geohash Signed-off-by: Fan Yang <yangfanlinux@gmail.com> * impl st_asewkb/st_aswkb/st_asewkt/st_aswkt/st_geohash Signed-off-by: Fan Yang <yangfanlinux@gmail.com> --------- Signed-off-by: Fan Yang <yangfanlinux@gmail.com>
1 parent 29407ce commit b8996f5

File tree

6 files changed

+411
-5
lines changed

6 files changed

+411
-5
lines changed

‎src/query/functions/src/scalars/geometry.rs

Lines changed: 206 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ use geo::MultiPoint;
3636
use geo::Point;
3737
use geo_types::Polygon;
3838
use geohash::decode_bbox;
39+
use geohash::encode;
3940
use geos::geo_types;
4041
use geos::geo_types::Coord;
4142
use geos::geo_types::LineString;
@@ -48,11 +49,14 @@ use geozero::CoordDimensions;
4849
use geozero::ToGeo;
4950
use geozero::ToJson;
5051
use geozero::ToWkb;
52+
use geozero::ToWkt;
5153
use jsonb::parse_value;
5254
use jsonb::to_string;
5355

5456
pub fn register(registry: &mut FunctionRegistry) {
5557
// aliases
58+
registry.register_aliases("st_aswkb", &["st_asbinary"]);
59+
registry.register_aliases("st_aswkt", &["st_astext"]);
5660
registry.register_aliases("st_makegeompoint", &["st_geom_point"]);
5761
registry.register_aliases("st_makepolygon", &["st_polygon"]);
5862
registry.register_aliases("st_makeline", &["st_make_line"]);
@@ -97,6 +101,124 @@ pub fn register(registry: &mut FunctionRegistry) {
97101
}),
98102
);
99103

104+
registry.register_passthrough_nullable_1_arg::<GeometryType, BinaryType, _, _>(
105+
"st_asewkb",
106+
|_, _| FunctionDomain::MayThrow,
107+
vectorize_with_builder_1_arg::<GeometryType, BinaryType>(|geometry, builder, ctx| {
108+
if let Some(validity) = &ctx.validity {
109+
if !validity.get_bit(builder.len()) {
110+
builder.commit_row();
111+
return;
112+
}
113+
}
114+
115+
let srid = match read_ewkb_srid(&mut io::Cursor::new(&geometry)) {
116+
Ok(srid) => srid,
117+
_ => {
118+
ctx.set_error(
119+
builder.len(),
120+
ErrorCode::GeometryError("input geometry must has the correct SRID")
121+
.to_string(),
122+
);
123+
builder.commit_row();
124+
return;
125+
}
126+
};
127+
128+
match Ewkb(geometry).to_ewkb(CoordDimensions::xy(), srid) {
129+
Ok(wkb) => builder.put_slice(wkb.as_slice()),
130+
Err(e) => {
131+
ctx.set_error(
132+
builder.len(),
133+
ErrorCode::GeometryError(e.to_string()).to_string(),
134+
);
135+
}
136+
};
137+
builder.commit_row();
138+
}),
139+
);
140+
141+
registry.register_passthrough_nullable_1_arg::<GeometryType, BinaryType, _, _>(
142+
"st_aswkb",
143+
|_, _| FunctionDomain::MayThrow,
144+
vectorize_with_builder_1_arg::<GeometryType, BinaryType>(|geometry, builder, ctx| {
145+
if let Some(validity) = &ctx.validity {
146+
if !validity.get_bit(builder.len()) {
147+
builder.commit_row();
148+
return;
149+
}
150+
}
151+
152+
match Ewkb(geometry).to_wkb(CoordDimensions::xy()) {
153+
Ok(wkb) => builder.put_slice(wkb.as_slice()),
154+
Err(e) => {
155+
ctx.set_error(
156+
builder.len(),
157+
ErrorCode::GeometryError(e.to_string()).to_string(),
158+
);
159+
}
160+
};
161+
builder.commit_row();
162+
}),
163+
);
164+
165+
registry.register_passthrough_nullable_1_arg::<GeometryType, StringType, _, _>(
166+
"st_asewkt",
167+
|_, _| FunctionDomain::MayThrow,
168+
vectorize_with_builder_1_arg::<GeometryType, StringType>(|geometry, builder, ctx| {
169+
if let Some(validity) = &ctx.validity {
170+
if !validity.get_bit(builder.len()) {
171+
builder.commit_row();
172+
return;
173+
}
174+
}
175+
176+
let srid = match read_ewkb_srid(&mut io::Cursor::new(&geometry)) {
177+
Ok(srid) => srid,
178+
_ => {
179+
ctx.set_error(
180+
builder.len(),
181+
ErrorCode::GeometryError("input geometry must has the correct SRID")
182+
.to_string(),
183+
);
184+
builder.commit_row();
185+
return;
186+
}
187+
};
188+
189+
match Ewkb(geometry).to_ewkt(srid) {
190+
Ok(ewkt) => builder.put_str(&ewkt),
191+
Err(e) => ctx.set_error(
192+
builder.len(),
193+
ErrorCode::GeometryError(e.to_string()).to_string(),
194+
),
195+
};
196+
builder.commit_row();
197+
}),
198+
);
199+
200+
registry.register_passthrough_nullable_1_arg::<GeometryType, StringType, _, _>(
201+
"st_aswkt",
202+
|_, _| FunctionDomain::MayThrow,
203+
vectorize_with_builder_1_arg::<GeometryType, StringType>(|geometry, builder, ctx| {
204+
if let Some(validity) = &ctx.validity {
205+
if !validity.get_bit(builder.len()) {
206+
builder.commit_row();
207+
return;
208+
}
209+
}
210+
211+
match Ewkb(geometry).to_wkt() {
212+
Ok(wkt) => builder.put_str(&wkt),
213+
Err(e) => ctx.set_error(
214+
builder.len(),
215+
ErrorCode::GeometryError(e.to_string()).to_string(),
216+
),
217+
};
218+
builder.commit_row();
219+
}),
220+
);
221+
100222
registry.register_passthrough_nullable_1_arg::<StringType, GeometryType, _, _>(
101223
"st_geomfromgeohash",
102224
|_, _| FunctionDomain::MayThrow,
@@ -109,8 +231,10 @@ pub fn register(registry: &mut FunctionRegistry) {
109231
}
110232

111233
if geohash.len() > 12 {
112-
ctx.set_error(builder.len(), "");
113-
builder.put_str("Currently the precision only implement within 12 digits!");
234+
ctx.set_error(
235+
builder.len(),
236+
"Currently the precision only implement within 12 digits!",
237+
);
114238
builder.commit_row();
115239
return;
116240
}
@@ -354,6 +478,67 @@ pub fn register(registry: &mut FunctionRegistry) {
354478
),
355479
);
356480

481+
registry.register_passthrough_nullable_1_arg::<GeometryType, StringType, _, _>(
482+
"st_geohash",
483+
|_, _| FunctionDomain::MayThrow,
484+
vectorize_with_builder_1_arg::<GeometryType, StringType>(|geometry, builder, ctx| {
485+
if let Some(validity) = &ctx.validity {
486+
if !validity.get_bit(builder.len()) {
487+
builder.commit_row();
488+
return;
489+
}
490+
}
491+
492+
match point_to_geohash(geometry, None) {
493+
Ok(hash) => builder.put_str(&hash),
494+
Err(e) => {
495+
ctx.set_error(
496+
builder.len(),
497+
ErrorCode::GeometryError(e.to_string()).to_string(),
498+
);
499+
return;
500+
}
501+
};
502+
builder.commit_row();
503+
}),
504+
);
505+
506+
registry.register_passthrough_nullable_2_arg::<GeometryType, Int32Type, StringType, _, _>(
507+
"st_geohash",
508+
|_, _, _| FunctionDomain::MayThrow,
509+
vectorize_with_builder_2_arg::<GeometryType, Int32Type, StringType>(
510+
|geometry, precision, builder, ctx| {
511+
if let Some(validity) = &ctx.validity {
512+
if !validity.get_bit(builder.len()) {
513+
builder.commit_row();
514+
return;
515+
}
516+
}
517+
518+
if precision > 12 {
519+
ctx.set_error(
520+
builder.len(),
521+
"Currently the precision only implement within 12 digits!",
522+
);
523+
builder.commit_row();
524+
return;
525+
}
526+
527+
match point_to_geohash(geometry, Some(precision)) {
528+
Ok(hash) => builder.put_str(&hash),
529+
Err(e) => {
530+
ctx.set_error(
531+
builder.len(),
532+
ErrorCode::GeometryError(e.to_string()).to_string(),
533+
);
534+
return;
535+
}
536+
};
537+
builder.commit_row();
538+
},
539+
),
540+
);
541+
357542
registry.register_passthrough_nullable_1_arg::<StringType, GeometryType, _, _>(
358543
"st_geometryfromwkb",
359544
|_, _| FunctionDomain::MayThrow,
@@ -1080,3 +1265,22 @@ fn json_to_geometry_impl(
10801265
Err(e) => Err(ErrorCode::GeometryError(e.to_string())),
10811266
}
10821267
}
1268+
1269+
fn point_to_geohash(
1270+
geometry: &[u8],
1271+
precision: Option<i32>,
1272+
) -> databend_common_exception::Result<String> {
1273+
let point = match Ewkb(geometry).to_geo() {
1274+
Ok(geo) => Point::try_from(geo),
1275+
Err(e) => return Err(ErrorCode::GeometryError(e.to_string())),
1276+
};
1277+
1278+
let hash = match point {
1279+
Ok(point) => encode(point.0, precision.map_or(12, |p| p as usize)),
1280+
Err(e) => return Err(ErrorCode::GeometryError(e.to_string())),
1281+
};
1282+
match hash {
1283+
Ok(hash) => Ok(hash),
1284+
Err(e) => Err(ErrorCode::GeometryError(e.to_string())),
1285+
}
1286+
}

‎src/query/functions/tests/it/scalars/geometry.rs

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,11 @@ use crate::scalars::run_ast;
2626
fn test_geometry() {
2727
let mut mint = Mint::new("tests/it/scalars/testdata");
2828
let file = &mut mint.new_goldenfile("geometry.txt").unwrap();
29+
test_st_asewkb(file);
30+
test_st_aswkb(file);
31+
test_st_asewkt(file);
32+
test_st_aswkt(file);
33+
test_st_geohash(file);
2934
test_st_asgeojson(file);
3035
test_st_geomfromgeohash(file);
3136
test_st_geompointfromgeohash(file);
@@ -40,6 +45,58 @@ fn test_geometry() {
4045
// test_st_transform(file);
4146
}
4247

48+
fn test_st_asewkb(file: &mut impl Write) {
49+
run_ast(
50+
file,
51+
"st_asewkb(to_geometry('SRID=4326;POINT(-122.35 37.55)'))",
52+
&[],
53+
);
54+
run_ast(
55+
file,
56+
"st_asewkb(to_geometry('SRID=0;LINESTRING(0.75 0.75, -10 20)'))",
57+
&[],
58+
);
59+
}
60+
61+
fn test_st_aswkb(file: &mut impl Write) {
62+
run_ast(file, "st_aswkb(to_geometry('POINT(-122.35 37.55)'))", &[]);
63+
run_ast(
64+
file,
65+
"st_aswkb(to_geometry('LINESTRING(0.75 0.75, -10 20)'))",
66+
&[],
67+
);
68+
}
69+
70+
fn test_st_asewkt(file: &mut impl Write) {
71+
run_ast(
72+
file,
73+
"st_asewkt(to_geometry('SRID=4326;POINT(-122.35 37.55)'))",
74+
&[],
75+
);
76+
run_ast(
77+
file,
78+
"st_asewkt(to_geometry('SRID=0;LINESTRING(0.75 0.75, -10 20)'))",
79+
&[],
80+
);
81+
}
82+
83+
fn test_st_aswkt(file: &mut impl Write) {
84+
run_ast(file, "st_asewkt(to_geometry('POINT(-122.35 37.55)'))", &[]);
85+
run_ast(
86+
file,
87+
"st_asewkt(to_geometry('LINESTRING(0.75 0.75, -10 20)'))",
88+
&[],
89+
);
90+
}
91+
92+
fn test_st_geohash(file: &mut impl Write) {
93+
run_ast(
94+
file,
95+
"st_geohash(to_geometry('POINT(-122.306100 37.554162)', 4326))",
96+
&[],
97+
);
98+
}
99+
43100
fn test_st_asgeojson(file: &mut impl Write) {
44101
run_ast(
45102
file,

‎src/query/functions/tests/it/scalars/testdata/function_list.txt

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ remove_nullable -> assume_not_null
3434
rlike -> regexp
3535
sha1 -> sha
3636
siphash -> siphash64
37+
st_asbinary -> st_aswkb
38+
st_astext -> st_aswkt
3739
st_geom_point -> st_makegeompoint
3840
st_geometryfromewkb -> st_geometryfromwkb
3941
st_geometryfromewkt -> st_geometryfromwkt
@@ -3504,8 +3506,20 @@ Functions overloads:
35043506
17 sqrt(Float32 NULL) :: Float64 NULL
35053507
18 sqrt(Float64) :: Float64
35063508
19 sqrt(Float64 NULL) :: Float64 NULL
3509+
0 st_asewkb(Geometry) :: Binary
3510+
1 st_asewkb(Geometry NULL) :: Binary NULL
3511+
0 st_asewkt(Geometry) :: String
3512+
1 st_asewkt(Geometry NULL) :: String NULL
35073513
0 st_asgeojson(Geometry) :: Variant
35083514
1 st_asgeojson(Geometry NULL) :: Variant NULL
3515+
0 st_aswkb(Geometry) :: Binary
3516+
1 st_aswkb(Geometry NULL) :: Binary NULL
3517+
0 st_aswkt(Geometry) :: String
3518+
1 st_aswkt(Geometry NULL) :: String NULL
3519+
0 st_geohash(Geometry) :: String
3520+
1 st_geohash(Geometry NULL) :: String NULL
3521+
2 st_geohash(Geometry, Int32) :: String
3522+
3 st_geohash(Geometry NULL, Int32 NULL) :: String NULL
35093523
0 st_geometryfromwkb(String) :: Geometry
35103524
1 st_geometryfromwkb(String NULL) :: Geometry NULL
35113525
2 st_geometryfromwkb(Binary) :: Geometry

0 commit comments

Comments
 (0)