PySpark Pandas API：如何将含列表的列拆分为多列？

问题背景 #

我在Databricks笔记本里尝试把包含列表（比如[599086.9706961295, 4503107.843920314]）的xy列拆分成“x”和“y”两列。之前在Jupyter笔记本里用两种方法都能成功拆分：

方法1：

# Jupyter中的代码
# 包含列表的列：xy
# 方法1
complete[['x', 'y']] = pd.Series(np.stack(complete['xy'].values).T.tolist())

方法2：

# 方法2
def sepXY(xy):
    return xy[0],xy[1]

complete['x'],complete['y'] = zip(*complete['xy'].apply(sepXY))

但在Databricks中使用PySpark Pandas（import pyspark.pandas as ps）时，这两种方法都报错：

• 方法1执行complete[['x', 'y']] = ps.Series(np.stack(complete['xy'].values).T.tolist())时抛出AssertionError，单独执行右侧代码能得到包含x、y列表的结果，但赋值时出错；
• 方法2执行时抛出ArrowInvalid: Could not convert (599086.9706961295, 4503107.843920314) with type tuple: did not recognize Python value type when inferring an Arrow data type；
• 尝试complete[['x','y']] = pd.DataFrame(complete.xy.tolist(), index= complete.index)时内核直接重启。

示例xy列数据如下：

xy
[599086.9706961295, 4503107.843920314]
[599088.5389507986, 4503112.7796745915]
[599072.8088083105, 4503064.139248001]
[599090.0996424126, 4503117.721156018]
[599074.3909188313, 4503068.925677084]

可行解决方案 #

方案1：直接通过索引提取元素 #

利用apply直接提取列表的第0、1位元素，避免返回tuple引发Arrow类型错误：

import pyspark.pandas as ps

# 提取x列（列表第一个元素）
complete['x'] = complete['xy'].apply(lambda x: x[0])
# 提取y列（列表第二个元素）
complete['y'] = complete['xy'].apply(lambda x: x[1])

方案2：借助PySpark原生API处理 #

PySpark Pandas底层依赖Spark，直接用Spark原生API更稳定高效，适合大数据场景：

# 转换为PySpark DataFrame
spark_df = complete.to_spark()

# 使用getItem提取列表中的元素
spark_df = spark_df.withColumn('x', spark_df['xy'].getItem(0)) \
                   .withColumn('y', spark_df['xy'].getItem(1))

# 转换回PySpark Pandas DataFrame
complete = spark_df.to_pandas_on_spark()

方案3：拆分后批量赋值 #

先将列表拆分为包含两列的临时PySpark Pandas DataFrame，再批量赋值：

# 将xy列的列表转换为两列的DataFrame
split_df = complete['xy'].apply(lambda x: ps.Series(x)).rename(columns={0: 'x', 1: 'y'})

# 赋值到原DataFrame
complete[['x', 'y']] = split_df

失败原因说明 #

• 方法1：PySpark Pandas的多列赋值逻辑与原生Pandas不同，直接用ps.Series赋值会触发断言错误，因为两者数据结构匹配规则不一致；
• 方法2：apply返回tuple时，Arrow无法识别该类型，PySpark Pandas依赖Arrow进行分布式数据传输，tuple不在其支持的类型范围内；
• 方法3：直接用本地Pandas DataFrame给分布式的PySpark Pandas DataFrame赋值，会引发序列化或内存溢出问题，导致内核重启。

内容的提问来源于stack exchange，提问作者Prathamesh Sawant