如何使用Pandas API on Spark对字符串列执行pivot操作

问题复现 #

Pandas 正常执行代码 #

import pandas as pd

pdf = pd.DataFrame({
    'id': [1, 1, 2, 2],
    'category': ['A', 'B', 'A', 'B'],
    'foo': ['x', 'y', 'z', 'w']
})
pandas_pivoted = pdf.pivot(index='id', columns='category', values='foo')
print(pandas_pivoted)

输出：

category  A  B
id            
1         x  y
2         z  w

Pandas API on Spark 报错代码 #

import pyspark.pandas as ps

psdf = ps.DataFrame({
    'id': [1, 1, 2, 2],
    'category': ['A', 'B', 'A', 'B'],
    'foo': ['x', 'y', 'z', 'w']
})
# 执行报错：values should be a numeric type
spark_pivoted = psdf.pivot(index='id', columns='category', values='foo')

解决方案 #

利用groupby+agg+unstack组合实现与Pandas一致的字符串列pivot效果，全程使用Pandas API on Spark语法，无需原生PySpark：

import pyspark.pandas as ps

psdf = ps.DataFrame({
    'id': [1, 1, 2, 2],
    'category': ['A', 'B', 'A', 'B'],
    'foo': ['x', 'y', 'z', 'w']
})

# 分组聚合后展开列，模拟pivot逻辑
pivoted = psdf.groupby(['id', 'category'])['foo'].first().unstack('category')
# 对齐Pandas输出的索引和列名格式
pivoted.columns.name = 'category'
pivoted.index.name = 'id'

print(pivoted)

输出与Pandas版本完全一致：

category  A  B
id            
1         x  y
2         z  w

原理说明 #

Pandas的字符串列pivot本质是对(index, columns)分组取唯一字符串值，而groupby+first()（或last()）可以安全聚合字符串列，再通过unstack将分组列转为表头，完美复刻Pivot的行为，同时避开Spark Pandas API对pivot值列类型的限制。

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