feat: add CorrT1METJet to NanoAOD schema and also add zero TrigObj mass

src/coffea/nanoevents/schemas/nanoaod.py

@@ -64,6 +64,7 @@ class NanoAODSchema(BaseSchema):
         "IsoTrack": "PtEtaPhiMCollection",
         "SoftActivityJet": "PtEtaPhiMCollection",
         "TrigObj": "PtEtaPhiMCollection",
+        "CorrT1METJet": "PtEtaPhiMCollection",
         # True lorentz: pt, eta, phi, mass
         "FatJet": "FatJet",
         "GenDressedLepton": "PtEtaPhiMCollection",
@@ -314,6 +315,22 @@ def _build_collections(self, field_names, input_contents):
         if "Photon_energy" in branch_forms:
             branch_forms["Photon_regrEnergy"] = branch_forms.pop("Photon_energy")
 
+        # Alias CorrT1METJet_rawPt to CorrT1METJet_pt and add zero mass for CorrT1METJet collection
+        if "oCorrT1METJet" in branch_forms:
+            if "CorrT1METJet_pt" not in branch_forms:
+                branch_forms["CorrT1METJet_pt"] = branch_forms["CorrT1METJet_rawPt"]
+            if "CorrT1METJet_mass" not in branch_forms:
+                branch_forms["CorrT1METJet_mass"] = transforms.zeros_from_offsets_form(
+                    branch_forms["oCorrT1METJet"]
+                )
+
+        # Add zero mass to the trigger objects
+        if "oTrigObj" in branch_forms:
+            if "TrigObj_mass" not in branch_forms:
+                branch_forms["TrigObj_mass"] = transforms.zeros_from_offsets_form(
+                    branch_forms["oTrigObj"]
+                )
+
         output = {}
         for name in collections:
             mixin = self.mixins.get(name, "NanoCollection")

src/coffea/nanoevents/trace.py

@@ -36,30 +36,39 @@ def _make_length_zero_one_tracer(
 ) -> tuple[ak.Array, list]:
     form = ak.forms.from_dict(events.attrs["@form"])
     buffer_key = events.attrs["@buffer_key"]
-    buffer_keys = form.expected_from_buffers(buffer_key=buffer_key).keys()
+    expected_buffer_keys = form.expected_from_buffers(buffer_key=buffer_key).keys()
 
     if length == 0:
-        buffers = ak.to_buffers(
-            form.length_zero_array(),
-            byteorder=ak._util.native_byteorder,
-        )[2].values()
+        tmp_array = form.length_zero_array()
     elif length == 1:
-        buffers = ak.to_buffers(
-            form.length_one_array(),
-            byteorder=ak._util.native_byteorder,
-        )[2].values()
+        tmp_array = form.length_one_array()
     else:
         raise ValueError("length must be 0 or 1")
 
+    def getkey(layout, form, attribute):
+        return buffer_key(
+            form_key=form.form_key,
+            attribute=attribute,
+            form=form,
+        )
+
+    container = {}
+    tmp_array._to_buffers(
+        form, getkey, container, events._layout.backend, ak._util.native_byteorder
+    )
+
     report = []
 
     def generate(buffer, report, buffer_key):
         report.append(buffer_key)
         return buffer
 
-    container = {}
-    for key, buffer in zip(buffer_keys, buffers):
+    for key, buffer in container.items():
         container[key] = partial(generate, buffer=buffer, report=report, buffer_key=key)
+
+    assert list(container.keys()) == list(
+        expected_buffer_keys
+    ), "length zero/one array buffer keys do not match the expected ones"
     array = ak.from_buffers(
         form=form,
         length=length,