Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 26 additions & 2 deletions python/cudf/cudf/core/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -437,7 +437,19 @@ def _setitem_tuple_arg(self, key, value):


class _DataFrameAtIndexer(_DataFrameLocIndexer):
pass
@_performance_tracking
def __getitem__(self, key):
indexing_utils.validate_scalar_key(
key, "Invalid call for scalar access (getting)!"
)
return super().__getitem__(key)

@_performance_tracking
def __setitem__(self, key, value):
indexing_utils.validate_scalar_key(
key, "Invalid call for scalar access (getting)!"
)
return super().__setitem__(key, value)


class _DataFrameIlocIndexer(_DataFrameIndexer):
Expand Down Expand Up @@ -507,7 +519,19 @@ def _setitem_tuple_arg(self, key, value):


class _DataFrameiAtIndexer(_DataFrameIlocIndexer):
pass
@_performance_tracking
def __getitem__(self, key):
indexing_utils.validate_scalar_key(
key, "iAt based indexing can only have integer indexers"
)
return super().__getitem__(key)

@_performance_tracking
def __setitem__(self, key, value):
indexing_utils.validate_scalar_key(
key, "iAt based indexing can only have integer indexers"
)
return super().__setitem__(key, value)


@_performance_tracking
Expand Down
25 changes: 25 additions & 0 deletions python/cudf/cudf/core/indexing_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from cudf.api.types import (
_is_scalar_or_zero_d_array,
is_integer,
is_list_like,
)
from cudf.core.column.column import as_column
from cudf.core.copy_types import BooleanMask, GatherMap
Expand Down Expand Up @@ -69,6 +70,30 @@ class ScalarIndexer:
)


def validate_scalar_key(key: Any, error_msg: str) -> None:
"""Validate that key contains only scalar values for .at/.iat indexers.

Parameters
----------
key : Any
The key to validate
error_msg : str
The error message to raise if validation fails

Raises
------
ValueError
If the key contains list-like indexers
"""
if not isinstance(key, tuple):
if is_list_like(key):
raise ValueError(error_msg)
else:
for k in key:
if is_list_like(k):
raise ValueError(error_msg)


# Helpers for code-sharing between loc and iloc paths
def expand_key(
key: Any, frame: DataFrame | Series, method_type: Literal["iloc", "loc"]
Expand Down
28 changes: 26 additions & 2 deletions python/cudf/cudf/core/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,19 @@ def __setitem__(self, key, value):


class _SeriesiAtIndexer(_SeriesIlocIndexer):
pass
@_performance_tracking
def __getitem__(self, key):
indexing_utils.validate_scalar_key(
key, "iAt based indexing can only have integer indexers"
)
return super().__getitem__(key)

@_performance_tracking
def __setitem__(self, key, value):
indexing_utils.validate_scalar_key(
key, "iAt based indexing can only have integer indexers"
)
return super().__setitem__(key, value)


class _SeriesLocIndexer(_FrameIndexer):
Expand Down Expand Up @@ -413,7 +425,19 @@ def _loc_to_iloc(self, arg):


class _SeriesAtIndexer(_SeriesLocIndexer):
pass
@_performance_tracking
def __getitem__(self, key):
indexing_utils.validate_scalar_key(
key, "Invalid call for scalar access (getting)!"
)
return super().__getitem__(key)

@_performance_tracking
def __setitem__(self, key, value):
indexing_utils.validate_scalar_key(
key, "Invalid call for scalar access (getting)!"
)
return super().__setitem__(key, value)


class Series(SingleColumnFrame, IndexedFrame):
Expand Down
18 changes: 1 addition & 17 deletions python/cudf/cudf/tests/dataframe/indexing/test_loc.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0
import re
import weakref
Expand Down Expand Up @@ -627,12 +627,6 @@ def test_dataframe_loc(scalar):
# Full slice
assert_eq(df.loc[:, "c"], pdf.loc[:, "c"])

# Repeat with at[]
assert_eq(df.loc[:, ["a"]], df.at[:, ["a"]])
assert_eq(df.loc[:, "d"], df.at[:, "d"])
assert_eq(df.loc[scalar], df.at[scalar])
assert_eq(df.loc[:, "c"], df.at[:, "c"])


@pytest.mark.parametrize("step", [1, 5])
def test_dataframe_loc_slice(step):
Expand Down Expand Up @@ -672,16 +666,6 @@ def test_dataframe_loc_slice(step):
df.loc[begin, "a":"a"], pdf.loc[begin, "a":"a"], check_dtype=False
)

