feat(query): implement ST_POINTN (#15373)

ST_POINTN

Signed-off-by: Fan Yang <yangfanlinux@gmail.com>

Author: kkk25641463

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

@@ -268,6 +268,74 @@ pub fn register(registry: &mut FunctionRegistry) {
         }),
     );
 
+    registry.register_passthrough_nullable_2_arg::<GeometryType, Int32Type, GeometryType, _, _>(
+        "st_pointn",
+        |_, _, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_2_arg::<GeometryType, Int32Type, GeometryType>(
+            |geometry, index, builder, ctx| {
+                if let Some(validity) = &ctx.validity {
+                    if !validity.get_bit(builder.len()) {
+                        builder.commit_row();
+                        return;
+                    }
+                }
+
+                let geo: geo_types::Geometry = match Ewkb(geometry).to_geo() {
+                    Ok(geo) => geo,
+                    Err(e) => {
+                        ctx.set_error(
+                            builder.len(),
+                            ErrorCode::GeometryError(e.to_string()).to_string(),
+                        );
+                        builder.commit_row();
+                        return;
+                    }
+                };
+
+                let point = match <geo_types::Geometry as TryInto<LineString>>::try_into(geo) {
+                    Ok(line_string) => {
+                        let len = line_string.0.len() as i32;
+                        if index >= -len && index < len && index != 0 {
+                            Point(
+                                line_string.0
+                                    [if index < 0 { len + index } else { index - 1 } as usize],
+                            )
+                        } else {
+                            ctx.set_error(
+                                builder.len(),
+                                ErrorCode::GeometryError(format!(
+                                    "Index { } is out of bounds",
+                                    index
+                                ))
+                                .to_string(),
+                            );
+                            builder.commit_row();
+                            return;
+                        }
+                    }
+                    Err(e) => {
+                        ctx.set_error(
+                            builder.len(),
+                            ErrorCode::GeometryError(e.to_string()).to_string(),
+                        );
+                        builder.commit_row();
+                        return;
+                    }
+                };
+
+                match geo_types::Geometry::from(point).to_wkb(CoordDimensions::xy()) {
+                    Ok(binary) => builder.put_slice(binary.as_slice()),
+                    Err(e) => ctx.set_error(
+                        builder.len(),
+                        ErrorCode::GeometryError(e.to_string()).to_string(),
+                    ),
+                };
+
+                builder.commit_row();
+            },
+        ),
+    );
+
     registry.register_combine_nullable_1_arg::<GeometryType, Int32Type, _, _>(
         "st_dimension",
         |_, _| FunctionDomain::MayThrow,

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

@@ -39,6 +39,7 @@ fn test_geometry() {
     test_st_makeline(file);
     test_st_makepoint(file);
     test_st_makepolygon(file);
+    test_st_pointn(file);
     test_to_geometry(file);
     test_to_string(file);
     test_try_to_geometry(file);
@@ -227,6 +228,19 @@ fn test_st_makepolygon(file: &mut impl Write) {
     )]);
 }
 
+fn test_st_pointn(file: &mut impl Write) {
+    run_ast(
+        file,
+        "ST_POINTN(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'), 1)",
+        &[],
+    );
+    run_ast(
+        file,
+        "ST_POINTN(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'), -2)",
+        &[],
+    );
+}
+
 fn test_to_geometry(file: &mut impl Write) {
     run_ast(file, "to_geometry('POINT(1820.12 890.56)')", &[]);
     run_ast(file, "to_geometry('SRID=4326;POINT(1820.12 890.56)')", &[]);

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

@@ -3546,6 +3546,8 @@ Functions overloads:
 1 st_makeline(Geometry NULL, Geometry NULL) :: Geometry NULL
 0 st_makepolygon(Geometry) :: Geometry
 1 st_makepolygon(Geometry NULL) :: Geometry NULL
+0 st_pointn(Geometry, Int32) :: Geometry
+1 st_pointn(Geometry NULL, Int32 NULL) :: Geometry NULL
 0 strcmp(String, String) :: Int8
 1 strcmp(String NULL, String NULL) :: Int8 NULL
 0 string_to_h3(String) :: UInt64

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

@@ -330,6 +330,24 @@ evaluation (internal):
 +--------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 
 
+ast            : ST_POINTN(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'), 1)
+raw expr       : ST_POINTN(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'), 1)
+checked expr   : st_pointn<Geometry, Int32>(to_geometry<String>("LINESTRING(1 1, 2 2, 3 3, 4 4)"), to_int32<UInt8>(1_u8))
+optimized expr : "POINT(1 1)"
+output type    : Geometry
+output domain  : Undefined
+output         : 'POINT(1 1)'
+
+
+ast            : ST_POINTN(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'), -2)
+raw expr       : ST_POINTN(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'), minus(2))
+checked expr   : st_pointn<Geometry, Int32>(to_geometry<String>("LINESTRING(1 1, 2 2, 3 3, 4 4)"), to_int32<Int16>(minus<UInt8>(2_u8)))
+optimized expr : "POINT(3 3)"
+output type    : Geometry
+output domain  : Undefined
+output         : 'POINT(3 3)'
+
+
 ast            : to_geometry('POINT(1820.12 890.56)')
 raw expr       : to_geometry('POINT(1820.12 890.56)')
 checked expr   : to_geometry<String>("POINT(1820.12 890.56)")

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

@@ -412,6 +412,17 @@ SELECT ST_DIMENSION(g) FROM t1;
 2
 2
 
+query T
+SELECT TO_STRING(ST_POINTN(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'), 2))
+----
+POINT(2 2)
+
+query T
+SELECT TO_STRING(ST_POINTN(TO_GEOMETRY('LINESTRING(1 1, 2 2, 3 3, 4 4)'), -2))
+----
+POINT(3 3)
+
+
 statement ok
 SET enable_geo_create_table=0