From 1e6d3979a886382bce53e3dd30c90ce1d4b4381c Mon Sep 17 00:00:00 2001 From: Nick Moore Date: Wed, 25 Sep 2024 07:59:36 +1000 Subject: [PATCH] fix for slightly broken pivotplugin --- countess/plugins/pivot.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/countess/plugins/pivot.py b/countess/plugins/pivot.py index 1c46ebd..98ab837 100644 --- a/countess/plugins/pivot.py +++ b/countess/plugins/pivot.py @@ -106,16 +106,22 @@ def finalize(self) -> Iterable[pd.DataFrame]: # series, and then `.agg()` applies the appropriate aggregation function. # XXX `mean` could get streamlined by recording (sum, count) tuples at the first stage, # and then adding those up here and dividing (but it's a bit tricky) + def _aggfunc(): if self.aggfunc in ("min", "max", "sum"): return self.aggfunc.value return lambda s: s.explode().agg(self.aggfunc.value) - dataframe = pd.concat(self.dataframes) - index_cols = [p.label for p in self.columns if p.value == "Index"] + try: + dataframe = pd.concat(self.dataframes) + columns = get_all_columns(dataframe) + index_cols = [p.label for p in self.columns if p.value == "Index" and p.label in columns] + + if index_cols: + dataframe = dataframe.groupby(by=index_cols, group_keys=True).agg(_aggfunc()) - if index_cols: - dataframe = dataframe.groupby(by=index_cols, group_keys=True).agg(_aggfunc()) + yield dataframe + except (KeyError, ValueError, TypeError) as exc: + logger.warning("Exception", exc_info=exc) - yield dataframe yield from super().finalize()