为 Spark 数据框中的每个组创建索引

Question

我在 Spark 中有一个包含 2 列的数据框，group_id 和 value，其中 value 是双列。我想根据group_id对数据进行分组，按value对每个组进行排序，然后添加第三列index，表示value在组值排序中的位置。

例如，考虑以下输入数据：

+--------+-----+
|group_id|value|
+--------+-----+
|1       |1.3  |
|2       |0.8  |
|1       |3.4  |
|1       |-1.7 |
|2       |2.3  |
|2       |5.9  |
|1       |2.7  |
|1       |0.0  |
+--------+-----+

然后输出将类似于

+--------+-----+-----+
|group_id|value|index|
+--------+-----+-----+
|1       |-1.7 |1    |
|1       |0.0  |2    |
|1       |1.3  |3    |
|1       |2.7  |4    |
|1       |3.4  |5    |
|2       |0.8  |1    |
|2       |2.3  |2    |
|2       |5.9  |3    |
+--------+-----+-----+

索引是否从0开始以及排序是升序还是降序并不重要。

作为后续，考虑在原始数据中存在第三列extra 的情况，该列在某些(group_id, value) 组合中采用多个值。一个例子是：

+--------+-----+-----+
|group_id|value|extra|
+--------+-----+-----+
|1       |1.3  |1    |
|1       |1.3  |2    |
|2       |0.8  |1    |
|1       |3.4  |1    |
|1       |3.4  |2    |
|1       |3.4  |3    |
|1       |-1.7 |1    |
|2       |2.3  |1    |
|2       |5.9  |1    |
|1       |2.7  |1    |
|1       |0.0  |1    |
+--------+-----+-----+

有没有办法添加一个index 列，这样extra 列不被考虑但仍然保留？这种情况下的输出是

+--------+-----+-----+-----+
|group_id|value|extra|index|
+--------+-----+-----+-----+
|1       |-1.7 |1    |1    |
|1       |0.0  |1    |2    |
|1       |1.3  |1    |3    |
|1       |1.3  |2    |3    |
|1       |2.7  |1    |4    |
|1       |3.4  |1    |5    |
|1       |3.4  |2    |5    |
|1       |3.4  |3    |5    |
|2       |0.8  |1    |1    |
|2       |2.3  |1    |2    |
|2       |5.9  |1    |3    |
+--------+-----+-----+-----+

我知道可以通过复制数据、删除 extra 列来做到这一点

1. 复制数据
2. 删除extra 列
3. 执行distinct 操作，这将导致原始示例中的数据
4. 使用原始解计算index 列
5. 将结果与第二个示例中的数据相结合

但是，这将涉及大量额外的计算和开销。

Answer 1

您可以使用Window 函数创建基于value 的排名列，由group_id 分区：

from pyspark.sql.window import Window
from pyspark.sql.functions import rank, dense_rank
# Define window
window = Window.partitionBy(df['group_id']).orderBy(df['value'])
# Create column
df.select('*', rank().over(window).alias('index')).show()
+--------+-----+-----+
|group_id|value|index|
+--------+-----+-----+
|       1| -1.7|    1|
|       1|  0.0|    2|
|       1|  1.3|    3|
|       1|  2.7|    4|
|       1|  3.4|    5|
|       2|  0.8|    1|
|       2|  2.3|    2|
|       2|  5.9|    3|
+--------+-----+-----+

因为您首先选择了'*'，所以您也可以使用上述代码保留所有其他变量。但是，您的第二个示例表明您正在寻找函数dense_rank()，它作为一个没有间隙的排名列：

df.select('*', dense_rank().over(window).alias('index'))