diff --git a/sdks/python/apache_beam/dataframe/io.py b/sdks/python/apache_beam/dataframe/io.py index 21eab0b82faf..7e7f53e0d485 100644 --- a/sdks/python/apache_beam/dataframe/io.py +++ b/sdks/python/apache_beam/dataframe/io.py @@ -565,6 +565,7 @@ def _read(self, size=-1): self._buffer = self._underlying.read(size) if not self._buffer: + self._tracker.try_claim(self._tracker.current_restriction().stop) self._done = True return self._empty diff --git a/sdks/python/apache_beam/dataframe/io_test.py b/sdks/python/apache_beam/dataframe/io_test.py index 4cd502d1b8d7..d35517ba690b 100644 --- a/sdks/python/apache_beam/dataframe/io_test.py +++ b/sdks/python/apache_beam/dataframe/io_test.py @@ -126,6 +126,12 @@ def test_wide_csv_with_dtypes(self): pcoll = p | beam.io.ReadFromCsv(f'{input}tmp.csv', dtype=str) assert_that(pcoll | beam.Map(max), equal_to(['99'])) + def test_empty_csv_read(self): + input = self.temp_dir({'empty.csv': 'col1,col2,col3\n'}) + with beam.Pipeline() as p: + pcoll = p | beam.io.ReadFromCsv(input + 'empty.csv') + assert_that(pcoll, equal_to([])) + def test_sharding_parameters(self): data = pd.DataFrame({'label': ['11a', '37a', '389a'], 'rank': [0, 1, 2]}) output = self.temp_dir()