Question

我是 Spark 新手，我正在尝试使用 Pyspark 执行以下操作：

我有一个包含 3 列的数据框，"id", "number1", "number2"。

对于 "id" 的每个值，我有多个行，我想要做的是创建一个元组列表，其中包含与每个 id 对应的所有行。

例如，对于以下数据框

id | number1 | number2 |
a  |       1 |       1 |
a  |       2 |       2 |
b  |       3 |       3 |
b  |       4 |       4 |

期望的结果是 2 个列表：

[(1, 1), (2, 2)]

和

[(3, 3), (4, 4)]

我不知道如何解决这个问题，因为我是新手。我设法获得了执行以下操作的不同 ID 的列表

distinct_ids = [x for x in df.select('id').distinct().collect()]

在我更熟悉的 Pandas 中，现在我会遍历每个不同 id 的数据框并为其收集所有行，但我确信这远非最佳。

你能给我一些想法吗？ Groupby 浮现在脑海中，但我不知道如何接近

Answer 1

您可以使用 groupby 并使用 collect_list 和 array 进行聚合：

import pyspark.sql.functions as F

df2 = df.groupBy('id').agg(F.collect_list(F.array('number1', 'number2')).alias('number'))

df2.show()
+---+----------------+
| id|          number|
+---+----------------+
|  b|[[3, 3], [4, 4]]|
|  a|[[1, 1], [2, 2]]|
+---+----------------+

如果你想取回元组列表，

result = [[tuple(j) for j in i] for i in [r[0] for r in df2.select('number').orderBy('number').collect()]]

其中 result 作为 [[(1, 1), (2, 2)], [(3, 3), (4, 4)]]

如果你想要一个 numpy 数组，你可以这样做

result = np.array([r[0] for r in df2.select('number').collect()])

给出

array([[[3, 3],
        [4, 4]],

       [[1, 1],
        [2, 2]]])

基于值过滤并在火花数据框中创建列表

1 个答案: