class Aggregator(object):
    """Perform a aggregations on datasets.

    Apply the `aggregation` to group columns by `groups` and the `columns`
    of the `dframe`. Store the resulting `dframe` as a linked dataset for
    `dataset`. If a linked dataset with the same groups already exists update
    this dataset.  Otherwise create a new linked dataset.
    """

    def __init__(self, dframe, groups, _type, name, columns):
        """Create an Aggregator.

        :param columns: The columns to aggregate over.
        :param dframe: The DataFrame to aggregate.
        :param groups: A list of columns to group on.
        :param _type: The aggregation to perform.
        :param name: The name of the aggregation.
        """
        self.columns = columns
        self.dframe = dframe
        self.groups = groups
        self.name = name
        aggregation = AGGREGATIONS.get(_type)
        self.aggregation = aggregation(self.name, self.groups, self.dframe)

    def save(self, dataset):
        """Save this aggregation.

        If an aggregated dataset for this aggregations group already exists
        store in this dataset, if not create a new aggregated dataset and store
        the aggregation in this new aggregated dataset.

        """
        new_dframe = self.aggregation.eval(self.columns)
        new_dframe = add_parent_column(new_dframe, dataset.dataset_id)

        a_dataset = dataset.aggregated_dataset(self.groups)

        if a_dataset is None:
            a_dataset = aggregated_dataset(dataset, new_dframe, self.groups)
        else:
            a_dframe = a_dataset.dframe()
            new_dframe = group_join(self.groups, a_dframe, new_dframe)
            a_dataset.replace_observations(new_dframe)

        self.new_dframe = new_dframe

    def update(self, dataset, child_dataset, formula, reducible):
        """Attempt to reduce an update and store."""
        parent_dataset_id = dataset.dataset_id

        # get dframe only including rows from this parent
        dframe = rows_for_parent_id(child_dataset.dframe(
            keep_parent_ids=True, reload_=True), parent_dataset_id)

        # remove rows in child from parent
        child_dataset.remove_parent_observations(parent_dataset_id)

        if reducible and self.__is_reducible():
            dframe = self.aggregation.reduce(dframe, self.columns)
        else:
            dframe = self.updated_dframe(dataset, formula, dframe)

        new_a_dframe = concat([child_dataset.dframe(), dframe])
        new_a_dframe = add_parent_column(new_a_dframe, parent_dataset_id)
        child_dataset.replace_observations(new_a_dframe)

        return child_dataset.dframe()

    def updated_dframe(self, dataset, formula, dframe):
        """Create a new aggregation and update return updated dframe."""
        # build column arguments from original dframe
        columns = parse_columns(dataset, formula, self.name, self.dframe)
        new_dframe = self.aggregation.eval(columns)

        new_columns = [x for x in new_dframe.columns if x not in self.groups]

        dframe = dframe.drop(new_columns, axis=1)
        dframe = group_join(self.groups, new_dframe, dframe)

        return dframe

    def __is_reducible(self):
        """If it is not grouped and a reduce is defined."""
        return not self.groups and 'reduce' in dir(self.aggregation) def add_parent_column(df, parent_dataset_id):
    """Add parent ID column to this DataFrame."""
    return add_constant_column(df, parent_dataset_id, PARENT_DATASET_ID) def join_dataset(left, other, on):
    """Left join an `other` dataset.

    :param other: Other dataset to join.
    :param on: Column or 2 comma seperated columns to join on.

    :returns: Joined DataFrame.

    :raises: `KeyError` if join columns not in datasets.
    """
    on_lhs, on_rhs = (on.split(',') * 2)[:2]

    right_dframe = other.dframe(padded=True)

    if on_lhs not in left.columns:
        raise KeyError('No item named "%s" in left hand side dataset' % on_lhs)

    if on_rhs not in right_dframe.columns:
        raise KeyError('No item named "%s" in right hand side dataset' %
                       on_rhs)

    right_dframe = right_dframe.set_index(on_rhs)

    if len(right_dframe.index) != len(right_dframe.index.unique()):
        msg = 'Join column "%s" of the right hand side dataset is not unique'
        raise NonUniqueJoinError(msg % on_rhs)

    shared_columns = left.columns.intersection(right_dframe.columns)

    if len(shared_columns):
        rename_map = [{c: '%s.%s' % (c, v) for c in shared_columns} for v
                      in ['x', 'y']]
        left.rename(columns=rename_map[0], inplace=True)
        right_dframe.rename(columns=rename_map[1], inplace=True)

    return left.join(right_dframe, on=on_lhs) class Parser(object):
    """Class for parsing and evaluating formula.

    Attributes:

    - aggregation: Aggregation parsed from formula.
    - aggregation_names: Possible aggregations.
    - bnf: Cached Backus-Naur Form of formula.
    - column_functions: Cached additional columns as aggregation parameters.
    - function_names: Names of possible functions in formulas.
    - operator_names: Names of possible operators in formulas.
    - parsed_expr: Cached parsed expression.
    - special_names: Names of possible reserved names in formulas.
    - reserved_words: List of all possible reserved words that may be used in
      formulas.
    """

    aggregation = None
    aggregation_names = AGGREGATIONS.keys()
    bnf = None
    column_functions = None
    function_names = ['date', 'percentile', 'today']
    operator_names = ['and', 'or', 'not', 'in']
    parsed_expr = None
    special_names = ['default']

    reserved_words = aggregation_names + function_names + operator_names +\
        special_names

    def __init__(self):
        self.bnf = self.__build_bnf()

    @classmethod
    def dependent_columns(cls, formula, dataset):
        functions, _ = cls.parse(formula)
        columns = [get_dependent_columns(f, dataset) for f in functions]

        return set.union(set(), *columns)

    @property
    def functions(self):
        return self.column_functions if self.aggregation else self.parsed_expr

    def store_aggregation(self, _, __, tokens):
        """Cached a parsed aggregation."""
        self.aggregation = tokens[0]
        self.column_functions = tokens[1:]

    def __build_bnf(self):
        """Parse formula to function based on language definition.

        Backus-Naur Form of formula language:

        =========   ==========
        Operation   Expression
        =========   ==========
        addop       '+' | '-'
        multop      '*' | '/'
        expop       '^'
        compop      '==' | '<' | '>' | '<=' | '>='
        notop       'not'
        andop       'and'
        orop        'or'
        real        \d+(.\d+)
        integer     \d+
        variable    \w+
        string      ".+"
        atom        real | integer | variable
        func        func ( atom )
        factor      atom [ expop factor]*
        term        factor [ multop factor ]*
        expr        term [ addop term ]*
        equation    expr [compop expr]*
        in          string in '[' "string"[, "string"]* ']'
        neg         [notop]* equation | in
        conj        neg [andop neg]*
        disj        conj [orop conj]*
        case        'case' disj: atom[, disj: atom]*[, 'default': atom]
        trans       trans ( case )
        agg         agg ( trans[, trans]* )
        =========   ==========

        """
        if self.bnf:
            return self.bnf

        # literal operators
        exp_op = Literal('^')
        sign_op = oneOf('+ -')
        mult_op = oneOf('* /')
        plus_op = oneOf('+ -')
        not_op = CaselessLiteral('not')
        and_op = CaselessLiteral('and')
        or_op = CaselessLiteral('or')
        in_op = CaselessLiteral('in').suppress()
        comparison_op = oneOf('< <= > >= != ==')
        case_op = CaselessLiteral('case').suppress()

        # aggregation functions
        aggregations = build_caseless_or_expression(self.aggregation_names)

        # literal syntactic
        open_bracket = Literal('[').suppress()
        close_bracket = Literal(']').suppress()
        open_paren = Literal('(').suppress()
        close_paren = Literal(')').suppress()
        comma = Literal(',').suppress()
        dquote = Literal('"').suppress()
        colon = Literal(':').suppress()

        # functions
        date_func = CaselessLiteral('date')
        percentile_func = CaselessLiteral('percentile')
        today_func = CaselessLiteral('today()').setParseAction(EvalToday)

        # case statment
        default = CaselessLiteral('default')

        reserved_words = MatchFirst(
            [Keyword(word) for word in self.reserved_words])

        # atoms
        integer = Word(nums)
        real = Combine(Word(nums) + '.' + Word(nums))
        variable = ~reserved_words + Word(alphanums + '_')
        atom = real | integer | variable
        atom.setParseAction(EvalConstant)

        # everything between pairs of double quotes is a string
        string = dquote + Regex('[^"]+') + dquote
        string.setParseAction(EvalString)

        # expressions
        in_list = open_bracket + string +\
            ZeroOrMore(comma + string) + close_bracket

        func_expr = operatorPrecedence(string, [
            (date_func, 1, opAssoc.RIGHT, EvalDate),
        ]) | today_func

        arith_expr = operatorPrecedence(atom | func_expr, [
            (sign_op, 1, opAssoc.RIGHT, EvalSignOp),
            (exp_op, 2, opAssoc.RIGHT, EvalExpOp),
            (mult_op, 2, opAssoc.LEFT, EvalMultOp),
            (plus_op, 2, opAssoc.LEFT, EvalPlusOp),
        ])

        comp_expr = operatorPrecedence(arith_expr, [
            (comparison_op, 2, opAssoc.LEFT, EvalComparisonOp),
        ])

        prop_expr = operatorPrecedence(comp_expr | in_list, [
            (in_op, 2, opAssoc.RIGHT, EvalInOp),
            (not_op, 1, opAssoc.RIGHT, EvalNotOp),
            (and_op, 2, opAssoc.LEFT, EvalAndOp),
            (or_op, 2, opAssoc.LEFT, EvalOrOp),
        ])

        default_statement = (default + colon + atom).setParseAction(EvalMapOp)
        map_statement = (prop_expr + colon + atom).setParseAction(EvalMapOp)

        case_list = map_statement + ZeroOrMore(
            comma + map_statement) + Optional(comma + default_statement)

        case_expr = operatorPrecedence(case_list, [
            (case_op, 1, opAssoc.RIGHT, EvalCaseOp),
        ]) | prop_expr

        trans_expr = operatorPrecedence(case_expr, [
            (percentile_func, 1, opAssoc.RIGHT, EvalPercentile),
        ])

        return ((aggregations + open_paren + Optional(
            trans_expr + ZeroOrMore(comma + trans_expr)))
            .setParseAction(self.store_aggregation) + close_paren)\
            | trans_expr

    @classmethod
    def parse(cls, formula):
        """Parse formula and return evaluation function.

        Parse `formula` into an aggregation name and functions.
        There will be multiple functions is the aggregation takes multiple
        arguments, e.g. ratio which takes a numerator and denominator formula.

        Examples:

        - constants
            - ``9 + 5``,
        - aliases
            - ``rating``,
            - ``gps``,
        - arithmetic
            - ``amount + gps_alt``,
            - ``amount - gps_alt``,
            - ``amount + 5``,
            - ``amount - gps_alt + 2.5``,
            - ``amount * gps_alt``,
            - ``amount / gps_alt``,
            - ``amount * gps_alt / 2.5``,
            - ``amount + gps_alt * gps_precision``,
        - precedence
            - ``(amount + gps_alt) * gps_precision``,
        - comparison
            - ``amount == 2``,
            - ``10 < amount``,
            - ``10 < amount + gps_alt``,
        - logical
            - ``not amount == 2``,
            - ``not(amount == 2)``,
            - ``amount == 2 and 10 < amount``,
            - ``amount == 2 or 10 < amount``,
            - ``not not amount == 2 or 10 < amount``,
            - ``not amount == 2 or 10 < amount``,
            - ``not amount == 2) or 10 < amount``,
            - ``not(amount == 2 or 10 < amount)``,
            - ``amount ^ 3``,
            - ``amount + gps_alt) ^ 2 + 100``,
            - ``amount``,
            - ``amount < gps_alt - 100``,
        - membership
            - ``rating in ["delectible"]``,
            - ``risk_factor in ["low_risk"]``,
            - ``amount in ["9.0", "2.0", "20.0"]``,
            - ``risk_factor in ["low_risk"]) and (amount in ["9.0", "20.0"])``,
        - dates
            - ``date("09-04-2012") - submit_date > 21078000``,
        - cases
            - ``case food_type in ["morning_food"]: 1, default: 3``
        - transformations: row-wise column based aggregations
            - ``percentile(amount)``

        :param formula: The string to parse.

        :returns: A tuple with the name of the aggregation in the formula, if
           any and a list of functions built from the input string.
        """
        parser = cls()

        try:
            parser.parsed_expr = parser.bnf.parseString(formula, parseAll=True)
        except ParseException, err:
            raise ParseError('Parse Failure for string "%s": %s' % (
                             formula, err))

        return [parser.functions, parser.aggregation]

    @classmethod
    def parse_aggregation(cls, formula):
        _, a = cls.parse(formula)
        return a

    @classmethod
    def parse_function(cls, formula):
        return cls.parse_functions(formula)[0]

    @classmethod
    def parse_functions(cls, formula):
        return [partial(f.eval) for f in cls.parse(formula)[0]]

    @classmethod
    def validate(cls, dataset, formula, groups):
        """Validate `formula` and `groups` for dataset.

        Validate the formula and group string by attempting to get a row from
        the dframe for the dataset and then running parser validation on this
        row. Additionally, ensure that the groups in the group string are
        columns in the dataset.

        :param dataset: The dataset to validate for.
        :param formula: The formula to validate.
        :param groups: A list of columns to group by.

        :returns: The aggregation (or None) for the formula.
        """
        cls.validate_formula(formula, dataset)

        for group in groups:
            if not group in dataset.schema.keys():
                raise ParseError(
                    'Group %s not in dataset columns.' % group)

    @classmethod
    def validate_formula(cls, formula, dataset):
        """Validate the *formula* on an example *row* of data.

        Rebuild the BNF then parse the `formula` given the sample `row`.

        :param formula: The formula to validate.
        :param dataset: The dataset to validate against.

        :returns: The aggregation for the formula.
        """
        # check valid formula
        cls.parse(formula)
        schema = dataset.schema

        if not schema:
            raise ParseError(
                'No schema for dataset, please add data or wait for it to '
                'finish processing')

        for column in cls.dependent_columns(formula, dataset):
            if column not in schema.keys():
                raise ParseError('Missing column reference: %s' % column)

    def __getstate__(self):
        """Get state for pickle."""
        return [
            self.aggregation,
            self.aggregation_names,
            self.function_names,
            self.operator_names,
            self.special_names,
            self.reserved_words,
            self.special_names,
        ]

    def __setstate__(self, state):
        """Set internal variables from pickled state."""
        self.aggregation, self.aggregation_names, self.function_names,\
            self.operator_names, self.special_names, self.reserved_words,\
            self.special_names = state
        self.__build_bnf() def recognize_dates(df, schema=None):
    """Convert data columns to datetimes.

    Check if object columns in a dataframe can be parsed as dates.
    If yes, rewrite column with values parsed as dates.

    If schema is passed, convert columes to datetime if column in *schema* is
    of type datetime.

    :param df: The DataFrame to convert columns in.
    :param schema: Schema to define columns of type datetime.

    :returns: A DataFrame with column values convert to datetime types.
    """
    return __parse_dates_schema(df, schema) if schema else __parse_dates(df) def df_to_jsondict(df):
    """Return DataFrame as a list of dicts for each row."""
    return [series_to_jsondict(series) for _, series in df.iterrows()] MONGO_ID = '_id' def parse_columns(dataset, formula, name, dframe=None, no_index=False):
    """Parse a formula and return columns resulting from its functions.

    Parse a formula into a list of functions then apply those functions to
    the Data Frame and return the resulting columns.

    :param formula: The formula to parse.
    :param name: Name of the formula.
    :param dframe: A DataFrame to apply functions to.
    :param no_index: Drop the index on result columns.
    """
    functions = Parser.parse_functions(formula)
    dependent_columns = Parser.dependent_columns(formula, dataset)

    # make select from dependent_columns
    if dframe is None:
        select = {col: 1 for col in dependent_columns or [MONGO_ID]}

        dframe = dataset.dframe(
            query_args=QueryArgs(select=select),
            keep_mongo_keys=True).set_index(MONGO_ID_ENCODED)

        if not dependent_columns:
            # constant column, use dummy
            dframe['dummy'] = 0

    return __build_columns(dataset, dframe, functions, name, no_index) class QueryArgs(object):
    def __init__(self, query=None, select=None, distinct=None, limit=0,
                 order_by=None, dataset=None):
        """A holder for query arguments.

        :param query: An optional query.
        :param select: An optional select to limit the fields in the dframe.
        :param distinct: Return distinct entries for this field.
        :param limit: Limit on the number of rows in the returned dframe.
        :param order_by: Sort resulting rows according to a column value and
            sign indicating ascending or descending.

        Example of `order_by`:

          - ``order_by='mycolumn'``
          - ``order_by='-mycolumn'``
        """
        self.query = parse_dates_from_query(query, dataset)
        self.select = select
        self.distinct = distinct
        self.limit = limit
        self.order_by = parse_order_by(order_by)

    def encode(self, encoding, query):
        """Encode query, order_by, and select given an encoding.

        The query will be combined with the existing query.

        :param encoding: A dict to encode the QueryArgs fields with.
        :param query: An additional dict to combine with the existing query.
        """
        self.query = replace_keys(combine_dicts(self.query, query), encoding)
        self.order_by = self.order_by and replace_keys(dict(self.order_by),
                                                       encoding).items()
        self.select = self.select and replace_keys(self.select, encoding)

    def __nonzero__(self):
        return bool(self.query or self.select or self.distinct or self.limit
                    or self.order_by) def combine_dicts(*dicts):
    """Combine dicts with keys in later dicts taking precedence."""
    return dict(chain(*[_dict.iteritems() for _dict in dicts])) def flatten(list_):
    return [item for sublist in list_ for item in sublist] def to_list(maybe_list):
    return maybe_list if isinstance(maybe_list, list) else [maybe_list]
