Repository: narwhals-dev/narwhals
Pull request #1124 (2024-10-03 16:07:08)

feat: add `nulls_last` kw in dataframe sort
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.

Required to have consistent behaviour across libraries (needed in plotly)


Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index efc34317..905ada59 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -176,6 +176,10 @@ class ArrowDataFrame:
         elif isinstance(item, tuple) and len(item) == 2:
             if isinstance(item[1], slice):
                 columns = self.columns
+                if item[1] == slice(None):
+                    if isinstance(item[0], Sequence) and len(item[0]) == 0:
+                        return self._from_native_frame(self._native_frame.slice(0, 0))
+                    return self._from_native_frame(self._native_frame.take(item[0]))
                 if isinstance(item[1].start, str) or isinstance(item[1].stop, str):
                     start, stop, step = convert_str_slice_to_int_slice(item[1], columns)
                     return self._from_native_frame(
@@ -391,7 +395,8 @@ class ArrowDataFrame:
         self,
         by: str | Iterable[str],
         *more_by: str,
-        descending: bool | Sequence[bool] = False,
+        descending: bool | Sequence[bool],
+        nulls_last: bool,
     ) -> Self:
         flat_keys = flatten([*flatten([by]), *more_by])
         df = self._native_frame
@@ -404,7 +409,10 @@ class ArrowDataFrame:
                 (key, "descending" if is_descending else "ascending")
                 for key, is_descending in zip(flat_keys, descending)
             ]
-        return self._from_native_frame(df.sort_by(sorting=sorting))
+
+        null_placement = "at_end" if nulls_last else "at_start"
+
+        return self._from_native_frame(df.sort_by(sorting, null_placement=null_placement))
 
     def to_pandas(self) -> Any:
         return self._native_frame.to_pandas()
diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py
index 916583ea..31949cf2 100644
--- a/narwhals/_dask/dataframe.py
+++ b/narwhals/_dask/dataframe.py
@@ -206,7 +206,8 @@ class DaskLazyFrame:
         self: Self,
         by: str | Iterable[str],
         *more_by: str,
-        descending: bool | Sequence[bool] = False,
+        descending: bool | Sequence[bool],
+        nulls_last: bool,
     ) -> Self:
         flat_keys = flatten([*flatten([by]), *more_by])
         df = self._native_frame
@@ -214,7 +215,10 @@ class DaskLazyFrame:
             ascending: bool | list[bool] = not descending
         else:
             ascending = [not d for d in descending]
