diff --git a/geopandas/tools/sjoin.py b/geopandas/tools/sjoin.py index ff5c966..5369c6c 100644 --- a/geopandas/tools/sjoin.py +++ b/geopandas/tools/sjoin.py @@ -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 diff --git a/geopandas/tools/tests/test_sjoin.py b/geopandas/tools/tests/test_sjoin.py index 83328b1..ee72862 100644 --- a/geopandas/tools/tests/test_sjoin.py +++ b/geopandas/tools/tests/test_sjoin.py @@ -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