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

Posted 2023-04-15

【中文标题】pyspark 相当于 pandas groupby(\'col1\').col2.head()

【英文标题】：pyspark equivalent of pandas groupby('col1').col2.head()pyspark 相当于 pandas groupby('col1').col2.head()

【发布时间】：2018-05-09 11:23:54

【问题描述】：

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

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

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

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

我想我可以循环遍历所有 col1 类型作为过滤器，但这似乎非常讨厌。

关于如何做到这一点的任何想法？谢谢。

【参考方案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）。您可以使用聚合函数，例如 min、max、average。但是得到一个 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

【讨论】：

完美。谢谢你。我认为在***上提问是我的大脑用来回答问题的过程——我在发布后不久就可以用udf来回答——但这更干净。我不知道 collect_list() 函数；显然还需要更多地查看 sql.functions。