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

Question

我是 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

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

Answer 1

当您在数组的索引和原始 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

注意，这里也是对的很好的参考答案。

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

Create a column in a PySpark dataframe using a list whose indices are present in one column of the dataframe

python

arrays

pyspark

spark-dataframe

pyspark-sql