iterator flag feature to determine return type

Author: ArturU043

servicex_analysis_utils/materialization.py

@@ -30,14 +30,15 @@
 import dask_awkward as dak 
 import logging 
 
-def to_awk(deliver_dict, dask=False, **kwargs):
+def to_awk(deliver_dict, dask=False, iterator=False, **kwargs):
     """
     Load an awkward array from the deliver() output with uproot or uproot.dask.
 
     Parameters:
         deliver_dict (dict): Returned dictionary from servicex.deliver()
                             (keys are sample names, values are file paths or URLs).
         dask (bool):        Optional. Flag to load as dask-awkward array. Default is False
+        iterator(bool):      Optional. Flag to materialize the data into arrays or to return iterables with uproot.iterate
         **kwargs :   Optional. Additional keyword arguments passed to uproot.dask, uproot.iterate and from_parquet
 
     
@@ -69,7 +70,12 @@ def to_awk(deliver_dict, dask=False, **kwargs):
             else:
                 if is_root==True:
                     # Use uproot.iterate to handle URLs and local paths files in chunks
-                    awk_arrays[sample]=uproot.iterate(paths, library="ak", **kwargs) # not an ak array but a generator
+                    iterators=uproot.iterate(paths, library="ak", **kwargs)
+                    if iterator==True:
+                        awk_arrays[sample]= iterators #return iterators
+                    else :
+                        awk_arrays[sample]=ak.concatenate(list(iterators)) #return array
+    
                 else:
                     #file is parquet 
                     awk_arrays[sample] = ak.from_parquet(paths, **kwargs)

tests/test_materialization.py

@@ -34,6 +34,7 @@
 import sys
 import numpy as np
 from servicex_analysis_utils.materialization import to_awk
+import types
 
 
 @pytest.fixture
@@ -69,8 +70,10 @@ def test_to_awk(build_test_samples):
 
     #Collecting all samples 
     assert list(result.keys())==["Test-Sample1", "Test-Sample2"]
-    arr1 = ak.concatenate(list(result["Test-Sample1"]))  # Materialize the generator from uproot.iterate  
-    arr2 = ak.concatenate(list(result["Test-Sample2"]))  
+    arr1 = result["Test-Sample1"]
+    arr2 = result["Test-Sample2"]
+ 
+  
 
     #Collecting all branches
     assert ak.fields(arr1) == ['branch1', 'branch2']
@@ -114,6 +117,25 @@ def test_to_awk_dask(build_test_samples):
     assert ak.all(arr1['branch2'].compute() == ak.from_numpy(np.zeros(100)))
     assert ak.all(arr2['branch1'].compute() == ak.from_numpy(np.ones(10)))
 
+def test_to_awk_delayed_and_kwargs(build_test_samples):
+    sx_dict = build_test_samples
+    result_delay = to_awk(sx_dict, iterator=True, expressions="branch1") #return iterable + selection kwarg
+        
+    #Checking iterator return type
+    assert isinstance(result_delay["Test-Sample1"], types.GeneratorType)  
+    assert isinstance(result_delay["Test-Sample2"], types.GeneratorType)  
+
+
+    arr1 = ak.concatenate(list(result_delay["Test-Sample1"]))  # Materialize the generator from uproot.iterate  
+    arr2 = ak.concatenate(list(result_delay["Test-Sample2"]))
+
+    #Checking materialization
+    assert isinstance(arr1, ak.Array), "to_awk(dask=True) does not produce an ak.Array instance"
+    assert isinstance(arr2, ak.Array), "to_awk(dask=True) does not produce an ak.Array instance"
+
+    #Checking only 1 branch selected
+    assert ak.fields(arr1) == ['branch1']
+    assert ak.fields(arr2) == ['branch1']
 
 def test_unsupported_file_format():
     fake_paths = {"fake-Sample": ["invalid_file.txt"]}