Python: How can I send an iterator to two different consumers without loading the entire thing into memory?

Question

I have an iterator that is consumed by two functions (mean_summarizer and std_summarizer in example below). I want both functions to process the iterator, WITHOUT ever having to load the entire iterator into memory at once.

Below is a minimal example (also in Colab) that provides the correct result, EXCEPT that it involves loading the entire input into memory at once. No need to understand the fancy code inside mean_summarizer, std_summarizer, and last - it's mainly like that for brevity.

Question is: What is the cleanest way to re-implement summarize_input_stream without changing the function signature (just the inside), such that its memory usage does not scale with length of the input stream?

I have a feeling coroutines are involved, but I don't know how to use them.

import numpy as np
from typing import Iterable, Mapping, Callable, Any

def summarize_input_stream(  # Run the input stream through multiple summarizers and collect results
        input_stream: Iterable[float],
        summarizers: Mapping[str, Callable[[Iterable[float]], float]]
) -> Mapping[str, float]:
    inputs = list(input_stream)  # PROBLEM IS HERE <-- We load entire stream into memory at once
    return {name: summarizer(inputs) for name, summarizer in summarizers.items()}

def last(iterable: Iterable[Any]) -> Any:  # Just returns last element of iterable
  return max(enumerate(iterable))[1]

def mean_summarizer(stream: Iterable[float]) -> float:  # Just computes mean online and returns final value
  return last(avg for avg in [0] for i, x in enumerate(stream) for avg in [avg*i/(i+1) + x/(i+1)])

def std_summarizer(stream: Iterable[float]) -> float:   # Just computes standard deviation online and returns final value
  return last(cumsum_of_sq/(i+1) - (cumsum/(i+1))**2 for cumsum_of_sq, cumsum in [(0, 0)] for i, x in enumerate(stream) for cumsum_of_sq, cumsum in [(cumsum_of_sq+x**2, cumsum+x)])**.5

summary_stats = summarize_input_stream(
    input_stream=(np.random.randn()*2+3 for _ in range(1000)),
    summarizers={'mean': mean_summarizer, 'std': std_summarizer}
)
print(summary_stats)
# e.g. {'mean': 3.020903422847062, 'std': 1.943724669289156}

Answer 1

I found a solution that does not involve changing the signature of summarize_input_stream. It launches one thread per summarizer and feeds each one incrementally via a separate blocking queue (link to Colab).

import numpy as np
from typing import Iterable, Mapping, Callable, Any
from threading import Thread
from queue import Queue
from functools import partial


def summarize_input_stream(  # Run the input stream through multiple summarizers and collect results
        input_stream: Iterable[float],
        summarizers: Mapping[str, Callable[[Iterable[float]], float]]
) -> Mapping[str, float]:
    POISON_PILL = object()
    def run_summarizer(summarizer: Callable[[Iterable[float]], float], queue: Queue) -> float:
        result = summarizer(iter(queue.get, POISON_PILL)) # Waits until the food is ready to eat
        queue.put(result)  # Use the queue the other way around to return the result
    queues = [Queue(maxsize=1) for _ in summarizers]  # <-- Note We could can probably be more time-efficient if we increase maxsize, which should cause less thread switching at the cost of more memory usage
    threads = [Thread(target=partial(run_summarizer, summarizer, queue)) for summarizer, queue in zip(summarizers.values(), queues)]
    for t in threads:
        t.start()
    for inp in input_stream:
        for queue in queues:
            queue.put(inp)  # Waits until the summarizer is hungry to feed it
    for queue in queues:
        queue.put(POISON_PILL)  # Stop the iteration
    for t in threads:
        t.join()
    results = [queue.get() for queue in queues]
    return {name: result for name, result in zip(summarizers, results)}

def last(iterable: Iterable[Any]) -> Any:  # Just returns last element of iterable
    return max(enumerate(iterable))[1]

def mean_summarizer(stream: Iterable[float]) -> float:  # Just computes mean online and returns final value
    return last(avg for avg in [0] for i, x in enumerate(stream) for avg in [avg * i / (i + 1) + x / (i + 1)])

