<commit_before>import io
import itertools
import math
import operator
import uuid
import warnings
from collections import defaultdict
from collections.abc import Iterable, Iterator
from functools import partial, reduce, wraps
from random import Random
from urllib.request import urlopen

import tlz as toolz
from fsspec.core import open_files
from tlz import (
    accumulate,
    compose,
    count,
    curry,
    first,
    frequencies,
    groupby,
    join,
    merge,
    merge_with,
    partition_all,
    peek,
    pluck,
    reduceby,
    remove,
    second,
    take,
    topk,
    unique,
    valmap,
)

from .. import config
from ..base import DaskMethodsMixin, dont_optimize, replace_name_in_key, tokenize
from ..context import globalmethod
from ..core import flatten, get_dependencies, istask, quote, reverse_dict
from ..delayed import Delayed, unpack_collections
from ..highlevelgraph import HighLevelGraph
from ..multiprocessing import get as mpget
from ..optimization import cull, fuse, inline
from ..sizeof import sizeof
from ..utils import (
    apply,
    digit,
    ensure_bytes,
    ensure_dict,
    ensure_unicode,
    funcname,
    insert,
    iter_chunks,
    key_split,
    parse_bytes,
    system_encoding,
    takes_multiple_arguments,
)
from . import chunk
from .avro import to_avro

no_default = "__no__default__"
no_result = type(
    "no_result", (object,), {"__slots__": (), "__reduce__": lambda self: "no_result"}
)


