Repository: dask/dask
File: dask/bag/core.py (Python)

Code before:

    Examples
    --------
    >>> def inc(x):
    ...     return x + 1
    >>> task = (sum, (list, (map, inc, [1, 2, 3])))
    >>> lazify_task(task)  # doctest: +ELLIPSIS
    (<built-in function sum>, (<class 'map'>, <function inc at ...>, [1, 2, 3]))
    """
    if isinstance(task, GraphNode):
        if isinstance(task, List) and len(task.args) < 50:
            return List(*[lazify_task(arg, False) for arg in task.args])
        if not isinstance(task, Task):
            return task
        if not start and task.func in (list, reify) and isinstance(task.args[0], Task):
            assert len(task.args) == 1
            task = task.args[0]
        if task.func is _execute_subgraph:
            subgraph = task.args[0]
            outkey = task.args[1]
            # If there is a reify at the start of this we don't want to act
            final_task = lazify_task(subgraph[outkey], True)
            subgraph = {
                k: lazify_task(v, False) for k, v in subgraph.items() if k != outkey
            }
            subgraph[outkey] = final_task
            return Task(
                task.key,
                _execute_subgraph,
                subgraph,
                outkey,
                *task.args[2:],
                **task.kwargs,
            )
        return Task(
            task.key,
            task.func,
            *[lazify_task(arg, False) for arg in task.args],
            **task.kwargs,
        )
    else:
        if type(task) is list and len(task) < 50:
            return [lazify_task(arg, False) for arg in task]
        if not istask(task):
            return task
        head, tail = task[0], task[1:]
        if not start and head in (list, reify):
            task = task[1]
            return lazify_task(*tail, start=False)
        else:
            return (head,) + tuple(lazify_task(arg, False) for arg in tail)

Review comment (performance):
This optimization is a little odd. We have to ensure that there is at least one reify at the very end since many bag operations are using the `_MapChunk` iterator such that inside of a task it is again able to iterate over chunks. If we don't have a reify at the end, the computation is just returning this iterator.

Now that I've understood this, the better approach would likely be to **add** a reify in this optimizer if we encounter such an iterator instead of removing all reifies. I will not do this in this PR but it may be something for later

Code after:

    Examples
    --------
    >>> def inc(x):
    ...     return x + 1
    >>> task = (sum, (list, (map, inc, [1, 2, 3])))
    >>> lazify_task(task)  # doctest: +ELLIPSIS
    (<built-in function sum>, (<class 'map'>, <function inc at ...>, [1, 2, 3]))
    """
    if isinstance(task, GraphNode):
        if isinstance(task, List) and len(task.args) < 50:
            return List(*[lazify_task(arg, False) for arg in task.args])
        if not isinstance(task, Task):
            return task
        if not start and task.func in (list, reify) and isinstance(task.args[0], Task):
            assert len(task.args) == 1
            task = task.args[0]
        if task.func is _execute_subgraph:
            subgraph = task.args[0]
            outkey = task.args[1]
            # If there is a reify at the output of the subgraph we don't want to act
            final_task = lazify_task(subgraph[outkey], True)
            subgraph = {
                k: lazify_task(v, False) for k, v in subgraph.items() if k != outkey
            }
            subgraph[outkey] = final_task
            return Task(
                task.key,
                _execute_subgraph,
                subgraph,
                outkey,
                *task.args[2:],
                **task.kwargs,
            )
        return Task(
            task.key,
            task.func,
            *[lazify_task(arg, False) for arg in task.args],
            **task.kwargs,
        )
    else:
        if type(task) is list and len(task) < 50:
            return [lazify_task(arg, False) for arg in task]
        if not istask(task):
            return task
        head, tail = task[0], task[1:]
        if not start and head in (list, reify):
            task = task[1]
            return lazify_task(*tail, start=False)
        else:
            return (head,) + tuple(lazify_task(arg, False) for arg in tail)