<commit_before>"""
Verb implementations for a :class:`pandas.DataFrame`
"""

import re
import warnings
from copy import copy

import numpy as np
import pandas as pd
import pandas.api.types as pdtypes

from .grouped_datatypes import GroupedDataFrame
from .options import get_option, options
from .utils import temporary_key, Q


def define(verb):
    if get_option('modify_input_data'):
        data = verb.data
    else:
        data = verb.data.copy()

    new_data = _evaluate_expressions_per_group(verb)
    for col in new_data:
        data[col] = new_data[col]
    return data


def create(verb):
    data = _get_base_dataframe(verb.data)
    new_data = _evaluate_expressions_per_group(verb)
    for col in new_data:
        data[col] = new_data[col]
    return data


def sample_n(verb):
    return verb.data.sample(**verb.kwargs)


def sample_frac(verb):
    return verb.data.sample(**verb.kwargs)


def select(verb):
    kw = verb.kwargs
    columns = verb.data.columns
    groups = _get_groups(verb)
    c0 = np.array([False]*len(columns))
    c1 = c2 = c3 = c4 = c5 = c6 = c0

    if verb.args:
        c1 = [x in set(verb.args) for x in columns]

    if kw['startswith']:
        c2 = [isinstance(x, str) and x.startswith(kw['startswith'])
              for x in columns]

    if kw['endswith']:
        c3 = [isinstance(x, str) and x.endswith(kw['endswith'])
              for x in columns]

    if kw['contains']:
        c4 = [isinstance(x, str) and kw['contains'] in x
              for x in columns]

    if kw['matches']:
        if hasattr(kw['matches'], 'match'):
            pattern = kw['matches']
        else:
            pattern = re.compile(kw['matches'])
        c5 = [isinstance(x, str) and bool(pattern.match(x))
              for x in columns]

    if groups:
        c6 = [x in set(groups) for x in columns]

    cond = np.logical_or.reduce((c1, c2, c3, c4, c5, c6))

    if kw['drop']:
        cond = ~cond

    data = verb.data.loc[:, cond]
    if data.is_copy:
        data.is_copy = None

    return data


def rename(verb):
    inplace = get_option('modify_input_data')
    data = verb.data.rename(columns=verb.lookup, inplace=inplace)
    return verb.data if inplace else data


def distinct(verb):
    if verb.new_columns:
        data = define(verb)
    else:
        data = verb.data
    return data.drop_duplicates(subset=verb.columns,
                                keep=verb.keep)


def arrange(verb):
    name_gen = ('col_{}'.format(x) for x in range(100))
    df = pd.DataFrame(index=verb.data.index)
    env = verb.env.with_outer_namespace({'Q': Q})
    for col, expr in zip(name_gen, verb.expressions):
        try:
            df[col] = verb.data[expr]
        except KeyError:
            df[col] = env.eval(expr, inner_namespace=verb.data)

    if len(df.columns):
        sorted_index = df.sort_values(by=list(df.columns)).index
        data = verb.data.loc[sorted_index, :]
    else:
        data = verb.data

    return data


def group_by(verb):
    if not all(g in verb.data for g in verb.groups):
        verb.data = define(verb)

    copy = not get_option('modify_input_data')
    return GroupedDataFrame(verb.data, verb.groups, copy=copy)


def ungroup(verb):
    return pd.DataFrame(verb.data)


def group_indices(verb):
    data = verb.data
    groups = verb.groups
    if isinstance(data, GroupedDataFrame):
        if groups:
            msg = "GroupedDataFrame ignored extra groups {}"
            warnings.warn(msg.format(groups))
        else:
            groups = data.plydata_groups
    else:
        data = create(verb)

    indices_dict = data.groupby(groups).indices
    indices = -np.ones(len(data), dtype=int)
    for i, (_, idx) in enumerate(sorted(indices_dict.items())):
        indices[idx] = i

    return indices


def summarize(verb):
    cols = verb.new_columns
    exprs = verb.expressions
    data = verb.data

    env = verb.env.with_outer_namespace(_aggregate_functions)
    env = env.with_outer_namespace({'Q': Q})

    try:
        grouper = data.groupby(data.plydata_groups)
    except AttributeError:
        data = _evaluate_summarize_expressions(exprs, cols, env, data)
    else:
        dfs = [_evaluate_summarize_expressions(exprs, cols, env, gdf)
               for _, gdf in grouper]
        data = pd.concat(dfs, axis=0, ignore_index=True)

    return data


