Uncovering and Solving Data Wrangling Issues with Property-Based Testing
Tags: data, datascience, testing, python
Published: 2024-05-02 13:52:51

_Article written as a member - and with the support - of the Philips' Software Excellence team. Curious about working in tech at Philips? [Find out more here.](https://www.careers.philips.com/global/en/philips-tech-careers)_



## Introduction

Testing data wrangling functions is notoriously challenging because it often involves      messy and unpredictable real-world data. This data can come in a wide variety of formats, with inconsistencies, missing values, and unexpected structures. 

Adding another layer of complexity are the wrangling operations themselves. These can range from simple filtering to complex transformations, and they often involve chaining multiple functions. This creates a cascading effect, where an error in one step can ripple through the entire process, making it difficult to pinpoint the source of the issue.

Despite the challenges, testing data wrangling functions remains crucial.
Complex data pipelines are inherently error-prone, and, without proper testing, errors can easily creep in and remain undetected for long periods, leading to downstream problems in analysis or reporting. 

Further, debugging data wrangling issues can be frustrating and time-consuming. Tests act as a safety net, allowing for quick identification and correction of problems before they cause downstream issues. 

Finally, well-written tests make it much easier to refactor code and add new functionality to the codebase. By ensuring that existing functionality remains intact, it allows for confident modification of the code without fear of unintended consequences.

## Property-based testing

Traditional testing involves checking a function’s behavior with a set of predefined      inputs. This works well for simple cases, but data wrangling functions often deal with a wide variety of messy, unpredictable data. 

Property-based testing flips the script. Instead of providing specific examples, it defines general rules—or properties—the function should always follow, regardless of the input. It then takes over, generating a large number of random inputs and checking if the function upholds these properties across all these different scenarios.

![Flowcharts comparing traditional testing approach with property-based testing](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/65gryygsry92mcasbbl8.png)

This approach proves particularly effective at catching bugs that traditional testing might overlook. By generating a vast array of random inputs, property-based testing significantly increases the chance of uncovering hidden bugs.

