ENH: Add a NpyAuxData c_metadata to PyArray_Descr

The purpose of this is to allow dtypes with metadata to
have faster C-level performance. With this, datetime64
metadata lookup can become a single C pointer dereference,
compared to a C-string to Python-string conversion +
a Python dict lookup as it is presently.

3 files changed, 190 insertions, 95 deletions

diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index ec890f87d..31a4bbfef 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -1,7 +1,7 @@
 #ifndef NDARRAYTYPES_H
 #define NDARRAYTYPES_H
 
-/* This is auto-generated by the installer */
+/* numpyconfig.h is auto-generated by the installer */
 #include "numpyconfig.h"
 
 #include "npy_common.h"
@@ -290,6 +290,41 @@ typedef enum {
     NPY_BUSDAY_RAISE
 } NPY_BUSDAY_ROLL;
 
+/************************************************************
+ * NumPy Auxiliary Data for inner loops, sort functions, etc.
+ ************************************************************/
+
+/*
+ * When creating an auxiliary data struct, this should always appear
+ * as the first member, like this:
+ *
+ * typedef struct {
+ *     NpyAuxData base;
+ *     double constant;
+ * } constant_multiplier_aux_data;
+ */
+typedef struct NpyAuxData_tag NpyAuxData;
+
+/* Function pointers for freeing or cloning auxiliary data */
+typedef void (NpyAuxData_FreeFunc) (NpyAuxData *);
+typedef NpyAuxData *(NpyAuxData_CloneFunc) (NpyAuxData *);
+
+struct NpyAuxData_tag {
+    NpyAuxData_FreeFunc *free;
+    NpyAuxData_CloneFunc *clone;
+    /* To allow for a bit of expansion without breaking the ABI */
+    void *reserved[2];
+};
+
+/* Macros to use for freeing and cloning auxiliary data */
+#define NPY_AUXDATA_FREE(auxdata) \
+    if ((auxdata) == NULL) \
+        ; \
+    else \
+        ((auxdata)->free(auxdata))
+#define NPY_AUXDATA_CLONE(auxdata) \
+    ((auxdata)->clone(auxdata))
+
 /*********************************************************************
  * NumPy functions for dealing with masks, such as in masked iteration
  *********************************************************************/
