Skip to content

Commit

Permalink
Reuse buffers zero-copy in from_shapely
Browse files Browse the repository at this point in the history
  • Loading branch information
@kylebarron
kylebarron committed Oct 7, 2024
1 parent 8575896 commit 3e4ddcf
Show file tree
Hide file tree
Showing 3 changed files with 109 additions and 262 deletions.
14 changes: 13 additions & 1 deletion python/geoarrow-core/python/geoarrow/rust/core/_constructors.pyi
View file
Original file line number Diff line number Diff line change
@@ -1,32 +1,44 @@
from typing import List, Tuple
from typing import Any, List, Tuple

from arro3.core.types import ArrayInput

from ._rust import NativeArray

def points(
coords: ArrayInput | Tuple[ArrayInput, ...] | List[ArrayInput],
*,
crs: Any | None = None,
) -> NativeArray: ...
def linestrings(
coords: ArrayInput | Tuple[ArrayInput, ...] | List[ArrayInput],
geom_offsets: ArrayInput,
*,
crs: Any | None = None,
) -> NativeArray: ...
def polygons(
coords: ArrayInput | Tuple[ArrayInput, ...] | List[ArrayInput],
geom_offsets: ArrayInput,
ring_offsets: ArrayInput,
*,
crs: Any | None = None,
) -> NativeArray: ...
def multipoints(
coords: ArrayInput | Tuple[ArrayInput, ...] | List[ArrayInput],
*,
crs: Any | None = None,
) -> NativeArray: ...
def multilinestrings(
coords: ArrayInput | Tuple[ArrayInput, ...] | List[ArrayInput],
geom_offsets: ArrayInput,
ring_offsets: ArrayInput,
*,
crs: Any | None = None,
) -> NativeArray: ...
def multipolygons(
coords: ArrayInput | Tuple[ArrayInput, ...] | List[ArrayInput],
geom_offsets: ArrayInput,
polygon_offsets: ArrayInput,
ring_offsets: ArrayInput,
*,
crs: Any | None = None,
) -> NativeArray: ...
57 changes: 40 additions & 17 deletions python/geoarrow-core/src/constructors.rs
View file
Original file line number Diff line number Diff line change
@@ -1,59 +1,75 @@
use std::sync::Arc;

use geoarrow::array::metadata::ArrayMetadata;
use geoarrow::array::{
LineStringArray, MultiLineStringArray, MultiPointArray, MultiPolygonArray, NativeArrayDyn,
PointArray, PolygonArray,
};
use pyo3::prelude::*;
use pyo3_geoarrow::{PyCoordBuffer, PyGeoArrowResult, PyNativeArray, PyOffsetBuffer};

use crate::crs::CRS;

fn create_array_metadata(crs: Option<CRS>) -> Arc<ArrayMetadata> {
Arc::new(ArrayMetadata {
crs: crs.map(|c| c.into_inner()),
..Default::default()
})
}

