batch_looping_profiler.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import great_expectations\n",
    "from great_expectations.profile.basic_dataset_profiler import BasicDatasetProfilerBase\n",
    "from great_expectations.data_context import DataContext\n",
    "from great_expectations.core import ExpectationSuite\n",
    "from great_expectations.core import ExpectationConfiguration\n",
    "from great_expectations.profile.base import (\n",
    "    DatasetProfiler,\n",
    "    ProfilerCardinality,\n",
    "    ProfilerDataType,\n",
    "    ProfilerTypeMapping,\n",
    ")\n",
    "import logging"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Metric Builders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "code_folding": [
     0,
     11,
     29,
     52
    ],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MyRowCountMetricBuilder():\n",
    "    \"\"\"\n",
    "    Leverages Welford's algorithm, a famous algorithm for \n",
    "    calculating a running variance. For numerical\n",
    "    stability, we don't keep track of variance as \n",
    "    we go, rather keeping track of the squares of \n",
    "    the distance from the mean. Then, at the end, \n",
    "    we can calculate the variance\n",
    "    \"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def initialize(cls, batch):\n",
    "        \"\"\"\n",
    "        Initializes an aggregate that will track:\n",
    "            count   the number of data points in the set\n",
    "            mean    the mean of the data set\n",
    "            M2      the squared distance from the mean\n",
    "        as a dictionary\n",
    "        \"\"\"\n",
    "        row_count = batch.get_row_count()\n",
    "        \n",
    "        batch_count = 1\n",
    "        mean = row_count\n",
    "        M2 = 0\n",
    "\n",
    "        return {'batch_count': batch_count, 'mean': mean, 'M2': M2}\n",
    "    \n",
    "    \n",
    "    @classmethod\n",
    "    def update(cls, current_aggregate, batch):\n",
    "        \"\"\"\n",
    "        Adds one data point to an aggregate, tracking:\n",
    "            count   the number of data points in the set\n",
    "            mean    the mean of the data set\n",
    "            M2      the squared distance from the mean\n",
    "        \"\"\"\n",
    "        row_count = batch.get_row_count()\n",
    "        \n",
    "        batch_count = current_aggregate['batch_count']\n",
    "        mean = current_aggregate['mean']\n",
    "        M2 = current_aggregate['M2']\n",
    "\n",
    "        batch_count += 1\n",
    "        delta = row_count - mean\n",
    "        mean += delta / batch_count\n",
    "        delta2 = row_count - mean\n",
    "        M2 += delta * delta2\n",
    "\n",
    "        return {'batch_count': batch_count, 'mean': mean, 'M2': M2}\n",
    "    \n",
    "    \n",
    "    @classmethod\n",
    "    def finalize(cls, current_aggregate):\n",
    "        \"\"\"\n",
    "        Uses an aggregate as defined in cls.init and cls.update\n",
    "        to retrieve mean, variance, and sample variance\n",
    "        of the data set\n",
    "        \"\"\"\n",
    "        import math\n",
    "        \n",
    "        batch_count = current_aggregate['batch_count']\n",
    "        mean = current_aggregate['mean']\n",
    "        M2 = current_aggregate['M2']\n",
    "\n",
    "        if batch_count < 2:\n",
    "            final_aggregate = {\n",
    "                'batch_count': batch_count, \n",
    "                'mean': mean, \n",
    "                'variance': float('nan'), \n",
    "                'sample_variance': float('nan'), \n",
    "                'standard_deviation': float('nan')\n",
    "            }\n",
    "        else: \n",
    "            mean = mean\n",
    "            variance = M2 / batch_count\n",
    "            standard_deviation = math.sqrt(variance)\n",
    "            final_aggregate = {\n",
    "                'batch_count': batch_count, \n",
    "                'mean': mean, \n",
    "                'standard_deviation': standard_deviation\n",
    "            }\n",
    "\n",
    "        return final_aggregate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "code_folding": [
     0,
     11,
     29,
     52
    ],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MyColumnMeanMetricBuilder():\n",
    "    \"\"\"\n",
    "    Leverages Welford's algorithm, a famous algorithm for \n",
    "    calculating a running variance. For numerical\n",
    "    stability, we don't keep track of variance as \n",
    "    we go, rather keeping track of the squares of \n",
    "    the distance from the mean. Then, at the end, \n",
    "    we can calculate the variance\n",
    "    \"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def initialize(cls, batch, column):\n",
    "        \"\"\"\n",
    "        Initializes an aggregate that will track:\n",
    "            count   the number of data points in the set\n",
    "            mean    the mean of the data set\n",
    "            M2      the squared distance from the mean\n",
    "        as a dictionary\n",
    "        \"\"\"\n",
    "        column_mean = batch.get_column_mean(column)\n",
    "        \n",
    "        batch_count = 1\n",
    "        mean = column_mean\n",
    "        M2 = 0\n",
    "\n",
    "        return {'batch_count': batch_count, 'mean': mean, 'M2': M2}\n",
    "    \n",
    "    \n",
    "    @classmethod\n",
    "    def update(cls, current_aggregate, batch, column):\n",
    "        \"\"\"\n",
    "        Adds one data point to an aggregate, tracking:\n",
    "            count   the number of data points in the set\n",
    "            mean    the mean of the data set\n",
    "            M2      the squared distance from the mean\n",
    "        \"\"\"\n",
    "        column_mean = batch.get_column_mean(column)\n",
    "        \n",
    "        batch_count = current_aggregate['batch_count']\n",
    "        mean = current_aggregate['mean']\n",
    "        M2 = current_aggregate['M2']\n",
    "\n",
    "        batch_count += 1\n",
    "        delta = column_mean - mean\n",
    "        mean += delta / batch_count\n",
    "        delta2 = column_mean - mean\n",
    "        M2 += delta * delta2\n",
    "\n",
    "        return {'batch_count': batch_count, 'mean': mean, 'M2': M2}\n",
    "    \n",
    "    \n",
    "    @classmethod\n",
    "    def finalize(cls, current_aggregate):\n",
    "        \"\"\"\n",
    "        Uses an aggregate as defined in cls.init and cls.update\n",
    "        to retrieve mean, variance, and sample variance\n",
    "        of the data set\n",
    "        \"\"\"\n",
    "        import math\n",
    "        \n",
    "        batch_count = current_aggregate['batch_count']\n",
    "        mean = current_aggregate['mean']\n",
    "        M2 = current_aggregate['M2']\n",
    "\n",
    "        if batch_count < 2:\n",
    "            final_aggregate = {\n",
    "                'batch_count': batch_count, \n",
    "                'mean': mean, \n",
    "                'variance': float('nan'), \n",
    "                'sample_variance': float('nan'), \n",
    "                'standard_deviation': float('nan')\n",
    "            }\n",
    "        else: \n",
    "            mean = mean\n",
    "            variance = M2 / batch_count\n",
    "            standard_deviation = math.sqrt(variance)\n",
    "            final_aggregate = {\n",
    "                'batch_count': batch_count, \n",
    "                'mean': mean, \n",
    "                'standard_deviation': standard_deviation\n",
    "            }\n",
    "\n",
    "        return final_aggregate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "code_folding": [
     0,
     22,
     41,
     62
    ],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "class MyColumnValueSetMetricBuilder():\n",
    "    \"\"\"\n",
    "    The aggregate we are building for each column\n",
    "    looks as follows:\n",
    "    \n",
    "    {\n",
    "        'batch_count': number_of_batches,\n",
    "        'total_row_count': aggregate_number_of_rows,\n",
    "        'values':{\n",
    "            'value_1': {\n",
    "                'total_instances': total_across_all_batches,\n",
    "                'batches_found_in': total_batches_found_in\n",
    "            },\n",
    "            'value_2': {\n",
    "                'total_instances': total_across_all_batches,\n",
    "                'batches_found_in': total_batches_found_in\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    \"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def initialize(cls, batch, column):\n",
    "        value_counts = batch.get_column_value_counts(column)\n",
    "        row_count = batch.get_row_count()\n",
    "        \n",
    "        aggregate = {}\n",
    "        aggregate['values'] = {}\n",
    "        for value in value_counts.index:\n",
    "            aggregate['values'][value] = {\n",
    "                'total_instances': value_counts[value],\n",
    "                'batches_found_in': 1\n",
    "            }\n",
    "\n",
    "        aggregate['batch_count'] = 1\n",
    "        aggregate['total_row_count'] = row_count\n",
    "        \n",
    "        return aggregate\n",
    "    \n",
    "    \n",
    "    @classmethod\n",
    "    def update(cls, current_aggregate, batch, column):\n",
    "        value_counts = batch.get_column_value_counts(column)\n",
    "        row_count = batch.get_row_count()\n",
    "        \n",
    "        current_aggregate['batch_count'] += 1\n",
    "        current_aggregate['total_row_count'] += row_count\n",
    "\n",
    "        for value in value_counts.index:\n",
    "            if value in current_aggregate['values']:\n",
    "                current_aggregate['values'][value]['total_instances'] += value_counts[value]\n",
    "                current_aggregate['values'][value]['batches_found_in'] += 1\n",
    "            else:\n",
    "                current_aggregate['values'][value] = {\n",
    "                    'total_instances': value_counts[value],\n",
    "                    'batches_found_in': 1\n",
    "                }\n",
    "        \n",
    "        return current_aggregate\n",
    "    \n",
    "    \n",
    "    @classmethod\n",
    "    def finalize(cls, current_aggregate):\n",
    "        batch_count = current_aggregate['batch_count']\n",
    "        total_row_count = current_aggregate['total_row_count']\n",
    "        \n",
    "        final_aggregate = current_aggregate.copy()\n",
    "        \n",
    "        for value in current_aggregate['values']:  \n",
    "            batches_found_in = current_aggregate['values'][value]['batches_found_in']\n",
    "            total_instances = current_aggregate['values'][value]['total_instances']\n",
    "\n",
    "            batch_frequency = batches_found_in / batch_count\n",
    "            row_frequency = total_instances / total_row_count\n",
    "            \n",
    "            final_aggregate['values'][value] = {\n",
    "                'batch_frequency': batch_frequency,\n",
    "                'row_frequency': row_frequency\n",
    "            }\n",
    "            \n",
    "        return final_aggregate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Expectation Builder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "code_folding": [
     2,
     20,
     48,
     55,
     81,
     111,
     150,
     188
    ],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# TODO: catch BatchKwargsError\n",
    "\n",
    "class BatchLoopingExpectationBuilder():\n",
    "    \"\"\"\n",
    "    Accepts expectation configurations, leverages MetricBuilders \n",
    "    while looping over batches, and returns an expectation\n",
    "    suite containing the specified expectations.\n",
    "    \"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def build_expectations(\n",
    "        cls, \n",
    "        context, \n",
    "        suite_name, \n",
    "        datasource_name, \n",
    "        generator_name, \n",
    "        data_asset_name,\n",
    "        row_count_config=None,\n",
    "        column_mean_configs=None,\n",
    "        column_value_set_configs=None\n",
    "    ):\n",
    "\n",
    "        # initialize suite\n",
    "        \n",
    "        suite = ExpectationSuite(suite_name)\n",
    "        \n",
    "        \n",
    "        # get partition ids\n",
    "        \n",
    "        datasource = context.get_datasource(datasource_name=datasource_name)\n",
    "        generator = datasource.get_batch_kwargs_generator(name=generator_name)\n",
    "        partition_ids = generator.get_available_partition_ids(data_asset_name=data_asset_name)\n",
    "        partition_ids.sort()\n",
    "                \n",
    "        \n",
    "        # loop batches\n",
    "        \n",
    "        for i, partition_id in enumerate(partition_ids):\n",
    "            batch_kwargs = context.build_batch_kwargs(\n",
    "                datasource=datasource_name, \n",
    "                batch_kwargs_generator=generator_name, \n",
    "                data_asset_name=data_asset_name,\n",
    "                partition_id=partition_id\n",
    "            )\n",
    "            batch = context.get_batch(batch_kwargs, suite)\n",
    "            columns = batch.get_table_columns()\n",
    "            \n",
    "            \n",
    "            if row_count_config:\n",
    "                if i == 0:\n",
    "                    row_count_aggregate = MyRowCountMetricBuilder().initialize(batch)\n",
    "                else:\n",
    "                    row_count_aggregate = MyRowCountMetricBuilder().update(row_count_aggregate, batch)\n",
    "                   \n",
    "                \n",
    "            if column_mean_configs:\n",
    "                if i == 0:\n",
    "                    column_mean_first_batch = {\n",
    "                            config['column']: True for config in column_mean_configs\n",
    "                        }\n",
    "                    column_mean_aggregates = {}\n",
    "                \n",
    "                for j, config in enumerate(column_mean_configs):\n",
    "                    column = config['column']\n",
    "                    \n",
    "                    if column in columns:\n",
    "                        if column_mean_first_batch[column]:\n",
    "                            column_mean_aggregates[column] = MyColumnMeanMetricBuilder().initialize(\n",
    "                                batch,\n",
    "                                column\n",
    "                            )\n",
    "                        else:\n",
    "                            column_mean_aggregates[column] = MyColumnMeanMetricBuilder().update(\n",
    "                                column_mean_aggregates[column],\n",
    "                                batch,\n",
    "                                column\n",
    "                            )\n",
    "                        \n",
    "                        column_mean_first_batch[column] = False\n",
    "                \n",
    "                \n",
    "            if column_value_set_configs:\n",
    "                if i == 0:\n",
    "                    column_value_set_first_batch = {\n",
    "                            config['column']: True for config in column_value_set_configs\n",
    "                        }\n",
    "                    column_value_set_aggregates = {}\n",
    "                \n",
    "                \n",
    "                for j, config in enumerate(column_value_set_configs):\n",
    "                    column = config['column']\n",
    "                    \n",
    "                    if column in columns:\n",
    "\n",
    "                        if column_value_set_first_batch[column]:\n",
    "                            column_value_set_aggregates[column] = MyColumnValueSetMetricBuilder().initialize(\n",
    "                                batch, \n",
    "                                column\n",
    "                            )\n",
    "                        else:\n",
    "                            column_value_set_aggregates[column] = MyColumnValueSetMetricBuilder().update(\n",
    "                                column_value_set_aggregates[column],\n",
    "                                batch,\n",
    "                                column\n",
    "                            )      \n",
    "                                    \n",
    "                        column_value_set_first_batch[column] = False\n",
    "        \n",
    "        \n",
    "        # finalize and add expectations\n",
    "        \n",
    "        if row_count_config:\n",
    "            number_of_stds = row_count_config['number_of_stds']\n",
    "            \n",
    "            final_aggregate = MyRowCountMetricBuilder().finalize(row_count_aggregate)\n",
    "        \n",
    "            rc_mean = final_aggregate['mean']\n",
    "            rc_std = final_aggregate['standard_deviation']\n",
    "            \n",
    "        \n",
    "            rc_min_value = round(rc_mean - (number_of_stds * rc_std))   # min value should be an integer\n",
    "            if 'min_value' in row_count_config:                         # check to see if a minimum min value has been specified\n",
    "                if row_count_config['min_value'] < 0:\n",
    "                    raise ValueError('Minimum value for row count should be at least 0')\n",
    "                if rc_min_value < row_count_config['min_value']:\n",
    "                    rc_min_value = row_count_config['min_value']\n",
    "            else:\n",
    "                rc_min_value *= int(rc_min_value >= 0)      # row count should always be >= 0\n",
    "        \n",
    "            rc_max_value = round(rc_mean + (number_of_stds * rc_std))   # max value should be an integer\n",
    "            if 'max_value' in row_count_config:                         # check to see if a maximum max value has been specified\n",
    "                if row_count_config['max_value'] < 0:\n",
    "                    raise ValueError('Maximum value for row count should be at least 0')\n",
    "                if rc_max_value > row_count_config['max_value']:\n",
    "                    rc_max_value = row_count_config['max_value']\n",
    "        \n",
    "            suite.add_expectation(\n",
    "                ExpectationConfiguration(\n",
    "                    expectation_type='expect_table_row_count_to_be_between',\n",
    "                    kwargs={\n",
    "                        'min_value':rc_min_value,\n",
    "                        'max_value':rc_max_value\n",
    "                    },\n",
    "                    meta={\n",
    "                        'BatchLoopingProfiler': final_aggregate\n",
    "                    }\n",
    "                )\n",
    "            )\n",
    "\n",
    "            \n",
    "        if column_mean_configs:\n",
    "            for j, config in enumerate(column_mean_configs):\n",
    "                column = config['column']\n",
    "                \n",
    "                number_of_stds = 2\n",
    "                if 'number_of_stds' in config:\n",
    "                    number_of_stds = config['number_of_stds']\n",
    "\n",
    "                final_aggregate = MyColumnMeanMetricBuilder().finalize(column_mean_aggregates[column])\n",
    "\n",
    "                cm_mean = final_aggregate['mean']\n",
    "                cm_std = final_aggregate['standard_deviation']\n",
    "\n",
    "                cm_min_value = cm_mean - (number_of_stds * cm_std)\n",
    "                if 'min_value' in config:                           # check to see if a minimum min value has been specified\n",
    "                    if cm_min_value < config['min_value']:\n",
    "                        cm_min_value = config['min_value']\n",
    "\n",
    "                cm_max_value = cm_mean + (number_of_stds * cm_std)\n",
    "                if 'max_value' in config:                           # check to see if a maximum max value has been specified\n",
    "                    if cm_max_value > config['max_value']:\n",
    "                        cm_max_value = config['max_value']\n",
    "\n",
    "                suite.add_expectation(\n",
    "                    ExpectationConfiguration(\n",
    "                        expectation_type='expect_column_mean_to_be_between',\n",
    "                        kwargs={\n",
    "                            'column': column,\n",
    "                            'min_value': cm_min_value,\n",
    "                            'max_value': cm_max_value\n",
    "                        },\n",
    "                        meta={\n",
    "                            'BatchLoopingProfiler': final_aggregate\n",
    "                        }\n",
    "                    )\n",
    "                )\n",
    "\n",
    "\n",
    "        if column_value_set_configs:\n",
    "            for j, config in enumerate(column_value_set_configs):\n",
    "                column = config['column']\n",
    "                if 'batch_frequency_threshold' not in config or 'row_frequency_threshold' not in config:\n",
    "                    raise ValueError('Please specify a batch_frequency_threshold and a row_frequency_threshold')\n",
    "                \n",
    "                batch_frequency_threshold = config['batch_frequency_threshold']\n",
    "                row_frequency_threshold = config['row_frequency_threshold']\n",
    "                \n",
    "                final_aggregate = MyColumnValueSetMetricBuilder().finalize(\n",
    "                    column_value_set_aggregates[column]\n",
    "                )\n",
    "                \n",
    "                batch_count = final_aggregate['batch_count']\n",
    "                total_row_count = final_aggregate['total_row_count']\n",
    "                \n",
    "                value_set = []\n",
    "                meta = {\n",
    "                    'batch_count': batch_count,\n",
    "                    'total_row_count': total_row_count,\n",
    "                    'values': {}\n",
    "                }\n",
    "                \n",
    "                for value in final_aggregate['values']:\n",
    "                    batch_frequency = final_aggregate['values'][value]['batch_frequency']\n",
    "                    row_frequency = final_aggregate['values'][value]['row_frequency']\n",
    "                    \n",
    "                    if batch_frequency >= batch_frequency_threshold:\n",
    "                        value_set.append(value)\n",
    "                        meta['values'][value] = {\n",
    "                            'batch_frequency': batch_frequency,\n",
    "                            'row_frequency': row_frequency\n",
    "                        }\n",
    "                    elif row_frequency >= row_frequency_threshold:\n",
    "                        value_set.append(value)\n",
    "                        meta['values'][value] = {\n",
    "                            'batch_frequency': batch_frequency,\n",
    "                            'row_frequency': row_frequency\n",
    "                        }\n",
    "                \n",
    "                suite.add_expectation(\n",
    "                    ExpectationConfiguration(\n",
    "                        expectation_type='expect_column_values_to_be_in_set',\n",
    "                        kwargs={\n",
    "                            'column': column,\n",
    "                            'value_set': value_set\n",
    "                        },\n",
    "                        meta={\n",
    "                            'BatchLoopingProfiler': meta\n",
    "                        }\n",
    "                    )\n",
    "                )\n",
    "        \n",
    "        \n",
    "        return suite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "## Profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "code_folding": [
     2,
     14
    ],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "# TODO: catch BatchKwargsError\n",
    "\n",
    "class BatchLoopingProfiler(BasicDatasetProfilerBase):\n",
    "    \"\"\"\n",
    "    Accepts profiler configurations\n",
    "    Learns column types and cardinalities\n",
    "    \n",
    "    Uses configurations, column types, and cardinalities together\n",
    "    to determine which expectations should be developed for each \n",
    "    column\n",
    "    \n",
    "    Leverages an ExpectationBuilder to build an expectation suite\n",
    "    \"\"\"\n",
    "    @classmethod\n",
    "    def _get_default_config(cls, expectation, column=None):\n",
    "        if expectation == 'expect_table_row_count_to_be_between':\n",
    "            config = {\n",
    "                'number_of_stds': 2,\n",
    "            }\n",
    "        \n",
    "        elif column:\n",
    "            if expectation == 'expect_column_mean_to_be_between':\n",
    "                config = {\n",
    "                    'column': column,\n",
    "                    'number_of_stds': 2,\n",
    "                }\n",
    "            elif expectation == 'expect_column_values_to_be_in_set':\n",
    "                config = {\n",
    "                    'column': column,\n",
    "                    'batch_frequency_threshold': 0.5,\n",
    "                    'row_frequency_threshold': 0.001\n",
    "                }\n",
    "            else:\n",
    "                raise NotImplementedError('The specified column level expectation has not been implemented')\n",
    "                \n",
    "        else:\n",
    "            raise NotImplementedError('The specified table level expectation has not been implemented')\n",
    "        \n",
    "        return config\n",
    "    \n",
    "    \n",
    "    @classmethod\n",
    "    def profile(\n",
    "        cls, \n",
    "        context, \n",
    "        suite_name, \n",
    "        datasource_name, \n",
    "        generator_name, \n",
    "        data_asset_name,\n",
    "        all_columns_config=None,\n",
    "        table_level_config=None,\n",
    "        column_subset_config=None,\n",
    "        individual_columns_config=None\n",
    "    ):\n",
    "        # build default configs\n",
    "        default_table_level_config = {\n",
    "            'expect_table_row_count_to_be_between': cls._get_default_config('expect_table_row_count_to_be_between')\n",
    "        }\n",
    "        \n",
    "        default_column_subset_config = {\n",
    "            ProfilerDataType.INT: [\n",
    "                'expect_column_mean_to_be_between'\n",
    "            ],\n",
    "            ProfilerDataType.FLOAT: [\n",
    "                'expect_column_mean_to_be_between'\n",
    "            ],\n",
    "            ProfilerDataType.STRING: [\n",
    "                'expect_column_values_to_be_in_set'\n",
    "            ],\n",
    "            ProfilerDataType.BOOLEAN: [\n",
    "                'expect_column_values_to_be_in_set'\n",
    "            ],\n",
    "            ProfilerDataType.DATETIME: [],\n",
    "            ProfilerDataType.UNKNOWN: []\n",
    "        }\n",
    "        \n",
    "        \n",
    "        # all_columns_config will be a way for the user to specify expectations\n",
    "        # to be included or excluded for all columns, as well as columns to be\n",
    "        # excluded by the profiler\n",
    "        excluded_columns = None\n",
    "        if all_columns_config:\n",
    "            for item in all_columns_config:\n",
    "                if item == 'excluded_columns':\n",
    "                    excluded_columns = all_columns_config['excluded_columns']\n",
    "                else:\n",
    "                    raise NotImplementedError\n",
    "        \n",
    "        # table_level_config will be a way for the user to specify included\n",
    "        # or excluded table level expectations\n",
    "        if table_level_config:\n",
    "            raise NotImplementedError\n",
    "        \n",
    "        \n",
    "        # column_subset_config will be a way for users to specify\n",
    "        # here's what I want for all columns of type INT, STRING, etc\n",
    "        if column_subset_config:\n",
    "            if 'semantic_types' in column_subset_config:\n",
    "                semantic_types = column_subset_config['semantic_types']\n",
    "                for semantic_type in semantic_types:\n",
    "                    if 'additional_expectations' in semantic_types[semantic_type]:\n",
    "                        raise NotImplementedError\n",
    "                    if 'excluded_expectations' in semantic_types[semantic_type]:\n",
    "                        raise NotImplementedError\n",
    "                    if 'included_expectations' in semantic_types[semantic_type]:\n",
    "                        raise NotImplementedError\n",
    "            if 'regex' in column_subset_config:\n",
    "                raise NotImplementedError\n",
    "                \n",
    "                \n",
    "        # we are going to need to do some sort of join between\n",
    "        # the user specified configs and the default configs\n",
    "        table_level_config = default_table_level_config\n",
    "        column_subset_config = default_column_subset_config\n",
    "        \n",
    "        \n",
    "        # initialize suite\n",
    "        suite = ExpectationSuite(suite_name)\n",
    "        \n",
    "        \n",
    "        # get partition ids\n",
    "        datasource = context.get_datasource(datasource_name=datasource_name)\n",
    "        generator = datasource.get_batch_kwargs_generator(name=generator_name)\n",
    "        partition_ids = generator.get_available_partition_ids(data_asset_name=data_asset_name)\n",
    "        partition_ids.sort()\n",
    "        \n",
    "        \n",
    "        # get final batch\n",
    "        batch_kwargs = context.build_batch_kwargs(\n",
    "            datasource=datasource_name, \n",
    "            batch_kwargs_generator=generator_name, \n",
    "            data_asset_name=data_asset_name,\n",
    "            partition_id=partition_ids[-1]\n",
    "        )\n",
    "        final_batch = context.get_batch(batch_kwargs, suite)\n",
    "        \n",
    "        \n",
    "        # gather column information\n",
    "        columns = final_batch.get_table_columns()\n",
    "        if excluded_columns:\n",
    "            for column in excluded_columns:\n",
    "                if column in columns:\n",
    "                    columns.remove(column)\n",
    "                else:\n",
    "                    raise ValueError('Specified column does not exist in the final batch')\n",
    "        column_types = {\n",
    "            column: cls._get_column_type(final_batch, column) for column in columns\n",
    "        }\n",
    "        \n",
    "        \n",
    "        # instantiate configs\n",
    "        for expectation in table_level_config:\n",
    "            if expectation == 'expect_table_row_count_to_be_between':\n",
    "                row_count_config = table_level_config[expectation]\n",
    "            else:\n",
    "                raise NotImplementedError\n",
    "                \n",
    "        column_mean_configs = []\n",
    "        column_value_set_configs = []\n",
    "        for column in columns:\n",
    "\n",
    "            if individual_columns_config and (column in individual_columns_config):\n",
    "                column_config = individual_columns_config[column]\n",
    "                if 'semantic_type' in column_config:\n",
    "                    column_types[column] = individual_columns_config[column]['semantic_type']\n",
    "                if 'additional_expectations' in column_config:\n",
    "                    raise NotImplementedError\n",
    "                if 'included_expectations' in column_config:\n",
    "                    raise NotImplementedError\n",
    "                if 'excluded_expectations' in column_config:\n",
    "                    raise NotImplementedError\n",
    "\n",
    "            column_type = column_types[column]\n",
    "            expectations = column_subset_config[column_type]\n",
    "            for expectation in expectations:\n",
    "                if expectation == 'expect_column_mean_to_be_between':\n",
    "                    column_mean_configs.append(cls._get_default_config(expectation, column))\n",
    "                elif expectation == 'expect_column_values_to_be_in_set':\n",
    "                    column_value_set_configs.append(cls._get_default_config(expectation, column))\n",
    "                else:\n",
    "                    raise NotImplementedError\n",
    "        \n",
    "        suite = BatchLoopingExpectationBuilder().build_expectations(\n",
    "            context, \n",
    "            expectation_suite_name, \n",
    "            datasource_name, \n",
    "            generator_name, \n",
    "            data_asset_name,\n",
    "            row_count_config=row_count_config,\n",
    "            column_mean_configs=column_mean_configs,\n",
    "            column_value_set_configs=column_value_set_configs\n",
    "        )\n",
    "        \n",
    "        return suite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Using BatchLoopingExpectationBuilder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "code_folding": [],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "context = DataContext()\n",
    "\n",
    "expectation_suite_name='test_suite'\n",
    "datasource_name='covid_data'\n",
    "generator_name='covid_generator'\n",
    "data_asset_name='covid'\n",
    "\n",
    "row_count_config = {\n",
    "    'number_of_stds': 2\n",
    "}\n",
    "\n",
    "column_mean_configs=[\n",
    "    {\n",
    "        'column': 'Confirmed',\n",
    "        'number_of_stds': 2,\n",
    "        'min_value': 0,\n",
    "    },\n",
    "    {\n",
    "        'column': 'Deaths',\n",
    "        'number_of_stds': 2,\n",
    "        'min_value': 0\n",
    "    },\n",
    "    {\n",
    "        'column': 'Recovered',\n",
    "        'number_of_stds': 2,\n",
    "        'min_value': 0\n",
    "    }\n",
    "]\n",
    "\n",
    "column_value_set_configs=[\n",
    "    {\n",
    "        'column': 'Country_Region',\n",
    "        'batch_frequency_threshold': 0.5,\n",
    "        'row_frequency_threshold': 0.001\n",
    "    }\n",
    "]\n",
    "\n",
    "\n",
    "\n",
    "suite = BatchLoopingExpectationBuilder().build_expectations(\n",
    "    context, \n",
    "    expectation_suite_name, \n",
    "    datasource_name, \n",
    "    generator_name, \n",
    "    data_asset_name,\n",
    "    row_count_config=row_count_config,\n",
    "    column_mean_configs=column_mean_configs,\n",
    "    column_value_set_configs=column_value_set_configs\n",
    ")\n",
    "\n",
    "\n",
    "context.save_expectation_suite(suite, expectation_suite_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Using BatchLoopingProfiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "context = DataContext()\n",
    "\n",
    "expectation_suite_name='test_suite_2'\n",
    "datasource_name='covid_data'\n",
    "generator_name='covid_generator'\n",
    "data_asset_name='covid'\n",
    "\n",
    "all_columns_config = {\n",
    "    'excluded_columns': [\n",
    "        'Admin2',\n",
    "        'Last_Update',\n",
    "        'Lat',\n",
    "        'Long_',\n",
    "        'Combined_Key'\n",
    "    ]\n",
    "}\n",
    "\n",
    "\n",
    "individual_columns_config = {\n",
    "    'FIPS': {\n",
    "        'semantic_type': ProfilerDataType.STRING\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "suite = BatchLoopingProfiler().profile(\n",
    "    context, \n",
    "    expectation_suite_name, \n",
    "    datasource_name, \n",
    "    generator_name, \n",
    "    data_asset_name,\n",
    "    all_columns_config=all_columns_config,\n",
    "    individual_columns_config=individual_columns_config\n",
    ")\n",
    "\n",
    "\n",
    "context.save_expectation_suite(suite, expectation_suite_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'local_site': 'file:///Users/rexboyce/profiler-project/great_expectations/uncommitted/data_docs/local_site/index.html'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context.build_data_docs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "# Sketching out config ideas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "domains: \n",
    "    domain1:\n",
    "        columns:\n",
    "            column1, column2, column3\n",
    "    domain2:\n",
    "        domain_builder: \n",
    "            MyDomainBuilder\n",
    "\n",
    "rulesets:\n",
    "    domain1:\n",
    "        rules:\n",
    "            expect_column_mean_to_be_between\n",
    "            rule2\n",
    "            rule3\n",
    "    domain2:\n",
    "        rule_builder: \n",
    "            class_name: MyRuleBuilder\n",
    "            expectation_type: expect_column_mean_to_be_between\n",
    "            previous_value: urn:ge:expectation_suites::expect_column_mean_to_be_between.min_value:domain=domain1\n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "table_level_config=[               # this list is a standard format that is used for all of the configs\n",
    "    'expectation_1',               # list can include just the name\n",
    "    {\n",
    "        'name': 'expectation_2',   # or an entire dict\n",
    "        'parameter1': 'value',\n",
    "        'parameter2': 'value'\n",
    "    },\n",
    "    'expectation_3'\n",
    "]\n",
    "\n",
    "\n",
    "included_columns = [column1, column7]\n",
    "\n",
    "\n",
    "all_columns_config={\n",
    "    'additional_expectations': [   # expectations to include in addition to the defaults\n",
    "        'expectation_1',               # list can include just the name\n",
    "        {\n",
    "            'name': 'expectation_2',   # or an entire dict\n",
    "            'parameter1': 'value',\n",
    "            'parameter2': 'value'\n",
    "        },\n",
    "        'expectation_3'\n",
    "    ],\n",
    "    'included_expectations': None, # might need a better name, this would be an exclusive list\n",
    "    'excluded_expectations': None  # no matter what, do not include this expectation on any columns\n",
    "}\n",
    "\n",
    "\n",
    "column_subset_rules={\n",
    "    'semantic_types': {\n",
    "        'text': {\n",
    "            'included_expectations': None,    # standard list format\n",
    "            'excluded_expectations': None     # standard list format\n",
    "        },\n",
    "        'numeric': config,\n",
    "        'datetime': config,\n",
    "        'boolean': config,\n",
    "    },\n",
    "    'regex': {\n",
    "        '.*(_ct)': {\n",
    "            'included_expectations': None,    # standard list format\n",
    "            'excluded_expectations': None     # standard list format\n",
    "        }\n",
    "    }\n",
    "}\n",
    "    \n",
    "    \n",
    "individual_columns_config={\n",
    "    'column1': {\n",
    "        'semantic_type': ProfilerDataType.STRING,              # user can override 'type' ex: this is an int but treat it as text\n",
    "        'additional_expectations': None,      # standard list format\n",
    "        'included_expectations': None,        # standard list format\n",
    "        'excluded_expectations': None,        # standard list format\n",
    "    },\n",
    "    'column2': {\n",
    "        'semantic_type': ProfilerDataType.FLOAT,\n",
    "        'additional_expectations': [          # standard list format\n",
    "            'expect_column_values_to_be_in_set', \n",
    "            {\n",
    "                'name':'expect_column_mean_to_be_between',\n",
    "                'parameter': value\n",
    "            }\n",
    "        ],\n",
    "        'included_expectations': None,\n",
    "        'excluded_expectations': None,\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Loading Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting an idea of the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = DataContext()\n",
    "\n",
    "suite_name='test_suite'\n",
    "datasource_name='covid_data'\n",
    "generator_name='covid_generator'\n",
    "data_asset_name='covid'\n",
    "\n",
    "suite = ExpectationSuite(suite_name)\n",
    "\n",
    "datasource = context.get_datasource(datasource_name=datasource_name)\n",
    "generator = datasource.get_batch_kwargs_generator(name=generator_name)\n",
    "partition_ids = generator.get_available_partition_ids(data_asset_name=data_asset_name)\n",
    "partition_ids.sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_kwargs = context.build_batch_kwargs(\n",
    "    datasource=datasource_name, \n",
    "    batch_kwargs_generator=generator_name, \n",
    "    data_asset_name=data_asset_name,\n",
    "    partition_id=partition_ids[-1]\n",
    ")\n",
    "batch = context.get_batch(batch_kwargs, suite)\n",
    "batch.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Finagling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "context = DataContext()\n",
    "\n",
    "suite_name='test_suite'\n",
    "datasource_name='covid_data'\n",
    "generator_name='covid_generator'\n",
    "data_asset_name='covid'\n",
    "\n",
    "suite = ExpectationSuite(suite_name)\n",
    "\n",
    "datasource = context.get_datasource(datasource_name=datasource_name)\n",
    "generator = datasource.get_batch_kwargs_generator(name=generator_name)\n",
    "partition_ids = generator.get_available_partition_ids(data_asset_name=data_asset_name)\n",
    "partition_ids.sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for partition_id in partition_ids[39:60]:\n",
    "    batch_kwargs = context.build_batch_kwargs(\n",
    "        datasource=datasource_name, \n",
    "        batch_kwargs_generator=generator_name, \n",
    "        data_asset_name=data_asset_name,\n",
    "        partition_id=partition_id\n",
    "    )\n",
    "    batch = context.get_batch(batch_kwargs, suite)\n",
    "    print(batch.columns)\n",
    "    batch = batch.rename({'Province/State':'Province_State', 'Country/Region':'Country_Region', 'Last Update':'Last_Update'}, axis=1)\n",
    "    print(batch.columns)\n",
    "    batch = batch.set_index('Province_State', drop=True)\n",
    "    batch.to_csv(f'../COVID-19/csse_covid_19_data/csse_covid_19_daily_reports/{partition_id}.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}