diff --git a/python/cudf/cudf/core/indexing.py b/python/cudf/cudf/core/indexing.py index aec931fefbf..a4e15890cbe 100755 --- a/python/cudf/cudf/core/indexing.py +++ b/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,17 @@ 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): + # TODO: switch to cudf.utils.dtypes.is_integer(arg) + if isinstance( + arg, cudf.Scalar + ) and pd.api.types.is_integer_dtype(arg.dtype): + found_index = arg.value + return found_index + elif pd.api.types.is_integer(arg): + found_index = arg + return found_index try: found_index = self._sr.index._values.find_first_value( arg, closest=False @@ -187,7 +195,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: diff --git a/python/cudf/cudf/tests/test_indexing.py b/python/cudf/cudf/tests/test_indexing.py index 086d59ab0f2..048e76ce009 100644 --- a/python/cudf/cudf/tests/test_indexing.py +++ b/python/cudf/cudf/tests/test_indexing.py @@ -102,13 +102,16 @@ def pdf_gdf_multi(): def test_series_indexing(i1, i2, i3): 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 @@ -123,6 +126,44 @@ def test_series_indexing(i1, i2, i3): assert series[i] == a1[i] +@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_get_item_iloc_defer(arg): + # Indexing for non-numeric 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) + + +def test_series_iloc_defer_cudf_scalar(): + ps = pd.Series([1, 2, 3], index=pd.Index(["a", "b", "c"])) + gs = cudf.from_pandas(ps) + + for t in index_dtypes: + arg = cudf.Scalar(1, dtype=t) + got = gs[arg] + expect = 2 + assert_eq(expect, got) + + def test_series_indexing_large_size(): n_elem = 100_000 gsr = cudf.Series(cupy.ones(n_elem))