def lazify_task(task, start=True):
    """
    Given a task, remove unnecessary calls to ``list`` and ``reify``.

    This traverses tasks and small lists.  We choose not to traverse down lists
    of size >= 50 because it is unlikely that sequences this long contain other
    sequences in practice.

    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 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])


def lazify(dsk):
    """
    Remove unnecessary calls to ``list`` in tasks.

    See Also
    --------
    dask.bag.core.lazify_task
    """
    return valmap(lazify_task, dsk)


def inline_singleton_lists(dsk, keys, dependencies=None):
    """Inline lists that are only used once.

    >>> d = {'b': (list, 'a'),
    ...      'c': (sum, 'b', 1)}
    >>> inline_singleton_lists(d, 'c')
    {'c': (<built-in function sum>, (<class 'list'>, 'a'), 1)}

    Pairs nicely with lazify afterwards.
    """
    if dependencies is None:
        dependencies = {k: get_dependencies(dsk, task=v) for k, v in dsk.items()}
    dependents = reverse_dict(dependencies)

    inline_keys = {
        k
        for k, v in dsk.items()
        if istask(v) and v and v[0] is list and len(dependents[k]) == 1
    }
    inline_keys.difference_update(flatten(keys))
    dsk = inline(dsk, inline_keys, inline_constants=False)
    for k in inline_keys:
        del dsk[k]
    return dsk


def optimize(dsk, keys, fuse_keys=None, rename_fused_keys=None, **kwargs):
    """Optimize a dask from a dask Bag."""
    dsk = ensure_dict(dsk)
    dsk2, dependencies = cull(dsk, keys)
    kwargs = {}
    if rename_fused_keys is not None:
        kwargs["rename_keys"] = rename_fused_keys
    dsk3, dependencies = fuse(dsk2, keys + (fuse_keys or []), dependencies, **kwargs)
    dsk4 = inline_singleton_lists(dsk3, keys, dependencies)
    dsk5 = lazify(dsk4)
    return dsk5


def _to_textfiles_chunk(data, lazy_file, last_endline):
    with lazy_file as f:
        if isinstance(f, io.TextIOWrapper):
            endline = "\n"
            ensure = ensure_unicode
        else:
            endline = b"\n"
            ensure = ensure_bytes
        started = False
        for d in data:
            if started:
                f.write(endline)
            else:
                started = True
            f.write(ensure(d))
        if last_endline:
            f.write(endline)


def to_textfiles(
    b,
    path,
    name_function=None,
    compression="infer",
    encoding=system_encoding,
    compute=True,
    storage_options=None,
    last_endline=False,
    **kwargs
):
    """Write dask Bag to disk, one filename per partition, one line per element.

    **Paths**: This will create one file for each partition in your bag. You
    can specify the filenames in a variety of ways.

    Use a globstring

    >>> b.to_textfiles('/path/to/data/*.json.gz')  # doctest: +SKIP

    The * will be replaced by the increasing sequence 1, 2, ...

    ::

        /path/to/data/0.json.gz
        /path/to/data/1.json.gz

    Use a globstring and a ``name_function=`` keyword argument.  The
    name_function function should expect an integer and produce a string.
    Strings produced by name_function must preserve the order of their
    respective partition indices.

    >>> from datetime import date, timedelta
    >>> def name(i):
    ...     return str(date(2015, 1, 1) + i * timedelta(days=1))

    >>> name(0)
    '2015-01-01'
    >>> name(15)
    '2015-01-16'

    >>> b.to_textfiles('/path/to/data/*.json.gz', name_function=name)  # doctest: +SKIP

    ::

        /path/to/data/2015-01-01.json.gz
        /path/to/data/2015-01-02.json.gz
        ...

    You can also provide an explicit list of paths.

    >>> paths = ['/path/to/data/alice.json.gz', '/path/to/data/bob.json.gz', ...]  # doctest: +SKIP
    >>> b.to_textfiles(paths) # doctest: +SKIP

    **Compression**: Filenames with extensions corresponding to known
    compression algorithms (gz, bz2) will be compressed accordingly.

    **Bag Contents**: The bag calling ``to_textfiles`` must be a bag of
    text strings. For example, a bag of dictionaries could be written to
    JSON text files by mapping ``json.dumps`` on to the bag first, and
    then calling ``to_textfiles`` :

    >>> b_dict.map(json.dumps).to_textfiles("/path/to/data/*.json")  # doctest: +SKIP

    **Last endline**: By default the last line does not end with a newline
    character. Pass ``last_endline=True`` to invert the default.
    """
    mode = "wb" if encoding is None else "wt"
    files = open_files(
        path,
        compression=compression,
        mode=mode,
        encoding=encoding,
        name_function=name_function,
        num=b.npartitions,
        **(storage_options or {})
    )

    name = "to-textfiles-" + uuid.uuid4().hex
    dsk = {
        (name, i): (_to_textfiles_chunk, (b.name, i), f, last_endline)
        for i, f in enumerate(files)
    }
    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[b])
    out = type(b)(graph, name, b.npartitions)

    if compute:
        out.compute(**kwargs)
        return [f.path for f in files]
    else:
        return out.to_delayed()


def finalize(results):
    if not results:
        return results
    if isinstance(results, Iterator):
        results = list(results)
    if isinstance(results[0], Iterable) and not isinstance(results[0], str):
        results = toolz.concat(results)
    if isinstance(results, Iterator):
        results = list(results)
    return results


def finalize_item(results):
    return results[0]


class StringAccessor:
    """String processing functions

    Examples
    --------

    >>> import dask.bag as db
    >>> b = db.from_sequence(['<NAME>', '<NAME>', '<NAME>'])
    >>> list(b.str.lower())
    ['<NAME>', '<NAME>', '<NAME>']

    >>> list(b.str.match('*Smith'))
    ['<NAME>', '<NAME>']

    >>> list(b.str.split(' '))
    [['Alice', 'Smith'], ['Bob', 'Jones'], ['Charlie', 'Smith']]
    """

    def __init__(self, bag):
        self._bag = bag

    def __dir__(self):
        return sorted(set(dir(type(self)) + dir(str)))

    def _strmap(self, key, *args, **kwargs):
        return self._bag.map(operator.methodcaller(key, *args, **kwargs))

    def __getattr__(self, key):
        try:
            return object.__getattribute__(self, key)
        except AttributeError:
            if key in dir(str):
                func = getattr(str, key)
                return robust_wraps(func)(partial(self._strmap, key))
            else:
                raise

    def match(self, pattern):
        """Filter strings by those that match a pattern.

        Examples
        --------

        >>> import dask.bag as db
        >>> b = db.from_sequence(['<NAME>', '<NAME>', '<NAME>'])
        >>> list(b.str.match('*Smith'))
        ['<NAME>', '<NAME>']

        See Also
        --------
        fnmatch.fnmatch
        """
        from fnmatch import fnmatch

        return self._bag.filter(partial(fnmatch, pat=pattern))


def robust_wraps(wrapper):
    """A weak version of wraps that only copies doc."""

    def _(wrapped):
        wrapped.__doc__ = wrapper.__doc__
        return wrapped

    return _


class Item(DaskMethodsMixin):
    def __init__(self, dsk, key):
        self.dask = dsk
        self.key = key
        self.name = key

    def __dask_graph__(self):
        return self.dask

    def __dask_keys__(self):
        return [self.key]

    def __dask_layers__(self):
        # Deal with instances created with Item.from_delayed()
        # e.g.: Item.from_delayed(da.ones(1).to_delayed()[0])
        # See Delayed.__dask_layers__
        if isinstance(self.dask, HighLevelGraph) and len(self.dask.layers) == 1:
            return tuple(self.dask.layers)
        else:
            return (self.key,)

    def __dask_tokenize__(self):
        return self.key

    __dask_optimize__ = globalmethod(optimize, key="bag_optimize", falsey=dont_optimize)
    __dask_scheduler__ = staticmethod(mpget)

    def __dask_postcompute__(self):
        return finalize_item, ()

    def __dask_postpersist__(self):
        return self._rebuild, ()

    def _rebuild(self, dsk, *, rename=None):
        key = replace_name_in_key(self.key, rename) if rename else self.key
        return Item(dsk, key)

    @staticmethod
    def from_delayed(value):
        """Create bag item from a dask.delayed value.

        See ``dask.bag.from_delayed`` for details
        """
        from dask.delayed import Delayed, delayed

        if not isinstance(value, Delayed) and hasattr(value, "key"):
            value = delayed(value)
        assert isinstance(value, Delayed)
        return Item(value.dask, value.key)

    @property
    def _args(self):
        return (self.dask, self.key)

    def __getstate__(self):
        return self._args

    def __setstate__(self, state):
        self.dask, self.key = state

    def apply(self, func):
        name = "{0}-{1}".format(funcname(func), tokenize(self, func, "apply"))
        dsk = {name: (func, self.key)}
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return Item(graph, name)

    __int__ = __float__ = __complex__ = __bool__ = DaskMethodsMixin.compute

    def to_delayed(self, optimize_graph=True):
        """Convert into a ``dask.delayed`` object.

        Parameters
        ----------
        optimize_graph : bool, optional
            If True [default], the graph is optimized before converting into
            ``dask.delayed`` objects.
        """
        from dask.delayed import Delayed

        dsk = self.__dask_graph__()
        if optimize_graph:
            dsk = self.__dask_optimize__(dsk, self.__dask_keys__())
        return Delayed(self.key, dsk)


class Bag(DaskMethodsMixin):
    """Parallel collection of Python objects

    Examples
    --------
    Create Bag from sequence

    >>> import dask.bag as db
    >>> b = db.from_sequence(range(5))
    >>> list(b.filter(lambda x: x % 2 == 0).map(lambda x: x * 10))
    [0, 20, 40]

    Create Bag from filename or globstring of filenames

    >>> b = db.read_text('/path/to/mydata.*.json.gz').map(json.loads)  # doctest: +SKIP

    Create manually (expert use)

    >>> dsk = {('x', 0): (range, 5),
    ...        ('x', 1): (range, 5),
    ...        ('x', 2): (range, 5)}
    >>> b = db.Bag(dsk, 'x', npartitions=3)

    >>> sorted(b.map(lambda x: x * 10))
    [0, 0, 0, 10, 10, 10, 20, 20, 20, 30, 30, 30, 40, 40, 40]

    >>> int(b.fold(lambda x, y: x + y))
    30
    """

    def __init__(self, dsk, name, npartitions):
        if not isinstance(dsk, HighLevelGraph):
            dsk = HighLevelGraph.from_collections(name, dsk, dependencies=[])
        self.dask = dsk
        self.name = name
        self.npartitions = npartitions

    def __dask_graph__(self):
        return self.dask

    def __dask_keys__(self):
        return [(self.name, i) for i in range(self.npartitions)]

    def __dask_layers__(self):
        return (self.name,)

    def __dask_tokenize__(self):
        return self.name

    __dask_optimize__ = globalmethod(optimize, key="bag_optimize", falsey=dont_optimize)
    __dask_scheduler__ = staticmethod(mpget)

    def __dask_postcompute__(self):
        return finalize, ()

    def __dask_postpersist__(self):
        return self._rebuild, ()

    def _rebuild(self, dsk, *, rename=None):
        name = self.name
        if rename:
            name = rename.get(name, name)
        return type(self)(dsk, name, self.npartitions)

    def __str__(self):
        return "dask.bag<%s, npartitions=%d>" % (key_split(self.name), self.npartitions)

    __repr__ = __str__

    str = property(fget=StringAccessor)

    def map(self, func, *args, **kwargs):
        """Apply a function elementwise across one or more bags.

        Note that all ``Bag`` arguments must be partitioned identically.

        Parameters
        ----------
        func : callable
        *args, **kwargs : Bag, Item, or object
            Extra arguments and keyword arguments to pass to ``func`` *after*
            the calling bag instance. Non-Bag args/kwargs are broadcasted
            across all calls to ``func``.

        Notes
        -----
        For calls with multiple `Bag` arguments, corresponding partitions
        should have the same length; if they do not, the call will error at
        compute time.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5), npartitions=2)
        >>> b2 = db.from_sequence(range(5, 10), npartitions=2)

        Apply a function to all elements in a bag:

        >>> b.map(lambda x: x + 1).compute()
        [1, 2, 3, 4, 5]

        Apply a function with arguments from multiple bags:

        >>> from operator import add
        >>> b.map(add, b2).compute()
        [5, 7, 9, 11, 13]

        Non-bag arguments are broadcast across all calls to the mapped
        function:

        >>> b.map(add, 1).compute()
        [1, 2, 3, 4, 5]

        Keyword arguments are also supported, and have the same semantics as
        regular arguments:

        >>> def myadd(x, y=0):
        ...     return x + y
        >>> b.map(myadd, y=b2).compute()
        [5, 7, 9, 11, 13]
        >>> b.map(myadd, y=1).compute()
        [1, 2, 3, 4, 5]

        Both arguments and keyword arguments can also be instances of
        ``dask.bag.Item``. Here we'll add the max value in the bag to each
        element:

        >>> b.map(myadd, b.max()).compute()
        [4, 5, 6, 7, 8]
        """
        return bag_map(func, self, *args, **kwargs)

    def starmap(self, func, **kwargs):
        """Apply a function using argument tuples from the given bag.

        This is similar to ``itertools.starmap``, except it also accepts
        keyword arguments. In pseudocode, this is could be written as:

        >>> def starmap(func, bag, **kwargs):
        ...     return (func(*args, **kwargs) for args in bag)

        Parameters
        ----------
        func : callable
        **kwargs : Item, Delayed, or object, optional
            Extra keyword arguments to pass to ``func``. These can either be
            normal objects, ``dask.bag.Item``, or ``dask.delayed.Delayed``.

        Examples
        --------
        >>> import dask.bag as db
        >>> data = [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]
        >>> b = db.from_sequence(data, npartitions=2)

        Apply a function to each argument tuple:

        >>> from operator import add
        >>> b.starmap(add).compute()
        [3, 7, 11, 15, 19]

        Apply a function to each argument tuple, with additional keyword
        arguments:

        >>> def myadd(x, y, z=0):
        ...     return x + y + z
        >>> b.starmap(myadd, z=10).compute()
        [13, 17, 21, 25, 29]

        Keyword arguments can also be instances of ``dask.bag.Item`` or
        ``dask.delayed.Delayed``:

        >>> max_second = b.pluck(1).max()
        >>> max_second.compute()
        10
        >>> b.starmap(myadd, z=max_second).compute()
        [13, 17, 21, 25, 29]
        """
        name = "{0}-{1}".format(
            funcname(func), tokenize(self, func, "starmap", **kwargs)
        )
        dependencies = [self]
        if kwargs:
            kwargs, collections = unpack_scalar_dask_kwargs(kwargs)
            dependencies.extend(collections)

        dsk = {
            (name, i): (reify, (starmap_chunk, func, (self.name, i), kwargs))
            for i in range(self.npartitions)
        }
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=dependencies)
        return type(self)(graph, name, self.npartitions)

    @property
    def _args(self):
        return (self.dask, self.name, self.npartitions)

    def __getstate__(self):
        return self._args

    def __setstate__(self, state):
        self.dask, self.name, self.npartitions = state

    def filter(self, predicate):
        """Filter elements in collection by a predicate function.

        >>> def iseven(x):
        ...     return x % 2 == 0

        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5))
        >>> list(b.filter(iseven))
        [0, 2, 4]
        """
        name = "filter-{0}-{1}".format(funcname(predicate), tokenize(self, predicate))
        dsk = dict(
            ((name, i), (reify, (filter, predicate, (self.name, i))))
            for i in range(self.npartitions)
        )
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return type(self)(graph, name, self.npartitions)

    def random_sample(self, prob, random_state=None):
        """Return elements from bag with probability of ``prob``.

        Parameters
        ----------
        prob : float
            A float between 0 and 1, representing the probability that each
            element will be returned.
        random_state : int or random.Random, optional
            If an integer, will be used to seed a new ``random.Random`` object.
            If provided, results in deterministic sampling.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5))
        >>> list(b.random_sample(0.5, 43))
        [0, 3, 4]
        >>> list(b.random_sample(0.5, 43))
        [0, 3, 4]
        """
        if not 0 <= prob <= 1:
            raise ValueError("prob must be a number in the interval [0, 1]")
        if not isinstance(random_state, Random):
            random_state = Random(random_state)

        name = "random-sample-%s" % tokenize(self, prob, random_state.getstate())
        state_data = random_state_data_python(self.npartitions, random_state)
        dsk = {
            (name, i): (reify, (random_sample, (self.name, i), state, prob))
            for i, state in zip(range(self.npartitions), state_data)
        }
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return type(self)(graph, name, self.npartitions)

    def remove(self, predicate):
        """Remove elements in collection that match predicate.

        >>> def iseven(x):
        ...     return x % 2 == 0

        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5))
        >>> list(b.remove(iseven))
        [1, 3]
        """
        name = "remove-{0}-{1}".format(funcname(predicate), tokenize(self, predicate))
        dsk = dict(
            ((name, i), (reify, (remove, predicate, (self.name, i))))
            for i in range(self.npartitions)
        )
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return type(self)(graph, name, self.npartitions)

    def map_partitions(self, func, *args, **kwargs):
        """Apply a function to every partition across one or more bags.

        Note that all ``Bag`` arguments must be partitioned identically.

        Parameters
        ----------
        func : callable
            The function to be called on every partition.
            This function should expect an ``Iterator`` or ``Iterable`` for
            every partition and should return an ``Iterator`` or ``Iterable``
            in return.
        *args, **kwargs : Bag, Item, Delayed, or object
            Arguments and keyword arguments to pass to ``func``.
            Partitions from this bag will be the first argument, and these will
            be passed *after*.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(range(1, 101), npartitions=10)
        >>> def div(nums, den=1):
        ...     return [num / den for num in nums]

        Using a python object:

        >>> hi = b.max().compute()
        >>> hi
        100
        >>> b.map_partitions(div, den=hi).take(5)
        (0.01, 0.02, 0.03, 0.04, 0.05)

        Using an ``Item``:

        >>> b.map_partitions(div, den=b.max()).take(5)
        (0.01, 0.02, 0.03, 0.04, 0.05)

        Note that while both versions give the same output, the second forms a
        single graph, and then computes everything at once, and in some cases
        may be more efficient.
        """
        return map_partitions(func, self, *args, **kwargs)

    def pluck(self, key, default=no_default):
        """Select item from all tuples/dicts in collection.

        >>> import dask.bag as db
        >>> b = db.from_sequence([{'name': 'Alice', 'credits': [1, 2, 3]},
        ...                       {'name': 'Bob',   'credits': [10, 20]}])
        >>> list(b.pluck('name'))
        ['Alice', 'Bob']
        >>> list(b.pluck('credits').pluck(0))
        [1, 10]
        """
        name = "pluck-" + tokenize(self, key, default)
        key = quote(key)
        if default == no_default:
            dsk = dict(
                ((name, i), (list, (pluck, key, (self.name, i))))
                for i in range(self.npartitions)
            )
        else:
            dsk = dict(
                ((name, i), (list, (pluck, key, (self.name, i), default)))
                for i in range(self.npartitions)
            )
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return type(self)(graph, name, self.npartitions)

    def unzip(self, n):
        """Transform a bag of tuples to ``n`` bags of their elements.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence([(i, i + 1, i + 2) for i in range(10)])
        >>> first, second, third = b.unzip(3)
        >>> isinstance(first, db.Bag)
        True
        >>> first.compute()
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

        Note that this is equivalent to:

        >>> first, second, third = (b.pluck(i) for i in range(3))
        """
        return tuple(self.pluck(i) for i in range(n))

    @wraps(to_textfiles)
    def to_textfiles(
        self,
        path,
        name_function=None,
        compression="infer",
        encoding=system_encoding,
        compute=True,
        storage_options=None,
        last_endline=False,
        **kwargs
    ):
        return to_textfiles(
            self,
            path,
            name_function,
            compression,
            encoding,
            compute,
            storage_options=storage_options,
            last_endline=last_endline,
            **kwargs
        )

    @wraps(to_avro)
    def to_avro(
        self,
        filename,
        schema,
        name_function=None,
        storage_options=None,
        codec="null",
        sync_interval=16000,
        metadata=None,
        compute=True,
        **kwargs
    ):
        return to_avro(
            self,
            filename,
            schema,
            name_function,
            storage_options,
            codec,
            sync_interval,
            metadata,
            compute,
            **kwargs
        )

    def fold(
        self, binop, combine=None, initial=no_default, split_every=None, out_type=Item
    ):
        """Parallelizable reduction

        Fold is like the builtin function ``reduce`` except that it works in
        parallel.  Fold takes two binary operator functions, one to reduce each
        partition of our dataset and another to combine results between
        partitions

        1.  ``binop``: Binary operator to reduce within each partition
        2.  ``combine``:  Binary operator to combine results from binop

        Sequentially this would look like the following:

        >>> intermediates = [reduce(binop, part) for part in partitions]  # doctest: +SKIP
        >>> final = reduce(combine, intermediates)  # doctest: +SKIP

        If only one function is given then it is used for both functions
        ``binop`` and ``combine`` as in the following example to compute the
        sum:

        >>> def add(x, y):
        ...     return x + y

        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5))
        >>> b.fold(add).compute()
        10

        In full form we provide both binary operators as well as their default
        arguments

        >>> b.fold(binop=add, combine=add, initial=0).compute()
        10

        More complex binary operators are also doable

        >>> def add_to_set(acc, x):
        ...     ''' Add new element x to set acc '''
        ...     return acc | set([x])
        >>> b.fold(add_to_set, set.union, initial=set()).compute()
        {0, 1, 2, 3, 4}

        See Also
        --------

        Bag.foldby
        """
        combine = combine or binop
        if initial is not no_default:
            return self.reduction(
                curry(_reduce, binop, initial=initial),
                curry(_reduce, combine),
                split_every=split_every,
                out_type=out_type,
            )
        else:
            from tlz.curried import reduce

            return self.reduction(
                reduce(binop),
                reduce(combine),
                split_every=split_every,
                out_type=out_type,
            )

    def frequencies(self, split_every=None, sort=False):
        """Count number of occurrences of each distinct element.

        >>> import dask.bag as db
        >>> b = db.from_sequence(['Alice', 'Bob', 'Alice'])
        >>> dict(b.frequencies())       # doctest: +SKIP
        {'Alice': 2, 'Bob', 1}
        """
        result = self.reduction(
            frequencies,
            merge_frequencies,
            out_type=Bag,
            split_every=split_every,
            name="frequencies",
        ).map_partitions(dictitems)
        if sort:
            result = result.map_partitions(sorted, key=second, reverse=True)
        return result

    def topk(self, k, key=None, split_every=None):
        """K largest elements in collection

        Optionally ordered by some key function

        >>> import dask.bag as db
        >>> b = db.from_sequence([10, 3, 5, 7, 11, 4])
        >>> list(b.topk(2))
        [11, 10]

        >>> list(b.topk(2, lambda x: -x))
        [3, 4]
        """
        if key:
            if callable(key) and takes_multiple_arguments(key):
                key = partial(apply, key)
            func = partial(topk, k, key=key)
        else:
            func = partial(topk, k)
        return self.reduction(
            func,
            compose(func, toolz.concat),
            out_type=Bag,
            split_every=split_every,
            name="topk",
        )

    def distinct(self, key=None):
        """Distinct elements of collection

        Unordered without repeats.

        Parameters
        ----------
        key: {callable,str}
            Defines uniqueness of items in bag by calling ``key`` on each item.
            If a string is passed ``key`` is considered to be ``lambda x: x[key]``.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(['Alice', 'Bob', 'Alice'])
        >>> sorted(b.distinct())
        ['Alice', 'Bob']
        >>> b = db.from_sequence([{'name': 'Alice'}, {'name': 'Bob'}, {'name': 'Alice'}])
        >>> b.distinct(key=lambda x: x['name']).compute()
        [{'name': 'Alice'}, {'name': 'Bob'}]
        >>> b.distinct(key='name').compute()
        [{'name': 'Alice'}, {'name': 'Bob'}]
        """
        func = chunk_distinct if key is None else partial(chunk_distinct, key=key)
        agg = merge_distinct if key is None else partial(merge_distinct, key=key)
        return self.reduction(func, agg, out_type=Bag, name="distinct")

    def reduction(
        self, perpartition, aggregate, split_every=None, out_type=Item, name=None
    ):
        """Reduce collection with reduction operators.

        Parameters
        ----------
        perpartition: function
            reduction to apply to each partition
        aggregate: function
            reduction to apply to the results of all partitions
        split_every: int (optional)
            Group partitions into groups of this size while performing reduction
            Defaults to 8
        out_type: {Bag, Item}
            The out type of the result, Item if a single element, Bag if a list
            of elements.  Defaults to Item.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(range(10))
        >>> b.reduction(sum, sum).compute()
        45
        """
        if split_every is None:
            split_every = 8
        if split_every is False:
            split_every = self.npartitions

        token = tokenize(self, perpartition, aggregate, split_every)
        a = "%s-part-%s" % (name or funcname(perpartition), token)
        is_last = self.npartitions == 1
        dsk = {
            (a, i): (empty_safe_apply, perpartition, (self.name, i), is_last)
            for i in range(self.npartitions)
        }
        k = self.npartitions
        b = a
        fmt = "%s-aggregate-%s" % (name or funcname(aggregate), token)
        depth = 0

        while k > split_every:
            c = fmt + str(depth)
            for i, inds in enumerate(partition_all(split_every, range(k))):
                dsk[(c, i)] = (
                    empty_safe_aggregate,
                    aggregate,
                    [(b, j) for j in inds],
                    False,
                )

            k = i + 1
            b = c
            depth += 1

        dsk[(fmt, 0)] = (
            empty_safe_aggregate,
            aggregate,
            [(b, j) for j in range(k)],
            True,
        )

        graph = HighLevelGraph.from_collections(fmt, dsk, dependencies=[self])
        if out_type is Item:
            dsk[fmt] = dsk.pop((fmt, 0))
            return Item(graph, fmt)
        else:
            return Bag(graph, fmt, 1)

    def sum(self, split_every=None):
        """Sum all elements"""
        return self.reduction(sum, sum, split_every=split_every)

    def max(self, split_every=None):
        """Maximum element"""
        return self.reduction(max, max, split_every=split_every)

    def min(self, split_every=None):
        """Minimum element"""
        return self.reduction(min, min, split_every=split_every)

    def any(self, split_every=None):
        """Are any of the elements truthy?"""
        return self.reduction(any, any, split_every=split_every)

    def all(self, split_every=None):
        """Are all elements truthy?"""
        return self.reduction(all, all, split_every=split_every)

    def count(self, split_every=None):
        """Count the number of elements."""
        return self.reduction(count, sum, split_every=split_every)

    def mean(self):
        """Arithmetic mean"""<commit_msg>Add some basic examples to doc strings on `dask.bag` `all`, `any`, `count` methods (#7630)<commit_after>import io
import itertools
import math
import operator
import uuid
import warnings
from collections import defaultdict
from collections.abc import Iterable, Iterator
from functools import partial, reduce, wraps
from random import Random
from urllib.request import urlopen

import tlz as toolz
from fsspec.core import open_files
from tlz import (
    accumulate,
    compose,
    count,
    curry,
    first,
    frequencies,
    groupby,
    join,
    merge,
    merge_with,
    partition_all,
    peek,
    pluck,
    reduceby,
    remove,
    second,
    take,
    topk,
    unique,
    valmap,
)

from .. import config
from ..base import DaskMethodsMixin, dont_optimize, replace_name_in_key, tokenize
from ..context import globalmethod
from ..core import flatten, get_dependencies, istask, quote, reverse_dict
from ..delayed import Delayed, unpack_collections
from ..highlevelgraph import HighLevelGraph
from ..multiprocessing import get as mpget
from ..optimization import cull, fuse, inline
from ..sizeof import sizeof
from ..utils import (
    apply,
    digit,
    ensure_bytes,
    ensure_dict,
    ensure_unicode,
    funcname,
    insert,
    iter_chunks,
    key_split,
    parse_bytes,
    system_encoding,
    takes_multiple_arguments,
)
from . import chunk
from .avro import to_avro

no_default = "__no__default__"
no_result = type(
    "no_result", (object,), {"__slots__": (), "__reduce__": lambda self: "no_result"}
)


def lazify_task(task, start=True):
    """
    Given a task, remove unnecessary calls to ``list`` and ``reify``.

    This traverses tasks and small lists.  We choose not to traverse down lists
    of size >= 50 because it is unlikely that sequences this long contain other
    sequences in practice.

    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 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])


