Merge pull request #986 from bmcdonald3/parquet-tests

Fix parquet tests for multilocale runs
Bears-R-Us · Dec 2, 2021 · 639faef · 639faef
2 parents 8b41dc5 + 7925ffd
commit 639faef
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 deletions.
diff --git a/src/Parquet.chpl b/src/Parquet.chpl
@@ -84,12 +84,12 @@ module Parquet {
     return ArrowTypes.notimplemented;
   }
 
-  proc writeDistArrayToParquet(A, filename, dsetname, rowGroupSize) {
+  proc writeDistArrayToParquet(A, filename, dsetname, rowGroupSize) throws {
     extern proc c_writeColumnToParquet(filename, chpl_arr, colnum,
                                      dsetname, numelems, rowGroupSize);
     var filenames: [0..#A.targetLocales().size] string;
     for i in 0..#A.targetLocales().size {
-      var suffix = i: string;
+      var suffix = '%04i'.format(i): string;
       filenames[i] = filename + "_LOCALE" + suffix + ".parquet";
     }
 
@@ -102,7 +102,7 @@ module Parquet {
       }
   }
 
-  proc write1DDistArrayParquet(filename: string, dsetname, A) {
+  proc write1DDistArrayParquet(filename: string, dsetname, A) throws {
     writeDistArrayToParquet(A, filename, dsetname, ROWGROUPS);
     return false;
   }

diff --git a/tests/parquet_test.py b/tests/parquet_test.py
@@ -1,9 +1,10 @@
 import glob, os
 from context import arkouda as ak
 from base_test import ArkoudaTest
+import numpy as np
 import pytest
 
-SIZE = 50
+SIZE = 100
 NUMFILES = 5
 verbose = True
 
@@ -13,23 +14,23 @@ def test_parquet(self):
         ak_arr = ak.randint(0, 2**32, SIZE)
         ak_arr.save_parquet("pq_testcorrect", "my-dset")
         pq_arr = ak.read_parquet("pq_testcorrect*", "my-dset")
-        # get the dset from the dictionary in multi-locale cases
+        self.assertTrue((ak_arr ==  pq_arr).all())
+
         for f in glob.glob('pq_test*'):
             os.remove(f)
-            self.assertTrue((ak_arr ==  pq_arr).all())
 
     def test_multi_file(self):
         adjusted_size = int(SIZE/NUMFILES)*NUMFILES
         test_arrs = []
+        elems = ak.randint(0, 2**32, adjusted_size)
+        per_arr = int(adjusted_size/NUMFILES)
         for i in range(NUMFILES):
-            test_arrs.append(ak.randint(0, 2**32, int(adjusted_size/NUMFILES)))
+            test_arrs.append(elems[(i*per_arr):(i*per_arr)+per_arr])
             test_arrs[i].save_parquet("pq_test" + str(i), "test-dset")
 
         pq_arr = ak.read_parquet("pq_test*", "test-dset")
-        self.assertTrue(len(pq_arr) == adjusted_size)
 
-        for i in range(NUMFILES):
-            sz = len(test_arrs[i])
-            self.assertTrue((test_arrs[i] == pq_arr[(i*sz):(i*sz)+sz]).all())
+        self.assertTrue((elems == pq_arr).all())
+
         for f in glob.glob('pq_test*'):
             os.remove(f)