[Python] Avoid ZeroDivisionError in ConcatRangeTracker for zero weights

CHANGES.md

@@ -104,6 +104,7 @@
 * (Python) `Duration` built from float seconds now rounds to the nearest microsecond instead of truncating, which could lose a microsecond ([#40263](https://github.com/apache/beam/issues/40263)).
 * (Python) `ReadFromText` with `escapechar` no longer skips a delimiter that directly follows an escaped delimiter, which merged two records into one ([#40459](https://github.com/apache/beam/issues/40459)).
 * (Python) DataFrame API: `agg` with a list that includes `'count'` or `'size'` now labels those rows `count`/`size` instead of `sum` ([#40504](https://github.com/apache/beam/issues/40504)).
+* (Python) `ConcatSource` no longer raises `ZeroDivisionError` when every source bundle reports a weight of 0, e.g. bundles smaller than 1 MB ([#21316](https://github.com/apache/beam/issues/21316)).
 * (Python) DataFrame API: `Series.value_counts(bins=...)` no longer always fails with a `TypeError` ([#40503](https://github.com/apache/beam/issues/40503)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 

sdks/python/apache_beam/io/concat_source.py

@@ -135,8 +135,8 @@ def _compute_cumulative_weights(source_bundles):
     min_diff = 1e-5
     # For the computation below, we need weights for all sources.
     # Substitute average weights for those whose weights are
-    # unspecified (or 1.0 for everything if none are known).
-    known = [s.weight for s in source_bundles if s.weight is not None]
+    # unspecified or zero (or 1.0 for everything if none are known).
+    known = [s.weight for s in source_bundles if s.weight]
     avg = sum(known) / len(known) if known else 1.0
     weights = [s.weight or avg for s in source_bundles]
 

sdks/python/apache_beam/io/concat_source_test.py

@@ -171,6 +171,15 @@ def test_position_at_fration(self):
     self.assertEqual(range_tracker.position_at_fraction(.601), (2, 17))
     self.assertEqual(range_tracker.position_at_fraction(1), (3, None))
 
+  def test_read_zero_weight_bundles_in_pipeline(self):
+    source = ConcatSource([
+        iobase.SourceBundle(0, RangeSource(0, 4), None, None),
+        iobase.SourceBundle(0, RangeSource(4, 8), None, None),
+    ])
+    with TestPipeline() as pipeline:
+      pcoll = pipeline | beam.io.Read(source)
+      assert_that(pcoll, equal_to(list(range(8))))
+
   def test_empty_source(self):
     read_all = source_test_utils.read_from_source