def lazify(dsk):
    """
    Remove unnecessary calls to ``list`` in tasks.

    See Also
    --------
    dask.bag.core.lazify_task
    """
    return valmap(lazify_task, dsk)


def inline_singleton_lists(dsk, keys, dependencies=None):
    """Inline lists that are only used once.

    >>> d = {'b': (list, 'a'),
    ...      'c': (sum, 'b', 1)}
    >>> inline_singleton_lists(d, 'c')
    {'c': (<built-in function sum>, (<class 'list'>, 'a'), 1)}

    Pairs nicely with lazify afterwards.
    """
    if dependencies is None:
        dependencies = {k: get_dependencies(dsk, task=v) for k, v in dsk.items()}
    dependents = reverse_dict(dependencies)

    inline_keys = {
        k
        for k, v in dsk.items()
        if istask(v) and v and v[0] is list and len(dependents[k]) == 1
    }
    inline_keys.difference_update(flatten(keys))
    dsk = inline(dsk, inline_keys, inline_constants=False)
    for k in inline_keys:
        del dsk[k]
    return dsk


def optimize(dsk, keys, fuse_keys=None, rename_fused_keys=None, **kwargs):
    """Optimize a dask from a dask Bag."""
    dsk = ensure_dict(dsk)
    dsk2, dependencies = cull(dsk, keys)
    kwargs = {}
    if rename_fused_keys is not None:
        kwargs["rename_keys"] = rename_fused_keys
    dsk3, dependencies = fuse(dsk2, keys + (fuse_keys or []), dependencies, **kwargs)
    dsk4 = inline_singleton_lists(dsk3, keys, dependencies)
    dsk5 = lazify(dsk4)
    return dsk5


