Repository: marcua/datools (language: python)

Issue:
Column statistics for enums (char/int) and ranges (%tiles for int/float)



Discussion / hints:

Fix patch:
diff --git a/datools/models.py b/datools/models.py
index de9896c..3bde11f 100644
--- a/datools/models.py
+++ b/datools/models.py
@@ -16,6 +16,12 @@
 class Column:
     name: str
 
+    def __hash__(self):
+        return self.name.__hash__()
+
+    def __eq__(self, other):
+        return self.name.__eq__(other.name)
+
 
 @dataclass
 class StringConstant:
diff --git a/datools/table_statistics.py b/datools/table_statistics.py
new file mode 100644
index 0000000..dcaea9f
--- /dev/null
+++ b/datools/table_statistics.py
@@ -0,0 +1,147 @@
+import sqlalchemy
+
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any
+from typing import Dict
+from typing import Generator
+from typing import List
+from typing import Tuple
+from typing import Set
+
+from datools.models import Column
+from datools.models import Table
+
+
+RANGE_BUCKETS = 3
+RANGE_VALUED_TYPES = {
+    sqlalchemy.sql.sqltypes.BIGINT,
+    sqlalchemy.sql.sqltypes.DATE,
+    sqlalchemy.sql.sqltypes.DATETIME,
+    sqlalchemy.sql.sqltypes.DECIMAL,
+    sqlalchemy.sql.sqltypes.FLOAT,
+    sqlalchemy.sql.sqltypes.INTEGER,
+    sqlalchemy.sql.sqltypes.NUMERIC,
+    sqlalchemy.sql.sqltypes.REAL,
+    sqlalchemy.sql.sqltypes.SMALLINT,
+    sqlalchemy.sql.sqltypes.TIME,
+    sqlalchemy.sql.sqltypes.TIMESTAMP,
+}
+SET_VALUED_TYPES = {
+    sqlalchemy.sql.sqltypes.BOOLEAN,
+    sqlalchemy.sql.sqltypes.CHAR,
+    sqlalchemy.sql.sqltypes.INTEGER,
+    sqlalchemy.sql.sqltypes.NCHAR,
+    sqlalchemy.sql.sqltypes.NVARCHAR,
+    sqlalchemy.sql.sqltypes.SMALLINT,
+    sqlalchemy.sql.sqltypes.TEXT,
+    sqlalchemy.sql.sqltypes.VARCHAR,
+}
+
+
+class ColumnStatistics:
+    pass
+
+
+@dataclass
+class SetValuedStatistics(ColumnStatistics):
+    distinct_values: int
+
+    def __repr__(self):
+        return f'SetValuedStatistics(distinct_values: {self.distinct_values})'
+
+
+@dataclass
+class RangeValuedStatistics(ColumnStatistics):
+    bucket_minimums: List[Any]
+
+    def __repr__(self):
+        return (
+            f'RangeValuedStatistics(bucket_minimums: '
+            f'{[minimum for minimum in self.bucket_minimums]})')
+
+
+def _set_valued_statistics(
+        conn: sqlalchemy.engine.Connection,
+        table: Table,
+        columns: List[Column]
+) -> Generator[Tuple[Column, SetValuedStatistics], None, None]:
+    clauses: list[str] = []
+    for column in columns:
+        clauses.append(
+            f'COUNT(DISTINCT {column.name})'
+            f'AS {column.name}_distinct_values')
+    results = list(conn.execute(
+        f'SELECT {", ".join(clauses)} FROM {table.name}'))[0]
+    for column in columns:
+        yield (
+            column,
+            SetValuedStatistics(results[
+                f'{column.name}_distinct_values']))
+
+
+def _range_valued_statistics(
+        conn: sqlalchemy.engine.Connection,
+        table: Table,
+        columns: List[Column]
+) -> Generator[Tuple[Column, SetValuedStatistics], None, None]:
+    bucket_clauses: list[str] = []
+    first_clauses: list[str] = []
+    for column in columns:
+        bucket_clauses.append(
+            f'{column.name}, '
+            f'NTILE({RANGE_BUCKETS}) OVER (ORDER BY {column.name} ASC) '
+            f'AS {column.name}_bucket')
+        first_clauses.append(
+            f'first_value({column.name}) OVER '
+            f'(PARTITION BY {column.name}_bucket '
+            f' ORDER BY {column.name} ASC '
+            f' RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) '
+            f'AS {column.name}_bucket_value')
+    results = conn.execute(
+        f'WITH buckets AS ( '
+        f'SELECT {", ".join(bucket_clauses)} '
+        f'FROM {table.name} '
+        f')'
+        f'SELECT {", ".join(first_clauses)} '
+        f'FROM buckets ')
+    # The output of the window function is the number of rows in the
+    # input rather than the number of range buckets. Because we have a
+    # single query across multiple columns, we can't `GROUP BY
+    # column_name_bucket` and get the `MIN` NTILE value for each
+    # bucket. Since sqlite doesn't support CUBE/ROLLUP/GROUPING SETS,
+    # we do this grouping manually by creating a `set` of the
+    # `first_value` for each bucket.
+    column_values: Dict[Column, List[Any]] = defaultdict(set)
+    for row in results:
+        for column in columns:
+            column_values[column].add(row[f'{column.name}_bucket_value'])
+    for column in columns:
+        yield (
+            column,
+            RangeValuedStatistics(sorted(column_values[column])))
+
+
+def column_statistics(
+        engine: sqlalchemy.engine.Engine,
+        table: Table,
+        columns_to_ignore: Set[Column],
+) -> Dict[Column, List[ColumnStatistics]]:
+    metadata = sqlalchemy.MetaData()
+    table_metadata = sqlalchemy.Table(
+        table.name, metadata, autoload_with=engine)
+    candidate_columns = [column for column in table_metadata.columns
+                         if column.name not in columns_to_ignore]
+    statistics: Dict[Column, List[ColumnStatistics]] = defaultdict(list)
+    conn = engine.connect()
+    for column, statistic in _set_valued_statistics(
+            conn, table,
+            [Column(column.name) for column in candidate_columns if
+             type(column.type) in SET_VALUED_TYPES]):
+        statistics[column].append(statistic)
+    for column, statistic in _range_valued_statistics(
+            conn, table,
+            [Column(column.name) for column in candidate_columns if
+             type(column.type) in RANGE_VALUED_TYPES]):
+        statistics[column].append(statistic)
+    return statistics


