Spark Scala中按年度统计Qty的特殊需求实现（含缺失月份处理）

原始数据 #

给定包含Month、Fruit、Qty列的DataFrame，示例如下：

需求 #

按Fruit分组，统计每一年的Qty总和，输出包含以下列的DataFrame：

• Year：年份
• Fruit：水果类别
• sum_of_qty_This_year：当前年份符合规则的Qty总和
• sum_of_qty_previous_year：前一年的Qty总和

输出示例：

特殊规则 #

若前一年存在缺失的月份（如下表中前一年缺失3月、9月），当前年份统计Qty时需排除这些缺失的月份，该规则适用于所有年份：

实现代码 #

import pandas as pd

# 假设原始数据已加载到df中，先将Month转为datetime类型
df['Month'] = pd.to_datetime(df['Month'])

# 提取年份和月份数字
df['Year'] = df['Month'].dt.year
df['Month_num'] = df['Month'].dt.month

# 按Fruit和Year分组，记录每个年份存在的月份集合、原始年份总Qty
grouped = df.groupby(['Fruit', 'Year'])
year_month_map = grouped['Month_num'].apply(set).reset_index(name='Exist_months')
year_total_qty = grouped['Qty'].sum().reset_index(name='total_qty')

# 合并月份集合与原始总Qty数据
merged = pd.merge(year_total_qty, year_month_map, on=['Fruit', 'Year'])

# 按水果、年份排序，获取前一年的月份集合
merged = merged.sort_values(['Fruit', 'Year'])
merged['prev_year_months'] = merged.groupby('Fruit')['Exist_months'].shift(1)

# 定义函数：根据前一年的月份集合，重新计算当前年份的合规Qty总和
def get_adjusted_sum(row):
    if pd.isna(row['prev_year_months']):
        # 最早年份无前置年份，直接用原始总Qty
        return row['total_qty']
    else:
        # 筛选当前年份中，前一年存在的月份数据求和
        filter_mask = (df['Fruit'] == row['Fruit']) & (df['Year'] == row['Year']) & (df['Month_num'].isin(row['prev_year_months']))
        return df.loc[filter_mask, 'Qty'].sum()

merged['sum_of_qty_This_year'] = merged.apply(get_adjusted_sum, axis=1)

# 获取前一年的合规总和
merged['sum_of_qty_previous_year'] = merged.groupby('Fruit')['sum_of_qty_This_year'].shift(1)

# 整理为目标输出格式
result = merged[['Year', 'Fruit', 'sum_of_qty_This_year', 'sum_of_qty_previous_year']]
print(result)

代码说明 #

1. 数据预处理：将Month转为datetime类型，拆分出年份和月份数字，便于后续分组筛选。
2. 分组统计基础数据：按水果和年份分组，记录每个年份实际存在的月份集合、原始年份总Qty。
3. 匹配前置年份规则：对每个水果的年份排序后，用shift方法获取前一年的月份集合。
4. 计算合规总和：根据前一年的月份集合，重新筛选当前年份的对应月份数据求和；最早年份直接使用原始总和。
5. 生成最终结果：再次用shift获取前一年的合规总和，整理得到目标结构的DataFrame。

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