Skip to content

Commit

Permalink
maybe work with struct wrapping
Browse files Browse the repository at this point in the history
  • Loading branch information
paleolimbot committed Jan 13, 2025
1 parent a2cbc7d commit a3e244a
Showing 1 changed file with 25 additions and 1 deletion.
26 changes: 25 additions & 1 deletion rust/geodatafusion/src/data_types.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
use std::sync::Arc;

use arrow_array::ArrayRef;
use arrow_schema::{DataType, Fields};
use datafusion::error::DataFusionError;
use datafusion::logical_expr::{Signature, Volatility};
use geoarrow::array::{CoordType, GeometryArray, PointArray, RectArray};
use geoarrow::array::{CoordType, GeometryArray, NativeArrayDyn, PointArray, RectArray};
use geoarrow::datatypes::{Dimension, NativeType};
use geoarrow::NativeArray;

Expand All @@ -24,6 +25,7 @@ pub(crate) fn any_single_geometry_type_input() -> Signature {
BOX2D_TYPE.into(),
BOX3D_TYPE.into(),
GEOMETRY_TYPE.into(),
struct_wrapped_type(POINT2D_TYPE),
],
Volatility::Immutable,
)
Expand All @@ -32,6 +34,22 @@ pub(crate) fn any_single_geometry_type_input() -> Signature {
/// This will not cast a PointArray to a GeometryArray
pub(crate) fn parse_to_native_array(array: ArrayRef) -> GeoDataFusionResult<Arc<dyn NativeArray>> {
let data_type = array.data_type();

match data_type {
DataType::Struct(fields) => {
if fields.size() != 1 {
return Err(DataFusionError::Execution(
"Extension workaround struct with fields.size() != 1".to_string(),
)
.into());
}

let native_array_dyn = NativeArrayDyn::from_arrow_array(&array, &fields[0])?;
return Ok(native_array_dyn.into());
}
_ => {}
}

if data_type.equals_datatype(&POINT2D_TYPE.into()) {
let point_array = PointArray::try_from((array.as_ref(), Dimension::XY))?;
Ok(Arc::new(point_array))
Expand All @@ -50,3 +68,9 @@ pub(crate) fn parse_to_native_array(array: ArrayRef) -> GeoDataFusionResult<Arc<
Err(DataFusionError::Execution(format!("Unexpected input data type: {}", data_type)).into())
}
}

fn struct_wrapped_type(native_type: NativeType) -> DataType {
return DataType::Struct(Fields::from(vec![
native_type.to_field(native_type.extension_name(), false)
]));
}

0 comments on commit a3e244a

Please sign in to comment.