Skip to content

Commit c18ddda

Browse files
authored
[GH-2039] Geopandas.GeoSeries: Implement from_wkb, from_wkt, from_xy + temp __init__ fix (#2040)
1 parent 6e8f733 commit c18ddda

3 files changed

Lines changed: 335 additions & 12 deletions

File tree

python/sedona/geopandas/geoseries.py

Lines changed: 269 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -123,9 +123,6 @@ def __init__(
123123
assert not copy
124124
assert not fastpath
125125

126-
self._anchor = data
127-
self._col_label = index
128-
129126
data_crs = None
130127
if hasattr(data, "crs"):
131128
data_crs = data.crs
@@ -136,6 +133,20 @@ def __init__(
136133
"allow_override=True)' to overwrite CRS or "
137134
"'GeoSeries.to_crs(crs)' to reproject geometries. "
138135
)
136+
# This is a temporary workaround since pyspark errors when creating a ps.Series from a ps.Series
137+
# This is NOT a scalable solution since we call to_pandas() on the data and is a hacky solution
138+
# but this should be resolved if/once https://github.com/apache/spark/pull/51300 is merged in.
139+
# For now, we reset self._anchor = data to have keep the geometry information (e.g crs) that's lost in to_pandas()
140+
super().__init__(
141+
data=data.to_pandas(),
142+
index=index,
143+
dtype=dtype,
144+
name=name,
145+
copy=copy,
146+
fastpath=fastpath,
147+
)
148+
149+
self._anchor = data
139150
else:
140151
if isinstance(data, pd.Series):
141152
assert index is None
@@ -1360,7 +1371,82 @@ def from_wkb(
13601371
on_invalid="raise",
13611372
**kwargs,
13621373
) -> "GeoSeries":
1363-
raise NotImplementedError("GeoSeries.from_wkb() is not implemented yet.")
1374+
r"""
1375+
Alternate constructor to create a ``GeoSeries``
1376+
from a list or array of WKB objects
1377+
1378+
Parameters
1379+
----------
1380+
data : array-like or Series
1381+
Series, list or array of WKB objects
1382+
index : array-like or Index
1383+
The index for the GeoSeries.
1384+
crs : value, optional
1385+
Coordinate Reference System of the geometry objects. Can be anything
1386+
accepted by
1387+
:meth:`pyproj.CRS.from_user_input() <pyproj.crs.CRS.from_user_input>`,
1388+
such as an authority string (eg "EPSG:4326") or a WKT string.
1389+
on_invalid: {"raise", "warn", "ignore"}, default "raise"
1390+
- raise: an exception will be raised if a WKB input geometry is invalid.
1391+
- warn: a warning will be raised and invalid WKB geometries will be returned
1392+
as None.
1393+
- ignore: invalid WKB geometries will be returned as None without a warning.
1394+
- fix: an effort is made to fix invalid input geometries (e.g. close
1395+
unclosed rings). If this is not possible, they are returned as ``None``
1396+
without a warning. Requires GEOS >= 3.11 and shapely >= 2.1.
1397+
1398+
kwargs
1399+
Additional arguments passed to the Series constructor,
1400+
e.g. ``name``.
1401+
1402+
Returns
1403+
-------
1404+
GeoSeries
1405+
1406+
See Also
1407+
--------
1408+
GeoSeries.from_wkt
1409+
1410+
Examples
1411+
--------
1412+
1413+
>>> wkbs = [
1414+
... (
1415+
... b"\x01\x01\x00\x00\x00\x00\x00\x00\x00"
1416+
... b"\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\xf0?"
1417+
... ),
1418+
... (
1419+
... b"\x01\x01\x00\x00\x00\x00\x00\x00\x00"
1420+
... b"\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x00@"
1421+
... ),
1422+
... (
1423+
... b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00"
1424+
... b"\x00\x08@\x00\x00\x00\x00\x00\x00\x08@"
1425+
... ),
1426+
... ]
1427+
>>> s = geopandas.GeoSeries.from_wkb(wkbs)
1428+
>>> s
1429+
0 POINT (1 1)
1430+
1 POINT (2 2)
1431+
2 POINT (3 3)
1432+
dtype: geometry
1433+
"""
1434+
if on_invalid != "raise":
1435+
raise NotImplementedError(
1436+
"GeoSeries.from_wkb(): only on_invalid='raise' is implemented"
1437+
)
1438+
1439+
from pyspark.sql.types import StructType, StructField, BinaryType
1440+
1441+
schema = StructType([StructField("data", BinaryType(), True)])
1442+
return cls._create_from_select(
1443+
f"ST_GeomFromWKB(`data`)",
1444+
data,
1445+
schema,
1446+
index,
1447+
crs,
1448+
**kwargs,
1449+
)
13641450

13651451
@classmethod
13661452
def from_wkt(
@@ -1371,11 +1457,154 @@ def from_wkt(
13711457
on_invalid="raise",
13721458
**kwargs,
13731459
) -> "GeoSeries":
1374-
raise NotImplementedError("GeoSeries.from_wkt() is not implemented yet.")
1460+
"""
1461+
Alternate constructor to create a ``GeoSeries``
1462+
from a list or array of WKT objects
1463+
1464+
Parameters
1465+
----------
1466+
data : array-like, Series
1467+
Series, list, or array of WKT objects
1468+
index : array-like or Index
1469+
The index for the GeoSeries.
1470+
crs : value, optional
1471+
Coordinate Reference System of the geometry objects. Can be anything
1472+
accepted by
1473+
:meth:`pyproj.CRS.from_user_input() <pyproj.crs.CRS.from_user_input>`,
1474+
such as an authority string (eg "EPSG:4326") or a WKT string.
1475+
on_invalid : {"raise", "warn", "ignore"}, default "raise"
1476+
- raise: an exception will be raised if a WKT input geometry is invalid.
1477+
- warn: a warning will be raised and invalid WKT geometries will be
1478+
returned as ``None``.
1479+
- ignore: invalid WKT geometries will be returned as ``None`` without a
1480+
warning.
1481+
- fix: an effort is made to fix invalid input geometries (e.g. close
1482+
unclosed rings). If this is not possible, they are returned as ``None``
1483+
without a warning. Requires GEOS >= 3.11 and shapely >= 2.1.
1484+
1485+
kwargs
1486+
Additional arguments passed to the Series constructor,
1487+
e.g. ``name``.
1488+
1489+
Returns
1490+
-------
1491+
GeoSeries
1492+
1493+
See Also
1494+
--------
1495+
GeoSeries.from_wkb
1496+
1497+
Examples
1498+
--------
1499+
1500+
>>> wkts = [
1501+
... 'POINT (1 1)',
1502+
... 'POINT (2 2)',
1503+
... 'POINT (3 3)',
1504+
... ]
1505+
>>> s = geopandas.GeoSeries.from_wkt(wkts)
1506+
>>> s
1507+
0 POINT (1 1)
1508+
1 POINT (2 2)
1509+
2 POINT (3 3)
1510+
dtype: geometry
1511+
"""
1512+
if on_invalid != "raise":
1513+
raise NotImplementedError(
1514+
"GeoSeries.from_wkt(): only on_invalid='raise' is implemented"
1515+
)
1516+
1517+
from pyspark.sql.types import StructType, StructField, StringType
1518+
1519+
schema = StructType([StructField("data", StringType(), True)])
1520+
return cls._create_from_select(
1521+
f"ST_GeomFromText(`data`)",
1522+
data,
1523+
schema,
1524+
index,
1525+
crs,
1526+
**kwargs,
1527+
)
13751528

13761529
@classmethod
13771530
def from_xy(cls, x, y, z=None, index=None, crs=None, **kwargs) -> "GeoSeries":
1378-
raise NotImplementedError("GeoSeries.from_xy() is not implemented yet.")
1531+
"""
1532+
Alternate constructor to create a :class:`~geopandas.GeoSeries` of Point
1533+
geometries from lists or arrays of x, y(, z) coordinates
1534+
1535+
In case of geographic coordinates, it is assumed that longitude is captured
1536+
by ``x`` coordinates and latitude by ``y``.
1537+
1538+
Parameters
1539+
----------
1540+
x, y, z : iterable
1541+
index : array-like or Index, optional
1542+
The index for the GeoSeries. If not given and all coordinate inputs
1543+
are Series with an equal index, that index is used.
1544+
crs : value, optional
1545+
Coordinate Reference System of the geometry objects. Can be anything
1546+
accepted by
1547+
:meth:`pyproj.CRS.from_user_input() <pyproj.crs.CRS.from_user_input>`,
1548+
such as an authority string (eg "EPSG:4326") or a WKT string.
1549+
**kwargs
1550+
Additional arguments passed to the Series constructor,
1551+
e.g. ``name``.
1552+
1553+
Returns
1554+
-------
1555+
GeoSeries
1556+
1557+
See Also
1558+
--------
1559+
GeoSeries.from_wkt
1560+
points_from_xy
1561+
1562+
Examples
1563+
--------
1564+
1565+
>>> x = [2.5, 5, -3.0]
1566+
>>> y = [0.5, 1, 1.5]
1567+
>>> s = geopandas.GeoSeries.from_xy(x, y, crs="EPSG:4326")
1568+
>>> s
1569+
0 POINT (2.5 0.5)
1570+
1 POINT (5 1)
1571+
2 POINT (-3 1.5)
1572+
dtype: geometry
1573+
"""
1574+
from pyspark.sql.types import StructType, StructField, DoubleType
1575+
1576+
schema = StructType(
1577+
[StructField("x", DoubleType(), True), StructField("y", DoubleType(), True)]
1578+
)
1579+
1580+
# Spark doesn't automatically cast ints to floats for us
1581+
x = [float(num) for num in x]
1582+
y = [float(num) for num in y]
1583+
z = [float(num) for num in z] if z else None
1584+
1585+
if z:
1586+
data = list(zip(x, y, z))
1587+
select = f"ST_PointZ(`x`, `y`, `z`)"
1588+
schema.add(StructField("z", DoubleType(), True))
1589+
else:
1590+
data = list(zip(x, y))
1591+
select = f"ST_Point(`x`, `y`)"
1592+
1593+
geoseries = cls._create_from_select(
1594+
select,
1595+
data,
1596+
schema,
1597+
index,
1598+
crs,
1599+
**kwargs,
1600+
)
1601+
1602+
if crs:
1603+
from pyproj import CRS
1604+
1605+
geoseries.crs = CRS.from_user_input(crs).to_epsg()
1606+
1607+
return geoseries
13791608

13801609
@classmethod
13811610
def from_shapely(
@@ -1387,6 +1616,40 @@ def from_shapely(
13871616
def from_arrow(cls, arr, **kwargs) -> "GeoSeries":
13881617
raise NotImplementedError("GeoSeries.from_arrow() is not implemented yet.")
13891618

1619+
@classmethod
1620+
def _create_from_select(
1621+
cls, select: str, data, schema, index, crs, **kwargs
1622+
) -> "GeoSeries":
1623+
1624+
from pyspark.pandas.utils import default_session
1625+
from pyspark.pandas.internal import InternalField
1626+
import numpy as np
1627+
1628+
if isinstance(data, list) and not isinstance(data[0], (tuple, list)):
1629+
data = [(obj,) for obj in data]
1630+
1631+
select = f"{select} as geometry"
1632+
1633+
spark_df = default_session().createDataFrame(data, schema=schema)
1634+
spark_df = spark_df.selectExpr(select)
1635+
1636+
internal = InternalFrame(
1637+
spark_frame=spark_df,
1638+
index_spark_columns=None,
1639+
column_labels=[("geometry",)],
1640+
data_spark_columns=[scol_for(spark_df, "geometry")],
1641+
data_fields=[
1642+
InternalField(np.dtype("object"), spark_df.schema["geometry"])
1643+
],
1644+
column_label_names=[("geometry",)],
1645+
)
1646+
return GeoSeries(
1647+
first_series(PandasOnSparkDataFrame(internal)),
1648+
index,
1649+
crs=crs,
1650+
name=kwargs.get("name", None),
1651+
)
1652+
13901653
def to_file(
13911654
self,
13921655
filename: Union[os.PathLike, typing.IO],

python/tests/geopandas/test_geoseries.py

Lines changed: 37 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,13 +119,47 @@ def test_from_file(self):
119119
pass
120120

121121
def test_from_wkb(self):
122-
pass
122+
wkbs = [
123+
(
124+
b"\x01\x01\x00\x00\x00\x00\x00\x00\x00"
125+
b"\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\xf0?"
126+
),
127+
(
128+
b"\x01\x01\x00\x00\x00\x00\x00\x00\x00"
129+
b"\x00\x00\x00@\x00\x00\x00\x00\x00\x00\x00@"
130+
),
131+
(
132+
b"\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00"
133+
b"\x00\x08@\x00\x00\x00\x00\x00\x00\x08@"
134+
),
135+
]
136+
s = sgpd.GeoSeries.from_wkb(wkbs)
137+
expected = gpd.GeoSeries([Point(1, 1), Point(2, 2), Point(3, 3)])
138+
self.check_sgpd_equals_gpd(s, expected)
123139

124140
def test_from_wkt(self):
125-
pass
141+
wkts = [
142+
"POINT (1 1)",
143+
"POINT (2 2)",
144+
"POINT (3 3)",
145+
]
146+
s = sgpd.GeoSeries.from_wkt(wkts)
147+
expected = gpd.GeoSeries([Point(1, 1), Point(2, 2), Point(3, 3)])
148+
self.check_sgpd_equals_gpd(s, expected)
126149

127150
def test_from_xy(self):
128-
pass
151+
x = [2.5, 5, -3.0]
152+
y = [0.5, 1, 1.5]
153+
s = sgpd.GeoSeries.from_xy(x, y, crs="EPSG:4326")
154+
expected = gpd.GeoSeries([Point(2.5, 0.5), Point(5, 1), Point(-3, 1.5)])
155+
self.check_sgpd_equals_gpd(s, expected)
156+
157+
z = [1, 2, 3]
158+
s = sgpd.GeoSeries.from_xy(x, y, z)
159+
expected = gpd.GeoSeries(
160+
[Point(2.5, 0.5, 1), Point(5, 1, 2), Point(-3, 1.5, 3)]
161+
)
162+
self.check_sgpd_equals_gpd(s, expected)
129163

130164
def test_from_shapely(self):
131165
pass

0 commit comments

Comments
 (0)