Add lexsort support for datetime and timedelta

7 files changed, 55 insertions, 10 deletions

diff --git a/numpy/core/src/multiarray/arraytypes.c.src b/numpy/core/src/multiarray/arraytypes.c.src
index dbbfa26b2..1bfbf160d 100644
--- a/numpy/core/src/multiarray/arraytypes.c.src
+++ b/numpy/core/src/multiarray/arraytypes.c.src
@@ -3720,7 +3720,7 @@ static PyArray_Descr @from@_Descr = {
  *         half, float, double, longdouble,
  *         cfloat, cdouble, clongdouble,
  *         object, datetime, timedelta#
- * #sort = 1*18, 0*3#
+ * #sort = 1*18, 0*1, 1*2#
  * #num = 1*15, 2*3, 1*3#
  * #fromtype = npy_bool,
  *             npy_byte, npy_ubyte, npy_short, npy_ushort, npy_int, npy_uint,
diff --git a/numpy/core/src/npysort/heapsort.c.src b/numpy/core/src/npysort/heapsort.c.src
index 45d9deb92..84c9d7bd4 100644
--- a/numpy/core/src/npysort/heapsort.c.src
+++ b/numpy/core/src/npysort/heapsort.c.src
@@ -50,14 +50,14 @@
  *
  * #TYPE = BOOL, BYTE, UBYTE, SHORT, USHORT, INT, UINT, LONG, ULONG,
  *         LONGLONG, ULONGLONG, HALF, FLOAT, DOUBLE, LONGDOUBLE,
- *         CFLOAT, CDOUBLE, CLONGDOUBLE#
+ *         CFLOAT, CDOUBLE, CLONGDOUBLE, DATETIME, TIMEDELTA#
  * #suff = bool, byte, ubyte, short, ushort, int, uint, long, ulong,
  *         longlong, ulonglong, half, float, double, longdouble,
- *         cfloat, cdouble, clongdouble#
+ *         cfloat, cdouble, clongdouble, datetime, timedelta#
  * #type = npy_bool, npy_byte, npy_ubyte, npy_short, npy_ushort, npy_int,
  *         npy_uint, npy_long, npy_ulong, npy_longlong, npy_ulonglong,
  *         npy_ushort, npy_float, npy_double, npy_longdouble, npy_cfloat,
- *         npy_cdouble, npy_clongdouble#
+ *         npy_cdouble, npy_clongdouble, npy_datetime, npy_timedelta#
  */
 
 int
diff --git a/numpy/core/src/npysort/mergesort.c.src b/numpy/core/src/npysort/mergesort.c.src
index 3bd6ee1a6..7f98c4016 100644
--- a/numpy/core/src/npysort/mergesort.c.src
+++ b/numpy/core/src/npysort/mergesort.c.src
@@ -50,14 +50,14 @@
  *
  * #TYPE = BOOL, BYTE, UBYTE, SHORT, USHORT, INT, UINT, LONG, ULONG,
  *         LONGLONG, ULONGLONG, HALF, FLOAT, DOUBLE, LONGDOUBLE,
- *         CFLOAT, CDOUBLE, CLONGDOUBLE#
+ *         CFLOAT, CDOUBLE, CLONGDOUBLE, DATETIME, TIMEDELTA#
  * #suff = bool, byte, ubyte, short, ushort, int, uint, long, ulong,
  *         longlong, ulonglong, half, float, double, longdouble,
- *         cfloat, cdouble, clongdouble#
+ *         cfloat, cdouble, clongdouble, datetime, timedelta#
  * #type = npy_bool, npy_byte, npy_ubyte, npy_short, npy_ushort, npy_int,
  *         npy_uint, npy_long, npy_ulong, npy_longlong, npy_ulonglong,
  *         npy_ushort, npy_float, npy_double, npy_longdouble, npy_cfloat,
- *         npy_cdouble, npy_clongdouble#
+ *         npy_cdouble, npy_clongdouble, npy_datetime, npy_timedelta#
  */
 
 static void
diff --git a/numpy/core/src/npysort/npysort_common.h b/numpy/core/src/npysort/npysort_common.h
index 318a80222..a3ce7664d 100644
--- a/numpy/core/src/npysort/npysort_common.h
+++ b/numpy/core/src/npysort/npysort_common.h
@@ -28,6 +28,8 @@
 #define CFLOAT_SWAP(a,b) {npy_cfloat tmp = (b); (b)=(a); (a) = tmp;}
 #define CDOUBLE_SWAP(a,b) {npy_cdouble tmp = (b); (b)=(a); (a) = tmp;}
 #define CLONGDOUBLE_SWAP(a,b) {npy_clongdouble tmp = (b); (b)=(a); (a) = tmp;}
+#define DATETIME_SWAP(a,b) {npy_datetime tmp = (b); (b)=(a); (a) = tmp;}
+#define TIMEDELTA_SWAP(a,b) {npy_timedelta tmp = (b); (b)=(a); (a) = tmp;}
 
 /* Need this for the argsort functions */
 #define INTP_SWAP(a,b) {npy_intp tmp = (b); (b)=(a); (a) = tmp;}
@@ -324,6 +326,20 @@ UNICODE_LT(npy_ucs4 *s1, npy_ucs4 *s2, size_t len)
 }
 
 
