Repository: geopandas/geopandas
Pull request #3347 (2024-06-22 00:08:39)

TST: Add test covering passing geojson string to read file
Adds a test for #3344


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4391142c..e64df83e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -62,6 +62,9 @@ New features and improvements:
 - `explore` now supports `GeoDataFrame`s with additional columns containing datetimes, uuids and
   other non JSON serializable objects (#3261).
 - The `GeoSeries.fillna` method now supports the `limit` keyword (#3290).
+- Added ``on_attribute`` option argument to the ``sjoin()``
+  method, allowing to restrict joins to the observations with 
+  matching attributes. (#3231)
 - Added support for `bbox` covering encoding in geoparquet. Can filter reading of parquet
 files based on a bounding box, and write out a bounding box column to parquet files (#3282).
 - `align` keyword in binary methods now defaults to `None`, treated as True. Explicit True
diff --git a/geopandas/geodataframe.py b/geopandas/geodataframe.py
index 4511f0b0..339315d6 100644
--- a/geopandas/geodataframe.py
+++ b/geopandas/geodataframe.py
@@ -2329,6 +2329,16 @@ properties': {'col1': 'name1'}, 'geometry': {'type': 'Point', 'coordinates': (1.
             Suffix to apply to overlapping column names (left GeoDataFrame).
         rsuffix : string, default 'right'
             Suffix to apply to overlapping column names (right GeoDataFrame).
+        distance : number or array_like, optional
+            Distance(s) around each input geometry within which to query the tree
+            for the 'dwithin' predicate. If array_like, must be
+            one-dimesional with length equal to length of left GeoDataFrame.
+            Required if ``predicate='dwithin'``.
+        on_attribute : string, list or tuple
+            Column name(s) to join on as an additional join restriction on top
+            of the spatial predicate. These must be found in both DataFrames.
+            If set, observations are joined only if the predicate applies
+            and values in specified columns match.
 
         Examples
         --------
diff --git a/geopandas/io/file.py b/geopandas/io/file.py
index dae1f46b..129f0f73 100644
--- a/geopandas/io/file.py
+++ b/geopandas/io/file.py
@@ -66,13 +66,11 @@ def _import_fiona():
 
 pyogrio = None
 pyogrio_import_error = None
-PYOGRIO_GE_081 = False
 
 
 def _import_pyogrio():
     global pyogrio
     global pyogrio_import_error
-    global PYOGRIO_GE_081
 
     if pyogrio is None:
         try:
@@ -81,10 +79,6 @@ def _import_pyogrio():
         except ImportError as err:
             pyogrio = False
             pyogrio_import_error = str(err)
-        else:
-            PYOGRIO_GE_081 = Version(
-                Version(pyogrio.__version__).base_version
-            ) >= Version("0.8.1")
 
 
 def _check_fiona(func):
diff --git a/geopandas/tools/sjoin.py b/geopandas/tools/sjoin.py
index 7a3c70c3..06d7ef74 100644
--- a/geopandas/tools/sjoin.py
+++ b/geopandas/tools/sjoin.py
@@ -18,6 +18,7 @@ def sjoin(
     lsuffix="left",
     rsuffix="right",
     distance=None,
+    on_attribute=None,
     **kwargs,
 ):
     """Spatial join of two GeoDataFrames.
@@ -50,6 +51,11 @@ def sjoin(
         for the 'dwithin' predicate. If array_like, must be
         one-dimesional with length equal to length of left GeoDataFrame.
         Required if ``predicate='dwithin'``.
+    on_attribute : string, list or tuple
+        Column name(s) to join on as an additional join restriction on top
+        of the spatial predicate. These must be found in both DataFrames.
+        If set, observations are joined only if the predicate applies
+        and values in specified columns match.
 
     Examples
     --------
@@ -103,18 +109,38 @@ def sjoin(
         first = next(iter(kwargs.keys()))
         raise TypeError(f"sjoin() got an unexpected keyword argument '{first}'")
 
-    _basic_checks(left_df, right_df, how, lsuffix, rsuffix)
+    on_attribute = _maybe_make_list(on_attribute)
+
+    _basic_checks(left_df, right_df, how, lsuffix, rsuffix, on_attribute=on_attribute),
 
-    indices = _geom_predicate_query(left_df, right_df, predicate, distance)
+    indices = _geom_predicate_query(
+        left_df, right_df, predicate, distance, on_attribute=on_attribute
+    )
 
     joined, _ = _frame_join(
-        left_df, right_df, indices, None, how, lsuffix, rsuffix, predicate
+        left_df,
+        right_df,
+        indices,
+        None,
+        how,
+        lsuffix,
+        rsuffix,
+        predicate,
+        on_attribute=on_attribute,
     )
 
     return joined
 
 
-def _basic_checks(left_df, right_df, how, lsuffix, rsuffix):
+def _maybe_make_list(obj):
+    if isinstance(obj, tuple):
+        return list(obj)
+    if obj is not None and not isinstance(obj, list):
+        return [obj]
+    return obj
+
+
+def _basic_checks(left_df, right_df, how, lsuffix, rsuffix, on_attribute=None):
     """Checks the validity of join input parameters.
 
     `how` must be one of the valid options.
@@ -131,6 +157,8 @@ def _basic_checks(left_df, right_df, how, lsuffix, rsuffix):
         left index suffix
     rsuffix : str
         right index suffix
+    on_attribute : list, default None
+        list of column names to merge on along with geometry
     """
     if not isinstance(left_df, GeoDataFrame):
         raise ValueError(
@@ -151,8 +179,28 @@ def _basic_checks(left_df, right_df, how, lsuffix, rsuffix):
     if not _check_crs(left_df, right_df):
         _crs_mismatch_warn(left_df, right_df, stacklevel=4)
 
+    if on_attribute:
+        for attr in on_attribute:
+            if (attr not in left_df) and (attr not in right_df):
+                raise ValueError(
+                    f"Expected column {attr} is missing from both of the dataframes."
+                )
+            if attr not in left_df:
+                raise ValueError(
+                    f"Expected column {attr} is missing from the left dataframe."
+                )
+            if attr not in right_df:
+                raise ValueError(
+                    f"Expected column {attr} is missing from the right dataframe."
+                )
+            if attr in (left_df.geometry.name, right_df.geometry.name):
+                raise ValueError(
+                    "Active geometry column cannot be used as an input "
+                    "for on_attribute parameter."
+                )
+
 
-def _geom_predicate_query(left_df, right_df, predicate, distance):
+def _geom_predicate_query(left_df, right_df, predicate, distance, on_attribute=None):
     """Compute geometric comparisons and get matching indices.
 
     Parameters
@@ -161,6 +209,9 @@ def _geom_predicate_query(left_df, right_df, predicate, distance):
     right_df : GeoDataFrame
     predicate : string
         Binary predicate to query.
+    on_attribute: list, default None
+        list of column names to merge on along with geometry
+
 
     Returns
     -------
@@ -200,6 +251,12 @@ def _geom_predicate_query(left_df, right_df, predicate, distance):
         l_idx = l_idx[indexer]
         r_idx = r_idx[indexer]
 
+    if on_attribute:
+        for attr in on_attribute:
+            (l_idx, r_idx), _ = _filter_shared_attribute(
+                left_df, right_df, l_idx, r_idx, attr
+            )
+
     return l_idx, r_idx
 
 
@@ -356,7 +413,15 @@ def _adjust_indexers(indices, distances, original_length, how, predicate):
 
 
 def _frame_join(
-    left_df, right_df, indices, distances, how, lsuffix, rsuffix, predicate
+    left_df,
+    right_df,
+    indices,
+    distances,
+    how,
+    lsuffix,
+    rsuffix,
+    predicate,
+    on_attribute=None,
 ):
     """Join the GeoDataFrames at the DataFrame level.
 
@@ -376,12 +441,18 @@ def _frame_join(
         Suffix to apply to overlapping column names (left GeoDataFrame).
     rsuffix : string
         Suffix to apply to overlapping column names (right GeoDataFrame).
+    on_attribute: list, default None
+        list of column names to merge on along with geometry
+
 
     Returns
     -------
     GeoDataFrame
         Joined GeoDataFrame.
     """
+    if on_attribute:  # avoid renaming or duplicating shared column
+        right_df = right_df.drop(on_attribute, axis=1)
+
     if how in ("inner", "left"):
         right_df = right_df.drop(right_df.geometry.name, axis=1)
     else:  # how == 'right':
@@ -407,7 +478,6 @@ def _frame_join(
     )
     left_df.columns = left_column_names
     right_df.columns = right_column_names
-
     left_index = left_df.columns[:left_nlevels]
     right_index = right_df.columns[:right_nlevels]
 
@@ -444,6 +514,7 @@ def _nearest_query(
     how: str,
     return_distance: bool,
     exclusive: bool,
+    on_attribute: Optional[list] = None,
 ):
     # use the opposite of the join direction for the index
     use_left_as_sindex = how == "right"
@@ -481,9 +552,32 @@ def _nearest_query(
             distances = np.array([], dtype=np.float64)
         else:
             distances = None
+
+    if on_attribute:
+        for attr in on_attribute:
+            (l_idx, r_idx), shared_attribute_rows = _filter_shared_attribute(
+                left_df, right_df, l_idx, r_idx, attr
+            )
+            distances = distances[shared_attribute_rows]
+
     return (l_idx, r_idx), distances
 
 
+def _filter_shared_attribute(left_df, right_df, l_idx, r_idx, attribute):
+    """
+    Returns the indices for the left and right dataframe that share the same entry
+    in the attribute column. Also returns a Boolean `shared_attribute_rows` for rows
+    with the same entry.
+    """
+    shared_attribute_rows = (
+        left_df[attribute].iloc[l_idx].values == right_df[attribute].iloc[r_idx].values
+    )
+
+    l_idx = l_idx[shared_attribute_rows]
+    r_idx = r_idx[shared_attribute_rows]
+    return (l_idx, r_idx), shared_attribute_rows
+
+
 def sjoin_nearest(
     left_df: GeoDataFrame,
     right_df: GeoDataFrame,
@@ -607,6 +701,7 @@ chicago_w_groceries[chicago_w_groceries["community"] == "UPTOWN"]
     Every operation in GeoPandas is planar, i.e. the potential third
     dimension is not taken into account.
     """
+
     _basic_checks(left_df, right_df, how, lsuffix, rsuffix)
 
     left_df.geometry.values.check_geographic_crs(stacklevel=1)
@@ -615,10 +710,22 @@ chicago_w_groceries[chicago_w_groceries["community"] == "UPTOWN"]
     return_distance = distance_col is not None
 
     indices, distances = _nearest_query(
-        left_df, right_df, max_distance, how, return_distance, exclusive
+        left_df,
+        right_df,
+        max_distance,
+        how,
+        return_distance,
+        exclusive,
     )
     joined, distances = _frame_join(
-        left_df, right_df, indices, distances, how, lsuffix, rsuffix, None
+        left_df,
+        right_df,
+        indices,
+        distances,
+        how,
+        lsuffix,
+        rsuffix,
+        None,
     )
 
     if return_distance:


Test patch:
diff --git a/geopandas/io/tests/test_file.py b/geopandas/io/tests/test_file.py
index 8072cf6c..d374d4d6 100644
--- a/geopandas/io/tests/test_file.py
+++ b/geopandas/io/tests/test_file.py
@@ -1,5 +1,6 @@
 import datetime
 import io
+import json
 import os
 import pathlib
 import shutil
@@ -17,7 +18,7 @@ from shapely.geometry import Point, Polygon, box, mapping
 import geopandas
 from geopandas import GeoDataFrame, read_file
 from geopandas._compat import HAS_PYPROJ, PANDAS_GE_20, PANDAS_GE_30
-from geopandas.io.file import _EXTENSION_TO_DRIVER, PYOGRIO_GE_081, _detect_driver
+from geopandas.io.file import _EXTENSION_TO_DRIVER, _detect_driver
 
 import pytest
 from geopandas.testing import assert_geodataframe_equal, assert_geoseries_equal
@@ -26,8 +27,15 @@ from pandas.testing import assert_frame_equal, assert_series_equal
 
 try:
     import pyogrio
+
+    # those version checks have to be defined here instead of imported from
+    # geopandas.io.file (those are only initialized lazily on first usage)
+    PYOGRIO_GE_090 = Version(Version(pyogrio.__version__).base_version) >= Version(
+        "0.9.0"
+    )
 except ImportError:
     pyogrio = False
+    PYOGRIO_GE_090 = False
 
 
 try:
@@ -606,6 +614,25 @@ def test_read_file_local_uri(file_path, engine):
     assert isinstance(gdf, geopandas.GeoDataFrame)
 
 
+@pytest.mark.skipif(not HAS_PYPROJ, reason="pyproj not installed")
+def test_read_file_geojson_string_path(engine):
+    if engine == "pyogrio" and not PYOGRIO_GE_090:
+        pytest.skip("fixed in pyogrio 0.9.0")
+    expected = GeoDataFrame({"val_with_hash": ["row # 0"], "geometry": [Point(0, 1)]})
+    features = {
+        "type": "FeatureCollection",
+        "features": [
+            {
+                "type": "Feature",
+                "properties": {"val_with_hash": "row # 0"},
+                "geometry": {"type": "Point", "coordinates": [0.0, 1.0]},
+            }
+        ],
+    }
+    df_read = read_file(json.dumps(features))
+    assert_geodataframe_equal(expected.set_crs("EPSG:4326"), df_read)
+
+
 def test_read_file_textio(file_path, engine):
     file_text_stream = open(file_path)
     file_stringio = io.StringIO(open(file_path).read())
@@ -700,7 +727,7 @@ def test_allow_legacy_gdal_path(engine, nybb_filename):
     assert isinstance(gdf, geopandas.GeoDataFrame)
 
 
-@pytest.mark.skipif(not PYOGRIO_GE_081, reason="bug fixed in pyogrio 0.8.1")
+@pytest.mark.skipif(not PYOGRIO_GE_090, reason="bug fixed in pyogrio 0.9.0")
 def test_read_file_with_hash_in_path(engine, nybb_filename, tmp_path):
     folder_with_hash = tmp_path / "path with # present"
     folder_with_hash.mkdir(exist_ok=True, parents=True)
diff --git a/geopandas/tools/tests/test_sjoin.py b/geopandas/tools/tests/test_sjoin.py
index 0453355b..0e44f87b 100644
--- a/geopandas/tools/tests/test_sjoin.py
+++ b/geopandas/tools/tests/test_sjoin.py
@@ -96,6 +96,52 @@ def dfs(request):
     return [request.param, df1, df2, expected]
 
 
+@pytest.fixture()
+def dfs_shared_attribute():
+    geo_left = [
+        Point(0, 0),
+        Point(1, 1),
+        Point(2, 2),
+        Point(3, 3),
+        Point(4, 4),
+        Point(5, 5),
+        Point(6, 6),
+        Point(7, 7),
+    ]
+    geo_right = [
+        Point(0, 0),
+        Point(1, 1),
+        Point(2, 2),
+        Point(3, 3),
+        Point(4, 4),
+        Point(5, 5),
+        Point(6, 6),
+        Point(7, 7),
+    ]
+    attr_tracker = ["A", "B", "C", "D", "E", "F", "G", "H"]
+
+    left_gdf = geopandas.GeoDataFrame(
+        {
+            "geometry": geo_left,
+            "attr_tracker": attr_tracker,
+            "duplicate_column": [0, 1, 2, 3, 4, 5, 6, 7],
+            "attr1": [True, True, True, True, True, True, True, True],
+            "attr2": [True, True, True, True, True, True, True, True],
+        }
+    )
+
+    right_gdf = geopandas.GeoDataFrame(
+        {
+            "geometry": geo_right,
+            "duplicate_column": [0, 1, 2, 3, 4, 5, 6, 7],
+            "attr1": [True, True, False, False, True, True, False, False],
+            "attr2": [True, True, False, False, False, False, False, False],
+        }
+    )
+
+    return left_gdf, right_gdf
+
+
 class TestSpatialJoin:
     @pytest.mark.parametrize(
         "how, lsuffix, rsuffix, expected_cols",
@@ -147,7 +193,6 @@ class TestSpatialJoin:
         index, df1, df2, expected = dfs
 
         res = sjoin(df1, df2, how="inner", predicate=predicate)
-
         exp = expected[predicate].dropna().copy()
         exp = exp.drop("geometry_y", axis=1).rename(columns={"geometry_x": "geometry"})
         exp[["df1", "df2"]] = exp[["df1", "df2"]].astype("int64")
@@ -168,7 +213,6 @@ class TestSpatialJoin:
             exp.index.names = df1.index.names
         if index == "named-multi-index":
             exp = exp.set_index(["df1_ix1", "df1_ix2"])
-
         assert_frame_equal(res, exp)
 
     @pytest.mark.parametrize(
@@ -424,6 +468,104 @@ class TestSpatialJoin:
         joined = sjoin(pts, polys, predicate=predicate, how="left")
         assert_index_equal(joined.index, pts.index)
 
+    def test_sjoin_shared_attribute(self, naturalearth_lowres, naturalearth_cities):
+        countries = read_file(naturalearth_lowres)
+        cities = read_file(naturalearth_cities)
+        countries = countries[["geometry", "name"]].rename(columns={"name": "country"})
+
+        # Add first letter of country/city as an attribute column to be compared
+        countries["firstLetter"] = countries["country"].astype(str).str[0]
+        cities["firstLetter"] = cities["name"].astype(str).str[0]
+
+        result = sjoin(cities, countries, on_attribute="firstLetter")
+        assert (
+            result["country"].astype(str).str[0] == result["name"].astype(str).str[0]
+        ).all()
+        assert result.shape == (23, 5)
+
+    @pytest.mark.parametrize(
+        "attr1_key_change_dict, attr2_key_change_dict",
+        [
+            pytest.param(
+                {True: "merge", False: "no_merge"},
+                {True: "merge", False: "no_merge"},
+                id="merge on string attributes",
+            ),
+            pytest.param(
+                {True: 2, False: 1},
+                {True: 2, False: 1},
+                id="merge on integer attributes",
+            ),
+            pytest.param(
+                {True: True, False: False},
+                {True: True, False: False},
+                id="merge on boolean attributes",
+            ),
+            pytest.param(
+                {True: True, False: False},
+                {True: "merge", False: "no_merge"},
+                id="merge on mixed attributes",
+            ),
+        ],
+    )
+    def test_sjoin_multiple_attributes_datatypes(
+        self, dfs_shared_attribute, attr1_key_change_dict, attr2_key_change_dict
+    ):
+        left_gdf, right_gdf = dfs_shared_attribute
+        left_gdf["attr1"] = left_gdf["attr1"].map(attr1_key_change_dict)
+        left_gdf["attr2"] = left_gdf["attr2"].map(attr2_key_change_dict)
+        right_gdf["attr1"] = right_gdf["attr1"].map(attr1_key_change_dict)
+        right_gdf["attr2"] = right_gdf["attr2"].map(attr2_key_change_dict)
+
+        joined = sjoin(left_gdf, right_gdf, on_attribute=("attr1", "attr2"))
+        assert (["A", "B"] == joined["attr_tracker"].values).all()
+
+    def test_sjoin_multiple_attributes_check_header(self, dfs_shared_attribute):
+        left_gdf, right_gdf = dfs_shared_attribute
+        joined = sjoin(left_gdf, right_gdf, on_attribute=["attr1"])
+
+        assert (["A", "B", "E", "F"] == joined["attr_tracker"].values).all()
+        assert {"attr2_left", "attr2_right", "attr1"}.issubset(joined.columns)
+        assert "attr1_left" not in joined
+
+    def test_sjoin_error_column_does_not_exist(self, dfs_shared_attribute):
+        left_gdf, right_gdf = dfs_shared_attribute
+        right_gdf_dropped_attr = right_gdf.drop("attr1", axis=1)
+        left_gdf_dropped_attr = left_gdf.drop("attr1", axis=1)
+
+        with pytest.raises(
+            ValueError,
+            match="Expected column attr1 is missing from the right dataframe.",
+        ):
+            sjoin(left_gdf, right_gdf_dropped_attr, on_attribute="attr1")
+
+        with pytest.raises(
+            ValueError,
+            match="Expected column attr1 is missing from the left dataframe.",
+        ):
+            sjoin(left_gdf_dropped_attr, right_gdf, on_attribute="attr1")
+
+        with pytest.raises(
+            ValueError,
+            match="Expected column attr1 is missing from both of the dataframes.",
+        ):
+            sjoin(left_gdf_dropped_attr, right_gdf_dropped_attr, on_attribute="attr1")
+
+    def test_sjoin_error_use_geometry_column(self, dfs_shared_attribute):
+        left_gdf, right_gdf = dfs_shared_attribute
+        with pytest.raises(
+            ValueError,
+            match="Active geometry column cannot be used as an input for "
+            "on_attribute parameter.",
+        ):
+            sjoin(left_gdf, right_gdf, on_attribute="geometry")
+        with pytest.raises(
+            ValueError,
+            match="Active geometry column cannot be used as an input for "
+            "on_attribute parameter.",
+        ):
+            sjoin(left_gdf, right_gdf, on_attribute=["attr1", "geometry"])
+
 
 class TestIndexNames:
     @pytest.mark.parametrize("how", ["inner", "left", "right"])