apache-spark - 创建包含 Spark 数据帧字段中数组中每个结构的第一个元素的数组

Question

如何在 PySpark 数据帧中从结构数组到每个结构的第一个元素的数组？

一个例子将更清楚地说明这一点。假设我有如下定义的数据框:

scoresheet = spark.createDataFrame([("Alice", [("Math",100),("English",80)]),("Bob", [("Math", 90)]),("Charlie", [])],["name","scores"])

上面定义的模式和数据框如下所示:

root
 |-- name: string (nullable = true)
 |-- scores: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- _1: string (nullable = true)
 |    |    |-- _2: long (nullable = true)

+-------+--------------------------+
|name   |scores                    |
+-------+--------------------------+
|Alice  |[[Math,100], [English,80]]|
|Bob    |[[Math,90]]               |
|Charlie|[]                        |
+-------+--------------------------+

您可以看到主题标记包含在 (Subject,Marks) 类型的有序结构中。对于每个学生。每个学生的科目数量不是恒定的，可能为零。

我想从这个开始生成一个新的数据框，它只包含每个学生数组中的主题，没有分数。它应该为没有科目的学生生成一个空数组。简而言之，它应该是这样的:

+-------+---------------+
|name   |scores         |
+-------+---------------+
|Alice  |[Math, English]|
|Bob    |[Math]         |
|Charlie|[]             |
+-------+---------------+

注意行数和之前一样；所以我不能为此使用爆炸，除非我之后重新组合，这在计算上似乎效率低下。

Answer 1

你能做的最好的是udf:

from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, StringType

take_first = udf(lambda rows: [row[0] for row in rows], ArrayType(StringType()))

scoresheet.withColumn("scores", take_first("scores"))