Merge pull request #413 from lincc-frameworks/fix_calc_nobs_divisions

add force flag to calc_nobs repartitioning
lincc-frameworks · Mar 28, 2024 · 19ff5ad · 19ff5ad
2 parents 750fe4b + fffa3f6
commit 19ff5ad
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/tape/ensemble.py b/src/tape/ensemble.py
@@ -816,7 +816,7 @@ def calc_nobs(self, by_band=False, label="nobs", temporary=True):
                     .reset_index()
                     .pivot_table(values=band_col, index=id_col, columns=band_col, aggfunc="sum"),
                     meta=meta,
-                ).repartition(divisions=self.object.divisions)
+                ).repartition(divisions=self.object.divisions, force=True)
             else:
                 band_counts = (
                     self.source.groupby([self._id_col])[self._band_col]  # group by each object
@@ -852,7 +852,7 @@ def calc_nobs(self, by_band=False, label="nobs", temporary=True):
                 # Map the groupby to each partition
                 counts = self.source.map_partitions(
                     lambda x: x.groupby([id_col])[[band_col]].aggregate("count")
-                ).repartition(divisions=self.object.divisions)
+                ).repartition(divisions=self.object.divisions, force=True)
             else:
                 # Just do a groupby on all source
                 counts = (