feat: support gis function ST_AsBinary and ST_GeomFromWKB

query_server/query/src/extension/expr/scalar_function/gis/mod.rs

@@ -1,10 +1,14 @@
+mod st_asbinary;
 mod st_distance;
+mod st_geomfromwkb;
 
 use spi::query::function::FunctionMetadataManager;
 use spi::Result;
 
 pub fn register_udfs(func_manager: &mut dyn FunctionMetadataManager) -> Result<()> {
     st_distance::register_udf(func_manager)?;
+    st_geomfromwkb::register_udf(func_manager)?;
+    st_asbinary::register_udf(func_manager)?;
     Ok(())
 }
 

query_server/query/src/extension/expr/scalar_function/gis/st_asbinary.rs

@@ -0,0 +1,44 @@
+use std::sync::Arc;
+
+use datafusion::arrow::array::{downcast_array, ArrayRef, BinaryArray, StringArray};
+use datafusion::arrow::datatypes::DataType;
+use datafusion::common::Result as DFResult;
+use datafusion::logical_expr::{ReturnTypeFunction, ScalarUDF, Signature, Volatility};
+use datafusion::physical_plan::functions::make_scalar_function;
+use geozero::wkt::WktStr;
+use geozero::{CoordDimensions, ToWkb};
+use spi::query::function::FunctionMetadataManager;
+use spi::Result;
+
+pub fn register_udf(func_manager: &mut dyn FunctionMetadataManager) -> Result<ScalarUDF> {
+    let udf = new();
+    func_manager.register_udf(udf.clone())?;
+    Ok(udf)
+}
+
+fn new() -> ScalarUDF {
+    let fun = make_scalar_function(func);
+
+    let signature = Signature::exact(vec![DataType::Utf8], Volatility::Immutable);
+    let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Binary)));
+
+    ScalarUDF::new("st_AsBinary", &signature, &return_type, &fun)
+}
+
+fn func(args: &[ArrayRef]) -> DFResult<ArrayRef> {
+    let wkt_arr = args[0].as_ref();
+    let wkt_arr = downcast_array::<StringArray>(wkt_arr);
+
+    let result: BinaryArray = wkt_arr
+        .iter()
+        .map(|opt| {
+            opt.and_then(|str| {
+                let wkb = WktStr(str);
+                // conversion failed to null
+                wkb.to_wkb(CoordDimensions::xy()).ok()
+            })
+        })
+        .collect();
+
+    Ok(Arc::new(result))
+}

query_server/query/src/extension/expr/scalar_function/gis/st_geomfromwkb.rs

@@ -0,0 +1,44 @@
+use std::sync::Arc;
+
+use datafusion::arrow::array::{downcast_array, ArrayRef, BinaryArray, StringArray};
+use datafusion::arrow::datatypes::DataType;
+use datafusion::common::Result as DFResult;
+use datafusion::logical_expr::{ReturnTypeFunction, ScalarUDF, Signature, Volatility};
+use datafusion::physical_plan::functions::make_scalar_function;
+use geozero::wkb::Wkb;
+use geozero::ToWkt;
+use spi::query::function::FunctionMetadataManager;
+use spi::Result;
+
+pub fn register_udf(func_manager: &mut dyn FunctionMetadataManager) -> Result<ScalarUDF> {
+    let udf = new();
+    func_manager.register_udf(udf.clone())?;
+    Ok(udf)
+}
+
+fn new() -> ScalarUDF {
+    let fun = make_scalar_function(func);
+
+    let signature = Signature::exact(vec![DataType::Binary], Volatility::Immutable);
+    let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Utf8)));
+
+    ScalarUDF::new("st_GeomFromWKB", &signature, &return_type, &fun)
+}
+
+fn func(args: &[ArrayRef]) -> DFResult<ArrayRef> {
+    let wkb_arr = args[0].as_ref();
+    let wkb_arr = downcast_array::<BinaryArray>(wkb_arr);
+
+    let result: StringArray = wkb_arr
+        .iter()
+        .map(|opt| {
+            opt.and_then(|bytes| {
+                let wkb = Wkb(bytes.to_vec());
+                // conversion failed to null
+                wkb.to_wkt().ok()
+            })
+        })
+        .collect();
+
+    Ok(Arc::new(result))
+}

