Filling missing values with pyspark using a probability distribution

Question

I want to fill missing values in my dataframe.

 In [1]: df = spark.createDataFrame([[1],[1],[2],[3],[3],[None],[3],[None],[3],[2],[None],[1],[4]], ['data'])     
 In [2]: df.show()
 +----+
 |data|
 +----+
 |   1|
 |   1|
 |   2|
 |   3|
 |   3|
 |null|
 |   3|
 |null|
 |   3|
 |   2|
 |null|
 |   1|
 |   4|
 +----+

I know I can use pyspark.ml Imputer to fill with the mean / median, or use this method to fill with the last valid value. These are fine options, but I would like to impute with a random sample from the data distribution. For example, in the data provided, nulls will be filled according to these probabilities:

 P(1) = .3
 P(2) = .2
 P(3) = .4
 P(4) = .1

What would be the best way to fill these values from a random sample?

Answer 1

You could try modeling it as a discrete distribution and then try obtaining the random samples. Try making a function p(x) and deriving the CDF from that. In the example you gave the CDF graph would look like this

Once you obtained your CDF you can try using Inverse Transform Sampling. This method allows you to obtain random samples for the distribution specified.

This link would be helpful for your understanding of that technique. These random samples can fill those missing values as per your requirement of probabilities.

Note:

There are other techniques as well, you could search and explore along the lines of random sample generation from discrete distributions. It might be the case that your actual data might fit for example something like Poisson's distribution etc. For all such cases, you can find various examples on how to generate random samples for those particular distributions

Answer 2

Here is a more concrete example, which sets missing values sampled at random from a Normal distribution, after estimating its parameters from the data.

If you want to sample from some other distribution based on the data you could modify the _fit method to set up the distribution.