Flatten itertools.chain objects?

[tim-one]

Feature or enhancement

Proposal:

If an iterator is repeatedly chained, speed slows down the longer the chain of chains gets. For example,

from itertools import count, chain, repeat
from time import perf_counter as now

it = chain([None], count(0))
expect = 0
while True:
    start = now()
    for _ in repeat(None, 1000):
        x = next(it)
        assert x is None
        x = next(it)
        assert x == expect
        expect += 1
        it = chain([None], it)
    finish = now()
    print(format(expect, '_'), finish - start)

on my box just now produced (Win64, Python 3.14.5):

1_000 0.0058611
2_000 0.0192824
3_000 0.026536800000000003
4_000 0.0335665
5_000 0.03840320000000001
6_000 0.049022399999999994
7_000 0.061446399999999984
8_000 0.07017669999999998
9_000 0.0710173
10_000 0.08084079999999999
11_000 0.09169510000000003
12_000 0.09909350000000006
13_000 0.1044349
14_000 0.12177879999999996
15_000 0.12580650000000015
16_000 0.13390919999999995
Traceback (most recent call last):
  ...
    x = next(it)
        ^^^^^^^^
RecursionError: maximum recursion depth exceeded

Could the chain constructor peek inside chain object arguments to skip over their exhausted iterators somehow at construction time?

In the example, it starts as chain([None], count(0)). By the time the next chain is done, the [None] part has already been exhausted. So the following chain([None], it) could conceivably notice that it is already a chain object, but that it's been reduced to wrapping only a single still-active iterator, and grab that remaining iterator itself rather than wrap all of the input it chain object.

I haven't thought through all this, but wanted to put this out here before I forget again.

Has this already been discussed elsewhere?

This is a minor feature, which does not need previous discussion elsewhere

Links to previous discussion of this feature:

No response

[tim-one]

Offline, @rhettinger invited me to assign this to him, but he's busy with "real life" for now. So I assigned him. There's no rush 😄.

[tim-one]

Two things: first, it doesn't always die with RecursionError on Windows. Sometimes it goes on longer, but the process crashes with a non-zero exit code (no message in the console window, just stops, with a %ERRORLEVEL% set to non-zero):

...
44_000 0.5022625999990851
45_000 0.5171336999628693
46_000 0.5790993999689817

C:\MyPy>echo %errorlevel%
-1073741571

Google says

This is the signed integer representation of the hexadecimal error code 0xC00000FD, which indicates a stack overflow or stack exhaustion error.

Second, here's a Python implementation of the idea fleshed out. The example can apparently "run forever" with this and without slowing down:

from collections import deque as _deque

class chain:
    __slots__ = "active",

    def __init__(self, *iterables):
        active = self.active = _deque()
        for it in map(iter, iterables):
            if isinstance(it, chain):
                active.extend(it.active)
            else:
                active.append(it.__next__)

    def __iter__(self):
        return self

    def __next__(self):
        active = self.active
        while active:
            try:
                return active[0]()
            except StopIteration:
                active.popleft()
        raise StopIteration

I'm not suggesting any particular implementation, though. I was just having fun 😄, hoping that code makes the idea clear.

[rhettinger]

FWIW, I've encountered deeply nested chains exactly once in the last two decades and the stacking cost was easily solved by switching to itertools.tee. See: more-itertools/more-itertools#909

I'll contemplate this more but want to express a general goal to keep most itertools as dirt simple as possible, ideally with a simple generator equivalents that are easily learned and remembered.

The suggested implementation adds buffering to a tool where that is unexpected. Also, the extra test slow downs the iterator creation. This matters in legit use cases such as flattening a long stream of 2-tuples. If it were implemented in C, it would have to contend with no-gil race conditions where the original iterator, its chain, and the outer chain could be advanced at the same time, so some kind of locking would be likely needed. Multiphase initialization adds complexity as well (see the get_random_state in Modules\_randommodule.c for a taste).

We have flattening for tee and partial, so this isn't without precedent. It is just a matter of how much optimization magic we're willing to add, how much complexity we're willing to add, and whether this is commonly needed. We did put in a hack for += with strings but that was only because it was such a common pitfall.

