fix: Make LumiMask compatible with dask-awkward

src/coffea/lumi_tools/lumi_tools.py

@@ -1,6 +1,7 @@
 import json
 
 import awkward as ak
+import dask_awkward as dak
 from numba import types
 from numba.typed import Dict
 
@@ -107,9 +108,9 @@ def __call__(self, runs, lumis):
 
         Parameters
         ----------
-            runs : numpy.ndarray
+            runs : numpy.ndarray or awkward.highlevel.Array or dask_awkward.Array
                 Vectorized list of run numbers
-            lumis : numpy.ndarray
+            lumis : numpy.ndarray or awkward.highlevel.Array or dask_awkward.Array
                 Vectorized list of lumiSection numbers
 
         Returns
@@ -118,18 +119,30 @@ def __call__(self, runs, lumis):
                 An array of dtype `bool` where valid (run, lumi) tuples
                 will have their corresponding entry set ``True``.
         """
-        # fill numba typed dict
-        _masks = Dict.empty(key_type=types.uint32, value_type=types.uint32[:])
-        for k, v in self._masks.items():
-            _masks[k] = v
-
-        if isinstance(runs, ak.highlevel.Array):
-            runs = ak.to_numpy(runs)
-        if isinstance(lumis, ak.highlevel.Array):
-            lumis = ak.to_numpy(lumis)
-        mask_out = np.zeros(dtype="bool", shape=runs.shape)
-        LumiMask._apply_run_lumi_mask_kernel(_masks, runs, lumis, mask_out)
-        return mask_out
+
+        def apply(runs, lumis):
+            # fill numba typed dict
+            _masks = Dict.empty(key_type=types.uint32, value_type=types.uint32[:])
+            for k, v in self._masks.items():
+                _masks[k] = v
+
+            runs_orig = runs
+            if isinstance(runs, ak.highlevel.Array):
+                runs = ak.to_numpy(ak.typetracer.length_zero_if_typetracer(runs))
+            if isinstance(lumis, ak.highlevel.Array):
+                lumis = ak.to_numpy(ak.typetracer.length_zero_if_typetracer(lumis))
+            mask_out = np.zeros(dtype="bool", shape=runs.shape)
+            LumiMask._apply_run_lumi_mask_kernel(_masks, runs, lumis, mask_out)
+            if isinstance(runs_orig, ak.Array):
+                mask_out = ak.Array(mask_out)
+            if ak.backend(runs_orig) == "typetracer":
+                mask_out = ak.Array(mask_out.layout.to_typetracer(forget_length=True))
+            return mask_out
+
+        if isinstance(runs, dak.Array):
+            return dak.map_partitions(apply, runs, lumis)
+        else:
+            return apply(runs, lumis)
 
     # This could be run in parallel, but windows does not support it
     @staticmethod

tests/test_lumi_tools.py

@@ -1,4 +1,7 @@
+import awkward as ak
 import cloudpickle
+import dask_awkward as dak
+from dask.distributed import Client
 
 from coffea.lumi_tools import LumiData, LumiList, LumiMask
 from coffea.util import numpy as np
@@ -56,6 +59,8 @@ def test_lumidata():
 
 
 def test_lumimask():
+    client = Client()
+
     lumimask = LumiMask(
         "tests/samples/Cert_294927-306462_13TeV_EOY2017ReReco_Collisions17_JSON.txt"
     )
@@ -86,6 +91,15 @@ def test_lumimask():
 
     assert np.all(lumimask(runs, lumis) == lumimask_pickle(runs, lumis))
 
+    runs_dak = dak.from_awkward(ak.Array(runs), 1)
+    lumis_dak = dak.from_awkward(ak.Array(lumis), 1)
+    assert np.all(
+        client.compute(lumimask(runs_dak, lumis_dak)).result()
+        == lumimask_pickle(runs, lumis)
+    )
+
+    client.close()
+
 
 def test_lumilist():
     lumidata = LumiData("tests/samples/lumi_small.csv")