def _to_textfiles_chunk(data, lazy_file, last_endline):
    with lazy_file as f:
        if isinstance(f, io.TextIOWrapper):
            endline = "\n"
            ensure = ensure_unicode
        else:
            endline = b"\n"
            ensure = ensure_bytes
        started = False
        for d in data:
            if started:
                f.write(endline)
            else:
                started = True
            f.write(ensure(d))
        if last_endline:
            f.write(endline)


def to_textfiles(
    b,
    path,
    name_function=None,
    compression="infer",
    encoding=system_encoding,
    compute=True,
    storage_options=None,
    last_endline=False,
    **kwargs
):
    """Write dask Bag to disk, one filename per partition, one line per element.

    **Paths**: This will create one file for each partition in your bag. You
    can specify the filenames in a variety of ways.

    Use a globstring

    >>> b.to_textfiles('/path/to/data/*.json.gz')  # doctest: +SKIP

    The * will be replaced by the increasing sequence 1, 2, ...

    ::

        /path/to/data/0.json.gz
        /path/to/data/1.json.gz

    Use a globstring and a ``name_function=`` keyword argument.  The
    name_function function should expect an integer and produce a string.
    Strings produced by name_function must preserve the order of their
    respective partition indices.

    >>> from datetime import date, timedelta
    >>> def name(i):
    ...     return str(date(2015, 1, 1) + i * timedelta(days=1))

    >>> name(0)
    '2015-01-01'
    >>> name(15)
    '2015-01-16'

    >>> b.to_textfiles('/path/to/data/*.json.gz', name_function=name)  # doctest: +SKIP

    ::

        /path/to/data/2015-01-01.json.gz
        /path/to/data/2015-01-02.json.gz
        ...

    You can also provide an explicit list of paths.

    >>> paths = ['/path/to/data/alice.json.gz', '/path/to/data/bob.json.gz', ...]  # doctest: +SKIP
    >>> b.to_textfiles(paths) # doctest: +SKIP

    **Compression**: Filenames with extensions corresponding to known
    compression algorithms (gz, bz2) will be compressed accordingly.

    **Bag Contents**: The bag calling ``to_textfiles`` must be a bag of
    text strings. For example, a bag of dictionaries could be written to
    JSON text files by mapping ``json.dumps`` on to the bag first, and
    then calling ``to_textfiles`` :

    >>> b_dict.map(json.dumps).to_textfiles("/path/to/data/*.json")  # doctest: +SKIP

    **Last endline**: By default the last line does not end with a newline
    character. Pass ``last_endline=True`` to invert the default.
    """
    mode = "wb" if encoding is None else "wt"
    files = open_files(
        path,
        compression=compression,
        mode=mode,
        encoding=encoding,
        name_function=name_function,
        num=b.npartitions,
        **(storage_options or {})
    )

    name = "to-textfiles-" + uuid.uuid4().hex
    dsk = {
        (name, i): (_to_textfiles_chunk, (b.name, i), f, last_endline)
        for i, f in enumerate(files)
    }
    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[b])
    out = type(b)(graph, name, b.npartitions)

    if compute:
        out.compute(**kwargs)
        return [f.path for f in files]
    else:
        return out.to_delayed()


