expression plugin!

Author: nickzoic

countess/plugins/expression.py

@@ -1,57 +1,94 @@
 import logging
+from typing import Optional
+import ast
+import re
 
-import pandas as pd
-
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from countess import VERSION
 from countess.core.parameters import BooleanParam, PerColumnArrayParam, TextParam
-from countess.core.plugins import PandasSimplePlugin
+from countess.core.plugins import DuckdbSimplePlugin
+from countess.utils.duckdb import duckdb_escape_identifier, duckdb_escape_literal
 
 logger = logging.getLogger(__name__)
 
-
-def process(df: pd.DataFrame, codes):
-    for code in codes:
-        if not code:
-            continue
-
-        try:
-            result = df.eval(code)
-        except Exception as exc:  # pylint: disable=W0718
-            logger.warning("Exception", exc_info=exc)
-            continue
-
-        if isinstance(result, pd.Series):
-            # this was a filter
-            df = df.copy()
-            df["__filter"] = result
-            df = df.query("__filter != 0").drop(columns="__filter")
-        else:
-            # this was a column assignment
-            df = result
-
-    return df
-
-
-class ExpressionPlugin(PandasSimplePlugin):
+UNOPS = { ast.UAdd: '+', ast.USub: '-' }
+BINOPS = { ast.Add: '+', ast.Mult: '*', ast.Div: '/', ast.Sub: '-', ast.FloorDiv: '//', ast.Mod: '%', ast.Pow: '**' }
+FUNCOPS = {'abs', 'len', 'sin', 'cos', 'tan', 'sqrt', 'log', 'log2', 'log10', 'pow', 'exp' }
+COMPOP = { ast.Eq: '=', ast.NotEq: '!=', ast.Lt: '<', ast.LtE: '<=', ast.Gt: '>', ast.GtE: '>=' }
+
+def _transmogrify(ast_node):
+    """Transform an AST node back into a string which can be parsed by DuckDB's expression
+    parser.  This is a pretty small subset of all the things you might write but on the
+    other hand it saved actually writing a parser."""
+    # XXX might have to write a parser anyway since the AST parser handles decimal
+    # literals badly.  Worry about that later.
+    if type(ast_node) is ast.Name:
+        return duckdb_escape_identifier(ast_node.id)
+    elif type(ast_node) is ast.Constant:
+        return duckdb_escape_literal(ast_node.value)
+    elif type(ast_node) is ast.UnaryOp and type(ast_node.op) in UNOPS:
+        return "(" + UNOPS[type(ast_node.op)] + _transmogrify(ast_node.operand) + ")"
+    elif type(ast_node) is ast.BinOp and type(ast_node.op) in BINOPS:
+        return "(" + _transmogrify(ast_node.left) +  BINOPS[type(ast_node.op)] + _transmogrify(ast_node.right) + ")"
+    elif type(ast_node) is ast.Compare and all(type(op) in COMPOP for op in ast_node.ops):
+        args = [ _transmogrify(x) for x in [ ast_node.left ] + ast_node.comparators ]
+        return "(" + (" AND ".join(
+            args[num] + COMPOP[type(op)] + args[num+1]
+            for num, op in enumerate(ast_node.ops)
+        )) + ")"
+    elif type(ast_node) is ast.Call and ast_node.func.id in FUNCOPS:
+        args = [ _transmogrify(x) for x in ast_node.args ]
+        return ast_node.func.id + "(" + (",".join(args)) + ")"
+    else:
+        raise NotImplementedError(f"Unknown Node {ast_node}")
+
+
+class ExpressionPlugin(DuckdbSimplePlugin):
     name = "Expression"
-    description = "Apply simple expressions"
+    description = "Apply simple expressions to each row"
     version = VERSION
 
     code = TextParam("Expressions")
     drop = PerColumnArrayParam("Drop Columns", BooleanParam("Drop"))
+    projection = None
 
-    def process_dataframe(self, dataframe: pd.DataFrame) -> pd.DataFrame:
-        codes = [c.replace("\n", " ").strip() for c in str(self.code).split("\n\n")]
-        df = process(dataframe, codes)
-
-        drop_names = [label for label, param in self.drop.get_column_params() if param.value]
-
-        drop_indexes = [col for col in drop_names if col in df.index.names]
-        if drop_indexes:
-            df = df.reset_index(drop_indexes, drop=True)
-
-        drop_columns = [col for col in drop_names if col in df.columns]
-        if drop_columns:
-            df.drop(columns=drop_columns, inplace=True)
+    def prepare(self, *a) -> None:
+        super().prepare(*a)
+        self.projection = []
+        try:
+            ast_root = ast.parse(self.code.value or "")
+        except SyntaxError as exc:
+            logger.debug("Syntax Error %s", exc)
+
+        for ast_node in ast_root.body:
+            try:
+                if type(ast_node) is ast.Assign:
+                    for ast_target in ast_node.targets:
+                        if type(ast_target) is ast.Name:
+                            expr = _transmogrify(ast_node.value)
+                            tgt = duckdb_escape_identifier(ast_target.id)
+                            self.projection.append(f"{expr} AS {tgt}")
+                elif type(ast_node) is ast.Expr:
+                    tgt = duckdb_escape_identifier(re.sub(r'_+$', '', re.sub(r'\W+', '_', ast.unparse(ast_node))))
+                    expr = _transmogrify(ast_node.value)
+                    self.projection.append(f"{expr} AS {tgt}")
+
+            except (NotImplementedError, KeyError) as exc:
+                logger.debug("Bad AST Node: %s %s", ast_node, exc)
+
+    def execute(self, ddbc: DuckDBPyConnection, source: Optional[DuckDBPyRelation]) -> Optional  [DuckDBPyRelation]:
+
+        column_params = dict(self.drop.get_column_params())
+        if any(v.value for k, v in column_params.items()):
+            projection = [
+                duckdb_escape_identifier(column)
+                for column in source.columns
+                if not column_params[column].value
+            ]
+        else:
+            projection = [ '*' ]
 
-        return df
+        sql = ', '.join(projection + self.projection)
+        print(">>> " + sql)
+        logger.debug("ExpressionPlugin.execute %s", sql)
+        return source.project(sql)

pyproject.toml

@@ -57,7 +57,7 @@ load_fastq = "countess.plugins.fastq:LoadFastqPlugin"
 load_fasta = "countess.plugins.fastq:LoadFastaPlugin"
 load_csv = "countess.plugins.csv:LoadCsvPlugin"
 group_by = "countess.plugins.group_by:GroupByPlugin"
-#expression = "countess.plugins.expression:ExpressionPlugin"
+expression = "countess.plugins.expression:ExpressionPlugin"
 pivot = "countess.plugins.pivot:PivotPlugin"
 join = "countess.plugins.join:JoinPlugin"
 #python = "countess.plugins.python:PythonPlugin"