Как сделать среднее (целевое) кодирование в 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

Но Koalas не разрешает обновление на уровне ячеек, то есть я не могу изменить значение на ks_df.at[i, col_name] = new_value

Так что янадеясь, что может быть какое-то решение pyspark для этой проблемы.