diff --git a/pyreadstat/pyreadstat.py b/pyreadstat/pyreadstat.py index 05fc0c4..993497d 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 745ded9..9e3299f 100644 --- a/tests/test_narwhalified.py +++ b/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)