def query(verb):
    data = verb.data.query(
        verb.expression,
        global_dict=verb.env.namespace,
        **verb.kwargs)
    data.is_copy = None
    return data


def do(verb):
    if verb.single_function:
        return _do_single_function(verb)
    else:
        return _do_functions(verb)


def head(verb):
    if isinstance(verb.data, GroupedDataFrame):
        grouper = verb.data.groupby(verb.data.plydata_groups)
        dfs = [gdf.head(verb.n) for _, gdf in grouper]
        data = pd.concat(dfs, axis=0, ignore_index=True)
        data.plydata_groups = list(verb.data.plydata_groups)
    else:
        data = verb.data.head(verb.n)

    return data


def tail(verb):
    if isinstance(verb.data, GroupedDataFrame):
        grouper = verb.data.groupby(verb.data.plydata_groups)
        dfs = [gdf.tail(verb.n) for _, gdf in grouper]
        data = pd.concat(dfs, axis=0, ignore_index=True)
        data.plydata_groups = list(verb.data.plydata_groups)
    else:
        data = verb.data.tail(verb.n)

    return data


def tally(verb):
    verb = copy(verb)
    # Prepare for summarize
    verb.new_columns = ['n']
    if verb.weights is not None:
        if isinstance(verb.weights, str):
            # Summarize will evaluate and sum up the weights
            verb.weights = 'sum({})'.format(verb.weights)
            verb.expressions = [verb.weights]
        else:
            # Do the summation here. The result does not depend
            # on the dataframe.
            verb.expressions = [np.sum(verb.weights)]
    else:
        verb.expressions = ['{n}']

    data = summarize(verb)
    if verb.sort:
        data = data.sort_values(by='n', ascending=False)
        data.reset_index(drop=True, inplace=True)

    return data


def count(verb):
    if (not isinstance(verb.data, GroupedDataFrame) and
            verb.groups):
        verb.data = GroupedDataFrame(verb.data, verb.groups, copy=True)
    return tally(verb)


def modify_where(verb):
    if get_option('modify_input_data'):
        data = verb.data
    else:
        data = verb.data.copy()
    idx = data.query(verb.where, global_dict=verb.env.namespace).index
    qdf = data.loc[idx, :]

    env = verb.env.with_outer_namespace({'Q': Q})
    for col, expr in zip(verb.columns, verb.expressions):
        # Do not create new columns, define does that
        if col not in data:
            raise KeyError("Column '{}' not in dataframe".format(col))
        if isinstance(expr, str):
            data.loc[idx, col] = env.eval(expr, inner_namespace=qdf)
        else:
            data.loc[idx, col] = expr

    return data


def define_where(verb):
    verb = copy(verb)
    if not get_option('modify_input_data'):
        verb.data = verb.data.copy()

    alt_expressions = [x[0] for x in verb.expressions]
    default_expressions = [x[1] for x in verb.expressions]

    with options(modify_input_data=True):
        verb.expressions = default_expressions
        verb.data = define(verb)

        verb.columns = verb.new_columns
        verb.expressions = alt_expressions
        data = modify_where(verb)

    return data


def dropna(verb):
    result = verb.data.dropna(
        axis=verb.axis,
        how=verb.how,
        thresh=verb.thresh,
        subset=verb.subset
    )
    return result


def fillna(verb):
    inplace = get_option('modify_input_data')
    result = verb.data.fillna(
        value=verb.value,
        method=verb.method,
        axis=verb.axis,
        limit=verb.limit,
        downcast=verb.downcast,
        inplace=inplace
    )
    return result if not inplace else verb.data


def call(verb):
    if isinstance(verb.func, str):
        func = getattr(verb.data, verb.func.lstrip('.'))
        return func(*verb.args, **verb.kwargs)
    else:
        return verb.func(verb.data, *verb.args, **verb.kwargs)


def inner_join(verb):
    verb.kwargs['how'] = 'inner'
    return _join(verb)


def outer_join(verb):
    verb.kwargs['how'] = 'outer'
    return _join(verb)


def left_join(verb):
    verb.kwargs['how'] = 'left'
    return _join(verb)


def right_join(verb):
    verb.kwargs['how'] = 'right'
    return _join(verb)


