limit get_loc to exact match for now

Author: jorisvandenbossche

pandas/_libs/index.pyx

@@ -561,15 +561,23 @@ cdef class StringObjectEngine(ObjectEngine):
 
     cdef:
         object na_value
+        bint uses_na
 
     def __init__(self, ndarray values, na_value):
         super().__init__(values)
         self.na_value = na_value
+        self.uses_na = na_value is C_NA
+
+    cdef bint _checknull(self, object val):
+        if self.uses_na:
+            return val is C_NA
+        else:
+            return util.is_nan(val)
 
     cdef _check_type(self, object val):
         if isinstance(val, str):
             return val
-        elif checknull(val):
+        elif self._checknull(val):
             return self.na_value
         else:
             raise KeyError(val)

pandas/core/indexes/base.py

@@ -5974,7 +5974,6 @@ def _should_fallback_to_positional(self) -> bool:
     def get_indexer_non_unique(
         self, target
     ) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]:
-        target = ensure_index(target)
         target = self._maybe_cast_listlike_indexer(target)
 
         if not self._should_compare(target) and not self._should_partial_index(target):
@@ -6222,10 +6221,6 @@ def _maybe_downcast_for_indexing(self, other: Index) -> tuple[Index, Index]:
                 # let's instead try with a straight Index
                 self = Index(self._values)
 
-        elif self.dtype == "string" and other.dtype == "object":
-            if lib.is_string_array(other._values, skipna=True):  # type: ignore[arg-type]
-                return self, other.astype(self.dtype)
-
         if not is_object_dtype(self.dtype) and is_object_dtype(other.dtype):
             # Reverse op so we dont need to re-implement on the subclasses
             other, self = other._maybe_downcast_for_indexing(self)

pandas/tests/frame/indexing/test_indexing.py

@@ -9,6 +9,8 @@
 import numpy as np
 import pytest
 
+from pandas._config import using_string_dtype
+
 from pandas._libs import iNaT
 from pandas.errors import InvalidIndexError
 
@@ -501,6 +503,7 @@ def test_setitem_ambig(self, using_infer_string):
         else:
             assert dm[2].dtype == np.object_
 
+    @pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string)")
     def test_setitem_None(self, float_frame):
         # GH #766
         float_frame[None] = float_frame["A"]

pandas/tests/frame/test_arithmetic.py

@@ -11,6 +11,8 @@
 import numpy as np
 import pytest
 
