Chain withColumn 用于在 PySpark 上多次更改一列

Question

我使用的是来自 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']))

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

Chain withColumn 用于在 PySpark 上多次更改一列

Chain withColumn for changing one column multiple times on PySpark

pyspark

pyspark-sql