如何从以不同列为条件的 PySpark 数据帧中提取数组元素? [英] How to extract array element from PySpark dataframe conditioned on different column?

问题描述

我有以下 PySpark 输入数据框:

+-------+------------+
| index | valuelist  |
+-------+------------+
| 1.0   | [10,20,30] |
| 2.0   | [11,21,31] |
| 0.0   | [14,12,15] |
+-------+------------+

地点:

• 索引:类型为 Double
• 值列表:输入矢量.(它是非数组)

从上面的输入数据帧中，我想在 PySpark

From the above Input Dataframe, I want to get the following Output Dataframe in PySpark

+-------+-------+
| index | value |
+-------+-------+
| 1.0   | 20    |
| 2.0   | 31    |
| 0.0   | 14    |
+-------+-------+

逻辑:

for each row:
  value = valuelist[index] 

推荐答案

Spark 1.5 及更高版本

您可以使用 pyspark.sql.functions.expr 来将列值作为输入传递给函数:

df.select("index", f.expr("valuelist[CAST(index AS integer)]").alias("value")).show()
#+-----+-----+
#|index|value|
#+-----+-----+
#|  1.0|   20|
#|  2.0|   31|
#|  0.0|   14|
#+-----+-----+

Spark 2.1 及更高版本

如果您有 spark 2.1 或更高版本，这里有一个使用 pyspark.sql.functions.posexplode 的替代方法:

Spark version 2.1 and higher

If you have spark version 2.1 or higher, here's an alternative using pyspark.sql.functions.posexplode:

import pyspark.sql.functions as f

df.select("index", f.posexplode("valuelist").alias("pos", "value"))\
    .where(f.col("index").cast("int") == f.col("pos"))\
    .select("index", "value")\
    .show()
#+-----+-----+
#|index|value|
#+-----+-----+
#|  1.0|   20|
#|  2.0|   31|
#|  0.0|   14|
#+-----+-----+

这篇关于如何从以不同列为条件的 PySpark 数据帧中提取数组元素?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！