Repository: cogent3/cogent3
Pull request #1439 (2023-05-13 01:03:05)

ENH: extend summary methods on annotation db's


Patch:
diff --git a/src/cogent3/core/annotation_db.py b/src/cogent3/core/annotation_db.py
index 2804b31d4..cc252b550 100644
--- a/src/cogent3/core/annotation_db.py
+++ b/src/cogent3/core/annotation_db.py
@@ -613,17 +613,45 @@ class SqliteAnnotationDbMixin:
             num += result
         return num
 
+    def count_distinct(
+        self,
+        *,
+        seqid: bool = False,
+        biotype: bool = False,
+        name: bool = False,
+    ) -> typing.Optional[Table]:
+        """return table of counts of distinct values"""
+        conditions = {k for k, v in locals().items() if k != "self" and v}
+        header = list(conditions)
+        columns = ", ".join(conditions)
+        if not columns:
+            return
+        placeholder = "{}"
+        if len(conditions) == 1:
+            sql_template = f"SELECT {columns}, COUNT(DISTINCT {columns}) FROM {placeholder} GROUP BY {columns};"
+        else:
+            sql_template = f"SELECT {columns}, COUNT(*) as count FROM {placeholder} GROUP BY {columns};"
+
+        data = []
+        for table in self.table_names:
+            sql = sql_template.format(table)
+            if result := self._execute_sql(sql).fetchall():
+                data.extend(tuple(r) for r in result)
+        return Table(header=header + ["count"], data=data)
+
     @property
     def describe(self) -> Table:
         """top level description of the annotation db"""
-        sql_template = "SELECT DISTINCT {} FROM {};"
+        sql_template = "SELECT {}, COUNT(DISTINCT {}) FROM {} GROUP BY {};"
         data = {}
         for column in ("seqid", "biotype"):
-            data[column] = set()
             for table in self.table_names:
-                sql = sql_template.format(column, table)
+                sql = sql_template.format(column, column, table, column)
                 if result := self._execute_sql(sql).fetchall():
-                    data[column] |= {r[column] for r in result}
+                    counts = dict(tuple(r) for r in result)
+                    for distinct, count in counts.items():
+                        key = f"{column}({distinct!r})"
+                        data[key] = data.get(key, 0) + count
 
         row_counts = []
         for table in self.table_names:
@@ -634,9 +662,8 @@ class SqliteAnnotationDbMixin:
 
         table = make_table(
             data={
-                "biotype": list(data.keys())
-                + [f"num_rows({t!r})" for t in self.table_names],
-                "count": [len(v) for v in data.values()] + row_counts,
+                "": list(data.keys()) + [f"num_rows({t!r})" for t in self.table_names],
+                "count": [v for v in data.values()] + row_counts,
             }
         )
         return table


Test patch:
diff --git a/tests/test_core/test_annotation_db.py b/tests/test_core/test_annotation_db.py
index 4fd146f8c..8fa1e14f8 100644
--- a/tests/test_core/test_annotation_db.py
+++ b/tests/test_core/test_annotation_db.py
@@ -62,6 +62,22 @@ def test_gff_describe(gff_db):
     assert isinstance(result, _Table)
 
 
+def test_count_distinct(gff_db):
+    # no arguments, returns None
+    assert gff_db.count_distinct() is None
+
+    # there are 8 biotypes in the c.elegans gff sample, 2 columns
+    # all arguments returns, from our example, all the rows
+    got = gff_db.count_distinct(biotype=True)
+    assert got.shape == (
+        8,
+        2,
+    )
+    # all names unique, 11 records, 4 columns
+    got = gff_db.count_distinct(biotype=True, seqid=True, name=True)
+    assert got.shape == (11, 4)
+
+
 def test_gff_features_matching(gff_db):
     result = list(gff_db.get_features_matching(biotype="CDS"))
     assert len(result)