From b62252eeca77b7be65e639e7544e1ad3ebdf6794 Mon Sep 17 00:00:00 2001 From: devteamaegis Date: Sun, 4 Oct 2026 18:54:38 -0400 Subject: [PATCH] Fix read_file_multiprocessing reading extra rows when a worker gets 0 rows When the number of rows to read is smaller than num_processes (for example row_limit=2 with 4 processes, or read_file_in_chunks with multiprocess=True and a chunksize smaller than num_processes), some workers were assigned a chunk of 0 rows. That 0 was passed as row_limit to the read function, where 0 means "no limit", so each of those workers read every remaining row of the file and the result contained duplicated rows. Chunks with 0 rows are now skipped. If all chunks are empty one is kept so that an empty dataframe with the right columns is still returned. --- pyreadstat/pyreadstat.py | 3 +++ tests/test_basic.py | 17 +++++++++++++++++ tests/test_narwhalified.py | 19 +++++++++++++++++++ 3 files changed, 39 insertions(+) diff --git a/pyreadstat/pyreadstat.py b/pyreadstat/pyreadstat.py index 97d5452..13bd34e 100644 --- a/pyreadstat/pyreadstat.py +++ b/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: diff --git a/tests/test_basic.py b/tests/test_basic.py index 84d7f7a..f92c89e 100644 --- a/tests/test_basic.py +++ b/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) diff --git a/tests/test_narwhalified.py b/tests/test_narwhalified.py index 4e51075..09193bd 100644 --- a/tests/test_narwhalified.py +++ b/tests/test_narwhalified.py @@ -781,6 +781,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)