Merge pull request #12658 from tylerjereddy/NaT_sort_order

BUG: NaT now sorts to ends of arrays

3 files changed, 44 insertions, 0 deletions

diff --git a/numpy/core/fromnumeric.py b/numpy/core/fromnumeric.py
index 6e5f3dabf..d454480a8 100644
--- a/numpy/core/fromnumeric.py
+++ b/numpy/core/fromnumeric.py
@@ -944,6 +944,10 @@ def sort(a, axis=-1, kind=None, order=None):
     'mergesort' and 'stable' are mapped to radix sort for integer data types. Radix sort is an
     O(n) sort instead of O(n log n).
 
+    .. versionchanged:: 1.17.0
+
+    NaT now sorts to the end of arrays for consistency with NaN.
+
     Examples
     --------
     >>> a = np.array([[1,4],[3,1]])
diff --git a/numpy/core/src/npysort/npysort_common.h b/numpy/core/src/npysort/npysort_common.h
index 5fd03b96f..30c0d47f3 100644
--- a/numpy/core/src/npysort/npysort_common.h
+++ b/numpy/core/src/npysort/npysort_common.h
@@ -329,6 +329,14 @@ UNICODE_LT(const npy_ucs4 *s1, const npy_ucs4 *s2, size_t len)
 NPY_INLINE static int
 DATETIME_LT(npy_datetime a, npy_datetime b)
 {
+    if (a == NPY_DATETIME_NAT) {
+        return 0;
+    }
+
+    if (b == NPY_DATETIME_NAT) {
+        return 1;
+    }
+
     return a < b;
 }
 
diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py
index 8bc6476da..41b84a69f 100644
--- a/numpy/core/tests/test_datetime.py
+++ b/numpy/core/tests/test_datetime.py
@@ -145,6 +145,38 @@ class TestDateTime(object):
         assert_(np.datetime64('NaT') != np.datetime64('NaT', 'us'))
         assert_(np.datetime64('NaT', 'us') != np.datetime64('NaT'))
 
+
+
+    @pytest.mark.parametrize("size", [
+        3, 21, 217, 1000])
+    def test_nat_argsort_stability(self, size):
+        # NaT < NaT should be False internally for
+        # sort stability
+        expected = np.arange(size)
+        arr = np.tile(np.datetime64('NaT'), size)
+        assert_equal(np.argsort(arr, kind='mergesort'), expected)
+
+    @pytest.mark.parametrize("arr, expected", [
+        # the example provided in gh-12629
+        (np.array(['NaT', 1, 2, 3], dtype='M8[ns]'),
+         np.array([1, 2, 3, 'NaT'], dtype='M8[ns]')),
+        # multiple NaTs
+        (np.array(['NaT', 9, 'NaT', -707], dtype='M8[s]'),
+         np.array([-707, 9, 'NaT', 'NaT'], dtype='M8[s]')),
+        # this sort explores another code path for NaT
+        (np.array([1, -2, 3, 'NaT'], dtype='M8[ns]'),
+         np.array([-2, 1, 3, 'NaT'], dtype='M8[ns]')),
+        # 2-D array
+        (np.array([[51, -220, 'NaT'],
+                   [-17, 'NaT', -90]], dtype='M8[us]'),
+         np.array([[-220, 51, 'NaT'],
+                   [-90, -17, 'NaT']], dtype='M8[us]')),
+        ])
+    def test_sort_nat(self, arr, expected):
+        # fix for gh-12629; NaT sorting to end of array
+        arr.sort()
+        assert_equal(arr, expected)
+
     def test_datetime_scalar_construction(self):
         # Construct with different units
         assert_equal(np.datetime64('1950-03-12', 'D'),