Imputing Missing Values Smartly with DataWig
By: Ansh Bordia
Tags: Missing Values, Feature Engineering, Data Science, Python, Machine Learning

Photo by Emily Morter on Unsplash

Data, and especially high quality data (alongside other things), plays a key role in the success of a machine learning model. Unfortunately, real world datasets are far from perfect and are often stained with noise and missing values. Missing values, in particular, can be problematic. Dropping them altogether results in a smaller dataset — not the ideal situation when your dataset is small. And even if you try to fill those annoying ‘NaNs’, the imputed data might not be representative of the missing data.

This is where DataWig comes in. As you might have guessed by the name, it is a wig for bald patches in your data. And look, while there are many methods to impute your data, DataWig belongs right at the top.

We’ll first see what DataWig is and how it works and then we will compare it against some most popular methods to see how well it performs.

DataWig was developed by AWS Labs around 3 years back. It tries to understand your data and uses that learning to do the imputation. So if you have 3 columns — ‘X’, ‘Y’ & ‘Z’ — and we want to impute ‘Z’, DataWig learns the contents of the other two columns to do its magic.

DataWig supports imputation of both categorical and numerical columns. A lot of imputation approaches are only catered towards numerical imputation, while those that cater to categorical aren’t so good or scalable in production. So this becomes extremely useful when you want to impute categorical columns.

Let’s dive a bit deeper into how DataWig works under the hood. This is based on its research paper, so if you don’t want to get into the technicalities, feel free to skip to the next section on how to use it. Though I’d recommend you read this part, its super interesting in my opinion.

The image shows several columns, with the column ‘Color’ to be imputed. DataWig first determines the type of each column. After that, each column is converted into a numeric representation (so that the machine understands it). Categorical columns are one hot encoded while sequential (text) columns are converted into a sequence based on the length and the characters in the string. Next comes the most important step — Featurizing. One-hot encoded data is passed through an embedding layer while sequential data is passed either via an LSTM layer or n-gram hashing is done on it. Finally, all the features are merged and passed through a logistic layer (since ‘Color’ is categorical) for the imputation.

Well that’s the theory, let’s see how you can use DataWig. It can be installed as a Python package.

pip install datawig

Note: When I last tried, there was an installation issue regarding the dependencies. Look at this link and replace ‘==’ with ‘>=’ for the dependencies.

Using DataWig is pretty straightforward as well. It provides you with two types of imputers — SimpleImputer & Imputer. Use the SimpleImputer when you don’t care how DataWig does the work underneath. But if you want to customize the working according to your requirements, use Imputer. For example, you can select if you want to use the LSTM or n-gram hashing for featurizing. Both imputers come with the easy to use fit() and predict() methods (similar to sklearn). Below we’ll do some coding where you’ll get to know more about these functionalities. But if you want to become a DataWig geek, refer to their documentation.