<commit_before>from itertools import count
from collections import Iterator
from math import ceil
from toolz import merge, accumulate, merge_sorted
import toolz
from operator import getitem, setitem
import pandas as pd
import numpy as np
from pframe import pframe

from .. import threaded
from .core import DataFrame, Series, get, names, _Frame
from ..compatibility import unicode
from ..utils import ignoring


tokens = ('-%d' % i for i in count(1))


def set_index(f, index, npartitions=None, **kwargs):
    """ Set DataFrame index to new column

    Sorts index and realigns Dataframe to new sorted order.  This shuffles and
    repartitions your data.
    """
    npartitions = npartitions or f.npartitions
    if not isinstance(index, Series):
        index2 = f[index]
    else:
        index2 = index

    divisions = (index2
                  .quantiles(np.linspace(0, 100, npartitions+1)[1:-1])
                  .compute())
    return f.set_partition(index, divisions, **kwargs)


partition_names = ('set_partition-%d' % i for i in count(1))

def set_partition(f, index, divisions, get=threaded.get, **kwargs):
    """ Set new partitioning along index given divisions """
    divisions = unique(divisions)
    name = next(names)
    if isinstance(index, Series):
        assert index.divisions == f.divisions
        dsk = dict(((name, i), (f._partition_type.set_index, block, ind))
                for i, (block, ind) in enumerate(zip(f._keys(), index._keys())))
        f2 = type(f)(merge(f.dask, index.dask, dsk), name,
                       f.column_info, f.divisions)
    else:
        dsk = dict(((name, i), (f._partition_type.set_index, block, index))
                for i, block in enumerate(f._keys()))
        f2 = type(f)(merge(f.dask, dsk), name, f.column_info, f.divisions)

    head = f2.head()
    pf = pframe(like=head, divisions=divisions, **kwargs)

    def append(block):
        pf.append(block)
        return 0

    f2.map_blocks(append).compute(get=get)
    pf.flush()

    return from_pframe(pf)


def from_pframe(pf):
    """ Load dask.array from pframe """
    name = next(names)
    dsk = dict(((name, i), (pframe.get_partition, pf, i))
                for i in range(pf.npartitions))

    return DataFrame(dsk, name, pf.columns, pf.divisions)


def unique(divisions):
    """ Polymorphic unique function

    >>> list(unique([1, 2, 3, 1, 2, 3]))
    [1, 2, 3]

    >>> unique(np.array([1, 2, 3, 1, 2, 3]))
    array([1, 2, 3])

    >>> unique(pd.Categorical(['Alice', 'Bob', 'Alice'], ordered=False))
    [Alice, Bob]
    Categories (2, object): [Alice, Bob]
    """
    if isinstance(divisions, np.ndarray):
        return np.unique(divisions)
    if isinstance(divisions, pd.Categorical):
        return pd.Categorical.from_codes(np.unique(divisions.codes),
                divisions.categories, divisions.ordered)
    if isinstance(divisions, (tuple, list, Iterator)):
        return tuple(toolz.unique(divisions))
    raise NotImplementedError()


def shuffle(df, index, npartitions=None, use_server=False):
        """ Group DataFrame by index

        Hash grouping of elements.  After this operation all elements that have
        the same index will be in the same partition.  Note that this requires
        full dataset read, serialization and shuffle.  This is expensive.  If
        possible you should avoid shuffles.

        This does not preserve a meaningful index/partitioning scheme.

        See Also
        --------
        partd
        """
        if isinstance(index, _Frame):
            assert df.divisions == index.divisions
        if npartitions is None:
            npartitions = df.npartitions

        import partd
        p = ('zpartd' + next(tokens),)
        if use_server:
            dsk1 = {p: partd.PandasBlocks(partd.Shared())}
        else:
            dsk1 = {p: (partd.PandasBlocks,)}

        # Partition data on disk
        name = next(names)
        if isinstance(index, _Frame):
            dsk2 = dict(((name, i),
                         (partition, part, ind, npartitions, p))
                         for i, (part, ind)
                         in enumerate(zip(df._keys(), index._keys())))
        else:<commit_msg>dataframe instantiates partds in groupby call<commit_after>from itertools import count
from collections import Iterator
from math import ceil
from toolz import merge, accumulate, merge_sorted
import toolz
from operator import getitem, setitem
import pandas as pd
import numpy as np
from pframe import pframe

