Merge pull request #1016 from CoffeaTeam/fix_serialize_taskgraph

fix: force application of to_packed in awkward binnings
CoffeaTeam · Jan 30, 2024 · 57c6441 · 57c6441
2 parents 757edc3 + 7d9fa9d
commit 57c6441
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/src/coffea/lookup_tools/txt_converters.py b/src/coffea/lookup_tools/txt_converters.py
@@ -192,7 +192,9 @@ def _build_standard_jme_lookup(
                     theBins = numpy.union1d(binMins, binMaxs[-1:])
                 allBins = numpy.append(allBins, theBins)
                 counts = numpy.append(counts, theBins.size)
-            bins[layout[i + offset_name]] = awkward.unflatten(allBins, counts)
+            bins[layout[i + offset_name]] = awkward.to_packed(
+                awkward.unflatten(allBins, counts)
+            )
         bin_order.append(layout[i + offset_name])
         offset_col += 1
 
@@ -214,11 +216,15 @@ def _build_standard_jme_lookup(
     for i in range(nEvalVars):
         var_order.append(layout[i + offset_name])
         if not interpolatedFunc:
-            clamp_mins[layout[i + offset_name]] = awkward.unflatten(
-                numpy.atleast_1d(pars[columns[i + offset_col]]), jagged_counts
+            clamp_mins[layout[i + offset_name]] = awkward.to_packed(
+                awkward.unflatten(
+                    numpy.atleast_1d(pars[columns[i + offset_col]]), jagged_counts
+                )
             )
-            clamp_maxs[layout[i + offset_name]] = awkward.unflatten(
-                numpy.atleast_1d(pars[columns[i + offset_col + 1]]), jagged_counts
+            clamp_maxs[layout[i + offset_name]] = awkward.to_packed(
+                awkward.unflatten(
+                    numpy.atleast_1d(pars[columns[i + offset_col + 1]]), jagged_counts
+                )
             )
             assert awkward.is_valid(clamp_mins[layout[i + offset_name]])
             assert awkward.is_valid(clamp_maxs[layout[i + offset_name]])
@@ -229,7 +235,9 @@ def _build_standard_jme_lookup(
     param_order = []
     offset_col = 2 * nBinnedVars + 1 + int(not interpolatedFunc) * 2 * nEvalVars
     for i in range(nParams):
-        jag = awkward.unflatten(pars[columns[i + offset_col]], jagged_counts)
+        jag = awkward.to_packed(
+            awkward.unflatten(pars[columns[i + offset_col]], jagged_counts)
+        )
         assert awkward.is_valid(jag)
         params.append(jag)
         param_order.append("p%i" % (i))
@@ -505,7 +513,9 @@ def convert_effective_area_file(eaFilePath):
                 theBins = numpy.union1d(binMins, binMaxs)
                 allBins = numpy.append(allBins, theBins)
                 counts = numpy.append(counts, theBins.size)
-            bins[layout[i + offset_name]] = awkward.unflatten(allBins, counts)
+            bins[layout[i + offset_name]] = awkward.to_packed(
+                awkward.unflatten(allBins, counts)
+            )
         bin_order.append(layout[i + offset_name])
         offset_col += 1
 

diff --git a/tests/test_jetmet_tools.py b/tests/test_jetmet_tools.py
@@ -715,12 +715,16 @@ def test_jet_resolution_sf_2d(optimization_enabled):
 def test_corrected_jets_factory(optimization_enabled):
     import os
 
+    from distributed import Client
+
     from coffea.jetmet_tools import CorrectedJetsFactory, CorrectedMETFactory, JECStack
 
     events = None
     from coffea.nanoevents import NanoEventsFactory
 
-    with dask.config.set({"awkward.optimization.enabled": True}):
+    with Client(), dask.config.set(
+        {"awkward.optimization.enabled": optimization_enabled}
+    ):
         events = NanoEventsFactory.from_root(
             {os.path.abspath("tests/samples/nano_dy.root"): "Events"},
             metadata={},