def finalize(results):
    if not results:
        return results
    if isinstance(results, Iterator):
        results = list(results)
    if isinstance(results[0], Iterable) and not isinstance(results[0], str):
        results = toolz.concat(results)
    if isinstance(results, Iterator):
        results = list(results)
    return results


def finalize_item(results):
    return results[0]


class StringAccessor:
    """String processing functions

    Examples
    --------

    >>> import dask.bag as db
    >>> b = db.from_sequence(['<NAME>', '<NAME>', '<NAME>'])
    >>> list(b.str.lower())
    ['<NAME>', '<NAME>', '<NAME>']

    >>> list(b.str.match('*Smith'))
    ['<NAME>', '<NAME>']

    >>> list(b.str.split(' '))
    [['Alice', 'Smith'], ['Bob', 'Jones'], ['Charlie', 'Smith']]
    """

    def __init__(self, bag):
        self._bag = bag

    def __dir__(self):
        return sorted(set(dir(type(self)) + dir(str)))

    def _strmap(self, key, *args, **kwargs):
        return self._bag.map(operator.methodcaller(key, *args, **kwargs))

    def __getattr__(self, key):
        try:
            return object.__getattribute__(self, key)
        except AttributeError:
            if key in dir(str):
                func = getattr(str, key)
                return robust_wraps(func)(partial(self._strmap, key))
            else:
                raise

    def match(self, pattern):
        """Filter strings by those that match a pattern.

        Examples
        --------

        >>> import dask.bag as db
        >>> b = db.from_sequence(['<NAME>', '<NAME>', '<NAME>'])
        >>> list(b.str.match('*Smith'))
        ['<NAME>', '<NAME>']

        See Also
        --------
        fnmatch.fnmatch
        """
        from fnmatch import fnmatch

        return self._bag.filter(partial(fnmatch, pat=pattern))


def robust_wraps(wrapper):
    """A weak version of wraps that only copies doc."""

    def _(wrapped):
        wrapped.__doc__ = wrapper.__doc__
        return wrapped

    return _


class Item(DaskMethodsMixin):
    def __init__(self, dsk, key):
        self.dask = dsk
        self.key = key
        self.name = key

    def __dask_graph__(self):
        return self.dask

    def __dask_keys__(self):
        return [self.key]

    def __dask_layers__(self):
        # Deal with instances created with Item.from_delayed()
        # e.g.: Item.from_delayed(da.ones(1).to_delayed()[0])
        # See Delayed.__dask_layers__
        if isinstance(self.dask, HighLevelGraph) and len(self.dask.layers) == 1:
            return tuple(self.dask.layers)
        else:
            return (self.key,)

    def __dask_tokenize__(self):
        return self.key

    __dask_optimize__ = globalmethod(optimize, key="bag_optimize", falsey=dont_optimize)
    __dask_scheduler__ = staticmethod(mpget)

    def __dask_postcompute__(self):
        return finalize_item, ()

    def __dask_postpersist__(self):
        return self._rebuild, ()

    def _rebuild(self, dsk, *, rename=None):
        key = replace_name_in_key(self.key, rename) if rename else self.key
        return Item(dsk, key)

    @staticmethod
    def from_delayed(value):
        """Create bag item from a dask.delayed value.

        See ``dask.bag.from_delayed`` for details
        """
        from dask.delayed import Delayed, delayed

        if not isinstance(value, Delayed) and hasattr(value, "key"):
            value = delayed(value)
        assert isinstance(value, Delayed)
        return Item(value.dask, value.key)

    @property
    def _args(self):
        return (self.dask, self.key)

    def __getstate__(self):
        return self._args

    def __setstate__(self, state):
        self.dask, self.key = state

    def apply(self, func):
        name = "{0}-{1}".format(funcname(func), tokenize(self, func, "apply"))
        dsk = {name: (func, self.key)}
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return Item(graph, name)

    __int__ = __float__ = __complex__ = __bool__ = DaskMethodsMixin.compute

    def to_delayed(self, optimize_graph=True):
        """Convert into a ``dask.delayed`` object.

        Parameters
        ----------
        optimize_graph : bool, optional
            If True [default], the graph is optimized before converting into
            ``dask.delayed`` objects.
        """
        from dask.delayed import Delayed

        dsk = self.__dask_graph__()
        if optimize_graph:
            dsk = self.__dask_optimize__(dsk, self.__dask_keys__())
        return Delayed(self.key, dsk)