+from pandas._config import using_string_dtype
+
 import pandas as pd
 from pandas import (
     DataFrame,
@@ -2099,21 +2101,12 @@ def test_enum_column_equality():
     tm.assert_series_equal(result, expected)
 
 
-def test_mixed_col_index_dtype(any_string_dtype):
+@pytest.mark.xfail(using_string_dtype(), reason="TODO(infer_string)")
+def test_mixed_col_index_dtype():
     # GH 47382
-    df1 = DataFrame(
-        columns=Index(list("abc"), dtype=any_string_dtype), data=1.0, index=[0]
-    )
-    df2 = DataFrame(columns=Index(list("abc"), dtype="object"), data=0.0, index=[0])
-
+    df1 = DataFrame(columns=list("abc"), data=1.0, index=[0])
+    df2 = DataFrame(columns=list("abc"), data=0.0, index=[0])
+    df1.columns = df2.columns.astype("string")
     result = df1 + df2
-    expected = DataFrame(
-        columns=Index(list("abc"), dtype=any_string_dtype), data=1.0, index=[0]
-    )
-    tm.assert_frame_equal(result, expected)
-
-    result = df2 + df1
-    expected = DataFrame(
-        columns=Index(list("abc"), dtype="object"), data=1.0, index=[0]
-    )
+    expected = DataFrame(columns=list("abc"), data=1.0, index=[0])
     tm.assert_frame_equal(result, expected)

pandas/tests/indexes/string/test_indexing.py

@@ -6,6 +6,13 @@
 import pandas._testing as tm
 
 
+def _isnan(val):
+    try:
+        return val is not pd.NA and np.isnan(val)
+    except TypeError:
+        return False
+
+
 class TestGetLoc:
     def test_get_loc(self, any_string_dtype):
         index = Index(["a", "b", "c"], dtype=any_string_dtype)
@@ -34,7 +41,14 @@ def test_get_loc_non_missing(self, any_string_dtype, nulls_fixture):
 
     def test_get_loc_missing(self, any_string_dtype, nulls_fixture):
         index = Index(["a", "b", nulls_fixture], dtype=any_string_dtype)
-        assert index.get_loc(nulls_fixture) == 2
+        if any_string_dtype == "string" and (
+            (any_string_dtype.na_value is pd.NA and nulls_fixture is not pd.NA)
+            or (_isnan(any_string_dtype.na_value) and not _isnan(nulls_fixture))
+        ):
+            with pytest.raises(KeyError):
+                index.get_loc(nulls_fixture)
+        else:
+            assert index.get_loc(nulls_fixture) == 2
 
 
 class TestGetIndexer:
@@ -83,32 +97,20 @@ def test_get_indexer_missing(self, any_string_dtype, null):
 
 class TestGetIndexerNonUnique:
     @pytest.mark.parametrize("null", [None, np.nan, float("nan"), pd.NA])
-    def test_get_indexer_non_unique_nas(self, request, any_string_dtype, null):
-        if (
-            any_string_dtype == "string"
-            and any_string_dtype.na_value is pd.NA
-            and isinstance(null, float)
-        ):
-            # TODO(infer_string)
-            request.applymarker(
-                pytest.mark.xfail(
-                    reason="NA-variant string dtype does not work with NaN"
-                )
-            )
-
+    def test_get_indexer_non_unique_nas(self, any_string_dtype, null):
         index = Index(["a", "b", null], dtype=any_string_dtype)
-        indexer, missing = index.get_indexer_non_unique([null])
+        indexer, missing = index.get_indexer_non_unique(["a", null])
 
-        expected_indexer = np.array([2], dtype=np.intp)
+        expected_indexer = np.array([0, 2], dtype=np.intp)
         expected_missing = np.array([], dtype=np.intp)
         tm.assert_numpy_array_equal(indexer, expected_indexer)
         tm.assert_numpy_array_equal(missing, expected_missing)
 
         # actually non-unique
         index = Index(["a", null, "b", null], dtype=any_string_dtype)
-        indexer, missing = index.get_indexer_non_unique([null])
+        indexer, missing = index.get_indexer_non_unique(["a", null])
 
-        expected_indexer = np.array([1, 3], dtype=np.intp)
+        expected_indexer = np.array([0, 1, 3], dtype=np.intp)
         tm.assert_numpy_array_equal(indexer, expected_indexer)
         tm.assert_numpy_array_equal(missing, expected_missing)
 

pandas/tests/reshape/test_pivot.py

@@ -2668,8 +2668,6 @@ def test_pivot_columns_not_given(self):
         with pytest.raises(TypeError, match="missing 1 required keyword-only argument"):
             df.pivot()
 
-    # this still fails because columns=None gets passed down to unstack as level=None
-    # while at that point None was converted to NaN
     @pytest.mark.xfail(
         using_string_dtype(), reason="TODO(infer_string) None is cast to NaN"
     )
@@ -2688,7 +2686,10 @@ def test_pivot_columns_is_none(self):
         expected = DataFrame({1: 3}, index=Index([2], name="b"))
         tm.assert_frame_equal(result, expected)
 
-    def test_pivot_index_is_none(self, using_infer_string):
+    @pytest.mark.xfail(
+        using_string_dtype(), reason="TODO(infer_string) None is cast to NaN"
+    )
+    def test_pivot_index_is_none(self):
         # GH#48293
         df = DataFrame({None: [1], "b": 2, "c": 3})
 
@@ -2699,10 +2700,11 @@ def test_pivot_index_is_none(self, using_infer_string):
 
         result = df.pivot(columns="b", index=None, values="c")
         expected = DataFrame(3, index=[1], columns=Index([2], name="b"))
-        if using_infer_string:
-            expected.index.name = np.nan
         tm.assert_frame_equal(result, expected)
 
+    @pytest.mark.xfail(
+        using_string_dtype(), reason="TODO(infer_string) None is cast to NaN"
+    )
     def test_pivot_values_is_none(self):
         # GH#48293
         df = DataFrame({None: [1], "b": 2, "c": 3})