將標籤分配給PySpark中的表中的分類數據

Question

我想使用pyspark sql將標籤分配給下面的數據框中的分類數字。

在婚姻欄1 =已婚，2 =未婚。在教育列1 =梯度和2 =本科生

 
Current Dataframe: 
+--------+---------+-----+ 
|MARRIAGE|EDUCATION|Total| 
+--------+---------+-----+ 
|  1|  2| 87| 
|  1|  1| 123| 
|  2|  2| 3| 
|  2|  1| 8| 
+--------+---------+-----+ 

 
Resulting Dataframe: 
+---------+---------+-----+ 
|MARRIAGE |EDUCATION|Total| 
+---------+---------+-----+ 
|Married |Grad  | 87| 
|Married |UnderGrad| 123| 
|UnMarried|Grad  | 3| 
|UnMarried|UnderGrad| 8| 
+---------+---------+-----+ 

是否有可能使用單個UDF和withColumn（）來分配標籤？有沒有什麼辦法通過傳遞整個數據框並保持列名不變，從而在單個UDF中分配？

我可以想出一個解決方案，通過使用單獨的udfs來完成每列的操作，如下所示。但無法弄清楚是否有辦法一起做。

from pyspark.sql import functions as F 

def assign_marital_names(record): 
    if record == 1: 
     return "Married" 
    elif record == 2: 
     return "UnMarried" 


def assign_edu_names(record): 
    if record == 1: 
     return "Grad" 
    elif record == 2: 
     return "UnderGrad" 

assign_marital_udf = F.udf(assign_marital_names) 
assign_edu_udf = F.udf(assign_edu_names) 
df.withColumn("MARRIAGE", assign_marital_udf("MARRIAGE")).\ 
withColumn("EDUCATION", assign_edu_udf("EDUCATION")).show(truncate=False) 

Answer 1

一個UDF只能生成一列。但是這可以是結構化的專欄，並且UDF可以在婚姻和教育上應用標籤。看到下面的代碼：

from pyspark.sql.types import * 
from pyspark.sql import Row 

udf_result = StructType([StructField('MARRIAGE', StringType()), StructField('EDUCATION', StringType())]) 

marriage_dict = {1: 'Married', 2: 'UnMarried'} 
education_dict = {1: 'Grad', 2: 'UnderGrad'} 
def assign_labels(marriage, education): 
    return Row(marriage_dict[marriage], education_dict[education]) 

assign_labels_udf = F.udf(assign_labels, udf_result) 
df.withColumn('labels', assign_labels_udf('MARRIAGE', 'EDUCATION')).printSchema() 
root 
|-- MARRIAGE: long (nullable = true) 
|-- EDUCATION: long (nullable = true) 
|-- Total: long (nullable = true) 
|-- labels: struct (nullable = true) 
| |-- MARRIAGE: string (nullable = true) 
| |-- EDUCATION: string (nullable = true) 

但是，正如你所看到的，它並不取代原來的列，它只是增加一個新的。要替換它們，您需要兩次使用withColumn，然後再使用labels。