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)