dadams/california-equity-git
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

library packages

Browse Source
This commit is contained in:
David Adams
2024-09-28 22:56:00 -07:00
parent 64d9b78b3a
commit 1973934e95
4893 changed files with 1184173 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

614
.venv/lib/python3.12/site-packages/geopandas/io/_geoarrow.py Normal file
View File

@@ -0,0 +1,614 @@
import json
from packaging.version import Version
from typing import Dict, Optional, Tuple
import numpy as np
import pandas as pd
import pyarrow as pa
from numpy.typing import NDArray
import shapely
from shapely import GeometryType
from geopandas import GeoDataFrame
from geopandas._compat import SHAPELY_GE_204
from geopandas.array import from_shapely, from_wkb
GEOARROW_ENCODINGS = [
"point",
"linestring",
"polygon",
"multipoint",
"multilinestring",
"multipolygon",
]
## GeoPandas -> GeoArrow
class ArrowTable:
"""
Wrapper class for Arrow data.
This class implements the `Arrow PyCapsule Protocol`_ (i.e. having an
``__arrow_c_stream__`` method). This object can then be consumed by
your Arrow implementation of choice that supports this protocol.
.. _Arrow PyCapsule Protocol: https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
Example
-------
>>> import pyarrow as pa
>>> pa.table(gdf.to_arrow()) # doctest: +SKIP
"""
def __init__(self, pa_table):
self._pa_table = pa_table
def __arrow_c_stream__(self, requested_schema=None):
return self._pa_table.__arrow_c_stream__(requested_schema=requested_schema)
class GeoArrowArray:
"""
Wrapper class for a geometry array as Arrow data.
This class implements the `Arrow PyCapsule Protocol`_ (i.e. having an
``__arrow_c_array/stream__`` method). This object can then be consumed by
your Arrow implementation of choice that supports this protocol.
.. _Arrow PyCapsule Protocol: https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html
Example
-------
>>> import pyarrow as pa
>>> pa.array(ser.to_arrow()) # doctest: +SKIP
"""
def __init__(self, pa_field, pa_array):
self._pa_array = pa_array
self._pa_field = pa_field
def __arrow_c_array__(self, requested_schema=None):
if requested_schema is not None:
raise NotImplementedError(
"Requested schema is not supported for geometry arrays"
)
return (
self._pa_field.__arrow_c_schema__(),
self._pa_array.__arrow_c_array__()[1],
)
def geopandas_to_arrow(
df,
index=None,
geometry_encoding="WKB",
interleaved=True,
include_z=None,
):
"""
Convert GeoDataFrame to a pyarrow.Table.
Parameters
----------
df : GeoDataFrame
The GeoDataFrame to convert.
index : bool, default None
If ``True``, always include the dataframe's index(es) as columns
in the file output.
If ``False``, the index(es) will not be written to the file.
If ``None``, the index(ex) will be included as columns in the file
output except `RangeIndex` which is stored as metadata only.
geometry_encoding : {'WKB', 'geoarrow' }, default 'WKB'
The GeoArrow encoding to use for the data conversion.
interleaved : bool, default True
Only relevant for 'geoarrow' encoding. If True, the geometries'
coordinates are interleaved in a single fixed size list array.
If False, the coordinates are stored as separate arrays in a
struct type.
include_z : bool, default None
Only relevant for 'geoarrow' encoding (for WKB, the dimensionality
of the individial geometries is preserved).
If False, return 2D geometries. If True, include the third dimension
in the output (if a geometry has no third dimension, the z-coordinates
will be NaN). By default, will infer the dimensionality from the
input geometries. Note that this inference can be unreliable with
empty geometries (for a guaranteed result, it is recommended to
specify the keyword).
"""
mask = df.dtypes == "geometry"
geometry_columns = df.columns[mask]
geometry_indices = np.asarray(mask).nonzero()[0]
df_attr = pd.DataFrame(df.copy(deep=False))
# replace geometry columns with dummy values -> will get converted to
# Arrow null column (not holding any memory), so we can afterwards
# fill the resulting table with the correct geometry fields
for col in geometry_columns:
df_attr[col] = None
table = pa.Table.from_pandas(df_attr, preserve_index=index)
geometry_encoding_dict = {}
if geometry_encoding.lower() == "geoarrow":
if Version(pa.__version__) < Version("10.0.0"):
raise ValueError("Converting to 'geoarrow' requires pyarrow >= 10.0.")
# Encode all geometry columns to GeoArrow
for i, col in zip(geometry_indices, geometry_columns):
field, geom_arr = construct_geometry_array(
np.array(df[col].array),
include_z=include_z,
field_name=col,
crs=df[col].crs,
interleaved=interleaved,
)
table = table.set_column(i, field, geom_arr)
geometry_encoding_dict[col] = (
field.metadata[b"ARROW:extension:name"]
.decode()
.removeprefix("geoarrow.")
)
elif geometry_encoding.lower() == "wkb":
# Encode all geometry columns to WKB
for i, col in zip(geometry_indices, geometry_columns):
field, wkb_arr = construct_wkb_array(
np.asarray(df[col].array), field_name=col, crs=df[col].crs
)
table = table.set_column(i, field, wkb_arr)
geometry_encoding_dict[col] = "WKB"
else:
raise ValueError(
f"Expected geometry encoding 'WKB' or 'geoarrow' got {geometry_encoding}"
)
return table, geometry_encoding_dict
def construct_wkb_array(
shapely_arr: NDArray[np.object_],
*,
field_name: str = "geometry",
crs: Optional[str] = None,
) -> Tuple[pa.Field, pa.Array]:
if shapely.geos_version > (3, 10, 0):
kwargs = {"flavor": "iso"}
else:
if shapely.has_z(shapely_arr).any():
raise ValueError("Cannot write 3D geometries with GEOS<3.10")
kwargs = {}
wkb_arr = shapely.to_wkb(shapely_arr, **kwargs)
extension_metadata = {"ARROW:extension:name": "geoarrow.wkb"}
if crs is not None:
extension_metadata["ARROW:extension:metadata"] = json.dumps(
{"crs": crs.to_json()}
)
else:
# In theory this should not be needed, but otherwise pyarrow < 17
# crashes on receiving such data through C Data Interface
# https://github.com/apache/arrow/issues/41741
extension_metadata["ARROW:extension:metadata"] = "{}"
field = pa.field(
field_name, type=pa.binary(), nullable=True, metadata=extension_metadata
)
parr = pa.array(np.asarray(wkb_arr), pa.binary())
return field, parr
def _convert_inner_coords(coords, interleaved, dims, mask=None):
if interleaved:
coords_field = pa.field(dims, pa.float64(), nullable=False)
typ = pa.list_(coords_field, len(dims))
if mask is None:
# mask keyword only added in pyarrow 15.0.0
parr = pa.FixedSizeListArray.from_arrays(coords.ravel(), type=typ)
else:
parr = pa.FixedSizeListArray.from_arrays(
coords.ravel(), type=typ, mask=mask
)
else:
if dims == "xy":
fields = [
pa.field("x", pa.float64(), nullable=False),
pa.field("y", pa.float64(), nullable=False),
]
parr = pa.StructArray.from_arrays(
[coords[:, 0].copy(), coords[:, 1].copy()], fields=fields, mask=mask
)
else:
fields = [
pa.field("x", pa.float64(), nullable=False),
pa.field("y", pa.float64(), nullable=False),
pa.field("z", pa.float64(), nullable=False),
]
parr = pa.StructArray.from_arrays(
[coords[:, 0].copy(), coords[:, 1].copy(), coords[:, 2].copy()],
fields=fields,
mask=mask,
)
return parr
def _linestring_type(point_type):
return pa.list_(pa.field("vertices", point_type, nullable=False))
def _polygon_type(point_type):
return pa.list_(
pa.field(
"rings",
pa.list_(pa.field("vertices", point_type, nullable=False)),
nullable=False,
)
)
def _multipoint_type(point_type):
return pa.list_(pa.field("points", point_type, nullable=False))
def _multilinestring_type(point_type):
return pa.list_(
pa.field("linestrings", _linestring_type(point_type), nullable=False)
)
def _multipolygon_type(point_type):
return pa.list_(pa.field("polygons", _polygon_type(point_type), nullable=False))
def construct_geometry_array(
shapely_arr: NDArray[np.object_],
include_z: Optional[bool] = None,
*,
field_name: str = "geometry",
crs: Optional[str] = None,
interleaved: bool = True,
) -> Tuple[pa.Field, pa.Array]:
# NOTE: this implementation returns a (field, array) pair so that it can set the
# extension metadata on the field without instantiating extension types into the
# global pyarrow registry
geom_type, coords, offsets = shapely.to_ragged_array(
shapely_arr, include_z=include_z
)
mask = shapely.is_missing(shapely_arr)
if mask.any():
if (
geom_type == GeometryType.POINT
and interleaved
and Version(pa.__version__) < Version("15.0.0")
):
raise ValueError(
"Converting point geometries with missing values is not supported "
"for interleaved coordinates with pyarrow < 15.0.0. Please "
"upgrade to a newer version of pyarrow."
)
mask = pa.array(mask, type=pa.bool_())
if geom_type == GeometryType.POINT and not SHAPELY_GE_204:
# bug in shapely < 2.0.4, see https://github.com/shapely/shapely/pull/2034
# this workaround only works if there are no empty points
indices = np.nonzero(mask)[0]
indices = indices - np.arange(len(indices))
coords = np.insert(coords, indices, np.nan, axis=0)
else:
mask = None
if coords.shape[-1] == 2:
dims = "xy"
elif coords.shape[-1] == 3:
dims = "xyz"
else:
raise ValueError(f"Unexpected coords dimensions: {coords.shape}")
extension_metadata: Dict[str, str] = {}
if crs is not None:
extension_metadata["ARROW:extension:metadata"] = json.dumps(
{"crs": crs.to_json()}
)
else:
# In theory this should not be needed, but otherwise pyarrow < 17
# crashes on receiving such data through C Data Interface
# https://github.com/apache/arrow/issues/41741
extension_metadata["ARROW:extension:metadata"] = "{}"
if geom_type == GeometryType.POINT:
parr = _convert_inner_coords(coords, interleaved, dims, mask=mask)
extension_metadata["ARROW:extension:name"] = "geoarrow.point"
field = pa.field(
field_name,
parr.type,
nullable=True,
metadata=extension_metadata,
)
return field, parr
elif geom_type == GeometryType.LINESTRING:
assert len(offsets) == 1, "Expected one offsets array"
(geom_offsets,) = offsets
_parr = _convert_inner_coords(coords, interleaved, dims)
parr = pa.ListArray.from_arrays(
pa.array(geom_offsets), _parr, _linestring_type(_parr.type), mask=mask
)
extension_metadata["ARROW:extension:name"] = "geoarrow.linestring"
field = pa.field(
field_name,
parr.type,
nullable=True,
metadata=extension_metadata,
)
return field, parr
elif geom_type == GeometryType.POLYGON:
assert len(offsets) == 2, "Expected two offsets arrays"
ring_offsets, geom_offsets = offsets
_parr = _convert_inner_coords(coords, interleaved, dims)
_parr1 = pa.ListArray.from_arrays(pa.array(ring_offsets), _parr)
parr = pa.ListArray.from_arrays(pa.array(geom_offsets), _parr1, mask=mask)
parr = parr.cast(_polygon_type(_parr.type))
extension_metadata["ARROW:extension:name"] = "geoarrow.polygon"
field = pa.field(
field_name,
parr.type,
nullable=True,
metadata=extension_metadata,
)
return field, parr
elif geom_type == GeometryType.MULTIPOINT:
assert len(offsets) == 1, "Expected one offsets array"
(geom_offsets,) = offsets
_parr = _convert_inner_coords(coords, interleaved, dims)
parr = pa.ListArray.from_arrays(
pa.array(geom_offsets), _parr, type=_multipoint_type(_parr.type), mask=mask
)
extension_metadata["ARROW:extension:name"] = "geoarrow.multipoint"
field = pa.field(
field_name,
parr.type,
nullable=True,
metadata=extension_metadata,
)
return field, parr
elif geom_type == GeometryType.MULTILINESTRING:
assert len(offsets) == 2, "Expected two offsets arrays"
ring_offsets, geom_offsets = offsets
_parr = _convert_inner_coords(coords, interleaved, dims)
_parr1 = pa.ListArray.from_arrays(pa.array(ring_offsets), _parr)
parr = pa.ListArray.from_arrays(pa.array(geom_offsets), _parr1, mask=mask)
parr = parr.cast(_multilinestring_type(_parr.type))
extension_metadata["ARROW:extension:name"] = "geoarrow.multilinestring"
field = pa.field(
field_name,
parr.type,
nullable=True,
metadata=extension_metadata,
)
return field, parr
elif geom_type == GeometryType.MULTIPOLYGON:
assert len(offsets) == 3, "Expected three offsets arrays"
ring_offsets, polygon_offsets, geom_offsets = offsets
_parr = _convert_inner_coords(coords, interleaved, dims)
_parr1 = pa.ListArray.from_arrays(pa.array(ring_offsets), _parr)
_parr2 = pa.ListArray.from_arrays(pa.array(polygon_offsets), _parr1)
parr = pa.ListArray.from_arrays(pa.array(geom_offsets), _parr2, mask=mask)
parr = parr.cast(_multipolygon_type(_parr.type))
extension_metadata["ARROW:extension:name"] = "geoarrow.multipolygon"
field = pa.field(
field_name,
parr.type,
nullable=True,
metadata=extension_metadata,
)
return field, parr
else:
raise ValueError(f"Unsupported type for geoarrow: {geom_type}")
## GeoArrow -> GeoPandas
def _get_arrow_geometry_field(field):
if (meta := field.metadata) is not None:
if (ext_name := meta.get(b"ARROW:extension:name", None)) is not None:
if ext_name.startswith(b"geoarrow."):
if (
ext_meta := meta.get(b"ARROW:extension:metadata", None)
) is not None:
ext_meta = json.loads(ext_meta.decode())
return ext_name.decode(), ext_meta
if isinstance(field.type, pa.ExtensionType):
ext_name = field.type.extension_name
if ext_name.startswith("geoarrow."):
ext_meta_ser = field.type.__arrow_ext_serialize__()
if ext_meta_ser:
ext_meta = json.loads(ext_meta_ser.decode())
else:
ext_meta = None
return ext_name, ext_meta
return None
def arrow_to_geopandas(table, geometry=None):
"""
Convert Arrow table object to a GeoDataFrame based on GeoArrow extension types.
Parameters
----------
table : pyarrow.Table
The Arrow table to convert.
geometry : str, default None
The name of the geometry column to set as the active geometry
column. If None, the first geometry column found will be used.
Returns
-------
GeoDataFrame
"""
if not isinstance(table, pa.Table):
table = pa.table(table)
geom_fields = []
for i, field in enumerate(table.schema):
geom = _get_arrow_geometry_field(field)
if geom is not None:
geom_fields.append((i, field.name, *geom))
if len(geom_fields) == 0:
raise ValueError("No geometry column found in the Arrow table.")
table_attr = table.drop([f[1] for f in geom_fields])
df = table_attr.to_pandas()
for i, col, ext_name, ext_meta in geom_fields:
crs = None
if ext_meta is not None and "crs" in ext_meta:
crs = ext_meta["crs"]
if ext_name == "geoarrow.wkb":
geom_arr = from_wkb(np.array(table[col]), crs=crs)
elif ext_name.split(".")[1] in GEOARROW_ENCODINGS:
geom_arr = from_shapely(
construct_shapely_array(table[col].combine_chunks(), ext_name), crs=crs
)
else:
raise TypeError(f"Unknown GeoArrow extension type: {ext_name}")
df.insert(i, col, geom_arr)
return GeoDataFrame(df, geometry=geometry or geom_fields[0][1])
def arrow_to_geometry_array(arr):
"""
Convert Arrow array object (representing single GeoArrow array) to a
geopandas GeometryArray.
Specifically for GeoSeries.from_arrow.
"""
if Version(pa.__version__) < Version("14.0.0"):
raise ValueError("Importing from Arrow requires pyarrow >= 14.0.")
schema_capsule, array_capsule = arr.__arrow_c_array__()
field = pa.Field._import_from_c_capsule(schema_capsule)
pa_arr = pa.Array._import_from_c_capsule(field.__arrow_c_schema__(), array_capsule)
geom_info = _get_arrow_geometry_field(field)
if geom_info is None:
raise ValueError("No GeoArrow geometry field found.")
ext_name, ext_meta = geom_info
crs = None
if ext_meta is not None and "crs" in ext_meta:
crs = ext_meta["crs"]
if ext_name == "geoarrow.wkb":
geom_arr = from_wkb(np.array(pa_arr), crs=crs)
elif ext_name.split(".")[1] in GEOARROW_ENCODINGS:
geom_arr = from_shapely(construct_shapely_array(pa_arr, ext_name), crs=crs)
else:
raise ValueError(f"Unknown GeoArrow extension type: {ext_name}")
return geom_arr
def _get_inner_coords(arr):
if pa.types.is_struct(arr.type):
if arr.type.num_fields == 2:
coords = np.column_stack(
[np.asarray(arr.field("x")), np.asarray(arr.field("y"))]
)
else:
coords = np.column_stack(
[
np.asarray(arr.field("x")),
np.asarray(arr.field("y")),
np.asarray(arr.field("z")),
]
)
return coords
else:
# fixed size list
return np.asarray(arr.values).reshape(len(arr), -1)
def construct_shapely_array(arr: pa.Array, extension_name: str):
"""
Construct a NumPy array of shapely geometries from a pyarrow.Array
with GeoArrow extension type.
"""
if isinstance(arr, pa.ExtensionArray):
arr = arr.storage
if extension_name == "geoarrow.point":
coords = _get_inner_coords(arr)
result = shapely.from_ragged_array(GeometryType.POINT, coords, None)
elif extension_name == "geoarrow.linestring":
coords = _get_inner_coords(arr.values)
offsets1 = np.asarray(arr.offsets)
offsets = (offsets1,)
result = shapely.from_ragged_array(GeometryType.LINESTRING, coords, offsets)
elif extension_name == "geoarrow.polygon":
coords = _get_inner_coords(arr.values.values)
offsets2 = np.asarray(arr.offsets)
offsets1 = np.asarray(arr.values.offsets)
offsets = (offsets1, offsets2)
result = shapely.from_ragged_array(GeometryType.POLYGON, coords, offsets)
elif extension_name == "geoarrow.multipoint":
coords = _get_inner_coords(arr.values)
offsets1 = np.asarray(arr.offsets)
offsets = (offsets1,)
result = shapely.from_ragged_array(GeometryType.MULTIPOINT, coords, offsets)
elif extension_name == "geoarrow.multilinestring":
coords = _get_inner_coords(arr.values.values)
offsets2 = np.asarray(arr.offsets)
offsets1 = np.asarray(arr.values.offsets)
offsets = (offsets1, offsets2)
result = shapely.from_ragged_array(
GeometryType.MULTILINESTRING, coords, offsets
)
elif extension_name == "geoarrow.multipolygon":
coords = _get_inner_coords(arr.values.values.values)
offsets3 = np.asarray(arr.offsets)
offsets2 = np.asarray(arr.values.offsets)
offsets1 = np.asarray(arr.values.values.offsets)
offsets = (offsets1, offsets2, offsets3)
result = shapely.from_ragged_array(GeometryType.MULTIPOLYGON, coords, offsets)
else:
raise ValueError(extension_name)
# apply validity mask
if arr.null_count:
mask = np.asarray(arr.is_null())
result = np.where(mask, None, result)
return result