如何从以不同列为条件的PySpark数据框中提取数组元素？

Question

我有以下 PySpark输入数据框：

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

位置：

• 索引：输入Double
• 值列表：输入 Vector 。 （它是 NOT阵列）

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

中获得以下输出数据框

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

逻辑：

for each row:
  value = valuelist[index] 

Answer 1

Spark 1.5及更高版本

您可以使用pyspark.sql.functions.expr至pass a column value as an input to a function：

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：

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

Answer 2

您可以创建一个新列，并将这两列作为输入传递。

from pyspark.sql import functions as F
columns = ['index', 'valuelist']
vals = [
     (0, [1,2]),
     (1, [1,2])
]

df = sqlContext.createDataFrame(vals, columns)
df = df.withColumn(
"value", udf(lambda index_and_list: index_and_list[0][index_and_list[1]], IntegerType())(
    F.struct(F.col("valuelist"), F.col("index")))
    )

获得以下输出：

> +-----+---------+-----+
|index|valuelist|value|
+-----+---------+-----+
|    0|   [1, 2]|    1|
|    1|   [1, 2]|    2|
+-----+---------+-----+