from collections import defaultdict
from celery.task import task
from pandas import concat, DataFrame
from bamboo.core.aggregator import Aggregator
from bamboo.core.frame import add_parent_column, join_dataset
from bamboo.core.parser import Parser
from bamboo.lib.datetools import recognize_dates
from bamboo.lib.jsontools import df_to_jsondict
from bamboo.lib.mongo import MONGO_ID
from bamboo.lib.parsing import parse_columns
from bamboo.lib.query_args import QueryArgs
from bamboo.lib.utils import combine_dicts, flatten, to_list

        offset += dataset.find_one(parent_id).num_rows

    return offset


def __propagate_column(dataset, parent_dataset):
    """Propagate columns in `parent_dataset` to `dataset`.

    When a new calculation is added to a dataset this will propagate the
    new column to all child (merged) datasets.

    :param dataset: THe child dataet.
    :param parent_dataset: The dataset to propagate.
    """
    # delete the rows in this dataset from the parent
    dataset.remove_parent_observations(parent_dataset.dataset_id)

    # get this dataset without the out-of-date parent rows
    dframe = dataset.dframe(keep_parent_ids=True)

    # create new dframe from the upated parent and add parent id
    parent_dframe = add_parent_column(parent_dataset.dframe(),
                                      parent_dataset.dataset_id)

    # merge this new dframe with the existing dframe
    updated_dframe = concat([dframe, parent_dframe])

    # save new dframe (updates schema)
    dataset.replace_observations(updated_dframe)
    dataset.clear_summary_stats()

    # recur into merged dataset
    [__propagate_column(x, dataset) for x in dataset.merged_datasets]


