Merge pull request #1047 from CoffeaTeam/fix-uproot-530-changes

fix: adjust to breaking changes in form mapping interface

pyproject.toml

@@ -38,7 +38,7 @@ classifiers = [
 ]
 dependencies = [
   "awkward>=2.5.2",
-  "uproot>=5.2.2",
+  "uproot>=5.3.0",
   "dask[array]<2024.2.0,>=2023.4.0",
   "dask-awkward>=2024.2.0",
   "dask-histogram>=2024.2.0",

src/coffea/nanoevents/factory.py

@@ -121,7 +121,16 @@ def __call__(self, form):
             self,
         )
 
-    def load_buffers(self, tree, keys, start, stop, interp_options):
+    def load_buffers(
+        self,
+        tree,
+        keys,
+        start,
+        stop,
+        decompression_executor,
+        interpretation_executor,
+        interp_options,
+    ):
         from functools import partial
 
         from coffea.nanoevents.util import tuple_to_key
@@ -139,6 +148,8 @@ def load_buffers(self, tree, keys, start, stop, interp_options):
             cache={},
             access_log=None,
             use_ak_forth=True,
+            decompression_executor=decompression_executor,
+            interpretation_executor=interpretation_executor,
         )
         mapping.preload_column_source(partition_key[0], partition_key[1], tree)
         buffer_key = partial(self._key_formatter, tuple_to_key(partition_key))
@@ -219,6 +230,8 @@ def from_root(
         use_ak_forth=True,
         delayed=True,
         known_base_form=None,
+        decompression_executor=None,
+        interpretation_executor=None,
     ):
         """Quickly build NanoEvents from a root file
 
@@ -256,6 +269,10 @@ def from_root(
                 Nanoevents will use dask as a backend to construct a delayed task graph representing your analysis.
             known_base_form:
                 If the base form of the input file is known ahead of time we can skip opening a single file and parsing metadata.
+            decompression_executor (None or Executor with a ``submit`` method):
+                see: https://github.com/scikit-hep/uproot5/blob/main/src/uproot/_dask.py#L109
+            interpretation_executor (None or Executor with a ``submit`` method):
+                see: https://github.com/scikit-hep/uproot5/blob/main/src/uproot/_dask.py#L113
         """
 
         if treepath is not uproot._util.unset and not isinstance(
@@ -304,6 +321,8 @@ def from_root(
                 filter_branch=_remove_not_interpretable,
                 steps_per_file=steps_per_file,
                 known_base_form=known_base_form,
+                decompression_executor=decompression_executor,
+                interpretation_executor=interpretation_executor,
                 **uproot_options,
             )
 

src/coffea/nanoevents/mapping/uproot.py

@@ -102,9 +102,23 @@ class UprootSourceMapping(BaseSourceMapping):
     _fix_awkward_form_of_iter = False
 
     def __init__(
-        self, fileopener, start, stop, cache=None, access_log=None, use_ak_forth=False
+        self,
+        fileopener,
+        start,
+        stop,
+        cache=None,
+        access_log=None,
+        use_ak_forth=False,
+        decompression_executor=None,
+        interpretation_executor=None,
     ):
         super().__init__(fileopener, start, stop, cache, access_log, use_ak_forth)
+        self.decompression_executor = (
+            decompression_executor or uproot.source.futures.TrivialExecutor()
+        )
+        self.interpretation_executor = (
+            interpretation_executor or uproot.source.futures.TrivialExecutor()
+        )
 
     @classmethod
     def _extract_base_form(cls, tree, iteritems_options={}):
@@ -195,8 +209,8 @@ def extract_column(
             interp,
             entry_start=start,
             entry_stop=stop,
-            decompression_executor=uproot.source.futures.TrivialExecutor(),
-            interpretation_executor=uproot.source.futures.TrivialExecutor(),
+            decompression_executor=self.decompression_executor,
+            interpretation_executor=self.interpretation_executor,
         )
 
         if allow_missing: