DOC+TST: Specify np.nan sort order in partition functions

Author: JuliaPoo

numpy/_core/fromnumeric.py

@@ -823,6 +823,8 @@ def partition(a, kth, axis=-1, kind='introselect', order=None):
     the real parts except when they are equal, in which case the order
     is determined by the imaginary parts.
 
+    The sort order of ``np.nan`` is bigger than ``np.inf``.
+
     Examples
     --------
     >>> a = np.array([7, 1, 7, 7, 1, 5, 7, 2, 3, 2, 6, 2, 3, 0])
@@ -923,7 +925,7 @@ def argpartition(a, kth, axis=-1, kind='introselect', order=None):
     is unstable, and hence the returned indices are not guaranteed
     to be the earliest/latest occurrence of the element.
 
-    The treatment of ``np.nan`` in the input array is undefined.
+    The sort order of ``np.nan`` is bigger than ``np.inf``.
 
     See `partition` for notes on the different selection algorithms.
 

numpy/_core/tests/test_multiarray.py

@@ -10211,6 +10211,16 @@ def test_partition_fp(N, dtype):
     assert_arr_partitioned(np.sort(arr)[k], k,
             arr[np.argpartition(arr, k, kind='introselect')])
 
+    # Check that `np.inf < np.nan`
+    # This follows np.sort
+    arr[0] = np.nan
+    arr[1] = np.inf
+    o1 = np.partition(arr, -2, kind='introselect')
+    o2 = arr[np.argpartition(arr, -2, kind='introselect')]
+    for out in [o1,o2]:
+        assert_(np.isnan(out[-1]))
+        assert_equal(out[-2], np.inf)
+
 def test_cannot_assign_data():
     a = np.arange(10)
     b = np.linspace(0, 1, 10)