GH-98363: Presize the list for batched()

[rhettinger]

Presizing the result lists saves space from list over-allocation and is faster than the PyList_Append() approach.

This is a short PR but it took a while to figure out that PyList_GetSlice() would make short work of the problem. Ideally, we need something like _PyTuple_Resize() to resize the list inplace using a realloc() and updating the allocated field.

• Issue: Add itertools.batched() #98363

[bedevere-bot]

🤖 New build scheduled with the buildbot fleet by @rhettinger for commit fd5b508 🤖

If you want to schedule another build, you need to add the ":hammer: test-with-buildbots" label again.

[sweeneyde]

Do we care how this handles strange iterators that keep going after having raised StopIteration?

>>> from itertools import batched
>>> class A:
...     def __init__(self):
...         self.i = 0
...     def __iter__(self):
...         return self
...     def __next__(self):
...         i = self.i
...         self.i += 1
...         if i == 7:
...             raise StopIteration()
...         return i
...
>>> x = batched(A(), 5)
>>> next(x)
[0, 1, 2, 3, 4]
>>> next(x)
[5, 6]
>>> next(x)
[8, 9, 10, 11, 12]

Maybe the Py_CLEAR(bo->it) could be moved to before this break;.

[rhettinger]

I'm thinking of leaving this as is. The current PR matches the behavior of the pure python version and the strange iterator works the same way with other tools:

>>> it = zip(range(10), A())
>>> list(it)
[(0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]
>>> list(it)
[(8, 8), (9, 9)]

>>> it = enumerate(A())
>>> list(it)
[(0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6)]
>>> next(it)
(7, 8)
>>> next(it)
(8, 9)

That said, I'm not really sure about this one and could easily be persuaded to add a Py_CLEAR(bo->it) to the code path for the short list.

[sweeneyde]

This does seem like an edge case that could go either way.

I'm noticing that some iterators do Py_CLEAR(obj->it) as soon as the next() call fails, while other objects have internal iterators that are never NULL. My concern is that batched() clears the iterator sometimes (when the first next() of a batch fails), but not other times (when a subsequent next() fails).

So I would think to do one or the other, but not both:

1. Clear the iterator immediately after any internal next() call fails, similar to pairwise, islice, chain, and cycle (sort of).

2. Never clear the iterator and remove the if (it == NULL) check altogether, similar to dropwhile, takewhile, accumulate, compress, zip and enumerate. Unfortunately, it would mean allocating and throwing away PyList_New(n), but that's in the rare case that someone is abusing the iterator.

Maybe removing the Py_CLEAR and the NULL check is most natural?

[sweeneyde]

Regardless of the decision about Py_CLEAR, I think this is a good improvement.

Out of curiosity, before this PR, I measured consume(batched(repeat(None, 1_000_000), 100)) to have a bimodal distribution of timings, with about half of runs at 41ms and half at 8ms (cache effects?). After this PR, I measured about 3ms consistently. This was on Windows, without PGO.

[rhettinger]

Thanks for the interesting line of thought. I like having this match the pure python equivalent. Also, I don't think that it would ever matter. Code like enumerate() and zip() has been deployed for 20 years with no ill effects. More-itertools has iterators than also succumb to your "strange iterator" and no one seems have had an issue.

The likely reason that you're getting a bimodal timing distribution is that PyList_Append() calls realloc() which uncontrollably runs in either of two modes, fast for an inplace resize when the memory can be extended inplace, and slow for a resize that has to move all the data because there isn't room for an inplace resize. The new code doesn't have that issue because PyList_New() makes a single call to malloc() which never moves data. The new code also has better memory utilization because it never returns an overallocated list.