如何在Dask中实现类似Pandas的explode功能？（适配低版本Pandas）

原始数据结构 #

我有如下结构的Pandas DataFrame：

import numpy as np
import pandas as pd

df = pd.DataFrame({"ph_number" : ['1234','2345','1234','1234','2345','1234','2345'],
                    "year": [2022,2022,2023,2022,2022,2022,2022],
                    "month": [9,10,1,10,8,11,12],
                    "device_set": ['vivo 1915:vivo','SM-A510F:samsung','1718:vivo vivo~^!vivo 1718:vivo','vivo 1915:vivo','SM-A510F:samsung~^!vivo 1718:vivo','vivo 1915:vivo','SM-A510F:samsung']
                  })

需求说明 #

需要将device_set列按分隔符~^!拆分后展开，对应每行重复其他列的数值。由于生产环境Pandas版本为0.23.4，无法使用pd.explode，我已经实现了对应的Pandas版本代码。

已实现的Pandas代码 #

df['device_set'] = df['device_set'].astype(str).apply(lambda x: x.split('~^!'))

pd.DataFrame({col:np.repeat(np.array(df[col].values), np.array(df["device_set"].str.len())) for col in df.columns.drop("device_set")}).assign(**{"device_set":np.concatenate(df["device_set"].values)})[df.columns]

错误的Dask实现代码 #

我自行编写的Dask版本代码无法正常运行，代码如下：

def dask_explode(df, lst_cols):
    with dask_session(time_to_use=2*60*60) as dask_client:
        mdd = dd.from_pandas(df, npartitions=54)
        mdd = pd.DataFrame({col:np.repeat(np.array(mdd[col].values), np.array(mdd[lst_cols].str.len())) for col in mdd.columns.drop(lst_cols)}).assign(**{lst_cols:np.concatenate(mdd[lst_cols].values)})[mdd.columns]
        return mdd.compute()

df['device_set'] = df['device_set'].astype(str).apply(lambda x: x.split('~^!'))

output = dask_explode(df, "device_set")

请求正确的Dask实现方案 #

内容的提问来源于stack exchange，提问作者Aayush Gupta