Цепочка withColumn для многократного изменения одного столбца в PySpark - PullRequest
Новая
       5

Цепочка withColumn для многократного изменения одного столбца в PySpark

0 голосов
/ 25 августа 2018

Я использую годовой доход взрослого населения от UCI.

У меня есть фрейм данных с категориальной переменной в одном столбце, который я хочу сгруппировать по разным категориям (некоторые общие характеристики). 

df.groupBy('education').count().show()

дает: 

+------------+-----+
|   education|count|
+------------+-----+
|        10th| 1223|
|     Masters| 2514|
|     5th-6th|  449|
|  Assoc-acdm| 1507|
|   Assoc-voc| 1959|
|     7th-8th|  823|
|         9th|  676|
|     HS-grad|14783|
|   Bachelors| 7570|
|        11th| 1619|
|     1st-4th|  222|
|   Preschool|   72|
|        12th|  577|
|   Doctorate|  544|
|Some-college| 9899|
| Prof-school|  785|
+------------+-----+

Я хочу поместить следующие категории в определенные группы следующим образом: 

dropout = ['Preschool', '1st-4th', '5th-6th', '7th-8th', '9th', '10th', '11th', '12th']
community_college = ['Assoc-acdm', 'Assoc-voc', 'Some-college']
masters = ['Prof-school']

И для этого я могу сделать следующее: 

from pyspark.sql.functions import when, col
df = df.withColumn('education', when(col('education').isin(dropout), 'Dropout').otherwise(df['education']))
df = df.withColumn('education', when(col('education').isin(community_college), 'Community_college').otherwise(df['education']))
df = df.withColumn('education', when(col('education') == 'Prof-school', 'Masters').otherwise(df['education']))

Получение: 

+-----------------+-----+
|        education|count|
+-----------------+-----+
|          Masters| 3299|
|          HS-grad|14783|
|        Bachelors| 7570|
|          Dropout| 5661|
|        Doctorate|  544|
|Community_college|13365|
+-----------------+-----+

Есть ли возможность связать эти withColumn?Я пробовал следующее безуспешно: 

df = df.withColumn('education', when(col('education').isin(dropout), 'Dropout').otherwise(df['education']))\
.withColumn('education', when(col('education').isin(community_college), 'Community_college').otherwise(df['education']))\
.withColumn('education', when(col('education') == 'Prof-school', 'Masters').otherwise(df['education']))

1 Ответ

0 голосов
/ 25 августа 2018

Да, есть цепочка когда (). 

df = df.withColumn('education', when(col('education').isin(dropout), 'Dropout')\
                               .when(col('education').isin(community_college), 'Community_college')\
                               .when(col('education') == 'Prof-school', 'Masters') \
                               .otherwise(df['education']))
...