在 for 循环中使用 udf 在 Pyspark 中创建多个列

Posted 2023-03-31

【中文标题】在 for 循环中使用 udf 在 Pyspark 中创建多个列

【英文标题】：use udf inside for loop to create multiple columns in Pyspark

【发布时间】：2020-09-02 07:57:16

【问题描述】：

我有一个带有一些列（col1、col2、col3、col4、col5...直到 32）的 spark 数据框，现在我创建了一个函数 (udf)，它接受 2 个输入参数并返回一些浮点值。

现在我想使用上述函数创建新列（按递增顺序，如 col33、col32、col33、col34..），其中一个参数增加，其他参数不变

def fun(col1,col2):
   if true:
      do someting
   else:
      do someting

我已将此函数转换为udf

udf_func = udf(fun,Floatype())

现在我想用这个函数在数据框中创建新列怎么做？

我试过了

for i in range(1,5):
   BS.withColumns("some_name with increasing number like abc_1,abc_2",udf_func(col1<this should be col1,col2..till 4>,col6<this is fixed>

如何在 PySpark 中实现这一点？

【问题讨论】：

你能举一个你开始使用的DataFrame的例子和预期的结果吗？

@kfkhalili 我已经添加了数据帧示例，我想在其中创建新列，就像我使用我创建的函数在第二个数据帧中显示的那样，也在参数之一的函数内将是来自 (col1-col5 和 fn 的第二个参数将只有 col5) 的列

我不确定我是否理解您的用例，但也许答案可以帮助您。

【参考方案1】：

您一次只能使用withColumn 创建一列，因此我们必须调用它多次。

# We set up the problem
columns = ["col1", "col2", "col3"]
data = [(1, 2, 3), (4, 5, 6), (7, 8, 9)]
rdd = spark.sparkContext.parallelize(data)
df = rdd.toDF(columns)

df.show()
#+----+----+----+
#|col1|col2|col3|
#+----+----+----+
#|   1|   2|   3|
#|   4|   5|   6|
#|   7|   8|   9|
#+----+----+----+

由于您的条件基于 if-else 条件，您可以使用 when 和 otherwise 在每次迭代中执行逻辑。因为我不知道你的用例，所以我检查一个简单的条件，如果colX 是偶数，我们将它添加到 col3，如果奇数，我们减去。

我们根据列名末尾的数字加上列数（在我们的例子中为 3）在每次迭代中创建一个新列，以生成 4、5、6。

# You'll need a function to extract the number at the end of the column name
import re
def get_trailing_number(s):
  m = re.search(r'\d+$', s)
  return int(m.group()) if m else None

from pyspark.sql.functions import col, when
from pyspark.sql.types import FloatType
rich_df = df
for i in df.columns:
  rich_df = rich_df.withColumn(f'colget_trailing_number(i) + 3', \
   when(col(i) % 2 == 0, col(i) + col("col3"))\
   .otherwise(col(i) - col("col3")).cast(FloatType()))

rich_df.show()
#+----+----+----+----+----+----+
#|col1|col2|col3|col4|col5|col6|
#+----+----+----+----+----+----+
#|   1|   2|   3|-2.0| 5.0| 0.0|
#|   4|   5|   6|10.0|-1.0|12.0|
#|   7|   8|   9|-2.0|17.0| 0.0|
#+----+----+----+----+----+----+

这是函数的 UDF 版本

def func(col, constant):
  if (col % 2 == 0):
    return float(col + constant)
  else:
    return float(col - constant)

func_udf = udf(lambda col, constant: func(col, constant), FloatType())

rich_df = df
for i in df.columns:
  rich_df = rich_df.withColumn(f'colget_trailing_number(i) + 3', \
                               func_udf(col(i), col("col3")))

rich_df.show()
#+----+----+----+----+----+----+
#|col1|col2|col3|col4|col5|col6|
#+----+----+----+----+----+----+
#|   1|   2|   3|-2.0| 5.0| 0.0|
#|   4|   5|   6|10.0|-1.0|12.0|
#|   7|   8|   9|-2.0|17.0| 0.0|
#+----+----+----+----+----+----+

如果不了解您要做什么，就很难说更多。