pyspark 相当于 pandas groupby('col1').col2.head()

Question

我有一个 Spark Dataframe，其中对于具有给定列值 (col1) 的每组行，我想获取 (col2) 中值的样本。 col1 的每个可能值的行数可能相差很大，所以我只是在寻找每种类型的固定数字，比如 10。

可能有更好的方法来做到这一点，但自然的方法似乎是 df.groupby('col1')

在 pandas 中，我可以做到 df.groupby('col1').col2.head()

我知道 spark 数据帧不是 pandas 数据帧，但这是一个很好的类比。

我想我可以循环遍历所有 col1 类型作为过滤器，但这看起来非常恶心。

关于如何做到这一点有什么想法吗？谢谢。

Answer 1

让我创建一个包含两列的示例 Spark 数据框。

df = SparkSQLContext.createDataFrame([[1, 'r1'],
 [1, 'r2'],
 [1, 'r2'],
 [2, 'r1'],
 [3, 'r1'],
 [3, 'r2'],
 [4, 'r1'],
 [5, 'r1'],
 [5, 'r2'],
 [5, 'r1']], schema=['col1', 'col2'])
df.show()

+----+----+
|col1|col2|
+----+----+
|   1|  r1|
|   1|  r2|
|   1|  r2|
|   2|  r1|
|   3|  r1|
|   3|  r2|
|   4|  r1|
|   5|  r1|
|   5|  r2|
|   5|  r1|
+----+----+

按 col1 分组后，我们得到 GroupedData 对象（而不是 Spark Dataframe）。您可以使用聚合函数，例如最小值、最大值、平均值。但是获取 head() 有点棘手。我们需要将 GroupedData 对象转换回 Spark Dataframe。这可以使用 pyspark collect_list() 聚合函数来完成。

from pyspark.sql import functions
df1 = df.groupBy(['col1']).agg(functions.collect_list("col2")).show(n=3)

输出为：

+----+------------------+
|col1|collect_list(col2)|
+----+------------------+
|   5|      [r1, r2, r1]|
|   1|      [r1, r2, r2]|
|   3|          [r1, r2]|
+----+------------------+
only showing top 3 rows

pyspark 相当于 pandas groupby('col1').col2.head()

pyspark equivalent of pandas groupby('col1').col2.head()

group-by

sample

pyspark

spark-dataframe