def std_summarizer(stream: Iterable[float]) -> float:  # Just computes standard deviation online and returns final value
    return last(cumsum_of_sq / (i + 1) - (cumsum / (i + 1)) ** 2 for cumsum_of_sq, cumsum in [(0, 0)] for i, x in enumerate(stream) for cumsum_of_sq, cumsum in
                [(cumsum_of_sq + x ** 2, cumsum + x)]) ** .5

summary_stats = summarize_input_stream(
    input_stream=(np.random.randn() * 2 + 3 for _ in range(1000)),
    summarizers={'mean': mean_summarizer, 'std': std_summarizer}
)
print(summary_stats)
# e.g. {'mean': 3.020903422847062, 'std': 1.943724669289156}

Answer 2

You can't do this. A generalized iterator can only be processed once, and to make it possible to process it twice, you need to store it in some way, either by listifying it as you're doing, or using itertools.tee (which, if one of the tee-d iterators is consumed completely before the other pulls any items, is morally equivalent; it has to store all of the data internally).

The only way to make this work is if you use a single summarizer that processes the input once and computes all relevant summaries at the same time.

Answer 3

As @ShadowRanger says, it seems (unless proven otherwise) that there is no solution to the question as stated that does not involve multi-threading.

However, with a minor change to the signature of summarize_input_stream (which admittedly violates the rules which I myself wrote), we can get the result without paying the (memory) price.

The trick (which I'll call "genfunctrification" unless it has a pre-existing name) is:

• We turn our summarizers from functions of type: Callable[[Iterable[float]], float] ...
• .. into generator-functions of type Callable[[Iterable[float]], Iterable[Callable[[], float]]]
• Note that we could have just made them return the results directly (Callable[[Iterable[float]], Iterable[float]), but this would involve wastefully recomputing things like (cumsum_of_sq/(i+1) - (cumsum/(i+1))**2)**.5 on each iteration, when we only actually need it on the last, so instead we make our iterators yield Callables, which can compute the result only when needed (after the last iteration).

The modified code (and Colab link).

import numpy as np
from typing import Iterable, Mapping, Callable, Any, Sequence
import itertools

def summarize_input_stream(  # Run the input stream through multiple summarizers and collect results
        input_stream: Iterable[float],
        summarizers: Mapping[str, Callable[[Iterable[float]], Iterable[Callable[[], float]]]]
) -> Mapping[str, float]:
    input_streams_teed = itertools.tee(input_stream, len(summarizers))
    result_getter_streams: Sequence[Iterable[Callable[[], float]]] = [summarizer(stream_copy) for summarizer, stream_copy in zip(summarizers.values(), input_streams_teed)]
    final_results = last([f() for f in func_tup] for func_tup in zip(*result_getter_streams))
    return {name: r for name, r in zip(summarizers, final_results)}

def last(iterable: Iterable[Any]) -> Any:  # Just returns last element of iterable
  return max(enumerate(iterable))[1]

def mean_summarizer(stream: Iterable[float]) -> Iterable[Callable[[], float]]:  # Just computes mean online and returns final value
  return ((lambda: avg) for avg in [0] for i, x in enumerate(stream) for avg in [avg*i/(i+1) + x/(i+1)])

def std_summarizer(stream: Iterable[float]) ->  Iterable[Callable[[], float]]:   # Just computes standard deviation online and returns final value
  return ((lambda: (cumsum_of_sq/(i+1) - (cumsum/(i+1))**2)**.5) for cumsum_of_sq, cumsum in [(0, 0)] for i, x in enumerate(stream) for cumsum_of_sq, cumsum in [(cumsum_of_sq+x**2, cumsum+x)])

summary_stats = summarize_input_stream(
    input_stream=(np.random.randn()*2+3 for _ in range(1000)),
    summarizers={'mean': mean_summarizer, 'std': std_summarizer}
)
print(summary_stats)
# e.g. {'mean': 3.020903422847062, 'std': 1.943724669289156}

• Note: since this particular implementation depends on dicts maintaining their order, it will only work reliably in Python 3.7+ - though it could easily be backported if needed.