def anti_join(verb):
    verb.kwargs['how'] = 'left'
    verb.kwargs['suffixes'] = ('', '_y')
    verb.kwargs['indicator'] = '_plydata_merge'
    df = _join(verb)
    data = df.query('_plydata_merge=="left_only"')[verb.x.columns]
    data.is_copy = None
    return data


def semi_join(verb):
    verb.kwargs['how'] = 'left'
    verb.kwargs['suffixes'] = ('', '_y')
    verb.kwargs['indicator'] = '_plydata_merge'
    df = _join(verb)
    data = df.query('_plydata_merge=="both"')[verb.x.columns]
    data.is_copy = None
    data.drop_duplicates(inplace=True)
    return data


# Helper functions

def _get_groups(verb):
    """
    Return groups
    """
    try:
        return verb.data.plydata_groups
    except AttributeError:
        return []


def _get_base_dataframe(df):
    """
    Remove all columns other than those grouped on
    """
    if isinstance(df, GroupedDataFrame):
        base_df = GroupedDataFrame(
            df.loc[:, df.plydata_groups], df.plydata_groups,
            copy=True)
    else:
        base_df = pd.DataFrame(index=df.index)
    return base_df


def _add_group_columns(data, gdf):
    """
    Add group columns to data with a value from the grouped dataframe

    It is assumed that the grouped dataframe contains a single group

    >>> data = pd.DataFrame({
    ...     'x': [5, 6, 7]})
    >>> gdf = GroupedDataFrame({
    ...     'g': list('aaa'),
    ...     'x': range(3)}, groups=['g'])
    >>> _add_group_columns(data, gdf)
       g  x
    0  a  5
    1  a  6
    2  a  7
    """
    n = len(data)
    if isinstance(gdf, GroupedDataFrame):
        for i, col in enumerate(gdf.plydata_groups):
            if col not in data:
                group_values = [gdf[col].iloc[0]] * n
                # Need to be careful and maintain the dtypes
                # of the group columns
                if pdtypes.is_categorical_dtype(gdf[col]):
                    col_values = pd.Categorical(
                        group_values,
                        categories=gdf[col].cat.categories,
                        ordered=gdf[col].cat.ordered
                    )
                else:
                    col_values = pd.Series(
                        group_values,
                        index=data.index,
                        dtype=gdf[col].dtype
                    )
                # Group columns come first
                data.insert(i, col, col_values)
    return data


def _evaluate_expressions_per_group(verb):
    """
    Evaluate Expressions and return the columns in a new dataframe
    """
    cols = verb.new_columns
    exprs = verb.expressions
    data = verb.data
    env = verb.env.with_outer_namespace({'Q': Q})

    try:
        grouper = data.groupby(data.plydata_groups)
    except AttributeError:
        data = _evaluate_expressions(exprs, cols, env, data)
    else:
        dfs = [_evaluate_expressions(exprs, cols, env, gdf)
               for _, gdf in grouper]
        data = pd.concat(dfs, axis=0, ignore_index=True)

    return data


def _evaluate_expressions(expressions, columns, env, gdf):
    """
    Evaluate Expressions and return the columns in a new dataframe
    """
    data = pd.DataFrame(index=gdf.index)
    env = env.with_outer_namespace({'Q': Q})
    for col, expr in zip(columns, expressions):
        if isinstance(expr, str):
            value = env.eval(expr, inner_namespace=gdf)
        else:
            value = expr

        data[col] = np.asarray(value)
    return _add_group_columns(data, gdf)


