fix: dataframe sort_values Series input keyerror. (#1285)

googleapis · Jan 14, 2025 · 5a2731b · 5a2731b
1 parent dc92e95
commit 5a2731b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git 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
@@ -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")