如何在Spark DataFrame中按ID计算满足条件的连续行温度差值

核心思路 #

针对每个ID分组，仅筛选ch_status=0的行，计算这些行的TEMPERATURE连续差值，将结果存入新增的diff列；不符合条件的行（ch_status≠0）的diff列保留缺失值。

代码实现 #

假设你的DataFrame名为df，可以使用以下代码完成需求：

import pandas as pd

# 初始化diff列为缺失值
df['diff'] = pd.NA

# 按ID分组处理每个组的数据
for id_group, group_data in df.groupby('ID'):
    # 筛选当前组内ch_status为0的行
    valid_mask = group_data['ch_status'] == 0
    # 计算符合条件行的TEMPERATURE连续差值
    temp_diff = group_data.loc[valid_mask, 'TEMPERATURE'].diff()
    # 将计算结果赋值回原DataFrame对应的位置
    df.loc[temp_diff.index, 'diff'] = temp_diff

补充说明 #

• diff()方法会计算当前行与前一行的TEMPERATURE差值，每个ID组内符合条件的第一行diff值为缺失值（因为没有前一行数据），如果需要将其设为0，可修改代码为：temp_diff = group_data.loc[valid_mask, 'TEMPERATURE'].diff().fillna(0)
• 如果你的DataFrame是按时间或其他顺序排序的，确保处理前已经按ID和排序字段（如时间）排好序，否则连续差值会不符合预期。

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