-        return self._from_native_frame(df.sort_values(flat_keys, ascending=ascending))
+        na_position = "last" if nulls_last else "first"
+        return self._from_native_frame(
+            df.sort_values(flat_keys, ascending=ascending, na_position=na_position)
+        )
 
     def join(
         self: Self,
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index aae86cef..b9788b53 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -186,6 +186,8 @@ class PandasLikeDataFrame:
 
         elif isinstance(item, tuple) and len(item) == 2 and isinstance(item[1], slice):
             columns = self._native_frame.columns
+            if item[1] == slice(None):
+                return self._from_native_frame(self._native_frame.iloc[item[0], :])
             if isinstance(item[1].start, str) or isinstance(item[1].stop, str):
                 start, stop, step = convert_str_slice_to_int_slice(item[1], columns)
                 return self._from_native_frame(
@@ -411,7 +413,8 @@ class PandasLikeDataFrame:
         self,
         by: str | Iterable[str],
         *more_by: str,
-        descending: bool | Sequence[bool] = False,
+        descending: bool | Sequence[bool],
+        nulls_last: bool,
     ) -> Self:
         flat_keys = flatten([*flatten([by]), *more_by])
         df = self._native_frame
@@ -419,7 +422,10 @@ class PandasLikeDataFrame:
             ascending: bool | list[bool] = not descending
         else:
             ascending = [not d for d in descending]
-        return self._from_native_frame(df.sort_values(flat_keys, ascending=ascending))
+        na_position = "last" if nulls_last else "first"
+        return self._from_native_frame(
+            df.sort_values(flat_keys, ascending=ascending, na_position=na_position)
+        )
 
     # --- convert ---
     def collect(self) -> PandasLikeDataFrame:
diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py
index a4e30ec6..0ef2f879 100644
--- a/narwhals/_polars/dataframe.py
+++ b/narwhals/_polars/dataframe.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Sequence
 
 from narwhals._polars.namespace import PolarsNamespace
 from narwhals._polars.utils import convert_str_slice_to_int_slice
@@ -117,6 +118,12 @@ class PolarsDataFrame:
 
             columns = self.columns
             if isinstance(item, tuple) and len(item) == 2 and isinstance(item[1], slice):
+                if item[1] == slice(None):
+                    if isinstance(item[0], Sequence) and not len(item[0]):
+                        return self._from_native_frame(self._native_frame[0:0])
+                    return self._from_native_frame(
+                        self._native_frame.__getitem__(item[0])
+                    )
                 if isinstance(item[1].start, str) or isinstance(item[1].stop, str):
                     start, stop, step = convert_str_slice_to_int_slice(item[1], columns)
                     return self._from_native_frame(
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 0ee778ab..e4ad31b3 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -170,9 +170,12 @@ class BaseFrame(Generic[FrameT]):
         by: str | Iterable[str],
         *more_by: str,
         descending: bool | Sequence[bool] = False,
+        nulls_last: bool = False,
     ) -> Self:
         return self._from_compliant_dataframe(
-            self._compliant_frame.sort(by, *more_by, descending=descending)
+            self._compliant_frame.sort(
+                by, *more_by, descending=descending, nulls_last=nulls_last
+            )
         )
 
     def join(
@@ -843,8 +846,6 @@ class DataFrame(BaseFrame[DataFrameT]):
         ):
             if item[1] == slice(None) and item[0] == slice(None):
                 return self
-            if item[1] == slice(None):
-                return self._from_compliant_dataframe(self._compliant_frame[item[0]])
             return self._from_compliant_dataframe(self._compliant_frame[item])
         if isinstance(item, str) or (isinstance(item, tuple) and len(item) == 2):
             return self._series(
@@ -1946,19 +1947,22 @@ class DataFrame(BaseFrame[DataFrameT]):
         by: str | Iterable[str],
         *more_by: str,
         descending: bool | Sequence[bool] = False,
+        nulls_last: bool = False,
     ) -> Self:
         r"""
         Sort the dataframe by the given columns.
 
         Arguments:
             by: Column(s) names to sort by.
+            *more_by: Additional columns to sort by, specified as positional arguments.
+            descending: Sort in descending order. When sorting by multiple columns, can be
+                specified per column by passing a sequence of booleans.
+            nulls_last: Place null values last.
 
-            *more_by: Additional columns to sort by, specified as positional
-                       arguments.
-
-            descending: Sort in descending order. When sorting by multiple
-                         columns, can be specified per column by passing a
-                         sequence of booleans.
+        Warning:
+            Unlike Polars, it is not possible to specify a sequence of booleans for
+            `nulls_last` in order to control per-column behaviour. Instead a single
+            boolean is applied for all `by` columns.
 
         Examples:
             >>> import narwhals as nw
@@ -1998,7 +2002,7 @@ class DataFrame(BaseFrame[DataFrameT]):
             │ 2    ┆ 5.0 ┆ c   │
             └──────┴─────┴─────┘
         """
-        return super().sort(by, *more_by, descending=descending)
+        return super().sort(by, *more_by, descending=descending, nulls_last=nulls_last)
 
     def join(
         self,
@@ -3860,20 +3864,23 @@ class LazyFrame(BaseFrame[FrameT]):
         by: str | Iterable[str],
         *more_by: str,
         descending: bool | Sequence[bool] = False,
+        nulls_last: bool = False,
     ) -> Self:
         r"""
         Sort the LazyFrame by the given columns.
 
         Arguments:
-            by: Column(s) to sort by. Accepts expression input. Strings are
-                 parsed as column names.
-
-            *more_by: Additional columns to sort by, specified as positional
-                       arguments.
+            by: Column(s) names to sort by.
+            *more_by: Additional columns to sort by, specified as positional arguments.
+            descending: Sort in descending order. When sorting by multiple columns, can be
+                specified per column by passing a sequence of booleans.
+            nulls_last: Place null values last; can specify a single boolean applying to
+                all columns or a sequence of booleans for per-column control.
 
-            descending: Sort in descending order. When sorting by multiple
-                         columns, can be specified per column by passing a
-                         sequence of booleans.
+        Warning:
+            Unlike Polars, it is not possible to specify a sequence of booleans for
+            `nulls_last` in order to control per-column behaviour. Instead a single
+            boolean is applied for all `by` columns.
 
         Examples:
             >>> import narwhals as nw
@@ -3913,7 +3920,7 @@ class LazyFrame(BaseFrame[FrameT]):
             │ 2    ┆ 5.0 ┆ c   │
             └──────┴─────┴─────┘
         """
-        return super().sort(by, *more_by, descending=descending)
+        return super().sort(by, *more_by, descending=descending, nulls_last=nulls_last)
 
     def join(
         self,


Test patch:
diff --git a/tests/frame/getitem_test.py b/tests/frame/getitem_test.py
index 894555b3..ce96c1b2 100644
--- a/tests/frame/getitem_test.py
+++ b/tests/frame/getitem_test.py
@@ -177,6 +177,15 @@ def test_slice_slice_columns(constructor_eager: Any) -> None:  # noqa: PLR0915
     result = df["a":"b"]  # type: ignore[misc]
     expected = {"a": [1, 2, 3], "b": [4, 5, 6]}
     compare_dicts(result, expected)
+    result = df[(0, 1), :]
+    expected = {"a": [1, 2], "b": [4, 5], "c": [7, 8], "d": [1, 4]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], :]
+    expected = {"a": [1, 2], "b": [4, 5], "c": [7, 8], "d": [1, 4]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], df.columns]
+    expected = {"a": [1, 2], "b": [4, 5], "c": [7, 8], "d": [1, 4]}
+    compare_dicts(result, expected)
 
 
 def test_slice_invalid(constructor_eager: Any) -> None:
diff --git a/tests/frame/sort_test.py b/tests/frame/sort_test.py
index 06f5d079..bea9177d 100644
--- a/tests/frame/sort_test.py
+++ b/tests/frame/sort_test.py
@@ -1,3 +1,7 @@
+from __future__ import annotations
+
+import pytest
+
 import narwhals.stable.v1 as nw
 from tests.utils import Constructor
 from tests.utils import compare_dicts
@@ -20,3 +24,19 @@ def test_sort(constructor: Constructor) -> None:
         "z": [8.0, 9.0, 7.0],
     }
     compare_dicts(result, expected)
+
+
+@pytest.mark.parametrize(
+    ("nulls_last", "expected"),
+    [
+        (True, {"a": [0, 2, 0, -1], "b": [3, 2, 1, float("nan")]}),
+        (False, {"a": [-1, 0, 2, 0], "b": [float("nan"), 3, 2, 1]}),
+    ],
+)
+def test_sort_nulls(
+    constructor: Constructor, *, nulls_last: bool, expected: dict[str, float]
+) -> None:
+    data = {"a": [0, 0, 2, -1], "b": [1, 3, 2, None]}
+    df = nw.from_native(constructor(data))
+    result = df.sort("b", descending=True, nulls_last=nulls_last)
+    compare_dicts(result, expected)