def _evaluate_summarize_expressions(expressions, columns, env, gdf):
    """
    Evaluate expressions to create a dataframe.

    Parameters
    ----------
    expressions : list of str
    columns : list of str
    env : environment
    gdf : dataframe
        Dataframe where all items are assumed to belong to the
        same group.

    This excutes an *apply* part of the *split-apply-combine*
    data manipulation strategy. Callers of the function do
    the *split* and the *combine*.

    A peak into the function

    >>> import pandas as pd
    >>> from .utils import get_empty_env

    A Dataframe with many groups

    >>> df = GroupedDataFrame(
    ...     {'x': list('aabbcc'),
    ...      'y': [0, 1, 2, 3, 4, 5]
    ...      }, groups=['x'])

    Do a groupby and obtain a dataframe with a single group,
    (i.e. the *split*)

    >>> grouper = df.groupby(df.plydata_groups)
    >>> group = 'b'      # The group we want to evalualte
    >>> gdf = grouper.get_group(group)
    >>> gdf
    groups: ['x']
      x  y
    2 b  2
    3 b  3

    Create the other parameters

    >>> env = get_empty_env()
    >>> columns = ['ysq', 'ycubed']
    >>> expressions = ['y**2', 'y**3']

    Finally, the *apply*

    >>> _evaluate_summarize_expressions(expressions, columns, env, gdf)
       x  ysq  ycubed
    0  b    4       8
    1  b    9      27

    The caller does the *combine*, for one or more of these
    results.
    """
    # Extra aggregate function that the user references with
    # the name `{n}`. It returns the length of the dataframe.
    def _plydata_n():
        return len(gdf)

    for col, expr in zip(columns, expressions):
        if isinstance(expr, str):
            expr = expr.format(n='_plydata_n()')
            with temporary_key(_aggregate_functions,
                               '_plydata_n', _plydata_n):
                value = env.eval(expr, inner_namespace=gdf)
        else:
            value = expr

        # Non-consecutive 0-n pd.series indices create gaps with
        # nans when inserted into a dataframe
        value = np.asarray(value)

        try:
            data[col] = value
        except NameError:
            try:
                n = len(value)
            except TypeError:
                n = 1
            data = pd.DataFrame({col: value}, index=range(n))

    return _add_group_columns(data, gdf)


def _evaluate_do_single_function(function, gdf):
    """
    Evaluate an expression to create a dataframe.

    Similar to :func:`_evaluate_summarize_expressions`, but for the
    ``do`` operation.
    """
    gdf.is_copy = None
    data = function(gdf)
    return _add_group_columns(data, gdf)


def _evaluate_do_functions(functions, columns, gdf):
    """
    Evaluate functions to create a dataframe.

    Similar to :func:`_evaluate_summarize_expressions`, but for the
    ``do`` operation.
    """
    gdf.is_copy = None
    for col, func in zip(columns, functions):
        value = np.asarray(func(gdf))

        try:
            data[col] = value
        except NameError:
            try:
                n = len(value)
            except TypeError:
                n = 1
            data = pd.DataFrame({col: value}, index=range(n))

    return _add_group_columns(data, gdf)


def _do_single_function(verb):
    func = verb.single_function
    data = verb.data

    try:
        groups = data.plydata_groups
    except AttributeError:
        data = _evaluate_do_single_function(func, data)
    else:
        dfs = [_evaluate_do_single_function(func, gdf)
               for _, gdf in data.groupby(groups)]
        data = pd.concat(dfs, axis=0, ignore_index=True)
        data.plydata_groups = list(groups)

    return data


def _do_functions(verb):
    cols = verb.columns
    funcs = verb.functions
    data = verb.data

    try:
        groups = data.plydata_groups
    except AttributeError:
        data = _evaluate_do_functions(funcs, cols, data)
    else:
        dfs = [_evaluate_do_functions(funcs, cols, gdf)
               for _, gdf in data.groupby(groups)]
        data = pd.concat(dfs, axis=0, ignore_index=True)
        data.plydata_groups = list(groups)

    return data


def _join(verb):
    """
    Join helper
    """
    data = pd.merge(verb.x, verb.y, **verb.kwargs)

    # Preserve x groups
    if isinstance(verb.x, GroupedDataFrame):
        data.plydata_groups = list(verb.x.plydata_groups)
    return data


# Aggregations functions

def _nth(arr, n):
    """
    Return the nth value of array

    If it is missing return NaN
    """
    try:
        return arr.iloc[n]
    except (KeyError, IndexError):
        return np.nan


def _n_distinct(arr):
    """
    Number of unique values in array
    """
    return len(pd.unique(arr))


_aggregate_functions = {
    'min': np.min,
    'max': np.max,
    'sum': np.sum,
    'cumsum': np.cumsum,
    'mean': np.mean,
    'median': np.median,
    'std': np.std,
    'first': lambda x: _nth(x, 0),
    'last': lambda x: _nth(x, -1),
    'nth': _nth,
    'n_distinct': _n_distinct,
    'n_unique': _n_distinct,
}
<commit_msg>Optimize select implementation<commit_after>"""
Verb implementations for a :class:`pandas.DataFrame`
"""

import re
import warnings
from copy import copy

import numpy as np
import pandas as pd
import pandas.api.types as pdtypes

