dask-parquet tests

Author: acordeir

tests/test_materialization.py

@@ -63,7 +63,7 @@ def build_test_samples(tmp_path):
 
 
 #Test functions
-def test_to_awk_collection(build_test_samples):
+def test_to_awk(build_test_samples):
     sx_dict = build_test_samples
     result = to_awk(sx_dict) #uproot.iterate expressions kwarg
 

tests/test_materialization_parquet.py

@@ -29,7 +29,7 @@ def build_test_samples(tmp_path):
     })
 
     # Write to Parquet files
-    ak.to_parquet(data1, test_path1)
+    ak.to_parquet(data1, test_path1, row_group_size=10000) #partions
     ak.to_parquet(data2, test_path2)
 
     # Dict simulating servicex.deliver() output
@@ -39,7 +39,7 @@ def build_test_samples(tmp_path):
 
 
 # Test function for to_awk with Parquet files
-def test_to_awk_collection(build_test_samples):
+def test_to_awk_parquet(build_test_samples):
     sx_dict = build_test_samples
     result = to_awk(sx_dict)  # Using ak.from_parquet internally
 
@@ -63,3 +63,28 @@ def test_to_awk_collection(build_test_samples):
     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
+
+def test_to_awk_dask_parquet(build_test_samples):
+    sx_dict = build_test_samples
+    result_da = to_awk(sx_dict, dask=True, split_row_groups=True) #split in partitions
+    
+    #Collecting all samples 
+    assert list(result_da.keys())==["Test-Sample1", "Test-Sample2"]
+    arr1=result_da["Test-Sample1"]
+    arr2=result_da["Test-Sample2"]
+
+    #Checking instance
+    assert isinstance(arr1, dak.Array), "to_awk(dask=True) does not produce an dak.Array instance"
+    assert isinstance(arr2, dak.Array), "to_awk(dask=True) does not produce an dak.Array instance"
+
+    #Testing partitionning kwarg
+    assert arr1.npartitions == 10
+    assert arr2.npartitions == 1
+
+    #Collecting all branches
+    assert ak.fields(arr1) == ['branch1', 'branch2']
+    assert ak.fields(arr2) == ['branch1']
+
+    #Collecting all elements per branch
+    assert ak.all(arr1['branch2'].compute() == ak.from_numpy(np.zeros(100)))
+    assert ak.all(arr2['branch1'].compute() == ak.from_numpy(np.ones(10)))