ENH: Refactor special zero-filling to be managed by the DType

6 files changed, 77 insertions, 48 deletions

diff --git a/numpy/core/src/multiarray/common.c b/numpy/core/src/multiarray/common.c
index da8d23a26..001d299c7 100644
--- a/numpy/core/src/multiarray/common.c
+++ b/numpy/core/src/multiarray/common.c
@@ -128,24 +128,6 @@ PyArray_DTypeFromObject(PyObject *obj, int maxdims, PyArray_Descr **out_dtype)
 }
 
 
-NPY_NO_EXPORT int
-_zerofill(PyArrayObject *ret)
-{
-    if (PyDataType_REFCHK(PyArray_DESCR(ret))) {
-        PyObject *zero = PyLong_FromLong(0);
-        PyArray_FillObjectArray(ret, zero);
-        Py_DECREF(zero);
-        if (PyErr_Occurred()) {
-            return -1;
-        }
-    }
-    else {
-        npy_intp n = PyArray_NBYTES(ret);
-        memset(PyArray_DATA(ret), 0, n);
-    }
-    return 0;
-}
-
 NPY_NO_EXPORT npy_bool
 _IsWriteable(PyArrayObject *ap)
 {
diff --git a/numpy/core/src/multiarray/common.h b/numpy/core/src/multiarray/common.h
index 4e067b22c..127c6250d 100644
--- a/numpy/core/src/multiarray/common.h
+++ b/numpy/core/src/multiarray/common.h
@@ -55,9 +55,6 @@ PyArray_DTypeFromObject(PyObject *obj, int maxdims,
 NPY_NO_EXPORT PyArray_Descr *
 _array_find_python_scalar_type(PyObject *op);
 
-NPY_NO_EXPORT int
-_zerofill(PyArrayObject *ret);
-
 NPY_NO_EXPORT npy_bool
 _IsWriteable(PyArrayObject *ap);
 
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index a6335c783..49296bd4d 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -783,6 +783,10 @@ PyArray_NewFromDescr_int(
     }
 
 
+    /* dtype-specific function to fill array with zero values, may be NULL */
+    fill_initial_function *fill_zero_value =
+        NPY_DT_SLOTS(NPY_DTYPE(descr))->fill_zero_value;
+
     if (data == NULL) {
         /* Store the handler in case the default is modified */
         fa->mem_handler = PyDataMem_GetHandler();
@@ -801,14 +805,29 @@ PyArray_NewFromDescr_int(
                 fa->strides[i] = 0;
             }
         }
-        /*
-         * It is bad to have uninitialized OBJECT pointers
-         * which could also be sub-fields of a VOID array
-         */
-        if (zeroed || PyDataType_FLAGCHK(descr, NPY_NEEDS_INIT)) {
+
+        if (
+            /*
+             * If NPY_NEEDS_INIT is set, always zero out the buffer, even if
+             * zeroed isn't set.
+             */
+            (PyDataType_FLAGCHK(descr, NPY_NEEDS_INIT)) ||
+            /*
+             * If fill_zero_value isn't set, the dtype definitely has no
+             * special zero value, so get a zero-filled buffer using
+             * calloc.  Otherwise, get the buffer with malloc and allow
+             * fill_zero_value to zero out the array with the appropriate
+             * special zero value for the dtype. VOID arrays have
+             * fill_zero_value set but may or may not contain objects, so
+             * always allocate with calloc if zeroed is set.
+             */
+            (zeroed &&
+             ((fill_zero_value == NULL) || (descr->type_num == NPY_VOID)))) {
+            /* allocate array buffer with calloc */
             data = PyDataMem_UserNEW_ZEROED(nbytes, 1, fa->mem_handler);
         }
         else {
+            /* allocate array buffer with malloc */
             data = PyDataMem_UserNEW(nbytes, fa->mem_handler);
         }
         if (data == NULL) {
@@ -846,6 +865,15 @@ PyArray_NewFromDescr_int(
     }
 
     /*
+     * If the array needs special dtype-specific zero-filling logic, do that
+     */
+    if (zeroed && (fill_zero_value != NULL)) {
+        if (fill_zero_value((PyArrayObject *)fa) < 0) {
+            goto fail;
+        }
+    }
+
+    /*
      * call the __array_finalize__ method if a subtype was requested.
      * If obj is NULL use Py_None for the Python callback.
      * For speed, we skip if __array_finalize__ is inherited from ndarray
@@ -3017,17 +3045,7 @@ PyArray_Zeros(int nd, npy_intp const *dims, PyArray_Descr *type, int is_f_order)
         return NULL;
     }
 
-    /* handle objects */
-    if (PyDataType_REFCHK(PyArray_DESCR(ret))) {
-        if (_zerofill(ret) < 0) {
-            Py_DECREF(ret);
-            return NULL;
-        }
-    }
-
-
     return (PyObject *)ret;
-
 }
 
 /*NUMPY_API
diff --git a/numpy/core/src/multiarray/dtypemeta.c b/numpy/core/src/multiarray/dtypemeta.c
index f268ba2cb..b896be3cb 100644
--- a/numpy/core/src/multiarray/dtypemeta.c
+++ b/numpy/core/src/multiarray/dtypemeta.c
@@ -698,6 +698,28 @@ object_common_dtype(
     return cls;
 }
 
+int
+object_fill_zero_value(PyArrayObject *arr)
+{
+    PyObject *zero = PyLong_FromLong(0);
+    PyArray_FillObjectArray(arr, zero);
+    Py_DECREF(zero);
+    if (PyErr_Occurred()) {
+        return -1;
+    }
+    return 0;
+}
+
+int
+void_fill_zero_value(PyArrayObject *arr)
+{
+    if (PyDataType_REFCHK(PyArray_DESCR(arr))) {
+        if (object_fill_zero_value(arr) < 0) {
+            return -1;
+        }
+    }
+    return 0;
+}
 
 /**
  * This function takes a PyArray_Descr and replaces its base class with
@@ -855,6 +877,7 @@ dtypemeta_wrap_legacy_descriptor(PyArray_Descr *descr)
     dt_slots->common_dtype = default_builtin_common_dtype;
     dt_slots->common_instance = NULL;
     dt_slots->ensure_canonical = ensure_native_byteorder;
+    dt_slots->fill_zero_value = NULL;
 
     if (PyTypeNum_ISSIGNED(dtype_class->type_num)) {
         /* Convert our scalars (raise on too large unsigned and NaN, etc.) */
@@ -866,6 +889,7 @@ dtypemeta_wrap_legacy_descriptor(PyArray_Descr *descr)
     }
     else if (descr->type_num == NPY_OBJECT) {
         dt_slots->common_dtype = object_common_dtype;
+        dt_slots->fill_zero_value = object_fill_zero_value;
     }
     else if (PyTypeNum_ISDATETIME(descr->type_num)) {
         /* Datetimes are flexible, but were not considered previously */
@@ -887,6 +911,7 @@ dtypemeta_wrap_legacy_descriptor(PyArray_Descr *descr)
                     void_discover_descr_from_pyobject);
             dt_slots->common_instance = void_common_instance;
             dt_slots->ensure_canonical = void_ensure_canonical;
+            dt_slots->fill_zero_value = void_fill_zero_value;
         }
         else {
             dt_slots->default_descr = string_and_unicode_default_descr;
diff --git a/numpy/core/src/multiarray/dtypemeta.h b/numpy/core/src/multiarray/dtypemeta.h
index 3b4dbad24..897d2a133 100644
--- a/numpy/core/src/multiarray/dtypemeta.h
+++ b/numpy/core/src/multiarray/dtypemeta.h
@@ -42,6 +42,16 @@ typedef struct {
      */
     get_traverse_loop_function *get_clear_loop;
     /*
+       Either NULL or a function that fills an array with zero values
+       appropriate for the dtype. If this function pointer is NULL, the initial
+       value is assumed to be zero. If this function is defined, the array
+       buffer is allocated with malloc and then this function is called to fill
+       the buffer. As a performance optimization, the array buffer is allocated
+       using calloc if this function is NULL, so it is best to avoid using this
+       function if zero-filling the array buffer makes sense for the dtype.
+    */
+    fill_initial_function *fill_zero_value;
+    /*
      * The casting implementation (ArrayMethod) to convert between two
      * instances of this DType, stored explicitly for fast access:
      */
@@ -63,7 +73,7 @@ typedef struct {
 
 // This must be updated if new slots before within_dtype_castingimpl
 // are added
-#define NPY_NUM_DTYPE_SLOTS 9
+#define NPY_NUM_DTYPE_SLOTS 10
 #define NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS 22
 #define NPY_DT_MAX_ARRFUNCS_SLOT \
   NPY_NUM_DTYPE_PYARRAY_ARRFUNCS_SLOTS + _NPY_DT_ARRFUNCS_OFFSET
diff --git a/numpy/core/src/multiarray/getset.c b/numpy/core/src/multiarray/getset.c
index ab35548ed..d019acbb5 100644
--- a/numpy/core/src/multiarray/getset.c
+++ b/numpy/core/src/multiarray/getset.c
@@ -797,20 +797,17 @@ array_imag_get(PyArrayObject *self, void *NPY_UNUSED(ignored))
     }
     else {
         Py_INCREF(PyArray_DESCR(self));
-        ret = (PyArrayObject *)PyArray_NewFromDescr(Py_TYPE(self),
-                                                    PyArray_DESCR(self),
-                                                    PyArray_NDIM(self),
-                                                    PyArray_DIMS(self),
-                                                    NULL, NULL,
-                                                    PyArray_ISFORTRAN(self),
-                                                    (PyObject *)self);
+        ret = (PyArrayObject *)PyArray_NewFromDescr_int(
+                Py_TYPE(self),
+                PyArray_DESCR(self),
+                PyArray_NDIM(self),
+                PyArray_DIMS(self),
+                NULL, NULL,
+                PyArray_ISFORTRAN(self),
+                (PyObject *)self, NULL, 1, 0);
         if (ret == NULL) {
             return NULL;
         }
-        if (_zerofill(ret) < 0) {
-            Py_DECREF(ret);
-            return NULL;
-        }
         PyArray_CLEARFLAGS(ret, NPY_ARRAY_WRITEABLE);
     }
     return (PyObject *) ret;