kwargs for parquet filter

Author: acordeir

servicex_analysis_utils/materialization.py

@@ -30,15 +30,15 @@
 import dask_awkward as dak 
 import logging 
 
-def to_awk(deliver_dict, dask=False, **uproot_kwargs):
+def to_awk(deliver_dict, dask=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
-        **uproot_kwargs :   Optional. Additional keyword arguments passed to uproot.dask or uproot.iterate
+        **kwargs :   Optional. Additional keyword arguments passed to uproot.dask, uproot.iterate and from_parquet
 
     
     Returns:
@@ -62,19 +62,19 @@ def to_awk(deliver_dict, dask=False, **uproot_kwargs):
             if dask:
                 if is_root==True:
                     # Use uproot.dask to handle URLs and local paths lazily 
-                    awk_arrays[sample] = uproot.dask(paths, library="ak", **uproot_kwargs)
+                    awk_arrays[sample] = uproot.dask(paths, library="ak", **kwargs)
                 else:
                     #file is parquet 
-                    awk_arrays[sample] = dak.from_parquet(paths)
+                    awk_arrays[sample] = dak.from_parquet(paths, **kwargs)
             else:
                 if is_root==True:
                     # Use uproot.iterate to handle URLs and local paths files in chunks
-                    tmp_arrays = list(uproot.iterate(paths, library="ak", **uproot_kwargs))
+                    tmp_arrays = list(uproot.iterate(paths, library="ak", **kwargs))
                     # Merge arrays
                     awk_arrays[sample] = ak.concatenate(tmp_arrays) 
                 else:
                     #file is parquet 
-                    awk_arrays[sample] = ak.from_parquet(paths)
+                    awk_arrays[sample] = ak.from_parquet(paths, **kwargs)
 
 
         except Exception as e:

tests/test_parquet.py renamed to tests/test_materialization_parquet.py

@@ -59,7 +59,7 @@ def test_to_awk_collection(build_test_samples):
     assert ak.all(arr1['branch2'] == ak.from_numpy(np.zeros(100)))
     assert ak.all(arr2['branch1'] == ak.from_numpy(np.ones(10)))
 
-    # Checking kwargs (Expressions equivalent in Parquet)
-    result_filtered = to_awk(sx_dict, filter_name="branch1")  # Now using filter_name instead of expressions
+    # Checking kwargs 
+    result_filtered = to_awk(sx_dict, columns="branch1")  
     arr1_filtered = result_filtered["Test-Sample1"]
     assert ak.fields(arr1_filtered) == ['branch1']  # branch2 should be filtered out