Question

我有一个火花数据框，其行为-

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]|
#+---+---------+---------+

Answer 2

我的pyspark技能变得生锈（我承认如今我不再磨练它们了），或者这确实是一个棘手的问题……我设法做到这一点的唯一方法是使用SQL语句：

user

对于以后的问题，最好遵循How to make good reproducible Apache Spark Dataframe examples上的建议准则。

从pyspark的dataframe ArrayType列中获取前N个元素

2 个答案: