From 972623aa64d0562b489d9a228d0500fb6ac3c042 Mon Sep 17 00:00:00 2001 From: Huan Chen Date: Mon, 13 Jan 2025 23:35:27 +0000 Subject: [PATCH] fix: dataframe sort_values Series input keyerror. --- bigframes/dataframe.py | 5 +++++ tests/system/small/test_dataframe.py | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index e7a0444af1..a4bdeb7fe1 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1978,6 +1978,11 @@ def sort_values( kind: str = "quicksort", na_position: typing.Literal["first", "last"] = "last", ) -> DataFrame: + if isinstance(by, (bigframes.series.Series, indexes.Index, DataFrame)): + raise KeyError( + f"Invalid key type: {type(by).__name__}. Please provide valid column name(s)." + ) + if na_position not in {"first", "last"}: raise ValueError("Param na_position must be one of 'first' or 'last'") diff --git a/tests/system/small/test_dataframe.py b/tests/system/small/test_dataframe.py index 4e0e5c2739..5b94df2446 100644 --- a/tests/system/small/test_dataframe.py +++ b/tests/system/small/test_dataframe.py @@ -2581,6 +2581,11 @@ def test_dataframe_sort_values( ) +def test_dataframe_sort_values_invalid_input(scalars_df_index): + with pytest.raises(KeyError): + scalars_df_index.sort_values(by=scalars_df_index["int64_col"]) + + def test_dataframe_sort_values_stable(scalars_df_index, scalars_pandas_df_index): bf_result = ( scalars_df_index.sort_values("int64_col", kind="stable")