cache the delayed heavy object so we always use the same key in the t…

…ask graph
scikit-hep · Jan 13, 2024 · 4b9147f · 4b9147f
1 parent b6192a7
commit 4b9147f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/coffea/lookup_tools/lookup_base.py b/src/coffea/lookup_tools/lookup_base.py
@@ -107,11 +107,12 @@ def __call__(self, *args, **kwargs):
 
             out_meta = tomap(self, *tuple([arg._meta for arg in actual_args]))
 
-            delayed_corr = dask.delayed(self)
+            if not hasattr(self, "_delayed_corr"):
+                setattr(self, "_delayed_corr", dask.delayed(self))
 
             return dask_awkward.map_partitions(
                 tomap,
-                delayed_corr,
+                self._delayed_corr,
                 *actual_args,
                 label=dask_label,
                 meta=out_meta,

diff --git a/src/coffea/ml_tools/helper.py b/src/coffea/ml_tools/helper.py
@@ -352,10 +352,11 @@ def __call__(self, wrapper, *args):
         packed_metas = repack(flattened_metas)
 
         wrap_meta = wrap(self, *packed_metas)
-        delayed_wrapper = dask.delayed(self)
+        if not hasattr(self, "_delayed_wrapper"):
+            setattr(self, "_delayed_wrapper", dask.delayed(self))
         arr = dask_awkward.lib.core.map_partitions(
             wrap,
-            delayed_wrapper,
+            self._delayed_wrapper,
             *packed_args,
             label=f"numpy_call_{self.__class__.__name__}_",
             meta=wrap_meta,