处理scala中spark数据集中的行组

Question

我正在寻找一种方法，将我的大型spark数据集划分为组/批，并在某些函数中处理这组行。所以基本上，这组行应该被输入到我的函数中，输出是我的单位，因为我不想聚合或更新输入记录，只是执行一些计算。

为了理解，假设我有以下输入。

假设我需要按col1和col2分组，这将给我以下分组

（1， A，1），（1， A，4），（1， A，5）---

（1，B，2）---

（1，C，3），（1，C，6）---

（2，X，7），（2，X，8）---

所以我想把这些组传递给我的函数来执行一些逻辑。现在，假设我在那个方法中求和Col3。（这不是我的要求，但是让我们假设我想在我单独的方法中求和）。生成以下o/p。

如何实现这一点，基于一些建议，我试图研究UDAF，但找不到使用它的方法。请注意，我的真实输入数据集有超过5亿条记录。谢谢

Answer 1

下面是一个基于您的输入的简单示例，让您开始：

    from pyspark.sql.types import IntegerType
    from pyspark.sql import SparkSession
    from pyspark.sql import functions as F
    
    spark = SparkSession.builder.getOrCreate()
    
    data = [
        (1, "A",  1),
        (1, "B",  2),
        (1, "C",  3),
        (1, "A",  4),
        (1, "A",  5),
        (1, "C",  6),
        (2, "X",  7),
        (2, "X",  8),
    ]
    
    df = spark.createDataFrame(data, ["col1", "col2", "col3"])
    df.show()
    
    +----+----+----+
    |col1|col2|col3|
    +----+----+----+
    |   1|   A|   1|
    |   1|   B|   2|
    |   1|   C|   3|
    |   1|   A|   4|
    |   1|   A|   5|
    |   1|   C|   6|
    |   2|   X|   7|
    |   2|   X|   8|
    +----+----+----+
    
    # define your function - pure Python here, no Spark needed
    def dummy_f(xs):
      return sum(xs)
    
    
    # apply your function as UDF - needs input function and return type (integer here)
    (
      df
      .groupBy(F.col("col1"), F.col("col2"))
      .agg(F.collect_list(F.col("col3").cast("int")).alias("col3"))
      .withColumn("col3sum", F.udf(dummy_f, IntegerType())(F.col("col3")))
    ).show()
    
    +----+----+---------+-------+
    |col1|col2|     col3|col3sum|
    +----+----+---------+-------+
    |   1|   A|[1, 4, 5]|     10|
    |   1|   B|      [2]|      2|
    |   1|   C|   [3, 6]|      9|
    |   2|   X|   [7, 8]|     15|
    +----+----+---------+-------+

根据输入函数的需要聚合列是关键。您可以使用create_map创建dict或collect_list，如下所示。