interpreter.py

import raco.myrial.groupby as groupby
import raco.myrial.multiway as multiway
from raco.myrial.cfg import ControlFlowGraph
from raco.myrial.emitarg import FullWildcardEmitArg, TableWildcardEmitArg
from raco.myrial.exceptions import *
import raco.algebra
import raco.expression
import raco.catalog
import raco.scheme
from raco.language.myrialang import (MyriaLeftDeepTreeAlgebra,
                                     MyriaHyperCubeAlgebra)
from raco.language.myrialang import compile_to_json
from raco.compile import optimize
from raco import relation_key
from raco.expression import StateVar

import collections
import copy


class DuplicateAliasException(Exception):
    """Bag comprehension arguments must have different alias names."""
    pass


class InvalidStatementException(Exception):
    pass


def get_unnamed_ref(column_ref, scheme, offset=0):
    """Convert a string or int into an attribute ref on the new table"""  # noqa
    if isinstance(column_ref, int):
        index = column_ref
    else:
        index = scheme.getPosition(column_ref)
    return raco.expression.UnnamedAttributeRef(index + offset)


def check_binop_compatability(op_name, left, right):
    """Check whether the arguments to an operation are compatible."""
    # Todo: check for type compatibility here?
    # https://github.com/uwescience/raco/issues/213
    if len(left.scheme()) != len(right.scheme()):
        raise SchemaMismatchException(op_name)


def check_assignment_compatability(before, after):
    """Check whether multiple assignments are compatible."""
    # TODO: check for exact schema match -- this is blocked by a general
    # cleanup of raco types.
    check_binop_compatability("assignment", before, after)


