Skip to content

Commit

Permalink
Fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@lewiszlw
lewiszlw committed Mar 28, 2024
1 parent 08a3922 commit a9ac22b
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 22 deletions.
24 changes: 16 additions & 8 deletions src/function/boundary.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
use crate::function::AsGeoJsonUdf;
use crate::geo::{GeometryArray, GeometryArrayBuilder};
use crate::DFResult;
use arrow_array::cast::AsArray;
use arrow_array::{GenericBinaryArray, OffsetSizeTrait};
use arrow_schema::DataType;
use datafusion_common::internal_datafusion_err;
use datafusion_common::{internal_datafusion_err, DataFusionError};
use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, TypeSignature, Volatility};
use geos::Geom;
use geozero::wkb::WkbDialect;
Expand Down Expand Up @@ -74,10 +73,14 @@ fn build_boundary_arr<O: OffsetSizeTrait>(
) -> DFResult<ColumnarValue> {
let mut builder = GeometryArrayBuilder::<O>::new(WkbDialect::Ewkb, wkb_arr.geom_len());
for i in 0..wkb_arr.geom_len() {
builder.append_geos_geometry(&wkb_arr.geos_value(i)?.map(|geom| {
geom.boundary()
.map_err(|e| internal_datafusion_err!("Failed to call boundary, e: {}", e))?
}))?;
if let Some(geom) = wkb_arr.geos_value(i)? {
builder
.append_geos_geometry(&Some(geom.boundary().map_err(|e| {
internal_datafusion_err!("Failed to call boundary, e: {}", e)
})?))?;
} else {
builder.append_null();
}
}

Ok(ColumnarValue::Array(Arc::new(builder.build())))
Expand All @@ -91,7 +94,7 @@ impl Default for BoundaryUdf {

#[cfg(test)]
mod tests {
use crate::function::{BoundaryUdf, GeomFromTextUdf};
use crate::function::{AsTextUdf, BoundaryUdf, GeomFromTextUdf};
use arrow::util::pretty::pretty_format_batches;
use datafusion::logical_expr::ScalarUDF;
use datafusion::prelude::SessionContext;
Expand All @@ -100,6 +103,7 @@ mod tests {
async fn boundary() {
let ctx = SessionContext::new();
ctx.register_udf(ScalarUDF::from(GeomFromTextUdf::new()));
ctx.register_udf(ScalarUDF::from(AsTextUdf::new()));
ctx.register_udf(ScalarUDF::from(BoundaryUdf::new()));
let df = ctx
.sql("SELECT ST_AsText(ST_Boundary(ST_GeomFromText('POLYGON((1 1,0 0, -1 1, 1 1))')));")
Expand All @@ -109,7 +113,11 @@ mod tests {
pretty_format_batches(&df.collect().await.unwrap())
.unwrap()
.to_string(),
""
"+--------------------------------------------------------------------------------+
| ST_AsText(ST_Boundary(ST_GeomFromText(Utf8(\"POLYGON((1 1,0 0, -1 1, 1 1))\")))) |
+--------------------------------------------------------------------------------+
| LINESTRING(1 1,0 0,-1 1,1 1) |
+--------------------------------------------------------------------------------+"
);
}
}
32 changes: 20 additions & 12 deletions src/function/buffer.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
use crate::function::AsGeoJsonUdf;
use crate::geo::{GeometryArray, GeometryArrayBuilder};
use crate::DFResult;
use arrow_array::cast::AsArray;
use arrow_array::{GenericBinaryArray, LargeStringArray, OffsetSizeTrait, StringArray};
use arrow_array::{GenericBinaryArray, OffsetSizeTrait};
use arrow_schema::DataType;
use datafusion_common::{internal_datafusion_err, internal_err, ScalarValue};
use datafusion_common::{internal_datafusion_err, internal_err, DataFusionError, ScalarValue};
use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, TypeSignature, Volatility};
use geos::Geom;
use geozero::wkb::WkbDialect;
Expand All @@ -24,12 +23,12 @@ impl BufferUdf {
vec![
TypeSignature::Exact(vec![
DataType::Binary,
DataType::Float32,
DataType::Float64,
DataType::Int32,
]),
TypeSignature::Exact(vec![
DataType::LargeBinary,
DataType::Float32,
DataType::Float64,
DataType::Int32,
]),
],
Expand Down Expand Up @@ -91,10 +90,14 @@ fn build_buffer_arr<O: OffsetSizeTrait>(
) -> DFResult<ColumnarValue> {
let mut builder = GeometryArrayBuilder::<O>::new(WkbDialect::Ewkb, wkb_arr.geom_len());
for i in 0..wkb_arr.geom_len() {
builder.append_geos_geometry(&wkb_arr.geos_value(i)?.map(|geom| {
geom.buffer(width, quadsegs)
.map_err(|e| internal_datafusion_err!("Failed to call buffer, e: {}", e))?
}))?;
if let Some(geom) = wkb_arr.geos_value(i)? {
builder.append_geos_geometry(&Some(
geom.buffer(width, quadsegs)
.map_err(|e| internal_datafusion_err!("Failed to call buffer, e: {}", e))?,
))?;
} else {
builder.append_null();
}
}

Ok(ColumnarValue::Array(Arc::new(builder.build())))
Expand All @@ -108,7 +111,7 @@ impl Default for BufferUdf {

#[cfg(test)]
mod tests {
use crate::function::{BufferUdf, GeomFromTextUdf};
use crate::function::{AsTextUdf, BufferUdf, GeomFromTextUdf};
use arrow::util::pretty::pretty_format_batches;
use datafusion::logical_expr::ScalarUDF;
use datafusion::prelude::SessionContext;
Expand All @@ -117,16 +120,21 @@ mod tests {
async fn buffer() {
let ctx = SessionContext::new();
ctx.register_udf(ScalarUDF::from(GeomFromTextUdf::new()));
ctx.register_udf(ScalarUDF::from(AsTextUdf::new()));
ctx.register_udf(ScalarUDF::from(BufferUdf::new()));
let df = ctx
.sql("SELECT ST_Buffer(ST_GeomFromText('POINT(100 90)'), 50.0, 8);")
.sql("SELECT ST_AsText(ST_Buffer(ST_GeomFromText('POINT(100 90)'), 50.0, 2::Integer));")
.await
.unwrap();
assert_eq!(
pretty_format_batches(&df.collect().await.unwrap())
.unwrap()
.to_string(),
""
"+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| ST_AsText(ST_Buffer(ST_GeomFromText(Utf8(\"POINT(100 90)\")),Float64(50),Int64(2))) |
+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| POLYGON((150 90,135.35533905932738 54.64466094067263,100 40,64.64466094067262 54.64466094067262,50 90,64.64466094067262 125.35533905932738,99.99999999999999 140,135.35533905932738 125.35533905932738,150 90)) |
+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+"
);
}
}
4 changes: 2 additions & 2 deletions src/function/split.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ fn split<O: OffsetSizeTrait, F: OffsetSizeTrait>(
#[cfg(test)]
mod tests {
use crate::function::{GeomFromTextUdf, SplitUdf};
use arrow::util::pretty::pretty_format_batches;
// use arrow::util::pretty::pretty_format_batches;
use datafusion::prelude::SessionContext;
use datafusion_expr::ScalarUDF;

Expand All @@ -141,7 +141,7 @@ mod tests {
let ctx = SessionContext::new();
ctx.register_udf(ScalarUDF::from(GeomFromTextUdf::new()));
ctx.register_udf(ScalarUDF::from(SplitUdf::new()));
let df = ctx
let _df = ctx
.sql("select ST_Split(ST_GeomFromText('LINESTRING ( 0 0, 1 1, 2 2 )'), ST_GeomFromText('POINT(1 1)'))")
.await
.unwrap();
Expand Down

0 comments on commit a9ac22b

Please sign in to comment.