[英]How to apply a function to a set of columns of a PySpark dataframe by rows?
[英]Pyspark - How to apply a function only to a subset of columns in a DataFrame?
我想用不同的方法将 function 应用于Spark
DataFrame 的某些列: fn
和fn1
。 我是这样做的:
def fn(column):
return(x*2)
udf_1 = udf(fn, DecimalType())
def fn1(column):
return(x*3)
udf_2 = udf(fn1, DecimalType())
def process_df1(df, col_name):
df1 = df.withColumn(col_name, udf_1(col_name))
return df1
def process_df2(df, col_name):
df2 = df.withColumn(col_name, udf_2(col_name))
return df2
对于单个列,它工作正常。 但是现在我得到了一个dict
list
,其中包含有关各个列的信息:
cols_info = [{'col_name': 'metric_1', 'process': 'True', 'method':'simple'}, {'col_name': 'metric_2', 'process': 'False', 'method':'hash'}]
我应该如何解析cols_info
列表并将上述逻辑仅应用于具有process:True
并使用所需method
的列?
首先想到的是用process:False
过滤掉列
list(filter(lambda col_info: col_info['process'] == 'True', cols_info))
但是我在这里仍然缺少一种更通用的方法。
selectExpr function 在这里会有用
import pyspark.sql.functions as F
from pyspark.sql.window import Window
#Test data
tst = sqlContext.createDataFrame([(1,2,3,4),(1,3,4,1),(1,4,5,5),(1,6,7,8),(2,1,9,2),(2,2,9,9)],schema=['col1','col2','col3','col4'])
def fn(x):
return(x*2)
def fn1(x):
return(x*3)
sqlContext.udf.register("fn1", fn)
sqlContext.udf.register("fn2", fn1)
cols_info =[{'col_name':'col1','encrypt':False,},{'col_name':'col2','encrypt':True,'method':'fn1'},{'col_name':'col3','encrypt':True,'method':'fn2'}]
# determine which columns have any of the encryption
modified_columns = [x['col_name'] for x in cols_info if x['encrypt']]
# select which colulmns have to be retained
columns_retain = list(set(tst.columns)-set(modified_columns))
#%
expr =columns_retain+[((x['method'])+'('+(x['col_name'])+') as '+ x['col_name']) for x in cols_info if x['encrypt']]
#%
tst_res = tst.selectExpr(*expr)
结果将是:
+----+----+----+----+
|col4|col1|col2|col3|
+----+----+----+----+
| 4| 1| 4| 9|
| 1| 1| 6| 12|
| 5| 1| 8| 15|
| 8| 1| 12| 21|
| 2| 2| 2| 27|
| 9| 2| 4| 27|
+----+----+----+----+
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.