基于模式分组列的PySpark实现：将价格数量列转为字典数组

问题描述 #

需要将输入DataFrame中的price_n和qty_n列按数字后缀分组，转换为包含price、qty、comments字段的字典数组，最终生成新的price_range列。

输入DataFrame #

期望输出DataFrame #

解决方案（Python Pandas） #

通过列名拆分+分组聚合的方式实现，具体代码如下：

import pandas as pd

# 构造输入DataFrame
df = pd.DataFrame({
    'item_name': ['Samsung Z', 'Moto G4', 'Mi 4i', 'Moto G3'],
    'price_1': [10000, 12000, 15000, 20000],
    'qty_1': [5, 10, 8, 5],
    'price_2': [9000, 10000, 12000, 18000],
    'qty_2': [10, 20, 20, 12],
    'price_3': [7000, 6000, 10000, 15000],
    'qty_3': [20, 50, 25, 30],
    'url': ['amazon.com', 'ebay.com', 'deals.com', 'ebay.com']
})

# 提取price和qty列的后缀数字
suffixes = [col.split('_')[1] for col in df.columns if col.startswith('price_')]

# 生成price_range列：逐行生成字典数组
df['price_range'] = df.apply(lambda row: [
    {'price': row[f'price_{s}'], 'qty': row[f'qty_{s}'], 'comments': ''}
    for s in suffixes
], axis=1)

# 保留目标列得到结果
result_df = df[['item_name', 'price_range', 'url']]
print(result_df)

代码说明 #

1. 从列名中提取price_n的后缀数字，确保price和qty的配对对应；
2. 使用apply逐行处理数据，为每个后缀数字配对对应的price和qty值，生成包含空comments字段的字典；
3. 将字典列表赋值给price_range列，最后筛选出需要的列即可得到目标DataFrame。

内容的提问来源于stack exchange，提问作者Ramineni Ravi Teja