class Bag(DaskMethodsMixin):
    """Parallel collection of Python objects

    Examples
    --------
    Create Bag from sequence

    >>> import dask.bag as db
    >>> b = db.from_sequence(range(5))
    >>> list(b.filter(lambda x: x % 2 == 0).map(lambda x: x * 10))
    [0, 20, 40]

    Create Bag from filename or globstring of filenames

    >>> b = db.read_text('/path/to/mydata.*.json.gz').map(json.loads)  # doctest: +SKIP

    Create manually (expert use)

    >>> dsk = {('x', 0): (range, 5),
    ...        ('x', 1): (range, 5),
    ...        ('x', 2): (range, 5)}
    >>> b = db.Bag(dsk, 'x', npartitions=3)

    >>> sorted(b.map(lambda x: x * 10))
    [0, 0, 0, 10, 10, 10, 20, 20, 20, 30, 30, 30, 40, 40, 40]

    >>> int(b.fold(lambda x, y: x + y))
    30
    """

    def __init__(self, dsk, name, npartitions):
        if not isinstance(dsk, HighLevelGraph):
            dsk = HighLevelGraph.from_collections(name, dsk, dependencies=[])
        self.dask = dsk
        self.name = name
        self.npartitions = npartitions

    def __dask_graph__(self):
        return self.dask

    def __dask_keys__(self):
        return [(self.name, i) for i in range(self.npartitions)]

    def __dask_layers__(self):
        return (self.name,)

    def __dask_tokenize__(self):
        return self.name

    __dask_optimize__ = globalmethod(optimize, key="bag_optimize", falsey=dont_optimize)
    __dask_scheduler__ = staticmethod(mpget)

    def __dask_postcompute__(self):
        return finalize, ()

    def __dask_postpersist__(self):
        return self._rebuild, ()

    def _rebuild(self, dsk, *, rename=None):
        name = self.name
        if rename:
            name = rename.get(name, name)
        return type(self)(dsk, name, self.npartitions)

    def __str__(self):
        return "dask.bag<%s, npartitions=%d>" % (key_split(self.name), self.npartitions)

    __repr__ = __str__

    str = property(fget=StringAccessor)

    def map(self, func, *args, **kwargs):
        """Apply a function elementwise across one or more bags.

        Note that all ``Bag`` arguments must be partitioned identically.

        Parameters
        ----------
        func : callable
        *args, **kwargs : Bag, Item, or object
            Extra arguments and keyword arguments to pass to ``func`` *after*
            the calling bag instance. Non-Bag args/kwargs are broadcasted
            across all calls to ``func``.

        Notes
        -----
        For calls with multiple `Bag` arguments, corresponding partitions
        should have the same length; if they do not, the call will error at
        compute time.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5), npartitions=2)
        >>> b2 = db.from_sequence(range(5, 10), npartitions=2)

        Apply a function to all elements in a bag:

        >>> b.map(lambda x: x + 1).compute()
        [1, 2, 3, 4, 5]

        Apply a function with arguments from multiple bags:

        >>> from operator import add
        >>> b.map(add, b2).compute()
        [5, 7, 9, 11, 13]

        Non-bag arguments are broadcast across all calls to the mapped
        function:

        >>> b.map(add, 1).compute()
        [1, 2, 3, 4, 5]

        Keyword arguments are also supported, and have the same semantics as
        regular arguments:

        >>> def myadd(x, y=0):
        ...     return x + y
        >>> b.map(myadd, y=b2).compute()
        [5, 7, 9, 11, 13]
        >>> b.map(myadd, y=1).compute()
        [1, 2, 3, 4, 5]

        Both arguments and keyword arguments can also be instances of
        ``dask.bag.Item``. Here we'll add the max value in the bag to each
        element:

        >>> b.map(myadd, b.max()).compute()
        [4, 5, 6, 7, 8]
        """
        return bag_map(func, self, *args, **kwargs)

    def starmap(self, func, **kwargs):
        """Apply a function using argument tuples from the given bag.

        This is similar to ``itertools.starmap``, except it also accepts
        keyword arguments. In pseudocode, this is could be written as:

        >>> def starmap(func, bag, **kwargs):
        ...     return (func(*args, **kwargs) for args in bag)

        Parameters
        ----------
        func : callable
        **kwargs : Item, Delayed, or object, optional
            Extra keyword arguments to pass to ``func``. These can either be
            normal objects, ``dask.bag.Item``, or ``dask.delayed.Delayed``.

        Examples
        --------
        >>> import dask.bag as db
        >>> data = [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]
        >>> b = db.from_sequence(data, npartitions=2)

        Apply a function to each argument tuple:

        >>> from operator import add
        >>> b.starmap(add).compute()
        [3, 7, 11, 15, 19]

        Apply a function to each argument tuple, with additional keyword
        arguments:

        >>> def myadd(x, y, z=0):
        ...     return x + y + z
        >>> b.starmap(myadd, z=10).compute()
        [13, 17, 21, 25, 29]

        Keyword arguments can also be instances of ``dask.bag.Item`` or
        ``dask.delayed.Delayed``:

        >>> max_second = b.pluck(1).max()
        >>> max_second.compute()
        10
        >>> b.starmap(myadd, z=max_second).compute()
        [13, 17, 21, 25, 29]
        """
        name = "{0}-{1}".format(
            funcname(func), tokenize(self, func, "starmap", **kwargs)
        )
        dependencies = [self]
        if kwargs:
            kwargs, collections = unpack_scalar_dask_kwargs(kwargs)
            dependencies.extend(collections)

        dsk = {
            (name, i): (reify, (starmap_chunk, func, (self.name, i), kwargs))
            for i in range(self.npartitions)
        }
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=dependencies)
        return type(self)(graph, name, self.npartitions)

    @property
    def _args(self):
        return (self.dask, self.name, self.npartitions)

    def __getstate__(self):
        return self._args

    def __setstate__(self, state):
        self.dask, self.name, self.npartitions = state

    def filter(self, predicate):
        """Filter elements in collection by a predicate function.

        >>> def iseven(x):
        ...     return x % 2 == 0

        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5))
        >>> list(b.filter(iseven))
        [0, 2, 4]
        """
        name = "filter-{0}-{1}".format(funcname(predicate), tokenize(self, predicate))
        dsk = dict(
            ((name, i), (reify, (filter, predicate, (self.name, i))))
            for i in range(self.npartitions)
        )
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return type(self)(graph, name, self.npartitions)

    def random_sample(self, prob, random_state=None):
        """Return elements from bag with probability of ``prob``.

        Parameters
        ----------
        prob : float
            A float between 0 and 1, representing the probability that each
            element will be returned.
        random_state : int or random.Random, optional
            If an integer, will be used to seed a new ``random.Random`` object.
            If provided, results in deterministic sampling.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5))
        >>> list(b.random_sample(0.5, 43))
        [0, 3, 4]
        >>> list(b.random_sample(0.5, 43))
        [0, 3, 4]
        """
        if not 0 <= prob <= 1:
            raise ValueError("prob must be a number in the interval [0, 1]")
        if not isinstance(random_state, Random):
            random_state = Random(random_state)

        name = "random-sample-%s" % tokenize(self, prob, random_state.getstate())
        state_data = random_state_data_python(self.npartitions, random_state)
        dsk = {
            (name, i): (reify, (random_sample, (self.name, i), state, prob))
            for i, state in zip(range(self.npartitions), state_data)
        }
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return type(self)(graph, name, self.npartitions)

    def remove(self, predicate):
        """Remove elements in collection that match predicate.

        >>> def iseven(x):
        ...     return x % 2 == 0

        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5))
        >>> list(b.remove(iseven))
        [1, 3]
        """
        name = "remove-{0}-{1}".format(funcname(predicate), tokenize(self, predicate))
        dsk = dict(
            ((name, i), (reify, (remove, predicate, (self.name, i))))
            for i in range(self.npartitions)
        )
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return type(self)(graph, name, self.npartitions)

    def map_partitions(self, func, *args, **kwargs):
        """Apply a function to every partition across one or more bags.

        Note that all ``Bag`` arguments must be partitioned identically.

        Parameters
        ----------
        func : callable
            The function to be called on every partition.
            This function should expect an ``Iterator`` or ``Iterable`` for
            every partition and should return an ``Iterator`` or ``Iterable``
            in return.
        *args, **kwargs : Bag, Item, Delayed, or object
            Arguments and keyword arguments to pass to ``func``.
            Partitions from this bag will be the first argument, and these will
            be passed *after*.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(range(1, 101), npartitions=10)
        >>> def div(nums, den=1):
        ...     return [num / den for num in nums]

        Using a python object:

        >>> hi = b.max().compute()
        >>> hi
        100
        >>> b.map_partitions(div, den=hi).take(5)
        (0.01, 0.02, 0.03, 0.04, 0.05)

        Using an ``Item``:

        >>> b.map_partitions(div, den=b.max()).take(5)
        (0.01, 0.02, 0.03, 0.04, 0.05)

        Note that while both versions give the same output, the second forms a
        single graph, and then computes everything at once, and in some cases
        may be more efficient.
        """
        return map_partitions(func, self, *args, **kwargs)

    def pluck(self, key, default=no_default):
        """Select item from all tuples/dicts in collection.

        >>> import dask.bag as db
        >>> b = db.from_sequence([{'name': 'Alice', 'credits': [1, 2, 3]},
        ...                       {'name': 'Bob',   'credits': [10, 20]}])
        >>> list(b.pluck('name'))
        ['Alice', 'Bob']
        >>> list(b.pluck('credits').pluck(0))
        [1, 10]
        """
        name = "pluck-" + tokenize(self, key, default)
        key = quote(key)
        if default == no_default:
            dsk = dict(
                ((name, i), (list, (pluck, key, (self.name, i))))
                for i in range(self.npartitions)
            )
        else:
            dsk = dict(
                ((name, i), (list, (pluck, key, (self.name, i), default)))
                for i in range(self.npartitions)
            )
        graph = HighLevelGraph.from_collections(name, dsk, dependencies=[self])
        return type(self)(graph, name, self.npartitions)

    def unzip(self, n):
        """Transform a bag of tuples to ``n`` bags of their elements.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence([(i, i + 1, i + 2) for i in range(10)])
        >>> first, second, third = b.unzip(3)
        >>> isinstance(first, db.Bag)
        True
        >>> first.compute()
        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

        Note that this is equivalent to:

        >>> first, second, third = (b.pluck(i) for i in range(3))
        """
        return tuple(self.pluck(i) for i in range(n))

    @wraps(to_textfiles)
    def to_textfiles(
        self,
        path,
        name_function=None,
        compression="infer",
        encoding=system_encoding,
        compute=True,
        storage_options=None,
        last_endline=False,
        **kwargs
    ):
        return to_textfiles(
            self,
            path,
            name_function,
            compression,
            encoding,
            compute,
            storage_options=storage_options,
            last_endline=last_endline,
            **kwargs
        )

    @wraps(to_avro)
    def to_avro(
        self,
        filename,
        schema,
        name_function=None,
        storage_options=None,
        codec="null",
        sync_interval=16000,
        metadata=None,
        compute=True,
        **kwargs
    ):
        return to_avro(
            self,
            filename,
            schema,
            name_function,
            storage_options,
            codec,
            sync_interval,
            metadata,
            compute,
            **kwargs
        )

    def fold(
        self, binop, combine=None, initial=no_default, split_every=None, out_type=Item
    ):
        """Parallelizable reduction

        Fold is like the builtin function ``reduce`` except that it works in
        parallel.  Fold takes two binary operator functions, one to reduce each
        partition of our dataset and another to combine results between
        partitions

        1.  ``binop``: Binary operator to reduce within each partition
        2.  ``combine``:  Binary operator to combine results from binop

        Sequentially this would look like the following:

        >>> intermediates = [reduce(binop, part) for part in partitions]  # doctest: +SKIP
        >>> final = reduce(combine, intermediates)  # doctest: +SKIP

        If only one function is given then it is used for both functions
        ``binop`` and ``combine`` as in the following example to compute the
        sum:

        >>> def add(x, y):
        ...     return x + y

        >>> import dask.bag as db
        >>> b = db.from_sequence(range(5))
        >>> b.fold(add).compute()
        10

        In full form we provide both binary operators as well as their default
        arguments

        >>> b.fold(binop=add, combine=add, initial=0).compute()
        10

        More complex binary operators are also doable

        >>> def add_to_set(acc, x):
        ...     ''' Add new element x to set acc '''
        ...     return acc | set([x])
        >>> b.fold(add_to_set, set.union, initial=set()).compute()
        {0, 1, 2, 3, 4}

        See Also
        --------

        Bag.foldby
        """
        combine = combine or binop
        if initial is not no_default:
            return self.reduction(
                curry(_reduce, binop, initial=initial),
                curry(_reduce, combine),
                split_every=split_every,
                out_type=out_type,
            )
        else:
            from tlz.curried import reduce

            return self.reduction(
                reduce(binop),
                reduce(combine),
                split_every=split_every,
                out_type=out_type,
            )

    def frequencies(self, split_every=None, sort=False):
        """Count number of occurrences of each distinct element.

        >>> import dask.bag as db
        >>> b = db.from_sequence(['Alice', 'Bob', 'Alice'])
        >>> dict(b.frequencies())       # doctest: +SKIP
        {'Alice': 2, 'Bob', 1}
        """
        result = self.reduction(
            frequencies,
            merge_frequencies,
            out_type=Bag,
            split_every=split_every,
            name="frequencies",
        ).map_partitions(dictitems)
        if sort:
            result = result.map_partitions(sorted, key=second, reverse=True)
        return result

    def topk(self, k, key=None, split_every=None):
        """K largest elements in collection

        Optionally ordered by some key function

        >>> import dask.bag as db
        >>> b = db.from_sequence([10, 3, 5, 7, 11, 4])
        >>> list(b.topk(2))
        [11, 10]

        >>> list(b.topk(2, lambda x: -x))
        [3, 4]
        """
        if key:
            if callable(key) and takes_multiple_arguments(key):
                key = partial(apply, key)
            func = partial(topk, k, key=key)
        else:
            func = partial(topk, k)
        return self.reduction(
            func,
            compose(func, toolz.concat),
            out_type=Bag,
            split_every=split_every,
            name="topk",
        )

    def distinct(self, key=None):
        """Distinct elements of collection

        Unordered without repeats.

        Parameters
        ----------
        key: {callable,str}
            Defines uniqueness of items in bag by calling ``key`` on each item.
            If a string is passed ``key`` is considered to be ``lambda x: x[key]``.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(['Alice', 'Bob', 'Alice'])
        >>> sorted(b.distinct())
        ['Alice', 'Bob']
        >>> b = db.from_sequence([{'name': 'Alice'}, {'name': 'Bob'}, {'name': 'Alice'}])
        >>> b.distinct(key=lambda x: x['name']).compute()
        [{'name': 'Alice'}, {'name': 'Bob'}]
        >>> b.distinct(key='name').compute()
        [{'name': 'Alice'}, {'name': 'Bob'}]
        """
        func = chunk_distinct if key is None else partial(chunk_distinct, key=key)
        agg = merge_distinct if key is None else partial(merge_distinct, key=key)
        return self.reduction(func, agg, out_type=Bag, name="distinct")

    def reduction(
        self, perpartition, aggregate, split_every=None, out_type=Item, name=None
    ):
        """Reduce collection with reduction operators.

        Parameters
        ----------
        perpartition: function
            reduction to apply to each partition
        aggregate: function
            reduction to apply to the results of all partitions
        split_every: int (optional)
            Group partitions into groups of this size while performing reduction
            Defaults to 8
        out_type: {Bag, Item}
            The out type of the result, Item if a single element, Bag if a list
            of elements.  Defaults to Item.

        Examples
        --------
        >>> import dask.bag as db
        >>> b = db.from_sequence(range(10))
        >>> b.reduction(sum, sum).compute()
        45
        """
        if split_every is None:
            split_every = 8
        if split_every is False:
            split_every = self.npartitions

        token = tokenize(self, perpartition, aggregate, split_every)
        a = "%s-part-%s" % (name or funcname(perpartition), token)
        is_last = self.npartitions == 1
        dsk = {
            (a, i): (empty_safe_apply, perpartition, (self.name, i), is_last)
            for i in range(self.npartitions)
        }
        k = self.npartitions
        b = a
        fmt = "%s-aggregate-%s" % (name or funcname(aggregate), token)
        depth = 0

        while k > split_every:
            c = fmt + str(depth)
            for i, inds in enumerate(partition_all(split_every, range(k))):
                dsk[(c, i)] = (
                    empty_safe_aggregate,
                    aggregate,
                    [(b, j) for j in inds],
                    False,
                )

            k = i + 1
            b = c
            depth += 1

        dsk[(fmt, 0)] = (
            empty_safe_aggregate,
            aggregate,
            [(b, j) for j in range(k)],
            True,
        )

        graph = HighLevelGraph.from_collections(fmt, dsk, dependencies=[self])
        if out_type is Item:
            dsk[fmt] = dsk.pop((fmt, 0))
            return Item(graph, fmt)
        else:
            return Bag(graph, fmt, 1)

    def sum(self, split_every=None):
        """Sum all elements"""
        return self.reduction(sum, sum, split_every=split_every)

    def max(self, split_every=None):
        """Maximum element"""
        return self.reduction(max, max, split_every=split_every)

    def min(self, split_every=None):
        """Minimum element"""
        return self.reduction(min, min, split_every=split_every)

    def any(self, split_every=None):
        """Are any of the elements truthy?

        Examples
        --------
        >>> import dask.bag as db
        >>> bool_bag = db.from_sequence([True, True, False])
        >>> bool_bag.any().compute()
        True
        """
        return self.reduction(any, any, split_every=split_every)

    def all(self, split_every=None):
        """Are all elements truthy?

        Examples
        --------
        >>> import dask.bag as db
        >>> bool_bag = db.from_sequence([True, True, False])
        >>> bool_bag.all().compute()
        False
        """
        return self.reduction(all, all, split_every=split_every)

    def count(self, split_every=None):
        """Count the number of elements.

        Examples
        --------
        >>> import dask.bag as db
        >>> numbers = db.from_sequence([1, 2, 3])
        >>> numbers.count().compute()
        3
        """
        return self.reduction(count, sum, split_every=split_every)

    def mean(self):
        """Arithmetic mean"""
{'fraction_non_alphanumeric': 0.127072495512762, 'fraction_numerical': 0.011651607800188616, 'mean_word_length': 2.0553980573500024, 'pattern_counts': {'":': 4, '<': 47, '<?xml version=': 0, '>': 708, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1497384', 'n_tokens_mistral': 20684, 'n_tokens_neox': 19556, 'n_words': 5777}