Merge pull request #79 from ngoldbaum/datetime-cast

Add a roundtrip cast for datetimes

Author: ngoldbaum

stringdtype/stringdtype/src/casts.c

@@ -907,17 +907,154 @@ STRING_TO_FLOAT_RESOLVE_DESCRIPTORS(float16, HALF)
 STRING_TO_FLOAT_CAST(float16, f16, npy_half_isinf, npy_double_to_half)
 FLOAT_TO_STRING_CAST(float16, f16, npy_half_to_double)
 
+// string to datetime
+
+static NPY_CASTING
+string_to_datetime_resolve_descriptors(
+        PyObject *NPY_UNUSED(self), PyArray_DTypeMeta *NPY_UNUSED(dtypes[2]),
+        PyArray_Descr *given_descrs[2], PyArray_Descr *loop_descrs[2],
+        npy_intp *NPY_UNUSED(view_offset))
+{
+    if (given_descrs[1] == NULL) {
+        PyErr_SetString(PyExc_TypeError,
+                        "Casting from StringDType to datetimes without a unit "
+                        "is not currently supported");
+        return (NPY_CASTING)-1;
+    }
+    else {
+        Py_INCREF(given_descrs[1]);
+        loop_descrs[1] = given_descrs[1];
+    }
+
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[0] = given_descrs[0];
+
+    return NPY_UNSAFE_CASTING;
+}
+
+static int
+string_to_datetime(PyArrayMethod_Context *context, char *const data[],
+                   npy_intp const dimensions[], npy_intp const strides[],
+                   NpyAuxData *NPY_UNUSED(auxdata))
+{
+    npy_intp N = dimensions[0];
+    char *in = data[0];
+    npy_datetime *out = (npy_datetime *)data[1];
+
+    npy_intp in_stride = strides[0];
+    npy_intp out_stride = strides[1] / sizeof(npy_datetime);
+
+    ss *s = NULL;
+    npy_datetimestruct dts;
+    NPY_DATETIMEUNIT in_unit = -1;
+    PyArray_DatetimeMetaData in_meta = {0, 1};
+    npy_bool out_special;
+
+    PyArray_Descr *dt_descr = context->descriptors[1];
+    PyArray_DatetimeMetaData *dt_meta =
+            &(((PyArray_DatetimeDTypeMetaData *)dt_descr->c_metadata)->meta);
+
+    while (N--) {
+        s = (ss *)in;
+        if (ss_isnull(s)) {
+            *out = NPY_DATETIME_NAT;
+        }
+        if (NpyDatetime_ParseISO8601Datetime(
+                    (const char *)s->buf, s->len, in_unit, NPY_UNSAFE_CASTING,
+                    &dts, &in_meta.base, &out_special) < 0) {
+            return -1;
+        }
+        if (NpyDatetime_ConvertDatetimeStructToDatetime64(dt_meta, &dts, out) <
+            0) {
+            return -1;
+        }
+
+        in += in_stride;
+        out += out_stride;
+    }
+
+    return 0;
+}
+
+static PyType_Slot s2dt_slots[] = {
+        {NPY_METH_resolve_descriptors,
+         &string_to_datetime_resolve_descriptors},
+        {NPY_METH_strided_loop, &string_to_datetime},
+        {0, NULL}};
+
+static char *s2dt_name = "cast_StringDType_to_Datetime";
+
+// datetime to string
+
+static int
+datetime_to_string(PyArrayMethod_Context *context, char *const data[],
+                   npy_intp const dimensions[], npy_intp const strides[],
+                   NpyAuxData *NPY_UNUSED(auxdata))
+{
+    npy_intp N = dimensions[0];
+    npy_datetime *in = (npy_datetime *)data[0];
+    char *out = data[1];
+
+    npy_intp in_stride = strides[0] / sizeof(npy_datetime);
+    npy_intp out_stride = strides[1];
+
+    npy_datetimestruct dts;
+    PyArray_Descr *dt_descr = context->descriptors[0];
+    PyArray_DatetimeMetaData *dt_meta =
+            &(((PyArray_DatetimeDTypeMetaData *)dt_descr->c_metadata)->meta);
+    // buffer passed to numpy to build datetime string
+    char datetime_buf[NPY_DATETIME_MAX_ISO8601_STRLEN];
+
+    while (N--) {
+        ss *out_ss = (ss *)out;
+        ssfree(out_ss);
+        if (*in == NPY_DATETIME_NAT) {
+            /* convert to NA */
+            out_ss = NULL;
+        }
+        else {
+            if (NpyDatetime_ConvertDatetime64ToDatetimeStruct(dt_meta, *in,
+                                                              &dts) < 0) {
+                return -1;
+            }
+
+            // zero out buffer
+            memset(datetime_buf, 0, NPY_DATETIME_MAX_ISO8601_STRLEN);
+
+            if (NpyDatetime_MakeISO8601Datetime(
+                        &dts, datetime_buf, NPY_DATETIME_MAX_ISO8601_STRLEN, 0,
+                        0, dt_meta->base, -1, NPY_UNSAFE_CASTING) < 0) {
+                return -1;
+            }
+
+            if (ssnewlen(datetime_buf, strlen(datetime_buf), out_ss) < 0) {
+                PyErr_SetString(PyExc_MemoryError, "ssnewlen failed");
+                return -1;
+            }
+        }
+
+        in += in_stride;
+        out += out_stride;
+    }
+
+    return 0;
+}
+
+static PyType_Slot dt2s_slots[] = {
+        {NPY_METH_resolve_descriptors,
+         &any_to_string_UNSAFE_resolve_descriptors},
+        {NPY_METH_strided_loop, &datetime_to_string},
+        {0, NULL}};
+
+static char *dt2s_name = "cast_Datetime_to_StringDType";
+
 // TODO: longdouble
 //        punting on this one because numpy's C routines for handling
 //        longdouble are not public (specifically NumPyOS_ascii_strtold)
 //        also this type is kinda niche and is not needed by pandas
 //
 //       cfloat, cdouble, and clongdouble
 //        not hard to do in principle but not needed by pandas.
