Merge pull request #20081 from charris/backport-19607

BUG: Fix NaT handling in the PyArray_CompareFunc for datetime and timedelta
numpy · Oct 10, 2021 · 222dc20 · 222dc20
2 parents 94c8a3c + 73e7a12
commit 222dc20
Showing 1 changed file with 33 additions and 4 deletions.
diff --git a/numpy/core/src/multiarray/arraytypes.c.src b/numpy/core/src/multiarray/arraytypes.c.src
@@ -2805,11 +2805,9 @@ BOOL_compare(npy_bool *ip1, npy_bool *ip2, PyArrayObject *NPY_UNUSED(ap))
 
 /**begin repeat
  * #TYPE = BYTE, UBYTE, SHORT, USHORT, INT, UINT,
- *         LONG, ULONG, LONGLONG, ULONGLONG,
- *         DATETIME, TIMEDELTA#
+ *         LONG, ULONG, LONGLONG, ULONGLONG#
  * #type = npy_byte, npy_ubyte, npy_short, npy_ushort, npy_int, npy_uint,
- *         npy_long, npy_ulong, npy_longlong, npy_ulonglong,
- *         npy_datetime, npy_timedelta#
+ *         npy_long, npy_ulong, npy_longlong, npy_ulonglong#
  */
 
 static int
@@ -2920,6 +2918,37 @@ C@TYPE@_compare(@type@ *pa, @type@ *pb)
 
 /**end repeat**/
 
+/**begin repeat
+ * #TYPE = DATETIME, TIMEDELTA#
+ * #type = npy_datetime, npy_timedelta#
+ */
+
+static int
+@TYPE@_compare(@type@ *pa, @type@ *pb)
+{
+    const @type@ a = *pa;
+    const @type@ b = *pb;
+    int ret;
+
+    if (a == NPY_DATETIME_NAT) {
+        if (b == NPY_DATETIME_NAT) {
+            ret = 0;
+        }
+        else {
+            ret = 1;
+        }
+    }
+    else if (b == NPY_DATETIME_NAT) {
+        ret = -1;
+    }
+    else {
+        ret = a < b ? -1 : a == b ? 0 : 1;
+    }
+    return ret;
+}
+
+/**end repeat**/
+
 static int
 HALF_compare (npy_half *pa, npy_half *pb, PyArrayObject *NPY_UNUSED(ap))
 {