python - Pyspark 爆炸列表创建列表中带有索引的列

Question

所以我有一个关于 pyspark 的问题。我有一个如下所示的数据框:

+---+------------+
| id|        list|
+---+------------+
|  2|[3, 5, 4, 2]|
+---+------------+
|  3|[4, 5, 3, 2]|
+---+------------+

我想将其分解为多行，并在单独的列中保留有关列表中每个元素的位置的信息。结果应该如下所示:

+---+------------+------------+
| id|    listitem|        rank|
+---+------------+------------+
|  2|           3|           1|
+---+------------+------------+
|  2|           5|           2|
+---+------------+------------+
|  2|           4|           3|
+---+------------+------------+
|  2|           2|           4|
+---+------------+------------+
|  3|           4|           1|
+---+------------+------------+
|  3|           5|           2|
+---+------------+------------+
|  3|           3|           3|
+---+------------+------------+
|  3|           2|           4|
+---+------------+------------+

排名列具有每个元素在列表中的位置的索引+1。关于实现它的最佳代码有什么建议吗？

Answer 1

您可以使用posexplode()或posexplode_outer()函数以获得所需的结果。

df = spark.createDataFrame([(2, [3, 5, 4, 2]), (3, [4, 5, 3, 2])], ["id", "list"])

df.select('id',posexplode_outer('list').alias('rank', 'listitem')) \
.withColumn('rank', col('rank') + 1).show()

+---+----+--------+
| id|rank|listitem|
+---+----+--------+
|  2|   1|       3|
|  2|   2|       5|
|  2|   3|       4|
|  2|   4|       2|
|  3|   1|       4|
|  3|   2|       5|
|  3|   3|       3|
|  3|   4|       2|
+---+----+--------+