apache-spark - 从pyspark中的数据框ArrayType列中获取前N个元素

Question

我有一个带有行的 Spark 数据框 -

1   |   [a, b, c]
2   |   [d, e, f]
3   |   [g, h, i]

现在我只想保留数组列中的前 2 个元素。

1   |   [a, b]
2   |   [d, e]
3   |   [g, h]

如何实现？

注意 - 请记住，我不是在这里提取单个数组元素，而是提取可能包含多个元素的数组的一部分。

Answer 1

以下是如何使用 API 函数执行此操作。

假设您的 DataFrame 如下:

df.show()
#+---+---------+
#| id|  letters|
#+---+---------+
#|  1|[a, b, c]|
#|  2|[d, e, f]|
#|  3|[g, h, i]|
#+---+---------+

df.printSchema()
#root
# |-- id: long (nullable = true)
# |-- letters: array (nullable = true)
# |    |-- element: string (containsNull = true)

您可以使用方括号访问 letters 中的元素按索引列，并将其包装在对 pyspark.sql.functions.array() 的调用中新建ArrayType柱子。

import pyspark.sql.functions as f

df.withColumn("first_two", f.array([f.col("letters")[0], f.col("letters")[1]])).show()
#+---+---------+---------+
#| id|  letters|first_two|
#+---+---------+---------+
#|  1|[a, b, c]|   [a, b]|
#|  2|[d, e, f]|   [d, e]|
#|  3|[g, h, i]|   [g, h]|
#+---+---------+---------+

或者，如果您有太多索引要列出，您可以使用列表推导式:

df.withColumn("first_two", f.array([f.col("letters")[i] for i in range(2)])).show()
#+---+---------+---------+
#| id|  letters|first_two|
#+---+---------+---------+
#|  1|[a, b, c]|   [a, b]|
#|  2|[d, e, f]|   [d, e]|
#|  3|[g, h, i]|   [g, h]|
#+---+---------+---------+

对于 pyspark 2.4+ 版本，您还可以使用 pyspark.sql.functions.slice() :

df.withColumn("first_two",f.slice("letters",start=1,length=2)).show()
#+---+---------+---------+
#| id|  letters|first_two|
#+---+---------+---------+
#|  1|[a, b, c]|   [a, b]|
#|  2|[d, e, f]|   [d, e]|
#|  3|[g, h, i]|   [g, h]|
#+---+---------+---------+

slice对于大型数组可能有更好的性能(注意起始索引是 1，而不是 0)