如何索引 arrays 数组中的每个元素？

Question

我有一个 ArrayType 列，其中每个元素也是一个恰好包含 2 个元素的数组。

from pyspark.sql import SparkSession

data = [
  {"u": ["apple", 23]},
  {"u": ["banana", 12]}
]

spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame(data)

df.show()

+------------+
|           u|
+------------+
| [apple, 23]|
|[banana, 12]|
+------------+

我想用它的第一个元素替换每个内部数组。 如果我一直在写香草 Python，它将是：

result = [ar[0] for ar in array_of_arrays]

借助 Spark，我可以使用 UDF：

from pyspark.sql import functions as f, types as t

fn = f.udf(lambda u: u[0], t.StringType())
new_df = df.select(fn(f.col("u")))

new_df.show()

+-----------+
|<lambda>(u)|
+-----------+
|      apple|
|     banana|
+-----------+

这是我想要的 output。 但是我怎么能用 PySpark 做到这一点，没有：

• 使用 UDF
• 展开数组并在之后重新聚合

Answer 1

df.withColumn('u_1', col('u')[0]).show()

+------------+------+
|           u|   u_1|
+------------+------+
| [apple, 23]| apple|
|[banana, 12]|banana|
+------------+------+