Fixed the renaming problem for the ttree, now using NanoAODSchema.

python/tree_reader.py

@@ -97,10 +97,6 @@ def getDataFrame(self, prefix, entry_block, fallback=None):
                     if br.startswith(f'{prefix}_') and
                     br != f'{prefix}_n']
         names = ['_'.join(br.split('_')[1:]) for br in branches]
-
-        # FIXME: name_map = dict(zip(names, branches, strict=False))
-        # strict=False does not work with the dask.
-
         name_map = dict(zip(names, branches))
         if len(branches) == 0:
             if fallback is not None:
@@ -109,35 +105,38 @@ def getDataFrame(self, prefix, entry_block, fallback=None):
             print(f'stored branch prefixes are: {prefs}')
             raise ValueError(f'[TreeReader::getDataFrame] No branches with prefix: {prefix}')
 
-        # FIXME: old code of eagerily reading ttree.
-        #akarray = self.tree.arrays(names,
-        #library='ak',
-        #aliases=name_map,
-        #entry_start=self.file_entry,
+        dask_akarray = NanoEventsFactory.from_root(
+                                   self.tree,
+                                   schemaclass=NanoAODSchema).events()
+
+        #akarray = self.tree.arrays(names, 
+        #library='ak', 
+        #aliases=name_map, 
+        #entry_start=self.file_entry, 
         #entry_stop=self.file_entry+entry_block)
+        #print("[0] prefix to select: ", prefix)
 
-        akarray = NanoEventsFactory.from_root(
-                                   self.tree,
-                                   schemaclass=BaseSchema).events()
+        dask_akarray = dask_akarray[prefix]     
+        #print("[1] Selected fields from prefix", dask_akarray.fields)   
+
+        dask_akarray = dask_akarray[names]
 
-        akarray = akarray[branches][self.file_entry : self.file_entry + entry_block]
+        #print("[2] specific fields with names", dask_akarray.fields)
+        dask_akarray = dask_akarray[self.file_entry : self.file_entry + entry_block]
 
-        # print(akarray)
         records = {}
-        for field in akarray.fields:
-            records[field] = akarray[field]
+        for field in dask_akarray.fields:
+            records[field] = dask_akarray[field]
 
-        # FIXME: Commented this part out, does not work with dask.
-        #if 'pt' in names and 'eta' in names and 'phi' in names:
-        #if 'mass' not in names and 'energy' not in names:
-        #records['mass'] = 0.*akarray['pt']
-        #return vector.zip(records)
+        if 'pt' in names and 'eta' in names and 'phi' in names:
+            if 'mass' not in names and 'energy' not in names:
+                records['mass'] = 0.*dask_akarray['pt']
+            return vector.zip(records)
 
         return ak.zip(records)
 
         # FIXME: we should probably do an ak.Record using sometjhing along the lines of:
         # ele_rec = ak.zip({'pt': tkele.pt, 'eta': tkele.eta, 'phi': tkele.phi}, with_name="pippo")
         # this would allow to handle the records and assign behaviours....
 
-        # return akarray
-
+        # return akarray