class ExpressionProcessor(object):
    """Convert syntactic expressions into relational algebra operations."""
    def __init__(self, symbols, catalog, use_dummy_schema=False):
        self.symbols = symbols
        self.catalog = catalog
        self.use_dummy_schema = use_dummy_schema

        # Variables accessed by the current operation
        self.uses_set = set()

    def get_and_clear_uses_set(self):
        """Retrieve the uses set and then clear its value."""
        try:
            return self.uses_set
        finally:
            self.uses_set = set()

    def evaluate(self, expr):
        method = getattr(self, expr[0].lower())
        return method(*expr[1:])

    def __lookup_symbol(self, _id):
        if _id not in self.symbols:
            raise NoSuchRelationException(_id)

        self.uses_set.add(_id)
        return copy.deepcopy(self.symbols[_id])

    def alias(self, _id):
        return self.__lookup_symbol(_id)

    def scan(self, rel_key):
        """Scan a database table."""
        assert isinstance(rel_key, relation_key.RelationKey)
        try:
            scheme = self.catalog.get_scheme(rel_key)
        except KeyError:
            if not self.use_dummy_schema:
                raise NoSuchRelationException(rel_key)

            # Create a dummy schema suitable for emitting plans
            scheme = raco.scheme.DummyScheme()

        return raco.algebra.Scan(rel_key, scheme,
                                 self.catalog.num_tuples(rel_key))

    def load(self, path, scheme, options):
        return raco.algebra.FileScan(path, scheme, options)

    def table(self, emit_clause):
        """Emit a single-row table literal."""
        emit_args = []
        for clause in emit_clause:
            emit_args.extend(clause.expand({}))

        from_args = collections.OrderedDict()
        from_args['$$SINGLETON$$'] = raco.algebra.SingletonRelation()

        # Add unbox relations to the from_args dictionary
        for name, sexpr in emit_args:
            self.extract_unbox_args(from_args, sexpr)

        op, info = multiway.merge(from_args)

        # rewrite clauses in terms of the new schema
        emit_args = [(name, multiway.rewrite_refs(sexpr, from_args, info))
                     for (name, sexpr) in emit_args]

        return raco.algebra.Apply(emitters=emit_args, input=op)

    @staticmethod
    def empty(_scheme):
        if not _scheme:
            _scheme = raco.scheme.Scheme()
        return raco.algebra.EmptyRelation(_scheme)

    def select(self, args):
        """Evaluate a select-from-where expression."""
        op = self.bagcomp(args.from_, args.where, args.select)
        if args.distinct:
            op = raco.algebra.Distinct(input=op)
        if args.limit is not None:
            op = raco.algebra.Limit(input=op, count=args.limit)

        return op

    def extract_unbox_args(self, from_args, sexpr):
        """Extract unbox arguments from a scalar expression.

        :param from_args: An ordered dictionary that maps from a
        relation alias (string) to an instance of raco.algebra.Operator.
        :param sexpr: A scalar expression (raco.expression.Expresssion)
        instance.
        """
        for sub_expr in sexpr.walk():
            if isinstance(sub_expr, raco.expression.Unbox):
                name = sub_expr.table_name
                assert isinstance(name, basestring)
                if name not in from_args:
                    from_args[name] = self.__lookup_symbol(name)

    def bagcomp(self, from_clause, where_clause, emit_clause):
        """Evaluate a bag comprehension.

        from_clause: A list of tuples of the form (id, expr).  expr can
        be None, which means "read the value from the symbol table".

        where_clause: An optional scalar expression (raco.expression).

        emit_clause: A list of EmitArg instances, each defining one or more
        output columns.
        """

        # Make sure no aliases were reused: [FROM X, X EMIT *] is illegal
        from_aliases = set([x[0] for x in from_clause])
        if len(from_aliases) != len(from_clause):
            raise DuplicateAliasException()

        # For each FROM argument, create a mapping from ID to operator
        # (id, raco.algebra.Operator)
        from_args = collections.OrderedDict()

        for _id, expr in from_clause:
            assert isinstance(_id, basestring)
            if expr:
                from_args[_id] = self.evaluate(expr)
            else:
                from_args[_id] = self.__lookup_symbol(_id)

        # Expand wildcards into a list of output columns
        assert emit_clause  # There should always be something to emit
        emit_args = []
        statemods = []
        for clause in emit_clause:
            emit_args.extend(clause.expand(from_args))
            statemods.extend(clause.get_statemods())

        orig_op, _info = multiway.merge(from_args)
        orig_schema_length = len(orig_op.scheme())

        # Add unbox relations to the from_args dictionary
        for name, sexpr in emit_args:
            self.extract_unbox_args(from_args, sexpr)
        if where_clause:
            self.extract_unbox_args(from_args, where_clause)

        # Create a single RA operation that is the cross of all targets
        op, info = multiway.merge(from_args)

        # HACK: calculate unboxed columns as implicit grouping columns,
        # so they can be used in grouping terms.
        new_schema_length = len(op.scheme())
        implicit_group_by_cols = range(orig_schema_length, new_schema_length)

        ################################################
        # Compile away unbox expressions in where, emit clauses
        ################################################

        if where_clause:
            where_clause = multiway.rewrite_refs(where_clause, from_args, info)
            # Extract the type of there where clause to force type safety
            # to be checked
            where_clause.typeof(op.scheme(), None)
            op = raco.algebra.Select(condition=where_clause, input=op)

        emit_args = [(name, multiway.rewrite_refs(sexpr, from_args, info))
                     for (name, sexpr) in emit_args]

        statemods = multiway.rewrite_statemods(statemods, from_args, info)

        if any(raco.expression.expression_contains_aggregate(ex)
               for name, ex in emit_args):
            return groupby.groupby(op, emit_args, implicit_group_by_cols,
                                   statemods)
        else:
            if statemods:
                return raco.algebra.StatefulApply(emit_args, statemods, op)
            if (len(from_args) == 1 and len(emit_clause) == 1 and
                isinstance(emit_clause[0],
                           (TableWildcardEmitArg, FullWildcardEmitArg))):
                return op
            return raco.algebra.Apply(emit_args, op)

    def distinct(self, expr):
        op = self.evaluate(expr)
        return raco.algebra.Distinct(input=op)

    def unionall(self, e1, e2):
        left = self.evaluate(e1)
        right = self.evaluate(e2)
        check_binop_compatability("unionall", left, right)
        return raco.algebra.UnionAll(left, right)

    def countall(self, expr):
        op = self.evaluate(expr)
        grouping_list = []
        agg_list = [raco.expression.COUNTALL()]
        return raco.algebra.GroupBy(grouping_list, agg_list, op)

    def intersect(self, e1, e2):
        left = self.evaluate(e1)
        right = self.evaluate(e2)
        check_binop_compatability("intersect", left, right)
        return raco.algebra.Intersection(left, right)

    def diff(self, e1, e2):
        left = self.evaluate(e1)
        right = self.evaluate(e2)
        check_binop_compatability("diff", left, right)
        return raco.algebra.Difference(left, right)

    def limit(self, expr, count):
        op = self.evaluate(expr)
        return raco.algebra.Limit(input=op, count=count)

    def cross(self, left_target, right_target):
        left = self.evaluate(left_target)
        right = self.evaluate(right_target)

        return raco.algebra.CrossProduct(left, right)

    def join(self, left_target, right_target):
        """Convert parser.JoinTarget arguments into a Join operation"""

        left = self.evaluate(left_target.expr)
        right = self.evaluate(right_target.expr)

        assert len(left_target.columns) == len(right_target.columns)

        left_scheme = left.scheme()
        left_refs = [get_unnamed_ref(c, left_scheme, 0)
                     for c in left_target.columns]

        right_scheme = right.scheme()
        right_refs = [get_unnamed_ref(c, right_scheme, len(left_scheme))
                      for c in right_target.columns]

        join_conditions = [raco.expression.EQ(x, y) for x, y in
                           zip(left_refs, right_refs)]

        # Merge the join conditions into a big AND expression

        def andify(x, y):
            """Merge two scalar expressions with an AND"""
            return raco.expression.AND(x, y)

        condition = reduce(andify, join_conditions)
        return raco.algebra.Join(condition, left, right)


