Apache spark 如何让Pyspark在两个级别聚合集合？

我需要通过将每个组中特定列中的值收集到一个集合中来聚合数据帧中的行

pyspark.sql.functions.collect\u set

完全满足我的需要

但是，我需要依次对两列执行此操作，因为我需要按一列对输入进行分组，按另一列将每个组划分为子组，并对每个子组进行聚合。我不知道如何获取

collect\u set

为每个组创建一个集合

例如：

df = spark.createDataFrame([('a', 'x', 11, 22), ('a', 'y', 33, 44), ('b', 'x', 55, 66), ('b', 'y', 77, 88),('a','x',12,23),('a','y',34,45),('b','x',56,67),('b','y',78,89)], ('col1', 'col2', 'col3', 'col4'))
df.show()
+----+----+----+----+
|col1|col2|col3|col4|
+----+----+----+----+
|   a|   x|  11|  22|
|   a|   y|  33|  44|
|   b|   x|  55|  66|
|   b|   y|  77|  88|
|   a|   x|  12|  23|
|   a|   y|  34|  45|
|   b|   x|  56|  67|
|   b|   y|  78|  89|
+----+----+----+----+

g1 = df.groupBy('col1', 'col2').agg(collect_set('col3'),collect_set('col4'))
g1.show()
+----+----+-----------------+-----------------+
|col1|col2|collect_set(col3)|collect_set(col4)|
+----+----+-----------------+-----------------+
|   a|   x|         [12, 11]|         [22, 23]|
|   b|   y|         [78, 77]|         [88, 89]|
|   a|   y|         [33, 34]|         [45, 44]|
|   b|   x|         [56, 55]|         [66, 67]|
+----+----+-----------------+-----------------+

g2 = g1.groupBy('col1').agg(collect_set('collect_set(col3)'),collect_set('collect_set(col4)'),count('col2'))
g2.show(truncate=False)
+----+--------------------------------------------+--------------------------------------------+-----------+
|col1|collect_set(collect_set(col3))              |collect_set(collect_set(col4))              |count(col2)|
+----+--------------------------------------------+--------------------------------------------+-----------+
|b   |[WrappedArray(56, 55), WrappedArray(78, 77)]|[WrappedArray(66, 67), WrappedArray(88, 89)]|2          |
|a   |[WrappedArray(33, 34), WrappedArray(12, 11)]|[WrappedArray(22, 23), WrappedArray(45, 44)]|2          |
+----+--------------------------+--------------------------------------------+-----------+

我希望结果看起来更像

+----+----------------+----------------+-----------+
|col1|   ...col3...   |   ...col4...   |count(col2)|
+----+----------------+----------------+-----------+
|b   |[56, 55, 78, 77]|[66, 67, 88, 89]|2          |
|a   |[33, 34, 12, 11]|[22, 23, 45, 44]|2          |
+----+----------------+----------------+-----------+

但是，我没有看到一个聚合函数来实现两个或多个集合的并集，也没有看到一个pyspark操作来扁平显示在

g2

中的“数组数组”结构

---

pyspark是否提供了一种简单的方法来实现这一点？或者我应该采取完全不同的方法吗？

您可以在之后使用自定义项展平列：

flatten = udf(lambda l: [x for i in l for x in i], ArrayType(IntegerType()))

我冒昧地将

g2

的列重命名为

col3

和

col4

，以保存键入。这使得：

g3 = g2.withColumn('col3flat', flatten('col3'))

>>> g3.show()
+----+--------------------+--------------------+-----+----------------+
|col1|                col3|                col4|count|        col3flat|
+----+--------------------+--------------------+-----+----------------+
|   b|[[78, 77], [56, 55]]|[[66, 67], [88, 89]]|    2|[78, 77, 56, 55]|
|   a|[[12, 11], [33, 34]]|[[22, 23], [45, 44]]|    2|[12, 11, 33, 34]|
+----+--------------------+--------------------+-----+----------------+

---

在PySpark 2.4.5中，您可以使用现在内置的。

您可以通过

from pyspark.sql.functions import collect_set, countDistinct
(
    df.
        groupby('col1').
        agg(
            collect_set('col3').alias('col3_vals'),
            collect_set('col4').alias('col4_vals'), 
            countDistinct('col2').alias('num_grps')
        ).
        show(truncate=False)
)
+----+----------------+----------------+--------+                               
|col1|col3_vals       |col4_vals       |num_grps|
+----+----------------+----------------+--------+
|b   |[78, 56, 55, 77]|[66, 88, 67, 89]|2       |
|a   |[33, 12, 34, 11]|[45, 22, 44, 23]|2       |
+----+----------------+----------------+--------+

---

你能写一个udf来在事实发生后变平吗？@hoyland这是一个想法，我正在研究它……它看起来有点野蛮。我想知道pyspark是否已经提供了一些东西来处理这种情况。但是这个很好用——谢谢。我想没有现成的东西。在Scala中，Spark为您操作/聚合组提供了更大的灵活性（例如，您可以编写自己的聚合器），但我认为Pyspark中没有其他东西。