spdl.source.utils.MergeIterator

class MergeIterator(iterables: Sequence[Iterable[T]], *, weights: Sequence[float] | None = None, stop_after: int = 0, seed: int = 0)[source]

Bases: Iterable[T]

Iterate over given iterables and yield one item from each iterator.

Parameters:

  • iterables – The source iterables

  • weights – The sampling weight used to choose the next iterable. Values are interpreted with binary-float semantics, and selection remains exact with respect to those converted values. Sources with zero weight are skipped. If not provided, the given iterables are visited in the given order repeatedly.

  • stop_after –

    Determines the stop criteria or the behavior when one of the input iterables gets exhausted, Available values are;

    • 0: The iteration continues until all the input iterables are exhausted. (default)

    • n > 0: The iteration stops when the specified number of items are yielded or all the input iterables are exhausted before yielding n items.

    • -1: The iteration stops when one of the iterator is exhausted.

  • seed – Used to seed the random generator when weights is provided.

Example

>>> iterables = [
...     [0, 1, 2],
...     [10, 11, 12],
...     [20, 21, 22],
... ]
>>>
>>> print(list(MergeIterator(iterables)))
[0, 10, 20, 1, 11, 21, 2, 12, 22]
>>>
>>> # By default, it stops after one iterable gets exhausted.
>>> iterables = [
...     [0, 1, 2],
...     [10, 11],
...     [20, 21, 22],
... ]
>>>
>>> print(list(MergeIterator(iterables)))
[0, 10, 20, 1, 11, 21, 2]  # 22 is not included
>>>
>>> # Stop after yielding the given number of items
>>> print(list(MergeIterator(iterables, stop_after=5)))
[0, 10, 20, 1, 11]
>>>
>>> # stop_after>1 ignores the exhaustion.
>>> print(list(MergeIterator(iterables, stop_after=9)))
[0, 10, 20, 1, 11, 21, 2, 22]
>>>
>>> # Providing weights will pick up the iterable stocastically.
>>> iterables = [
...     [0, 1, 2],
...     [10, 11, 12],
...     [20, 21, 22],
... ]
>>> print(sorted(MergeIterator(iterables, stop_after=9, weights=[1, 1, 1])))
[0, 1, 2, 10, 11, 12, 20, 21, 22]

Changed in version 0.8.0: Exhausted weighted sources are removed from future draws, and invalid weight totals now raise ValueError. Weighted selection preserves the exact ratios of the converted binary-float weights.

__iter__() → Iterator[T][source]