如何在pyspark中进行平均（目标）编码

Question

我需要对数据集中的所有分类列进行平均（目标）编码。为了简化这个问题，假设我的数据集中有 2 列，第一列是标签列，第二列是分类列。

例如

label | cate1   
  0   |  abc    
  1   |  abc    
  0   |  def    
  0   |  def    
  1   |  ghi

所以按照均值编码策略：https://towardsdatascience.com/why-you-should-try-mean-encoding-17057262cd0

输出应该是这样的

label | cate1    
  0   |  0.5   
  1   |  0.5    
  0   |  0.0    
  0   |  0.0    
  1   |  1.0

我试过考拉来解决这个问题，但失败了。这是我尝试过的：

for col_name in convert_cols:


    cat_mean_dict = dict()
    # get category name <-> count dictionary
    cur_col_cate_count_ = ks_df[col_name].value_counts().to_dict()
    print(cur_col_cate_count_)

    # calculate all different categories positive result count and mean value
    start_time = time.time()
    for key in cur_col_cate_count_:

        current_col_positive_count = ks_df.loc[(ks_df['0'] == 1) & (ks_df[col_name] == key)].shape[0]
        key_mean = current_col_positive_count / cur_col_cate_count_[key]
        cat_mean_dict[key] = key_mean


    for i in range(ks_df.shape[0]):
        cate_origin_hash = ks_df.at[i, col_name]
        if cate_origin_hash in cat_mean_dict:
            ks_df.at[i, col_name] = cat_mean_dict[cate_origin_hash]
        else:
            ks_df.at[i, col_name] = -1

但是考拉不允许单元格级别的更新，这意味着我不能通过ks_df.at[i, col_name] = new_value修改值

所以我希望 pyspark 可以解决这个问题。

Answer 1

请在下面找到一个 pyspark 解决方案：

# spark inputs
spark_data = [Row(label=0, cate1='abc'),
              Row(label=1, cate1='abc'),
              Row(label=0, cate1='def'),
              Row(label=0, cate1='def'),
              Row(label=1, cate1='ghi')]

df = spark.createDataFrame(spark_data)

df.show()
>>>
+-----+-----+
|cate1|label|
+-----+-----+
|  abc|    0|
|  abc|    1|
|  def|    0|
|  def|    0|
|  ghi|    1|
+-----+-----+


# function
def target_mean_encoding(df, col, target):
    """
    :param df: pyspark.sql.dataframe
        dataframe to apply target mean encoding
    :param col: str list
        list of columns to apply target encoding
    :param target: str
        target column
    :return:
        dataframe with target encoded columns
    """
    target_encoded_columns_list = []
    for c in col:
        means = df.groupby(F.col(c)).agg(F.mean(target).alias(f"{c}_mean_encoding"))
        dict_ = means.toPandas().to_dict()
        target_encoded_columns = [F.when(F.col(c) == v, encoder)
                                  for v, encoder in zip(dict_[c].values(),
                                                        dict_[f"{c}_mean_encoding"].values())]
        target_encoded_columns_list.append(F.coalesce(*target_encoded_columns).alias(f"{c}_mean_encoding"))
    return df.select(target, *target_encoded_columns_list)


# function apply on spark inputs
df_target_encoded = target_mean_encoding(df, col=['cate1'], target='label')

df_target_encoded.show()
>>> 
+-----+-------------------+
|label|cate1_mean_encoding|
+-----+-------------------+
|    0|                0.5|
|    1|                0.5|
|    0|                0.0|
|    0|                0.0|
|    1|                1.0|
+-----+-------------------+


# if you want to keep the same column name after target mean encoder
df_target_encoded.withColumnRenamed('cate1_mean_encoding', 'cate1')

df_target_encoded.show()
>>>
+-----+-----+
|label|cate1|
+-----+-----+
|    0|  0.5|
|    1|  0.5|
|    0|  0.0|
|    0|  0.0|
|    1|  1.0|
+-----+-----+