<commit_before>from itertools import count
from math import ceil
from toolz import merge, accumulate, unique, merge_sorted
from operator import getitem, setitem
import pandas as pd
import numpy as np
from chest import Chest

from .core import Frame, get
from ..compatibility import unicode
from ..utils import ignoring


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

store_names = ('store-%d' % i for i in count(1))
sort_names = ('sort-%d' % i for i in count(1))
index_names = ('index-%d' % i for i in count(1))
length_names = ('len-%d' % i for i in count(1))

def getitem2(a, b, default=None):
    try:
        return getitem(a, b)
    except KeyError:
        return default


def set_partition(f, column, blockdivs, cache=Chest):
    if callable(cache):
        cache = cache()

    set_index = 'set-index' + next(tokens)
    store = 'store-block' + next(tokens)

    # Set index on each block
    _set_index = {(set_index, i): (pd.DataFrame.set_index,
                                    (f.name, i), column)
                for i in range(f.npartitions)}

    # Store each block in cache
    _stores = {(store, i): (setitem, cache,
                                (tuple, [set_index, i]),
                                (set_index, i))
                for i in range(f.npartitions)}

    # Set new local indexes and store to disk
    get(merge(f.dask, _set_index, _stores), _stores.keys())

    # Do shuffle in cache
    old_keys = [(set_index, i) for i in range(f.npartitions)]
    new_keys = shuffle(cache, old_keys, blockdivs, delete=True)

    dsk = {k: (getitem, cache, (tuple, list(k))) for k in new_keys}

    return Frame(dsk, new_keys[0][0], f.columns, blockdivs)

def set_index(f, index, npartitions=None, cache=Chest, sortsize=2**24,
        chunksize=2**20, out_chunksize=2**16, empty=np.empty):
    """ Set Frame index to new column

    Sorts index and realigns frame to new sorted order.  This shuffles and
    repartitions your data.
    """
    npartitions = npartitions or f.npartitions
    if callable(cache):
        cache = cache()

    """
    We now use dask to compute indexes and reindex blocks

    1.  Compute each block of the dask
    2.  Set its index to match index
    3.  Pull out the index to separate data
    4.  Store all blocks
    5.  Store all indexes
    6.  Compute the lengths of all blocks
    """
    set_index = 'set-index' + next(tokens)
    indexname = 'index' + next(tokens)
    store = 'store-block' + next(tokens)
    store_index = 'store-index' + next(tokens)
    length = 'len-of-index' + next(tokens)

    if isinstance(index, Frame) and index.blockdivs == f.blockdivs:
        _set_index = {(set_index, i): (pd.DataFrame.set_index,
                                        (f.name, i), (index.name, i))
                    for i in range(f.npartitions)}
    elif isinstance(index, (str, unicode, int)):
        _set_index = {(set_index, i): (pd.DataFrame.set_index,
                                        (f.name, i), index)
                    for i in range(f.npartitions)}
    else:
        raise ValueError("Invalid index")
    _indexes = {(indexname, i): (getattr,
                                  (getattr, (set_index, i), 'index'),
                                  'values')
                for i in range(f.npartitions)}
    _stores = {(store, i): (setitem, cache,
                                (tuple, [set_index, i]),
                                (set_index, i))
                for i in range(f.npartitions)}
    _store_indexes = {(store_index, i): (setitem, cache,
                                         (tuple, [indexname, i]),
                                         (indexname, i))
                for i in range(f.npartitions)}
    _lengths = {(length, i): (len, (indexname, i))
                for i in range(f.npartitions)}

    dsk = merge(f.dask, _set_index, _indexes, _stores, _store_indexes, _lengths)

    if isinstance(index, Frame):
        dsk.update(index.dask)
    keys = [sorted(_lengths.keys()),
            sorted(_stores.keys()),
            sorted(_store_indexes.keys())]

    # Compute the frame and store blocks and index-blocks into cache
    lengths = get(dsk, keys)[0]

    blockdivs = blockdivs_by_approximate_percentiles(cache, indexname,
                    lengths, out_chunksize)

    old_keys = [(set_index, i) for i in range(f.npartitions)]
    new_keys = shuffle(cache, old_keys, blockdivs, delete=True)

    dsk = {k: (getitem, cache, (tuple, list(k))) for k in new_keys}

    return Frame(dsk, new_keys[0][0], f.columns, blockdivs)


def shard_df_on_index(df, blockdivs):
    """ Shard a DataFrame by ranges on its index

    Example
    -------

    >>> df = pd.DataFrame({'a': [0, 10, 20, 30, 40], 'b': [5, 4 ,3, 2, 1]})
    >>> df
        a  b
    0   0  5
    1  10  4
    2  20  3
    3  30  2
    4  40  1

    >>> shards = list(shard_df_on_index(df, [2, 4]))
    >>> shards[0]
        a  b
    0   0  5
    1  10  4

    >>> shards[1]
        a  b
    2  20  3
    3  30  2

    >>> shards[2]
        a  b
    4  40  1
    """
    blockdivs = list(blockdivs)
    df = df.sort()
    indices = df.index.searchsorted(blockdivs)
    yield df.iloc[:indices[0]]
    for i in range(len(indices) - 1):
        yield df.iloc[indices[i]: indices[i+1]]<commit_msg>don't unnecessarily turn blockdivs into list<commit_after>from itertools import count