from .grouped_datatypes import GroupedDataFrame
from .options import get_option, options
from .utils import temporary_key, Q


def define(verb):
    if get_option('modify_input_data'):
        data = verb.data
    else:
        data = verb.data.copy()

    new_data = _evaluate_expressions_per_group(verb)
    for col in new_data:
        data[col] = new_data[col]
    return data


def create(verb):
    data = _get_base_dataframe(verb.data)
    new_data = _evaluate_expressions_per_group(verb)
    for col in new_data:
        data[col] = new_data[col]
    return data


def sample_n(verb):
    return verb.data.sample(**verb.kwargs)


def sample_frac(verb):
    return verb.data.sample(**verb.kwargs)


def select(verb):
    kw = verb.kwargs
    columns = verb.data.columns
    groups = _get_groups(verb)
    conds = []

    if verb.args:
        _args = set(verb.args)
        c1 = [x in _args for x in columns]
        conds.append(c1)

    if kw['startswith']:
        c2 = [isinstance(x, str) and x.startswith(kw['startswith'])
              for x in columns]
        conds.append(c2)

    if kw['endswith']:
        c3 = [isinstance(x, str) and x.endswith(kw['endswith'])
              for x in columns]
        conds.append(c3)

    if kw['contains']:
        c4 = [isinstance(x, str) and kw['contains'] in x
              for x in columns]
        conds.append(c4)

    if kw['matches']:
        if hasattr(kw['matches'], 'match'):
            pattern = kw['matches']
        else:
            pattern = re.compile(kw['matches'])
        c5 = [isinstance(x, str) and bool(pattern.match(x))
              for x in columns]
        conds.append(c5)

    if groups:
        _groups = set(groups)
        c6 = [x in _groups for x in columns]
        conds.append(c6)

    if conds:
        cond = np.logical_or.reduce(conds)
    else:
        cond = np.array([False]*len(columns))

    if kw['drop']:
        cond = ~cond

    data = verb.data.loc[:, cond]
    if data.is_copy:
        data.is_copy = None

    return data


def rename(verb):
    inplace = get_option('modify_input_data')
    data = verb.data.rename(columns=verb.lookup, inplace=inplace)
    return verb.data if inplace else data


def distinct(verb):
    if verb.new_columns:
        data = define(verb)
    else:
        data = verb.data
    return data.drop_duplicates(subset=verb.columns,
                                keep=verb.keep)


def arrange(verb):
    name_gen = ('col_{}'.format(x) for x in range(100))
    df = pd.DataFrame(index=verb.data.index)
    env = verb.env.with_outer_namespace({'Q': Q})
    for col, expr in zip(name_gen, verb.expressions):
        try:
            df[col] = verb.data[expr]
        except KeyError:
            df[col] = env.eval(expr, inner_namespace=verb.data)

    if len(df.columns):
        sorted_index = df.sort_values(by=list(df.columns)).index
        data = verb.data.loc[sorted_index, :]
    else:
        data = verb.data

    return data


def group_by(verb):
    if not all(g in verb.data for g in verb.groups):
        verb.data = define(verb)

    copy = not get_option('modify_input_data')
    return GroupedDataFrame(verb.data, verb.groups, copy=copy)


def ungroup(verb):
    return pd.DataFrame(verb.data)


def group_indices(verb):
    data = verb.data
    groups = verb.groups
    if isinstance(data, GroupedDataFrame):
        if groups:
            msg = "GroupedDataFrame ignored extra groups {}"
            warnings.warn(msg.format(groups))
        else:
            groups = data.plydata_groups
    else:
        data = create(verb)

    indices_dict = data.groupby(groups).indices
    indices = -np.ones(len(data), dtype=int)
    for i, (_, idx) in enumerate(sorted(indices_dict.items())):
        indices[idx] = i

    return indices


def summarize(verb):
    cols = verb.new_columns
    exprs = verb.expressions
    data = verb.data

    env = verb.env.with_outer_namespace(_aggregate_functions)
    env = env.with_outer_namespace({'Q': Q})

    try:
        grouper = data.groupby(data.plydata_groups)
    except AttributeError:
        data = _evaluate_summarize_expressions(exprs, cols, env, data)
    else:
        dfs = [_evaluate_summarize_expressions(exprs, cols, env, gdf)
               for _, gdf in grouper]
        data = pd.concat(dfs, axis=0, ignore_index=True)

    return data


