update

Author: JanvanderLinden

scripts/das.py

@@ -6,6 +6,8 @@
 parser = optparse.OptionParser()
 parser.add_option("--get-parent", dest="get_parent", action="store_true", default=False)
 parser.add_option("--get-files", dest="get_files", action="store_true", default=False)
+parser.add_option("--get-meta", dest="get_meta", action="store_true", default=False)
+parser.add_option("--get-cmsdb", dest="get_cmsdb", action="store_true", default=False)
 parser.add_option("--get-xs", dest="get_xs", action="store_true", default=False)
 (opts, args) = parser.parse_args()
 
@@ -46,6 +48,51 @@ def get_xs( dataset, nf=1 ):
     xs_unit = xs_line.split(" ")[-1]
 
     return xs_val, xs_unit
+
+def get_meta( dataset ):
+    # first get datasset info and id
+    cmd = f"dasgoclient -query='dataset={dataset}' -json"
+    infos = call(cmd)
+    
+    dataset_id = infos[0]["dataset"][0].get("dataset_id", -1)
+    dataset_name = infos[0]["dataset"][0]["name"]
+    dataset_era = infos[0]["dataset"][0].get("acquisition_era_name", " ")[-1]
+    is_data = infos[0]["dataset"][0].get("datatype", "") == "data"
+    era_aux = ""
+    procs = ""
+    if is_data: 
+        era_aux = {"era": dataset_era}
+        procs = "procs.data"
+
+
+    # then get number of files
+    cmd = f"dasgoclient -query='summary dataset={dataset}' -json"
+    infos = call(cmd)
+
+    nfiles = infos[0]["summary"][0]["nfiles"]
+    nevents = infos[0]["summary"][0]["nevents"]
+
+    return {"id": dataset_id, "name": dataset_name, "nfiles": nfiles, "nevents": nevents, "is_data": is_data, "aux": era_aux, "procs": procs}
+    
+    
+def format_cmsdb( dataset ):
+    template = """
+cpn.add_dataset(
+    name="",
+    id={id},
+    is_data={is_data},
+    processes=[{procs}],
+    keys=[
+        "{name}",
+    ],
+    n_files={nfiles},
+    n_events={nevents},
+    aux={aux}
+)
+"""
+    return template.format( **get_meta(dataset) )
+
+    
 
 # loop over all arguments
 out_dict = {}
@@ -76,6 +123,10 @@ def get_xs( dataset, nf=1 ):
             ret = get_parent(dataset)
         elif opts.get_xs:
             ret = get_xs(dataset)
+        elif opts.get_meta:
+            ret = get_meta(dataset)
+        elif opts.get_cmsdb:
+            ret = format_cmsdb(dataset)
         else:
             # call dasgoclient command
             cmd = f"dasgoclient -query='dataset={dataset}' -json"

scripts/pq.py

@@ -1,10 +1,9 @@
-print("using pq.py")
 import awkward as ak
 import pickle
 import numpy as np
 import sys, os
+import pyarrow.parquet as pq
 
-print("call 'help()' to show features")
 def help():
     print("usage:")
     print("\tpq [parquet/pickle files]")
@@ -16,14 +15,6 @@ def load_file(f):
     else:
         return ak.from_parquet(f)
 
-infiles = sys.argv[1:]
-files = [
-    load_file(f)
-    for f in infiles
-]
-data = files[0]
-
-
 def to_np(arr):
     if arr.ndim > 1:
         return ak.flatten(arr).to_numpy()
@@ -47,3 +38,55 @@ def percentiles(arr, percentiles=[0, 5, 10, 25, 50, 75, 90, 95, 100], ret=False)
         print("percentiles:")
         for a, b in zip(percentiles, x):
             print(f"{a}%: {b}")
+
+def validate(f):
+    try:
+        if f.endswith(".parquet"):
+            data = pq.ParquetFile(f)
+        else:
+            with open(f, "rb") as pf:
+                pickle.load(pf)
+            
+    except:
+        print(f"Broken file {f}")
+        return False
+    
+    return True
+
+if __name__ == "__main__":
+    import optparse
+    import glob
+    parser = optparse.OptionParser()
+    parser.add_option("-v", "--validate", dest="validate")
+    parser.add_option("--del", dest="delete", action="store_true", default=False)
+    (opts, args) = parser.parse_args()
+    
+    infiles = []
+    for f in args:
+        if "*" in f:
+            infiles += glob.glob(args)
+        else:
+            infiles.append(f)
+
+    if opts.validate:
+        print(f"validating {len(infiles)} pq files...")
+        i = 0
+        for f in infiles:
+            if not validate(f):
+                i += 1
+                if opts.delete:
+                    print(f"\t--> Removing file")
+                    os.remove(f)
+        print(f"\n--> {i}/{len(infiles)} broken")
+                
+        
+    else:
+        print("using pq.py")
+        print("call 'help()' to show features")
+        files = [
+            load_file(f)
+            for f in infiles
+        ]
+        data = files[0]
+
+

setup.sh

@@ -11,6 +11,7 @@ alias condorSubmit='python3 $TOOLBOX/scripts/condorSubmit.py'
 alias summary='python3 -m correctionlib.cli summary'
 
 alias pq='python3 -i $TOOLBOX/scripts/pq.py'
+alias pq_validate='python3 $TOOLBOX/scripts/pq.py -v'
 alias das='python3 $TOOLBOX/scripts/das.py'
 
 PYTHONPATH=$TOOLBOX:$PYTHONPATH