pyspark - pyspark udf 的参数数量可变

Question

我有大约 275 列，我想在 25 列中搜索正则表达式字符串 "^D(410|412)。如果此搜索字符串出现在 25 列中的任何一列中，我想将 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"]))

Answer 1

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

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|
+---+---+---+---+---+---+---+---+---+---+-----------+

希望这能解决您的问题