force compute before adding to agg dataframe

ks905383 · Feb 20, 2025 · c6bbd0f · c6bbd0f
1 parent 04bb543
commit c6bbd0f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/xagg/core.py b/xagg/core.py
@@ -862,6 +862,10 @@ def aggregate(ds,wm,impl=None,silent=None):
                                    output_dtypes = [float]).to_dataset(name=var)
                                 for var in ds if (('bnds' not in ds[var].sizes) and ('loc' in ds[var].sizes))])
 
+                # Trigger computation, otherwise it separately triggers computation
+                # for every poly_idx below, with tremendous overhead 
+                agg = agg.compute()
+
                 # Put in `aggregated` format (TODO: should reform this to remove some of the 
                 # nested lists that are here for some reason)
                 wm.agg[var] = [[[agg[var].sel(poly_idx = idx).values]] for idx in wm.agg.index]