|
| 1 | +from servicex import query as q, deliver, dataset |
| 2 | + |
| 3 | + |
| 4 | +def uproot_pythonfunction_dict(): |
| 5 | + |
| 6 | + def run_query(input_filenames=None): |
| 7 | + import uproot # type: ignore |
| 8 | + with uproot.open({input_filenames: "reco"}) as o: |
| 9 | + br = o.arrays("el_pt_NOSYS") |
| 10 | + return br |
| 11 | + |
| 12 | + query1 = q.PythonFunction().with_uproot_function(run_query) |
| 13 | + |
| 14 | + query2 = q.FuncADL_ATLASr22() # type: ignore |
| 15 | + jets_per_event = query2.Select(lambda e: e.Jets('AnalysisJets')) |
| 16 | + jet_info_per_event = jets_per_event.Select( |
| 17 | + lambda jets: { |
| 18 | + 'pt': jets.Select(lambda j: j.pt()), |
| 19 | + 'eta': jets.Select(lambda j: j.eta()) |
| 20 | + } |
| 21 | + ) |
| 22 | + |
| 23 | + spec = { |
| 24 | + 'Sample': [{ |
| 25 | + 'Name': "Uproot_PythonFunction_Dict", |
| 26 | + 'RucioDID': "user.mtost:user.mtost.singletop.p6026.Jun13", |
| 27 | + 'Query': query1 |
| 28 | + }, |
| 29 | + { |
| 30 | + 'Name': "func_adl_xAOD_simple", |
| 31 | + 'Dataset': dataset.FileList( |
| 32 | + [ |
| 33 | + "root://eospublic.cern.ch//eos/opendata/atlas/rucio/mc20_13TeV/DAOD_PHYSLITE.37622528._000013.pool.root.1", # noqa: E501 |
| 34 | + ] |
| 35 | + ), |
| 36 | + 'Query': jet_info_per_event |
| 37 | + }] |
| 38 | + } |
| 39 | + return deliver(spec, servicex_name="servicex-uc-af") |
| 40 | + |
| 41 | +if __name__ == "__main__": |
| 42 | + files = uproot_pythonfunction_dict() |
| 43 | + assert len(files["Uproot_PythonFunction_Dict"])==27 |
| 44 | + assert len(files["func_adl_xAOD_simple"])==1 |
0 commit comments