将标签分配给PySpark中的表中的分类数据

问题描述:

我想使用pyspark sql将标签分配给下面的数据框中的分类数字。将标签分配给PySpark中的表中的分类数据

在婚姻栏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) 

一个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