Explode an array using pandas

explode_array.py

import pandas as pd

def explode_array(df, column_to_explode, axis=1, pop=False, keep=None):
    new_cols = []
    if axis == 0:
        rows = np.repeat(range(df.shape[0]), df[column_to_explode].apply(len))
        cols = [i for i, _ in enumerate(df.columns) if _ != column_to_explode]
        new_df = df.iloc[rows, cols].copy()
        new_df[column_to_explode] = [_ for row in df[column_to_explode] for _ in row]
        new_cols.append(column_to_explode)
    if axis == 1:
        new_df = df.copy()
        for i, col in enumerate(zip(*new_df[column_to_explode])):
            new_col = "{}_{}".format(column_to_explode, i)
            new_df[new_col] = col
            new_cols.append(new_col)
        if pop:
            new_df = new_df.drop([column_to_explode], axis=1)
    if keep:
        if isinstance(keep, list):
            new_df = new_df[keep + new_cols]
        else:
            new_df = new_df[[keep] + new_cols]
    return new_df