import pyspark.ml.feature as MF

imputer = MF.Imputer(strategy='mean', inputCols=['a', 'b'], outputCols=['out_a', 'out_b'])
model = imputer.fit(df)
model.transform(df).show()

# Calculate the mean for each numeric column
mean_values = vehicles_df_spark.agg(*[F.mean(c).alias(c) for c in vehicles_df_spark.columns if isinstance(vehicles_df_spark.schema[c].dataType, NumericType)]).first()

# Create a dictionary for filling missing values
mean_dict = {c: mean_values[c] for c in mean_values.asDict()}

# Fill missing values with the calculated mean
imputed_df = vehicles_df_spark.fillna(mean_dict)

Median

# Calculate the median for each numeric column
# Since Spark does not have a built-in median function, we use the approxQuantile method
median_values = {c: vehicles_df_spark.approxQuantile(c, [0.5], 0)[0] 
                 for c in vehicles_df_spark.columns if isinstance(vehicles_df_spark.schema[c].dataType, NumericType)}

# Fill missing values with the calculated median
imputed_df = vehicles_df_spark.fillna(median_values)