#[pyfunction]
pub fn points(coords: PyCoordBuffer) -> PyGeoArrowResult<PyNativeArray> {
#[pyo3(signature = (coords, *, crs = None))]
pub fn points(coords: PyCoordBuffer, crs: Option<CRS>) -> PyGeoArrowResult<PyNativeArray> {
let metadata = create_array_metadata(crs);
match coords {
PyCoordBuffer::TwoD(coords) => {
let array = PointArray::new(coords, None, Default::default());
let array = PointArray::new(coords, None, metadata);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
PyCoordBuffer::ThreeD(coords) => {
let array = PointArray::new(coords, None, Default::default());
let array = PointArray::new(coords, None, metadata);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
}
}

#[pyfunction]
#[pyo3(signature = (coords, geom_offsets, *, crs = None))]
pub fn linestrings(
coords: PyCoordBuffer,
geom_offsets: PyOffsetBuffer,
crs: Option<CRS>,
) -> PyGeoArrowResult<PyNativeArray> {
let metadata = create_array_metadata(crs);
match coords {
PyCoordBuffer::TwoD(coords) => {
let array =
LineStringArray::new(coords, geom_offsets.into_inner(), None, Default::default());
let array = LineStringArray::new(coords, geom_offsets.into_inner(), None, metadata);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
PyCoordBuffer::ThreeD(coords) => {
let array =
LineStringArray::new(coords, geom_offsets.into_inner(), None, Default::default());
let array = LineStringArray::new(coords, geom_offsets.into_inner(), None, metadata);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
}
}

#[pyfunction]
#[pyo3(signature = (coords, geom_offsets, ring_offsets, *, crs = None))]
pub fn polygons(
coords: PyCoordBuffer,
geom_offsets: PyOffsetBuffer,
ring_offsets: PyOffsetBuffer,
crs: Option<CRS>,
) -> PyGeoArrowResult<PyNativeArray> {
let metadata = create_array_metadata(crs);
match coords {
PyCoordBuffer::TwoD(coords) => {
let array = PolygonArray::new(
coords,
geom_offsets.into_inner(),
ring_offsets.into_inner(),
None,
Default::default(),
metadata,
);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
Expand All @@ -63,46 +79,50 @@ pub fn polygons(
geom_offsets.into_inner(),
ring_offsets.into_inner(),
None,
Default::default(),
metadata,
);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
}
}

#[pyfunction]
#[pyo3(signature = (coords, geom_offsets, *, crs = None))]
pub fn multipoints(
coords: PyCoordBuffer,
geom_offsets: PyOffsetBuffer,
crs: Option<CRS>,
) -> PyGeoArrowResult<PyNativeArray> {
let metadata = create_array_metadata(crs);
match coords {
PyCoordBuffer::TwoD(coords) => {
let array =
MultiPointArray::new(coords, geom_offsets.into_inner(), None, Default::default());
let array = MultiPointArray::new(coords, geom_offsets.into_inner(), None, metadata);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
PyCoordBuffer::ThreeD(coords) => {
let array =
MultiPointArray::new(coords, geom_offsets.into_inner(), None, Default::default());
let array = MultiPointArray::new(coords, geom_offsets.into_inner(), None, metadata);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
}
}

#[pyfunction]
#[pyo3(signature = (coords, geom_offsets, ring_offsets, *, crs = None))]
pub fn multilinestrings(
coords: PyCoordBuffer,
geom_offsets: PyOffsetBuffer,
ring_offsets: PyOffsetBuffer,
crs: Option<CRS>,
) -> PyGeoArrowResult<PyNativeArray> {
let metadata = create_array_metadata(crs);
match coords {
PyCoordBuffer::TwoD(coords) => {
let array = MultiLineStringArray::new(
coords,
geom_offsets.into_inner(),
ring_offsets.into_inner(),
None,
Default::default(),
metadata,
);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
Expand All @@ -112,20 +132,23 @@ pub fn multilinestrings(
geom_offsets.into_inner(),
ring_offsets.into_inner(),
None,
Default::default(),
metadata,
);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
}
}

#[pyfunction]
#[pyo3(signature = (coords, geom_offsets, polygon_offsets, ring_offsets, *, crs = None))]
pub fn multipolygons(
coords: PyCoordBuffer,
geom_offsets: PyOffsetBuffer,
polygon_offsets: PyOffsetBuffer,
ring_offsets: PyOffsetBuffer,
crs: Option<CRS>,
) -> PyGeoArrowResult<PyNativeArray> {
let metadata = create_array_metadata(crs);
match coords {
PyCoordBuffer::TwoD(coords) => {
let array = MultiPolygonArray::new(
Expand All @@ -134,7 +157,7 @@ pub fn multipolygons(
polygon_offsets.into_inner(),
ring_offsets.into_inner(),
None,
Default::default(),
metadata,
);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
Expand All @@ -145,7 +168,7 @@ pub fn multipolygons(
polygon_offsets.into_inner(),
ring_offsets.into_inner(),
None,
Default::default(),
metadata,
);
Ok(PyNativeArray::new(NativeArrayDyn::new(Arc::new(array))))
}
Expand Down
Loading

0 comments on commit 3e4ddcf

Please sign in to comment.