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

问题描述

我有一个带有一些列（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中实现这一目标？

解决方法

您只能使用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'col{get_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'col{get_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|
#+----+----+----+----+----+----+

在不了解您要做什么的情况下很难说更多话。