如何解决在循环中使用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|
#+----+----+----+----+----+----+
在不了解您要做什么的情况下很难说更多话。
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。