Merge pull request #23652 from aurivus-ph/faster-structured-to-unstru…

…ctured

ENH: structured_to_unstructured: view more often

doc/release/upcoming_changes/23652.improvement.rst

@@ -0,0 +1,7 @@
+``numpy.lib.recfunctions.structured_to_unstructured`` returns views in more cases
+---------------------------------------------------------------------------------
+``structured_to_unstructured`` now returns a view, if the stride between the
+fields is constant. Prior, padding between the fields or a reversed field
+would lead to a copy.
+This change only applies to ``ndarray``, ``memmap`` and ``recarray``. For all
+other array subclasses, the behavior remains unchanged.

numpy/lib/recfunctions.py

@@ -885,6 +885,50 @@ def count_elem(dt):
                     fields.extend([(d, c, o + i*size) for d, c, o in subfields])
     return fields
 
+def _common_stride(offsets, counts, itemsize):
+    """
+    Returns the stride between the fields, or None if the stride is not
+    constant. The values in "counts" designate the lengths of
+    subarrays. Subarrays are treated as many contiguous fields, with
+    always positive stride.
+    """
+    if len(offsets) <= 1:
+        return itemsize
+
+    negative = offsets[1] < offsets[0]  # negative stride
+    if negative:
+        # reverse, so offsets will be ascending
+        it = zip(reversed(offsets), reversed(counts))
+    else:
+        it = zip(offsets, counts)
+
+    prev_offset = None
+    stride = None
+    for offset, count in it:
+        if count != 1:  # subarray: always c-contiguous
+            if negative:
+                return None  # subarrays can never have a negative stride
+            if stride is None:
+                stride = itemsize
+            if stride != itemsize:
+                return None
+            end_offset = offset + (count - 1) * itemsize
+        else:
+            end_offset = offset
+
+        if prev_offset is not None:
+            new_stride = offset - prev_offset
+            if stride is None:
+                stride = new_stride
+            if stride != new_stride:
+                return None
+
+        prev_offset = end_offset
+
+    if negative:
+        return -stride
+    return stride
+
 
 def _structured_to_unstructured_dispatcher(arr, dtype=None, copy=None,
                                            casting=None):
@@ -910,9 +954,15 @@ def structured_to_unstructured(arr, dtype=None, copy=False, casting='unsafe'):
     dtype : dtype, optional
        The dtype of the output unstructured array.
     copy : bool, optional
-        See copy argument to `numpy.ndarray.astype`. If true, always return a
-        copy. If false, and `dtype` requirements are satisfied, a view is
-        returned.
+        If true, always return a copy. If false, a view is returned if
+        possible, such as when the `dtype` and strides of the fields are
+        suitable and the array subtype is one of `np.ndarray`, `np.recarray`
+        or `np.memmap`.
+
+        .. versionchanged:: 1.25.0
+            A view can now be returned if the fields are separated by a
+            uniform stride.
+
     casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
         See casting argument of `numpy.ndarray.astype`. Controls what kind of
         data casting may occur.
@@ -960,7 +1010,7 @@ def structured_to_unstructured(arr, dtype=None, copy=False, casting='unsafe'):
     if dtype is None:
         out_dtype = np.result_type(*[dt.base for dt in dts])
     else:
-        out_dtype = dtype
+        out_dtype = np.dtype(dtype)
 
     # Use a series of views and casts to convert to an unstructured array:
 
@@ -972,6 +1022,39 @@ def structured_to_unstructured(arr, dtype=None, copy=False, casting='unsafe'):
                                  'itemsize': arr.dtype.itemsize})
     arr = arr.view(flattened_fields)
 
+    # we only allow a few types to be unstructured by manipulating the
+    # strides, because we know it won't work with, for example, np.matrix nor
+    # np.ma.MaskedArray.
+    can_view = type(arr) in (np.ndarray, np.recarray, np.memmap)
+    if (not copy) and can_view and all(dt.base == out_dtype for dt in dts):
+        # all elements have the right dtype already; if they have a common
+        # stride, we can just return a view
+        common_stride = _common_stride(offsets, counts, out_dtype.itemsize)
+        if common_stride is not None:
+            wrap = arr.__array_wrap__
+
+            new_shape = arr.shape + (sum(counts), out_dtype.itemsize)
+            new_strides = arr.strides + (abs(common_stride), 1)
+
+            arr = arr[..., np.newaxis].view(np.uint8)  # view as bytes
+            arr = arr[..., min(offsets):]  # remove the leading unused data
+            arr = np.lib.stride_tricks.as_strided(arr,
+                                                  new_shape,
+                                                  new_strides,
+                                                  subok=True)
+
+            # cast and drop the last dimension again
+            arr = arr.view(out_dtype)[..., 0]
+
+            if common_stride < 0:
+                arr = arr[..., ::-1]  # reverse, if the stride was negative
+            if type(arr) is not type(wrap.__self__):
+                # Some types (e.g. recarray) turn into an ndarray along the
+                # way, so we have to wrap it again in order to match the
+                # behavior with copy=True.
+                arr = wrap(arr)
+            return arr
+
     # next cast to a packed format with all fields converted to new dtype
     packed_fields = np.dtype({'names': names,
                               'formats': [(out_dtype, dt.shape) for dt in dts]})

