From 8862fbfa9ae3d1c03fbcb57395fa2d797443ff87 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 25 Aug 2026 15:17:07 -0400 Subject: [PATCH] Fix ZeroDivisionError at initial invocation of monitoring_info when there is no work yet --- sdks/python/apache_beam/io/iobase.py | 12 ++++++++---- sdks/python/apache_beam/io/iobase_test.py | 21 +++++++++++++++++++++ 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/io/iobase.py b/sdks/python/apache_beam/io/iobase.py index aa03280050fa..b7be85935997 100644 --- a/sdks/python/apache_beam/io/iobase.py +++ b/sdks/python/apache_beam/io/iobase.py @@ -1716,15 +1716,19 @@ def total_work(self) -> float: def fraction_completed(self) -> float: if self._fraction is not None: return self._fraction - else: - return float(self._completed) / self.total_work + total_work = self.total_work + if total_work == 0.: + return 1.0 + return float(self._completed) / total_work @property def fraction_remaining(self) -> float: if self._fraction is not None: return 1 - self._fraction - else: - return float(self._remaining) / self.total_work + total_work = self.total_work + if total_work == 0.: + return 0.0 + return float(self._remaining) / total_work def with_completed(self, completed: int) -> 'RestrictionProgress': return RestrictionProgress( diff --git a/sdks/python/apache_beam/io/iobase_test.py b/sdks/python/apache_beam/io/iobase_test.py index 60c261563155..4621fa45bb15 100644 --- a/sdks/python/apache_beam/io/iobase_test.py +++ b/sdks/python/apache_beam/io/iobase_test.py @@ -261,5 +261,26 @@ def test_read_unbounded_serializes_as_expanded_composite(self): self.assertTrue(read_transforms[0].subtransforms) +class RestrictionProgressTest(unittest.TestCase): + def test_restriction_progress(self): + # Total work == 0 edge cases (avoids ZeroDivisionError) + progress_zero_int = iobase.RestrictionProgress(completed=0, remaining=0) + self.assertEqual(progress_zero_int.fraction_completed, 1.0) + self.assertEqual(progress_zero_int.fraction_remaining, 0.0) + + # Progress with completed and remaining + progress_work = iobase.RestrictionProgress(completed=25, remaining=75) + self.assertEqual(progress_work.completed_work, 25) + self.assertEqual(progress_work.remaining_work, 75) + self.assertEqual(progress_work.total_work, 100) + self.assertEqual(progress_work.fraction_completed, 0.25) + self.assertEqual(progress_work.fraction_remaining, 0.75) + + # Progress with fraction + progress_frac = iobase.RestrictionProgress(fraction=0.4) + self.assertEqual(progress_frac.fraction_completed, 0.4) + self.assertAlmostEqual(progress_frac.fraction_remaining, 0.6) + + if __name__ == '__main__': unittest.main()