Also, as a matter of "clean programming practices" my spidey-sense tingles at the thought of reaching inside another object like this. In the case of tee, share state was unavoidable. In this case, it feels like a hack. I occasionally put up with hacks like reusing objects when the refcnt is one, but have paid a huge price for it over the years as a point of contention with other developers who want to eliminate it, undermine its reliability, or need better atomicity in a non-gil environment. If we did implement this suggestion, I suspect that it will be a problem over and over again in the future.

[tim-one]

Don't much care to argue. I don't see any fundamentally new "synch problems". People can already do:

it = chain(it2, it3)

and iterate it2 and/or it3 "at the same time" it is poking at them. I expect that typically it2 and it3 won't be used directly again. But if they aren't, then flattening chain would be just as benign. If it's not flattened, it's still the case that it2 and/or it3 can be iterated at the same time it is poking at them. There is no sense in which the chain object "takes over" ownership of its arguments. It just holds its own references to them. Nothing unsavory about that to my eyes.

The project where this most recently came up was using tee instead, but gave up on it, because it was slowing dramatically as time went on. Not due to lack of sharing the same data stream, but due to that the project created and used so very many (thousands) of them simultaneously that they appeared to be kicking everything else out of HW data caches.

I didn't write the following; the person digging into this did:

from itertools import *
from time import perf_counter as time

n = 10**6
batchsize = 20
xss = list(batched(range(n), batchsize))

for _ in range(5):

    it = zip(*[xs for xs in xss])
    t = time()
    for _ in it:
        del _
    t_without = time() - t
    del it

    it = zip(*[tee(xs, 1)[0] for xs in xss])
    t = time()
    for _ in it:
        del _
    t_with = time() - t
    del it

    print(*(round(t / n * 1e9) for t in [t_without, t_with]))

I'm not sure what it's actually showing 😦, but the extra nanoseconds per element after wrapping things in tee() goes up as n increases. But I'm inclined to believe that it's not just the extra memory burden of tee driving it. It's not "cache-friendly" regardless.

So the project switched to using chain instead. And bumped into worse slowdowns for the reasons I opened this issue. They do appear to be "shallow" reasons to me,

I'm sure your experience differs from mine. In mine, people shy away from tee because they just plain don't understand it, but do understand chain. Used to be worse when the docs cautioned they should never use the argument to tee() directly again.

In fact, the more-itertools issue report you linked to seems thoroughly typical to me: when people want to peek ahead, the first thing they think of is chain(). Even people with lots of relevant experience.

Happily, they usually don't get into trouble, because they don't usually need to peek ahead over and over and over and ... again.

Relatively few people actually complain when they do get in trouble. Instead they live with it, move to something else, or give up.

Anyway, thanks for the consideration! In the turnabout-is-fair-play department, in over 3 decades I've never had a need to flatten a long stream of 2-tuples 😜 .

BTW, I don't really care about the RecursionError or "silent crash" symptoms. I never experienced them in real life, just noted them because the contrived demonstration happened to trigger them.

[serhiy-storchaka]

I think this is pretty common problem. Since iterators do not support peek or put back operations, users need to use chain() and tee(). The latter is less obvious, and depends heavily on the fact that objects are destroyed immediately after the refcount drops to zero in CPython. chain() is used for these purposes in tokenize and glob, but not recursively, so this doesn't cause an issue.

There is another related issue. When you create chain([large_object], iterator), the reference to large_object is kept even after advancing the chain iterator. This is because chain(*args) is implemented as chain.from_iterable(args). Implementing it as an equivalent to chain.from_iterable(iter(deque(args).popleft)) would help, but it would complicate the implementation.

If the same code can solve both problems, I'd say it's worth it. If not, then no.

[pochmann3]

I've seen this a couple times, for example in this Stack Overflow answer with 34 upvotes for flattening an arbitrarily nested lists structure:

remainder = itertools.chain(first, remainder)

As you'd expect, it did quite poorly in my benchmark with deeply nested data.

I think that's the common (ab)use case: people want to insert another iterable (or just one element) at the current position of the chain. So maybe a cleaner way to support that would be to offer exactly that: an insert(iterable) method. Then the above could be done efficiently like this:

