[SNOW-1487754]: Add support for inplace=True in sort_index (#1835)

snowflakedb · Jun 26, 2024 · 197afbf · 197afbf
1 parent 09f4108
commit 197afbf
Show file tree

Hide file tree

Showing 7 changed files with 11 additions and 27 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -91,6 +91,7 @@
 - Added support for `Index.copy()`
 - Added support for Index APIs: `dtype`, `values`, `item()`, `tolist()`, `to_series()` and `to_frame()`
 - Expand support for DataFrames with no rows in `pd.pivot_table` and `DataFrame.pivot_table`.
+- Added support for `inplace` parameter in `DataFrame.sort_index` and `Series.sort_index`.
 
 ## 1.18.0 (2024-05-28)
 

diff --git a/docs/source/modin/supported/dataframe_supported.rst b/docs/source/modin/supported/dataframe_supported.rst
@@ -383,7 +383,7 @@ Methods
 |                             |                                 |                                  | or ``numeric_only=False``                          |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``sort_index``              | P                               |                                  | ``N`` if given the ``key`` param. ``N`` if         |
-|                             |                                 |                                  | ``axis == 1``, ``inplace == True``, or MultiIndex. |
+|                             |                                 |                                  | ``axis == 1``, or MultiIndex.                      |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``sort_values``             | P                               |                                  | ``N`` if given the ``key`` param or ``axis == 1``  |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+

diff --git a/docs/source/modin/supported/series_supported.rst b/docs/source/modin/supported/series_supported.rst
@@ -372,7 +372,7 @@ Methods
 |                             |                                 |                                  | or ``numeric_only=False``                          |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``sort_index``              | P                               |                                  | ``N`` if given the ``key`` param,                  |
-|                             |                                 |                                  | or if ``inplace == True`` or MultiIndex            |
+|                             |                                 |                                  | or MultiIndex                                      |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+
 | ``sort_values``             | P                               |                                  | ``N`` if given the ``key`` parameter               |
 +-----------------------------+---------------------------------+----------------------------------+----------------------------------------------------+

diff --git a/src/snowflake/snowpark/modin/pandas/base.py b/src/snowflake/snowpark/modin/pandas/base.py
@@ -3135,7 +3135,6 @@ def sort_index(
             axis=axis,
             level=level,
             ascending=ascending,
-            inplace=inplace,
             kind=kind,
             na_position=na_position,
             sort_remaining=sort_remaining,

diff --git a/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py b/src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py
@@ -2389,7 +2389,6 @@ def sort_index(
         axis: int,
         level: list[Union[str, int]],
         ascending: Union[bool, list[bool]],
-        inplace: bool,
         kind: SortKind,
         na_position: NaPosition,
         sort_remaining: bool,
@@ -2404,7 +2403,6 @@ def sort_index(
             level: If not None, sort on values in specified index level(s).
             ascending: A list of bools to represent ascending vs descending sort. Defaults to True.
                 When the index is a MultiIndex the sort direction can be controlled for each level individually.
-            inplace: Whether to modify the DataFrame rather than creating a new one.
             kind: Choice of sorting algorithm. Perform stable sort if 'stable'. Defaults to unstable sort.
                 Snowpark pandas ignores choice of sorting algorithm except 'stable'.
             na_position: Puts NaNs at the beginning if 'first'; 'last' puts NaNs at the end. Defaults to 'last'
@@ -2445,10 +2443,6 @@ def sort_index(
             ErrorMessage.not_implemented(
                 "sort_index is not supported yet on axis=1 in Snowpark pandas."
             )
-        if inplace:
-            ErrorMessage.not_implemented(
-                "sort_index is not supported yet with inplace=True in Snowpark pandas."
-            )
         if key:
             ErrorMessage.not_implemented(
                 "Snowpark pandas sort_index API doesn't yet support 'key' parameter"

diff --git a/tests/integ/modin/frame/test_sort_index.py b/tests/integ/modin/frame/test_sort_index.py
@@ -14,8 +14,9 @@
 @pytest.mark.parametrize("ascending", [True, False])
 @pytest.mark.parametrize("na_position", ["first", "last"])
 @pytest.mark.parametrize("ignore_index", [True, False])
+@pytest.mark.parametrize("inplace", [True, False])
 @sql_count_checker(query_count=1)
-def test_sort_index_dataframe(ascending, na_position, ignore_index):
+def test_sort_index_dataframe(ascending, na_position, ignore_index, inplace):
     native_df = native_pd.DataFrame(
         [1, 2, np.nan, 4, 5], index=[np.nan, 29, 234, 1, 150], columns=["A"]
     )
@@ -27,7 +28,9 @@ def test_sort_index_dataframe(ascending, na_position, ignore_index):
             ascending=ascending,
             na_position=na_position,
             ignore_index=ignore_index,
+            inplace=inplace,
         ),
+        inplace=inplace,
     )
 
 
@@ -40,15 +43,6 @@ def test_sort_index_dataframe_axis_1_unsupported():
         snow_df.sort_index(axis=1)
 
 
-@sql_count_checker(query_count=0)
-def test_sort_index_dataframe_inplace_unsupported():
-    snow_df = pd.DataFrame(
-        [1, 2, np.nan, 4, 5], index=[np.nan, 29, 234, 1, 150], columns=["A"]
-    )
-    with pytest.raises(NotImplementedError):
-        snow_df.sort_index(inplace=True)
-
-
 @sql_count_checker(query_count=0)
 def test_sort_index_dataframe_multiindex_unsupported():
     arrays = [

diff --git a/tests/integ/modin/series/test_sort_index.py b/tests/integ/modin/series/test_sort_index.py
@@ -14,8 +14,9 @@
 @pytest.mark.parametrize("ascending", [True, False])
 @pytest.mark.parametrize("na_position", ["first", "last"])
 @pytest.mark.parametrize("ignore_index", [True, False])
+@pytest.mark.parametrize("inplace", [True, False])
 @sql_count_checker(query_count=1)
-def test_sort_index_series(ascending, na_position, ignore_index):
+def test_sort_index_series(ascending, na_position, ignore_index, inplace):
     native_series = native_pd.Series(["a", "b", np.nan, "d"], index=[3, 2, 1, np.nan])
     snow_series = pd.Series(native_series)
     eval_snowpark_pandas_result(
@@ -25,17 +26,12 @@ def test_sort_index_series(ascending, na_position, ignore_index):
             ascending=ascending,
             na_position=na_position,
             ignore_index=ignore_index,
+            inplace=inplace,
         ),
+        inplace=inplace,
     )
 
 
-@sql_count_checker(query_count=0)
-def test_sort_index_series_inplace_unsupported():
-    snow_series = pd.Series(["a", "b", np.nan, "d"], index=[3, 2, 1, np.nan])
-    with pytest.raises(NotImplementedError):
-        snow_series.sort_index(inplace=True)
-
-
 @sql_count_checker(query_count=0)
 def test_sort_index_series_multiindex_unsupported():
     arrays = [