PySpark中对两列连续区间值进行分组的实现求助

你的核心问题在于原始数据的区间并非按连续顺序排列（比如索引12的区间位于末尾但属于前面的连续组），直接用上下行shift对比无法识别跨位置的连续关系。以下是可行的实现步骤及代码：

实现代码 #

import pandas as pd

# 构造示例DataFrame
data = {
    'index': [1,2,3,4,5,6,7,8,9,10,11,12],
    'initial_range': [1000000,6000000,6300000,6400000,6600000,6700000,6750000,7000000,7600000,7700000,7750000,6500000],
    'final_range': [5999999,6299999,6399999,6499999,6699999,6749999,6799999,7399999,7699999,7749999,7799999,6549999]
}
df = pd.DataFrame(data).set_index('index')

# 1. 复制数据并按initial_range排序，保留原始索引
sorted_df = df.reset_index().sort_values('initial_range').set_index('index')

# 2. 标记新分组起始点：当前区间与前一区间不连续时为新组
sorted_df['new_group'] = sorted_df['initial_range'] != sorted_df['final_range'].shift() + 1

# 3. 生成唯一分组ID
sorted_df['group_id'] = sorted_df['new_group'].cumsum()

# 4. 计算每组的最小起始值和最大结束值，拼接分组字符串
group_mapping = sorted_df.groupby('group_id').agg(
    min_initial=('initial_range', 'min'),
    max_final=('final_range', 'max')
).apply(lambda row: f"{row['min_initial']}-{row['max_final']}", axis=1)

# 5. 将分组信息映射回原DataFrame
df['grouping'] = sorted_df['group_id'].map(group_mapping)

# 输出按原始索引排序的结果
print(df.sort_index())

步骤说明 #

1. 排序数据：先按initial_range排序，确保连续的区间在序列中相邻，解决原始数据中区间乱序的问题。
2. 标记分组边界：判断当前区间的起始值是否等于前一区间的结束值+1，不等则标记为新分组起点。
3. 生成分组ID：对新分组标记进行累加，为每个连续区间组分配唯一ID。
4. 生成分组名称：按分组ID聚合，取每组的最小起始值和最大结束值拼接成要求的格式。
5. 映射回原数据：将分组名称对应到原始DataFrame的每一行，保留原始索引顺序。

运行代码后将得到你需要的结果，其中索引12的区间会正确归到1000000-6549999分组。

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