Rapid creation of nD array from generator yielding numpy arrays?

Question

I have a generator that yields NumPy arrays, and need a way to rapidly construct another NumPy array from the results of the generator (array of arrays) by taking a specific number of yields from the generator. Speed is the critical aspect in my problem. I've tried np.fromiter but it seems it doesn't support constructing from arrays:

import numpy as np

def generator():
    for i in range(5):
        yield np.array([i]*10)


arr = np.fromiter(iter(generator()), dtype=np.ndarray, count=3)

This throws an error, as described in several other SO posts:

Calling np.sum(np.fromiter(generator))

Numpy ValueError: setting an array element with a sequence

However, I haven't found any answer that offers a rapid way to source arrays from the generator without having to do:

it = iter(generator())
arr = np.array([next(it) for _ in range(3)])

Here it is indeed shown that np.fromiter is much faster: Faster way to convert list of objects to numpy array

Is it possible to rapidly source numpy arrays from the generator without having use the slow list to array conversion? I specifically want to avoid the np.array(list(...)) construct, because I will be calling it hundreds of thousands of times, and the delay will eventually add up and make a big difference in execution time.

Answer 1

What about using itertools.islice?

from itertools import islice
g = generator()
arr = np.array(list(islice(g, 3)))

# or in one line:
# arr = np.array(list(islice(generator(), 3

output:

array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]])