Janky profiler experiment.

excellent_profiler.py

import os
from typing import Any, Tuple
import pandas as pd
import sqlalchemy
from great_expectations import DataContext

from great_expectations.core import ExpectationSuite
from great_expectations.dataset import Dataset
from great_expectations.profile.base import Profiler, ProfilerTypeMapping


class ExcellentProfiler(Profiler):
    """This profiler is interesting because it uses both data and metadata as a source of truth."""

    def validate(self, item_to_validate: Any) -> None:
        super().validate(item_to_validate)
        if not isinstance(item_to_validate, tuple):
            # TODO this tuple workaround stinks
            raise TypeError("This profiler requries a tuple of a (batch, csv_filename)")
        batch = item_to_validate[0]
        if not isinstance(batch, Dataset):
            raise TypeError("This batch isn't valid.")

        csv_filename = item_to_validate[1]
        if not isinstance(csv_filename, str):
            raise TypeError("The csv_filename must be a string")
        if not csv_filename.endswith(".csv"):
            raise TypeError("This profiler requires a csv")
        if not os.path.isfile(csv_filename):
            raise FileNotFoundError(
                "The csv could not be found. Please check your path."
            )

    def _profile(
        self, terribly_named_tuple: Tuple[Dataset, str], suite_name: str = None
    ) -> ExpectationSuite:
        batch = terribly_named_tuple[0]
        csv_filename = terribly_named_tuple[1]
        print(f"Yo dawg I'm profiling this csv {csv_filename}")
        metadata_df = pd.read_csv(csv_filename)
        print(metadata_df.head())
        print(batch.head())
        print(batch.dtypes)

        for index, row in metadata_df.iterrows():
            col_type = row["type"]
            col_name = row["name"]
            col_semantic_type = row["semantic_type"]

            batch = self._create_type_list_expectation(batch, col_name, col_type)

            # handle value sets
            if col_semantic_type == "value_set":
                batch = self._create_value_set_expectation(batch, col_name)

            # text means we don't really have anything interesting other than completeness, uniqueness, data type
            elif (
                col_semantic_type in ProfilerTypeMapping.STRING_TYPE_NAMES
                # TODO these may not be needed
                or col_semantic_type in ["text", "str", "string"]
                or isinstance(col_type, sqlalchemy.types.String)
            ):
                #         print('string', col_name, col_type, col_semantic_type)
                pass

            # numeric types
            elif (
                col_semantic_type in ProfilerTypeMapping.INT_TYPE_NAMES
                or col_semantic_type in ProfilerTypeMapping.FLOAT_TYPE_NAMES
                # TODO these may not be needed
                or isinstance(col_type, sqlalchemy.types.Numeric)
                or isinstance(col_type, sqlalchemy.types.Integer)
            ):
                batch = self._create_numeric_expectations(batch, col_name)

            # date types
            elif (
                col_semantic_type in ProfilerTypeMapping.DATETIME_TYPE_NAMES
                # TODO these may not be needed
                or isinstance(col_type, sqlalchemy.types.Date)
                or isinstance(col_type, sqlalchemy.types.DateTime)
            ):
                batch = self._create_date_expectations(batch, col_name)

            else:
                print("unknown type", col_name, col_type, col_semantic_type)

        return batch.get_expectation_suite(discard_failed_expectations=False)

    def _create_value_set_expectation(self, batch, col_name):
        #         print('value set', col_name, col_type, col_semantic_type)
        try:
            value_counts = batch.get_column_value_counts(col_name)
            distinct_values = list(pd.DataFrame(value_counts).index)

            # profile distinct values
            batch.expect_column_values_to_be_in_set(col_name, value_set=distinct_values)

            # profile distribution of values
            # this actually works really well!
            weights = [1.0 / len(distinct_values)] * len(distinct_values)
            batch.expect_column_kl_divergence_to_be_less_than(
                column=col_name,
                partition_object={"values": distinct_values, "weights": weights,},
            )
        except KeyError:
            print(f"Column {col_name} was not found in this dataset. Skipping.")
        return batch

    def _create_type_list_expectation(self, batch, col_name, col_type):
        # I'm only looking at data type for all columns since completeness and uniqueness
        # are handled in a different suite, but we could also move those two here if we wanted.
        # profile datatype
        if col_type in ProfilerTypeMapping.BOOLEAN_TYPE_NAMES:
            type_list = ProfilerTypeMapping.BOOLEAN_TYPE_NAMES
        elif col_type in ProfilerTypeMapping.FLOAT_TYPE_NAMES:
            type_list = ProfilerTypeMapping.FLOAT_TYPE_NAMES
        elif col_type in ProfilerTypeMapping.INT_TYPE_NAMES:
            type_list = ProfilerTypeMapping.INT_TYPE_NAMES
        elif col_type in ProfilerTypeMapping.STRING_TYPE_NAMES:
            type_list = ProfilerTypeMapping.STRING_TYPE_NAMES
        elif col_type in ProfilerTypeMapping.DATETIME_TYPE_NAMES:
            type_list = ProfilerTypeMapping.DATETIME_TYPE_NAMES
        else:
            type_list = None
        if type_list is not None:
            batch.expect_column_values_to_be_in_type_list(col_name, type_list=type_list)
        return batch

    def _create_numeric_expectations(self, batch, col_name):
        col_min = batch.get_column_min(col_name)
        col_max = batch.get_column_max(col_name)
        batch.expect_column_values_to_be_between(
            col_name, min_value=col_min, max_value=col_max
        )
        # distribution - doesn't work yet

        # increment = max(1, math.floor((max_value-min_value)/10))
        # values = list(range(min_value,max_value+1, increment))
        # weights = [1.0/len(values) ] * len(values)
        #
        # batch.expect_column_kl_divergence_to_be_less_than(
        #     column=col_name, partition_object={'values': values, 'weights': weights}
        # )

        return batch

    def _create_date_expectations(self, batch, col_name):
        col_min = batch.get_column_min(col_name)
        col_max = batch.get_column_max(col_name)
        batch.expect_column_values_to_be_between(
            col_name, min_value=col_min, max_value=col_max
        )
        # print('date', col_name, col_type, col_semantic_type)
        return batch


if __name__ == "__main__":
    profiler = ExcellentProfiler()
    csv_filename = "opioid-metadata.csv"
    context = DataContext()
    suite_name = "wow_a_really_cool_suite"
    batch_kwargs = batch_kwargs = {
        "data_asset_name": None,
        "datasource": "datas",
        "path": "data/opioids.csv",
    }
    suite = context.create_expectation_suite(suite_name, overwrite_existing=True)
    batch = context.get_batch(batch_kwargs, suite)
    janky_tuple_of_stuff = (batch, csv_filename)
    suite = profiler.profile(janky_tuple_of_stuff, suite_name=suite_name)
    print(suite)
    context.save_expectation_suite(suite)

    batch = context.get_batch(batch_kwargs, suite)
    # NOTE!!! If the run_name is "profiling" it shows up in it's own tab!
    context.run_validation_operator("action_list_operator", [batch], run_name="profiling")
    context.run_validation_operator("action_list_operator", [batch], run_name="look a validation!")
    context.build_data_docs()
    context.open_data_docs()

opioid-metadata.csv

opioids.csv