Skip to content

Commit 81ecce8

Browse files
kkk25641463kkk25641463
authored
feat(query): implement ST_X/ST_Y/ST_STARTPOINT (#15384)
* ST_X/ST_Y/ST_STARTPOINT Signed-off-by: Fan Yang <yangfanlinux@gmail.com> * ST_X/ST_Y/ST_STARTPOINT Signed-off-by: Fan Yang <yangfanlinux@gmail.com> --------- Signed-off-by: Fan Yang <yangfanlinux@gmail.com>
1 parent 7d0cd72 commit 81ecce8

File tree

5 files changed

+222
-0
lines changed

5 files changed

+222
-0
lines changed

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

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ use geozero::ToWkb;
5454
use geozero::ToWkt;
5555
use jsonb::parse_value;
5656
use jsonb::to_string;
57+
use num_traits::AsPrimitive;
5758

5859
pub fn register(registry: &mut FunctionRegistry) {
5960
// aliases
@@ -865,6 +866,122 @@ pub fn register(registry: &mut FunctionRegistry) {
865866
),
866867
);
867868

869+
registry.register_passthrough_nullable_1_arg::<GeometryType, GeometryType, _, _>(
870+
"st_startpoint",
871+
|_, _| FunctionDomain::MayThrow,
872+
vectorize_with_builder_1_arg::<GeometryType, GeometryType>(|geometry, builder, ctx| {
873+
if let Some(validity) = &ctx.validity {
874+
if !validity.get_bit(builder.len()) {
875+
builder.commit_row();
876+
return;
877+
}
878+
}
879+
880+
let geo: geo_types::Geometry = match Ewkb(geometry).to_geo() {
881+
Ok(geo) => geo,
882+
Err(e) => {
883+
ctx.set_error(
884+
builder.len(),
885+
ErrorCode::GeometryError(e.to_string()).to_string(),
886+
);
887+
builder.commit_row();
888+
return;
889+
}
890+
};
891+
892+
let point = match <geo_types::Geometry as TryInto<LineString>>::try_into(geo) {
893+
Ok(line_string) => line_string.points().next().unwrap(),
894+
Err(e) => {
895+
ctx.set_error(
896+
builder.len(),
897+
ErrorCode::GeometryError(e.to_string()).to_string(),
898+
);
899+
builder.commit_row();
900+
return;
901+
}
902+
};
903+
904+
match geo_types::Geometry::from(point).to_wkb(CoordDimensions::xy()) {
905+
Ok(binary) => builder.put_slice(binary.as_slice()),
906+
Err(e) => ctx.set_error(
907+
builder.len(),
908+
ErrorCode::GeometryError(e.to_string()).to_string(),
909+
),
910+
};
911+
builder.commit_row();
912+
}),
913+
);
914+
915+
registry.register_combine_nullable_1_arg::<GeometryType, NumberType<F64>, _, _>(
916+
"st_x",
917+
|_, _| FunctionDomain::MayThrow,
918+
vectorize_with_builder_1_arg::<GeometryType, NullableType<NumberType<F64>>>(
919+
|geometry, builder, ctx| {
920+
if let Some(validity) = &ctx.validity {
921+
if !validity.get_bit(builder.len()) {
922+
builder.push_null();
923+
return;
924+
}
925+
}
926+
927+
let geo: geo_types::Geometry = match Ewkb(geometry).to_geo() {
928+
Ok(geo) => geo,
929+
Err(e) => {
930+
ctx.set_error(
931+
builder.len(),
932+
ErrorCode::GeometryError(e.to_string()).to_string(),
933+
);
934+
builder.push_null();
935+
return;
936+
}
937+
};
938+
939+
match <geo_types::Geometry as TryInto<Point>>::try_into(geo) {
940+
Ok(point) => builder.push(F64::from(AsPrimitive::<f64>::as_(point.x()))),
941+
Err(e) => ctx.set_error(
942+
builder.len(),
943+
ErrorCode::GeometryError(e.to_string()).to_string(),
944+
),
945+
};
946+
},
947+
),
948+
);
949+
950+
registry.register_combine_nullable_1_arg::<GeometryType, NumberType<F64>, _, _>(
951+
"st_y",
952+
|_, _| FunctionDomain::MayThrow,
953+
vectorize_with_builder_1_arg::<GeometryType, NullableType<NumberType<F64>>>(
954+
|geometry, builder, ctx| {
955+
if let Some(validity) = &ctx.validity {
956+
if !validity.get_bit(builder.len()) {
957+
builder.push_null();
958+
return;
959+
}
960+
}
961+
962+
let geo: geo_types::Geometry = match Ewkb(geometry).to_geo() {
963+
Ok(geo) => geo,
964+
Err(e) => {
965+
ctx.set_error(
966+
builder.len(),
967+
ErrorCode::GeometryError(e.to_string()).to_string(),
968+
);
969+
builder.push_null();
970+
return;
971+
}
972+
};
973+
974+
match <geo_types::Geometry as TryInto<Point>>::try_into(geo) {
975+
Ok(point) => builder.push(F64::from(AsPrimitive::<f64>::as_(point.y()))),
976+
Err(e) => ctx.set_error(
977+
builder.len(),
978+
ErrorCode::GeometryError(e.to_string()).to_string(),
979+
),
980+
};
981+
},
982+
),
983+
);
984+
868985
registry.register_combine_nullable_1_arg::<GeometryType, Int32Type, _, _>(
869986
"st_srid",
870987
|_, _| FunctionDomain::MayThrow,

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,9 @@ fn test_geometry() {
4141
test_st_makepolygon(file);
4242
test_st_pointn(file);
4343
test_st_srid(file);
44+
test_st_startpoint(file);
45+
test_st_x(file);
46+
test_st_y(file);
4447
test_to_geometry(file);
4548
test_to_string(file);
4649
test_try_to_geometry(file);
@@ -253,6 +256,24 @@ fn test_st_srid(file: &mut impl Write) {
253256
run_ast(file, "st_srid(NULL)", &[]);
254257
}
255258

259+
fn test_st_startpoint(file: &mut impl Write) {
260+
run_ast(
261+
file,
262+
"st_startpoint(to_geometry('LINESTRING(1 1, 2 2, 3 3, 4 4)'))",
263+
&[],
264+
);
265+
}
266+
fn test_st_x(file: &mut impl Write) {
267+
run_ast(file, "st_x(st_makegeompoint(37.5, 45.5))", &[]);
268+
run_ast(file, "st_x(st_makegeompoint(NULL, NULL))", &[]);
269+
run_ast(file, "st_x(NULL)", &[]);
270+
}
271+
fn test_st_y(file: &mut impl Write) {
272+
run_ast(file, "st_y(st_makegeompoint(37.5, 45.5))", &[]);
273+
run_ast(file, "st_y(st_makegeompoint(NULL, NULL))", &[]);
274+
run_ast(file, "st_y(NULL)", &[]);
275+
}
276+
256277
fn test_to_geometry(file: &mut impl Write) {
257278
run_ast(file, "to_geometry('POINT(1820.12 890.56)')", &[]);
258279
run_ast(file, "to_geometry('SRID=4326;POINT(1820.12 890.56)')", &[]);

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3550,6 +3550,12 @@ Functions overloads:
35503550
1 st_pointn(Geometry NULL, Int32 NULL) :: Geometry NULL
35513551
0 st_srid(Geometry) :: Int32 NULL
35523552
1 st_srid(Geometry NULL) :: Int32 NULL
3553+
0 st_startpoint(Geometry) :: Geometry
3554+
1 st_startpoint(Geometry NULL) :: Geometry NULL
3555+
0 st_x(Geometry) :: Float64 NULL
3556+
1 st_x(Geometry NULL) :: Float64 NULL
3557+
0 st_y(Geometry) :: Float64 NULL
3558+
1 st_y(Geometry NULL) :: Float64 NULL
35533559
0 strcmp(String, String) :: Int8
35543560
1 strcmp(String NULL, String NULL) :: Int8 NULL
35553561
0 string_to_h3(String) :: UInt64

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

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -384,6 +384,69 @@ output domain : {NULL}
384384
output : NULL
385385

386386

387+
ast : st_startpoint(to_geometry('LINESTRING(1 1, 2 2, 3 3, 4 4)'))
388+
raw expr : st_startpoint(to_geometry('LINESTRING(1 1, 2 2, 3 3, 4 4)'))
389+
checked expr : st_startpoint<Geometry>(to_geometry<String>("LINESTRING(1 1, 2 2, 3 3, 4 4)"))
390+
optimized expr : "POINT(1 1)"
391+
output type : Geometry
392+
output domain : Undefined
393+
output : 'POINT(1 1)'
394+
395+
396+
ast : st_x(st_makegeompoint(37.5, 45.5))
397+
raw expr : st_x(st_makegeompoint(37.5, 45.5))
398+
checked expr : st_x<Geometry>(st_makegeompoint<Float64, Float64>(to_float64<Decimal(3, 1)>(37.5_d128(3,1)), to_float64<Decimal(3, 1)>(45.5_d128(3,1))))
399+
optimized expr : 37.5_f64
400+
output type : Float64 NULL
401+
output domain : {37.5..=37.5}
402+
output : 37.5
403+
404+
405+
ast : st_x(st_makegeompoint(NULL, NULL))
406+
raw expr : st_x(st_makegeompoint(NULL, NULL))
407+
checked expr : st_x<Geometry NULL>(st_makegeompoint<Float64 NULL, Float64 NULL>(CAST(NULL AS Float64 NULL), CAST(NULL AS Float64 NULL)))
408+
optimized expr : NULL
409+
output type : Float64 NULL
410+
output domain : {NULL}
411+
output : NULL
412+
413+
414+
ast : st_x(NULL)
415+
raw expr : st_x(NULL)
416+
checked expr : st_x<Geometry NULL>(CAST(NULL AS Geometry NULL))
417+
optimized expr : NULL
418+
output type : Float64 NULL
419+
output domain : {NULL}
420+
output : NULL
421+
422+
423+
ast : st_y(st_makegeompoint(37.5, 45.5))
424+
raw expr : st_y(st_makegeompoint(37.5, 45.5))
425+
checked expr : st_y<Geometry>(st_makegeompoint<Float64, Float64>(to_float64<Decimal(3, 1)>(37.5_d128(3,1)), to_float64<Decimal(3, 1)>(45.5_d128(3,1))))
426+
optimized expr : 45.5_f64
427+
output type : Float64 NULL
428+
output domain : {45.5..=45.5}
429+
output : 45.5
430+
431+
432+
ast : st_y(st_makegeompoint(NULL, NULL))
433+
raw expr : st_y(st_makegeompoint(NULL, NULL))
434+
checked expr : st_y<Geometry NULL>(st_makegeompoint<Float64 NULL, Float64 NULL>(CAST(NULL AS Float64 NULL), CAST(NULL AS Float64 NULL)))
435+
optimized expr : NULL
436+
output type : Float64 NULL
437+
output domain : {NULL}
438+
output : NULL
439+
440+
441+
ast : st_y(NULL)
442+
raw expr : st_y(NULL)
443+
checked expr : st_y<Geometry NULL>(CAST(NULL AS Geometry NULL))
444+
optimized expr : NULL
445+
output type : Float64 NULL
446+
output domain : {NULL}
447+
output : NULL
448+
449+
387450
ast : to_geometry('POINT(1820.12 890.56)')
388451
raw expr : to_geometry('POINT(1820.12 890.56)')
389452
checked expr : to_geometry<String>("POINT(1820.12 890.56)")

‎tests/sqllogictests/suites/query/functions/02_0060_function_geometry.test

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -443,6 +443,21 @@ SELECT ST_SRID(NULL)
443443
----
444444
NULL
445445

446+
query T
447+
SELECT ST_STARTPOINT(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'))
448+
----
449+
POINT(1 1)
450+
451+
query T
452+
SELECT ST_X(ST_MAKEGEOMPOINT(37.5, 45.5)), ST_Y(ST_MAKEGEOMPOINT(37.5, 45.5))
453+
----
454+
37.5 45.5
455+
456+
query T
457+
SELECT ST_X(ST_MAKEGEOMPOINT(NULL, NULL)), ST_Y(ST_MAKEGEOMPOINT(NULL, NULL)),ST_X(NULL),ST_Y(NULL)
458+
----
459+
NULL NULL NULL NULL
460+
446461
statement ok
447462
SET enable_geo_create_table=0
448463

0 commit comments

Comments
 (0)