PySpark 将索引列传递给 element_at()

Question

您可以将带有索引值的 pyspark 列传递给 element_at 函数吗？

我有一个类似于以下内容的 df：

------------------
 foo   |   bar
------------------
  1    |   [3, 2, 1]
  2    |   [3, 3, 1]
  1    |   [8, 2, 1]

我正在尝试使用 foo 作为索引来创建以下内容：

 baz  |
--------
  3   |
  3   |
  8   |

我正在尝试使用以下内容但没有成功：

df.withColumn('baz', F.element_at('bar', F.col('foo'))

Answer 1

此代码将为您提供所需的结果：

df.withColumn("baz", F.col("bar")[F.col("foo") - F.lit(1)]).show()

+---+---------+---+
|foo|      bar|baz|
+---+---------+---+
|  1|[3, 2, 1]|  3|
|  2|[3, 3, 1]|  3|
|  1|[8, 2, 1]|  8|
+---+---------+---+

我不得不从中减去一个，foo因为看起来您正在使用 1-indexing。PySpark 数组是零索引的。

这是更标准的解决方案。

+---+---------+---+
|foo|      bar|baz|
+---+---------+---+
|  1|[3, 2, 1]|  3|
|  2|[3, 3, 1]|  3|
|  1|[8, 2, 1]|  8|
+---+---------+---+

+---+---------+---+
|foo|      bar|baz|
+---+---------+---+
|  1|[3, 2, 1]|  2|
|  2|[3, 3, 1]|  1|
|  1|[8, 2, 1]|  2|
+---+---------+---+

这是制作示例数据的代码：

df.withColumn("baz", F.col("bar")[F.col("foo")]).show()