def query(verb):
    data = verb.data.query(
        verb.expression,
        global_dict=verb.env.namespace,
        **verb.kwargs)
    data.is_copy = None
    return data


def do(verb):
    if verb.single_function:
        return _do_single_function(verb)
    else:
        return _do_functions(verb)


def head(verb):
    if isinstance(verb.data, GroupedDataFrame):
        grouper = verb.data.groupby(verb.data.plydata_groups)
        dfs = [gdf.head(verb.n) for _, gdf in grouper]
        data = pd.concat(dfs, axis=0, ignore_index=True)
        data.plydata_groups = list(verb.data.plydata_groups)
    else:
        data = verb.data.head(verb.n)

    return data


def tail(verb):
    if isinstance(verb.data, GroupedDataFrame):
        grouper = verb.data.groupby(verb.data.plydata_groups)
        dfs = [gdf.tail(verb.n) for _, gdf in grouper]
        data = pd.concat(dfs, axis=0, ignore_index=True)
        data.plydata_groups = list(verb.data.plydata_groups)
    else:
        data = verb.data.tail(verb.n)

    return data


def tally(verb):
    verb = copy(verb)
    # Prepare for summarize
    verb.new_columns = ['n']
    if verb.weights is not None:
        if isinstance(verb.weights, str):
            # Summarize will evaluate and sum up the weights
            verb.weights = 'sum({})'.format(verb.weights)
            verb.expressions = [verb.weights]
        else:
            # Do the summation here. The result does not depend
            # on the dataframe.
            verb.expressions = [np.sum(verb.weights)]
    else:
        verb.expressions = ['{n}']

    data = summarize(verb)
    if verb.sort:
        data = data.sort_values(by='n', ascending=False)
        data.reset_index(drop=True, inplace=True)

    return data


def count(verb):
    if (not isinstance(verb.data, GroupedDataFrame) and
            verb.groups):
        verb.data = GroupedDataFrame(verb.data, verb.groups, copy=True)
    return tally(verb)


def modify_where(verb):
    if get_option('modify_input_data'):
        data = verb.data
    else:
        data = verb.data.copy()
    idx = data.query(verb.where, global_dict=verb.env.namespace).index
    qdf = data.loc[idx, :]

    env = verb.env.with_outer_namespace({'Q': Q})
    for col, expr in zip(verb.columns, verb.expressions):
        # Do not create new columns, define does that
        if col not in data:
            raise KeyError("Column '{}' not in dataframe".format(col))
        if isinstance(expr, str):
            data.loc[idx, col] = env.eval(expr, inner_namespace=qdf)
        else:
            data.loc[idx, col] = expr

    return data


def define_where(verb):
    verb = copy(verb)
    if not get_option('modify_input_data'):
        verb.data = verb.data.copy()

    alt_expressions = [x[0] for x in verb.expressions]
    default_expressions = [x[1] for x in verb.expressions]

    with options(modify_input_data=True):
        verb.expressions = default_expressions
        verb.data = define(verb)

        verb.columns = verb.new_columns
        verb.expressions = alt_expressions
        data = modify_where(verb)

    return data


def dropna(verb):
    result = verb.data.dropna(
        axis=verb.axis,
        how=verb.how,
        thresh=verb.thresh,
        subset=verb.subset
    )
    return result


def fillna(verb):
    inplace = get_option('modify_input_data')
    result = verb.data.fillna(
        value=verb.value,
        method=verb.method,
        axis=verb.axis,
        limit=verb.limit,
        downcast=verb.downcast,
        inplace=inplace
    )
    return result if not inplace else verb.data


def call(verb):
    if isinstance(verb.func, str):
        func = getattr(verb.data, verb.func.lstrip('.'))
        return func(*verb.args, **verb.kwargs)
    else:
        return verb.func(verb.data, *verb.args, **verb.kwargs)


def inner_join(verb):
    verb.kwargs['how'] = 'inner'
    return _join(verb)


def outer_join(verb):
    verb.kwargs['how'] = 'outer'
    return _join(verb)


def left_join(verb):
    verb.kwargs['how'] = 'left'
    return _join(verb)


def right_join(verb):
    verb.kwargs['how'] = 'right'
    return _join(verb)


