PySpark - 根据另一列引用的列名创建列

Question

PySpark - 根据另一列引用的列名创建列

3

我对Spark非常陌生，无法使其正常工作... 希望有一种简单的方法可以做到这一点... 我尝试的最佳方式是通过下表进行描述：（我需要获取“所需”列）

    colA    colB    colC    ref     required
1   a1        b1    c1      colA     a1
2   a2        b2    c2      colA     a2
3   a3        b3    c3      colB     b3
4   a4        b4    c4      colB     b4
5   a5        b5    c5      colC     c5
6   a6        b6    c6      colC     c6

上面只是一个示例 - 在实际的例子中我有超过50列，因此使用条件语句并不可行...

我知道这可以很容易地在pandas中完成，例如：

df['required'] = df.apply(lambda x: x.loc[x.ref], axis=1)

或者

df['required'] = df.lookup(df.index, df.ref)

有什么建议可以在PySpark中完成这个任务吗？

- morienor

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Psidom · Accepted Answer

一种方法是使用when和coalesce函数：

import pyspark.sql.functions as F

cols = ['colA', 'colB', 'colC']
df.withColumn('required', F.coalesce(*[F.when(df.ref == c, df[c]) for c in cols])).show()
+----+----+----+----+--------+
|colA|colB|colC| ref|required|
+----+----+----+----+--------+
|  a1|  b1|  c1|colA|      a1|
|  a2|  b2|  c2|colA|      a2|
|  a3|  b3|  c3|colB|      b3|
|  a4|  b4|  c4|colB|      b4|
|  a5|  b5|  c5|colC|      c5|
|  a6|  b6|  c6|colC|      c6|
+----+----+----+----+--------+

基本上，您需要检查ref列的名称与哪一列相等，并从该列中取值--F.when(df.ref == c, df[c])；这将创建一个列对象列表，其值在其名称出现在ref列时保留，否则其值为NULL；然后通过合并列对象列表，将NULL值填充为有效列值。