feat!: Add column statistics to query inspection (#18)

Author: borchero

sqlcompyre/analysis/query_inspection.py

@@ -1,8 +1,11 @@
 # Copyright (c) QuantCo 2024-2025
 # SPDX-License-Identifier: BSD-3-Clause
 
+from __future__ import annotations
 
+from dataclasses import dataclass
 from functools import cached_property, lru_cache
+from typing import Any
 
 import sqlalchemy as sa
 
@@ -15,21 +18,21 @@ class QueryInspection:
         or :meth:`~sqlcompyre.api.inspect_table` functions instead.
     """
 
-    def __init__(self, engine: sa.Engine, selectable: sa.Select):
+    def __init__(self, engine: sa.Engine, query: sa.FromClause):
         """
         Args:
             engine: The engine to use for connecting to the database.
             query: The query whose results to inspect.
         """
         self.engine = engine
-        self.query = selectable
+        self.query = query
 
     @cached_property
     def row_count(self) -> int:
         """Get the number of rows returned by the query."""
         with self.engine.connect() as conn:
             return conn.execute(
-                sa.select(sa.func.count()).select_from(self.query.subquery())
+                sa.select(sa.func.count()).select_from(self.query)
             ).scalar_one()
 
     @lru_cache
@@ -45,14 +48,53 @@ def distinct_row_count(self, *columns: str) -> int:
         """
 
         if len(columns) == 0:
-            data_query = self.query.distinct()
+            data_query = sa.select(self.query).distinct()
 
         else:
-            subquery = self.query.subquery()
             data_query = (
-                sa.select(sa.text(", ".join(columns))).distinct().select_from(subquery)
+                sa.select(sa.text(", ".join(columns)))
+                .distinct()
+                .select_from(self.query)
             )
 
         count_query = sa.select(sa.func.count()).select_from(data_query.subquery())
         with self.engine.connect() as conn:
             return conn.execute(count_query).scalar_one()
+
+    @lru_cache
+    def column_stats(self, column: str) -> ColumnStats:
+        """Obtain statistics about a single column.
+
+        Args:
+            column: The name of the column to obtain information about.
+
+        Returns:
+            An object providing access to column statistics.
+        """
+        return ColumnStats(self.engine, self.query.c[column])
+
+
+# ----------------------------------------- COLUMN STATS ---------------------------------------- #
+
+
+@dataclass
+class ColumnStats:
+    """Obtain statistics about column values in a table."""
+
+    def __init__(self, engine: sa.Engine, column: sa.ColumnElement):
+        self.engine = engine
+        self.column = column
+
+    @cached_property
+    def min(self) -> Any | None:
+        """The minimum value in the column."""
+        query = sa.select(sa.func.min(self.column))
+        with self.engine.connect() as conn:
+            return conn.execute(query).scalar()
+
+    @cached_property
+    def max(self) -> Any | None:
+        """The maximum value in the column."""
+        query = sa.select(sa.func.max(self.column))
+        with self.engine.connect() as conn:
+            return conn.execute(query).scalar()

sqlcompyre/api.py

@@ -12,16 +12,14 @@
 # ---------------------------------------------------------------------------------------------
 
 
-def inspect(
-    engine: sa.Engine, query: sa.Select | sa.FromClause | str
-) -> QueryInspection:
+def inspect(engine: sa.Engine, query: sa.Select | sa.FromClause) -> QueryInspection:
     """Inspect the results of a query in the database.
 
     Args:
         engine: The engine to use to access the database.
         query: The query whose results to inspect. This can either be a SQLAlchemy ``SELECT``
-            statement, a ``FROM`` clause (which includes plain :class:`sqlalchemy.Table` objects),
-            or a SQL query specified as string.
+            statement or a ``FROM`` clause (which includes plain :class:`sqlalchemy.Table`
+            objects).
 
     Returns:
         A query inspection object that can be used to easily gain insights into the query
@@ -31,10 +29,8 @@ def inspect(
         :meth:`inspect_table` if you want to inspect the results of ``SELECT * FROM table`` and
         specify the table as a string.
     """
-    if isinstance(query, str):
-        query = sa.select("*").select_from(sa.text(f"({query}) as anon_1"))
-    elif isinstance(query, sa.FromClause):
-        query = sa.select(query)
+    if isinstance(query, sa.Select):
+        return QueryInspection(engine, query.subquery())
     return QueryInspection(engine, query)
 
 
@@ -63,7 +59,7 @@ def inspect_table(engine: sa.Engine, table: sa.Table | str) -> QueryInspection:
         sa_table = meta.tables[table]
     else:
         sa_table = table
-    return inspect(engine, sa.select(sa_table))
+    return inspect(engine, sa_table)
 
 
 # ---------------------------------------------------------------------------------------------

tests/analysis/query_inspection/test_column_stats.py

@@ -0,0 +1,18 @@
+# Copyright (c) QuantCo 2024-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import sqlalchemy as sa
+
+import sqlcompyre as sc
+
+
+def test_column_stats_min(engine: sa.Engine, table_characters: sa.Table):
+    inspection = sc.inspect_table(engine, table_characters)
+    assert inspection.column_stats("age").min == 6
+    assert inspection.column_stats("first_name").min == "Daisy"
+
+
+def test_column_stats_max(engine: sa.Engine, table_characters: sa.Table):
+    inspection = sc.inspect_table(engine, table_characters)
+    assert inspection.column_stats("age").max == 65
+    assert inspection.column_stats("first_name").max == "Scrooge"

tests/analysis/query_inspection/test_row_counts.py

@@ -28,22 +28,6 @@ def test_row_count_query(engine: sa.Engine, table_characters: sa.Table):
     assert inspection.row_count == 6
 
 
-def test_row_count_raw_query(engine: sa.Engine, table_characters: sa.Table):
-    inspection = sc.inspect(
-        engine,
-        f"""
-        SELECT characters.first_name, characters.last_name, characters.age
-        FROM {str(table_characters)}
-        WHERE characters.last_name = 'Duck'
-        """,
-    )
-
-    assert inspection.row_count == 6
-    assert inspection.distinct_row_count() == 5
-    assert inspection.distinct_row_count("last_name") == 1
-    assert inspection.distinct_row_count("last_name", "age") == 3
-
-
 @pytest.mark.parametrize(
     ("columns", "expected"),
     [([], 7), (["last_name"], 3), (["last_name", "age"], 5)],