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

1つのUDFで1つの列しか作成できません。しかし、これは構造化された列にすることができ、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) 

しかし、わかりましたが、それは元の列を置き換えるものではなく、新しいものを追加するだけです。それらを置き換えるにはを2回使用してからlabelsを落とす必要があります。