data table plugin :check:

Author: nickzoic

countess/core/plugins.py

@@ -136,6 +136,19 @@ def execute(self, ddbc: DuckDBPyConnection, source: Optional[DuckDBPyRelation])
         raise NotImplementedError(f"{self.__class__}.execute")
 
 
+class DuckdbInputPlugin(DuckdbPlugin):
+    num_inputs = 0
+
+    def execute_multi(
+        self, ddbc: DuckDBPyConnection, sources: Mapping
+    ) -> Optional[DuckDBPyRelation]:
+        assert len(sources) == 0
+        return self.execute(ddbc, None)
+
+    def execute(self, ddbc: DuckDBPyConnection, source: None) -> Optional[DuckDBPyRelation]:
+        raise NotImplementedError(f"{self.__class__}.execute")
+
+
 class DuckdbStatementPlugin(DuckdbSimplePlugin):
     def statement(self, ddbc: DuckDBPyConnection, source_table_name: str) -> str:
         raise NotImplementedError(f"{self.__class__}.statement")
@@ -156,10 +169,9 @@ class LoadFileMultiParam(MultiParam):
     filename = FileParam("Filename")
 
 
-class DuckdbLoadFilePlugin(DuckdbSimplePlugin):
+class DuckdbLoadFilePlugin(DuckdbInputPlugin):
     files = FileArrayParam("Files", LoadFileMultiParam("File"))
     file_types: Sequence[tuple[str, Union[str, list[str]]]] = [("Any", "*")]
-    num_inputs = 0
 
     def __init__(self, *a, **k):
         super().__init__(*a, **k)
@@ -170,9 +182,7 @@ def filenames_and_params(self):
             for filename in glob.iglob(file_param.filename.value):
                 yield filename, file_param
 
-    def execute(self, ddbc: DuckDBPyConnection, source: Optional[DuckDBPyRelation]) -> Optional[DuckDBPyRelation]:
-        assert source is None
-
+    def execute(self, ddbc: DuckDBPyConnection, source: None) -> Optional[DuckDBPyRelation]:
         filenames_and_params = list(self.filenames_and_params())
 
         cursor = ddbc
@@ -238,19 +248,25 @@ def dropped_columns(self) -> set[str]:
         return set()
 
     def input_columns(self) -> dict[str, str]:
-        raise NotImplementedError(f"{self.__class__}.input_columns")
+        return None
 
     def output_columns(self) -> dict[str, str]:
         raise NotImplementedError(f"{self.__class__}.output_columns")
 
     def execute(self, ddbc, source):
         """Perform a query which calls `self.transform` for every row."""
 
-        escaped_input_columns = {
-            duckdb_escape_identifier(k): str(v).upper()
-            for k, v in self.input_columns().items()
-            if k is not None and v is not None
-        }
+        if self.input_columns() is None:
+            escaped_input_columns = {
+                duckdb_escape_identifier(k): str(v).upper()
+                for k, v in zip(source.columns, source.dtypes)
+            }
+        else:
+            escaped_input_columns = {
+                duckdb_escape_identifier(k): str(v).upper()
+                for k, v in self.input_columns().items()
+                if k is not None and v is not None
+            }
 
         escaped_output_columns = {
             duckdb_escape_identifier(k): str(v).upper()

countess/plugins/data_table.py

@@ -1,6 +1,6 @@
-from typing import Iterable
+import logging
 
-import pandas as pd
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
 
 from countess import VERSION
 from countess.core.parameters import (
@@ -11,16 +11,18 @@
     StringParam,
     TabularMultiParam,
 )
-from countess.core.plugins import PandasInputPlugin
+from countess.core.plugins import DuckdbInputPlugin
+from countess.utils.duckdb import duckdb_escape_literal, duckdb_escape_identifier
 
+logger = logging.getLogger(__name__)
 
 class _ColumnsMultiParam(MultiParam):
     name = StringParam("Name")
-    type = DataTypeChoiceParam("Type", "string")
+    type = DataTypeChoiceParam("Type")
     index = BooleanParam("Index?")
 
 
-class DataTablePlugin(PandasInputPlugin):
+class DataTablePlugin(DuckdbInputPlugin):
     """DataTable"""
 
     name = "DataTable"
@@ -32,7 +34,7 @@ class DataTablePlugin(PandasInputPlugin):
     columns = ArrayParam("Columns", _ColumnsMultiParam("Column"))
     rows = ArrayParam("Rows", TabularMultiParam("Row"))
 
-    show_preview = False
+    #show_preview = False
 
     def fix_columns(self):
         old_rows = self.rows.params
@@ -61,17 +63,22 @@ def set_parameter(self, key: str, *a, **k):
             self.fix_columns()
         super().set_parameter(key, *a, **k)
 
-    def finalize(self) -> Iterable[pd.DataFrame]:
+    def execute(self, ddbc: DuckDBPyConnection, source: None) -> DuckDBPyRelation:
         self.fix_columns()
-        values = []
-        for row in self.rows:
-            values.append({str(col.name): row[str(col.name)].value for col in self.columns})
 
-        df = pd.DataFrame(values)
-
-        index_cols = [str(col.name) for col in self.columns if col.index]
+        if len(self.rows) == 0:
+            return None
+
+        sql = ("SELECT * FROM (VALUES " + 
+            (", ".join(
+                "(" + (", ".join(duckdb_escape_literal(val.value) for val in row.values())) + ")"
+                for row in self.rows
+            )) +
+            ") _(" +
+            (", ".join(duckdb_escape_identifier(col.name.value) for col in self.columns)) +
+            ")"
+        )
 
-        if index_cols:
-            df = df.set_index(index_cols)
+        logger.debug("DataTablePlugin.execute sql %s", sql)
 
-        yield df
+        return ddbc.sql(sql)