class StatementProcessor(object):
    """Evaluate a list of statements"""

    def __init__(self, catalog, use_dummy_schema=False):
        # Map from identifiers (aliases) to raco.algebra.Operation instances
        self.symbols = {}

        assert isinstance(catalog, raco.catalog.Catalog)
        self.catalog = catalog
        self.ep = ExpressionProcessor(self.symbols, catalog, use_dummy_schema)

        self.cfg = ControlFlowGraph()

    def evaluate(self, statements):
        """Evaluate a list of statements"""
        for statement in statements:
            # Switch on the first tuple entry
            method = getattr(self, statement[0].lower())
            method(*statement[1:])

    def __evaluate_expr(self, expr, _def):
        """Evaluate an expression; add a node to the control flow graph.

        :param expr: An expression to evaluate
        :type expr: Myrial AST tuple
        :param _def: The variable defined by the expression, or None for
                     non-statements
        :type _def: string
        """

        op = self.ep.evaluate(expr)
        uses_set = self.ep.get_and_clear_uses_set()
        self.cfg.add_op(op, _def, uses_set)
        return op

    def __do_assignment(self, _id, expr):
        """Process an assignment statement; add a node to the control flow
        graph.

        :param _id: The target variable name.
        :type _id: string
        :param expr: The relational expression to evaluate
        :type expr: A Myrial expression AST node tuple
        """

        child_op = self.ep.evaluate(expr)
        if _id in self.symbols:
            check_assignment_compatability(child_op, self.symbols[_id])

        op = raco.algebra.StoreTemp(_id, child_op)
        uses_set = self.ep.get_and_clear_uses_set()
        self.cfg.add_op(op, _id, uses_set)

        # Point future references of this symbol to a scan of the materialized
        # table. Note that this assumes there is no scoping in Myrial.
        self.symbols[_id] = raco.algebra.ScanTemp(_id, child_op.scheme())

    def assign(self, _id, expr):
        """Map a variable to the value of an expression."""
        self.__do_assignment(_id, expr)

    def store(self, _id, rel_key, how_partitioned):
        assert isinstance(rel_key, relation_key.RelationKey)

        alias_expr = ("ALIAS", _id)
        child_op = self.ep.evaluate(alias_expr)

        if how_partitioned:
            scheme = child_op.scheme()
            col_list = [get_unnamed_ref(a, scheme) for a in how_partitioned]
            child_op = raco.algebra.Shuffle(child_op, col_list)
        op = raco.algebra.Store(rel_key, child_op)

        uses_set = self.ep.get_and_clear_uses_set()
        self.cfg.add_op(op, None, uses_set)

    def sink(self, _id):
        alias_expr = ("ALIAS", _id)
        child_op = self.ep.evaluate(alias_expr)
        op = raco.algebra.Sink(child_op)

        uses_set = self.ep.get_and_clear_uses_set()
        self.cfg.add_op(op, None, uses_set)

    def dump(self, _id):
        alias_expr = ("ALIAS", _id)
        child_op = self.ep.evaluate(alias_expr)
        op = raco.algebra.Dump(child_op)
        uses_set = self.ep.get_and_clear_uses_set()
        self.cfg.add_op(op, None, uses_set)

    def dowhile(self, statement_list, termination_ex):
        first_op_id = self.cfg.next_op_id  # op ID of the top of the loop

        for _type, _id, expr in statement_list:
            if _type != 'ASSIGN':
                # TODO: Better error message
                raise InvalidStatementException('%s not allowed in do/while' %
                                                _type.lower())
            self.__do_assignment(_id, expr)

        last_op_id = self.cfg.next_op_id

        self.__evaluate_expr(termination_ex, None)

        # Add a control flow edge from the loop condition to the top of the
        # loop
        self.cfg.add_edge(last_op_id, first_op_id)

    def get_logical_plan(self, **kwargs):
        """Return an operator representing the logical query plan."""
        return self.cfg.get_logical_plan(
            dead_code_elimination=kwargs.get('dead_code_elimination', True),
            apply_chaining=kwargs.get('apply_chaining', True))

    def __get_physical_plan_for__(self, target_phys_algebra, **kwargs):
        logical_plan = self.get_logical_plan(**kwargs)

        kwargs['target'] = target_phys_algebra
        return optimize(logical_plan, **kwargs)

    def get_physical_plan(self, **kwargs):
        """Return an operator representing the physical query plan."""
        target_phys_algebra = kwargs.get('target_alg')
        if target_phys_algebra is None:
            if kwargs.get('multiway_join', False):
                target_phys_algebra = MyriaHyperCubeAlgebra(self.catalog)
            else:
                target_phys_algebra = MyriaLeftDeepTreeAlgebra()

        return self.__get_physical_plan_for__(target_phys_algebra, **kwargs)

    def get_json(self, **kwargs):
        lp = self.get_logical_plan()
        pps = self.get_physical_plan(**kwargs)
        # TODO This is not correct. The first argument is the raw query string,
        # not the string representation of the logical plan
        return compile_to_json(str(lp), pps, pps, "myrial")