如何在Pyspark中使用具有列功能的数据框中的函数？

Question

我有一些字典和函数定义：

dict_TEMPERATURE = {(0, 70): 'Low', (70.01, 73.99): 'Normal-Low',(74, 76): 'Normal', (76.01, 80): 'Normal-High', (80.01, 300): 'High'} 
... 
hierarchy_dict = {'TEMP': dict_TEMPERATURE, 'PRESS': dict_PRESSURE, 'SH_SP': dict_SHAFT_SPEED, 'POI': dict_POI, 'TRIG': dict_TRIGGER} 



def function_definition(valor, atributo): 

    dict_atributo = hierarchy_dict[atributo] 
    valor_generalizado = None 

    if isinstance(valor, (int, long, float, complex)): 

     for key, value in dict_atributo.items(): 

      if(isinstance(key, tuple)): 
       lista = list(key) 

       if (valor > key[0] and valor < key[1]): 
        valor_generalizado = value 

    else: # if it is not numeric 
     valor_generalizado = dict_atributo.get(valor) 


    return valor_generalizado 

什么这个功能基本上做的是：检查被作为参数传递给了“function_definition”功能传递的价值，并根据代替它的价值它的字典的引用。因此，如果我调用“function_definition（60，'TEMP'）”，它将返回'LOW'。

在另一方面，我有下一个结构的数据框（这是一个例子）：

+----+-----+-----+---+----+ 
|TEMP|SH_SP|PRESS|POI|TRIG| 
+----+-----+-----+---+----+ 
| 0| 1| 2| 0| 0| 
| 0| 2| 3| 1| 1| 
| 0| 3| 4| 2| 1| 
| 0| 4| 5| 3| 1| 
| 0| 5| 6| 4| 1| 
| 0| 1| 2| 5| 1| 
+----+-----+-----+---+----+ 

我想要做的是替换数据帧中的一列的基础上，该值功能如上所定义的，所以我有下一个代码行：

dataframe_new = dataframe.withColumn(atribute_name, function_definition(dataframe[atribute_name], atribute_name)) 

但在执行时，它我获得下一个错误消息：

AssertionError: col should be Column 

我的代码有什么问题？怎么可能做到这一点？

Answer 1

你function_definition（勇武，atributo）返回一个英勇一个String（valor_generalizado）。

AssertionError：col应为列意味着您传递的参数不是列的WithColumn(colName,col)。 所以你必须改变你的数据，为了有列，例如你可以看到下面。

数据框例如（相同的结构，你的）：

a = [(10.0,1.2),(73.0,4.0)] # like your dataframe, this is only an example 

dataframe = spark.createDataFrame(a,["tp", "S"]) # tp and S are random names for these columns 

dataframe.show() 
+----+---+ 
| tp| S| 
+----+---+ 
|10.0|1.2| 
|73.0|4.0| 
+----+---+ 

正如你可以看到here

udf Creates a Column expression representing a user defined function (UDF).

解决方案：

from pyspark.sql.functions import udf 

attr = 'TEMP' 
udf_func = udf(lambda x: function_definition(x,attr),returnType=StringType()) 

dataframe_new = dataframe.withColumn("newCol",udf_func(dataframe.tp)) 
dataframe_new.show() 

+----+---+----------+ 
| tp| S| newCol| 
+----+---+----------+ 
|10.0|1.2|  Low| 
|73.0|4.0|Normal-Low| 
+----+---+----------+