如何在pyspark中用均值替换异常值？

from pyspark.sql import SparkSession
from pyspark.sql.functions import mean, stddev
from pyspark.sql.functions import col

spark = SparkSession.builder.appName("ReplaceOutliers").getOrCreate()

data = spark.read.csv("data.csv", header=True, inferSchema=True)

stats = data.select(mean(col("column_name")).alias("mean"), stddev(col("column_name")).alias("stddev")).collect()
mean_value = stats[0]["mean"]
stddev_value = stats[0]["stddev"]

def replace_outliers(value):
    if abs(value - mean_value) > 3 * stddev_value:
        return mean_value
    else:
        return value

replace_outliers_udf = udf(replace_outliers)
data = data.withColumn("column_name", replace_outliers_udf(col("column_name")))

data.show()