test_transformers.py

from collections import defaultdict

import numpy as np
import pandas as pd
import pytest

from rdt import HyperTransformer
from rdt.performance.datasets import BaseDatasetGenerator
from rdt.transformers import BaseTransformer

DATA_SIZE = 1000
TEST_COL = 'test_col'

PRIMARY_SDTYPES = ['boolean', 'categorical', 'datetime', 'numerical']

INT64_MIN = np.iinfo(np.int64).min

# Additional arguments for transformers
TRANSFORMER_ARGS = {
    'BinaryEncoder': {
        'missing_value_replacement': -1,
        'missing_value_generation': 'from_column',
    },
    'UnixTimestampEncoder': {'missing_value_generation': 'from_column'},
    'OptimizedTimestampEncoder': {'missing_value_generation': 'from_column'},
    'FloatFormatter': {'missing_value_generation': 'from_column'},
    'GaussianNormalizer': {'missing_value_generation': 'from_column'},
    'ClusterBasedNormalizer': {'missing_value_generation': 'from_column'},
    'LogScaler': {'constant': float(INT64_MIN), 'missing_value_generation': 'from_column'},
}

# Mapping of rdt sdtype to dtype
SDTYPE_TO_DTYPES = {
    'boolean': ['b', 'O'],
    'categorical': ['O', 'i', 'f'],
    'datetime': ['M'],
    'float': ['f', 'i'],
    'id': ['O', 'i', 'f'],
    'integer': ['i'],
    'numerical': ['f', 'i'],
    'pii': ['O', 'i', 'f'],
    'text': ['O', 'i', 'f'],
}


def _validate_helper(validator_function, args, steps):
    """Wrap around validation functions to either return a boolean or assert.

    Args:
        validator_function(function):
            The function to validate.
        args (list):
            The args to pass into the function.
        steps (list):
            List of steps that the validation has completed.
    """
    if steps is not None:
        steps.append(validator_function.__name__)

    validator_function(*args)


def _is_valid_transformer(transformer_name):
    """Determine if transformer should be tested or not."""
    invalid_names = [
        'IdentityTransformer',
        'Dummy',
        'OrderedLabelEncoder',
        'CustomLabelEncoder',
        'OrderedUniformEncoder',
        'BaseMultiColumnTransformer',
    ]
    return all(invalid_name not in transformer_name for invalid_name in invalid_names)


def _get_all_transformers():
    """Get all transformers to be tested."""
    all_transformers = BaseTransformer.get_subclasses()
    return [
        transformer
        for transformer in all_transformers
        if _is_valid_transformer(transformer.__name__)
    ]


def _build_generator_map():
    """Build a map of sdtype to data generator.

    Output:
        dict:
            A mapping of sdtype (str) to a list of data
            generators (rdt.tests.datasets.BaseDatasetGenerator).
    """
    generators = defaultdict(list)

    for generator in BaseDatasetGenerator.get_subclasses():
        generators[generator.SDTYPE].append(generator)

    return generators


def _find_dataset_generators(sdtype, generators):
    """Find the dataset generators for the given sdtype."""
    if sdtype is None:
        primary_generators = []
        for primary_sdtype in PRIMARY_SDTYPES:
            primary_generators.extend(_find_dataset_generators(primary_sdtype, generators))

        return primary_generators

    return generators.get(sdtype, [])


def _validate_dataset_generators(dataset_generators):
    """Check that the number of dataset generators is greater than zero."""
    assert len(dataset_generators) > 0, 'There are no associated dataset generators.'


def _validate_transformed_data(transformer, transformed_data):
    """Check that the transformed data is the expected dtype."""
    expected_sdtypes = transformer.get_output_sdtypes()
    transformed_dtypes = transformed_data.dtypes

    for column, expected_sdtype in expected_sdtypes.items():
        message = f'Column {column} is expected but not found in transformed data.'
        assert column in transformed_data, message
        message = f'Column {column} is not the expected sdtype {expected_sdtype}'
        assert transformed_dtypes[column].kind in SDTYPE_TO_DTYPES[expected_sdtype], message


def _validate_reverse_transformed_data(transformer, reversed_data, input_dtype):
    """Check that the reverse transformed data is the expected dtype.

    Expect that the dtype is equal to the dtype of the input data.
    """
    expected_sdtype = transformer.get_supported_sdtypes()[0]
    message = f'Reverse transformed data is not the expected sdtype {expected_sdtype}'
    assert reversed_data.dtypes[TEST_COL].kind in SDTYPE_TO_DTYPES[expected_sdtype], message


