Closes Bears-R-Us#3339: Add multi-batch parquet read tests (Bears-R-U…

…s#3350)

* Closes Bears-R-Us#3339: Add multi-batch parquet read tests

This PR (closes Bears-R-Us#3339) adds testing for parquet reads or arrays and strings large enough to trigger more than one batch. We also add testing of a segarray of segstrings containing empty segs and empty strings

* add proto

---------

Co-authored-by: Tess Hayes <stress-tess@users.noreply.github.com>

PROTO_tests/tests/io_test.py

@@ -527,11 +527,13 @@ def test_ipv4_columns(self, comp):
         df["a"] = df["a"].export_uint()
         assert ak.arange(10).to_list() == df["a"].to_list()
 
-    def test_empty_segs_segarray(self):
+    def test_multi_batch_reads(self):
         # verify reproducer for #3074 is resolved
+        # seagarray w/ empty segs multi-batch pq reads
 
         # bug seemed to consistently appear for val_sizes
-        # exceeding 700000, round up to ensure we'd hit it
+        # exceeding 700000 (likely due to this requiring more than one batch)
+        # we round up to ensure we'd hit it
         val_size = 1000000
 
         df_dict = dict()
@@ -544,7 +546,8 @@ def test_empty_segs_segarray(self):
             rng.integers(0, 2**32, size=val_size, dtype="uint"),
             rng.integers(0, 1, size=val_size, dtype="bool"),
             rng.integers(-(2**32), 2**32, size=val_size, dtype="int"),
-            some_nans,
+            some_nans,  # contains nans
+            ak.random_strings_uniform(0, 4, val_size, seed=seed),  # contains empty strings
         ]
 
         for vals in vals_list:
@@ -571,6 +574,15 @@ def test_empty_segs_segarray(self):
                 print("\nseed: ", seed)
                 assert_series_equal(pddf["rand"], to_pd, check_names=False, rtol=1e-05, atol=1e-08)
 
+                # test writing multi-batch non-segarrays
+                file_path = f"{tmp_dirname}/multi_batch_vals"
+                vals.to_parquet(file_path, dataset="my_vals")
+                read = ak.read_parquet(file_path + "*")["my_vals"]
+                if isinstance(vals, ak.pdarray) and vals.dtype == ak.float64:
+                    assert np.allclose(read.to_list(), vals.to_list(), equal_nan=True)
+                else:
+                    assert (read == vals).all()
+
 
 class TestHDF5:
     hdf_test_base_tmp = f"{os.getcwd()}/hdf_io_test"

tests/parquet_test.py

@@ -584,11 +584,13 @@ def test_decimal_reads(self):
             for i in range(1, 39):
                 self.assertTrue(np.allclose(ak_data["decCol" + str(i)].to_ndarray(), data[i - 1]))
 
-    def test_empty_segs_segarray(self):
+    def test_multi_batch_reads(self):
         # verify reproducer for #3074 is resolved
+        # seagarray w/ empty segs multi-batch pq reads
 
         # bug seemed to consistently appear for val_sizes
-        # exceeding 700000, round up to ensure we'd hit it
+        # exceeding 700000 (likely due to this requiring more than one batch)
+        # we round up to ensure we'd hit it
         val_size = 1000000
 
         df_dict = dict()
@@ -601,7 +603,8 @@ def test_empty_segs_segarray(self):
             rng.integers(0, 2**32, size=val_size, dtype="uint"),
             rng.integers(0, 1, size=val_size, dtype="bool"),
             rng.integers(-(2**32), 2**32, size=val_size, dtype="int"),
-            some_nans,
+            some_nans,  # contains nans
+            ak.random_strings_uniform(0, 4, val_size, seed=seed),  # contains empty strings
         ]
 
         for vals in vals_list:
@@ -628,6 +631,15 @@ def test_empty_segs_segarray(self):
                 print("\nseed: ", seed)
                 assert_series_equal(pddf["rand"], to_pd, check_names=False, rtol=1e-05, atol=1e-08)
 
+                # test writing multi-batch non-segarrays
+                file_path = f"{tmp_dirname}/multi_batch_vals"
+                vals.to_parquet(file_path, dataset="my_vals")
+                read = ak.read_parquet(file_path + "*")["my_vals"]
+                if isinstance(vals, ak.pdarray) and vals.dtype == ak.float64:
+                    assert np.allclose(read.to_list(), vals.to_list(), equal_nan=True)
+                else:
+                    assert (read == vals).all()
+
     @pytest.mark.optional_parquet
     def test_against_standard_files(self):
         datadir = "resources/parquet-testing"