ENH: support no-copy pickling for any array that can be transposed to a C-contiguous array (numpy#28105)

* add support for transposed contiguous array and corresponding test

* fix lint

* rm bad includes

* fix

* fix

* discard vla

* accept legacy pkl

* check return

* fix ci

* check after transpose

* discard auto include

* use N format in pybuildvalue

* use pytuple pack

* clean up

* add comment on fall back behavior

* add mroe test

* use byte str instead of pkl to test

* get shape with pyarray api

* fmt

* better comment

* fix as suggested

* fix lint

* discard auto include

* shorter tested strs

* fmt

* last correct commit

* last correct commit

* [skip ci] memory leak

* add cleanup

* [skip ci] add release doc

* [skip ci] add check for pylong_fromlong

* typo

* fix memory leak

* more clean up

* use same random seed

* use xdecref

* check PyArray_Transpose

* check PyArray_IntTupleFromIntp

* fix

* once-over

* add additional test

* once over again

* fmt

* pickle format consistency

* fix incorrect usage of cpy api

* fmt

Author: IndifferentArea

doc/release/upcoming_changes/28105.improvement.rst

@@ -0,0 +1,2 @@
+* No-copy pickling is now supported for any 
+  array that can be transposed to a C-contiguous array.

numpy/_core/numeric.py

@@ -1920,7 +1920,9 @@ def fromfunction(function, shape, *, dtype=float, like=None, **kwargs):
 _fromfunction_with_like = array_function_dispatch()(fromfunction)
 
 
-def _frombuffer(buf, dtype, shape, order):
+def _frombuffer(buf, dtype, shape, order, axis_order=None):
+    if order == 'K' and axis_order is not None:
+        return frombuffer(buf, dtype=dtype).reshape(shape, order='C').transpose(axis_order)
     return frombuffer(buf, dtype=dtype).reshape(shape, order=order)
 
 

numpy/_core/src/multiarray/methods.c

@@ -1845,77 +1845,116 @@ array_reduce_ex_regular(PyArrayObject *self, int NPY_UNUSED(protocol))
 static PyObject *
 array_reduce_ex_picklebuffer(PyArrayObject *self, int protocol)
 {
-    PyObject *numeric_mod = NULL, *from_buffer_func = NULL;
-    PyObject *pickle_module = NULL, *picklebuf_class = NULL;
-    PyObject *picklebuf_args = NULL;
+    PyObject *from_buffer_func = NULL;
+    PyObject *picklebuf_class = NULL;
     PyObject *buffer = NULL, *transposed_array = NULL;
     PyArray_Descr *descr = NULL;
+    PyObject *rev_perm = NULL;  // only used in 'K' order
     char order;
 
     descr = PyArray_DESCR(self);
 
     /* we expect protocol 5 to be available in Python 3.8 */
-    pickle_module = PyImport_ImportModule("pickle");
-    if (pickle_module == NULL){
-        return NULL;
-    }
-    picklebuf_class = PyObject_GetAttrString(pickle_module, "PickleBuffer");
-    Py_DECREF(pickle_module);
-    if (picklebuf_class == NULL) {
+    if (npy_cache_import_runtime("pickle", "PickleBuffer", &picklebuf_class) == -1) {
         return NULL;
     }
 
     /* Construct a PickleBuffer of the array */
-
-    if (!PyArray_IS_C_CONTIGUOUS((PyArrayObject*) self) &&
-         PyArray_IS_F_CONTIGUOUS((PyArrayObject*) self)) {
+    if (PyArray_IS_C_CONTIGUOUS((PyArrayObject *)self)) {
+        order = 'C';
+    }
+    else if (PyArray_IS_F_CONTIGUOUS((PyArrayObject *)self)) {
         /* if the array if Fortran-contiguous and not C-contiguous,
          * the PickleBuffer instance will hold a view on the transpose
          * of the initial array, that is C-contiguous. */
         order = 'F';
-        transposed_array = PyArray_Transpose((PyArrayObject*)self, NULL);
-        picklebuf_args = Py_BuildValue("(N)", transposed_array);
+        transposed_array = PyArray_Transpose((PyArrayObject *)self, NULL);
+        if (transposed_array == NULL) {
+            return NULL;
+        }
     }
     else {
-        order = 'C';
-        picklebuf_args = Py_BuildValue("(O)", self);
-    }
-    if (picklebuf_args == NULL) {
-        Py_DECREF(picklebuf_class);
-        return NULL;
+        order = 'K';
+        const int n = PyArray_NDIM(self);
+        npy_stride_sort_item items[NPY_MAXDIMS];
+        // sort (strde, perm) as descending = transpose to C
+        PyArray_CreateSortedStridePerm(n, PyArray_STRIDES(self), items);
+        rev_perm = PyTuple_New(n);
+        if (rev_perm == NULL) {
+            return NULL;
+        }
+        PyArray_Dims perm;
+        npy_intp dims[NPY_MAXDIMS];
+        for (int i = 0; i < n; i++) {
+            dims[i] = items[i].perm;
+            PyObject *idx = PyLong_FromLong(i);
+            if (idx == NULL) {
+                Py_DECREF(rev_perm);
+                return NULL;
+            }
+            PyTuple_SET_ITEM(rev_perm, items[i].perm, idx);
+        }
+        perm.ptr = dims;
+        perm.len = n;
+        transposed_array = PyArray_Transpose((PyArrayObject *)self, &perm);
+        if (transposed_array == NULL) {
+            Py_DECREF(rev_perm);
+            return NULL;
+        }
+        if (!PyArray_IS_C_CONTIGUOUS((PyArrayObject *)transposed_array)) {
+            // self is non-contiguous
+            Py_DECREF(rev_perm);
+            Py_DECREF(transposed_array);
+            return array_reduce_ex_regular(self, protocol);
+        }
     }
-
-    buffer = PyObject_CallObject(picklebuf_class, picklebuf_args);
-    Py_DECREF(picklebuf_class);
-    Py_DECREF(picklebuf_args);
+    buffer = PyObject_CallOneArg(picklebuf_class, transposed_array == NULL ? (PyObject*) self: transposed_array);
     if (buffer == NULL) {
         /* Some arrays may refuse to export a buffer, in which case
          * just fall back on regular __reduce_ex__ implementation
          * (gh-12745).
          */
+        Py_XDECREF(rev_perm);
+        Py_XDECREF(transposed_array);
         PyErr_Clear();
         return array_reduce_ex_regular(self, protocol);
     }
 
     /* Get the _frombuffer() function for reconstruction */
-
-    numeric_mod = PyImport_ImportModule("numpy._core.numeric");
-    if (numeric_mod == NULL) {
+    if (npy_cache_import_runtime("numpy._core.numeric", "_frombuffer",
+                                 &from_buffer_func) == -1) {
+        Py_XDECREF(rev_perm);
+        Py_XDECREF(transposed_array);
         Py_DECREF(buffer);
         return NULL;
     }
-    from_buffer_func = PyObject_GetAttrString(numeric_mod,
-                                              "_frombuffer");
-    Py_DECREF(numeric_mod);
-    if (from_buffer_func == NULL) {
+
+    PyObject *shape = NULL;
+    if (order == 'K') {
+        shape = PyArray_IntTupleFromIntp(
+                PyArray_NDIM((PyArrayObject *)transposed_array),
+                PyArray_SHAPE((PyArrayObject *)transposed_array));
+    }
+    else {
+        shape = PyArray_IntTupleFromIntp(PyArray_NDIM(self),
+                                         PyArray_SHAPE(self));
+    }
+    Py_XDECREF(transposed_array);
+    if (shape == NULL) {
+        Py_XDECREF(rev_perm);
         Py_DECREF(buffer);
         return NULL;
     }
-
-    return Py_BuildValue("N(NONN)",
-                         from_buffer_func, buffer, (PyObject *)descr,
-                         PyObject_GetAttrString((PyObject *)self, "shape"),
-                         PyUnicode_FromStringAndSize(&order, 1));
+    if (order == 'K') {
+        return Py_BuildValue("N(NONNN)", from_buffer_func, buffer,
+                             (PyObject *)descr, shape,
+                             PyUnicode_FromStringAndSize(&order, 1), rev_perm);
+    }
+    else {
+        return Py_BuildValue("N(NONN)", from_buffer_func, buffer,
+                             (PyObject *)descr, shape,
+                             PyUnicode_FromStringAndSize(&order, 1));
+    }
 }
 
 static PyObject *
@@ -1930,8 +1969,6 @@ array_reduce_ex(PyArrayObject *self, PyObject *args)
 
     descr = PyArray_DESCR(self);
     if ((protocol < 5) ||
-        (!PyArray_IS_C_CONTIGUOUS((PyArrayObject*)self) &&
-         !PyArray_IS_F_CONTIGUOUS((PyArrayObject*)self)) ||
         PyDataType_FLAGCHK(descr, NPY_ITEM_HASOBJECT) ||
         (PyType_IsSubtype(((PyObject*)self)->ob_type, &PyArray_Type) &&
          ((PyObject*)self)->ob_type != &PyArray_Type) ||
@@ -1943,6 +1980,11 @@ array_reduce_ex(PyArrayObject *self, PyObject *args)
         return array_reduce_ex_regular(self, protocol);
     }
     else {
+        /* The func will check internally
+         * if the array isn't backed by a contiguous data buffer or
+         * if the array refuses to export a buffer
+         * In either case, fall back to `array_reduce_ex_regular`
+         */
         return array_reduce_ex_picklebuffer(self, protocol);
     }
 }

numpy/_core/tests/test_multiarray.py

@@ -4404,17 +4404,55 @@ def test_f_contiguous_array(self):
 
         assert_equal(f_contiguous_array, depickled_f_contiguous_array)
 
+    @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL < 5, reason="requires pickle protocol 5")
+    @pytest.mark.parametrize('transposed_contiguous_array', 
+        [np.random.default_rng(42).random((2, 3, 4)).transpose((1, 0, 2)), 
+         np.random.default_rng(42).random((2, 3, 4, 5)).transpose((1, 3, 0, 2))] + 
+        [np.random.default_rng(42).random(np.arange(2, 7)).transpose(np.random.permutation(5)) for _ in range(3)])
+    def test_transposed_contiguous_array(self, transposed_contiguous_array):
+        buffers = []
+        # When using pickle protocol 5, arrays which can be transposed to c_contiguous
+        # can be serialized using out-of-band buffers
+        bytes_string = pickle.dumps(transposed_contiguous_array, protocol=5,
+                                    buffer_callback=buffers.append)
+
+        assert len(buffers) > 0
+
+        depickled_transposed_contiguous_array = pickle.loads(bytes_string,
+                                                    buffers=buffers)
+
+        assert_equal(transposed_contiguous_array, depickled_transposed_contiguous_array)
+
+    @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL < 5, reason="requires pickle protocol 5")
+    def test_load_legacy_pkl_protocol5(self):
+        # legacy byte strs are dumped in 2.2.1
+        c_contiguous_dumped = b'\x80\x05\x95\x90\x00\x00\x00\x00\x00\x00\x00\x8c\x13numpy._core.numeric\x94\x8c\x0b_frombuffer\x94\x93\x94(\x96\x18\x00\x00\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x94\x8c\x05numpy\x94\x8c\x05dtype\x94\x93\x94\x8c\x02u1\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94bK\x03K\x04K\x02\x87\x94\x8c\x01C\x94t\x94R\x94.' # noqa
+        f_contiguous_dumped = b'\x80\x05\x95\x90\x00\x00\x00\x00\x00\x00\x00\x8c\x13numpy._core.numeric\x94\x8c\x0b_frombuffer\x94\x93\x94(\x96\x18\x00\x00\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x94\x8c\x05numpy\x94\x8c\x05dtype\x94\x93\x94\x8c\x02u1\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94bK\x03K\x04K\x02\x87\x94\x8c\x01F\x94t\x94R\x94.' # noqa
+        transposed_contiguous_dumped = b'\x80\x05\x95\xa5\x00\x00\x00\x00\x00\x00\x00\x8c\x16numpy._core.multiarray\x94\x8c\x0c_reconstruct\x94\x93\x94\x8c\x05numpy\x94\x8c\x07ndarray\x94\x93\x94K\x00\x85\x94C\x01b\x94\x87\x94R\x94(K\x01K\x04K\x03K\x02\x87\x94h\x03\x8c\x05dtype\x94\x93\x94\x8c\x02u1\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94b\x89C\x18\x00\x01\x08\t\x10\x11\x02\x03\n\x0b\x12\x13\x04\x05\x0c\r\x14\x15\x06\x07\x0e\x0f\x16\x17\x94t\x94b.' # noqa
+        no_contiguous_dumped = b'\x80\x05\x95\x91\x00\x00\x00\x00\x00\x00\x00\x8c\x16numpy._core.multiarray\x94\x8c\x0c_reconstruct\x94\x93\x94\x8c\x05numpy\x94\x8c\x07ndarray\x94\x93\x94K\x00\x85\x94C\x01b\x94\x87\x94R\x94(K\x01K\x03K\x02\x86\x94h\x03\x8c\x05dtype\x94\x93\x94\x8c\x02u1\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94b\x89C\x06\x00\x01\x04\x05\x08\t\x94t\x94b.' # noqa
+        x = np.arange(24, dtype='uint8').reshape(3, 4, 2)
+        assert_equal(x, pickle.loads(c_contiguous_dumped))
+        x = np.arange(24, dtype='uint8').reshape(3, 4, 2, order='F')
+        assert_equal(x, pickle.loads(f_contiguous_dumped))
+        x = np.arange(24, dtype='uint8').reshape(3, 4, 2).transpose((1, 0, 2))
+        assert_equal(x, pickle.loads(transposed_contiguous_dumped))
+        x = np.arange(12, dtype='uint8').reshape(3, 4)[:, :2]
+        assert_equal(x, pickle.loads(no_contiguous_dumped))
+
     def test_non_contiguous_array(self):
         non_contiguous_array = np.arange(12).reshape(3, 4)[:, :2]
         assert not non_contiguous_array.flags.c_contiguous
         assert not non_contiguous_array.flags.f_contiguous
 
         # make sure non-contiguous arrays can be pickled-depickled
         # using any protocol
+        buffers = []
         for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
             depickled_non_contiguous_array = pickle.loads(
-                    pickle.dumps(non_contiguous_array, protocol=proto))
+                    pickle.dumps(non_contiguous_array, protocol=proto, 
+                                 buffer_callback=buffers.append if proto >= 5 else None))
 
+            assert_equal(len(buffers), 0)
             assert_equal(non_contiguous_array, depickled_non_contiguous_array)
 
     def test_roundtrip(self):