Popular property-based testing libraries to consider are:
- Python: [Hypothesis](https://hypothesis.readthedocs.io/)
- R: [Hedgehog](https://rdrr.io/cran/hedgehog/man/hedgehog.html)
- Scala: [ScalaCheck](https://scalacheck.org/documentation.html)
- Julia: [Supposition.jl](https://seelengrab.github.io/Supposition.jl/)

For a more detailed view of property-based testing, consider reading Fred Hebert’s book [PropEr Testing](https://propertesting.com/toc.html).

## A quick note about Fuzzing

Fuzzing, in simple terms, is a test that, considering a large variety of inputs, checks if the code crashes. Think of it as a brute-force approach to test code resilience.

The [hypothesis](https://hypothesis.readthedocs.io/) library, used in this article, blurs the line that separates property-based testing from fuzzing, and some tests may be closer to one concept or the other. The author and main maintainer of _hypothesis_ has an [article](https://hypothesis.works/articles/what-is-property-based-testing/) that explains his views on both property-based testing and fuzzing. I recommend reading it for further clarification.

## Case study

### Problem

I want to assess and predict trends in stroke mortality in U.S. states based on environmental quality indexes. This analysis relies on data wrangling functions that clean and join datasets from several sources.

Since these datasets are periodically updated, the data wrangling functions must adapt to handle new and potentially unpredictable data characteristics. This makes it difficult to test these functions. Focusing solely on the current data may not be sufficient. Ideally, tests should anticipate a wide range of issues that new datasets might introduce, such as unexpected values, missing entries, or structural changes. Property-based testing excels in this scenario. By generating a broad variety of test cases, it can help identify potential problems in the data wrangling logic regardless of the specific format of future datasets. This ensures the functions remain robust and adaptable even as the underlying data evolves.

For this study, data is expected to adhere to the following format:

| us_state | metric_1  | metric_2  | ... | metric_n  | target_variable |
| -------- | --------- | --------- | --- | --------- | --------------- |
| state_1  | value_1_1 | value_2_1 | ... | value_n_1 | target_1        |
| ...      | ...       | ...       | ... | ...       | ...             |
| state_n  | value_1_n | value_2_n | ... | value_n_n | target_n        |

Where all the values, other than “us_state”, are numerical and scaled (between 0 and 1).

#### Datasets

Metrics can be found on the dataset from the U.S. Environmental Protection Agency: 
- Dataset: [https://catalog.data.gov/dataset/usepa-environmental-quality-index-eqi-air-water-land-built-and-sociodemographic-domains-non-tra4](https://catalog.data.gov/dataset/usepa-environmental-quality-index-eqi-air-water-land-built-and-sociodemographic-domains-non-tra4)
- Data dictionary: [https://catalog.data.gov/dataset/usepa-environmental-quality-index-eqi-air-water-land-built-and-sociodemographic-domains-non-tra4](https://catalog.data.gov/dataset/usepa-environmental-quality-index-eqi-air-water-land-built-and-sociodemographic-domains-non-tra4)

The raw data has the following structure:

|stfips|State|County_Name|cat_RUCC|sociod_EQI_2Jan2018_VC|RUCC1_sociod_EQI_2Jan2018_VC|...|RUCC4_EQI_2Jan2018_VC_5|
|-|-|-|-|-|-|-|-|
|1001|AL|Autaga County|1|-0.525910|0.225516|...|NaN|
|1003|AL|Baldwin County|1|-0.702319|0.185723|...|NaN|
|...|...|...|...|...|...|...|...|
|56045|WY|Weston County|3|0.275503|NaN|...|NaN|

The target variable (stroke mortality) can be found on the dataset from the U.S. Department of Health & Human Services. 
- Dataset: [https://catalog.data.gov/dataset/stroke-mortality-data-among-us-adults-35-by-state-territory-and-county-2019-2021](https://catalog.data.gov/dataset/stroke-mortality-data-among-us-adults-35-by-state-territory-and-county-2019-2021)

This dataset has the following structure:

|Year|LocationAbbr|LocationDesc|GeographicLevel|DataSource|Class|Topic|Data_Value|...|Georeference|
|-|-|-|-|-|-|-|-|-|-|
|2020|AK|Nome|County|NVSS|Cardiovascular Diseases|Stroke Mortality|110.7|...|POINT (-163.9462296 64.903977039)|
|2020|CT|Tolland County|County|NVSS|Cardiovascular Diseases|Stroke Mortality|63.4|...|POINT (-72.337294 41.852989582)|
|...|...|...|...|...|...|...|...|...|...|
|2020|VT|Caledonia County|County|NVSS|Cardiovascular Diseases|Stroke Mortality|60.4|...|POINT (-72.10338989 44.460163015)|

Both datasets are available for free under open licenses.
- [Environmental Quality Index (EQI) dataset license](https://edg.epa.gov/EPA_Data_License.html)
- [Stroke Mortality Data Among U.S. Adults dataset license](http://opendefinition.org/licenses/odc-odbl/)

### Data wrangling functions

The following data wrangling functions were written to clean and prepare the data for training machine learning models.

```{python}
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler


def get_target(df, location_column, target_column):
    return df[[location_column, target_column]]\
        .groupby(location_column)\
        .sum()\
        .reset_index()


def get_metrics(df, location_column, metrics_columns):
    return df[[location_column] + metrics_columns]\
        .groupby(location_column)\
        .mean()\
        .reset_index()


def _scale_columns(df):
    aux = df.copy()

    scaler = MinMaxScaler()

    columns = aux.select_dtypes(include=np.number)\
        .columns\
        .tolist()

    aux[columns] = scaler.fit_transform(aux[columns])

    return aux


def get_df(target_df,
           metrics_df,
           target_location_column,
           metrics_location_column):
    aux = pd.merge(metrics_df,
                   target_df,
                   left_on=metrics_location_column,
                   right_on=target_location_column,
                   how="inner")

    if target_location_column != metrics_location_column:
        aux.drop(target_location_column, axis=1, inplace=True)

    return _scale_columns(aux)

```

Code explanation:

- _get_target()_: takes the dataset with the target variable (Stroke Mortality), groups and sums the target variable by state and returns the results. Once through the function, the data looks like this:

|LocationAbbr|Data_Value|
|-|-|
|AK|12707.8|
|AL|71147.8|
|...|...|
|WY|9418.8|

- _get_metrics()_: collects the metrics from the Environmental Quality Index dataset, calculates the average of each metric by state, and returns the results.

|State|sociod_EQI_2Jan2018_VC|RUCC1_sociod_EQI_2Jan2018_VC|...|RUCC4_EQI_2Jan2018_VC_5|
|-|-|-|-|-|
|AK|0.277039|-0.271424|...|1.235294|
|AL|0.441253|0.606241|...|4.727273|
|...|...|...|...|...|
|WY|-0.396936|-0.158441|...|2.000000|

- __scale_columns()_: is a helper function that scales all numerical variables to a number between 0 and 1.
- _get_df()_: creates the final data frame by combining both resulting datasets from the previous functions by state and scaling the results with the __scale_columns()_ helper function.

### Test code

To test any of the functions, we need to pass them a data frame. We can create simple data frames containing combinations of edge cases for each variable. A better alternative, however, is to use property-based testing techniques to generate random combinations of cases for each variable, run through the functions and test if the results fit the expected output.

In Python this can be done with the [hypothesis](https://hypothesis.readthedocs.io/) library.

The Python testing framework _pytest_ handles the tests. To get the output of _hypothesis_ with _pytest_ we can use the command `pytest --hypothesis-run-statistics`.

The following were the imports used for testing:

```{python}
from math import isnan

import numpy as np
import pytest
from datatest import validate
from hypothesis import given, settings
from hypothesis.extra.pandas import column, data_frames
from hypothesis.strategies import text, shared

from data_cleaning import get_target, get_metrics, get_df

```

Quick explanation of each import from the _hypothesis_ library: 
- _data_frames()_ handles the generation of pandas data frames. 
- _column()_ works within _data_frames()_ to create the columns of the generated data frame. It takes the column name, may take a dtype argument specifying the column data type and may take an elements parameter that specifies the strategy to use to generate each of the elements in the column (using other generators from the library). If elements are missing _hypothesis_ tries to infer the strategy using the dtype argument.
- _text()_ defines the data generation strategy for texts. 
- _shared()_ allows two or more generators to use the same data, provided by another generator. In simple terms, each _shared()_ using the same data generation strategy and the same key argument results in the same data each time the generators are sampled. 
- _given()_ is a function annotation that provides the data from the generators to the actual tests. 
- _settings()_ overrides _hypotesis_ settings for individual tests.

First step is to code the generators:

```{python}
states = text(min_size=1)

stroke_df_gen = data_frames([
    column("state", dtype=object, elements=shared(states, key="states")),
    column("target", dtype=float),
    column("garbage1", dtype=object, elements=text()),
    column("garbage2", dtype=int),
    column("garbage3", dtype=float)
])

eqi_df_gen = data_frames([
    column("state", dtype=object, elements=shared(states, key="states")),
    column("metric1", dtype=float),
    column("metric2", dtype=float),
    column("garbage1", dtype=object, elements=text()),
    column("garbage2", dtype=int),
    column("garbage3", dtype=float)
])

```

Code explanation:
- _states()_ generates random texts to be used as “state names”.
- _stroke_df_gen_ contains the generator for the data frame containing the target variable. Since the U.S. state column is used to merge the data frame of the target variable with the one with the metrics, the values generated by these columns need to be consistent with one another. This is achieved with the shared strategy from the _hypothesis_ library. The target variable must be a floating point number (stroke mortality per 100,000 population). Additional “garbage” columns were added to represent the extra data in the dataset. Hopefully, the functions should ignore these.
- _eqi_df_gen_ generates the metrics data frame. Like _stroke_df_gen_, it uses the shared strategy to keep the consistency of the U.S. state column between the two data frames. At a glance, all metrics columns seem to be floating point numbers on the original dataset, so that is represented in the data frame generator. Also, like in _stroke_df_gen_, “garbage” columns were added to represent the extra data found in the original dataset.

For the actual tests, I use the [datatest](https://datatest.readthedocs.io/) library, as it provides a straightforward way to test the properties of pandas data frames (as well as other common data types). Its _validate()_ function can be used both to test the types of each column/variable in a data frame and to test whether the values of a column pass a check from a helper function (for instance `lambda x: x > 0`, to check if all values are positive numbers).

Test code:

```{python}
@given(stroke_df_gen)
def test_get_target(s):
    res = get_target(s, "state", "target")
    validate(res, requirement=(str, float))
    validate(res["target"], requirement=lambda x: x >= 0)


def _not_nan(x):
    return not isnan(x)


@given(eqi_df_gen)
def test_get_metrics(s):
    res = get_metrics(s, "state", ["metric1", "metric2"])

    validate(res, requirement=(str, float, float))
    validate(res["metric1"], requirement=_not_nan)
    validate(res["metric1"], requirement=_not_nan)


def _check_scaled(x):
    return 0 <= x <= 1


@settings(deadline=None)
@given(stroke_df_gen, eqi_df_gen)
def test_get_df(strokes, eqi):
    strokes = get_target(strokes, "state", "target")
    eqi = get_metrics(eqi, "state", ["metric1", "metric2"])

    res = get_df(strokes, eqi, "state", "state")

    validate(res, requirement=(str, float, float, float))

    num_columns = res.select_dtypes(include=np.number)\
        .columns\
        .tolist()

    for c in num_columns:
        validate(res[c], requirement=_check_scaled)

```

Code explanation:
- _test_get_target()_ tests if the resulting data frames of the _get_target()_ function, having received the samples from _stroke_df_gen_ (_hypothesis_ defaults to 100 samples), conforms to column types str and float (note that this indirectly tests the number of columns, as it expects exactly 1 string column and 1 float, in that order) and if the target variable is positive (stroke mortality by 100,000 population should always be a positive number).
- _test_get_metrics()_ tests if _get_metrics()_, after receiving the samples from _eqi_df_gen_, returns data frames with 1 column with str type followed by 2 columns with float type. It also checks if the float columns contain missing data (as that can significantly harm the effectiveness of a machine learning model).
- _test_get_df()_ tests if _get_df()_, receiving samples from both _stroke_df_gen_ and _eqi_df_gen_ samples, returns data frames with 1 column with type str followed by 3 with type float and if each numerical column is scaled between 0 and 1 (important since many machine learning algorithms are sensitive to variables scales).

### Test results

Running the tests shows that all functions fail to meet the requirements in certain scenarios. In other words, the property-based tests succeeded in finding edge cases that break the code.

Analyzing the tests output makes it easier to identify what needs to be fixed/addressed in the code.

Tests summary:

```
============================================================ short test summary info ============================================================ 
FAILED tests/test_data_cleaning.py::test_get_target - datatest.ValidationError: does not satisfy <lambda> (1 difference): [
FAILED tests/test_data_cleaning.py::test_get_metrics - datatest.ValidationError: does not satisfy _not_nan() (1 difference): [
FAILED tests/test_data_cleaning.py::test_get_df - exceptiongroup.ExceptionGroup: Hypothesis found 3 distinct failures. (3 sub-exceptions)
```

- _test_get_target()_ fails due to the `lambda x: x >= 0` validation. This happens because _get_target()_ does not handle cases where the value is not a positive number. When building datasets, input errors may occur that lead to nonsensical values (like in this case, where a positive value is expected but a negative one is generated) or missing values. As it stands, _get_target()_ cannot handle these cases.
- _test_get_metrics()_ fails due to the __not_nan()_ validation. This means _get_metrics()_ does not handle missing values in its current form. This is critical, as the presence of missing values is a common occurrence on this dataset.
- _test_get_df()_ fails in 3 instances, but the test summary does not provide enough detail to understand why. To understand the reason behind these failures, we need to examine the test output more closely.

The output for _test_get_df()_ contains the following error:

```
    | datatest.ValidationError: does not satisfy _check_scaled() (1 difference): [
    |     Invalid(nan),
    | ]
    | Falsifying example: test_get_df(
    |     strokes=
    |           state  target garbage1  garbage2  garbage3
    |         0     0     0.0                  0       0.0
    |     ,
    |     eqi=
    |           state  metric1  metric2 garbage1  garbage2  garbage3
    |         0     0      0.0      NaN                  0       0.0
    |     ,
    | )
```

The original function uses _MinMaxScaler()_ from the _scikit-learn_ library to do the scaling on the metrics. This function returns a missing value (_NaN_) if one is inputted, causing the issue (_NaN_ is not a number between 0 and 1). This error gets propagated from _get_target()_ and _get_metrics()_. Correcting these functions should be enough to fix this error.

This was the second error found:

```
    | ValueError: Input X contains infinity or a value too large for dtype('float64').
    | Falsifying example: test_get_df(
    |     strokes=
    |           state  target garbage1  garbage2  garbage3
    |         0     0     0.0                  0       0.0
    |     ,
    |     eqi=
    |           state  metric1  metric2 garbage1  garbage2  garbage3
    |         0     0      0.0      inf                  0       0.0
    |     ,
    | )
```

This indicates that _get_df()_ cannot handle numbers that exceed the representable range of the float64 data type. While these extreme values are uncommon, data wrangling operations can sometimes produce them, making it advisable to address the issue.

The third error:

```
    | ValueError: Found array with 0 sample(s) (shape=(0, 3)) while a minimum of 1 is required by MinMaxScaler.
    | Falsifying example: test_get_df(
    |     strokes=
    |         Empty DataFrame
    |         Columns: [state, target, garbage1, garbage2, garbage3]
    |         Index: []
    |     ,
    |     eqi=
    |         Empty DataFrame
    |         Columns: [state, metric1, metric2, garbage1, garbage2, garbage3]
    |         Index: []
    |     ,  # or any other generated value
    | )
```

This shows that _get_df()_ cannot handle empty data frames, which can occur due to issues during data collection/extraction or because of filtering operations.

### Fixing the functions

Strategies: 
- _get_target()_: negative values should be replaced with 0 (should not count towards the sum total). 
- _get_metrics()_: U.S. states missing a metric should use the country average in its place. If the metric is missing for all states replace values with 0. 
- _get_df()_: 
    - 1st error: should be fixed by addressing the previous functions errors.
    - 2nd error: this could be addressed in the previous functions, specifically _get_target()_ and _get_metrics()_, so that they do not include extreme values (for instance, infinity) in their calculations.
    - 3rd error: the function should throw an error if the input data frames are empty. A test to check if it properly throws the error in this case should also be written.

Corrected functions:

```{python}
def get_target(df, location_column, target_column):
    aux = df[[location_column, target_column]]

    aux[target_column] = aux[target_column]\
        .apply(lambda x: x if 0 <= x < 3e8 else .0)

    return aux\
        .groupby(location_column)\
        .sum()\
        .reset_index()


def get_metrics(df, location_column, metrics_columns):
    aux = df[[location_column] + metrics_columns]
    aux[metrics_columns] = aux[metrics_columns]\
        .applymap(lambda x: x if -3e8 < x < 3e8 else np.NaN)

    aux = aux\
        .groupby(location_column)\
        .mean()\
        .reset_index()

    aux[metrics_columns] = aux[metrics_columns]\
        .fillna(aux[metrics_columns].mean())\
        .fillna(.0)

    return aux


def _scale_columns(df):
    aux = df.copy()

    scaler = MinMaxScaler()

    columns = aux.select_dtypes(include=np.number)\
        .columns\
        .tolist()

    aux[columns] = scaler.fit_transform(aux[columns])

    return aux


def get_df(target_df,
           metrics_df,
           target_location_column,
           metrics_location_column):
    if target_df.empty or metrics_df.empty:
        raise Exception("One or more empty DataFrames")

    aux = pd.merge(metrics_df,
                   target_df,
                   left_on=metrics_location_column,
                   right_on=target_location_column,
                   how="inner")

    if target_location_column != metrics_location_column:
        aux.drop(target_location_column, axis=1, inplace=True)

    return _scale_columns(aux)

```

And the adjustment for _test_get_df()_:

```{python}
@settings(deadline=None)
@given(stroke_df_gen, eqi_df_gen)
def test_get_df(strokes, eqi):
    strokes = get_target(strokes, "state", "target")
    eqi = get_metrics(eqi, "state", ["metric1", "metric2"])

    if strokes.empty or eqi.empty:
        with pytest.raises(Exception) as excinfo:
            get_df(strokes, eqi, "state", "state")
            assert str(excinfo.value) == "One or more empty DataFrames"
    else:
        res = get_df(strokes, eqi, "state", "state")

        validate(res, requirement=(str, float, float, float))

        num_columns = res.select_dtypes(include=np.number)\
            .columns\
            .tolist()

        for c in num_columns:
            validate(res[c], requirement=_check_scaled)

```

With these changes, the code passes the tests.

## Conclusion

Property-based testing for data wrangling functions has significant advantages in ensuring the robustness and reliability of data pipelines. By automatically generating a wide range of test cases with edge combinations, property-based testing goes beyond the limitations of traditional testing methods. It uncovers unexpected scenarios and potential issues that might be missed by hand-written tests focused on specific use cases.

The provided example showcased how property-based testing identified critical shortcomings in data wrangling functions, such as handling missing values, extreme numbers, and empty data frames. Addressing these issues proactively prevents errors from propagating downstream in the data pipeline, leading to more reliable and trustworthy results.

Property-based testing is a valuable tool for data scientists, engineers and analysts by streamlining the testing process and enhancing the quality of data wrangling functions. Its ability to generate a diverse set of test cases and uncover hidden issues makes it a powerful ally in building robust and reliable data pipelines.

Furthermore, the potential integration of property-based testing with continuous integration/continuous delivery (CI/CD) pipelines can further automate the testing process and ensure consistent code quality throughout the development lifecycle. As data wrangling becomes increasingly complex, property-based testing will undoubtedly play a crucial role in maintaining data integrity and delivering reliable results for data-driven projects.