PySpark UDF 到多列

Question

所以我有一个这样的 UDF：

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

我通常只是像这样传递 UDF：

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

我想知道是否有任何方法可以做到这一点，但不必一一列出。

Answer 1

使用理解：

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
])