add a "collate" plugin

countess/plugins/collate.py

@@ -0,0 +1,84 @@
+from typing import Iterable, List
+
+import numpy as np
+import pandas as pd
+
+from countess import VERSION
+from countess.core.logger import Logger
+from countess.core.parameters import BooleanParam, ChoiceParam, ColumnOrNoneChoiceParam, IntegerParam, PerColumnArrayParam, TabularMultiParam
+from countess.core.plugins import PandasProcessPlugin
+from countess.utils.pandas import flatten_columns, get_all_columns
+
+
+class CollatePlugin(PandasProcessPlugin):
+ """Collates, sorts and selects data"""
+
+ name = "Collate"
+ description = "Collate and sort records by column(s), taking the first N"
+ version = VERSION
+
+ input_columns: dict[str, np.dtype]
+
+ parameters = {
+ "columns": PerColumnArrayParam(
+ "Columns",
+ ChoiceParam("Role", choices=["—", "Group", "Sort (Asc)", "Sort (Desc)"])
+ ),
+ "limit": IntegerParam("First N records", 0),
+ }
+
+ dataframes: List[pd.DataFrame]
+
+ def __init__(self, *a, **k):
+ super().__init__(*a, **k)
+ self.prepare()
+
+ def prepare(self, *_):
+ self.dataframes = []
+ self.input_columns = {}
+
+ def process(self, data: pd.DataFrame, source: str, logger: Logger) -> Iterable:
+ # XXX should do this in two stages: group each dataframe and then combine.
+ # that can wait for a more general MapReduceFinalizePlugin class though.
+ assert self.dataframes is not None
+
+ self.dataframes.append(data)
+ self.input_columns.update(get_all_columns(data))
+ return []
+
+ def sort_and_limit(self, df):
+ column_parameters = list(zip(self.input_columns, self.parameters["columns"]))
+ sort_cols = [col for col, param in column_parameters if param.value.startswith("Sort")]
+ sort_dirs = [
+ param.value.endswith("(Asc)")
+ for param in self.parameters["columns"]
+ if param.value.startswith("Sort")
+ ]
+
+ df = df.sort_values(by=sort_cols, ascending=sort_dirs)
+ if self.parameters["limit"].value > 0:
+ df = df.head(self.parameters["limit"].value)
+ return df
+
+ def finalize(self, logger: Logger) -> Iterable[pd.DataFrame]:
+ assert isinstance(self.parameters["columns"], PerColumnArrayParam)
+ assert self.dataframes
+
+ self.parameters["columns"].set_column_choices(self.input_columns.keys())
+
+ column_parameters = list(zip(self.input_columns, self.parameters["columns"]))
+ group_cols = [col for col, param in column_parameters if param.value == "Group"]
+
+ df = pd.concat(self.dataframes)
+
+ try:
+ if group_cols:
+ df = df.groupby(group_cols, group_keys=False).apply(self.sort_and_limit)
+ else:
+ df = self.sort_and_limit(df)
+
+ print(f">>>>> {type(df)}\n{df}")
+
+ yield df
+ except ValueError as exc:
+ logger.exception(exc)

pyproject.toml

@@ -47,6 +47,7 @@ hdf = [
 ]
 
 [project.entry-points.countess_plugins]
+collate = "countess.plugins.collate:CollatePlugin"
 load_fastq = "countess.plugins.fastq:LoadFastqPlugin"
 load_hdf = "countess.plugins.hdf5:LoadHdfPlugin"
 load_csv = "countess.plugins.csv:LoadCsvPlugin"