from math import ceil
from toolz import merge, accumulate, unique, merge_sorted
from operator import getitem, setitem
import pandas as pd
import numpy as np
from chest import Chest

from .core import Frame, get
from ..compatibility import unicode
from ..utils import ignoring


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

store_names = ('store-%d' % i for i in count(1))
sort_names = ('sort-%d' % i for i in count(1))
index_names = ('index-%d' % i for i in count(1))
length_names = ('len-%d' % i for i in count(1))

def getitem2(a, b, default=None):
    try:
        return getitem(a, b)
    except KeyError:
        return default


def set_partition(f, column, blockdivs, cache=Chest):
    if callable(cache):
        cache = cache()

    set_index = 'set-index' + next(tokens)
    store = 'store-block' + next(tokens)

    # Set index on each block
    _set_index = {(set_index, i): (pd.DataFrame.set_index,
                                    (f.name, i), column)
                for i in range(f.npartitions)}

    # Store each block in cache
    _stores = {(store, i): (setitem, cache,
                                (tuple, [set_index, i]),
                                (set_index, i))
                for i in range(f.npartitions)}

    # Set new local indexes and store to disk
    get(merge(f.dask, _set_index, _stores), _stores.keys())

    # Do shuffle in cache
    old_keys = [(set_index, i) for i in range(f.npartitions)]
    new_keys = shuffle(cache, old_keys, blockdivs, delete=True)

    dsk = {k: (getitem, cache, (tuple, list(k))) for k in new_keys}

    return Frame(dsk, new_keys[0][0], f.columns, blockdivs)

def set_index(f, index, npartitions=None, cache=Chest, sortsize=2**24,
        chunksize=2**20, out_chunksize=2**16, empty=np.empty):
    """ Set Frame index to new column

    Sorts index and realigns frame to new sorted order.  This shuffles and
    repartitions your data.
    """
    npartitions = npartitions or f.npartitions
    if callable(cache):
        cache = cache()

    """
    We now use dask to compute indexes and reindex blocks

    1.  Compute each block of the dask
    2.  Set its index to match index
    3.  Pull out the index to separate data
    4.  Store all blocks
    5.  Store all indexes
    6.  Compute the lengths of all blocks
    """
    set_index = 'set-index' + next(tokens)
    indexname = 'index' + next(tokens)
    store = 'store-block' + next(tokens)
    store_index = 'store-index' + next(tokens)
    length = 'len-of-index' + next(tokens)

    if isinstance(index, Frame) and index.blockdivs == f.blockdivs:
        _set_index = {(set_index, i): (pd.DataFrame.set_index,
                                        (f.name, i), (index.name, i))
                    for i in range(f.npartitions)}
    elif isinstance(index, (str, unicode, int)):
        _set_index = {(set_index, i): (pd.DataFrame.set_index,
                                        (f.name, i), index)
                    for i in range(f.npartitions)}
    else:
        raise ValueError("Invalid index")
    _indexes = {(indexname, i): (getattr,
                                  (getattr, (set_index, i), 'index'),
                                  'values')
                for i in range(f.npartitions)}
    _stores = {(store, i): (setitem, cache,
                                (tuple, [set_index, i]),
                                (set_index, i))
                for i in range(f.npartitions)}
    _store_indexes = {(store_index, i): (setitem, cache,
                                         (tuple, [indexname, i]),
                                         (indexname, i))
                for i in range(f.npartitions)}
    _lengths = {(length, i): (len, (indexname, i))
                for i in range(f.npartitions)}

    dsk = merge(f.dask, _set_index, _indexes, _stores, _store_indexes, _lengths)

    if isinstance(index, Frame):
        dsk.update(index.dask)
    keys = [sorted(_lengths.keys()),
            sorted(_stores.keys()),
            sorted(_store_indexes.keys())]

    # Compute the frame and store blocks and index-blocks into cache
    lengths = get(dsk, keys)[0]

    blockdivs = blockdivs_by_approximate_percentiles(cache, indexname,
                    lengths, out_chunksize)

    old_keys = [(set_index, i) for i in range(f.npartitions)]
    new_keys = shuffle(cache, old_keys, blockdivs, delete=True)

    dsk = {k: (getitem, cache, (tuple, list(k))) for k in new_keys}

    return Frame(dsk, new_keys[0][0], f.columns, blockdivs)


def shard_df_on_index(df, blockdivs):
    """ Shard a DataFrame by ranges on its index

    Example
    -------

    >>> df = pd.DataFrame({'a': [0, 10, 20, 30, 40], 'b': [5, 4 ,3, 2, 1]})
    >>> df
        a  b
    0   0  5
    1  10  4
    2  20  3
    3  30  2
    4  40  1

    >>> shards = list(shard_df_on_index(df, [2, 4]))
    >>> shards[0]
        a  b
    0   0  5
    1  10  4

    >>> shards[1]
        a  b
    2  20  3
    3  30  2

    >>> shards[2]
        a  b
    4  40  1
    """
    df = df.sort()
    indices = df.index.searchsorted(blockdivs)
    yield df.iloc[:indices[0]]
    for i in range(len(indices) - 1):
        yield df.iloc[indices[i]: indices[i+1]]
{'fraction_non_alphanumeric': 0.12062588904694167, 'fraction_numerical': 0.016311047889995257, 'mean_word_length': 2.052387843704776, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3585544', 'n_tokens_mistral': 3584, 'n_tokens_neox': 3420, 'n_words': 1044}