wassname/geopandas
1
0
Fork 0
mirror of https://github.com/wassname/geopandas.git synced 2026-09-21 13:00:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

BUG: Retain primary index name in sjoin (#1150)

Browse Source 
* Partial solution toward #846

* And now with more _style_

* TST: merge tests from #1159
This commit is contained in:
Brendan Ward
2019-10-15 09:28:37 +01:00
committed by Martin Fleischmann
parent e0b7205808
commit 32609cb4cc
2 changed files with 81 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
geopandas/tools/sjoin.py
+29 -18
View File
@@ -90,9 +90,11 @@ def sjoin(
# and store references to the original indices, to be reaffixed later.
# GH 352
left_df = left_df.copy(deep=True)
left_index_name = left_df.index.name
left_df.index = left_df.index.rename(index_left)
left_df = left_df.reset_index()
right_df = right_df.copy(deep=True)
right_index_name = right_df.index.name
right_df.index = right_df.index.rename(index_right)
right_df = right_df.reset_index()
@@ -171,27 +173,35 @@ def sjoin(
if how == "inner":
result = result.set_index("_key_left")
joined = left_df.merge(result, left_index=True, right_index=True).merge(
right_df.drop(right_df.geometry.name, axis=1),
left_on="_key_right",
right_index=True,
suffixes=("_%s" % lsuffix, "_%s" % rsuffix),
joined = (
left_df.merge(result, left_index=True, right_index=True)
.merge(
right_df.drop(right_df.geometry.name, axis=1),
left_on="_key_right",
right_index=True,
suffixes=("_%s" % lsuffix, "_%s" % rsuffix),
)
.set_index(index_left)
.drop(["_key_right"], axis=1)
)
joined = joined.set_index(index_left).drop(["_key_right"], axis=1)
joined.index.name = None
joined.index.name = left_index_name
elif how == "left":
result = result.set_index("_key_left")
joined = left_df.merge(
result, left_index=True, right_index=True, how="left"
).merge(
right_df.drop(right_df.geometry.name, axis=1),
how="left",
left_on="_key_right",
right_index=True,
suffixes=("_%s" % lsuffix, "_%s" % rsuffix),
joined = (
left_df.merge(result, left_index=True, right_index=True, how="left")
.merge(
right_df.drop(right_df.geometry.name, axis=1),
how="left",
left_on="_key_right",
right_index=True,
suffixes=("_%s" % lsuffix, "_%s" % rsuffix),
)
.set_index(index_left)
.drop(["_key_right"], axis=1)
)
joined = joined.set_index(index_left).drop(["_key_right"], axis=1)
joined.index.name = None
joined.index.name = left_index_name
else: # how == 'right':
joined = (
left_df.drop(left_df.geometry.name, axis=1)
@@ -204,7 +214,8 @@ def sjoin(
how="right",
)
.set_index(index_right)
.drop(["_key_left", "_key_right"], axis=1)
)
joined = joined.drop(["_key_left", "_key_right"], axis=1)
joined.index.name = right_index_name
return joined
geopandas/tools/tests/test_sjoin.py
+52 -13
View File
@@ -36,10 +36,15 @@ def dfs(request):
df1 = GeoDataFrame({"geometry": polys1, "df1": [0, 1, 2]})
df2 = GeoDataFrame({"geometry": polys2, "df2": [3, 4, 5]})
if request.param == "string-index":
df1.index = ["a", "b", "c"]
df2.index = ["d", "e", "f"]
if request.param == "named-index":
df1.index.name = "df1_ix"
df2.index.name = "df2_ix"
# construction expected frames
expected = {}
@@ -79,7 +84,9 @@ class TestSpatialJoin:
with pytest.warns(UserWarning):
sjoin(df1, df2)
@pytest.mark.parametrize("dfs", ["default-index", "string-index"], indirect=True)
@pytest.mark.parametrize(
"dfs", ["default-index", "string-index", "named-index"], indirect=True
)
@pytest.mark.parametrize("op", ["intersects", "contains", "within"])
def test_inner(self, op, dfs):
index, df1, df2, expected = dfs
@@ -93,27 +100,40 @@ class TestSpatialJoin:
exp[["index_left", "index_right"]] = exp[
["index_left", "index_right"]
].astype("int64")
exp = exp.set_index("index_left")
exp.index.name = None
if index == "named-index":
exp[["df1_ix", "df2_ix"]] = exp[["df1_ix", "df2_ix"]].astype("int64")
exp = exp.set_index("df1_ix").rename(columns={"df2_ix": "index_right"})
if index in ["default-index", "string-index"]:
exp = exp.set_index("index_left")
exp.index.name = None
assert_frame_equal(res, exp)
@pytest.mark.parametrize("dfs", ["default-index", "string-index"], indirect=True)
@pytest.mark.parametrize(
"dfs", ["default-index", "string-index", "named-index"], indirect=True
)
@pytest.mark.parametrize("op", ["intersects", "contains", "within"])
def test_left(self, op, dfs):
index, df1, df2, expected = dfs
res = sjoin(df1, df2, how="left", op=op)
exp = expected[op].dropna(subset=["index_left"]).copy()
if index in ["default-index", "string-index"]:
exp = expected[op].dropna(subset=["index_left"]).copy()
elif index == "named-index":
exp = expected[op].dropna(subset=["df1_ix"]).copy()
exp = exp.drop("geometry_y", axis=1).rename(columns={"geometry_x": "geometry"})
exp["df1"] = exp["df1"].astype("int64")
if index == "default-index":
exp["index_left"] = exp["index_left"].astype("int64")
# TODO: in result the dtype is object
res["index_right"] = res["index_right"].astype(float)
exp = exp.set_index("index_left")
exp.index.name = None
elif index == "named-index":
exp[["df1_ix"]] = exp[["df1_ix"]].astype("int64")
exp = exp.set_index("df1_ix").rename(columns={"df2_ix": "index_right"})
if index in ["default-index", "string-index"]:
exp = exp.set_index("index_left")
exp.index.name = None
assert_frame_equal(res, exp)
@@ -146,21 +166,31 @@ class TestSpatialJoin:
with pytest.raises(ValueError, match="'right_df' should be GeoDataFrame"):
sjoin(df1, df2.geometry)
@pytest.mark.parametrize("dfs", ["default-index", "string-index"], indirect=True)
@pytest.mark.parametrize(
"dfs", ["default-index", "string-index", "named-index"], indirect=True
)
@pytest.mark.parametrize("op", ["intersects", "contains", "within"])
def test_right(self, op, dfs):
index, df1, df2, expected = dfs
res = sjoin(df1, df2, how="right", op=op)
exp = expected[op].dropna(subset=["index_right"]).copy()
if index in ["default-index", "string-index"]:
exp = expected[op].dropna(subset=["index_right"]).copy()
elif index == "named-index":
exp = expected[op].dropna(subset=["df2_ix"]).copy()
exp = exp.drop("geometry_x", axis=1).rename(columns={"geometry_y": "geometry"})
exp["df2"] = exp["df2"].astype("int64")
if index == "default-index":
exp["index_right"] = exp["index_right"].astype("int64")
res["index_left"] = res["index_left"].astype(float)
exp = exp.set_index("index_right")
exp = exp.reindex(columns=res.columns)
elif index == "named-index":
exp[["df2_ix"]] = exp[["df2_ix"]].astype("int64")
exp = exp.set_index("df2_ix").rename(columns={"df1_ix": "index_left"})
if index in ["default-index", "string-index"]:
exp = exp.set_index("index_right")
exp = exp.reindex(columns=res.columns)
exp.index.name = None
assert_frame_equal(res, exp, check_index_type=False)
@@ -244,9 +274,18 @@ class TestSpatialJoinNYBB:
# original index names should be unchanged
pointdf2 = self.pointdf.copy()
pointdf2.index.name = "pointid"
sjoin(pointdf2, self.polydf, how=how)
polydf = self.polydf.copy()
polydf.index.name = "polyid"
res = sjoin(pointdf2, polydf, how=how)
assert pointdf2.index.name == "pointid"
assert self.polydf.index.name is None
assert polydf.index.name == "polyid"
# original index name should pass through to result
if how == "right":
assert res.index.name == "polyid"
else: # how == "left", how == "inner"
assert res.index.name == "pointid"
def test_sjoin_values(self):
# GH190