geoarrow / geoarrow-rs

GeoArrow in Rust, Python, and JavaScript (WebAssembly) with vectorized geometry operations
http://geoarrow.org/geoarrow-rs/
Apache License 2.0
259 stars 17 forks source link

Generic geometry accessor traits #804

Closed kylebarron closed 1 month ago

kylebarron commented 1 month ago

We want to be able to wrap functions very quickly. This enables that for GEOS... now we can get one-line wrappers:

pub fn area<'a>(array: &'a dyn NativeGEOSGeometryAccessor<'a>) -> Result<Float64Array> {
    Ok(try_unary_primitive(array, |geom| geom.area())?)
}

Perhaps we still want to implement the trait on each concrete array. I'm not sure.

Where here try_unary_primitive is the following:

// Note: This is derived from arrow-rs here:
// https://github.com/apache/arrow-rs/blob/3ed7cc61d4157263ef2ab5c2d12bc7890a5315b3/arrow-array/src/array/primitive_array.rs#L806-L830
fn try_unary_primitive<'a, F, O, E>(
    array: &'a dyn NativeGEOSGeometryAccessor<'a>,
    op: F,
) -> std::result::Result<PrimitiveArray<O>, E>
where
    O: ArrowPrimitiveType,
    F: Fn(geos::Geometry) -> std::result::Result<O::Native, E>,
    E: std::convert::From<geos::Error>,
{
    let len = array.len();

    let nulls = array.nulls().cloned();
    let mut buffer = BufferBuilder::<O::Native>::new(len);
    buffer.append_n_zeroed(len);
    let slice = buffer.as_slice_mut();

    let f = |idx| {
        unsafe { *slice.get_unchecked_mut(idx) = op(array.value_as_geometry_unchecked(idx)?)? };
        Ok::<_, E>(())
    };

    match &nulls {
        Some(nulls) => nulls.try_for_each_valid_idx(f)?,
        None => (0..len).try_for_each(f)?,
    }

    let values = buffer.finish().into();
    Ok(PrimitiveArray::new(values, nulls))
}