-//
-//       datetime
-//        numpy's utilities for parsing a string into a datetime
-//        are not public (specifically parse_iso_8601_datetime).
 
 PyArrayMethod_Spec *
 get_cast_spec(const char *name, NPY_CASTING casting,
@@ -961,7 +1098,7 @@ get_casts()
             get_cast_spec(t2t_name, NPY_NO_CASTING,
                           NPY_METH_SUPPORTS_UNALIGNED, t2t_dtypes, s2s_slots);
 
-    int num_casts = 27;
+    int num_casts = 29;
 
 #if NPY_SIZEOF_BYTE == NPY_SIZEOF_SHORT
     num_casts += 4;
@@ -1033,6 +1170,22 @@ get_casts()
     DTYPES_AND_CAST_SPEC(f32, Float)
     DTYPES_AND_CAST_SPEC(f16, Half)
 
+    PyArray_DTypeMeta **s2dt_dtypes = get_dtypes(
+            (PyArray_DTypeMeta *)&StringDType, &PyArray_DatetimeDType);
+
+    PyArrayMethod_Spec *StringToDatetimeCastSpec = get_cast_spec(
+            s2dt_name, NPY_UNSAFE_CASTING,
+            NPY_METH_NO_FLOATINGPOINT_ERRORS | NPY_METH_REQUIRES_PYAPI,
+            s2dt_dtypes, s2dt_slots);
+
+    PyArray_DTypeMeta **dt2s_dtypes = get_dtypes(
+            &PyArray_DatetimeDType, (PyArray_DTypeMeta *)&StringDType);
+
+    PyArrayMethod_Spec *DatetimeToStringCastSpec = get_cast_spec(
+            dt2s_name, NPY_UNSAFE_CASTING,
+            NPY_METH_NO_FLOATINGPOINT_ERRORS | NPY_METH_REQUIRES_PYAPI,
+            dt2s_dtypes, dt2s_slots);
+
     PyArrayMethod_Spec **casts =
             malloc((num_casts + 1) * sizeof(PyArrayMethod_Spec *));
 
@@ -1089,6 +1242,8 @@ get_casts()
     casts[cast_i++] = FloatToStringCastSpec;
     casts[cast_i++] = StringToHalfCastSpec;
     casts[cast_i++] = HalfToStringCastSpec;
+    casts[cast_i++] = StringToDatetimeCastSpec;
+    casts[cast_i++] = DatetimeToStringCastSpec;
     casts[cast_i++] = NULL;
 
     assert(casts[num_casts] == NULL);

stringdtype/stringdtype/src/casts.h

@@ -1,15 +1,15 @@
 #ifndef _NPY_CASTS_H
 #define _NPY_CASTS_H
 
+// needed for Py_UCS4
 #include <Python.h>
 
+// need these defines and includes for PyArrayMethod_Spec
 #define PY_ARRAY_UNIQUE_SYMBOL stringdtype_ARRAY_API
-#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#define NPY_NO_DEPRECATED_API NPY_2_0_API_VERSION
+#define NPY_TARGET_VERSION NPY_2_0_API_VERSION
 #define NO_IMPORT_ARRAY
-#include "numpy/arrayobject.h"
 #include "numpy/experimental_dtype_api.h"
-#include "numpy/halffloat.h"
-#include "numpy/ndarraytypes.h"
 
 PyArrayMethod_Spec **
 get_casts();

stringdtype/stringdtype/src/dtype.h

@@ -7,12 +7,15 @@
 // clang-format on
 
 #define PY_ARRAY_UNIQUE_SYMBOL stringdtype_ARRAY_API
-#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#define NPY_NO_DEPRECATED_API NPY_2_0_API_VERSION
+#define NPY_TARGET_VERSION NPY_2_0_API_VERSION
 #define NO_IMPORT_ARRAY
 #include "numpy/arrayobject.h"
 #include "numpy/experimental_dtype_api.h"
+#include "numpy/halffloat.h"
 #include "numpy/ndarraytypes.h"
 #include "numpy/npy_math.h"
+#include "numpy/ufuncobject.h"
 
 typedef struct {
     PyArray_Descr base;

stringdtype/stringdtype/src/main.c

@@ -1,7 +1,8 @@
 #include <Python.h>
 
 #define PY_ARRAY_UNIQUE_SYMBOL stringdtype_ARRAY_API
-#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
+#define NPY_NO_DEPRECATED_API NPY_2_0_API_VERSION
+#define NPY_TARGET_VERSION NPY_2_0_API_VERSION
 #include "numpy/arrayobject.h"
 #include "numpy/experimental_dtype_api.h"
 
@@ -88,9 +89,8 @@ static struct PyModuleDef moduledef = {
 PyMODINIT_FUNC
 PyInit__main(void)
 {
-    if (_import_array() < 0) {
-        return NULL;
-    }
+    import_array();
+
     if (import_experimental_dtype_api(13) < 0) {
         return NULL;
     }

stringdtype/stringdtype/src/umath.c

@@ -1,17 +1,10 @@
 #include <Python.h>
 
-#define PY_ARRAY_UNIQUE_SYMBOL stringdtype_ARRAY_API
-#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
-#define NO_IMPORT_ARRAY
-#include "numpy/arrayobject.h"
-#include "numpy/experimental_dtype_api.h"
-#include "numpy/ndarraytypes.h"
-#include "numpy/ufuncobject.h"
+#include "umath.h"
 
 #include "dtype.h"
 #include "static_string.h"
 #include "string.h"
-#include "umath.h"
 
 static NPY_CASTING
 multiply_resolve_descriptors(

stringdtype/tests/test_stringdtype.py

@@ -515,3 +515,29 @@ def test_create_with_na(dtype):
         == f"array(['hello', {dtype.na_object}, 'world'], dtype={dtype})"
     )
     assert arr[1] is dtype.na_object
+
+
+def test_datetime_cast(dtype):
+    a = np.array(
+        [
+            np.datetime64("1923-04-14T12:43:12"),
+            np.datetime64("1994-06-21T14:43:15"),
+            np.datetime64("2001-10-15T04:10:32"),
+            np.datetime64("NaT"),
+            np.datetime64("1995-11-25T16:02:16"),
+            np.datetime64("2005-01-04T03:14:12"),
+            np.datetime64("2041-12-03T14:05:03"),
+        ]
+    )
+    sa = a.astype(dtype)
+    assert sa[3] is dtype.na_object
+
+    ra = sa.astype(a.dtype)
+    assert np.isnat(ra[3])
+
+    np.testing.assert_array_equal(a, ra)
+
+    # don't worry about comparing how NaT is converted
+    sa = np.delete(sa, 3)
+    a = np.delete(a, 3)
+    np.testing.assert_array_equal(sa, a.astype("U"))