query_server/sqllogicaltests/cases/function/gis/st_asbinary.slt

@@ -0,0 +1,35 @@
+include ./setup.slt
+
+##########
+## Query
+##########
+
+query 
+SELECT time, loc, st_AsBinary(loc) FROM gis_loc order by time, loc;
+----
+1999-12-31T00:00:00 POINT(0 0) 010100000000000000000000000000000000000000
+1999-12-31T00:00:00.005 POINT(0 1) 01010000000000000000000000000000000000f03f
+1999-12-31T00:00:00.010 POINT(0 2) 010100000000000000000000000000000000000040
+1999-12-31T00:00:10.015 POINT(0 3) 010100000000000000000000000000000000000840
+1999-12-31T00:00:10.020 POINT(0 4) 010100000000000000000000000000000000001040
+1999-12-31T00:10:00.025 POINT(0 5) 010100000000000000000000000000000000001440
+1999-12-31T00:10:00.030 POINT(0 6) 010100000000000000000000000000000000001840
+1999-12-31T01:00:00.035 POINT(0 7) 010100000000000000000000000000000000001c40
+
+query 
+select st_AsBinary(null);
+----
+NULL
+
+query 
+select st_AsBinary('POINT(0 0)');
+----
+010100000000000000000000000000000000000000
+
+query 
+select st_AsBinary('POINT(0, 0)');
+----
+NULL
+
+query error Arrow error: Io error: Status \{ code: Internal, message: "Build logical plan: Datafusion: Error during planning: No function matches the given name and argument types 'st_AsBinary\(Utf8, Null\)'\. You might need to add explicit type casts\.\\n\\tCandidate functions:\\n\\tst_AsBinary\(Utf8\)", .*
+select st_AsBinary('POINT(0, 0)', null);

query_server/sqllogicaltests/cases/function/gis/st_geomfromwkb.slt

@@ -0,0 +1,22 @@
+query 
+select ST_GeomFromWKB(st_AsBinary('POINT(0 0)'));
+----
+POINT(0 0)
+
+query 
+select ST_GeomFromWKB(st_AsBinary('MULTIPOINT (10 40, 40 30, 20 20, 30 10)'));
+----
+MULTIPOINT(10 40,40 30,20 20,30 10)
+
+query 
+select ST_GeomFromWKB(st_AsBinary('MULTIPOLYGON (((30 20, 45 40, 10 40, 30 20)),((15 5, 40 10, 10 20, 5 10, 15 5)))'));
+----
+MULTIPOLYGON(((30 20,45 40,10 40,30 20)),((15 5,40 10,10 20,5 10,15 5)))
+
+query 
+select ST_GeomFromWKB(st_AsBinary('MULTIPOLYGON (((40 40, 20 45, 45 30, 40 40)),((20 35, 10 30, 10 10, 30 5, 45 20, 20 35),(30 20, 20 15, 20 25, 30 20)))'));
+----
+MULTIPOLYGON(((40 40,20 45,45 30,40 40)),((20 35,10 30,10 10,30 5,45 20,20 35),(30 20,20 15,20 25,30 20)))
+
+query error Arrow error: Io error: Status \{ code: Internal, message: "Build logical plan: Datafusion: Error during planning: No function matches the given name and argument types 'st_GeomFromWKB\(Utf8\)'\. You might need to add explicit type casts\.\\n\\tCandidate functions:\\n\\tst_GeomFromWKB\(Binary\)", .*
+select ST_GeomFromWKB('invalid');