4
\$\begingroup\$

I wanted to write an iterator that would split items up into groups, so i can bulk insert into mysql without crashing or overwhelming the server.

This is the implementation I came up with. I am hoping that when yielding from the list its not creating a copy / clone of the list. I need more experience with python internals to be able to debug / performance profile this method i wrote.

def iter_group(iterable, batch_size:int):
    iterable_type = type(iterable)
    length = len(iterable)
    start = batch_size*-1
    end = 0
    while(end < length):
        start += batch_size
        end += batch_size
        if iterable_type == list:
            yield (iterable[i] for i in range(start,min(length-1,end)))
        else:
            yield iterable[start:end]

#test data

items = range(0,10001)
for item_group in iter_group(items, 1000):
    for item in item_group:
        print(item,end=' ')
\$\endgroup\$
4
\$\begingroup\$

We can improve readability of your code as follows:

  • using type annotations partially looks incomplete, when we can be more descriptive:

    from typing import (Iterable, 
                        TypeVar)
    
    ElementType = TypeVar('ElementType')
    
    
    def iter_group(iterable: Iterable[ElementType],
                   batch_size: int) -> Iterable[Iterable[ElementType]]:
        ...
    
  • checking if object has given type should be done using isinstance built-in function, so instead of

    iterable_type = type(iterable)
    ...
    if iterable_type == list:
        ...
    

    we can simply write

    if isinstance(iterable, list):
        ...
    

    (btw I don't understand why are you treating list as a special case)

  • if you want negate a number then we can use - unary operator, no need to multiply by -1, so instead of

    start = batch_size * -1
    

    we can write

    start = -batch_size
    
  • no need in parentheses for while-loop condition, we are not in C/C++/JAVA, we can simply write:

    while end < length:
        ...
    
  • setting start to -batch_size and end to 0 and change them right after that seems redundant when we can set them to 0 and batch_size respectfully and increment at the end of the while loop body, so instead

    start = batch_size * -1
    end = 0
    while end < length:
        start += batch_size
        end += batch_size
        ...
    

    we can write

    start = 0
    end = batch_size
    while end < length:
        ...
        start += batch_size
        end += batch_size
    

If I understood correctly what you are trying to achieve is to split iterable into evenly sized chunks, which is a well-known StackOverflow question, but accepted answer works only for sequences (str/list/tuple) and so does your solution. For example it won't work for potentially infinite sequences like classic Fibonacci numbers generator

>>> def fibonacci():
        a, b = 0, 1
        while True:
            yield a
            a, b = b, a + b
>>> next(iter_group(fibonacci(), 5))
Traceback (most recent call last):
    ...
        length = len(iterable)
TypeError: object of type 'generator' has no len()

If you want to work with arbitrary iterables (e.g. generators), we may use itertools module and this brilliant solution which uses iter built-in function form with sentinel value

from itertools import islice
from typing import (Iterable,
                    Tuple,
                    TypeVar)

ElementType = TypeVar('ElementType')


def chunks(iterable: Iterable[ElementType],
           batch_size: int) -> Iterable[Tuple[ElementType, ...]]:
    iterator = iter(iterable)
    return iter(lambda: tuple(islice(iterator, batch_size)), ())

but if we make a simple benchmark with

def sequence_chunks(iterable: Sequence[ElementType],
                    batch_size: int) -> Iterable[Sequence[ElementType]]:
    for start in range(0, len(iterable), batch_size):
        yield iterable[start:start + batch_size]

like

import timeit

...
print('original solution',
      min(timeit.repeat('list(iter_group(iterable, 1000))',
                        'from __main__ import iter_group\n'
                        'iterable = range(0, 10001)',
                        number=10000)))
print('Ned Batchelder\'s solution',
      min(timeit.repeat('list(sequence_chunks(iterable, 1000))',
                        'from __main__ import sequence_chunks\n'
                        'iterable = range(0, 10001)',
                        number=10000)))
print('senderle\'s solution',
      min(timeit.repeat('list(chunks(iterable, 1000))',
                        'from __main__ import chunks\n'
                        'iterable = range(0, 10001)',
                        number=10000)))

on my laptop with Windows 10 and Python 3.5 we'll have

original solution 0.07320549999999999
Ned Batchelder's solution 0.06249870000000002
senderle's solution 2.6072023999999994

so how can we have both speed and handle cases with non-sequence iterables?

Here comes functools module with singledispatch function decorator. We can use it like

import timeit
from collections import abc
from functools import singledispatch
from itertools import islice
from typing import (Iterable,
                    Sequence,
                    TypeVar)

ElementType = TypeVar('ElementType')


@singledispatch
def chunks(iterable: Iterable[ElementType],
           batch_size: int) -> Iterable[Iterable[ElementType]]:
    iterator = iter(iterable)
    return iter(lambda: tuple(islice(iterator, batch_size)), ())


@chunks.register(abc.Sequence)
def sequence_chunks(iterable: Sequence[ElementType],
                    batch_size: int) -> Iterable[Iterable[ElementType]]:
    for start in range(0, len(iterable), batch_size):
        yield iterable[start:start + batch_size]

after that call to chunks will end up in sequence_chunks for sequences and in general chunks for all other cases.

But

print('single-dispatched solution',
      min(timeit.repeat('list(chunks(iterable, 1000))',
                        'from __main__ import chunks\n'
                        'iterable = range(0, 10001)',
                        number=10000)))

gives

single-dispatched solution 0.0737681

so we lose some time during dispatching, but we can save some time and space if we will use itertools.islice for sequences as well like

@chunks.register(abc.Sequence)
def sequence_chunks(iterable: Sequence[ElementType],
                    batch_size: int) -> Iterable[Iterable[ElementType]]:
    iterator = iter(iterable)
    for _ in range(ceil_division(len(iterable), batch_size)):
        yield islice(iterator, batch_size)


def ceil_division(left_number: int, right_number: int) -> int:
    """
    Divides given numbers with ceiling.
    """
    # based on https://stackoverflow.com/a/17511341/5997596
    return -(-left_number // right_number)

which gives

complete single-dispatched solution 0.03895900000000002

P.S.

@AJNeufeld's solution gives

0.0647120000000001

on my laptop

\$\endgroup\$
  • \$\begingroup\$ Note: list(sequence_chunks(range(0, 10001), 1000)) return [ <itertools.islice object>, <itertools.islice object>, ..., <itertools.islice object> ] where list(iter_group(range(10001), 1000)) returns [range(0, 1000), range(1000, 2000), ..., range(10000, 10001)]. And list(chunks(range(0, 10001), 1000)) returns [(0, 1, 2, ..., 998, 999), (1000, 1001, ..., 1999), ... , (10000, ) ]! These are very, very different animals. \$\endgroup\$ – AJNeufeld Dec 10 '18 at 23:04
  • \$\begingroup\$ Note: itertools.islice() has not advanced the iterator (singular) in the returned list of slices. If we assign x = list(sequence_chunks(range(0, 10001), 1000)), and then ask for the next() values of the third, second and first iterators, [ next(x[2]), next(x[1]), next(x[0])] you might expect [2000, 1000, 0], but you will actually get [0, 1, 2]! \$\endgroup\$ – AJNeufeld Dec 10 '18 at 23:21
  • \$\begingroup\$ @AJNeufeld: I know, but in OP there is a statement (iterable[i] for i in range(start,min(length-1,end))), so I thought iterators are satisfactory choice, especially if author wants to insert records in batches \$\endgroup\$ – Azat Ibrakov Dec 11 '18 at 8:23
  • \$\begingroup\$ ``` start = 0 end = batch_size while end < length: ... start += batch_size end += batch_size ``` \$\endgroup\$ – Ben Dec 12 '18 at 1:21
0
\$\begingroup\$

Your code cannot work with iterables of unknown length:

def unknown_length():
    i = 0
    while i < 12:
        yield i
        i += 1

items = iter(unknown_length())
for item_group in iter_group(items, 5):
    for item in item_group:
        print(item,end=' ')
    print()

TypeError: object of type 'generator' has no len()


This code (iterable[i] for i in range(start,min(length-1,end))) is creating a new tuple containing a shallow copy of each group of batch_size items, although only 1 at a time so you aren't ending up with a complete copy of the iterable all at once, requiring double the memory.


The range(start,min(length-1,end)) omits the your last data item! The end limit is non-inclusive, so the last item you get is iterable[length-2].


The following iter_group yields a nested generator which yield items until batch_size items have been returned. Then, the next nested generator will be yielded, until the list item has been produced.

def iter_group(iterable, batch_size:int):
    it = iter(iterable)

    def group(first):
        yield first
        try:
            for _ in range(batch_size - 1):
                yield next(it)
        except StopIteration:
            pass

    try:
        while True:
            yield group(next(it))
    except StopIteration:
        pass

If you want to try to optimize this for sliceable iterables, you could try:

def iter_group(iterable, batch_size:int):

    try:
        # Maybe we can use slices?
        start = 0
        n = len(iterable)
        while start < n:
            yield iterable[start:start+batch_size]
            start += batch_size

    except TypeError:
        # Nope!  Couldn't get either the length or a slice.

        it = iter(iterable)

        def group(first):
            yield first

            try:
                for _ in range(batch_size - 1):
                    yield next(it)
            except StopIteration:
                pass

        try:
            while True:
                yield group(next(it))
        except StopIteration:
            pass
\$\endgroup\$
  • \$\begingroup\$ statement (iterable[i] for i in range(start,min(length-1,end))) does not create a new tuple, but a generator expression \$\endgroup\$ – Azat Ibrakov Dec 10 '18 at 22:26
  • \$\begingroup\$ @AzatIbrakov You're right - good catch. \$\endgroup\$ – AJNeufeld Dec 10 '18 at 22:30
  • \$\begingroup\$ @AzatIbrakov does it return a NEW generator expression, or does it set the state of the existing expression? \$\endgroup\$ – Ben Dec 12 '18 at 1:16
  • 1
    \$\begingroup\$ @Ben It will yield a new generator expression each time. \$\endgroup\$ – AJNeufeld Dec 12 '18 at 2:07

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.