PySpark按id_A分组并按Distance升序生成id_B列表

原始DataFrame #

+----------+----------+------------------+
|      id_A|      id_B|   Distance       | 
+----------+----------+------------------+
| 120745612| 122913167|0.6142857142857143|
|1243257970| 370926553|0.8061224489795918|
|1305652409| 253051944|0.8252427184466019|
|1350805455| 311286173|0.5789473684210527|
|1544864070| 390580289|0.7894736842105263|
| 164533143| 763751752|0.8153846153846154|
|1683553267| 787287056|0.9117647058823529|
| 175951349| 175951349|               0.0|

实现步骤 #

1. 先对DataFrame按id_A分组、Distance升序排序，确保组内id_B按距离从小到大排列
2. 分组后将每个组的id_B聚合为列表

代码实现 #

import pandas as pd

# 假设你的DataFrame名为df
# 按id_A和Distance升序排序
df_sorted = df.sort_values(by=['id_A', 'Distance'], ascending=[True, True])

# 分组聚合生成id_B列表
result = df_sorted.groupby('id_A')['id_B'].agg(list).reset_index()

# 打印结果
print(result)

输出结果 #

id_A               id_B
0  120745612        [122913167]
1 1243257970        [370926553]
2 1305652409        [253051944]
3 1350805455        [311286173]
4 1544864070        [390580289]
5  164533143        [763751752]
6 1683553267        [787287056]
7  175951349        [175951349]

注：你的示例预期中列表包含多个元素，但当前原始数据里每个id_A仅对应一个id_B，所以每个列表只有一个元素。如果后续有同一id_A对应多个id_B的情况，上述代码会自动按Distance升序排列列表内的元素。

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