def anti_join(verb):
    verb.kwargs['how'] = 'left'
    verb.kwargs['suffixes'] = ('', '_y')
    verb.kwargs['indicator'] = '_plydata_merge'
    df = _join(verb)
    data = df.query('_plydata_merge=="left_only"')[verb.x.columns]
    data.is_copy = None
    return data


def semi_join(verb):
    verb.kwargs['how'] = 'left'
    verb.kwargs['suffixes'] = ('', '_y')
    verb.kwargs['indicator'] = '_plydata_merge'
    df = _join(verb)
    data = df.query('_plydata_merge=="both"')[verb.x.columns]
    data.is_copy = None
    data.drop_duplicates(inplace=True)
    return data


# Helper functions

def _get_groups(verb):
    """
    Return groups
    """
    try:
        return verb.data.plydata_groups
    except AttributeError:
        return []


def _get_base_dataframe(df):
    """
    Remove all columns other than those grouped on
    """
    if isinstance(df, GroupedDataFrame):
        base_df = GroupedDataFrame(
            df.loc[:, df.plydata_groups], df.plydata_groups,
            copy=True)
    else:
        base_df = pd.DataFrame(index=df.index)
    return base_df


def _add_group_columns(data, gdf):
    """
    Add group columns to data with a value from the grouped dataframe

    It is assumed that the grouped dataframe contains a single group

    >>> data = pd.DataFrame({
    ...     'x': [5, 6, 7]})
    >>> gdf = GroupedDataFrame({
    ...     'g': list('aaa'),
    ...     'x': range(3)}, groups=['g'])
    >>> _add_group_columns(data, gdf)
       g  x
    0  a  5
    1  a  6
    2  a  7
    """
    n = len(data)
    if isinstance(gdf, GroupedDataFrame):
        for i, col in enumerate(gdf.plydata_groups):
            if col not in data:
                group_values = [gdf[col].iloc[0]] * n
                # Need to be careful and maintain the dtypes
                # of the group columns
                if pdtypes.is_categorical_dtype(gdf[col]):
                    col_values = pd.Categorical(
                        group_values,
                        categories=gdf[col].cat.categories,
                        ordered=gdf[col].cat.ordered
                    )
                else:
                    col_values = pd.Series(
                        group_values,
                        index=data.index,
                        dtype=gdf[col].dtype
                    )
                # Group columns come first
                data.insert(i, col, col_values)
    return data


def _evaluate_expressions_per_group(verb):
    """
    Evaluate Expressions and return the columns in a new dataframe
    """
    cols = verb.new_columns
    exprs = verb.expressions
    data = verb.data
    env = verb.env.with_outer_namespace({'Q': Q})

    try:
        grouper = data.groupby(data.plydata_groups)
    except AttributeError:
        data = _evaluate_expressions(exprs, cols, env, data)
    else:
        dfs = [_evaluate_expressions(exprs, cols, env, gdf)
               for _, gdf in grouper]
        data = pd.concat(dfs, axis=0, ignore_index=True)

    return data


def _evaluate_expressions(expressions, columns, env, gdf):
    """
    Evaluate Expressions and return the columns in a new dataframe
    """
    data = pd.DataFrame(index=gdf.index)
    env = env.with_outer_namespace({'Q': Q})
    for col, expr in zip(columns, expressions):
        if isinstance(expr, str):
            value = env.eval(expr, inner_namespace=gdf)
        else:
            value = expr

        data[col] = np.asarray(value)
    return _add_group_columns(data, gdf)


