limit parquet file dimension

bfonta · Sep 27, 2023 · 44e1f6d · 44e1f6d
1 parent 2d169dd
commit 44e1f6d
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 5 deletions.
diff --git a/bye_splits/data_handle/data_process.py b/bye_splits/data_handle/data_process.py
@@ -142,8 +142,7 @@ def EventDataParticle(particles, PU, tag, reprocess, debug=False, logger=None):
         indata.adir = cfg["io"][PU][particles]["dir"]
         indata.tree = cfg["io"][PU][particles]["tree"]
 
-    tag = particles + "_" + tag
+    tag = particles + "_" + PU + "_" + tag 
     tag += "_debug" * debug
 
-
     return EventData(indata, tag, defevents, reprocess, logger)
diff --git a/bye_splits/data_handle/event.py b/bye_splits/data_handle/event.py
@@ -158,7 +158,17 @@ def select(self):
         with up.open(self.indata.path, array_cache='550 MB', num_workers=8) as f:
             tree = f[self.indata.tree_path]
             allvars = set([y for x in self.var.values() for y in x.values()])
-            data = tree.arrays(filter_name='/' + '|'.join(allvars) + '/', entry_stop=200, library='ak')
+
+            threshold_size = 0.1 
+            threshold_size_bytes = threshold_size * 8e+9
+            data = ak.Array([])
+            for array in tree.iterate(filter_name='/' + '|'.join(allvars) + '/', step_size='20 MB', library='ak'):
+                if (data.layout.nbytes + array.layout.nbytes) <= threshold_size_bytes:
+                    data = ak.concatenate([data, array], axis=0)
+                else:
+                    break
+
+            #data = tree.arrays(filter_name='/' + '|'.join(allvars) + '/', entry_stop=5000, library='ak')
         # data[self.var.v] = data.waferv
         # data[self.newvar.vs] = -1 * data.waferv
         # data[self.newvar.c] = "#8a2be2"

diff --git a/bye_splits/plot/display_plotly/processing.py b/bye_splits/plot/display_plotly/processing.py
@@ -59,7 +59,7 @@ def get_data(self, pars, particles_str, event = ''):
             self.list_events = [key.split('ev_')[1] for key in file.keys()]
 
         event = event or self.random_event(self.filename)
-        if event not in self.list_events:
+        if str(event) not in self.list_events:
             dict_ev, gen_info = run_radii_chain(pars, particles, PU, self.coefs, event)
             self.process_event(dict_ev, gen_info)
 

diff --git a/bye_splits/scripts/run_radii_chain.py b/bye_splits/scripts/run_radii_chain.py
@@ -24,7 +24,7 @@
 import pandas as pd
 
 def run_radii_chain(pars, particles, PU, coefs, event=None):
-    df_gen, df_cl, df_tc = get_data_reco_chain_start(nevents=30, reprocess=True, particles=particles, PU=PU, event=event)
+    df_gen, df_cl, df_tc = get_data_reco_chain_start(nevents=30, reprocess=False, particles=particles, PU=PU, event=event)
 
     fill_d = params.read_task_params("fill")
     tasks.fill.fill(pars, df_gen, df_cl, df_tc, **fill_d)