BUG: nlargest/nsmallest can now consider nan values like sort_values(ascending=True).head(n)

pandas/core/algorithms.py

@@ -1247,6 +1247,8 @@ class SelectNSeries(SelectN):
 
     def compute(self, method: str) -> Series:
 
+        from pandas.core.reshape.concat import concat
+
         n = self.n
         dtype = self.obj.dtype
         if not self.is_valid_dtype_n_method(dtype):
@@ -1256,6 +1258,7 @@ def compute(self, method: str) -> Series:
             return self.obj[[]]
 
         dropped = self.obj.dropna()
+        nan_index = self.obj.drop(dropped.index)
 
         if is_extension_array_dtype(dropped.dtype):
             # GH#41816 bc we have dropped NAs above, MaskedArrays can use the
@@ -1272,7 +1275,7 @@ def compute(self, method: str) -> Series:
         # slow method
         if n >= len(self.obj):
             ascending = method == "nsmallest"
-            return dropped.sort_values(ascending=ascending).head(n)
+            return self.obj.sort_values(ascending=ascending).head(n)
 
         # fast method
         new_dtype = dropped.dtype
@@ -1290,6 +1293,8 @@ def compute(self, method: str) -> Series:
         if self.keep == "last":
             arr = arr[::-1]
 
+        nbase = n
+        findex = len(self.obj)
         narr = len(arr)
         n = min(n, narr)
 
@@ -1301,12 +1306,13 @@ def compute(self, method: str) -> Series:
 
         if self.keep != "all":
             inds = inds[:n]
+            findex = nbase
 
         if self.keep == "last":
             # reverse indices
             inds = narr - 1 - inds
 
-        return dropped.iloc[inds]
+        return concat([dropped.iloc[inds], nan_index])[:findex]
 
 
 class SelectNFrame(SelectN):

pandas/tests/frame/methods/test_nlargest.py

@@ -209,3 +209,10 @@ def test_nlargest_multiindex_column_lookup(self):
         result = df.nlargest(3, ("x", "b"))
         expected = df.iloc[[3, 2, 1]]
         tm.assert_frame_equal(result, expected)
+
+    def test_nlargest_nan(self):
+        # GH#43060
+        df = pd.DataFrame([np.nan, np.nan, 0, 1, 2, 3])
+        result = df.nlargest(5, 0)
+        expected = df.sort_values(0, ascending=False).head(5)
+        tm.assert_frame_equal(result, expected)

pandas/tests/series/methods/test_nlargest.py

@@ -127,8 +127,8 @@ def test_nsmallest_nlargest(self, s_main_dtypes_split):
     def test_nlargest_misc(self):
 
         ser = Series([3.0, np.nan, 1, 2, 5])
-        tm.assert_series_equal(ser.nlargest(), ser.iloc[[4, 0, 3, 2]])
-        tm.assert_series_equal(ser.nsmallest(), ser.iloc[[2, 3, 0, 4]])
+        tm.assert_series_equal(ser.nlargest(), ser.iloc[[4, 0, 3, 2, 1]])
+        tm.assert_series_equal(ser.nsmallest(), ser.iloc[[2, 3, 0, 4, 1]])
 
         msg = 'keep must be either "first", "last"'
         with pytest.raises(ValueError, match=msg):