numpy/lib/tests/test_recfunctions.py

@@ -228,7 +228,7 @@ def test_repack_fields(self):
         dt = np.dtype((np.record, dt))
         assert_(repack_fields(dt).type is np.record)
 
-    def test_structured_to_unstructured(self):
+    def test_structured_to_unstructured(self, tmp_path):
         a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
         out = structured_to_unstructured(a)
         assert_equal(out, np.zeros((4,5), dtype='f8'))
@@ -263,8 +263,13 @@ def test_structured_to_unstructured(self):
                      dtype=[('x', 'i4'), ('y', 'i4'), ('z', 'i4')])
         dd = structured_to_unstructured(d)
         ddd = unstructured_to_structured(dd, d.dtype)
-        assert_(dd.base is d)
-        assert_(ddd.base is d)
+        assert_(np.shares_memory(dd, d))
+        assert_(np.shares_memory(ddd, d))
+
+        # check that reversing the order of attributes works
+        dd_attrib_rev = structured_to_unstructured(d[['z', 'x']])
+        assert_equal(dd_attrib_rev, [[5, 1], [7, 4], [11, 7], [12, 10]])
+        assert_(np.shares_memory(dd_attrib_rev, d))
 
         # including uniform fields with subarrays unpacked
         d = np.array([(1, [2,  3], [[ 4,  5], [ 6,  7]]),
@@ -273,8 +278,30 @@ def test_structured_to_unstructured(self):
                             ('x2', ('i4', (2, 2)))])
         dd = structured_to_unstructured(d)
         ddd = unstructured_to_structured(dd, d.dtype)
-        assert_(dd.base is d)
-        assert_(ddd.base is d)
+        assert_(np.shares_memory(dd, d))
+        assert_(np.shares_memory(ddd, d))
+
+        # check that reversing with sub-arrays works as expected
+        d_rev = d[::-1]
+        dd_rev = structured_to_unstructured(d_rev)
+        assert_equal(dd_rev, [[8, 9, 10, 11, 12, 13, 14],
+                              [1, 2, 3, 4, 5, 6, 7]])
+
+        # check that sub-arrays keep the order of their values
+        d_attrib_rev = d[['x2', 'x1', 'x0']]
+        dd_attrib_rev = structured_to_unstructured(d_attrib_rev)
+        assert_equal(dd_attrib_rev, [[4, 5, 6, 7, 2, 3, 1],
+                                     [11, 12, 13, 14, 9, 10, 8]])
+
+        # with ignored field at the end
+        d = np.array([(1, [2,  3], [[4, 5], [6, 7]], 32),
+                      (8, [9, 10], [[11, 12], [13, 14]], 64)],
+                     dtype=[('x0', 'i4'), ('x1', ('i4', 2)),
+                            ('x2', ('i4', (2, 2))), ('ignored', 'u1')])
+        dd = structured_to_unstructured(d[['x0', 'x1', 'x2']])
+        assert_(np.shares_memory(dd, d))
+        assert_equal(dd, [[1, 2, 3, 4, 5, 6, 7],
+                          [8, 9, 10, 11, 12, 13, 14]])
 
         # test that nested fields with identical names don't break anything
         point = np.dtype([('x', int), ('y', int)])
@@ -318,6 +345,35 @@ def inspect(dt, dtype=None):
         assert_raises(NotImplementedError, unstructured_to_structured,
                                            np.zeros((3,0), dtype=np.int32))
 
+        # test supported ndarray subclasses
+        d_plain = np.array([(1, 2), (3, 4)], dtype=[('a', 'i4'), ('b', 'i4')])
+        dd_expected = structured_to_unstructured(d_plain, copy=True)
+
+        # recarray
+        d = d_plain.view(np.recarray)
+
+        dd = structured_to_unstructured(d, copy=False)
+        ddd = structured_to_unstructured(d, copy=True)
+        assert_(np.shares_memory(d, dd))
+        assert_(type(dd) is np.recarray)
+        assert_(type(ddd) is np.recarray)
+        assert_equal(dd, dd_expected)
+        assert_equal(ddd, dd_expected)
+
+        # memmap
+        d = np.memmap(tmp_path / 'memmap',
+                      mode='w+',
+                      dtype=d_plain.dtype,
+                      shape=d_plain.shape)
+        d[:] = d_plain
+        dd = structured_to_unstructured(d, copy=False)
+        ddd = structured_to_unstructured(d, copy=True)
+        assert_(np.shares_memory(d, dd))
+        assert_(type(dd) is np.memmap)
+        assert_(type(ddd) is np.memmap)
+        assert_equal(dd, dd_expected)
+        assert_equal(ddd, dd_expected)
+
     def test_unstructured_to_structured(self):
         # test if dtype is the args of np.dtype
         a = np.zeros((20, 2))