ChainwithColumn用于在PySpark上多次更改一列

时间:2018-08-25 11:18:24

标签: pyspark pyspark-sql

我正在使用来自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 :(得分:1)

是的,通过链接when()来实现。

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']))