diff --git a/CHANGES.md b/CHANGES.md index d1910513da58..ed4f05533b1d 100644 --- a/CHANGES.md +++ b/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)). diff --git a/sdks/python/apache_beam/io/concat_source.py b/sdks/python/apache_beam/io/concat_source.py index 35ae6fe27317..bd6b7396d5d3 100644 --- a/sdks/python/apache_beam/io/concat_source.py +++ b/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] diff --git a/sdks/python/apache_beam/io/concat_source_test.py b/sdks/python/apache_beam/io/concat_source_test.py index efa24b3975fc..3cd169df5069 100644 --- a/sdks/python/apache_beam/io/concat_source_test.py +++ b/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