Test patch:
diff --git a/tests/fixtures.py b/tests/fixtures.py
index 4f8056d..2ad1dcd 100644
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -11,7 +11,7 @@
 from sqlalchemy.schema import Table
 
 
-def generate_testdb() -> Engine:
+def generate_scorpion_testdb() -> Engine:
     """Create a test DB from Table 1 of the Scorpion paper.
 
     Rather than make sensor_id a foreign key, we make it a character
@@ -42,11 +42,66 @@ def generate_testdb() -> Engine:
         (datetime(2021, 5, 5, 13), '3', 2.3, 0.5, 80),
     ]
     conn.execute(sensor_readings.insert(), [
-        {'created_at': created_at,
-         'sensor_id': sensor_id,
-         'voltage': voltage,
-         'humidity': humidity,
-         'temperature': temperature}
-        for (created_at, sensor_id, voltage, humidity, temperature) in values
+        dict(zip(
+            ('created_at', 'sensor_id', 'voltage', 'humidity', 'temperature'),
+            value))
+        for value in values
+    ])
+    return engine
+
+
+def generate_synthetic_testdb() -> Engine:
+    """Create a synthetic database with several data types and
+    distributions.
+    For various data types (datetime, integer, string, float), we
+    generate distributions that are:
+    * Unique overall (every value is different)
+    * The same overall (every value is the same)
+    * Unique within a bucket (but repeats across buckets).
+    """
+    engine = create_engine('sqlite://')
+    conn = engine.connect()
+    metadata = MetaData()
+    sensor_readings = Table(
+        'synthetic_data', metadata,
+        Column('id', Integer, primary_key=True),
+        Column('same_datetime', DateTime, nullable=False),
+        Column('unique_datetime', DateTime, nullable=False),
+        Column('bucket_unique_datetime', DateTime, nullable=False),
+        Column('same_string', String, nullable=False),
+        Column('unique_string', String, nullable=False),
+        Column('bucket_unique_string', String, nullable=False),
+        Column('same_float', Float, nullable=False),
+        Column('unique_float', Float, nullable=False),
+        Column('bucket_unique_float', Float, nullable=False),
+        Column('same_int', Integer, nullable=False),
+        Column('unique_int', Integer, nullable=False),
+        Column('bucket_unique_int', Integer, nullable=False),
+    )
+    metadata.create_all(engine)
+    values = []
+    for bucket in range(1, 20):
+        for row in range(1, 10):
+            values.append((
+                datetime(2021, 5, 5, 11),
+                datetime(2021, 5, bucket, row),
+                datetime(2021, 5, 5, row),
+                'hi',
+                f'{bucket}-{row}',
+                f'{row}',
+                1.1,
+                bucket * 1.0 + (row / 100.0),
+                row * 1.0,
+                1,
+                bucket * 1000 + row,
+                row
+            ))
+    conn.execute(sensor_readings.insert(), [
+        dict(zip(
+            ('same_datetime', 'unique_datetime', 'bucket_unique_datetime',
+             'same_string', 'unique_string', 'bucket_unique_string',
+             'same_float', 'unique_float', 'bucket_unique_float',
+             'same_int', 'unique_int', 'bucket_unique_int'), value))
+        for value in values
     ])
     return engine
diff --git a/tests/test_scorpion.py b/tests/test_scorpion.py
index e4e4cca..8ee90ab 100644
--- a/tests/test_scorpion.py
+++ b/tests/test_scorpion.py
@@ -8,17 +8,17 @@
 from datools.models import StringConstant
 from datools.models import Table
 from datools.scorpion.explanations import generate_explanations
-from .fixtures import generate_testdb
+from .fixtures import generate_scorpion_testdb
 
 
 def test_scorpion():
-    engine = generate_testdb()
+    engine = generate_scorpion_testdb()
     assert(
         generate_explanations(
             engine,
-            Table('table'),
+            Table('sensor_readings'),
             (Column('group'), ),
-            (Aggregate(AggregateFunction.AVERAGE, Column('agg')), ),
+            Aggregate(AggregateFunction.AVERAGE, Column('agg')),
             (Predicate(Column('filter'),
                        Operator.EQUALS,
                        StringConstant('one')), ),
diff --git a/tests/test_table_statistics.py b/tests/test_table_statistics.py
new file mode 100644
index 0000000..31201c6
--- /dev/null
+++ b/tests/test_table_statistics.py
@@ -0,0 +1,53 @@
+#!/usr/bin/env python
+
+from datools.models import Column
+from datools.models import Table
+from datools.table_statistics import column_statistics
+from datools.table_statistics import RangeValuedStatistics
+from datools.table_statistics import SetValuedStatistics
+from .fixtures import generate_synthetic_testdb
+
+
+def test_table_statistics():
+    """Tests `column_statistics` on synthetic data in 19 buckets with 9
+    values each.
+    """
+    engine = generate_synthetic_testdb()
+    statistics = column_statistics(engine, Table('synthetic_data'), {})
+    assert {
+        Column('id'):
+            [SetValuedStatistics(171), RangeValuedStatistics([1, 58, 115])],
+        Column('same_datetime'):
+            [RangeValuedStatistics([
+                '2021-05-05 11:00:00.000000'])],
+        Column('unique_datetime'):
+            [RangeValuedStatistics([
+                '2021-05-01 01:00:00.000000',
+                '2021-05-07 04:00:00.000000',
+                '2021-05-13 07:00:00.000000'])],
+        Column('bucket_unique_datetime'):
+            [RangeValuedStatistics([
+                '2021-05-05 01:00:00.000000',
+                '2021-05-05 04:00:00.000000',
+                '2021-05-05 07:00:00.000000'])],
+        Column('same_string'):
+            [SetValuedStatistics(1)],
+        Column('unique_string'):
+            [SetValuedStatistics(171)],
+        Column('bucket_unique_string'):
+            [SetValuedStatistics(9)],
+        Column('same_float'):
+            [RangeValuedStatistics([1.1])],
+        Column('unique_float'):
+            [RangeValuedStatistics([1.01, 7.04, 13.07])],
+        Column('bucket_unique_float'):
+            [RangeValuedStatistics([1.0, 4.0, 7.0])],
+        Column('same_int'):
+            [SetValuedStatistics(1), RangeValuedStatistics([1])],
+        Column('unique_int'):
+            [SetValuedStatistics(171), RangeValuedStatistics([
+                1001, 7004, 13007])],
+        Column('bucket_unique_int'):
+            [SetValuedStatistics(9), RangeValuedStatistics([
+                1, 4, 7])],
+    } == statistics