<commit_before>"""
Expression splitting for chunked computation

To evaluate an expression on a large dataset we may need to chunk that dataset
into pieces and evaluate on each of the pieces individually.  This module
contains logic to break up an expression-to-be-evaluated-on-the-entire-array
into

1.  An expression to be evaluated on each chunk of the data
2.  An expression to be evaluated on the concatenated intermediate results

As an example, consider counting the number of non-null elements in a large
list.  We have the following recipe

1.  Break the large collection into chunks, each of which fits comfortably into
    memory
2.  For each chunk, call compute chunk.count()
3.  Gather all of these results into a single list (which hopefully fits in
    memory)
4.  for this aggregate, compute aggregate.sum()

And so, given this expression

    expr -> expr.count()

We needed the following expressions

    chunk -> chunk.count()
    agg -> agg.sum()

This module performs this transformation for a wide array of chunkable
expressions.  It supports elementwise operations, reductions,
split-apply-combine, and selections.  It notably does not support sorting,
joining, or slicing.

If explicit chunksizes are given it can also reason about the size and shape of
the intermediate aggregate.  It can also do this in N-Dimensions.
"""
from __future__ import absolute_import, division, print_function

from toolz import concat
import datashape
from datashape.predicates import isscalar, isrecord
from math import floor

from .core import *
from .expressions import *
from .expressions import ndim, shape
from .math import sqrt
from .reductions import *
from .split_apply_combine import *
from .collections import *
from .table import *
from ..dispatch import dispatch
from ..compatibility import builtins

good_to_split = (Reduction, Summary, By, Distinct)
can_split = good_to_split + (Selection, ElemWise)

__all__ = ['path_split', 'split']

def path_split(leaf, expr):
    """ Find the right place in the expression tree/line to parallelize

    >>> t = Symbol('t', 'var * {name: string, amount: int, id: int}')

    >>> path_split(t, t.amount.sum() + 1)
    sum(t.amount)

    >>> path_split(t, t.amount.distinct().sort())
    distinct(t.amount)
    """
    last = None
    for node in list(path(expr, leaf))[:-1][::-1]:
        if isinstance(node, good_to_split):
            return node
        elif not isinstance(node, can_split):
            return last
        last = node
    if not node:
        raise ValueError()
    return node


def split(leaf, expr, chunk=None, agg=None, **kwargs):
    """ Split expression for chunked computation

    Break up a computation ``leaf -> expr`` so that it can be run in chunks.
    This returns two computations, one to perform on each chunk and then one to
    perform on the union of these intermediate results

        chunk -> chunk-expr
        aggregate -> aggregate-expr

    The chunk_expr will have the same dimesions as the input
    (reductions set keepdims=True) so that this should work cleanly with
    concatenation functions like ``np.concatenate``.

    Returns
    -------

    Pair of (Symbol, Expr) pairs

        (chunk, chunk_expr), (aggregate, aggregate_expr)

    >>> t = Symbol('t', 'var * {name: string, amount: int, id: int}')
    >>> expr = t.id.count()
    >>> split(t, expr)
    ((chunk, count(chunk.id, keepdims=True)), (aggregate, sum(aggregate)))
    """
    center = path_split(leaf, expr)
    if not chunk:
        if leaf.ndim > 1:
            raise ValueError("Please provide a chunk symbol")
        else:
            chunk = Symbol('chunk', datashape.var * leaf.dshape.measure)

    chunk_expr = _split_chunk(center, leaf=leaf, chunk=chunk, **kwargs)

    if not agg:
        agg_shape = aggregate_shape(leaf, expr, chunk, chunk_expr)
        agg_dshape = DataShape(*(agg_shape + (chunk_expr.dshape.measure,)))
        agg = Symbol('aggregate', agg_dshape)

    agg_expr = _split_agg(center, leaf=leaf, agg=agg)

    return ((chunk, chunk_expr),
            (agg, expr._subs({center: agg})._subs({agg: agg_expr})))


reductions = {sum: (sum, sum), count: (count, sum),
              min: (min, min), max: (max, max),
              any: (any, any), all: (all, all)}

@dispatch(Expr)
def _split(expr, leaf=None, chunk=None, agg=None, keepdims=True):
    return ((chunk, _split_chunk(expr, leaf=leaf, chunk=chunk,
                                 keepdims=keepdims)),
            (agg, _split_agg(expr, leaf=leaf, agg=agg)))


