Make random_state_data() faster (dask#2379)

* Make random_state_data() faster

Also speed up test_random_state_data

* Fix doctest

Author: pitrou

dask/array/random.py

@@ -37,7 +37,7 @@ class RandomState(object):
     >>> state = da.random.RandomState(1234)  # a seed
     >>> x = state.normal(10, 0.1, size=3, chunks=(2,))
     >>> x.compute()
-    array([ 10.06307943,   9.91493648,  10.0822082 ])
+    array([ 10.01867852,  10.04812289,   9.89649746])
 
     See Also:
         np.random.RandomState

dask/tests/test_utils.py

@@ -59,11 +59,10 @@ class Bar(object):
     assert foo((1, 2.0, b)) == (2, 1.0, b)
 
 
-@pytest.mark.slow
 def test_random_state_data():
     seed = 37
     state = np.random.RandomState(seed)
-    n = 100000
+    n = 10000
 
     # Use an integer
     states = random_state_data(n, seed)
@@ -72,6 +71,7 @@ def test_random_state_data():
     # Use RandomState object
     states2 = random_state_data(n, state)
     for s1, s2 in zip(states, states2):
+        assert s1.shape == (624,)
         assert (s1 == s2).all()
 
     # Consistent ordering

dask/utils.py

@@ -263,7 +263,7 @@ def random_state_data(n, random_state=None):
     Parameters
     ----------
     n : int
-        Number of tuples to return.
+        Number of arrays to return.
     random_state : int or np.random.RandomState, optional
         If an int, is used to seed a new ``RandomState``.
     """
@@ -272,9 +272,10 @@ def random_state_data(n, random_state=None):
     if not isinstance(random_state, np.random.RandomState):
         random_state = np.random.RandomState(random_state)
 
-    maxuint32 = np.iinfo(np.uint32).max
-    return [(random_state.rand(624) * maxuint32).astype('uint32')
-            for i in range(n)]
+    random_data = random_state.bytes(624 * n * 4)  # `n * 624` 32-bit integers
+    l = list(np.frombuffer(random_data, dtype=np.uint32).reshape((n, -1)))
+    assert len(l) == n
+    return l
 
 
 def is_integer(i):