from .. import threaded
from .core import DataFrame, Series, get, names, _Frame
from ..compatibility import unicode
from ..utils import ignoring


tokens = ('-%d' % i for i in count(1))


def set_index(f, index, npartitions=None, **kwargs):
    """ Set DataFrame index to new column

    Sorts index and realigns Dataframe to new sorted order.  This shuffles and
    repartitions your data.
    """
    npartitions = npartitions or f.npartitions
    if not isinstance(index, Series):
        index2 = f[index]
    else:
        index2 = index

    divisions = (index2
                  .quantiles(np.linspace(0, 100, npartitions+1)[1:-1])
                  .compute())
    return f.set_partition(index, divisions, **kwargs)


partition_names = ('set_partition-%d' % i for i in count(1))

def set_partition(f, index, divisions, get=threaded.get, **kwargs):
    """ Set new partitioning along index given divisions """
    divisions = unique(divisions)
    name = next(names)
    if isinstance(index, Series):
        assert index.divisions == f.divisions
        dsk = dict(((name, i), (f._partition_type.set_index, block, ind))
                for i, (block, ind) in enumerate(zip(f._keys(), index._keys())))
        f2 = type(f)(merge(f.dask, index.dask, dsk), name,
                       f.column_info, f.divisions)
    else:
        dsk = dict(((name, i), (f._partition_type.set_index, block, index))
                for i, block in enumerate(f._keys()))
        f2 = type(f)(merge(f.dask, dsk), name, f.column_info, f.divisions)

    head = f2.head()
    pf = pframe(like=head, divisions=divisions, **kwargs)

    def append(block):
        pf.append(block)
        return 0

    f2.map_blocks(append).compute(get=get)
    pf.flush()

    return from_pframe(pf)


def from_pframe(pf):
    """ Load dask.array from pframe """
    name = next(names)
    dsk = dict(((name, i), (pframe.get_partition, pf, i))
                for i in range(pf.npartitions))

    return DataFrame(dsk, name, pf.columns, pf.divisions)


def unique(divisions):
    """ Polymorphic unique function

    >>> list(unique([1, 2, 3, 1, 2, 3]))
    [1, 2, 3]

    >>> unique(np.array([1, 2, 3, 1, 2, 3]))
    array([1, 2, 3])

    >>> unique(pd.Categorical(['Alice', 'Bob', 'Alice'], ordered=False))
    [Alice, Bob]
    Categories (2, object): [Alice, Bob]
    """
    if isinstance(divisions, np.ndarray):
        return np.unique(divisions)
    if isinstance(divisions, pd.Categorical):
        return pd.Categorical.from_codes(np.unique(divisions.codes),
                divisions.categories, divisions.ordered)
    if isinstance(divisions, (tuple, list, Iterator)):
        return tuple(toolz.unique(divisions))
    raise NotImplementedError()


def shuffle(df, index, npartitions=None, use_server=False):
        """ Group DataFrame by index

        Hash grouping of elements.  After this operation all elements that have
        the same index will be in the same partition.  Note that this requires
        full dataset read, serialization and shuffle.  This is expensive.  If
        possible you should avoid shuffles.

        This does not preserve a meaningful index/partitioning scheme.

        See Also
        --------
        partd
        """
        if isinstance(index, _Frame):
            assert df.divisions == index.divisions
        if npartitions is None:
            npartitions = df.npartitions

        import partd
        p = ('zpartd' + next(tokens),)
        if use_server:
            dsk1 = {}
            p = partd.PandasBlocks(partd.Shared())
        else:
            dsk1 = {}
            p = partd.PandasBlocks()
            # dsk1 = {p: (partd.PandasBlocks,)}

        # Partition data on disk
        name = next(names)
        if isinstance(index, _Frame):
            dsk2 = dict(((name, i),
                         (partition, part, ind, npartitions, p))
                         for i, (part, ind)
                         in enumerate(zip(df._keys(), index._keys())))
        else:
{'fraction_non_alphanumeric': 0.1180058651026393, 'fraction_numerical': 0.009266862170087977, 'mean_word_length': 2.4771615008156607, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4978595', 'n_tokens_mistral': 2670, 'n_tokens_neox': 2544, 'n_words': 773}