feat: support explicit bulk preloading (uproot-only)

src/coffea/nanoevents/factory.py

@@ -267,6 +267,7 @@ def from_root(
         decompression_executor=None,
         interpretation_executor=None,
         delayed=uproot._util.unset,
+        preload=None,
     ):
         """Quickly build NanoEvents from a root file
 
@@ -310,6 +311,9 @@ def from_root(
                 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
+            preload (None or Callable):
+                A function to call to preload specific branches/columns in bulk. Only works in eager and virtual mode.
+                Passed to ``tree.arrays`` as the ``filter_branch`` argument to filter branches to be preloaded.
 
         Returns
         -------
@@ -415,6 +419,24 @@ def from_root(
             f"{entry_start}-{entry_stop}",
         )
         uuidpfn = {partition_key[0]: tree.file.file_path}
+
+        preloaded_arrays = None
+        if preload is not None:
+            preloaded_arrays = tree.arrays(
+                filter_branch=preload,
+                entry_start=entry_start,
+                entry_stop=entry_stop,
+                ak_add_doc=True,
+                decompression_executor=decompression_executor,
+                interpretation_executor=interpretation_executor,
+                how=dict,
+            )
+            # this ensures that the preloaded arrays are only sliced as they are supposed to be
+            preloaded_arrays = {
+                k: _OnlySliceableAs(v, slice(entry_start, entry_stop))
+                for k, v in preloaded_arrays.items()
+            }
+
         mapping = UprootSourceMapping(
             TrivialUprootOpener(uuidpfn, uproot_options),
             entry_start,
@@ -423,6 +445,7 @@ def from_root(
             access_log=access_log,
             use_ak_forth=use_ak_forth,
             virtual=mode == "virtual",
+            preloaded_arrays=preloaded_arrays,
         )
         mapping.preload_column_source(partition_key[0], partition_key[1], tree)
 

src/coffea/nanoevents/mapping/uproot.py

@@ -122,6 +122,7 @@ def __init__(
         virtual=False,
         decompression_executor=None,
         interpretation_executor=None,
+        preloaded_arrays=None,
     ):
         super().__init__(
             fileopener=fileopener,
@@ -138,6 +139,7 @@ def __init__(
         self.interpretation_executor = (
             interpretation_executor or uproot.source.futures.TrivialExecutor()
         )
+        self.preloaded_arrays = preloaded_arrays
 
     @classmethod
     def _extract_base_form(cls, tree, iteritems_options={}):
@@ -211,16 +213,22 @@ def extract_column(
                 "Received columnhandle of None when missing column in file is not allowed!"
             )
 
-        interp = columnhandle.interpretation
-        interp._forth = use_ak_forth
-
-        the_array = columnhandle.array(
-            interp,
-            entry_start=start,
-            entry_stop=stop,
-            decompression_executor=self.decompression_executor,
-            interpretation_executor=self.interpretation_executor,
-        )
+        if (
+            self.preloaded_arrays is not None
+            and columnhandle.name in self.preloaded_arrays
+        ):
+            the_array = self.preloaded_arrays[columnhandle.name][start:stop]
+        else:
+            interp = columnhandle.interpretation
+            interp._forth = use_ak_forth
+
+            the_array = columnhandle.array(
+                interp,
+                entry_start=start,
+                entry_stop=stop,
+                decompression_executor=self.decompression_executor,
+                interpretation_executor=self.interpretation_executor,
+            )
         if isinstance(the_array.layout, awkward.contents.ListOffsetArray):
             the_array = awkward.Array(the_array.layout.to_ListOffsetArray64(True))