Pandas to_numpy() results in an array of lists. How do I get a 2D numpy array from this?

Question

I have a column of lists, and each list contains the same number of values.

If I do

df['column'].to_numpy()

I get an array of lists:

array([list([0, 4688, 11, 43486, 40508, 13, 5,...
       list([0, 40928, 17707, 22705, 9, 38312, 2..
       list([0, 6766, 368, 3551, 28837,..
      dtype=object)

How do I get a 2D array instead?

Answer 1

You can do this:

np.array(df['column'].tolist())

Or you can simply stack them:

np.stack(df['column'].to_numpy())

This will stack your lists on top of each other and output is a 2-D array. You have to make sure lists are of the same length. Numpy arrays are rectangular.

Answer 2

You can use list list comprehension

>>> df
                 A
0  [0, 1, 2, 3, 4]
1  [0, 1, 2, 3, 4]
2  [0, 1, 2, 3, 4]

>>> np.array([x for x in df['A']])
array([[0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4],
       [0, 1, 2, 3, 4]])