Skip to content

Commit 63c373c

Browse files
Implement drivers for geoparquet.
1 parent 9c0be50 commit 63c373c

File tree

2 files changed

+49
-0
lines changed

2 files changed

+49
-0
lines changed

pins/drivers.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,16 @@ def _assert_is_pandas_df(x):
2323
)
2424

2525

26+
def _assert_is_geopandas_df(x):
27+
# Assume we have already protected against uninstalled geopandas
28+
import geopandas as gpd
29+
30+
if not isinstance(x, gpd.GeoDataFrame):
31+
raise NotImplementedError(
32+
"Currently only geopandas.GeoDataFrame can be saved to a GeoParquet."
33+
)
34+
35+
2636
def load_path(meta, path_to_version):
2737
# Check that only a single file name was given
2838
fnames = [meta.file] if isinstance(meta.file, str) else meta.file
@@ -105,6 +115,17 @@ def load_data(
105115

106116
return pd.read_csv(f)
107117

118+
elif meta.type == "geoparquet":
119+
try:
120+
import geopandas as gpd
121+
except ModuleNotFoundError:
122+
raise ModuleNotFoundError(
123+
'The "geopandas" package is required to read "geoparquet" type '
124+
"files."
125+
) from None
126+
127+
return gpd.read_parquet(f)
128+
108129
elif meta.type == "joblib":
109130
import joblib
110131

@@ -147,6 +168,8 @@ def save_data(
147168
if apply_suffix:
148169
if type == "file":
149170
suffix = "".join(Path(obj).suffixes)
171+
elif type == "geoparquet":
172+
suffix = ".parquet"
150173
else:
151174
suffix = f".{type}"
152175
else:
@@ -178,6 +201,11 @@ def save_data(
178201

179202
obj.to_parquet(final_name)
180203

204+
elif type == "geoparquet":
205+
_assert_is_geopandas_df(obj)
206+
207+
obj.to_parquet(final_name)
208+
181209
elif type == "joblib":
182210
import joblib
183211

pins/tests/test_drivers.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,27 @@ def test_driver_roundtrip(tmp_dir2, type_):
7777
assert df.equals(obj)
7878

7979

80+
def test_driver_geoparquet_roundtrip(tmp_dir2):
81+
import geopandas as gpd
82+
83+
gdf = gpd.GeoDataFrame(
84+
{"x": [1, 2, 3], "geometry": gpd.points_from_xy([1, 2, 3], [1, 2, 3])}
85+
)
86+
87+
fname = "some_gdf"
88+
full_file = f"{fname}.parquet"
89+
90+
p_obj = tmp_dir2 / fname
91+
res_fname = save_data(gdf, p_obj, "geoparquet")
92+
93+
assert Path(res_fname).name == full_file
94+
95+
meta = MetaRaw(full_file, "geoparquet", "my_pin")
96+
obj = load_data(meta, fsspec.filesystem("file"), tmp_dir2, allow_pickle_read=True)
97+
98+
assert gdf.equals(obj)
99+
100+
80101
@pytest.mark.parametrize(
81102
"type_",
82103
[

0 commit comments

Comments
 (0)