@@ -571,44 +606,60 @@ typedef struct {
 
 typedef struct _PyArray_Descr {
         PyObject_HEAD
-        PyTypeObject *typeobj;  /*
-                                 * the type object representing an
-                                 * instance of this type -- should not
-                                 * be two type_numbers with the same type
-                                 * object.
-                                 */
-        char kind;              /* kind for this type */
-        char type;              /* unique-character representing this type */
-        char byteorder;         /*
-                                 * '>' (big), '<' (little), '|'
-                                 * (not-applicable), or '=' (native).
-                                 */
-        char flags;             /* flags describing data type */
-        int type_num;           /* number representing this type */
-        int elsize;             /* element size for this type */
-        int alignment;          /* alignment needed for this type */
-        struct _arr_descr                                       \
-        *subarray;              /*
-                                 * Non-NULL if this type is
-                                 * is an array (C-contiguous)
-                                 * of some other type
-                                 */
-        PyObject *fields;       /* The fields dictionary for this type
-                                 * For statically defined descr this
-                                 * is always Py_None
-                                 */
-
-        PyObject *names;        /*
-                                 * An ordered tuple of field names or NULL
-                                 * if no fields are defined
-                                 */
-
-        PyArray_ArrFuncs *f;     /*
-                                  * a table of functions specific for each
-                                  * basic data descriptor
-                                  */
-
-        PyObject *metadata;     /* Metadata about this dtype */
+        /*
+         * the type object representing an
+         * instance of this type -- should not
+         * be two type_numbers with the same type
+         * object.
+         */
+        PyTypeObject *typeobj;
+        /* kind for this type */
+        char kind;
+        /* unique-character representing this type */
+        char type;
+        /*
+         * '>' (big), '<' (little), '|'
+         * (not-applicable), or '=' (native).
+         */
+        char byteorder;
+        /* flags describing data type */
+        char flags;
+        /* number representing this type */
+        int type_num;
+        /* element size (itemsize) for this type */
+        int elsize;
+        /* alignment needed for this type */
+        int alignment;
+        /*
+         * Non-NULL if this type is
+         * is an array (C-contiguous)
+         * of some other type
+         */
+        struct _arr_descr *subarray;
+        /*
+         * The fields dictionary for this type
+         * For statically defined descr this
+         * is always Py_None
+         */
+        PyObject *fields;
+        /*
+         * An ordered tuple of field names or NULL
+         * if no fields are defined
+         */
+        PyObject *names;
+        /*
+         * a table of functions specific for each
+         * basic data descriptor
+         */
+        PyArray_ArrFuncs *f;
+        /* Metadata about this dtype */
+        PyObject *metadata;
+        /*
+         * Metadata specific to the C implementation
+         * of the particular dtype. This was added
+         * for NumPy 1.7.0.
+         */
+        NpyAuxData *c_metadata;
 } PyArray_Descr;
 
 typedef struct _arr_descr {
@@ -1748,41 +1799,6 @@ PyArray_HASMASKNA(PyArrayObject *arr)
 #define PyDataType_ISBYTESWAPPED(d) (!PyDataType_ISNOTSWAPPED(d))
 
 /************************************************************
- * NumPy Auxiliary Data for inner loops, sort functions, etc.
- ************************************************************/
-
-/*
- * When creating an auxiliary data struct, this should always appear
- * as the first member, like this:
- *
- * typedef struct {
- *     NpyAuxData base;
- *     double constant;
- * } constant_multiplier_aux_data;
- */
-typedef struct NpyAuxData_tag NpyAuxData;
-
-/* Function pointers for freeing or cloning auxiliary data */
-typedef void (NpyAuxData_FreeFunc) (NpyAuxData *);
-typedef NpyAuxData *(NpyAuxData_CloneFunc) (NpyAuxData *);
-
-struct NpyAuxData_tag {
-    NpyAuxData_FreeFunc *free;
-    NpyAuxData_CloneFunc *clone;
-    /* To allow for a bit of expansion without breaking the ABI */
-    void *reserved[2];
-};
-
-/* Macros to use for freeing and cloning auxiliary data */
-#define NPY_AUXDATA_FREE(auxdata) \
-    if ((auxdata) == NULL) \
-        ; \
-    else \
-        ((auxdata)->free(auxdata))
-#define NPY_AUXDATA_CLONE(auxdata) \
-    ((auxdata)->clone(auxdata))
-
-/************************************************************
  * A struct used by PyArray_CreateSortedStridePerm, new in 1.7.
  ************************************************************/
 
diff --git a/numpy/core/src/multiarray/arraytypes.c.src b/numpy/core/src/multiarray/arraytypes.c.src
index 9dbcb8439..d9e461a15 100644
--- a/numpy/core/src/multiarray/arraytypes.c.src
+++ b/numpy/core/src/multiarray/arraytypes.c.src
@@ -3699,18 +3699,33 @@ static PyArray_ArrFuncs _Py@NAME@_ArrFuncs = {
  */
 static PyArray_Descr @from@_Descr = {
     PyObject_HEAD_INIT(&PyArrayDescr_Type)
+    /* typeobj */
     &Py@NAME@ArrType_Type,
+    /* kind */
     NPY_@from@LTR,
+    /* type */
     NPY_@from@LTR,
+    /* byteorder */
     '@endian@',
+    /* flags */
     0,
+    /* type_num */
     NPY_@from@,
+    /* elsize */
     0,
+    /* alignment */
     _ALIGN(@align@),
+    /* subarray */
     NULL,
+    /* fields */
     NULL,
+    /* names */
     NULL,
+    /* f */
     &_Py@NAME@_ArrFuncs,
+    /* metadata */
+    NULL,
+    /* c_metadata */
     NULL,
 };
 
@@ -3824,18 +3839,33 @@ static PyArray_ArrFuncs _Py@NAME@_ArrFuncs = {
  */
 NPY_NO_EXPORT PyArray_Descr @from@_Descr = {
     PyObject_HEAD_INIT(&PyArrayDescr_Type)
+    /* typeobj */
     &Py@NAME@ArrType_Type,
+    /* kind */
     NPY_@kind@LTR,
+    /* type */
     NPY_@from@LTR,
+    /* byteorder */
     '@endian@',
+    /* flags */
     @isobject@,
+    /* type_num */
     NPY_@from@,
+    /* elsize */
     @num@*sizeof(@fromtype@),
+    /* alignment */
     _ALIGN(@fromtype@),
+    /* subarray */
     NULL,
+    /* fields */
     NULL,
+    /* names */
     NULL,
+    /* f */
     &_Py@NAME@_ArrFuncs,
+    /* metadata */
+    NULL,
+    /* c_metadata */
     NULL,
 };
 
diff --git a/numpy/core/src/multiarray/descriptor.c b/numpy/core/src/multiarray/descriptor.c
index 724c2eb07..e58b268c3 100644
--- a/numpy/core/src/multiarray/descriptor.c
+++ b/numpy/core/src/multiarray/descriptor.c
@@ -34,38 +34,35 @@ static PyObject *typeDict = NULL;   /* Must be explicitly loaded */
 static PyArray_Descr *
 _use_inherit(PyArray_Descr *type, PyObject *newobj, int *errflag);
 
-NPY_NO_EXPORT PyArray_Descr *
-_arraydescr_fromobj(PyObject *obj)
+/*
+ * Creates a dtype object from ctypes inputs.
+ *
+ * Returns NULL if this is not possible, but does not
+ * set a Python exception.
+ */
+static PyArray_Descr *
+_arraydescr_fromctypes(PyObject *obj)
 {
     PyObject *dtypedescr;
-    PyArray_Descr *new;
+    PyArray_Descr *newdescr;
     int ret;
 
-    dtypedescr = PyObject_GetAttrString(obj, "dtype");
-    PyErr_Clear();
-    if (dtypedescr) {
-        ret = PyArray_DescrConverter(dtypedescr, &new);
-        Py_DECREF(dtypedescr);
-        if (ret == NPY_SUCCEED) {
-            return new;
-        }
-        PyErr_Clear();
-    }
     /* Understand basic ctypes */
     dtypedescr = PyObject_GetAttrString(obj, "_type_");
     PyErr_Clear();
     if (dtypedescr) {
-        ret = PyArray_DescrConverter(dtypedescr, &new);
+        ret = PyArray_DescrConverter(dtypedescr, &newdescr);
         Py_DECREF(dtypedescr);
         if (ret == NPY_SUCCEED) {
             PyObject *length;
+            /* Check for ctypes arrays */
             length = PyObject_GetAttrString(obj, "_length_");
             PyErr_Clear();
             if (length) {
                 /* derived type */
                 PyObject *newtup;
                 PyArray_Descr *derived;
-                newtup = Py_BuildValue("NO", new, length);
+                newtup = Py_BuildValue("NO", newdescr, length);
                 ret = PyArray_DescrConverter(newtup, &derived);
                 Py_DECREF(newtup);
                 if (ret == NPY_SUCCEED) {
@@ -74,7 +71,7 @@ _arraydescr_fromobj(PyObject *obj)
                 PyErr_Clear();
                 return NULL;
             }
-            return new;
+            return newdescr;
         }
         PyErr_Clear();
         return NULL;
@@ -85,16 +82,52 @@ _arraydescr_fromobj(PyObject *obj)
     dtypedescr = PyObject_GetAttrString(obj, "_fields_");
     PyErr_Clear();
     if (dtypedescr) {
-        ret = PyArray_DescrAlignConverter(dtypedescr, &new);
+        ret = PyArray_DescrAlignConverter(dtypedescr, &newdescr);
         Py_DECREF(dtypedescr);
         if (ret == NPY_SUCCEED) {
-            return new;
+            return newdescr;
         }
         PyErr_Clear();
     }
+
     return NULL;
 }
 
+/*
+ * This function creates a dtype object when:
+ *  - The object has a "dtype" attribute, and it can be converted
+ *    to a dtype object.
+ *  - The object is a ctypes type object, including array
+ *    and structure types.
+ *
+ * Returns NULL if this is not possible, but does not
+ * set a Python exception.
+ */
+NPY_NO_EXPORT PyArray_Descr *
+_arraydescr_fromobj(PyObject *obj)
+{
+    PyObject *dtypedescr;
+    PyArray_Descr *newdescr = NULL;
+    int ret;
+
+    /* For arbitrary objects that have a "dtype" attribute */
+    dtypedescr = PyObject_GetAttrString(obj, "dtype");
+    PyErr_Clear();
+    if (dtypedescr != NULL) {
+        ret = PyArray_DescrConverter(dtypedescr, &newdescr);
+        Py_DECREF(dtypedescr);
+        if (ret == NPY_SUCCEED) {
+            return newdescr;
+        }
+        PyErr_Clear();
+    }
+    return _arraydescr_fromctypes(obj);
+}
+
+/*
+ * Sets the global typeDict object, which is a dictionary mapping
+ * dtype names to numpy scalar types.
+ */
 NPY_NO_EXPORT PyObject *
 array_set_typeDict(PyObject *NPY_UNUSED(ignored), PyObject *args)
 {
@@ -140,7 +173,10 @@ _check_for_commastring(char *type, Py_ssize_t len)
                 && type[2] == ')'))) {
         return 1;
     }
-    /* Check for presence of commas outside square [] brackets */
+    /*
+     * Check for presence of commas outside square [] brackets. This
+     * allows commas inside of [], for parameterized dtypes to use.
+     */
     sqbracket = 0;
     for (i = 1; i < len; i++) {
         switch (type[i]) {
@@ -310,9 +346,9 @@ _convert_from_tuple(PyObject *obj)
  * shape parameter).
  *
  * field-name can be a string or a 2-tuple
- * data-type can now be a list, string, or 2-tuple (string, metadata dictionary))
+ * data-type can now be a list, string, or 2-tuple
+ *          (string, metadata dictionary)
  */
-
 static PyArray_Descr *
 _convert_from_array_descr(PyObject *obj, int align)
 {
@@ -834,6 +870,9 @@ _use_fields_dict(PyObject *obj, int align)
     return res;
 }
 
+/*
+ * Creates a struct dtype object from a Python dictionary.
+ */
 static PyArray_Descr *
 _convert_from_dict(PyObject *obj, int align)
 {
@@ -1464,6 +1503,16 @@ PyArray_DescrNew(PyArray_Descr *base)
            (char *)base + sizeof(PyObject),
            sizeof(PyArray_Descr) - sizeof(PyObject));
 
+    /* The c_metadata has a by-value ownership model, need to clone it */
+    if (new->c_metadata != NULL) {
+        new->c_metadata = NPY_AUXDATA_CLONE(new->c_metadata);
+        if (new->c_metadata == NULL) {
+            PyErr_NoMemory();
+            Py_DECREF(new);
+            return NULL;
+        }
+    }
+
     if (new->fields == Py_None) {
         new->fields = NULL;
     }