diff --git a/query_server/query/src/extension/expr/scalar_function/gis/mod.rs b/query_server/query/src/extension/expr/scalar_function/gis/mod.rs index 55d600148..536afe80a 100644 --- a/query_server/query/src/extension/expr/scalar_function/gis/mod.rs +++ b/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(()) } diff --git a/query_server/query/src/extension/expr/scalar_function/gis/st_asbinary.rs b/query_server/query/src/extension/expr/scalar_function/gis/st_asbinary.rs new file mode 100644 index 000000000..108080818 --- /dev/null +++ b/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 { + 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 { + let wkt_arr = args[0].as_ref(); + let wkt_arr = downcast_array::(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)) +} diff --git a/query_server/query/src/extension/expr/scalar_function/gis/st_geomfromwkb.rs b/query_server/query/src/extension/expr/scalar_function/gis/st_geomfromwkb.rs new file mode 100644 index 000000000..455862b71 --- /dev/null +++ b/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 { + 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 { + let wkb_arr = args[0].as_ref(); + let wkb_arr = downcast_array::(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)) +} diff --git a/query_server/sqllogicaltests/cases/function/gis/st_asbinary.slt b/query_server/sqllogicaltests/cases/function/gis/st_asbinary.slt new file mode 100644 index 000000000..6c5d7e2a9 --- /dev/null +++ b/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); diff --git a/query_server/sqllogicaltests/cases/function/gis/st_geomfromwkb.slt b/query_server/sqllogicaltests/cases/function/gis/st_geomfromwkb.slt new file mode 100644 index 000000000..eb1ca3d42 --- /dev/null +++ b/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');