python之使用其索引存在于数据帧的一列中的列表在 PySpark 数据帧中创建一列

我是 Python 和 PySpark 的新手。我在 PySpark 中有一个数据框，如下所示:

## +---+---+------+ 
## | x1| x2|   x3 | 
## +---+---+------+ 
## |  0| a |  13.0| 
## |  2| B | -33.0| 
## |  1| B | -63.0| 
## +---+---+------+ 

我有一个数组: arr = [10, 12, 13]

我想在数据框中创建一个 x4 列，这样它应该具有基于 x1 的值作为索引的列表中的相应值。最终数据集应如下所示:

## +---+---+------+-----+ 
## | x1| x2|   x3 |  x4 | 
## +---+---+------+-----+ 
## |  0| a |  13.0| 10  | 
## |  2| B | -33.0| 13  | 
## |  1| B | -63.0| 12  | 
## +---+---+------+-----+ 

我试过使用下面的代码来实现:

df.withColumn("x4", lit(arr[col('x1')])).show()

但是，我收到一个错误:

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices 

有什么方法可以有效地实现这一点？

请您参考如下方法:

当您在数组的索引和原始 DataFrame 之间进行连接时，一种方法是将数组转换为 DataFrame，生成 rownumber()-1(变成你的索引)，然后将两个 DataFrame 连接在一起。

from pyspark.sql import Row 
 
# Create original DataFrame `df` 
df = sqlContext.createDataFrame( 
    [(0, "a", 13.0), (2, "B", -33.0), (1, "B", -63.0)], ("x1", "x2", "x3")) 
df.createOrReplaceTempView("df") 
 
# Create column "x4" 
row = Row("x4") 
 
# Take the array 
arr = [10, 12, 13] 
 
# Convert Array to RDD, and then create DataFrame 
rdd = sc.parallelize(arr) 
df2 = rdd.map(row).toDF() 
df2.createOrReplaceTempView("df2") 
 
# Create indices via row number 
df3 = spark.sql("SELECT (row_number() OVER (ORDER by x4))-1 as indices, * FROM df2") 
df3.createOrReplaceTempView("df3") 

现在您有了两个 DataFrame:df 和 df3，您可以运行下面的 SQL 查询将两个 DataFrame 连接在一起。

select a.x1, a.x2, a.x3, b.x4 from df a join df3 b on b.indices = a.x1 

请注意，这里也是对 adding columns to DataFrames 的很好的引用答案.