feat(query): implement st_geometryfromewkb function (#15118)

* impl st_geometryfromewkb

* impl st_geometryfromewkb

Author: kkk25641463

src/query/functions/Cargo.toml

@@ -37,7 +37,7 @@ geo = { workspace = true }
 geo-types = "0.7.13"
 geohash = "0.13.0"
 geos = { version = "8.3", features = ["static", "geo", "geo-types"] }
-geozero = { workspace = true }
+geozero = { workspace = true, features = ["with-geos"] }
 h3o = "0.4.0"
 hex = "0.4.3"
 itertools = { workspace = true }

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

@@ -14,6 +14,7 @@
 
 use databend_common_exception::ErrorCode;
 use databend_common_expression::types::geometry::GeometryType;
+use databend_common_expression::types::BinaryType;
 use databend_common_expression::types::Int32Type;
 use databend_common_expression::types::NumberType;
 use databend_common_expression::types::StringType;
@@ -32,6 +33,8 @@ use geos::Geom;
 use geos::Geometry;
 use geozero::wkb::Ewkb;
 use geozero::CoordDimensions;
+use geozero::GeozeroGeometry;
+use geozero::ToGeos;
 use geozero::ToWkb;
 use geozero::ToWkt;
 
@@ -41,6 +44,12 @@ pub fn register(registry: &mut FunctionRegistry) {
     // aliases
     registry.register_aliases("st_makegeompoint", &["st_geom_point"]);
     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",
@@ -164,6 +173,130 @@ pub fn register(registry: &mut FunctionRegistry) {
         ),
     );
 
+    registry.register_passthrough_nullable_1_arg::<StringType, GeometryType, _, _>(
+        "st_geometryfromwkb",
+        |_, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_1_arg::<StringType, GeometryType>(|ewkb, builder, ctx| {
+            if let Some(validity) = &ctx.validity {
+                if !validity.get_bit(builder.len()) {
+                    builder.commit_row();
+                    return;
+                }
+            }
+
+            let b_ewkb = match hex::decode(ewkb) {
+                Ok(b) => b,
+                Err(e) => {
+                    ctx.set_error(
+                        builder.len(),
+                        ErrorCode::GeometryError(e.to_string()).to_string(),
+                    );
+                    return builder.put_str("");
+                }
+            };
+            match Ewkb(&b_ewkb).to_geos() {
+                Ok(_) => builder.put_slice(b_ewkb.as_slice()),
+                Err(e) => ctx.set_error(builder.len(), e.to_string()),
+            }
+            builder.commit_row();
+        }),
+    );
+
+    registry.register_passthrough_nullable_1_arg::<BinaryType, GeometryType, _, _>(
+        "st_geometryfromwkb",
+        |_, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_1_arg::<BinaryType, GeometryType>(|ewkb, builder, ctx| {
+            if let Some(validity) = &ctx.validity {
+                if !validity.get_bit(builder.len()) {
+                    builder.commit_row();
+                    return;
+                }
+            }
+
+            match Ewkb(&ewkb).to_geos() {
+                Ok(_) => builder.put_slice(ewkb),
+                Err(e) => ctx.set_error(builder.len(), e.to_string()),
+            }
+            builder.commit_row();
+        }),
+    );
+
+    registry.register_passthrough_nullable_2_arg::<StringType, Int32Type, GeometryType, _, _>(
+        "st_geometryfromwkb",
+        |_, _, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_2_arg::<StringType, Int32Type, GeometryType>(
+            |ewkb, srid, builder, ctx| {
+                if let Some(validity) = &ctx.validity {
+                    if !validity.get_bit(builder.len()) {
+                        builder.commit_row();
+                        return;
+                    }
+                }
+
+                let result = {
+                    hex::decode(ewkb)
+                        .map_err(|e| ErrorCode::from(e.to_string()))
+                        .and_then(|binary| {
+                            Ewkb(binary)
+                                .to_geos()
+                                .map_err(|e| ErrorCode::GeometryError(e.to_string()))
+                                .and_then(|geos| {
+                                    geos.to_ewkb(geos.dims(), Some(srid))
+                                        .map_err(|e| ErrorCode::GeometryError(e.to_string()))
+                                })
+                        })
+                };
+
+                match result {
+                    Ok(data) => {
+                        builder.put_slice(data.as_slice());
+                    }
+                    Err(e) => {
+                        ctx.set_error(builder.len(), e.to_string());
+                    }
+                }
+
+                builder.commit_row();
+            },
+        ),
+    );
+
+    registry.register_passthrough_nullable_2_arg::<BinaryType, Int32Type, GeometryType, _, _>(
+        "st_geometryfromwkb",
+        |_, _, _| FunctionDomain::MayThrow,
+        vectorize_with_builder_2_arg::<BinaryType, Int32Type, GeometryType>(
+            |ewkb, srid, builder, ctx| {
+                if let Some(validity) = &ctx.validity {
+                    if !validity.get_bit(builder.len()) {
+                        builder.commit_row();
+                        return;
+                    }
+                }
+
+                let result = {
+                    Ewkb(ewkb)
+                        .to_geos()
+                        .map_err(|e| ErrorCode::GeometryError(e.to_string()))
+                        .and_then(|geos| {
+                            geos.to_ewkb(geos.dims(), Some(srid))
+                                .map_err(|e| ErrorCode::GeometryError(e.to_string()))
+                        })
+                };
+
+                match result {
+                    Ok(data) => {
+                        builder.put_slice(data.as_slice());
+                    }
+                    Err(e) => {
+                        ctx.set_error(builder.len(), e.to_string());
+                    }
+                }
+
+                builder.commit_row();
+            },
+        ),
+    );
+
     registry.register_passthrough_nullable_1_arg::<StringType, GeometryType, _, _>(
         "st_geometryfromwkt",
         |_, _| FunctionDomain::MayThrow,

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

@@ -29,6 +29,7 @@ fn test_geometry() {
     test_st_makeline(file);
     test_st_makepoint(file);
     test_to_string(file);
+    test_st_geometryfromwkb(file);
     test_st_geometryfromwkt(file);
     // test_st_transform(file);
 }
@@ -74,6 +75,44 @@ fn test_to_string(file: &mut impl Write) {
     ]);
 }
 