@dispatch(tuple(reductions))
def _split_chunk(expr, leaf=None, chunk=None, keepdims=True):<commit_msg>expr: split nelements<commit_after>"""
Expression splitting for chunked computation

To evaluate an expression on a large dataset we may need to chunk that dataset
into pieces and evaluate on each of the pieces individually.  This module
contains logic to break up an expression-to-be-evaluated-on-the-entire-array
into

1.  An expression to be evaluated on each chunk of the data
2.  An expression to be evaluated on the concatenated intermediate results

As an example, consider counting the number of non-null elements in a large
list.  We have the following recipe

1.  Break the large collection into chunks, each of which fits comfortably into
    memory
2.  For each chunk, call compute chunk.count()
3.  Gather all of these results into a single list (which hopefully fits in
    memory)
4.  for this aggregate, compute aggregate.sum()

And so, given this expression

    expr -> expr.count()

We needed the following expressions

    chunk -> chunk.count()
    agg -> agg.sum()

This module performs this transformation for a wide array of chunkable
expressions.  It supports elementwise operations, reductions,
split-apply-combine, and selections.  It notably does not support sorting,
joining, or slicing.

If explicit chunksizes are given it can also reason about the size and shape of
the intermediate aggregate.  It can also do this in N-Dimensions.
"""
from __future__ import absolute_import, division, print_function

from toolz import concat
import datashape
from datashape.predicates import isscalar, isrecord
from math import floor

from .core import *
from .expressions import *
from .expressions import ndim, shape
from .math import sqrt
from .reductions import *
from .split_apply_combine import *
from .collections import *
from .table import *
from ..dispatch import dispatch
from ..compatibility import builtins

good_to_split = (Reduction, Summary, By, Distinct)
can_split = good_to_split + (Selection, ElemWise)

__all__ = ['path_split', 'split']

def path_split(leaf, expr):
    """ Find the right place in the expression tree/line to parallelize

    >>> t = Symbol('t', 'var * {name: string, amount: int, id: int}')

    >>> path_split(t, t.amount.sum() + 1)
    sum(t.amount)

    >>> path_split(t, t.amount.distinct().sort())
    distinct(t.amount)
    """
    last = None
    for node in list(path(expr, leaf))[:-1][::-1]:
        if isinstance(node, good_to_split):
            return node
        elif not isinstance(node, can_split):
            return last
        last = node
    if not node:
        raise ValueError()
    return node


def split(leaf, expr, chunk=None, agg=None, **kwargs):
    """ Split expression for chunked computation

    Break up a computation ``leaf -> expr`` so that it can be run in chunks.
    This returns two computations, one to perform on each chunk and then one to
    perform on the union of these intermediate results

        chunk -> chunk-expr
        aggregate -> aggregate-expr

    The chunk_expr will have the same dimesions as the input
    (reductions set keepdims=True) so that this should work cleanly with
    concatenation functions like ``np.concatenate``.

    Returns
    -------

    Pair of (Symbol, Expr) pairs

        (chunk, chunk_expr), (aggregate, aggregate_expr)

    >>> t = Symbol('t', 'var * {name: string, amount: int, id: int}')
    >>> expr = t.id.count()
    >>> split(t, expr)
    ((chunk, count(chunk.id, keepdims=True)), (aggregate, sum(aggregate)))
    """
    center = path_split(leaf, expr)
    if not chunk:
        if leaf.ndim > 1:
            raise ValueError("Please provide a chunk symbol")
        else:
            chunk = Symbol('chunk', datashape.var * leaf.dshape.measure)

    chunk_expr = _split_chunk(center, leaf=leaf, chunk=chunk, **kwargs)

    if not agg:
        agg_shape = aggregate_shape(leaf, expr, chunk, chunk_expr)
        agg_dshape = DataShape(*(agg_shape + (chunk_expr.dshape.measure,)))
        agg = Symbol('aggregate', agg_dshape)

    agg_expr = _split_agg(center, leaf=leaf, agg=agg)

    return ((chunk, chunk_expr),
            (agg, expr._subs({center: agg})._subs({agg: agg_expr})))


reductions = {sum: (sum, sum), count: (count, sum),
              min: (min, min), max: (max, max),
              any: (any, any), all: (all, all),
              nelements: (nelements, sum)}

@dispatch(Expr)
def _split(expr, leaf=None, chunk=None, agg=None, keepdims=True):
    return ((chunk, _split_chunk(expr, leaf=leaf, chunk=chunk,
                                 keepdims=keepdims)),
            (agg, _split_agg(expr, leaf=leaf, agg=agg)))


@dispatch(tuple(reductions))
def _split_chunk(expr, leaf=None, chunk=None, keepdims=True):
{'fraction_non_alphanumeric': 0.10192452424470487, 'fraction_numerical': 0.0021503064186646596, 'mean_word_length': 3.3104726598702503, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3580996', 'n_tokens_mistral': 2749, 'n_tokens_neox': 2677, 'n_words': 1032}