diff --git a/bigframes/dataframe.py b/bigframes/dataframe.py index 01e9bd6308..aae081349b 100644 --- a/bigframes/dataframe.py +++ b/bigframes/dataframe.py @@ -1976,6 +1976,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")