Add in a DuckdbSqlPlugin for all the plugins which just do some SQL

Author: nickzoic

countess/core/plugins.py

@@ -156,6 +156,23 @@ def execute(
         raise NotImplementedError(f"{self.__class__}.execute")
 
 
+class DuckdbSqlPlugin(DuckdbSimplePlugin):
+    def execute(
+        self, ddbc: DuckDBPyConnection, source: DuckDBPyRelation, row_limit: Optional[int] = None
+    ) -> Optional[DuckDBPyRelation]:
+        sql = self.sql(source.alias, source.columns)
+        logger.debug(f"{self.__class__}.execute sql %s", sql)
+        if sql:
+            try:
+                return ddbc.sql(sql)
+            except duckdb.duckdb.DatabaseError as exc:
+                logger.warning(exc)
+        return None
+
+    def sql(self, table_name: str, columns: Iterable[str]) -> Optional[str]:
+        raise NotImplementedError(f"{self.__class__}.sql")
+
+
 class DuckdbInputPlugin(DuckdbPlugin):
     num_inputs = 0
 

countess/plugins/correlation.py

@@ -1,16 +1,15 @@
-from typing import Optional
 import logging
-
-from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from typing import Iterable, Optional
 
 from countess import VERSION
-from countess.core.parameters import PerNumericColumnArrayParam, BooleanParam, ColumnOrNoneChoiceParam
-from countess.core.plugins import DuckdbSimplePlugin
+from countess.core.parameters import BooleanParam, ColumnOrNoneChoiceParam, PerNumericColumnArrayParam
+from countess.core.plugins import DuckdbSqlPlugin
 from countess.utils.duckdb import duckdb_escape_identifier, duckdb_escape_literal
 
 logger = logging.getLogger(__name__)
 
-class CorrelationPlugin(DuckdbSimplePlugin):
+
+class CorrelationPlugin(DuckdbSqlPlugin):
     """Correlations"""
 
     name = "Correlation Tool"
@@ -21,28 +20,24 @@ class CorrelationPlugin(DuckdbSimplePlugin):
     columns = PerNumericColumnArrayParam("Columns", BooleanParam("Correlate?", False))
     group = ColumnOrNoneChoiceParam("Group")
 
-    def execute(
-        self, ddbc: DuckDBPyConnection, source: DuckDBPyRelation, row_limit: Optional[int] = None
-    ) -> Optional[DuckDBPyRelation]:
-
+    def sql(self, table_name: str, columns: Iterable[str]) -> Optional[str]:
         grp = duckdb_escape_identifier(self.group.value) if self.group.is_not_none() else None
 
         if sum(1 for c in self.columns.params if c.value) < 2:
             return None
 
-        sql = " union all ".join(f"""
-            select {(grp + ", ") if grp else ""}
-            {duckdb_escape_literal(c1.label)} as column_x,
-            {duckdb_escape_literal(c2.label)} as column_y,
-            corr({duckdb_escape_identifier(c2.label)},{duckdb_escape_identifier(c1.label)}) as correlation_coefficient,
-            covar_pop({duckdb_escape_identifier(c2.label)},{duckdb_escape_identifier(c1.label)}) as covariance_population,
-            regr_r2({duckdb_escape_identifier(c2.label)},{duckdb_escape_identifier(c1.label)}) as pearsons_r2
-            from {source.alias}
-            {("group by "+grp) if grp else ""}
-        """
+        return " union all ".join(
+            f"""
+select {(grp + ", ") if grp else ""}
+{duckdb_escape_literal(c1.label)} as column_x,
+{duckdb_escape_literal(c2.label)} as column_y,
+corr({duckdb_escape_identifier(c2.label)},{duckdb_escape_identifier(c1.label)}) as correlation_coefficient,
+covar_pop({duckdb_escape_identifier(c2.label)},{duckdb_escape_identifier(c1.label)}) as covariance_population,
+regr_r2({duckdb_escape_identifier(c2.label)},{duckdb_escape_identifier(c1.label)}) as pearsons_r2
+from {table_name}
+{("group by "+grp) if grp else ""}
+            """
             for c1 in self.columns.params
             for c2 in self.columns.params
             if c1.value and c2.value and c1.label < c2.label
         )
-        logger.debug("CorrelationPlugin.execute sql %s", sql)
-        return ddbc.sql(sql)

countess/plugins/group_by.py

@@ -1,11 +1,11 @@
 import logging
-from typing import Optional
+from typing import Iterable, Optional
 
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 
 from countess import VERSION
 from countess.core.parameters import BooleanParam, PerColumnArrayParam, TabularMultiParam
-from countess.core.plugins import DuckdbSimplePlugin
+from countess.core.plugins import DuckdbSqlPlugin
 from countess.utils.duckdb import duckdb_escape_identifier
 
 logger = logging.getLogger(__name__)
@@ -30,7 +30,7 @@ def _op(op_name, col_name):
     return f"{op_call}{col_ident}) AS {col_output}"
 
 
-class GroupByPlugin(DuckdbSimplePlugin):
+class GroupByPlugin(DuckdbSqlPlugin):
     """Groups by an arbitrary column and rolls up rows"""
 
     name = "Group By"
@@ -42,9 +42,7 @@ class GroupByPlugin(DuckdbSimplePlugin):
     columns = PerColumnArrayParam("Columns", ColumnMultiParam("Column"))
     join = BooleanParam("Join Back?")
 
-    def execute(
-        self, ddbc: DuckDBPyConnection, source: DuckDBPyRelation, row_limit: Optional[int] = None
-    ) -> Optional[DuckDBPyRelation]:
+    def sql(self, table_name: str, columns: Iterable[str]) -> Optional[str]:
         column_params = list(self.columns.get_column_params())
         columns = (
             ", ".join(
@@ -60,16 +58,16 @@ def execute(
             for col_name, col_param in column_params
             if col_param.params["index"].value
         )
-        if group_by:
-            sql = f"SELECT {group_by}, {columns} FROM {source.alias} GROUP BY {group_by}"
-        else:
-            sql = f"SELECT {columns} FROM {source.alias}"
-
         if self.join:
             if group_by:
-                sql = f"SELECT * FROM {source.alias} JOIN ({sql}) USING ({group_by})"
+                return (
+                    f"SELECT * FROM {table_name} JOIN (SELECT {group_by}, {columns} "
+                    "FROM {table_name} GROUP BY {group_by}) USING ({group_by})"
+                )
             else:
-                sql = f"SELECT * FROM {source.alias} CROSS JOIN ({sql})"
-
-        logger.debug("GroupByPlugin.execute sql %s", sql)
-        return ddbc.sql(sql)
+                return "SELECT * FROM {table_name} CROSS JOIN (SELECT {columns} " "FROM {table_name}"
+        else:
+            if group_by:
+                return f"SELECT {group_by}, {columns} FROM {table_name} GROUP BY {group_by}"
+            else:
+                return f"SELECT {columns} FROM {table_name}"

countess/plugins/join.py

@@ -85,10 +85,15 @@ def execute_multi(
             )
         if row_limit is not None:
             query += f" LIMIT {row_limit}"
-        logger.debug(query)
+
+        logger.debug("JoinPlugin.execute_multi tables[0] %s %d", tables[0].alias, len(tables[0]))
+        logger.debug("JoinPlugin.execute_multi tables[1] %s %d", tables[1].alias, len(tables[1]))
+        logger.debug("JoinPlugin.execute_multi query %s", query)
 
         try:
-            return ddbc.sql(query)
+            rel = ddbc.sql(query)
+            logger.debug("JoinPlugin.execute_multi output %d", len(rel))
+            return rel
         except duckdb.ConversionException as exc:
             logger.info(exc)
             return None

countess/plugins/score_scale.py

@@ -1,7 +1,5 @@
 import logging
-from typing import Optional
-
-from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from typing import Iterable, Optional
 
 from countess import VERSION
 from countess.core.parameters import (
@@ -13,7 +11,7 @@
     StringParam,
     TabularMultiParam,
 )
-from countess.core.plugins import DuckdbSimplePlugin
+from countess.core.plugins import DuckdbSqlPlugin
 from countess.utils.duckdb import duckdb_escape_identifier, duckdb_escape_literal
 
 logger = logging.getLogger(__name__)
@@ -41,7 +39,7 @@ def filter(self):
             raise NotImplementedError()
 
 
-class ScoreScalingPlugin(DuckdbSimplePlugin):
+class ScoreScalingPlugin(DuckdbSqlPlugin):
     name = "Score Scaling"
     description = "Scaled Scores using variant classification"
     version = VERSION
@@ -58,13 +56,11 @@ def __init__(self, *a, **k):
         self.classifiers[0].label = "Scale to 0.0"
         self.classifiers[1].label = "Scale to 1.0"
 
-    def execute(
-        self, ddbc: DuckDBPyConnection, source: DuckDBPyRelation, row_limit: Optional[int] = None
-    ) -> Optional[DuckDBPyRelation]:
+    def sql(self, table_name: str, columns: Iterable[str]) -> Optional[str]:
         score_col_id = duckdb_escape_identifier(self.score_col.value)
         scaled_col_id = duckdb_escape_identifier(self.scaled_col.value)
 
-        all_columns = ",".join("T0." + duckdb_escape_identifier(c) for c in source.columns if self.scaled_col != c)
+        all_columns = ",".join("T0." + duckdb_escape_identifier(c) for c in columns if self.scaled_col != c)
 
         if self.group_col.is_not_none():
             group_col_id = "T0." + duckdb_escape_identifier(self.group_col.value)
@@ -73,17 +69,13 @@ def execute(
 
         c0, c1 = self.classifiers
 
-        sql = f"""
+        return f"""
             select {all_columns}, ({score_col_id} - T1.score_0) / (T1.score_1 - T1.score_0) as {scaled_col_id}
-            from {source.alias} T0 join (
+            from {table_name} T0 join (
                 select {group_col_id} as score_group,
                     median({score_col_id}) filter ({c0.filter()}) as score_0,
                     median({score_col_id}) filter ({c1.filter()}) as score_1
-                from {source.alias} T0
+                from {table_name} T0
                 group by score_group
             ) T1 on ({group_col_id} = T1.score_group)
         """
-
-        logger.debug("ScoreScalingPlugin sql %s", sql)
-
-        return ddbc.sql(sql)

countess/plugins/vampseq.py

@@ -1,11 +1,11 @@
 import logging
-from typing import Optional
+from typing import Iterable, Optional
 
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 
 from countess import VERSION
 from countess.core.parameters import ColumnOrNoneChoiceParam, FloatParam, PerNumericColumnArrayParam, TabularMultiParam
-from countess.core.plugins import DuckdbSimplePlugin
+from countess.core.plugins import DuckdbSqlPlugin
 from countess.utils.duckdb import duckdb_escape_identifier, duckdb_escape_literal
 
 logger = logging.getLogger(__name__)
@@ -15,7 +15,7 @@ class CountColumnParam(TabularMultiParam):
     weight = FloatParam("Weight")
 
 
-class VampSeqScorePlugin(DuckdbSimplePlugin):
+class VampSeqScorePlugin(DuckdbSqlPlugin):
     name = "VAMP-seq Scoring"
     description = "Calculate scores from weighed bin counts"
     version = VERSION
@@ -32,17 +32,15 @@ def prepare(self, ddbc: DuckDBPyConnection, source: Optional[DuckDBPyRelation])
             for n, c in enumerate(count_cols):
                 c.weight.value = (n + 1) / len(count_cols)
 
-    def execute(
-        self, ddbc: DuckDBPyConnection, source: DuckDBPyRelation, row_limit: Optional[int] = None
-    ) -> Optional[DuckDBPyRelation]:
+    def sql(self, table_name: str, columns: Iterable[str]) -> Optional[str]:
         weighted_columns = {
             duckdb_escape_identifier(name): duckdb_escape_literal(param.weight.value)
             for name, param in self.columns.get_column_params()
             if param.weight.value is not None
         }
 
         if not weighted_columns:
-            return source
+            return None
 
         if self.group_col.is_not_none():
             group_col_id = "T0." + duckdb_escape_identifier(self.group_col.value)
@@ -53,14 +51,11 @@ def execute(
         weighted_counts = " + ".join(f"T0.{k} * {v} / T1.{k}" for k, v in weighted_columns.items())
         total_counts = " + ".join(f"T0.{k} / T1.{k}" for k in weighted_columns.keys())
 
-        sql = f"""
+        return f"""
             select T0.*, ({weighted_counts}) / ({total_counts}) as score
-            from {source.alias} T0 join (
+            from {table_name} T0 join (
                 select {group_col_id} as score_group, {sums}
-                from {source.alias} T0
+                from {table_name} T0
                 group by score_group
             ) T1 on ({group_col_id} = T1.score_group)
         """
-
-        logger.debug("VampseqScorePlugin sql %s", sql)
-        return ddbc.sql(sql)

countess/plugins/variant.py

@@ -1,9 +1,7 @@
 import logging
 import string
 from functools import lru_cache
-from typing import Any, Optional
-
-from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from typing import Any, Iterable, Optional
 
 from countess import VERSION
 from countess.core.parameters import (
@@ -17,7 +15,7 @@
     StringCharacterSetParam,
     StringParam,
 )
-from countess.core.plugins import DuckdbParallelTransformPlugin, DuckdbSimplePlugin
+from countess.core.plugins import DuckdbParallelTransformPlugin, DuckdbSqlPlugin
 from countess.utils.duckdb import duckdb_escape_identifier
 from countess.utils.variant import TooManyVariationsException, find_variant_string
 
@@ -135,16 +133,14 @@ def transform(self, data: dict[str, Any]) -> Optional[dict[str, Any]]:
         return data
 
 
-class VariantClassifier(DuckdbSimplePlugin):
+class VariantClassifier(DuckdbSqlPlugin):
     name = "Protein Variant Classifier"
     description = "Classifies protein variants into simple types"
     version = VERSION
 
     variant_col = ColumnChoiceParam("Protein variant Column", "variant")
 
-    def execute(
-        self, ddbc: DuckDBPyConnection, source: DuckDBPyRelation, row_limit: Optional[int] = None
-    ) -> Optional[DuckDBPyRelation]:
+    def sql(self, table_name: str, columns: Iterable[str]) -> Optional[str]:
         variant_col_id = duckdb_escape_identifier(self.variant_col.value)
         output_col_id = duckdb_escape_identifier(self.variant_col + "_type")
 
@@ -154,23 +150,20 @@ def execute(
         # once for each distinct variant string.  Then the cases
         # in the outer select use the parts of the regex match to
         # classify the variant.
-        sql = rf"""
+        return rf"""
             select S.*, case when T.a != '' or T.c == '' and T.e == '=' then 'W'
                when T.c != '' and (T.c = T.e or T.e = '=') then 'S'
                when T.e = 'Ter' or T.e = '*' then 'N'
                when T.c != '' and T.d != '' and T.e != '' then 'M'
                else '?'
             end as {output_col_id}
-            from {source.alias} S join (
+            from {table_name} S join (
                 select {variant_col_id} as z, unnest(regexp_extract(
                     {variant_col_id},
                     '(_?[Ww][Tt])|(p.)?([A-Z][a-z]*)?(\d+)?([A-Z][a-z]*|[=*])?',
                     ['a','b','c','d','e']
                 ))
-                from {source.alias}
+                from {table_name}
                 group by z
             ) T on S.{variant_col_id} = T.z
         """
-
-        logger.debug("VariantClassifier sql %s", sql)
-        return ddbc.sql(sql)