BUG: fix StringDType distinct-allocator bugs and add tests (numpy#31609)

Author: ngoldbaum

numpy/_core/_internal.py

@@ -13,6 +13,7 @@
 from numpy import _NoValue
 from numpy.exceptions import DTypePromotionError
 
+from ._multiarray_umath import _is_view_safe_cast
 from .multiarray import StringDType, array, dtype, promote_types
 
 try:
@@ -484,7 +485,7 @@ def _getfield_is_safe(oldtype, newtype, offset):
 
     """
     if newtype.hasobject or oldtype.hasobject:
-        if offset == 0 and newtype == oldtype:
+        if offset == 0 and _is_view_safe_cast(oldtype, newtype):
             return
         if oldtype.names is not None:
             for name in oldtype.names:
@@ -514,9 +515,10 @@ def _view_is_safe(oldtype, newtype):
 
     """
 
-    # if the types are equivalent, there is no problem.
-    # for example: dtype((np.record, 'i4,i4')) == dtype((np.void, 'i4,i4'))
-    if oldtype == newtype:
+    # more precise than ``oldtype == newtype``: e.g. dtype((np.record, 'i4,i4'))
+    # views safely as dtype((np.void, 'i4,i4')), while two equal StringDType
+    # instances with separate allocators do not
+    if _is_view_safe_cast(oldtype, newtype):
         return
 
     if newtype.hasobject or oldtype.hasobject:

numpy/_core/src/multiarray/compiled_base.c

@@ -16,6 +16,7 @@
 #include "ctors.h"
 #include "common.h"
 #include "dtypemeta.h"
+#include "dtype_transfer.h"
 #include "simd/simd.h"
 
 #include <string.h>
@@ -393,18 +394,58 @@ arr_place(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
     j = 0;
 
     copyswap = PyDataType_GetArrFuncs(PyArray_DESCR(array))->copyswap;
-    NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(array));
-    for (i = 0; i < ni; i++) {
-        if (mask_data[i]) {
-            if (j >= nv) {
-                j = 0;
+    if (copyswap == NULL) {
+        NPY_cast_info cast_info;
+        NPY_ARRAYMETHOD_FLAGS flags;
+        const npy_intp one = 1;
+        const npy_intp elsize = chunk;
+        const npy_intp strides[2] = {elsize, elsize};
+
+        NPY_cast_info_init(&cast_info);
+        if (PyArray_GetDTypeTransferFunction(
+                PyArray_ISALIGNED(values) && PyArray_ISALIGNED(array),
+                strides[0], strides[1],
+                PyArray_DESCR(values), PyArray_DESCR(array), 0,
+                &cast_info, &flags) < 0) {
+            goto fail;
+        }
+        if (!(flags & NPY_METH_REQUIRES_PYAPI)) {
+            NPY_BEGIN_THREADS;
+        }
+        for (i = 0; i < ni; i++) {
+            if (mask_data[i]) {
+                if (j >= nv) {
+                    j = 0;
+                }
+
+                char *data[2] = {src + j*chunk, dest + i*chunk};
+                if (cast_info.func(
+                        &cast_info.context, data, &one, strides,
+                        cast_info.auxdata) < 0) {
+                    NPY_END_THREADS;
+                    NPY_cast_info_xfree(&cast_info);
+                    goto fail;
+                }
+                j++;
             }
+        }
+        NPY_END_THREADS;
+        NPY_cast_info_xfree(&cast_info);
+    }
+    else {
+        NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(array));
+        for (i = 0; i < ni; i++) {
+            if (mask_data[i]) {
+                if (j >= nv) {
+                    j = 0;
+                }
 
-            copyswap(dest + i*chunk, src + j*chunk, 0, array);
-            j++;
+                copyswap(dest + i*chunk, src + j*chunk, 0, array);
+                j++;
+            }
         }
+        NPY_END_THREADS;
     }
-    NPY_END_THREADS;
 
     Py_XDECREF(values);
     Py_XDECREF(mask);

numpy/_core/src/multiarray/convert_datatype.c

@@ -648,6 +648,36 @@ PyArray_SafeCast(PyArray_Descr *type1, PyArray_Descr *type2,
 }
 
 
+/*
+ * Python-level helper answering whether reinterpreting data of dtype *from*
+ * as dtype *to* is a view. Unlike descriptor equality or `np.can_cast` with
+ * "no" casting, this distinguishes equivalent descriptor instances with
+ * separate internal state (e.g. StringDType allocators).
+ */
+NPY_NO_EXPORT PyObject *
+_is_view_safe_cast(PyObject *NPY_UNUSED(module), PyObject *const *args,
+                   Py_ssize_t len_args)
+{
+    if (len_args != 2) {
+        PyErr_SetString(PyExc_TypeError,
+                "_is_view_safe_cast() takes exactly two arguments");
+        return NULL;
+    }
+    if (!PyArray_DescrCheck(args[0]) || !PyArray_DescrCheck(args[1])) {
+        PyErr_SetString(PyExc_TypeError,
+                "_is_view_safe_cast() arguments must be dtype instances");
+        return NULL;
+    }
+    PyArray_Descr *from = (PyArray_Descr *)args[0];
+    PyArray_Descr *to = (PyArray_Descr *)args[1];
+    npy_intp view_offset = NPY_MIN_INTP;
+    /* ignore_error=1: dtype pairs with no resolvable cast are simply not views */
+    npy_intp is_safe = PyArray_SafeCast(from, to, &view_offset,
+                                        NPY_NO_CASTING, 1);
+    return PyBool_FromLong(is_safe && view_offset == 0);
+}
+
+
 /* Provides an ordering for the dtype 'kind' character codes */
 NPY_NO_EXPORT int
 dtype_kind_to_ordering(char kind)

numpy/_core/src/multiarray/convert_datatype.h

@@ -15,6 +15,10 @@ PyArray_GetCastingImpl(PyArray_DTypeMeta *from, PyArray_DTypeMeta *to);
 NPY_NO_EXPORT PyObject *
 _get_castingimpl(PyObject *NPY_UNUSED(module), PyObject *args);
 
+NPY_NO_EXPORT PyObject *
+_is_view_safe_cast(PyObject *NPY_UNUSED(module), PyObject *const *args,
+                   Py_ssize_t len_args);
+
 NPY_NO_EXPORT PyArray_VectorUnaryFunc *
 PyArray_GetCastFunc(PyArray_Descr *descr, int type_num);
 

numpy/_core/src/multiarray/getset.c

@@ -12,10 +12,13 @@
 
 #include "npy_import.h"
 
+#include "array_assign.h"
 #include "common.h"
 #include "conversion_utils.h"
 #include "ctors.h"
+#include "dtype_transfer.h"
 #include "dtypemeta.h"
+#include "lowlevel_strided_loops.h"
 #include "scalartypes.h"
 #include "descriptor.h"
 #include "flagsobject.h"
@@ -671,26 +674,43 @@ array_flat_set(PyArrayObject *self, PyObject *val, void *NPY_UNUSED(ignored))
         retval = 0;
         goto exit;
     }
-    swap = PyArray_ISNOTSWAPPED(self) != PyArray_ISNOTSWAPPED(arr);
     copyswap = PyDataType_GetArrFuncs(PyArray_DESCR(self))->copyswap;
-    if (PyDataType_REFCHK(PyArray_DESCR(self))) {
+    if (copyswap == NULL || PyDataType_REFCHK(PyArray_DESCR(self))) {
+        /* reference dtypes have copyswap, but the transfer path handles
+           refcounts and is better for structured dtypes */
+        NPY_cast_info cast_info;
+        NPY_ARRAYMETHOD_FLAGS transfer_flags = 0;
+        npy_intp one = 1;
+        npy_intp itemsize = PyArray_ITEMSIZE(self);
+        npy_intp transfer_strides[2] = {itemsize, itemsize};
+
+        NPY_cast_info_init(&cast_info);
+        if (PyArray_GetDTypeTransferFunction(
+                IsUintAligned(self) && IsUintAligned(arr),
+                itemsize, itemsize,
+                PyArray_DESCR(arr), PyArray_DESCR(self), 0,
+                &cast_info, &transfer_flags) < 0) {
+            goto exit;
+        }
         while (selfit->index < selfit->size) {
-            PyArray_Item_XDECREF(selfit->dataptr, PyArray_DESCR(self));
-            PyArray_Item_INCREF(arrit->dataptr, PyArray_DESCR(arr));
-            memmove(selfit->dataptr, arrit->dataptr, sizeof(PyObject **));
-            if (swap) {
-                copyswap(selfit->dataptr, NULL, swap, self);
+            char *args[2] = {arrit->dataptr, selfit->dataptr};
+            if (cast_info.func(&cast_info.context, args, &one,
+                               transfer_strides, cast_info.auxdata) < 0) {
+                NPY_cast_info_xfree(&cast_info);
+                goto exit;
             }
             PyArray_ITER_NEXT(selfit);
             PyArray_ITER_NEXT(arrit);
             if (arrit->index == arrit->size) {
                 PyArray_ITER_RESET(arrit);
             }
         }
+        NPY_cast_info_xfree(&cast_info);
         retval = 0;
         goto exit;
     }
 
+    swap = PyArray_ISNOTSWAPPED(self) != PyArray_ISNOTSWAPPED(arr);
     while(selfit->index < selfit->size) {
         copyswap(selfit->dataptr, arrit->dataptr, swap, self);
         PyArray_ITER_NEXT(selfit);

numpy/_core/src/multiarray/item_selection.c

@@ -457,9 +457,10 @@ PyArray_PutTo(PyArrayObject *self, PyObject* values0, PyObject *indices0,
         NPY_BEGIN_THREADS_THRESHOLDED(ni);
     }
     else {
-        PyArray_Descr *dtype = PyArray_DESCR(self);
         if (PyArray_GetDTypeTransferFunction(
-                PyArray_ISALIGNED(self), itemsize, itemsize, dtype, dtype, 0,
+                PyArray_ISALIGNED(self) && PyArray_ISALIGNED(values),
+                itemsize, itemsize,
+                PyArray_DESCR(values), PyArray_DESCR(self), 0,
                 &cast_info, &flags) < 0) {
             goto fail;
         }
@@ -755,7 +756,9 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
 
         NPY_cast_info_init(&cast_info);
         if (PyArray_GetDTypeTransferFunction(
-                PyArray_ISALIGNED(self), itemsize, itemsize, dtype, dtype, 0,
+                PyArray_ISALIGNED(self) && PyArray_ISALIGNED(values),
+                itemsize, itemsize,
+                PyArray_DESCR(values), PyArray_DESCR(self), 0,
                 &cast_info, &flags) < 0) {
             goto fail;
         }
@@ -778,15 +781,15 @@ PyArray_PutMask(PyArrayObject *self, PyObject* values0, PyObject* mask0)
                 }
             }
         }
+        NPY_END_THREADS;
         NPY_cast_info_xfree(&cast_info);
     }
     else {
         NPY_BEGIN_THREADS;
         npy_fastputmask(dest, src, mask_data, ni, nv, itemsize);
+        NPY_END_THREADS;
     }
 
-    NPY_END_THREADS;
-
     Py_XDECREF(values);
     Py_XDECREF(mask);
     if (copied) {
@@ -1024,7 +1027,10 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
     PyArrayObject **mps, *ap;
     PyArrayMultiIterObject *multi = NULL;
     npy_intp mi;
-    NPY_cast_info cast_info = {.func = NULL};
+    /* PyArray_MultiIterFromObjects below bounds n by NPY_MAXARGS */
+    NPY_cast_info cast_infos[NPY_MAXARGS];
+    int needs_transfer = 0;
+    NPY_BEGIN_THREADS_DEF;
     ap = NULL;
 
     /*
@@ -1118,23 +1124,35 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
     npy_intp transfer_strides[2] = {elsize, elsize};
     npy_intp one = 1;
     NPY_ARRAYMETHOD_FLAGS transfer_flags = 0;
-    if (PyDataType_REFCHK(dtype)) {
-        int is_aligned = IsUintAligned(obj);
+    needs_transfer = PyDataType_REFCHK(dtype);
+    if (needs_transfer) {
         PyArray_Descr *obj_dtype = PyArray_DESCR(obj);
-        PyArray_GetDTypeTransferFunction(
-                    is_aligned,
-                    dtype->elsize,
-                    obj_dtype->elsize,
-                    dtype,
-                    obj_dtype, 0, &cast_info,
-                    &transfer_flags);
+        for (i = 0; i < n; i++) {
+            NPY_cast_info_init(&cast_infos[i]);
+        }
+        for (i = 0; i < n; i++) {
+            int is_aligned = IsUintAligned(obj) && IsUintAligned(mps[i]);
+            if (PyArray_GetDTypeTransferFunction(
+                        is_aligned,
+                        PyArray_DESCR(mps[i])->elsize,
+                        obj_dtype->elsize,
+                        PyArray_DESCR(mps[i]),
+                        obj_dtype, 0, &cast_infos[i],
+                        &transfer_flags) < 0) {
+                goto fail;
+            }
+        }
     }
 
+    if (!(transfer_flags & NPY_METH_REQUIRES_PYAPI)) {
+        NPY_BEGIN_THREADS_THRESHOLDED(multi->size);
+    }
     while (PyArray_MultiIter_NOTDONE(multi)) {
         mi = *((npy_intp *)PyArray_MultiIter_DATA(multi, n));
         if (mi < 0 || mi >= n) {
             switch(clipmode) {
             case NPY_RAISE:
+                NPY_END_THREADS;
                 PyErr_SetString(PyExc_ValueError,
                         "invalid entry in choice "\
                         "array");
@@ -1161,22 +1179,28 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
                 break;
             }
         }
-        if (cast_info.func == NULL) {
+        if (!needs_transfer) {
             /* We ensure memory doesn't overlap, so can use memcpy */
             memcpy(ret_data, PyArray_MultiIter_DATA(multi, mi), elsize);
         }
         else {
             char *args[2] = {PyArray_MultiIter_DATA(multi, mi), ret_data};
-            if (cast_info.func(&cast_info.context, args, &one,
-                                transfer_strides, cast_info.auxdata) < 0) {
+            if (cast_infos[mi].func(&cast_infos[mi].context, args, &one,
+                                    transfer_strides,
+                                    cast_infos[mi].auxdata) < 0) {
                 goto fail;
             }
         }
         ret_data += elsize;
         PyArray_MultiIter_NEXT(multi);
     }
+    NPY_END_THREADS;
 
-    NPY_cast_info_xfree(&cast_info);
+    if (needs_transfer) {
+        for (i = 0; i < n; i++) {
+            NPY_cast_info_xfree(&cast_infos[i]);
+        }
+    }
     Py_DECREF(multi);
     for (i = 0; i < n; i++) {
         Py_XDECREF(mps[i]);
@@ -1194,7 +1218,12 @@ PyArray_Choose(PyArrayObject *ip, PyObject *op, PyArrayObject *out,
     return (PyObject *)obj;
 
  fail:
-    NPY_cast_info_xfree(&cast_info);
+    NPY_END_THREADS;
+    if (needs_transfer) {
+        for (i = 0; i < n; i++) {
+            NPY_cast_info_xfree(&cast_infos[i]);
+        }
+    }
     Py_XDECREF(multi);
     for (i = 0; i < n; i++) {
         Py_XDECREF(mps[i]);