title: Pyspark Dataframe Imputations -- Replace Unknown & Missing Values with Column Mean based on specified condition
body: <p>Given a Spark dataframe, I would like to compute a column mean based on the non-missing and non-unknown values for that column. I would then like to take this mean and use it to replace the column's missing &amp; unknown values. </p>
  
  <p>For example, assuming I'm working with a: </p>
  
  <ul>
  <li>Dataframe named df, where each record represents one individual and all columns are integer or numeric</li>
  <li>Column named age (ages for each record) </li>
  <li>Column named missing_age (which equals 1 if that individual has no age, 0 otherwise)</li>
  <li>Column named unknown_age (which equals 1 if that individual has unknown age, 0 otherwise)</li>
  </ul>
  
  <p>Then I can compute this mean as shown below. </p>
  
  <pre><code>calc_mean = df.where((col("unknown_age") == 0) &amp; (col("missing_age") == 0))
  .agg(avg(col("age")))
  </code></pre>
  
  <p>OR via SQL and windows functions, </p>
  
  <pre><code>mean_compute = hiveContext.sql("select avg(age) over() as mean from df 
  where missing_age = 0 and unknown_age = 0")
  </code></pre>
  
  <p>I <strong>don't</strong> want to use SQL/windows functions if I can help it. My challenge has been taking this mean and replacing the unknown/missing values with it using non-SQL methods. </p>
  
  <p>I've tried using when(), where(), replace(), withColumn, UDFs, and combinations... Regardless of what I do, I either get errors or the results aren't what I expect. Here's an example of one of many things I've tried that didn't work. </p>
  
  <pre><code>imputed = df.when((col("unknown_age") == 1) | (col("missing_age") == 1),
  calc_mean).otherwise("age")
  </code></pre>
  
  <p>I've scoured the web, but haven't found similar imputation type questions so any help is much appreciated. It could be something very simple that I've missed.</p>
  
  <p>A side note -- I'm trying to apply this code to all columns in the Spark Dataframe that don't have unknown_ or missing_ in the column names. Can I just wrap the Spark related code in a Python 'for loop' and loop through all of the applicable columns to do this? </p>
  
  <p><strong>UPDATE:</strong></p>
  
  <p>Also figured out how to loop through columns... Here's an example. </p>
  
  <pre><code>for x in df.columns:
      if 'unknown_' not in x and 'missing_' not in x:
          avg_compute = df.where(df['missing_' + x] != 1).agg(avg(x)).first()[0]
          df = df.withColumn(x + 'mean_miss_imp', when((df['missing_' + x] == 1), 
          avg_compute).otherwise(df[x]))
  </code></pre>