pyspark collect_set或collect_list with groupby

时间:2016-06-02 00:17:59

标签: list group-by set pyspark collect

如何在collect_set之后的数据框中使用collect_listgroupby。例如:df.groupby('key').collect_set('values')。我收到错误:AttributeError: 'GroupedData' object has no attribute 'collect_set'

2 个答案:

答案 0 :(得分:57)

你需要使用agg。例如:

from pyspark import SparkContext
from pyspark.sql import HiveContext
from pyspark.sql import functions as F

sc = SparkContext("local")

sqlContext = HiveContext(sc)

df = sqlContext.createDataFrame([
    ("a", None, None),
    ("a", "code1", None),
    ("a", "code2", "name2"),
], ["id", "code", "name"])

df.show()

+---+-----+-----+
| id| code| name|
+---+-----+-----+
|  a| null| null|
|  a|code1| null|
|  a|code2|name2|
+---+-----+-----+

请注意,上面你必须创建一个HiveContext。有关处理不同Spark版本的信息,请参阅https://stackoverflow.com/a/35529093/690430

(df
  .groupby("id")
  .agg(F.collect_set("code"),
       F.collect_list("name"))
  .show())

+---+-----------------+------------------+
| id|collect_set(code)|collect_list(name)|
+---+-----------------+------------------+
|  a|   [code1, code2]|           [name2]|
+---+-----------------+------------------+

答案 1 :(得分:0)

如果数据帧较大,则可以尝试使用pandas udf(GROUPED_AGG)以避免内存错误。它也快得多。

  

分组聚合熊猫UDF与Spark聚合函数相似。分组的聚合熊猫UDF与groupBy()。agg()和pyspark.sql.Window一起使用。它定义了一个或多个pandas.Series到标量值的聚合,其中每个pandas.Series代表组或窗口中的一列。 pandas udf

示例:

import pyspark.sql.functions as F

@F.pandas_udf('string', F.PandasUDFType.GROUPED_AGG)
def collect_list(name):
    return ', '.join(name)

grouped_df = df.groupby('id').agg(collect_list(df["name"]).alias('names'))