是否可以从 PySpark 中的数组中获取任意数量的元素？

Question

我的数据框有两个数组列。我想从第一列中获取索引位于第二列中的元素。例如，我有以下数据集

df = spark.createDataFrame(
   [
      {
         'text': ['0', '1', '2', '3', '4', '5']
         'indices': [0, 2, 4],
      },
   ]
)

所以我想要值为“['0', '2', '4]”的列。

不写UDF可以实现吗？

Answer 1

您可以尝试使用

expr

函数与

TRANSFORM

和

element_at

根据第二个数组中提供的索引从第一个数组中选择元素。

例如：

from pyspark.sql import SparkSession
from pyspark.sql.functions import expr

df = df.withColumn(
    "selected_text",
    expr("TRANSFORM(indices, i -> element_at(text, i + 1))")
)
df.show()