PySpark：用一列来索引另一列（两列的udf？）-Java 学习之路

(Edited Feb 14th)

假设我有一个带有以下模式的Spark（PySpark）数据帧：

root
 |-- myarray: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- myindices: array (nullable = true)
 |    |-- element: integer (containsNull = true)

看起来像：

+--------------------+----------+
|          myarray   | myindices|
+--------------------+----------+
|                 [A]|    [0]   |
|              [B, C]|    [1]   |
|        [D, E, F, G]|   [0,2]  |
+--------------------+----------+

如何使用第二个数组来索引第一个？

我的目标是创建一个新的数据框，如下所示：

+--------------------+----------+------+
|          myarray   | myindices|result|
+--------------------+----------+------+
|                 [A]|    [0]   |  [A] |
|              [B, C]|    [1]   |  [C] |
|        [D, E, F, G]|   [0,2]  | [D,F]|
+--------------------+----------+------+

（可以安全地假设 myindices 的内容始终保证在所讨论的行的 myarray 的基数范围内，因此没有越界问题 . ）

似乎 .getItem() 方法仅适用于单个参数，因此我可能需要一个UDF，但我知道无法创建一个具有多个列作为输入的UDF . 有或没有UDF的任何解决方案？

1 回答

from pyspark.sql import functions as f

rdd = spark.sparkContext.parallelize([(['A'], [0]), (['B', 'C'], [1]), (['D', 'E', 'F'], [0, 2])])
df = spark.createDataFrame(rdd, ['myarray', 'myindices'])
my_UDF = f.UserDefinedFunction(lambda x, y: map(lambda z: x[z], y), returnType=ArrayType(StringType()))
res = df.withColumn('result', my_UDF(df['myarray'], df['myindices']))
res.show(truncate=False)

output:
+---------+---------+------+
|myarray  |myindices|result|
+---------+---------+------+
|[A]      |[0]      |[A]   |
|[B, C]   |[1]      |[C]   |
|[D, E, F]|[0, 2]   |[D, F]|
+---------+---------+------+

回复于 2024-04-30T05:28:12+08:00

PySpark：用一列来索引另一列（两列的udf？）

1 回答

相关问题