blosc2.random: add permutation, permuted, and shuffle

Unlike the rest of the module these are not chunk-parallel: whole-array
shuffling is inherently sequential, so each call materializes the full
array into memory and shuffles it single-threaded with one spawned
SeedSequence stream, then writes the result back out. shuffle additionally
requires its argument to already be a blosc2.NDArray, since it mutates in
place and returns None to match numpy's contract.

Brings coverage of numpy.random.Generator's public methods to 42/43 (only
bytes, which returns raw bytes rather than an NDArray, remains out of
scope). Docs, coverage table, and RELEASE_NOTES updated to match.

Author: FrancescAlted

RELEASE_NOTES.md

@@ -11,7 +11,7 @@ XXX version-specific blurb XXX
   stream and is generated concurrently in a thread pool, giving full `PCG64`
   quality with genuinely parallel generation (measured ~3x faster than
   `asarray(np.random.default_rng(...).random(...))` on a 100M-element array).
-  Covers 39 of `numpy.random.Generator`'s 43 public methods (full
+  Covers 42 of `numpy.random.Generator`'s 43 public methods (full
   compatibility table in `doc/reference/random.rst`):
     - Core: `random`, `integers`, `normal`, `uniform`, `choice`
       (`replace=True` only).
@@ -25,9 +25,13 @@ XXX version-specific blurb XXX
     - 4 vector-valued distributions (output shape `shape + (k,)`, one draw
       per trailing vector): `dirichlet`, `multinomial`,
       `multivariate_hypergeometric`, `multivariate_normal`.
-    - Not implemented: `permutation`/`permuted`/`shuffle` (whole-array
-      shuffling needs cross-chunk coordination this module doesn't do) and
-      `bytes` (returns raw `bytes`, not an `NDArray`).
+    - `permutation`, `permuted`, `shuffle`: unlike the rest of the module,
+      these are not chunk-parallel — whole-array shuffling is inherently
+      sequential, so they materialize the full array and shuffle it
+      single-threaded. `shuffle` additionally requires its argument to
+      already be an `NDArray`, since it mutates in place and returns `None`,
+      matching numpy.
+    - Not implemented: `bytes` (returns raw `bytes`, not an `NDArray`).
 
 ## Changes from 4.9.0 to 4.9.1
 

doc/reference/random.rst

@@ -43,6 +43,12 @@ consistently across the module:
   :meth:`~Generator.multivariate_normal`) draw one length-``k`` vector per element:
   output shape is ``shape + (k,)``, and that trailing dimension is always kept whole
   within a chunk.
+- :meth:`~Generator.permutation`, :meth:`~Generator.permuted`, and
+  :meth:`~Generator.shuffle` are **not** chunk-parallel: whole-array shuffling is
+  inherently sequential, so these load the full array into memory and shuffle it
+  single-threaded rather than generating chunks independently. :meth:`~Generator.shuffle`
+  additionally requires its argument to already be a :class:`~blosc2.NDArray`, since it
+  mutates it in place and returns ``None``, matching numpy.
 
 Coverage of :class:`numpy.random.Generator`'s public methods:
 
@@ -129,11 +135,11 @@ Coverage of :class:`numpy.random.Generator`'s public methods:
      - ✅
      -
    * - ``permutation``
-     - ❌
-     - whole-array shuffle is inherently sequential, not per-chunk independent
+     - ✅
+     - single-threaded, full-materialization — see above
    * - ``permuted``
-     - ❌
-     - same as ``permutation``
+     - ✅
+     - single-threaded, full-materialization — see above
    * - ``poisson``
      - ✅
      -
@@ -147,8 +153,8 @@ Coverage of :class:`numpy.random.Generator`'s public methods:
      - ✅
      -
    * - ``shuffle``
-     - ❌
-     - same as ``permutation``
+     - ✅
+     - single-threaded, full-materialization; requires an ``NDArray`` — see above
    * - ``standard_cauchy``
      - ✅
      -

src/blosc2/random.py

@@ -32,6 +32,10 @@ def _chunk_slices(a: blosc2.NDArray):
         )
 
 
+def _materialize(x) -> np.ndarray:
+    return x[:] if isinstance(x, blosc2.NDArray) else np.asarray(x)
+
+
 def _bounded_map(fn, n: int, max_workers: int):
     # ponytail: window bounds peak memory to O(workers * chunk); a plain
     # executor.map would submit all n chunks upfront and could buffer O(n).