remainder.insert(first)

Although it would be unusual for an iterator to have anything other than the minimum iterator methods. Then again, generator iterators have send().

(Maybe "prepend" would be a better name, although that feels a bit odd to me when the chain has already partially been iterated, as you'd then "prepend" in the middle.)

[pochmann3]

And with that method, the whole flattening function could neatly be written like this and would likely be the fastest:

def flatten(lst):
    xs = chain(lst)
    for x in xs:
        if isinstance(x, list):
            xs.insert(x)
        else:
            yield x

[efimov-mikhail]

What do you think about optional keyword parameter for itertools.chain that force flattening and releasing, if possible?

[tim-one]

In case it helps, I'm going to sketch the heart of what triggered this report. Some of us are working on a "lazy" mergesort. which yields the sorted values of its argument one at a time. In fact, the pure Python implementation advanced to the stage where it's a little faster than list.sort() on randomly ordered data when to run to completion, and much faster than heapq.nsmallest() to find the smallest n when n isn't "very small".

But in cases of partially ordered input, ,sort()'s "galloping mode" gives it an enormous advantage. I didn't at first think it was possible to adapt that to an iterator-based approach.

But it is! Where it is the iterable producing a run, exponential search can be done like so:

buffer = list(islice(it, power_of_2))
use bisect on the buffer to see whether we've
    found the right spot to insert the other run's value
if not, double `power_of_2`, `yield from buffer` and try again
but if the other run's value belongs at interior index `i` in the buffer:
    yield from islice(buffer, i)
    yield the other run's value
    and then the mystery!
         how to push the rest of buffer back on to `it`

The buffer can become arbitrarily large. In the case of using tee, the data is also stored in the tee object's hidden deque-like gimmick, and that alone is uncomfortable. But it's easy to live with. We work with a tee'd clone of it instead, and

del buffer # don't yield from the buffer in this case
yield from islice(original_it, i) # replay the values from the hidden deque

and that's all. The rest of the buffer is still waiting to be yielded, and the calling code is none the wiser.

The alternative is to del buffer[:i] and then it = chain(buffer, it). This can lead to an unbounded number of nested chains, and the buffers being chained can be arbitrarily large.

So the pure-Python implementation of chain I gave addressed both: nested chains are flattened out of existence, and references to exhausted iterables are dropped ASAP. In context, the original it becomes trash as soon as the new chain object is bound to it, and can be garbage-collected at once.

[pochmann3]

it = chain(buffer, it)

Which would likewise be it.insert(buffer) with my alternative.

A Python demo (also dropping exhausted iterables):

class chain supporting inserts

class chain:

    def __init__(self, *iterables):
        self.iterables = list(reversed(iterables))

    def __iter__(self):
        return self

    def __next__(self):
        while self.iterables:
            self.iterables[-1] = iter(self.iterables[-1])
            for x in self.iterables[-1]:
                return x
            self.iterables.pop()
        raise StopIteration

    def insert(self, iterable):
        self.iterables.append(iterable)


def flatten(lst):
    xs = chain(lst)
    for x in xs:
        if isinstance(x, list):
            xs.insert(x)
        else:
            yield x

print(*flatten([[[1, [2], 3], [4, 5]], 6]))

Attempt This Online!

And out of curiosity a generator version, using send(iterable) to insert:

generator chain supporting inserts

def chain(*iterables):
    def gen():
        insert = yield
        while True:
            if insert is None:
                while iterables:
                    iterator = iter(iterables.pop())
                    for x in iterator:
                        insert = yield x
                        if insert is not None:
                            iterables.append(iterator)
                            break
                    else:
                        continue
                    break
                else:
                    return
            else:
                iterables.append(insert)
                insert = yield
    iterables = list(reversed(iterables))
    it = gen()
    next(it)
    return it


def flatten(lst):
    xs = chain(lst)
    for x in xs:
        if isinstance(x, list):
            xs.send(x)
        else:
            yield x

print(*flatten([[[1, [2], 3], [4, 5]], 6]))

Attempt This Online!