# Repeat with at[]
assert_eq(
df.loc[begin:end:step, ["c", "d", "a"]],
df.at[begin:end:step, ["c", "d", "a"]],
)
assert_eq(df.loc[begin:end, ["c", "d"]], df.at[begin:end, ["c", "d"]])
assert_eq(df.loc[begin:end:step, "a":"c"], df.at[begin:end:step, "a":"c"])
assert_eq(df.loc[begin:begin, "a"], df.at[begin:begin, "a"])
assert_eq(df.loc[begin, "a":"a"], df.at[begin, "a":"a"], check_dtype=False)


def test_dataframe_loc_arraylike():
size = 123
Expand Down
66 changes: 66 additions & 0 deletions python/cudf/cudf/tests/dataframe/test_at_iat.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0

import pytest

import cudf


@pytest.fixture
def df_with_index():
return cudf.DataFrame(
{"A": [1, 2, 3], "B": [4, 5, 6]}, index=["x", "y", "z"]
)


@pytest.fixture
def df_without_index():
return cudf.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})


def test_dataframe_at_scalar_getitem(df_with_index):
assert df_with_index.at["x", "A"] == 1
assert df_with_index.at["y", "B"] == 5


def test_dataframe_at_scalar_setitem(df_with_index):
df_with_index.at["x", "A"] = 10
assert df_with_index.at["x", "A"] == 10


@pytest.mark.parametrize("key", [[["x"], "A"], ["x", ["A"]], [["x"], ["A"]]])
def test_dataframe_at_rejects_list_like(df_with_index, key):
with pytest.raises(ValueError, match="Invalid call for scalar access"):
df_with_index.at[key[0], key[1]]


@pytest.mark.parametrize("key", [[["x"], "A"], ["x", ["A"]]])
def test_dataframe_at_setitem_rejects_list_like(df_with_index, key):
with pytest.raises(ValueError, match="Invalid call for scalar access"):
df_with_index.at[key[0], key[1]] = 10


def test_dataframe_iat_scalar_getitem(df_without_index):
assert df_without_index.iat[0, 0] == 1
assert df_without_index.iat[1, 1] == 5


def test_dataframe_iat_scalar_setitem(df_without_index):
df_without_index.iat[0, 0] = 10
assert df_without_index.iat[0, 0] == 10


@pytest.mark.parametrize("key", [[[0], 0], [0, [0]], [[0], [0]]])
def test_dataframe_iat_rejects_list_like(df_without_index, key):
with pytest.raises(
ValueError, match="iAt based indexing can only have integer indexers"
):
df_without_index.iat[key[0], key[1]]


@pytest.mark.parametrize("key", [[[0], 0], [0, [0]]])
def test_dataframe_iat_setitem_rejects_list_like(df_without_index, key):
with pytest.raises(
ValueError, match="iAt based indexing can only have integer indexers"
):
df_without_index.iat[key[0], key[1]] = 10
62 changes: 62 additions & 0 deletions python/cudf/cudf/tests/series/test_at_iat.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0

import pytest

import cudf


@pytest.fixture
def sr_with_index():
return cudf.Series([1, 2, 3], index=["x", "y", "z"])


@pytest.fixture
def sr_without_index():
return cudf.Series([1, 2, 3])


def test_series_at_scalar_getitem(sr_with_index):
assert sr_with_index.at["x"] == 1
assert sr_with_index.at["y"] == 2


def test_series_at_scalar_setitem(sr_with_index):
sr_with_index.at["x"] = 10
assert sr_with_index.at["x"] == 10


@pytest.mark.parametrize("key", [[["x"]], [["x", "y"]]])
def test_series_at_rejects_list_like(sr_with_index, key):
with pytest.raises(ValueError, match="Invalid call for scalar access"):
sr_with_index.at[key[0]]


def test_series_at_setitem_rejects_list_like(sr_with_index):
with pytest.raises(ValueError, match="Invalid call for scalar access"):
sr_with_index.at[["x"]] = 10


def test_series_iat_scalar_getitem(sr_without_index):
assert sr_without_index.iat[0] == 1
assert sr_without_index.iat[1] == 2


def test_series_iat_scalar_setitem(sr_without_index):
sr_without_index.iat[0] = 10
assert sr_without_index.iat[0] == 10


@pytest.mark.parametrize("key", [[[0]], [[0, 1]]])
def test_series_iat_rejects_list_like(sr_without_index, key):
with pytest.raises(
ValueError, match="iAt based indexing can only have integer indexers"
):
sr_without_index.iat[key[0]]


def test_series_iat_setitem_rejects_list_like(sr_without_index):
with pytest.raises(
ValueError, match="iAt based indexing can only have integer indexers"
):
sr_without_index.iat[[0]] = 10