@@ -314,6 +318,42 @@ def multivariate_normal(self, mean, cov, *, shape, **kwargs: Any) -> blosc2.NDAr
         """
         return self._fill_vector("multivariate_normal", (mean, cov), {}, shape, len(mean), **kwargs)
 
+    def permutation(self, x, *, axis: int = 0, **kwargs: Any) -> blosc2.NDArray:
+        """Return a shuffled copy of ``x`` (or of ``arange(x)`` if ``x`` is an int).
+
+        Mirrors :meth:`numpy.random.Generator.permutation`.
+
+        Unlike the rest of this module, permuting is inherently sequential: it loads
+        ``x`` fully into memory and shuffles it single-threaded, rather than generating
+        chunks independently in parallel.
+        """
+        rng = np.random.default_rng(self._root.spawn(1)[0])
+        arr = np.arange(x) if isinstance(x, int | np.integer) else _materialize(x)
+        return blosc2.asarray(rng.permutation(arr, axis=axis), **kwargs)
+
+    def permuted(self, x, *, axis: int | None = None, **kwargs: Any) -> blosc2.NDArray:
+        """Return a copy of ``x`` with elements permuted along ``axis`` (flattened if ``None``).
+
+        Mirrors :meth:`numpy.random.Generator.permuted`. See :meth:`permutation` for the
+        single-threaded, full-materialization caveat.
+        """
+        rng = np.random.default_rng(self._root.spawn(1)[0])
+        return blosc2.asarray(rng.permuted(_materialize(x), axis=axis), **kwargs)
+
+    def shuffle(self, x: blosc2.NDArray, *, axis: int = 0) -> None:
+        """Shuffle ``x`` in place along ``axis``. Mirrors :meth:`numpy.random.Generator.shuffle`.
+
+        ``x`` must be a :class:`~blosc2.NDArray`: its full contents are read into memory,
+        shuffled single-threaded (see :meth:`permutation`), and written back. Returns
+        ``None``, matching numpy.
+        """
+        if not isinstance(x, blosc2.NDArray):
+            raise TypeError("blosc2.random.Generator.shuffle requires a blosc2.NDArray to shuffle in place")
+        rng = np.random.default_rng(self._root.spawn(1)[0])
+        arr = x[:]
+        rng.shuffle(arr, axis=axis)
+        x[:] = arr
+
 
 def default_rng(seed=None) -> Generator:
     """Construct a chunk-parallel :class:`Generator`, mirroring :func:`numpy.random.default_rng`.

tests/test_random.py

@@ -213,3 +213,57 @@ def test_vector_dist_chunks_mismatch_raises():
 def test_vector_dist_scalar_leading_shape():
     a = blosc2.random.default_rng(0).dirichlet([1.0, 2.0, 3.0], shape=())
     assert a.shape == (3,)
+
+
+def test_permutation_int_is_a_permutation_and_reproducible():
+    a = blosc2.random.default_rng(0).permutation(10)
+    b = blosc2.random.default_rng(0).permutation(10)
+    np.testing.assert_array_equal(a[:], b[:])
+    assert sorted(a[:].tolist()) == list(range(10))
+
+
+def test_permutation_successive_calls_differ():
+    rng = blosc2.random.default_rng(0)
+    a = rng.permutation(1000)
+    b = rng.permutation(1000)
+    assert not np.array_equal(a[:], b[:])
+
+
+def test_permutation_array_axis():
+    arr = blosc2.arange(12).reshape((3, 4))
+    p = blosc2.random.default_rng(1).permutation(arr, axis=1)
+    assert p.shape == (3, 4)
+    for row_before, row_after in zip(arr[:], p[:], strict=True):
+        assert sorted(row_after.tolist()) == sorted(row_before.tolist())
+
+
+def test_permutation_kwargs_passthrough():
+    p = blosc2.random.default_rng(0).permutation(1000, chunks=(100,), cparams={"clevel": 5})
+    assert p.chunks == (100,)
+    assert p.schunk.cparams.clevel == 5
+
+
+def test_permuted_is_a_permutation():
+    a = blosc2.random.default_rng(2).permuted(np.arange(10))
+    assert sorted(a[:].tolist()) == list(range(10))
+
+
+def test_permuted_flattens_when_axis_none():
+    a = blosc2.random.default_rng(2).permuted(blosc2.arange(12).reshape((3, 4)))
+    assert a.shape == (3, 4)
+    assert sorted(a[:].flatten().tolist()) == list(range(12))
+
+
+def test_shuffle_in_place_preserves_multiset():
+    a = blosc2.arange(10)
+    before = a[:].copy()
+    result = blosc2.random.default_rng(3).shuffle(a)
+    after = a[:]
+    assert result is None
+    assert sorted(before.tolist()) == sorted(after.tolist())
+    assert not np.array_equal(before, after)
+
+
+def test_shuffle_requires_ndarray():
+    with pytest.raises(TypeError):
+        blosc2.random.default_rng(3).shuffle(np.arange(5))