+NPY_INLINE static int
+DATETIME_LT(npy_datetime a, npy_datetime b)
+{
+    return a < b;
+}
+
+
+NPY_INLINE static int
+TIMEDELTA_LT(npy_timedelta a, npy_timedelta b)
+{
+    return a < b;
+}
+
+
 NPY_INLINE static void
 GENERIC_COPY(char *a, char *b, size_t len)
 {
diff --git a/numpy/core/src/npysort/quicksort.c.src b/numpy/core/src/npysort/quicksort.c.src
index 2225887b1..272615ab3 100644
--- a/numpy/core/src/npysort/quicksort.c.src
+++ b/numpy/core/src/npysort/quicksort.c.src
@@ -50,14 +50,14 @@
  *
  * #TYPE = BOOL, BYTE, UBYTE, SHORT, USHORT, INT, UINT, LONG, ULONG,
  *         LONGLONG, ULONGLONG, HALF, FLOAT, DOUBLE, LONGDOUBLE,
- *         CFLOAT, CDOUBLE, CLONGDOUBLE#
+ *         CFLOAT, CDOUBLE, CLONGDOUBLE, DATETIME, TIMEDELTA#
  * #suff = bool, byte, ubyte, short, ushort, int, uint, long, ulong,
  *         longlong, ulonglong, half, float, double, longdouble,
- *         cfloat, cdouble, clongdouble#
+ *         cfloat, cdouble, clongdouble, datetime, timedelta#
  * #type = npy_bool, npy_byte, npy_ubyte, npy_short, npy_ushort, npy_int,
  *         npy_uint, npy_long, npy_ulong, npy_longlong, npy_ulonglong,
  *         npy_ushort, npy_float, npy_double, npy_longdouble, npy_cfloat,
- *         npy_cdouble, npy_clongdouble#
+ *         npy_cdouble, npy_clongdouble, npy_datetime, npy_timedelta#
  */
 
 int
diff --git a/numpy/core/src/private/npy_sort.h b/numpy/core/src/private/npy_sort.h
index 120e7efc2..46825a0c5 100644
--- a/numpy/core/src/private/npy_sort.h
+++ b/numpy/core/src/private/npy_sort.h
@@ -171,6 +171,22 @@ int aheapsort_unicode(npy_ucs4 *vec, npy_intp *ind, npy_intp cnt, PyArrayObject
 int amergesort_unicode(npy_ucs4 *vec, npy_intp *ind, npy_intp cnt, PyArrayObject *arr);
 
 
+int quicksort_datetime(npy_datetime *vec, npy_intp cnt, void *null);
+int heapsort_datetime(npy_datetime *vec, npy_intp cnt, void *null);
+int mergesort_datetime(npy_datetime *vec, npy_intp cnt, void *null);
+int aquicksort_datetime(npy_datetime *vec, npy_intp *ind, npy_intp cnt, void *null);
+int aheapsort_datetime(npy_datetime *vec, npy_intp *ind, npy_intp cnt, void *null);
+int amergesort_datetime(npy_datetime *vec, npy_intp *ind, npy_intp cnt, void *null);
+
+
+int quicksort_timedelta(npy_timedelta *vec, npy_intp cnt, void *null);
+int heapsort_timedelta(npy_timedelta *vec, npy_intp cnt, void *null);
+int mergesort_timedelta(npy_timedelta *vec, npy_intp cnt, void *null);
+int aquicksort_timedelta(npy_timedelta *vec, npy_intp *ind, npy_intp cnt, void *null);
+int aheapsort_timedelta(npy_timedelta *vec, npy_intp *ind, npy_intp cnt, void *null);
+int amergesort_timedelta(npy_timedelta *vec, npy_intp *ind, npy_intp cnt, void *null);
+
+
 int npy_quicksort(void *base, size_t num, size_t size, npy_comparator cmp);
 int npy_heapsort(void *base, size_t num, size_t size, npy_comparator cmp);
 int npy_mergesort(void *base, size_t num, size_t size, npy_comparator cmp);
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index 7b9cf0932..68bdaf97a 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -2282,6 +2282,19 @@ class TestLexsort(TestCase):
 
         assert_array_equal(x[1][idx], np.sort(x[1]))
 
+    def test_datetime(self):
+        a = np.array([0,0,0], dtype='datetime64[D]')
+        b = np.array([2,1,0], dtype='datetime64[D]')
+        idx = np.lexsort((b, a))
+        expected_idx = np.array([2, 1, 0])
+        assert_array_equal(idx, expected_idx)
+
+        a = np.array([0,0,0], dtype='timedelta64[D]')
+        b = np.array([2,1,0], dtype='timedelta64[D]')
+        idx = np.lexsort((b, a))
+        expected_idx = np.array([2, 1, 0])
+        assert_array_equal(idx, expected_idx)
+
 
 class TestIO(object):
     """Test tofile, fromfile, tostring, and fromstring"""