Выброс лечения с использованием среднего значения в Pyspark

Question

Мой фрейм данных выглядит как -

id       gender      age
1          m         27
2          m         39
3          f         99
4          f         11
5          m         46
6          f         60

Я хочу, чтобы мой окончательный фрейм данных был похож -

id       gender      age       new_age
1          m         27          27
2          m         39          39
3          f         99          43
4          f         11          43
5          m         46          46
6          f         60          60

Мой код -

from pyspark.sql.functions import mean as _mean, stddev as _stddev, col

condition = ((df['age'] >= 18 & df['age'] <=60))
df = df.withColumn("new_age", when(condition, (col("age"))).otherwise(_mean(col('age')))

, но я хочу среднее значение только 27,39,46 и 60 ... не часть выброса. Как это сделать в pyspark?

Answer 1

Вот способ, которым вы можете сделать:

from operator import add

# convert outliers into a list of strings
outliers = [11,99]
outliers_str = '|'.join([str(i) for i in outliers])

# calculate mean without outlier values
mean_val = df.select("age").rdd.flatMap(lambda x: [i for i in x if i not in outliers]).mean()

# replace mean with outlier values
df = df.withColumn('new_age', F.regexp_replace('age', outliers_str, f'{mean_val}').cast('int'))

+---+------+---+-------+
| id|gender|age|new_age|
+---+------+---+-------+
|  1|     m| 27|     27|
|  2|     m| 39|     39|
|  3|     f| 99|     43|
|  4|     f| 11|     43|
|  5|     m| 46|     46|
|  6|     f| 60|     60|
+---+------+---+-------+