pythonjsonpandasdataframepandas-explode

Efficient way to unnest (explode) multiple list columns in a pandas DataFrame


I am reading multiple JSON objects into one DataFrame. The problem is that some of the columns are lists. Also, the data is very big and because of that I cannot use the available solutions on the internet. They are very slow and memory-inefficient

Here is how my data looks like:

df = pd.DataFrame({'A': ['x1','x2','x3', 'x4'], 'B':[['v1','v2'],['v3','v4'],['v5','v6'],['v7','v8']], 'C':[['c1','c2'],['c3','c4'],['c5','c6'],['c7','c8']],'D':[['d1','d2'],['d3','d4'],['d5','d6'],['d7','d8']], 'E':[['e1','e2'],['e3','e4'],['e5','e6'],['e7','e8']]})
    A       B          C           D           E
0   x1  [v1, v2]    [c1, c2]    [d1, d2]    [e1, e2]
1   x2  [v3, v4]    [c3, c4]    [d3, d4]    [e3, e4]
2   x3  [v5, v6]    [c5, c6]    [d5, d6]    [e5, e6]
3   x4  [v7, v8]    [c7, c8]    [d7, d8]    [e7, e8]

And this is the shape of my data: (441079, 12)

My desired output is:

    A       B          C           D           E
0   x1      v1         c1         d1          e1
0   x1      v2         c2         d2          e2
1   x2      v3         c3         d3          e3
1   x2      v4         c4         d4          e4
.....

EDIT: After being marked as duplicate, I would like to stress on the fact that in this question I was looking for an efficient method of exploding multiple columns. Therefore the approved answer is able to explode an arbitrary number of columns on very large datasets efficiently. Something that the answers to the other question failed to do (and that was the reason I asked this question after testing those solutions).


Solution

  • pandas >= 1.3

    In more recent versions, pandas allows you to explode multiple columns at once using DataFrame.explode, provided all values have lists of equal size. Thus, you are able to use this:

    df.explode(['B', 'C', 'D', 'E']).reset_index(drop=True)
    
        A   B   C   D   E
    0  x1  v1  c1  d1  e1
    1  x1  v2  c2  d2  e2
    2  x2  v3  c3  d3  e3
    3  x2  v4  c4  d4  e4
    4  x3  v5  c5  d5  e5
    5  x3  v6  c6  d6  e6
    6  x4  v7  c7  d7  e7
    7  x4  v8  c8  d8  e8
    

    pandas >= 0.25

    For slightly older versions, you can apply Series.explode on each column.

    df.set_index(['A']).apply(pd.Series.explode).reset_index()
    
        A   B   C   D   E
    0  x1  v1  c1  d1  e1
    1  x1  v2  c2  d2  e2
    2  x2  v3  c3  d3  e3
    3  x2  v4  c4  d4  e4
    4  x3  v5  c5  d5  e5
    5  x3  v6  c6  d6  e6
    6  x4  v7  c7  d7  e7
    7  x4  v8  c8  d8  e8
    

    The idea is to set as the index all columns that must NOT be exploded first, then reset the index after.

    Funnily enough, this happens to be faster than calling df.explode, according to my tests. YMMV.


    explode methods are quite performant in general:

    df2 = pd.concat([df] * 100, ignore_index=True)
    
    %timeit df2.explode(['B', 'C', 'D', 'E']).reset_index(drop=True)
    %timeit df2.set_index(['A']).apply(pd.Series.explode).reset_index() # fastest
    %%timeit
    (df2.set_index('A')
        .apply(lambda x: x.apply(pd.Series).stack())
        .reset_index()
        .drop('level_1', axis=1))
    
    
    2.59 ms ± 112 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    1.27 ms ± 239 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
    120 ms ± 9.48 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)