如何在pyspark数据框中创建包含列表值总和的列

Question

我是 pyspark 的新手，我正在处理一个复杂的数据框。经过一些过滤后，我试图将 N 行从 list 放入我的 df.column。

我有以下 df.struct：

root
 |-- struct1: struct (nullable = true)
 |    |-- array1: array (nullable = true)
 |    |    |-- element: struct (containsNull = true)
 |    |    |    |-- struct2 : struct (nullable = true)
 |    |    |    |    |-- date: string (nullable = true)
 |    |    |    |    |-- value: string (nullable = true)
 |    |    |    |-- struct3 : struct (nullable = true)
 |    |    |    |    |-- date: string (nullable = true)
 |    |    |    |    |-- value: string (nullable = true)
 |    |    |    |-- property: string (nullable = true)

我想要实现的是当属性为 Good 时得到 所有 struct2.values 的总和。因为我可以为 array1 设置多个 (N) 个值。

现在，我得到了第一个属性的一个小句子。但我不能以成功的方式将它传递给 udf 以迭代所有可能的行： df.withColumn("Sum", (col('struct1.array1')[0])['property'])

我想到的一些步骤是：

• 当 property=Good 时过滤列表中的每个元素

• 在 udf 中返回一个带有 struct3.value 总和的 lambda 值

期望的输出应该是这样的：

None
+---------------------------------------------------------------------------------------------------------+
|Struct1                                                                                            |Sum|
+---------------------------------------------------------------------------------------------------------+
|[[[[2020-01-01, 10], [2020-02-02, 15], Good], [[2020-01-01, 20], [2020-02-02, 25], Good]]]         |20|
+---------------------------------------------------------------------------------------------------------+

任何帮助将不胜感激

Answer 1

在这种情况下，您不一定需要 UDF。当使用 Spark >= 2.4.0 时，您可以通过使用内置的 high-order 函数来实现相同的功能，如下所示：

from pyspark.sql.functions import expr  

df.withColumn("good_elements", expr("""transform( \
                                         filter(struct1.array1, e -> e.property == 'Good'), 
                                         e -> cast(e.struct2.value as int)
                                    )""")) \
  .withColumn("sum", expr("aggregate(good_elements, 0, (sum, e) -> sum + e)"))

• filter(struct1.array1, e -> e.property == 'Good')：首先我们过滤具有property == 'Good'的项目

• transform(..., e -> cast(e.struct2.value as int)：接下来我们将每个项目转换为整数并将它们存储到名为good_elements的新列中

• aggregate(good_elements, 0, (sum, e) -> sum + e)：最后我们通过计算good_elements的和来创建列sum