Repository: modin-project/modin

Issue:
nlargest function implementation



Discussion:

Resolving patch:
diff --git a/modin/backends/pandas/query_compiler.py b/modin/backends/pandas/query_compiler.py
--- a/modin/backends/pandas/query_compiler.py
+++ b/modin/backends/pandas/query_compiler.py
@@ -352,6 +352,8 @@ def transpose(self, *args, **kwargs):
     idxmin = ReductionFunction.register(pandas.DataFrame.idxmin)
     median = ReductionFunction.register(pandas.DataFrame.median)
     nunique = ReductionFunction.register(pandas.DataFrame.nunique)
+    nsmallest = ReductionFunction.register(pandas.DataFrame.nsmallest)
+    nlargest = ReductionFunction.register(pandas.DataFrame.nlargest)
     skew = ReductionFunction.register(pandas.DataFrame.skew)
     std = ReductionFunction.register(pandas.DataFrame.std)
     var = ReductionFunction.register(pandas.DataFrame.var)
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -1215,11 +1215,55 @@ def ne(self, other, axis="columns", level=None):
         )
 
     def nlargest(self, n, columns, keep="first"):
-        return self._default_to_pandas(pandas.DataFrame.nlargest, n, columns, keep=keep)
+        """
+        Return the first `n` rows ordered by `columns` in descending order.
+        Return the first `n` rows with the largest values in `columns`, in
+        descending order. The columns that are not specified are returned as
+        well, but not used for ordering.
+        This method is equivalent to
+        ``df.sort_values(columns, ascending=False).head(n)``, but more
+        performant.
+        Parameters
+        ----------
+        n : int
+            Number of rows to return.
+        columns : label or list of labels
+            Column label(s) to order by.
+        keep : {'first', 'last', 'all'}, default 'first'
+            Where there are duplicate values:
+            - `first` : prioritize the first occurrence(s)
+            - `last` : prioritize the last occurrence(s)
+            - ``all`` : do not drop any duplicates, even it means
+                        selecting more than `n` items.
+            .. versionadded:: 0.24.0
+        """
+        return DataFrame(query_compiler=self._query_compiler.nlargest(n, columns, keep))
 
     def nsmallest(self, n, columns, keep="first"):
-        return self._default_to_pandas(
-            pandas.DataFrame.nsmallest, n, columns, keep=keep
+        """
+        Return the first `n` rows ordered by `columns` in ascending order.
+        Return the first `n` rows with the smallest values in `columns`, in
+        ascending order. The columns that are not specified are returned as
+        well, but not used for ordering.
+        This method is equivalent to
+        ``df.sort_values(columns, ascending=True).head(n)``, but more
+        performant.
+        Parameters
+        ----------
+        n : int
+            Number of items to retrieve.
+        columns : list or str
+            Column name or names to order by.
+        keep : {'first', 'last', 'all'}, default 'first'
+            Where there are duplicate values:
+            - ``first`` : take the first occurrence.
+            - ``last`` : take the last occurrence.
+            - ``all`` : do not drop any duplicates, even it means
+              selecting more than `n` items.
+            .. versionadded:: 0.24.0
+        """
+        return DataFrame(
+            query_compiler=self._query_compiler.nsmallest(n, columns, keep)
         )
 
     def pivot(self, index=None, columns=None, values=None):


Test patch:
diff --git a/modin/pandas/test/test_dataframe.py b/modin/pandas/test/test_dataframe.py
--- a/modin/pandas/test/test_dataframe.py
+++ b/modin/pandas/test/test_dataframe.py
@@ -3933,68 +3933,70 @@ def test_mode(self, request, data, axis, numeric_only):
             df_equals(modin_result, pandas_result)
 
     def test_nlargest(self):
-        df = pd.DataFrame(
-            {
-                "population": [
-                    59000000,
-                    65000000,
-                    434000,
-                    434000,
-                    434000,
-                    337000,
-                    11300,
-                    11300,
-                    11300,
-                ],
-                "GDP": [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311],
-                "alpha-2": ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"],
-            },
-            index=[
-                "Italy",
-                "France",
-                "Malta",
-                "Maldives",
-                "Brunei",
-                "Iceland",
-                "Nauru",
-                "Tuvalu",
-                "Anguilla",
+        data = {
+            "population": [
+                59000000,
+                65000000,
+                434000,
+                434000,
+                434000,
+                337000,
+                11300,
+                11300,
+                11300,
             ],
+            "GDP": [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311],
+            "alpha-2": ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"],
+        }
+        index = [
+            "Italy",
+            "France",
+            "Malta",
+            "Maldives",
+            "Brunei",
+            "Iceland",
+            "Nauru",
+            "Tuvalu",
+            "Anguilla",
+        ]
+        modin_df = pd.DataFrame(data=data, index=index)
+        pandas_df = pandas.DataFrame(data=data, index=index)
+        df_equals(
+            modin_df.nlargest(3, "population"), pandas_df.nlargest(3, "population")
         )
-        with pytest.warns(UserWarning):
-            df.nlargest(3, "population")
 
     def test_nsmallest(self):
-        df = pd.DataFrame(
-            {
-                "population": [
-                    59000000,
-                    65000000,
-                    434000,
-                    434000,
-                    434000,
-                    337000,
-                    11300,
-                    11300,
-                    11300,
-                ],
-                "GDP": [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311],
-                "alpha-2": ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"],
-            },
-            index=[
-                "Italy",
-                "France",
-                "Malta",
-                "Maldives",
-                "Brunei",
-                "Iceland",
-                "Nauru",
-                "Tuvalu",
-                "Anguilla",
+        data = {
+            "population": [
+                59000000,
+                65000000,
+                434000,
+                434000,
+                434000,
+                337000,
+                11300,
+                11300,
+                11300,
             ],
+            "GDP": [1937894, 2583560, 12011, 4520, 12128, 17036, 182, 38, 311],
+            "alpha-2": ["IT", "FR", "MT", "MV", "BN", "IS", "NR", "TV", "AI"],
+        }
+        index = [
+            "Italy",
+            "France",
+            "Malta",
+            "Maldives",
+            "Brunei",
+            "Iceland",
+            "Nauru",
+            "Tuvalu",
+            "Anguilla",
+        ]
+        modin_df = pd.DataFrame(data=data, index=index)
+        pandas_df = pandas.DataFrame(data=data, index=index)
+        df_equals(
+            modin_df.nsmallest(3, "population"), pandas_df.nsmallest(3, "population")
         )
-        with pytest.warns(UserWarning):
-            df.nsmallest(3, "population")
 
     @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
     @pytest.mark.parametrize("axis", axis_values, ids=axis_keys)