How to I join arrays that are in a list or from a generator?

Numba Support: How do I do ...?
1

I am having trouble neatly joining numpy arrays that are in a list or from a generator in numba jited functions. The simplest example is as follows:

import numpy as np
from numba import njit

@njit
def my_program():
    array_1 = np.array([0,3])
    array_2 = np.array([0,4,2,3])
    array_3 = np.array([9,1,3,3,5,9])
    
    list_of_arrays = [array_1, array_2, array_3]
    
    return join_arrays(list_of_arrays)

my_program()

Where join_arrays should return a 1-D numpy array like np.array([0,3,0,4,2,3,9,1,3,3,5,9]).

I have tried:


@njit
def join_arrays_a(list_of_arrays):
    return np.hstack(list_of_arrays)

@njit
def join_arrays_b(list_of_arrays):
    return np.hstack(tuple(list_of_arrays))

@njit
def join_arrays_c(arrays):
    tot_len = sum(list(map(len, arrays)))
    new_array = np.zeros(shape=(tot_len,), dtype=np.int64)
    prev_array_lens = 0
    for array in arrays:
        for i in range(len(array)):
            new_array[prev_array_lens+i] = array[i]
        prev_array_lens += len(array)
    return new_array

join_arrays_c works but I would hope that there was a better way to do this, or some utility functions somewhere so that I do not have to write so much code like this. Am I missing something?
join_arrays_a and join_arrays_b don’t work when used in the my_program instead of join_arrays_c.

2

Is this the real problem? Numba seems like an odd solution here and likely won’t add much speed (if any)

It looks like you’re trying to re-implement numpy,concatenate, which seems like the right tool for this job.

3

Hi Nelson, much like hstack, numpy.concatenate has the same restriction that the arrays must be in a tuple when called within a jitted function and so does not solve this problem. And indeed, I am not expecting numba to speed up this function when there are only a small number of arrays, but for a large number I would think that this would be much slower than avoiding any python object creation.

4

Thanks for the clarification. I don’t know of anything better than what you’ve demonstrated