From e830977976930fd46aa85a6a3f5af8d68fc452b8 Mon Sep 17 00:00:00 2001 From: Marshall Crumiller Date: Fri, 17 Jan 2025 16:03:30 -0500 Subject: [PATCH] Add height check for single index --- py-polars/polars/_utils/getitem.py | 3 +++ py-polars/tests/unit/dataframe/test_df.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/py-polars/polars/_utils/getitem.py b/py-polars/polars/_utils/getitem.py index 05404437c0ce..364da0ff91f8 100644 --- a/py-polars/polars/_utils/getitem.py +++ b/py-polars/polars/_utils/getitem.py @@ -291,6 +291,9 @@ def _select_rows( ) -> DataFrame | Series: """Select one or more rows from the DataFrame.""" if isinstance(key, int): + if key >= len(df): + msg = f"index {key} is out of bounds for DataFrame of height {len(df)}" + raise IndexError(msg) return df.slice(key, 1) if isinstance(key, slice): diff --git a/py-polars/tests/unit/dataframe/test_df.py b/py-polars/tests/unit/dataframe/test_df.py index 8c2fafcdd279..6fd66aa5bb3d 100644 --- a/py-polars/tests/unit/dataframe/test_df.py +++ b/py-polars/tests/unit/dataframe/test_df.py @@ -2995,3 +2995,21 @@ def test_get_column_after_drop_20119() -> None: df.drop_in_place("a") c = df.get_column("c") assert_series_equal(c, pl.Series("c", ["C"])) + + +def test_select_oob_row_20775() -> None: + df = pl.DataFrame({"a": [1, 2, 3]}) + with pytest.raises( + IndexError, + match="index 99 is out of bounds for DataFrame of height 3", + ): + df[99] + + +def test_select_oob_element_20775() -> None: + df = pl.DataFrame({"a": [1, 2, 3]}) + with pytest.raises( + IndexError, + match="index 99 is out of bounds for sequence of length 3", + ): + df[99, "a"]