def _test_transformer_with_dataset(transformer_class, input_data, steps):
    """Test the given transformer with the given input data.

    This method verifies the transformed and reverse transformed data's dtype

    Args:
        transformer_class (rdt.transformers.BaseTransformer):
            The transformer class to test.
        input_data (pandas.Series):
            The data to test on.
        steps (list):
            List of steps that the validation has completed.
    """

    transformer_args = TRANSFORMER_ARGS.get(transformer_class.__name__, {})
    transformer = transformer_class(**transformer_args)
    # Fit
    transformer.fit(input_data, [TEST_COL])

    # Transform
    transformed = transformer.transform(input_data)
    _validate_helper(
        _validate_transformed_data,
        [transformer, transformed],
        steps,
    )

    # Reverse transform
    out = transformer.reverse_transform(transformed)
    _validate_helper(
        _validate_reverse_transformed_data,
        [transformer, out, input_data.dtypes[TEST_COL]],
        steps,
    )


def _validate_hypertransformer_transformed_data(transformed_data):
    """Check that the transformed data is not null and of type float."""
    assert transformed_data.notna().all(axis=None), 'Transformed data has nulls.'

    for dtype in transformed_data.dtypes:
        assert dtype.kind in SDTYPE_TO_DTYPES['numerical'], 'Transformed data is not numerical.'


def _validate_hypertransformer_reverse_transformed_data(transformer, reversed_data):
    """Check that the reverse transformed data has the same dtype as the input."""
    expected_sdtype = transformer().get_supported_sdtypes()[0]
    message = f'Reversed transformed data is not the expected sdtype {expected_sdtype}'
    assert reversed_data.dtype.kind in SDTYPE_TO_DTYPES[expected_sdtype], message


def _test_transformer_with_hypertransformer(transformer_class, input_data, steps):
    """Test the given transformer in the hypertransformer.

    Run the provided transformer using the hypertransformer using the provided
    input data. Verify that the expected dtypes are returned by transform
    and reverse_transform.

    Args:
        transformer_class (rdt.transformers.BaseTransformer):
            The transformer class to test.
        input_data (pandas.Series):
            The data to test on.
        steps (list):
            List of steps that the validation has completed.
    """
    transformer_args = TRANSFORMER_ARGS.get(transformer_class.__name__, {})
    hypertransformer = HyperTransformer()
    if transformer_args:
        field_transformers = {TEST_COL: transformer_class(**transformer_args)}

    else:
        field_transformers = {TEST_COL: transformer_class()}

    sdtypes = {}
    for field, transformer in field_transformers.items():
        sdtypes[field] = transformer.get_supported_sdtypes()[0]

    config = {'sdtypes': sdtypes, 'transformers': field_transformers}
    hypertransformer.set_config(config)
    hypertransformer.fit(input_data)

    transformed = hypertransformer.transform(input_data)
    _validate_helper(_validate_hypertransformer_transformed_data, [transformed], steps)

    out = hypertransformer.reverse_transform(transformed)
    _validate_helper(
        _validate_hypertransformer_reverse_transformed_data,
        [transformer_class, out[TEST_COL]],
        steps,
    )


def validate_transformer(transformer, steps=None, subtests=None):
    """Validate that the transformer passes all integration checks.

    Args:
        transformer (rdt.transformer.BaseTransformer):
            The transformer to validate.
        steps (list):
            List of steps that the validation has completed.
        subtests:
            Whether or not to test with subtests.
    """
    input_sdtype = transformer.get_supported_sdtypes()[0]

    dataset_generators = _find_dataset_generators(input_sdtype, generators)
    _validate_helper(_validate_dataset_generators, [dataset_generators], steps)

    for dg in dataset_generators:
        data = pd.DataFrame({TEST_COL: dg.generate(DATA_SIZE)})

        if subtests:
            with subtests.test(msg=f'test_transformer_with_dataset_{dg}', generator=dg):
                _test_transformer_with_dataset(transformer, data, steps)
                _test_transformer_with_hypertransformer(transformer, data, steps)
        else:
            _test_transformer_with_dataset(transformer, data, steps)
            _test_transformer_with_hypertransformer(transformer, data, steps)


transformers = _get_all_transformers()
generators = _build_generator_map()


@pytest.mark.parametrize('transformer', transformers)
def test_transformer(subtests, transformer):
    """Test the transformer end-to-end.

    Test the transformer end-to-end with at least one generated dataset. Test
    both the transformer by itself, and by running in the hypertransformer.

    Args:
        transformer (rdt.transformers.BaseTransformer):
            The transformer to test.
    """
    validate_transformer(transformer, subtests=subtests)