Fix read_file_multiprocessing returning duplicated rows when fewer rows than processes are read

pyreadstat/pyreadstat.py

@@ -1180,6 +1180,9 @@ def read_file_multiprocessing(
         prev_offset = offset
         prev_div = div
         offsets.append((offset, div))
+    # a chunk with 0 rows would be passed as row_limit=0, which means no limit, and that worker would read
+    # all the remaining rows. Drop those chunks, but keep one if all are empty to get an empty dataframe.
+    offsets = [x for x in offsets if x[1] > 0] or offsets[:1]
     jobs = [(read_function, file_path, offset, chunksize, kwargs) for offset, chunksize in offsets]
     pool = mp.Pool(processes=num_processes)
     try:

tests/test_basic.py

@@ -668,6 +668,23 @@ def test_chunk_reader_multiprocess_dict(self):
         self.assertTrue(df_multi.equals(df_single))
 
 
+    def test_multiprocess_reader_fewer_rows_than_processes(self):
+        # some workers get 0 rows to read, they should not read the rest of the file
+        fpath = os.path.join(self.basic_data_folder, "sample_large.sav")
+        df_multi, meta_multi = pyreadstat.read_file_multiprocessing(pyreadstat.read_sav, fpath, num_processes=4, row_offset=10, row_limit=3)
+        df_single, meta_single = pyreadstat.read_sav(fpath, row_offset=10, row_limit=3)
+        self.assertTrue(df_multi.equals(df_single))
+
+    def test_chunk_reader_multiprocess_chunksize_smaller_than_processes(self):
+        fpath = os.path.join(self.basic_data_folder, "sample_large.sav")
+        reader = pyreadstat.read_file_in_chunks(pyreadstat.read_sav, fpath, chunksize=2, offset=0, limit=4, multiprocess=True, num_processes=4)
+        alldfs = list()
+        for df, meta in reader:
+            alldfs.append(df)
+        df_multi = pd.concat(alldfs, axis=0, ignore_index=True)
+        df_single, meta_single = pyreadstat.read_sav(fpath, row_offset=0, row_limit=4)
+        self.assertTrue(df_multi.equals(df_single))
+
     def test_multiprocess_reader_xport(self):
         fpath = os.path.join(self.basic_data_folder, "sample.xpt")
         df_multi, meta_multi = pyreadstat.read_file_multiprocessing(pyreadstat.read_xport, fpath, num_rows=1000) 

tests/test_narwhalified.py

@@ -785,6 +785,25 @@ def test_chunk_reader_multiprocess_dict(self):
         df_single, meta_single = pyreadstat.read_sav(fpath, output_format=self.backend)
         self.assertTrue(df_multi.equals(df_single))
 
+    def test_multiprocess_reader_fewer_rows_than_processes(self):
+        # some workers get 0 rows to read, they should not read the rest of the file
+        fpath = os.path.join(self.basic_data_folder, "sample_large.sav")
+        df_multi, meta_multi = pyreadstat.read_file_multiprocessing(pyreadstat.read_sav, fpath, num_processes=4, row_offset=10, row_limit=3, output_format=self.backend)
+        df_single, meta_single = pyreadstat.read_sav(fpath, row_offset=10, row_limit=3, output_format=self.backend)
+        self.assertTrue(df_multi.equals(df_single))
+
+    def test_chunk_reader_multiprocess_chunksize_smaller_than_processes(self):
+        fpath = os.path.join(self.basic_data_folder, "sample_large.sav")
+        reader = pyreadstat.read_file_in_chunks(pyreadstat.read_sav, fpath, chunksize=2, offset=0, limit=4, multiprocess=True, num_processes=4, output_format=self.backend)
+        alldfs = list()
+        for df, meta in reader:
+            alldfs.append(nw.from_native(df))
+        df_multi = nw.concat(alldfs, how='vertical').to_native()
+        if self.backend == "pandas":
+            df_multi = df_multi.reset_index(drop=True)
+        df_single, meta_single = pyreadstat.read_sav(fpath, row_offset=0, row_limit=4, output_format=self.backend)
+        self.assertTrue(df_multi.equals(df_single))
+
     def test_multiprocess_reader_xport(self):
         fpath = os.path.join(self.basic_data_folder, "sample.xpt")
         df_multi, meta_multi = pyreadstat.read_file_multiprocessing(pyreadstat.read_xport, fpath, num_rows=1000, output_format=self.backend)