Repository: marcua/datools (language: python)

Issue:
Wrapper to support grouping stats on SQLite/DuckDB/databases that don't natively support cubes and grouping sets



Discussion / hints:

Fix patch:
diff --git a/datools/errors.py b/datools/errors.py
new file mode 100644
index 0000000..0681b06
--- /dev/null
+++ b/datools/errors.py
@@ -0,0 +1,2 @@
+class DatoolsError(Exception):
+    pass
diff --git a/datools/scorpion/__init__.py b/datools/explanation/__init__.py
similarity index 100%
rename from datools/scorpion/__init__.py
rename to datools/explanation/__init__.py
diff --git a/datools/explanation/algorithms.py b/datools/explanation/algorithms.py
new file mode 100644
index 0000000..a579372
--- /dev/null
+++ b/datools/explanation/algorithms.py
@@ -0,0 +1,244 @@
+import sqlalchemy
+
+from math import floor
+from textwrap import dedent
+from textwrap import indent
+from typing import Dict
+from typing import Generator
+from typing import List
+from typing import Set
+from typing import Tuple
+
+from datools.errors import DatoolsError
+from datools.models import Aggregate
+from datools.models import Column
+from datools.models import Constant
+from datools.models import Explanation
+from datools.models import Operator
+from datools.models import Predicate
+from datools.models import Table
+from datools.sqlalchemy_utils import GROUP_COLUMNS_KEY
+from datools.sqlalchemy_utils import GROUPING_ID_KEY
+from datools.sqlalchemy_utils import GROUPING_SETS_KEY
+from datools.sqlalchemy_utils import INDENT
+from datools.sqlalchemy_utils import grouping_sets_query
+from datools.sqlalchemy_utils import query_columns
+from datools.sqlalchemy_utils import query_rows
+from datools.table_statistics import column_statistics
+from datools.table_statistics import RangeValuedStatistics
+from datools.table_statistics import SetValuedStatistics
+
+
+def _single_column_candidate_predicates(
+        engine: sqlalchemy.engine.Engine,
+        table: Table,
+        group_bys: Tuple[Column, ...],
+        aggregate: Aggregate,
+) -> Generator[Tuple[Predicate, ...], None, None]:
+    """Based on column statistics for range- and set-valued columns,
+    returns a list of predicates that either filter on ranges (i.e.,
+    lower <= column AND column < higher) or equality (i.e., column =
+    constant). These candidates can be tested as outlier explanations.
+
+    :returns: A tuple representing a conjunction of
+              predicates (e.g., predicate1 AND predicate2).
+    """
+    columns_to_ignore: Set[Column] = {group_by for group_by in group_bys}
+    columns_to_ignore.add(aggregate.column)
+    statistics = column_statistics(engine, table, columns_to_ignore)
+    for column, statistics_list in statistics.items():
+        for statistic in statistics_list:
+            if isinstance(statistic, SetValuedStatistics):
+                yield from (Predicate(
+                    column,
+                    Operator.EQUALS,
+                    Constant(value)) for value in statistic.popular_values)
+            elif isinstance(statistic, RangeValuedStatistics):
+                if len(statistic.bucket_minimums) == 1:
+                    continue
+                pairs = zip(statistic.bucket_minimums,
+                            statistic.bucket_minimums[1:] + [None])
+                for first, second in pairs:
+                    first_predicate = Predicate(
+                        column, Operator.GTEQ, Constant(first))
+                    if second is None:
+                        yield (first_predicate, )
+                    yield (
+                        first_predicate,
+                        Predicate(column, Operator.LT, Constant(second)))
+
+
+def _explanation_counts_query(
+        engine: sqlalchemy.engine.Engine,
+        relation: str,
+        on_columns: Tuple[Column, ...],
+        min_support_rows: int = None
+) -> Tuple[str, Dict[int, Tuple[str, ...]]]:
+    explanation_query = dedent(
+        f'''
+        SELECT
+            {GROUPING_ID_KEY} AS grouping_id,
+            {GROUP_COLUMNS_KEY},
+            1.0 * COUNT(*) AS explanation_size
+        FROM query
+        GROUP BY {GROUPING_SETS_KEY}
+        ''')
+    if min_support_rows is not None:
+        explanation_query += f'HAVING explanation_size > {min_support_rows}\n'
+
+    group_explanations_query, grouping_set_index = grouping_sets_query(
+        engine,
+        explanation_query,
+        tuple((column, ) for column in on_columns))
+    return dedent(
+        f'''
+        WITH query AS (
+            {relation}
+        )
+        {group_explanations_query}
+        '''), grouping_set_index
+
+
+def _diff_query(
+        test_explanations_query: str,
+        control_explanations_query: str,
+        num_test_rows: float,
+        num_control_rows: float,
+        on_columns: Tuple[Column, ...],
+        min_risk_ratio: float
+) -> str:
+    join_conditions = (
+        ['test.grouping_id = control.grouping_id']
+        + [f'((test.{column.name} = control.{column.name}) '
+           f' OR ((test.{column.name} IS NULL) '
+           f'     AND (control.{column.name} IS NULL)))'
+           for column in on_columns])
+    join_statement = ' AND '.join(join_conditions)
+
+    # TODO(marcua): Consult with someone better at statistics on how
+    # to avoid division by 0 in the risk ratio when a group encompases
+    # the entire relation. For now, make the relation size one larger
+    # than it actually is.
+    adjusted_test_rows = num_test_rows + 1
+    adjusted_control_rows = num_control_rows + 1
+
+    diff_query = dedent(
+        f'''
+        WITH
+        test AS (
+            {indent(test_explanations_query, 3 * INDENT)}
+        ),
+        control AS (
+            {indent(control_explanations_query, 3 * INDENT)}
+        )
+        SELECT
+            test.grouping_id,
+            {', '.join(f'test.{column.name}' for column in on_columns)},
+            test.explanation_size AS test_explanation_size,
+            control.explanation_size AS control_explanation_size,
+            (test.explanation_size / (test.explanation_size
+                                      + COALESCE(control.explanation_size, 0)))
+            /
+            (({adjusted_test_rows} - test.explanation_size)
+             / (({adjusted_test_rows} - test.explanation_size)
+                + ({adjusted_control_rows}
+                   - COALESCE(control.explanation_size, 0)))
+            ) AS risk_ratio
+        FROM test
+        LEFT JOIN control ON {join_statement}
+        WHERE risk_ratio > {min_risk_ratio}
+        ORDER BY risk_ratio DESC
+        ''')
+    return diff_query
+
+
+def diff(
+        engine: sqlalchemy.engine.Engine,
+        test_relation: str,
+        control_relation: str,
+        on_columns: Tuple[Column, ...],
+        min_support: float,
+        min_risk_ratio: float,
+        max_order: int
+) -> List[Explanation]:
+    """
+    Generates candidate explanations for why records are more likely to appear
+    in `test_relation` than `control_relation`. For example, if you provided
+    a test set of customers who are more likely to churn and a control set of
+    customers who are less likely to churn, the algorithm might list predicates
+    related to business category or degree of service.
+
+    This is an implementatin of "DIFF: A Relational Interface
+    for Large-Scale Data Explanation"
+    by Firas Abuzaid, Peter Kraft, Sahaana Suri, Edward Gan,
+    Eric Xu, Atul Shenoy, Asvin Ananthanarayan, John Sheu,
+    Erik Meijer, Xi Wu, Jeff Naughton, Peter Bailis, and Matei Zaharia.
+    http://www.bailis.org/papers/diff-vldb2019.pdf
+
+    Figure 1 of that paper contains a declarative description of the
+    implemented algorithm.
+
+    :param engine: A SQLAlchemy engine.
+    :param test_relation: A SQL query resulting in rows whose presence
+                          you would like to explain.
+    :param control_relation: A SQL query (with the same schema
+                             as `test_relation`) that contains a control
+                             group of rows that are different from/a
+                             baseline for `test_relation`.
+    :param on_columns: A list of columns that might contain an explanation
+                       for the difference between test and control.
+    :param min_support: A value in the range [0, 1] with the minimum fraction
+                        of the test set that an explanation should contain.
+    :param min_risk_ratio: A positive floating point value with the minimum
+                           risk ratio of an explanation. A high risk ratio
+                           means an explanation is more likely to describe a
+                           row in `test_relation` than `control_relation`.
+    :param max_order: The largest number of columns on which to consider
+                      an explanation. For example, a 2-column explanation could
+                      be that (category='dog walker' AND signup_day='sunday').
+                      The current implementation only supports 1-column
+                      explanations.
+    """
+    if max_order != 1:
+        raise DatoolsError('Only one-column predicates are supported for now')
+
+    # Get all column names from test_relation and control_relation,
+    # ensure they are the same.
+    # TODO(marcua): compare types.
+    test_column_names = query_columns(engine, test_relation)
+    control_column_names = query_columns(engine, control_relation)
+    if test_column_names != control_column_names:
+        raise DatoolsError(
+            'test_relation and control_relation have different schemas')
+
+    # Ensure on_columns are a subset of the test/control columns.
+    on_column_names = tuple(column.name for column in on_columns)
+    if set(on_column_names) - set(test_column_names):
+        raise DatoolsError('on_columns is not a subset of test_relation')
+
+    # Get size of test_relation, control_relation.
+    num_test_rows = 1.0 * query_rows(engine, test_relation)
+    num_control_rows = 1.0 * query_rows(engine, control_relation)
+    min_support_rows = floor(num_test_rows * min_support)
+
+    # GROUP BY all test_relation columns, remove ones with a size less
+    # than min_support_rows.
+    test_explanations_query, grouping_set_index = _explanation_counts_query(
+        engine, test_relation, on_columns, min_support_rows)
+    control_explanations_query, _ = _explanation_counts_query(
+        engine, control_relation, on_columns, min_support_rows=None)
+    diff_query = _diff_query(
+        test_explanations_query, control_explanations_query,
+        num_test_rows, num_control_rows,
+        on_columns, min_risk_ratio)
+
+    result = engine.execute(diff_query)
+    explanations = []
+    for row in result:
+        row = dict(row)
+        predicates = tuple(
+            Predicate(column, Operator.EQUALS, row[column.name])
+            for column in grouping_set_index[row['grouping_id']])
+        explanations.append(Explanation(predicates, row['risk_ratio']))
+    result.close()
+    return explanations
diff --git a/datools/models.py b/datools/models.py
index d708d43..8553767 100644
--- a/datools/models.py
+++ b/datools/models.py
@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any
+from typing import Tuple
 
 
 AggregateFunction = Enum(
@@ -45,3 +46,9 @@ class Predicate:
 @dataclass
 class Table:
     name: str
+
+
+@dataclass
+class Explanation:
+    predicates: Tuple[Predicate, ...]
+    risk_ratio: float
diff --git a/datools/scorpion/explanations.py b/datools/scorpion/explanations.py
deleted file mode 100644
index 767ba65..0000000
--- a/datools/scorpion/explanations.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import sqlalchemy
-
-from typing import Generator
-from typing import Set
-from typing import Tuple
-
-from datools.models import Aggregate
-from datools.models import Column
-from datools.models import Constant
-from datools.models import Operator
-from datools.models import Predicate
-from datools.models import Table
-from datools.table_statistics import column_statistics
-from datools.table_statistics import RangeValuedStatistics
-from datools.table_statistics import SetValuedStatistics
-
-
-def _single_column_candidate_predicates(
-        engine: sqlalchemy.engine.Engine,
-        table: Table,
-        group_bys: Tuple[Column, ...],
-        aggregate: Aggregate,
-) -> Generator[Tuple[Predicate, ...], None, None]:
-    """Based on column statistics for range- and set-valued columns,
-    returns a list of predicates that either filter on ranges (i.e.,
-    lower <= column AND column < higher) or equality (i.e., column =
-    constant). These candidates can be tested as outlier explanations.
-
-    :returns: A tuple representing a conjunction of
-              predicates (e.g., predicate1 AND predicate2).
-    """
-    columns_to_ignore: Set[Column] = {group_by for group_by in group_bys}
-    columns_to_ignore.add(aggregate.column)
-    statistics = column_statistics(engine, table, columns_to_ignore)
-    for column, statistics_list in statistics.items():
-        for statistic in statistics_list:
-            if isinstance(statistic, SetValuedStatistics):
-                yield from (Predicate(
-                    column,
-                    Operator.EQUALS,
-                    Constant(value)) for value in statistic.popular_values)
-            elif isinstance(statistic, RangeValuedStatistics):
-                if len(statistic.bucket_minimums) == 1:
-                    continue
-                pairs = zip(statistic.bucket_minimums,
-                            statistic.bucket_minimums[1:] + [None])
-                for first, second in pairs:
-                    first_predicate = Predicate(
-                        column, Operator.GTEQ, Constant(first))
-                    if second is None:
-                        yield (first_predicate, )
-                    yield (
-                        first_predicate,
-                        Predicate(column, Operator.LT, Constant(second)))
-
-
-def generate_explanations(
-        engine: sqlalchemy.engine.Engine,
-        table: Table,
-        group_bys: Tuple[Column, ...],
-        aggregate: Aggregate,
-        outlier_predicates: Tuple[Predicate, ...],
-        holdout_predicates: Tuple[Predicate, ...]
-) -> Tuple[Tuple[Predicate, ...]]:
-    """Generates explanation predicates based on Scorpion (Wu & Madden,
-    VLDB 2013) that cause an aggregate to be different in an outlier
-    set than a holdout set.
-
-    :param connection: A SQLite connection.
-    :param table: The name of a table.
-    :param group_bys: A list of columns to group on.
-    :param aggregate: An aggregate on which we identify an outlier.
-    :param outlier_predicates: Filters on the result that
-        exhibit unexpected conditions.
-    :param holdout_predicates: Filters on the result that
-        exhibit normal conditions.
-    """
-    candidates = _single_column_candidate_predicates(
-        engine, table, group_bys, aggregate)
-    return tuple(candidates)
diff --git a/datools/sqlalchemy_utils.py b/datools/sqlalchemy_utils.py
new file mode 100644
index 0000000..5cb5c7c
--- /dev/null
+++ b/datools/sqlalchemy_utils.py
@@ -0,0 +1,99 @@
+import sqlalchemy
+from tabulate import tabulate
+
+from typing import Dict
+from typing import Tuple
+
+from datools.models import Column
+
+
+GROUP_COLUMNS_KEY = '%%%GROUP_COLUMNS%%%'
+GROUPING_SETS_KEY = '%%%GROUPING_SETS%%%'
+GROUPING_ID_KEY = '%%%GROUPING_ID%%%'
+INDENT = '    '
+
+
+def query_columns(engine: sqlalchemy.engine.Engine, query: str) -> Tuple[str]:
+    results = engine.execute(query)
+    columns = tuple(column[0] for column in results.cursor.description)
+    results.close()
+    return columns
+
+
+def query_results_pretty_print(
+        engine: sqlalchemy.engine.Engine, query: str, label: str = None
+) -> None:
+    if label:
+        print(f'*** {label} ***')
+    result = engine.execute(query)
+    all_rows = (dict(row) for row in result)
+    print(tabulate(all_rows, headers='keys', tablefmt='psql'))
+    result.close()
+
+
+def query_rows(engine: sqlalchemy.engine.Engine, query: str) -> int:
+    count_query = (
+        f'WITH query AS ({query}) '
+        f'SELECT COUNT(*) AS num_rows FROM query')
+    results = engine.execute(count_query)
+    rows = results.first().num_rows
+    results.close()
+    return rows
+
+
+def grouping_sets_query(
+        engine: sqlalchemy.engine.Engine,
+        query: str,
+        sets: Tuple[Tuple[Column, ...]],
+        group_columns_key: str = GROUP_COLUMNS_KEY,
+        grouping_sets_key: str = GROUPING_SETS_KEY,
+        grouping_id_key: str = GROUPING_ID_KEY
+) -> Tuple[str, Dict[int, Tuple[str, ...]]]:
+    """Takes a `query` like
+
+    SELECT
+        {grouping_id_key} AS grouping_id,
+        {group_columns_key},
+        AGGREGATE1(args), AGGREGATE2(args)
+    ...
+    GROUP BY {grouping_sets_key}
+    ...
+
+    and returns a tuple with
+      * a grouping sets query for the grouping sets in `sets`, and
+      * a dictionary mapping set IDs to grouping set columns.
+
+    If the database that `engine` is connected to natively supports
+    grouping sets, utilize the standard SQL syntax for them. If it doesn't,
+    implement the query by capturing the UNION ALL output of multiple
+    GROUP BY subqueries.
+
+    The `grouping_id` will be set to a unique value for each grouping set,
+    and will be deterministic on the iteration order of `sets`.
+    """
+    if engine.url.get_backend_name() == 'postgresql':
+        # TODO(marcua): Implement grouping sets using SQL syntax in
+        # databases that support it.
+        pass
+    column_indices = {}
+    for grouping_set in sets:
+        for column in grouping_set:
+            index = column_indices.get(column)
+            if index is None:
+                column_indices[column] = len(column_indices)
+
+    queries = []
+    set_index: Dict[int, Tuple[str, ...]] = {}
+    for set_id, grouping_set in enumerate(sets):
+        set_index[set_id] = grouping_set
+        group_columns = [f'NULL AS {column.name}'
+                         for column in column_indices.keys()]
+        for column in grouping_set:
+            group_columns[column_indices[column]] = column.name
+        grouping_set_names = tuple(column.name for column in grouping_set)
+        queries.append(
+            query
+            .replace(grouping_id_key, str(set_id))
+            .replace(group_columns_key, ', '.join(group_columns))
+            .replace(grouping_sets_key, ', '.join(grouping_set_names)))
+    return '\nUNION ALL\n'.join(queries), set_index
diff --git a/requirements_dev.txt b/requirements_dev.txt
index 283f5d5..80365cf 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -9,4 +9,5 @@ Sphinx==1.8.5
 twine==1.14.0
 Click==7.0
 pytest==4.6.5
-pytest-runner==5.1
\ No newline at end of file
+pytest-runner==5.1
+tabulate==0.8.9


Test patch:
diff --git a/tests/fixtures.py b/tests/fixtures.py
index 2ad1dcd..bb9ac83 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -12,7 +12,8 @@
 
 
 def generate_scorpion_testdb() -> Engine:
-    """Create a test DB from Table 1 of the Scorpion paper.
+    """Create a test DB from Table 1 of the Scorpion paper (Wu & Madden,
+    VLDB 2013).
 
     Rather than make sensor_id a foreign key, we make it a character
     type to test a wider variety of column types.
diff --git a/tests/test_diff.py b/tests/test_diff.py
new file mode 100644
index 0000000..2ed065c
--- /dev/null
+++ b/tests/test_diff.py
@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+
+from datools.models import Column
+from datools.models import Explanation
+from datools.models import Operator
+from datools.models import Predicate
+from datools.explanation.algorithms import diff
+from .fixtures import generate_scorpion_testdb
+
+
+def test_diff():
+    engine = generate_scorpion_testdb()
+    candidates = diff(
+        engine,
+        'SELECT * FROM sensor_readings WHERE temperature > 50',
+        'SELECT * FROM sensor_readings WHERE temperature <= 50',
+        (Column('created_at'), Column('sensor_id'), Column('voltage'),
+         Column('humidity')),
+        .05,
+        2.0,
+        1)
+    assert(candidates == [
+        Explanation((Predicate(Column('voltage'), Operator.EQUALS, 2.3), ),
+                    risk_ratio=9.0),
+        Explanation((Predicate(Column('sensor_id'), Operator.EQUALS, '3'), ),
+                    risk_ratio=5 + (1.0 / 3))])
diff --git a/tests/test_scorpion.py b/tests/test_scorpion.py
deleted file mode 100644
index bb10073..0000000
--- a/tests/test_scorpion.py
+++ /dev/null
@@ -1,29 +0,0 @@
-#!/usr/bin/env python
-
-from datetime import datetime
-
-from datools.models import Aggregate
-from datools.models import AggregateFunction
-from datools.models import Column
-from datools.models import Operator
-from datools.models import Predicate
-from datools.models import Constant
-from datools.models import Table
-from datools.scorpion.explanations import generate_explanations
-from .fixtures import generate_scorpion_testdb
-
-
-def test_scorpion():
-    engine = generate_scorpion_testdb()
-    candidates = generate_explanations(
-        engine,
-        Table('sensor_readings'),
-        (Column('created_at'), ),  # TODO(marcua): Round to hours.
-        Aggregate(AggregateFunction.AVERAGE, Column('temperature')),
-        (Predicate(Column('created_at'),
-                   Operator.EQUALS,
-                   Constant(datetime(2021, 5, 5, 11))), ),
-        (Predicate(Column('created_at'),
-                   Operator.NOT_EQUALS,
-                   Constant(datetime(2021, 5, 5, 11))), ))
-    assert(len(candidates) == 24)