Apache spark 创建包含Spark dataframe字段中数组中每个结构的第一个元素的数组

如何在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|[]                        |
+-------+--------------------------+

上面定义的模式和数据帧如下所示：

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-wise标记包含在类型为

（Subject，marks）

的有序结构中。每个学生的科目数量不是恒定的，可能是零

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

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

---

请注意，行数与之前相同；因此，除非我随后重新组合，否则我不能使用explode进行此操作，这在计算上似乎效率低下。

您最好使用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"))

---

以下是带有explode、group by和aggregate的版本，仅供参考

import pyspark.sql.functions as f

scoresheet.select('name').join(
    scoresheet
    .withColumn('score', f.explode('scores'))
    .withColumn('subject', f.col('score').getField('_1'))
    .groupBy('name')
    .agg(f.collect_list('subject').alias('subjects'))
    , on='name'
    , how='left'
)

由于这完全是在PySpark中实现的，因此如果在某些情况下比UDF版本快，我不会感到惊讶，但我没有做任何评测。左连接是为了确保没有结果的学生不会在最终结果中被丢弃