Defer to iloc when indexer is an integer-like object and the Index is non-numeric dtype

python/cudf/cudf/core/indexing.py

@@ -16,6 +16,7 @@
     is_categorical_dtype,
     is_column_like,
     is_list_like,
+    is_numerical_dtype,
     is_scalar,
     to_cudf_compatible_scalar,
 )
@@ -164,10 +165,18 @@ def __setitem__(self, key, value):
         self._sr.iloc[key] = value
 
     def _loc_to_iloc(self, arg):
-        from cudf.core.column import column
-        from cudf.core.series import Series
-
         if is_scalar(arg):
+            if not is_numerical_dtype(self._sr.index.dtype):
+                if pd.api.types.is_integer(arg):
+                    found_index = arg
+                    return found_index
+
+                elif isinstance(
+                    arg, cudf.Scalar
+                ) and pd.api.types.is_integer_dtype(arg.dtype):
+                    found_index = arg.value
+                    return found_index
+
             try:
                 found_index = self._sr.index._values.find_first_value(
                     arg, closest=False
@@ -187,7 +196,7 @@ def _loc_to_iloc(self, arg):
             return indices_from_labels(self._sr, arg)
 
         else:
-            arg = Series(column.as_column(arg))
+            arg = cudf.core.series.Series(cudf.core.column.as_column(arg))
             if arg.dtype in (bool, np.bool_):
                 return arg
             else:

python/cudf/cudf/tests/test_indexing.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2021, NVIDIA CORPORATION.
 
+import re
 from itertools import combinations
 
 import cupy
@@ -99,16 +100,35 @@ def pdf_gdf_multi():
         + ["numpy.array[%s]" % np.dtype(t).type.__name__ for t in index_dtypes]
     ),
 )
-def test_series_indexing(i1, i2, i3):
+@pytest.mark.parametrize(
+    "arg",
+    [
+        1,
+        -1,
+        "b",
+        np.int32(1),
+        np.uint32(1),
+        np.int8(1),
+        np.uint8(1),
+        np.int16(1),
+        np.uint16(1),
+        np.int64(1),
+        np.uint64(1),
+    ],
+)
+def test_series_indexing(i1, i2, i3, arg):
     a1 = np.arange(20)
     series = cudf.Series(a1)
+
     # Indexing
     sr1 = series.iloc[i1]
     assert sr1.null_count == 0
     np.testing.assert_equal(sr1.to_array(), a1[:12])
+
     sr2 = sr1.iloc[i2]
     assert sr2.null_count == 0
     np.testing.assert_equal(sr2.to_array(), a1[3:12])
+
     # Index with stride
     sr3 = sr2.iloc[i3]
     assert sr3.null_count == 0
@@ -122,6 +142,33 @@ def test_series_indexing(i1, i2, i3):
         for i in i1:  # numpy integers
             assert series[i] == a1[i]
 
+    # Indexing for non-integer dtype Index
+    ps = pd.Series([1, 2, 3], index=pd.Index(["a", "b", "c"]))
+    gs = cudf.from_pandas(ps)
+
+    expect = ps[arg]
+    got = gs[arg]
+
+    assert_eq(expect, got)
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    ["int32", "int16", "int8", "int64", "uint32", "uint16", "uint8", "uint64"],
+)
+def test_series_indexing_cudf_Scalar(dtype):
+    ps = pd.Series([1, 2, 3], index=pd.Index(["a", "b", "c"]))
+    gs = cudf.from_pandas(ps)
+
+    arg = cudf.Scalar(1, dtype=dtype)
+    got = gs[arg]
+
+    with pytest.raises(
+        TypeError, match=re.escape(f"'{arg}' is an invalid key"),
+    ):
+        expect = ps[arg]
+        assert_eq(expect, got)
+
 
 def test_series_indexing_large_size():
     n_elem = 100_000