SNOW-2398216: Add support for sort_values in faster pandas (#3862)

Author: sfc-gh-helmeleegy

CHANGELOG.md

@@ -46,6 +46,7 @@
 - Set `cte_optimization_enabled` to True for all Snowpark pandas sessions.
 - Add support for `isna`, `isnull`, `notna`, `notnull` in faster pandas.
 - Add support for `str.contains`, `str.startswith`, `str.endswith`, and `str.slice` in faster pandas.
+- Add support for `sort_values` in faster pandas.
 
 ## 1.40.0 (2025-10-02)
 

src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py

@@ -4058,6 +4058,46 @@ def sort_rows_by_column_values(
         key: Optional[IndexKeyFunc] = None,
         include_indexer: bool = False,
         include_index: bool = True,
+    ) -> "SnowflakeQueryCompiler":
+        """
+        Wrapper around _sort_rows_by_column_values_internal to be supported in faster pandas.
+        """
+        relaxed_query_compiler = None
+        if self._relaxed_query_compiler is not None:
+            relaxed_query_compiler = (
+                self._relaxed_query_compiler._sort_rows_by_column_values_internal(
+                    columns=columns,
+                    ascending=ascending,
+                    kind=kind,
+                    na_position=na_position,
+                    ignore_index=ignore_index,
+                    key=key,
+                    include_indexer=include_indexer,
+                    include_index=include_index,
+                )
+            )
+        qc = self._sort_rows_by_column_values_internal(
+            columns=columns,
+            ascending=ascending,
+            kind=kind,
+            na_position=na_position,
+            ignore_index=ignore_index,
+            key=key,
+            include_indexer=include_indexer,
+            include_index=include_index,
+        )
+        return self._maybe_set_relaxed_qc(qc, relaxed_query_compiler)
+
+    def _sort_rows_by_column_values_internal(
+        self,
+        columns: list[Hashable],
+        ascending: list[bool],
+        kind: SortKind,
+        na_position: NaPosition,
+        ignore_index: bool,
+        key: Optional[IndexKeyFunc] = None,
+        include_indexer: bool = False,
+        include_index: bool = True,
     ) -> "SnowflakeQueryCompiler":
         """
         Reorder the rows based on the lexicographic order of the given columns.

tests/integ/modin/test_faster_pandas.py

@@ -370,6 +370,35 @@ def test_str_slice(session):
     assert_series_equal(snow_result, native_result)
 
 
+@sql_count_checker(query_count=3)
+def test_sort_values(session):
+    # create tables
+    table_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+    session.create_dataframe(
+        native_pd.DataFrame([[2, 12], [1, 11], [3, 13]], columns=["A", "B"])
+    ).write.save_as_table(table_name, table_type="temp")
+
+    # create snow dataframes
+    df = pd.read_snowflake(table_name)
+    snow_result = df.sort_values(by="A")
+
+    # verify that the input dataframe has a populated relaxed query compiler
+    assert df._query_compiler._relaxed_query_compiler is not None
+    assert df._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    # verify that the output dataframe also has a populated relaxed query compiler
+    assert snow_result._query_compiler._relaxed_query_compiler is not None
+    assert (
+        snow_result._query_compiler._relaxed_query_compiler._dummy_row_pos_mode is True
+    )
+
+    # create pandas dataframes
+    native_df = df.to_pandas()
+    native_result = native_df.sort_values(by="A")
+
+    # compare results
+    assert_frame_equal(snow_result, native_result)
+
+
 @sql_count_checker(query_count=0)
 def test_dummy_row_pos_optimization_enabled_on_session(db_parameters):
     with Session.builder.configs(db_parameters).create() as new_session: