naïve pipeline support and update plugins to new arch

countess/core/pipeline.py

@@ -1,13 +1,15 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional
 
+from more_itertools import interleave_longest
+
 from queue import Empty
 import multiprocessing
 import time
 import psutil
 
 from countess.core.logger import Logger
-from countess.core.plugins import BasePlugin, get_plugin_classes
+from countess.core.plugins import BasePlugin, get_plugin_classes, FileInputMixin
 
 PRERUN_ROW_LIMIT = 100000
 
@@ -57,7 +59,6 @@ def __next__(self):
         raise StopIteration
 
 
-
 @dataclass
 class PipelineNode:
     name: str
@@ -85,22 +86,38 @@ def is_descendant_of(self, node):
 
     def execute(self, logger: Logger, row_limit: Optional[int] = None):
         assert row_limit is None or isinstance(row_limit, int)
-        inputs = {p.name: p.result for p in self.parent_nodes}
-        self.result = []
-        if self.plugin:
-            try:
-                self.result = self.plugin.process_inputs(inputs, logger, row_limit)
-                if isinstance(self.result, (bytes, str)):
-                    pass
-                elif row_limit is not None:
-                    self.result = list(self.result)
-                #if not isinstance(self.result, (bytes, str)) and (row_limit is not None or len(self.child_nodes) > 1):
-                    # XXX freeze to handle fan-out and reloading.
-                #    self.result = list(self.result)
-                elif hasattr(self.result, '__iter__'):
-                    self.result = MultiprocessingProxy(self.result, self.name)
-            except Exception as exc:  # pylint: disable=broad-exception-caught
-                logger.exception(exc)
+
+        if self.plugin is None:
+            self.result = []
+        elif self.result and not self.is_dirty:
+            return
+        elif isinstance(self.plugin, FileInputMixin):
+            num_files = self.plugin.num_files()
+            gg = [
+                    self.plugin.load_file(n, logger, row_limit // num_files)
+                    for n in range(0, num_files)
+            ]
+            self.result = list(interleave_longest(*gg))
+        else:
+            self.plugin.prepare([p.name for p in self.parent_nodes])
+            input_source_data = interleave_longest(*(
+                [ (p.name, r) for r in p.result ]
+                for p in self.parent_nodes
+            ))
+            self.result = list(self.plugin.collect(
+                d
+                for source, data in input_source_data
+                for d in self.plugin.process(data, source, logger)
+            ))
+        for p in self.parent_nodes:
+            r = self.plugin.finished(p.name, logger)
+            if r:
+                self.result += list(r)
+        r = self.plugin.finalize(logger)
+        if r:
+            self.result += list(r)
+
+        self.is_dirty = False
 
     def load_config(self, logger: Logger):
         assert isinstance(self.plugin, BasePlugin)

countess/plugins/csv.py

@@ -143,6 +143,13 @@ def read_file_to_dataframe(self, file_params, logger, row_limit=None):
 
         return df
 
+    def num_files(self):
+        return len(self.parameters["files"])
+
+    def load_file(self, file_number: int, logger: Logger, row_limit: Optional[int] = None) -> Iterable:
+        file_params = self.parameters["files"][file_number]
+        yield self.read_file_to_dataframe(file_params, logger, row_limit)
+
 
 class SaveCsvPlugin(PandasOutputPlugin):
     name = "CSV Save"

countess/plugins/expression.py

@@ -10,6 +10,7 @@ def process(df: pd.DataFrame, codes, logger: Logger):
     for code in codes:
         if not code:
             continue
+
         try:
             result = df.eval(code)
         except Exception as exc:  # pylint: disable=W0718

countess/plugins/fastq.py

@@ -1,13 +1,14 @@
 import gzip
 from itertools import islice
+from typing import Optional, Iterable
 
 import pandas as pd
 from fqfa.fastq.fastq import parse_fastq_reads  # type: ignore
 
 from countess import VERSION
 from countess.core.parameters import BooleanParam, FloatParam
 from countess.core.plugins import PandasInputPlugin
-
+from countess.core.logger import Logger
 
 def _file_reader(file_handle, min_avg_quality, row_limit=None):
     for fastq_read in islice(parse_fastq_reads(file_handle), 0, row_limit):
@@ -54,3 +55,11 @@ def read_file_to_dataframe(self, file_params, logger, row_limit=None):
                 dataframe = dataframe.groupby("sequence").count()
 
         return dataframe
+
+    def num_files(self):
+        return len(self.parameters["files"])
+
+    def load_file(self, file_number: int, logger: Logger, row_limit: Optional[int] = None) -> Iterable:
+        file_params = self.parameters["files"][file_number]
+        yield self.read_file_to_dataframe(file_params, logger, row_limit)
+

countess/plugins/join.py

@@ -2,12 +2,11 @@
 from typing import Iterable, Optional, Union
 
 import pandas as pd
-from moore_itertools import product
 
 from countess import VERSION
 from countess.core.logger import Logger
 from countess.core.parameters import ArrayParam, BooleanParam, ColumnOrIndexChoiceParam, MultiParam
-from countess.core.plugins import PandasBasePlugin
+from countess.core.plugins import PandasBasePlugin, PandasProductPlugin
 from countess.utils.pandas import get_all_columns
 
 
@@ -18,7 +17,7 @@ def _join_how(left_required: bool, right_required: bool) -> str:
         return "right" if right_required else "outer"
 
 
-class JoinPlugin(PandasBasePlugin):
+class JoinPlugin(PandasProductPlugin):
     """Joins Pandas Dataframes"""
 
     name = "Join"
@@ -41,22 +40,51 @@ class JoinPlugin(PandasBasePlugin):
             max_size=2,
         ),
     }
+    join_params = None
+    input_columns_1 = None
+    input_columns_2 = None
+
+    def prepare(self, sources: list[str]):
+        super().prepare(sources)
+
+        assert isinstance(self.parameters["inputs"], ArrayParam)
+        assert len(self.parameters["inputs"]) == 2
+        ip1, ip2 = self.parameters["inputs"]
+        ip1.label = f"Input 1: {sources[0]}"
+        ip2.label = f"Input 2: {sources[1]}"
+
+        self.join_params = {
+            "how": _join_how(ip1.required.value, ip2.required.value),
+            "left_index": ip1.join_on.is_index(),
+            "right_index": ip2.join_on.is_index(),
+            "left_on": None if ip1.join_on.is_index() else ip1.join_on.value,
+            "right_on": None if ip2.join_on.is_index() else ip2.join_on.value,
+        }
+        self.input_columns_1 = {}
+        self.input_columns_2 = {}
+
+    def process_dataframes(self, dataframe1: pd.DataFrame, dataframe2: pd.DataFrame, logger: Logger) -> pd.DataFrame:
+
+        # update columns on inputs, these won't propagate back in the case of multiprocess runs but
+        # they will work in preview mode where we only run this in a single thread.
+
+        self.input_columns_1.update(get_all_columns(dataframe1))
+        self.input_columns_2.update(get_all_columns(dataframe2))
 
-    def join_dataframes(self, dataframe1: pd.DataFrame, dataframe2: pd.DataFrame, join_params) -> pd.DataFrame:
         # "left_on" and "right_on" don't seem to mind if the column
         # is an index, but don't seem to work correctly if the column
         # is part of a multiindex: the other multiindex columns go missing.
-        join1 = join_params.get("left_on")
+        join1 = self.join_params.get("left_on")
         if join1 and dataframe1.index.name != join1:
             drop_index = dataframe1.index.name is None and dataframe1.index.names[0] is None
             dataframe1 = dataframe1.reset_index(drop=drop_index)
 
-        join2 = join_params.get("right_on")
+        join2 = self.join_params.get("right_on")
         if join2 and dataframe2.index.name != join2:
             drop_index = dataframe2.index.name is None and dataframe2.index.names[0] is None
             dataframe2 = dataframe2.reset_index(drop=drop_index)
 
-        dataframe = dataframe1.merge(dataframe2, **join_params)
+        dataframe = dataframe1.merge(dataframe2, **self.join_params)
 
         if self.parameters["inputs"][0]["drop"].value and join1 in dataframe.columns:
             dataframe.drop(columns=join1, inplace=True)
@@ -65,36 +93,9 @@ def join_dataframes(self, dataframe1: pd.DataFrame, dataframe2: pd.DataFrame, jo
 
         return dataframe
 
-    def process_inputs(
-        self, inputs: Mapping[str, Iterable[pd.DataFrame]], logger: Logger, row_limit: Optional[int]
-    ) -> Iterable[pd.DataFrame]:
-        try:
-            input1, input2 = inputs.values()
-        except ValueError:
-            raise NotImplementedError("Only two-way joins implemented at this time")  # pylint: disable=raise-missing-from
-
-        inputs_param = self.parameters["inputs"]
-        assert isinstance(inputs_param, ArrayParam)
-        assert len(inputs_param) == 2
-        ip1, ip2 = inputs_param
-
-        join_params = {
-            "how": _join_how(ip1.required.value, ip2.required.value),
-            "left_index": ip1.join_on.is_index(),
-            "right_index": ip2.join_on.is_index(),
-            "left_on": None if ip1.join_on.is_index() else ip1.join_on.value,
-            "right_on": None if ip2.join_on.is_index() else ip2.join_on.value,
-        }
-
-        input_columns_1 = {}
-        input_columns_2 = {}
-
-        for df_in1, df_in2 in product(input1, input2):
-            input_columns_1.update(get_all_columns(df_in1))
-            input_columns_2.update(get_all_columns(df_in2))
-            df_out = self.join_dataframes(df_in1, df_in2, join_params)
-            if len(df_out):
-                yield df_out
-
-        ip1.set_column_choices(input_columns_1.keys())
-        ip2.set_column_choices(input_columns_2.keys())
+    def finalize(self, logger: Logger) -> None:
+        assert isinstance(self.parameters["inputs"], ArrayParam)
+        assert len(self.parameters["inputs"]) == 2
+        ip1, ip2 = self.parameters["inputs"]
+        ip1.set_column_choices(self.input_columns_1.keys())
+        ip2.set_column_choices(self.input_columns_2.keys())

countess/plugins/python.py

@@ -14,6 +14,9 @@
 
 SIMPLE_TYPES = set((bool, int, float, str, tuple, list))
 
+# XXX should probably actually be based on
+# PandasTransformDictToDictPlugin
+# which is a bit more efficient.
 
 class PythonPlugin(PandasTransformRowToDictPlugin):
     name = "Python Code"
@@ -37,6 +40,10 @@ def process_row(self, row: pd.Series, logger: Logger):
         return dict((k, v) for k, v in row_dict.items() if type(v) in SIMPLE_TYPES)
 
     def process_dataframe(self, dataframe: pd.DataFrame, logger: Logger) -> pd.DataFrame:
+        """Override parent class because we a) want to reset
+        the indexes so we can use their values easily and 
+        b) we don't need to merge afterwards"""
+
         dataframe = dataframe.reset_index(drop=False)
         series = self.dataframe_to_series(dataframe, logger)
         dataframe = self.series_to_dataframe(series)

countess/plugins/regex.py

@@ -46,6 +46,10 @@ class RegexToolPlugin(PandasTransformSingleToTuplePlugin):
 
     compiled_re = None
 
+    def prepare(self, sources: list[str]):
+
+        self.compiled_re = re.compile(self.parameters["regex"].value)
+
     def process_dataframe(self, dataframe: pd.DataFrame, logger: Logger) -> pd.DataFrame:
         df = super().process_dataframe(dataframe, logger)
 
@@ -66,15 +70,6 @@ def process_dataframe(self, dataframe: pd.DataFrame, logger: Logger) -> pd.DataF
 
         return df
 
-    def process_inputs(
-        self, inputs: Mapping[str, Iterable[pd.DataFrame]], logger: Logger, row_limit: Optional[int]
-    ) -> Iterable[pd.DataFrame]:
-        self.compiled_re = re.compile(self.parameters["regex"].value)
-        while self.compiled_re.groups > len(self.parameters["output"].params):
-            self.parameters["output"].add_row()
-
-        return super().process_inputs(inputs, logger, row_limit)
-
     def process_value(self, value: str, logger: Logger) -> Tuple[str]:
         if value is not None:
             try:

pyproject.toml

@@ -22,7 +22,6 @@ classifiers = [
 dependencies = [
     'fqfa~=1.2.3',
     'more_itertools~=9.1.0',
-    'moore-itertools',
     'numpy~=1.24.2',
     'pandas~=2.0.0',
     'rapidfuzz~=2.15.1',