pyspark udf

时间:2017-03-21 18:40:40

标签: pyspark udf

我有大约275个列,我想在25列中搜索正则表达式字符串"^D(410|412)。如果此搜索字符串出现在任何2 5列中,我想将true添加到MyNewColumn

使用下面我可以为2列执行此操作。无论如何传递可变数量的列?

以下代码适用于2列

def moreThanTwoArgs(col1,col2): 
return bool((re.search("^D(410|412)",col1) or re.search("^D(410|412)",col2)))

twoUDF= udf(moreThanTwoArgs,BooleanType())
df = df.withColumn("MyNewColumn", twoUDF(df["X1"], df["X2"]))

1 个答案:

答案 0 :(得分:4)

我尝试了一些类似的示例代码试试这个并继续: -

df1 = sc.parallelize(
        [
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], 
    ]).toDF(['c1', 'c2', 'c3', 'c4', 'c5', 'c6', 'c7', 'c8', 'c9', 'c10'])
df1.show()
+---+---+---+---+---+---+---+---+---+---+
| c1| c2| c3| c4| c5| c6| c7| c8| c9|c10|
+---+---+---+---+---+---+---+---+---+---+
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|
+---+---+---+---+---+---+---+---+---+---+


import pyspark.sql.functions as F
import pyspark.sql.types as T
import re

def booleanFindFunc(*args):
    return sum(args)

udfBoolean = F.udf(booleanFindFunc, T.StringType())


#Below is Sum of three columns (c1+c2+c2)
df1.withColumn("MyNewColumn", booleanFindFunc(F.col("c1"), F.col("c2"), F.col("c2"))).show()
+---+---+---+---+---+---+---+---+---+---+-----------+
| c1| c2| c3| c4| c5| c6| c7| c8| c9|c10|MyNewColumn|
+---+---+---+---+---+---+---+---+---+---+-----------+
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|          5|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|          5|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|          5|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|          5|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|          5|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|          5|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|          5|
+---+---+---+---+---+---+---+---+---+---+-----------+



#Below is Sum of All Columns (c1+c2+c3---+c10)
df1.withColumn("MyNewColumn", booleanFindFunc(*[F.col(i) for i in df1.columns])).show()

+---+---+---+---+---+---+---+---+---+---+-----------+
| c1| c2| c3| c4| c5| c6| c7| c8| c9|c10|MyNewColumn|
+---+---+---+---+---+---+---+---+---+---+-----------+
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         55|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         55|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         55|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         55|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         55|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         55|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         55|
+---+---+---+---+---+---+---+---+---+---+-----------+

#Below is Sum of All odd Columns (c1+c3+c5--+c9)
df1.withColumn("MyNewColumn", booleanFindFunc(*[F.col(i) for i in df1.columns if int(i[1:])%2])).show()

+---+---+---+---+---+---+---+---+---+---+-----------+
| c1| c2| c3| c4| c5| c6| c7| c8| c9|c10|MyNewColumn|
+---+---+---+---+---+---+---+---+---+---+-----------+
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         25|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         25|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         25|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         25|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         25|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         25|
|  1|  2|  3|  4|  5|  6|  7|  8|  9| 10|         25|
+---+---+---+---+---+---+---+---+---+---+-----------+

希望这将解决您的问题