+fn test_st_geometryfromwkb(file: &mut impl Write) {
+    run_ast(
+        file,
+        "st_geometryfromwkb('0101000020797f000066666666a9cb17411f85ebc19e325641')",
+        &[],
+    );
+
+    run_ast(
+        file,
+        "st_geometryfromwkb(unhex('0101000020797f000066666666a9cb17411f85ebc19e325641'))",
+        &[],
+    );
+
+    run_ast(
+        file,
+        "st_geometryfromwkb('0101000020797f000066666666a9cb17411f85ebc19e325641', 4326)",
+        &[],
+    );
+
+    run_ast(
+        file,
+        "st_geometryfromwkb(unhex('0101000020797f000066666666a9cb17411f85ebc19e325641'), 4326)",
+        &[],
+    );
+
+    run_ast(file, "st_geometryfromwkb(a, b)", &[
+        (
+            "a",
+            StringType::from_data(vec![
+                "0101000020797f000066666666a9cb17411f85ebc19e325641",
+                "0101000020797f000066666666a9cb17411f85ebc19e325641",
+                "0101000020797f000066666666a9cb17411f85ebc19e325641",
+            ]),
+        ),
+        ("b", Int32Type::from_data(vec![32633, 4326, 3857])),
+    ]);
+}
+
 fn test_st_geometryfromwkt(file: &mut impl Write) {
     // without srid
     run_ast(

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

@@ -35,10 +35,13 @@ rlike -> regexp
 sha1 -> sha
 siphash -> siphash64
 st_geom_point -> st_makegeompoint
+st_geometryfromewkb -> st_geometryfromwkb
 st_geometryfromewkt -> st_geometryfromwkt
 st_geometryfromtext -> st_geometryfromwkt
+st_geomfromewkb -> st_geometryfromwkb
 st_geomfromewkt -> st_geometryfromwkt
 st_geomfromtext -> st_geometryfromwkt
+st_geomfromwkb -> st_geometryfromwkb
 st_geomfromwkt -> st_geometryfromwkt
 st_make_line -> st_makeline
 str_to_date -> to_date
@@ -3495,6 +3498,14 @@ Functions overloads:
 17 sqrt(Float32 NULL) :: Float64 NULL
 18 sqrt(Float64) :: Float64
 19 sqrt(Float64 NULL) :: Float64 NULL
+0 st_geometryfromwkb(String) :: Geometry
+1 st_geometryfromwkb(String NULL) :: Geometry NULL
+2 st_geometryfromwkb(Binary) :: Geometry
+3 st_geometryfromwkb(Binary NULL) :: Geometry NULL
+4 st_geometryfromwkb(String, Int32) :: Geometry
+5 st_geometryfromwkb(String NULL, Int32 NULL) :: Geometry NULL
+6 st_geometryfromwkb(Binary, Int32) :: Geometry
+7 st_geometryfromwkb(Binary NULL, Int32 NULL) :: Geometry NULL
 0 st_geometryfromwkt(String) :: Geometry
 1 st_geometryfromwkt(String NULL) :: Geometry NULL
 2 st_geometryfromwkt(String, Int32) :: Geometry

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

@@ -104,6 +104,66 @@ evaluation (internal):
 +--------+-----------------------------------------------------------------------------------------------------------------+
 
 
+ast            : st_geometryfromwkb('0101000020797f000066666666a9cb17411f85ebc19e325641')
+raw expr       : st_geometryfromwkb('0101000020797f000066666666a9cb17411f85ebc19e325641')
+checked expr   : st_geometryfromwkb<String>("0101000020797f000066666666a9cb17411f85ebc19e325641")
+optimized expr : "SRID=32633;POINT(389866.35 5819003.03)"
+output type    : Geometry
+output domain  : Undefined
+output         : 'SRID=32633;POINT(389866.35 5819003.03)'
+
+
+ast            : st_geometryfromwkb(unhex('0101000020797f000066666666a9cb17411f85ebc19e325641'))
+raw expr       : st_geometryfromwkb(unhex('0101000020797f000066666666a9cb17411f85ebc19e325641'))
+checked expr   : st_geometryfromwkb<Binary>(from_hex<String>("0101000020797f000066666666a9cb17411f85ebc19e325641"))
+optimized expr : "SRID=32633;POINT(389866.35 5819003.03)"
+output type    : Geometry
+output domain  : Undefined
+output         : 'SRID=32633;POINT(389866.35 5819003.03)'
+
+
+ast            : st_geometryfromwkb('0101000020797f000066666666a9cb17411f85ebc19e325641', 4326)
+raw expr       : st_geometryfromwkb('0101000020797f000066666666a9cb17411f85ebc19e325641', 4326)
+checked expr   : st_geometryfromwkb<String, Int32>("0101000020797f000066666666a9cb17411f85ebc19e325641", to_int32<UInt16>(4326_u16))
+optimized expr : "SRID=4326;POINT(389866.35 5819003.03)"
+output type    : Geometry
+output domain  : Undefined
+output         : 'SRID=4326;POINT(389866.35 5819003.03)'
+
+
+ast            : st_geometryfromwkb(unhex('0101000020797f000066666666a9cb17411f85ebc19e325641'), 4326)
+raw expr       : st_geometryfromwkb(unhex('0101000020797f000066666666a9cb17411f85ebc19e325641'), 4326)
+checked expr   : st_geometryfromwkb<Binary, Int32>(from_hex<String>("0101000020797f000066666666a9cb17411f85ebc19e325641"), to_int32<UInt16>(4326_u16))
+optimized expr : "SRID=4326;POINT(389866.35 5819003.03)"
+output type    : Geometry
+output domain  : Undefined
+output         : 'SRID=4326;POINT(389866.35 5819003.03)'
+
+
+ast            : st_geometryfromwkb(a, b)
+raw expr       : st_geometryfromwkb(a::String, b::Int32)
+checked expr   : st_geometryfromwkb<String, Int32>(a, b)
+optimized expr : st_geometryfromwkb<String, Int32>("0101000020797f000066666666a9cb17411f85ebc19e325641", b)
+evaluation:
++--------+---------------------------------------------------------------------------------------------------------------+----------------+------------------------------------------+
+|        | a                                                                                                             | b              | Output                                   |
++--------+---------------------------------------------------------------------------------------------------------------+----------------+------------------------------------------+
+| Type   | String                                                                                                        | Int32          | Geometry                                 |
+| Domain | {"0101000020797f000066666666a9cb17411f85ebc19e325641"..="0101000020797f000066666666a9cb17411f85ebc19e325641"} | {3857..=32633} | Unknown                                  |
+| Row 0  | '0101000020797f000066666666a9cb17411f85ebc19e325641'                                                          | 32633          | 'SRID=32633;POINT(389866.35 5819003.03)' |
+| Row 1  | '0101000020797f000066666666a9cb17411f85ebc19e325641'                                                          | 4326           | 'SRID=4326;POINT(389866.35 5819003.03)'  |
+| Row 2  | '0101000020797f000066666666a9cb17411f85ebc19e325641'                                                          | 3857           | 'SRID=3857;POINT(389866.35 5819003.03)'  |
++--------+---------------------------------------------------------------------------------------------------------------+----------------+------------------------------------------+
+evaluation (internal):
++--------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| Column | Data                                                                                                                                                                                                                                                                                                                                                              |
++--------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| a      | StringColumn { data: 0x303130313030303032303739376630303030363636363636363661396362313734313166383565626331396533323536343130313031303030303230373937663030303036363636363636366139636231373431316638356562633139653332353634313031303130303030323037393766303030303636363636363636613963623137343131663835656263313965333235363431, offsets: [0, 50, 100, 150] } |
+| b      | Int32([32633, 4326, 3857])                                                                                                                                                                                                                                                                                                                                        |
+| Output | BinaryColumn { data: 0x0101000020797f000066666666a9cb17411f85ebc19e3256410101000020e610000066666666a9cb17411f85ebc19e3256410101000020110f000066666666a9cb17411f85ebc19e325641, offsets: [0, 25, 50, 75] }                                                                                                                                                         |
++--------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+
+
 ast            : st_geometryfromwkt('POINT(389866.35 5819003.03)')
 raw expr       : st_geometryfromwkt('POINT(389866.35 5819003.03)')
 checked expr   : st_geometryfromwkt<String>("POINT(389866.35 5819003.03)")

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

@@ -67,8 +67,46 @@ SRID=32633;LINESTRING(1 2,3.5 4.5)
 SRID=4326;LINESTRING(1 2,1 2,10.1 5.5)
 LINESTRING(1 2,10.1 5.5,3.5 4.5,6.1 7.9)
 
+statement ok
+DROP TABLE IF EXISTS t1;
+
+statement ok
+CREATE TABLE t1 (b binary);
+
+statement ok
+INSERT INTO t1 VALUES(UNHEX('0101000020797f000066666666a9cb17411f85ebc19e325641'));
+
+query T
+SELECT TO_STRING(ST_GEOMETRYFROMWKB(b)) FROM t1;
+----
+SRID=32633;POINT(389866.35 5819003.03)
+
+query T
+SELECT TO_STRING(ST_GEOMETRYFROMWKB(b, 4326)) FROM t1;
+----
+SRID=4326;POINT(389866.35 5819003.03)
+
+statement ok
+DROP TABLE IF EXISTS t1;
+
+statement ok
+CREATE TABLE t1 (b STRING);
+
+statement ok
+INSERT INTO t1 VALUES('0101000020797f000066666666a9cb17411f85ebc19e325641');
+
+query T
+SELECT TO_STRING(ST_GEOMETRYFROMWKB(b)) FROM t1;
+----
+SRID=32633;POINT(389866.35 5819003.03)
+
+query T
+SELECT TO_STRING(ST_GEOMETRYFROMWKB(b, 4326)) FROM t1;
+----
+SRID=4326;POINT(389866.35 5819003.03)
+
 statement ok
 SET enable_geo_create_table=0
 
 statement ok
-DROP TABLE IF EXISTS t1
+DROP TABLE IF EXISTS t1;