From 8b702023e787839845a816dd106eaaaf7801dfed Mon Sep 17 00:00:00 2001 From: ManvithPanyam <250704031+ManvithPanyam@users.noreply.github.com> Date: Sun, 2 Aug 2026 22:01:00 +0530 Subject: [PATCH] fix(dataframe): claim remaining restriction range on empty/header-only CSV reads Signed-off-by: ManvithPanyam <250704031+ManvithPanyam@users.noreply.github.com> --- sdks/python/apache_beam/dataframe/io.py | 1 + sdks/python/apache_beam/dataframe/io_test.py | 6 ++++++ 2 files changed, 7 insertions(+) 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()