ENH: Improve performance for np.result_type (#28710)

eendebakpt · web-flow · commit 9ead596ce4f8 · 2025-04-23T09:55:22.000+02:00
Improve performance of  np.result_type

* Use NPY_ALLOC_WORKSPACE in array_result_type
* Avoid increfs/decrefs in PyArray_ResultType
* Fast path for identical arguments to np.result_type
diff --git a/numpy/_core/src/multiarray/convert_datatype.c b/numpy/_core/src/multiarray/convert_datatype.c
@@ -1597,13 +1597,12 @@ PyArray_ResultType(
         return NULL;
     }
 
-    PyArray_DTypeMeta **all_DTypes = (PyArray_DTypeMeta **)workspace;
+    PyArray_DTypeMeta **all_DTypes = (PyArray_DTypeMeta **)workspace; // borrowed references
     PyArray_Descr **all_descriptors = (PyArray_Descr **)(&all_DTypes[narrs+ndtypes]);
 
     /* Copy all dtypes into a single array defining non-value-based behaviour */
     for (npy_intp i=0; i < ndtypes; i++) {
         all_DTypes[i] = NPY_DTYPE(descrs[i]);
-        Py_INCREF(all_DTypes[i]);
         all_descriptors[i] = descrs[i];
     }
 
@@ -1628,14 +1627,10 @@ PyArray_ResultType(
             all_descriptors[i_all] = PyArray_DTYPE(arrs[i]);
             all_DTypes[i_all] = NPY_DTYPE(all_descriptors[i_all]);
         }
-        Py_INCREF(all_DTypes[i_all]);
     }
 
     PyArray_DTypeMeta *common_dtype = PyArray_PromoteDTypeSequence(
             narrs+ndtypes, all_DTypes);
-    for (npy_intp i=0; i < narrs+ndtypes; i++) {
-        Py_DECREF(all_DTypes[i]);
-    }
     if (common_dtype == NULL) {
         goto error;
     }
diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c
@@ -3583,24 +3583,28 @@ static PyObject *
 array_result_type(PyObject *NPY_UNUSED(dummy), PyObject *const *args, Py_ssize_t len)
 {
     npy_intp i, narr = 0, ndtypes = 0;
-    PyArrayObject **arr = NULL;
-    PyArray_Descr **dtypes = NULL;
     PyObject *ret = NULL;
 
     if (len == 0) {
         PyErr_SetString(PyExc_ValueError,
                         "at least one array or dtype is required");
-        goto finish;
+        return NULL;
     }
 
-    arr = PyArray_malloc(2 * len * sizeof(void *));
+    NPY_ALLOC_WORKSPACE(arr, PyArrayObject *, 2 * 3, 2 * len);
     if (arr == NULL) {
-        return PyErr_NoMemory();
+        return NULL;
     }
-    dtypes = (PyArray_Descr**)&arr[len];
+    PyArray_Descr **dtypes = (PyArray_Descr**)&arr[len];
+
+    PyObject *previous_obj = NULL;
 
     for (i = 0; i < len; ++i) {
         PyObject *obj = args[i];
+        if (obj == previous_obj) {
+            continue;
+        }
+
         if (PyArray_Check(obj)) {
             Py_INCREF(obj);
             arr[narr] = (PyArrayObject *)obj;
@@ -3636,7 +3640,7 @@ array_result_type(PyObject *NPY_UNUSED(dummy), PyObject *const *args, Py_ssize_t
     for (i = 0; i < ndtypes; ++i) {
         Py_DECREF(dtypes[i]);
     }
-    PyArray_free(arr);
+    npy_free_workspace(arr);
     return ret;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -1597,13 +1597,12 @@ PyArray_ResultType(`
`1597`	`1597`	`return NULL;`
`1598`	`1598`	`}`
`1599`	`1599`
`1600`		`- PyArray_DTypeMeta all_DTypes = (PyArray_DTypeMeta )workspace;`
	`1600`	`+ PyArray_DTypeMeta all_DTypes = (PyArray_DTypeMeta )workspace; // borrowed references`
`1601`	`1601`	`PyArray_Descr all_descriptors = (PyArray_Descr )(&all_DTypes[narrs+ndtypes]);`
`1602`	`1602`
`1603`	`1603`	`/* Copy all dtypes into a single array defining non-value-based behaviour */`
`1604`	`1604`	`for (npy_intp i=0; i < ndtypes; i++) {`
`1605`	`1605`	`all_DTypes[i] = NPY_DTYPE(descrs[i]);`
`1606`		`- Py_INCREF(all_DTypes[i]);`
`1607`	`1606`	`all_descriptors[i] = descrs[i];`
`1608`	`1607`	`}`
`1609`	`1608`
`@@ -1628,14 +1627,10 @@ PyArray_ResultType(`
`1628`	`1627`	`all_descriptors[i_all] = PyArray_DTYPE(arrs[i]);`
`1629`	`1628`	`all_DTypes[i_all] = NPY_DTYPE(all_descriptors[i_all]);`
`1630`	`1629`	`}`
`1631`		`- Py_INCREF(all_DTypes[i_all]);`
`1632`	`1630`	`}`
`1633`	`1631`
`1634`	`1632`	`PyArray_DTypeMeta *common_dtype = PyArray_PromoteDTypeSequence(`
`1635`	`1633`	`narrs+ndtypes, all_DTypes);`
`1636`		`- for (npy_intp i=0; i < narrs+ndtypes; i++) {`
`1637`		`- Py_DECREF(all_DTypes[i]);`
`1638`		`- }`
`1639`	`1634`	`if (common_dtype == NULL) {`
`1640`	`1635`	`goto error;`
`1641`	`1636`	`}`