def simple_call():
dataset_uproot = "user.kchoi:user.kchoi.ttHML_80fb_ttbar"
uproot_transformer_image = "sslhep/servicex_func_adl_uproot_transformer:issue6"
sx_dataset = ServiceXDataset(dataset_uproot,
image=uproot_transformer_image)
ds = ServiceXDatasetSource(sx_dataset, "nominal")
data = ds.Select("lambda e: {'lep_pt_1': e.lep_Pt_1, 'lep_pt_2': e.lep_Pt_2}") \
.AsParquetFiles('junk.parquet') \
.value()
print(data)
def test_bad_wrong_call_name():
'A call needs to be vs a Name node, not something else?'
ds = ServiceXDatasetSource("file://junk.root")
next = ast.Call(func=ast.Name(id='ResultBogus'), args=[ds.query_ast])
with pytest.raises(FuncADLServerException) as e:
ObjectStream(next) \
.value()
assert "ResultBogus" in str(e.value)
def test_bad_wrong_call_type():
'A call needs to be vs a Name node, not something else?'
ds = ServiceXDatasetSource("file://junk.root")
next = ast.Call(func=ast.Attribute(value=ds.query_ast, attr='dude'))
with pytest.raises(FuncADLServerException) as e:
ObjectStream(next) \
.value()
assert "fetch a call from" in str(e.value)
def test_bad_call():
'Normally expect ast.Call - what if not?'
ds = ServiceXDatasetSource("file://junk.root")
next = ast.BinOp(left=ds.query_ast, op=ast.Add(), right=ast.Num(n=10))
with pytest.raises(FuncADLServerException) as e:
ObjectStream(next) \
.value()
assert "Unable" in str(e.value)
async def test_xaod_scoped_dataset_name(simple_Servicex_fe_watcher):
'Simple pandas based query'
f_ds = ServiceXDatasetSource(r'user.fork:bogus_ds')
_ = await f_ds \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \
.Select('lambda j: j.pt()/1000.0') \
.AsPandasDF('JetPt') \
.value_async()
simple_Servicex_fe_watcher[1].assert_called_once()
assert simple_Servicex_fe_watcher[1].call_args[0][0] \
== 'user.fork:bogus_ds'
async def test_xaod_awkward_query(simple_Servicex_fe_watcher):
'Simple pandas based query'
f_ds = ServiceXDatasetSource(r'bogus_ds')
await f_ds \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \
.Select('lambda j: j.pt()/1000.0') \
.AsAwkwardArray('JetPt') \
.value_async()
caller = simple_Servicex_fe_watcher[0].get_data_awkward_async
caller.assert_called_once()
args = caller.call_args[0]
assert len(args) == 1
assert args[0].find('SelectMany') >= 0
assert args[0].startswith('(call ResultTTree')
async def test_xaod_pandas_query(simple_Servicex_fe_watcher):
'Simple pandas based query'
f_ds = ServiceXDatasetSource(r'bogus_ds')
r = await f_ds \
.SelectMany('lambda e: e.Jets("AntiKt4EMTopoJets")') \
.Select('lambda j: j.pt()/1000.0') \
.AsPandasDF('JetPt') \
.value_async()
assert r is not None
assert isinstance(r, pd.DataFrame)
assert len(r) == 0
caller = simple_Servicex_fe_watcher[0].get_data_pandas_df_async
caller.assert_called_once()
args = caller.call_args[0]
assert len(args) == 1
assert args[0].find('SelectMany') >= 0
assert args[0].startswith('(call ResultTTree')
async def test_uproot_parquet_query(simple_Servicex_fe_watcher):
'Simple parquet based query going to xAOD'
f_ds = ServiceXDatasetSource(r'bogus_ds', 'tree-to-scan')
r = await f_ds \
.SelectMany("lambda e: e.jet_pt") \
.Select("lambda j: j / 1000.0") \
.AsParquetFiles('junk.parquet', ['JetPt']) \
.value_async()
assert r is not None
assert isinstance(r, list)
assert len(r) == 1
caller = simple_Servicex_fe_watcher[0].get_data_parquet_async
caller.assert_called_once()
args = caller.call_args[0]
assert len(args) == 1
assert args[0].find('SelectMany') >= 0
assert 'ResultParquet' not in args[0]
def test_as_qastle_uproot():
a = ServiceXDatasetSource("junk.root", 'MainTree')
from qastle import python_ast_to_text_ast
q = python_ast_to_text_ast(a.query_ast)
assert q == "(call EventDataset 'ServiceXDatasetSource' 'MainTree')"
def test_find_EventDataSet_good():
a = ServiceXDatasetSource("file://junk.root") \
.value(executor=do_exe)
assert isinstance(a, ast.Call)
def test_sx_dataset(mocker):
dummy_ds = mocker.MagicMock(spec=ServiceXDataset)
a = ServiceXDatasetSource(dummy_ds) \
.value(executor=do_exe)
assert isinstance(a, ast.Call)