PySpark UDF到多个列

时间:2018-01-03 22:07:41

标签: apache-spark pyspark

所以我有一个像这样的UDF:

 tudf = udf(lambda value: 1 if value>=1 else 0,IntegerType())

我通常只是传递这样的UDF:

 df = fdf.withColumn('COLUMN1',tudf(df.COLUMN1))

我想知道是否有任何方法可以做到这一点但是有多个列而不必一个接一个地去。

1 个答案:

答案 0 :(得分:3)

使用理解:

fdf.select([
  tudf(c).alias(c) if c in cols_to_transform else c for c in fdf.columns
])

虽然这里不推荐udf

from pyspark.sql.functions import when, col

fdf.select([
  when(col(c) >= 1, 1).otherwise(0).alias(c) if c in cols_to_transform else c 
  for c in fdf.columns
])