Spark:在对列的子集执行插补时如何保留数据框的所有列?

Spark: How to keep all columns of a dataframe when performing imputations on a subset of columns?

我正在尝试对我的数据框 (df_data) 中的列子集进行下面的插补,但在此过程中,我生成的数据框 (imputeDF) 只有插补列。我希望保留所有原始列并添加回估算列(并删除原始非估算列)。估算的数据框没有 id 可以重新加入 df_data 并且我已经研究了 withColumn 函数但无法让它为我工作。任何帮助,将不胜感激。谢谢

removeAllDF = df_data.na.drop()
imputeDF=df_data['exact_age','lnght_of_resd','acct_tenure_mnth_nbr','acct_ttce_mnth_nbr','tot_promo_amt',               'tot_rev_amt','int_base_pkg_speed','int_mx_elgbl_speed']
for x in imputeDF.columns:
    meanValue = removeAllDF.agg(avg(x)).first()[0]
    print(x, meanValue)
    imputeDF = imputeDF.na.fill(meanValue, (x))
    return imputeDF
removeAllDF = df_data.na.drop()
imputeDF=df_data['exact_age','lnght_of_resd','acct_tenure_mnth_nbr','acct_ttce_mnth_nbr','tot_promo_amt',               'tot_rev_amt','int_base_pkg_speed','int_mx_elgbl_speed']
for x in imputeDF.columns:
    meanValue = removeAllDF.agg(avg(x)).first()[0]
    print(x, meanValue)
    df_data = df_data.na.fill(meanValue, (x))

df_data.show