pyspark 将 one-hot 编码向量分解为具有正确名称的每一列

Question

将 one-hot 编码应用于多个分类列

X_cat = X.select(cat_cols)

str_indexer = [StringIndexer(inputCol=col, outputCol=col+"_si", handleInvalid="skip") for col in cat_cols]
ohe = [OneHotEncoder(inputCol=f"{col}_si", outputCol=f"{col}_ohe", dropLast=True) for col in cat_cols]
# ohe.setDropLast(False) # older version

pl = Pipeline(stages=str_indexer + ohe).fit(X_cat)
X_cat = pl.transform(X_cat)

si_cols = [col_nm for col_nm in X_cat.columns if col_nm.endswith("_si")]
ohe_cols = [col_nm for col_nm in X_cat.columns if col_nm.endswith("ohe")]
X_cat_ohe = X_cat.select(ohe_cols)

给我

root
 |-- workclass_ohe: vector (nullable = true)
 |-- education_ohe: vector (nullable = true)
 |-- marital-status_ohe: vector (nullable = true)
 |-- occupation_ohe: vector (nullable = true)

+-------------+---------------+
|workclass_ohe|  education_ohe|
+-------------+---------------+
|(8,[4],[1.0])| (15,[2],[1.0])|
|(8,[1],[1.0])| (15,[2],[1.0])|
|(8,[0],[1.0])| (15,[0],[1.0])|
|(8,[0],[1.0])| (15,[5],[1.0])|

这基本上是

    workclass_ohe                                   education_ohe
0   (0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0)    (0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...
1   (0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)    (0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...
2   (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)    (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...
3   (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)    (0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, ...
4   (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)    (0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...

我想将向量中的值分解为具有正确名称的新列。

所需 output

 workclass_state-gov  workclass_selfemp  workclass_private  workclass_middle education_1   education_2  education_3
         0                   0                     0                1               0              0            1
         0                   1                     0                1               0              0            1
         1                   0                     0                0               1              0            0
        ...

From pyspark - 将一次热编码后获得的稀疏向量转换为列

我可以从X_cat_ohe添加新列，但是我无法确定哪个值（例如：state-gov）对应于第 0 个向量、第一个向量等等...

Answer 1

感谢使用 Pyspark 的虚拟编码，我可以将其扩展到多列。

for col_nm in cat_cols:
    category = X.select(col_nm).distinct().rdd.flatMap(lambda x:x).collect()
    category = [col_nm + "_" + ct for ct in category]
    exprs = [f.when(f.col(col_nm) == ct, 1).otherwise(0)\
              .alias(str(ct)) for ct in category]
    X = X.select(exprs+X.columns)