代码之家 › 专栏 › 技术社区 › Steven

在UDF中使用函数dict

user-defined-functions pyspark

0

Steven · 技术社区 · 8 年前

在pySpark中,我有一个数据帧 df 具体如下:

Site    A   B
1       3   83
1       16  26
1       98  46
1       80  14
1       83  54
2       0   83
2       75  67
2       72  24
2       60  13
6       40  50
6       34  60
6       36  39
6       68  6
6       91  51
6       81  82

另一方面,我有一个措辞 myDict 函数的数量 g = {1 : f1, 2 : f2, 3: f3, 6:f6}

我想用字典生成一个新的专栏。类似于: df.withColumn("MyCol", myDict[df.Site](df.A, df.B))

但当我这样做时,我收到了错误:

不可损坏类型:“列”

回溯(最近一次呼叫最后一次):

我该怎么写?

1 回复 | 直到 6 年前

1

TMichel 8 年前

您想要使用 Currying .

这个 withColumn 函数通过 lit() 功能( lit 实际返回一列)。

为了传入额外参数,必须使用 higher-order function 返回一个 udf :

class MyUDFs():
    @staticmethod
    def trans(myDict):
        def cb(Site,A,B):
            return myDict[Site](A, B)
        return udf(cb, StringType())

df = df.withColumn("MyCol",MyUDFs.trans(myDict)(df["Site"],df["A"],df["B"]))