def _evaluate_summarize_expressions(expressions, columns, env, gdf):
    """
    Evaluate expressions to create a dataframe.

    Parameters
    ----------
    expressions : list of str
    columns : list of str
    env : environment
    gdf : dataframe
        Dataframe where all items are assumed to belong to the
        same group.

    This excutes an *apply* part of the *split-apply-combine*
    data manipulation strategy. Callers of the function do
    the *split* and the *combine*.

    A peak into the function

    >>> import pandas as pd
    >>> from .utils import get_empty_env

    A Dataframe with many groups

    >>> df = GroupedDataFrame(
    ...     {'x': list('aabbcc'),
    ...      'y': [0, 1, 2, 3, 4, 5]
    ...      }, groups=['x'])

    Do a groupby and obtain a dataframe with a single group,
    (i.e. the *split*)

    >>> grouper = df.groupby(df.plydata_groups)
    >>> group = 'b'      # The group we want to evalualte
    >>> gdf = grouper.get_group(group)
    >>> gdf
    groups: ['x']
      x  y
    2 b  2
    3 b  3

    Create the other parameters

    >>> env = get_empty_env()
    >>> columns = ['ysq', 'ycubed']
    >>> expressions = ['y**2', 'y**3']

    Finally, the *apply*

    >>> _evaluate_summarize_expressions(expressions, columns, env, gdf)
       x  ysq  ycubed
    0  b    4       8
    1  b    9      27

    The caller does the *combine*, for one or more of these
    results.
    """
    # Extra aggregate function that the user references with
    # the name `{n}`. It returns the length of the dataframe.
    def _plydata_n():
        return len(gdf)

    for col, expr in zip(columns, expressions):
        if isinstance(expr, str):
            expr = expr.format(n='_plydata_n()')
            with temporary_key(_aggregate_functions,
                               '_plydata_n', _plydata_n):
                value = env.eval(expr, inner_namespace=gdf)
        else:
            value = expr

        # Non-consecutive 0-n pd.series indices create gaps with
        # nans when inserted into a dataframe
        value = np.asarray(value)

        try:
            data[col] = value
        except NameError:
            try:
                n = len(value)
            except TypeError:
                n = 1
            data = pd.DataFrame({col: value}, index=range(n))

    return _add_group_columns(data, gdf)


def _evaluate_do_single_function(function, gdf):
    """
    Evaluate an expression to create a dataframe.

    Similar to :func:`_evaluate_summarize_expressions`, but for the
    ``do`` operation.
    """
    gdf.is_copy = None
    data = function(gdf)
    return _add_group_columns(data, gdf)


def _evaluate_do_functions(functions, columns, gdf):
    """
    Evaluate functions to create a dataframe.

    Similar to :func:`_evaluate_summarize_expressions`, but for the
    ``do`` operation.
    """
    gdf.is_copy = None
    for col, func in zip(columns, functions):
        value = np.asarray(func(gdf))

        try:
            data[col] = value
        except NameError:
            try:
                n = len(value)
            except TypeError:
                n = 1
            data = pd.DataFrame({col: value}, index=range(n))

    return _add_group_columns(data, gdf)


def _do_single_function(verb):
    func = verb.single_function
    data = verb.data

    try:
        groups = data.plydata_groups
    except AttributeError:
        data = _evaluate_do_single_function(func, data)
    else:
        dfs = [_evaluate_do_single_function(func, gdf)
               for _, gdf in data.groupby(groups)]
        data = pd.concat(dfs, axis=0, ignore_index=True)
        data.plydata_groups = list(groups)

    return data


def _do_functions(verb):
    cols = verb.columns
    funcs = verb.functions
    data = verb.data

    try:
        groups = data.plydata_groups
    except AttributeError:
        data = _evaluate_do_functions(funcs, cols, data)
    else:
        dfs = [_evaluate_do_functions(funcs, cols, gdf)
               for _, gdf in data.groupby(groups)]
        data = pd.concat(dfs, axis=0, ignore_index=True)
        data.plydata_groups = list(groups)

    return data


def _join(verb):
    """
    Join helper
    """
    data = pd.merge(verb.x, verb.y, **verb.kwargs)

    # Preserve x groups
    if isinstance(verb.x, GroupedDataFrame):
        data.plydata_groups = list(verb.x.plydata_groups)
    return data


# Aggregations functions

def _nth(arr, n):
    """
    Return the nth value of array

    If it is missing return NaN
    """
    try:
        return arr.iloc[n]
    except (KeyError, IndexError):
        return np.nan


def _n_distinct(arr):
    """
    Number of unique values in array
    """
    return len(pd.unique(arr))


_aggregate_functions = {
    'min': np.min,
    'max': np.max,
    'sum': np.sum,
    'cumsum': np.cumsum,
    'mean': np.mean,
    'median': np.median,
    'std': np.std,
    'first': lambda x: _nth(x, 0),
    'last': lambda x: _nth(x, -1),
    'nth': _nth,
    'n_distinct': _n_distinct,
    'n_unique': _n_distinct,
}

{'fraction_non_alphanumeric': 0.11332713106175137, 'fraction_numerical': 0.003605070357018258, 'mean_word_length': 2.273099248263393, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 87, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2930853', 'n_tokens_mistral': 11357, 'n_tokens_neox': 10961, 'n_words': 2756}