def __remapped_data(dataset_id, mapping, slugified_data):
    column_map = mapping.get(dataset_id) if mapping else None

    if column_map:
        slugified_data = [{column_map.get(k, k): v for k, v in row.items()}
                          for row in slugified_data]

    return slugified_data


def __slugify_data(new_data, labels_to_slugs):
    slugified_data = []
    new_data = to_list(new_data)

    for row in new_data:
        for key, value in row.iteritems():
            if labels_to_slugs.get(key) and key != MONGO_ID:
                del row[key]
                row[labels_to_slugs[key]] = value

        slugified_data.append(row)

    return slugified_data


def __update_aggregate_datasets(dataset, new_dframe, update=None):
    calcs_to_data = __calculation_data(dataset)

    for formula, slug, groups, a_dataset in calcs_to_data:
        __update_aggregate_dataset(dataset, formula, new_dframe, slug, groups,
                                   a_dataset, update is None)


def __update_aggregate_dataset(dataset, formula, new_dframe, name, groups,
                               a_dataset, reducible):
    """Update the aggregated dataset built for `dataset` with `calculation`.

    Proceed with the following steps:

        - delete the rows in this dataset from the parent
        - recalculate aggregated dataframe from aggregation
        - update aggregated dataset with new dataframe and add parent id
        - recur on all merged datasets descending from the aggregated
          dataset

    :param formula: The formula to execute.
    :param new_dframe: The DataFrame to aggregate on.
    :param name: The name of the aggregation.
    :param groups: A column or columns to group on.
    :type group: String, list of strings, or None.
    :param a_dataset: The DataSet to store the aggregation in.
    """
    # parse aggregation and build column arguments
    aggregator = __create_aggregator(
        dataset, formula, name, groups, dframe=new_dframe)
    new_agg_dframe = aggregator.update(dataset, a_dataset, formula, reducible)

    # jsondict from new dframe
    new_data = df_to_jsondict(new_agg_dframe)

    for merged_dataset in a_dataset.merged_datasets:
        # remove rows in child from this merged dataset
        merged_dataset.remove_parent_observations(a_dataset.dataset_id)

        # calculate updates for the child
        calculate_updates(merged_dataset, new_data,
                          parent_dataset_id=a_dataset.dataset_id)


def __update_joined_datasets(dataset, update):
    """Update any joined datasets."""
    if 'add' in update:
        new_dframe = update['add']

    for direction, other_dataset, on, j_dataset in dataset.joined_datasets:
        if 'add' in update:
            if direction == 'left':
                # only proceed if on in new dframe
                if on in new_dframe.columns:
                    left_dframe = other_dataset.dframe(padded=True)

                    # only proceed if new on value is in on column in lhs
                    if len(set(new_dframe[on]).intersection(
                            set(left_dframe[on]))):

                        merged_dframe = join_dataset(left_dframe, dataset, on)