![](/img/trans.png)
[英]PySpark: How to apply UDF to multiple columns to create multiple new columns?
[英]use udf inside for loop to create multiple columns in Pyspark
我有一個帶有一些列(col1、col2、col3、col4、col5...直到 32)的 spark dataframe 現在我創建了一個 function (udf),它采用 2 個輸入參數並返回一些浮點值。
現在我想使用上面的 function 創建新列(按升序排列,如 col33、col32、col33、col34..),其中一個參數遞增,其他參數不變
def fun(col1,col2):
if true:
do someting
else:
do someting
我已將此 function 轉換為 udf
udf_func = udf(fun,Floatype())
現在我想使用這個 function 在 dataframe 中創建新列,該怎么做?
我試過
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|
#+----+----+----+----+----+----+
這是 function 的 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|
#+----+----+----+----+----+----+
在不了解您要做什么的情況下很難說更多。
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.