odps/df/backends/sqlalchemy/compiler.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2022 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import itertools
from datetime import timedelta

from ..core import Backend
from ..errors import CompileError
from ...expr.reduction import *
from ...expr.arithmetic import *
from ...expr.datetimes import *
from ...expr.window import *
from ...expr.merge import *
from ...expr.utils import highest_precedence_data_type
from ... import types as df_types
from ...utils import is_constant_scalar, traverse_until_source, is_source_collection
from ....compat import lzip
from .... import utils
from . import types

try:
    from sqlalchemy import Table as SATable, select, func, case, \
        extract, desc, distinct, literal, join, union, union_all
    from sqlalchemy.sql.expression import Alias, text

    # define the compile ext
    from .ext import *

    has_sqlalchemy = True
except ImportError:
    has_sqlalchemy = False


BINARY_OP = {
    'Add': operator.add,
    'Substract': operator.sub,
    'Multiply': operator.mul,
    'Divide': operator.div if six.PY2 else operator.truediv,
    'Mod': operator.mod,
    'FloorDivide': operator.floordiv,
    'Greater': operator.gt,
    'GreaterEqual': operator.ge,
    'Less': operator.lt,
    'LessEqual': operator.le,
    'Equal': operator.eq,
    'NotEqual': operator.ne,
    'And': operator.and_,
    'Or': operator.or_
}

UNARY_OP = {
    'Negate': operator.neg,
    'Invert': operator.inv,
}

DATE_KEY_DIC = {
    'Day': 'days',
    'Hour': 'hours',
    'Minute': 'minutes',
    'Second': 'seconds',
    'MilliSecond': 'milliseconds',
    'MicroSecond': 'microseconds',
}

MATH_COMPILE_DIC = {
    'Abs': 'abs',
    'Sqrt': 'sqrt',
    'Sin': 'sin',
    'Cos': 'cos',
    'Tan': 'tan',
    'Exp': 'exp',
    'Arccos': 'acos',
    'Arcsin': 'asin',
    'Arctan': 'atan',
    'Ceil': 'ceil',
    'Floor': 'floor'
}

WINDOW_COMPILE_DIC = {
    'CumSum': 'sum',
    'CumMean': 'avg',
    'CumStd': 'stddev',
    'CumMax': 'max',
    'CumMin': 'min',
    'CumCount': 'count',
    'Lag': 'lag',
    'Lead': 'lead',
    'Rank': 'rank',
    'DenseRank': 'dense_rank',
    'PercentRank': 'percent_rank',
    'RowNumber': 'row_number'
}

DATE_PARTS_DIC = {
    'Year': 'year',
    'Month': 'month',
    'Day': 'day',
    'Hour': 'hour',
    'Minute': 'minute',
    'Second': 'second',
    'WeekOfYear': 'week',
    'DayOfYear': 'doy',
    'MicroSecond': 'microseconds',
    'UnixTimestamp': 'epoch',
}


class SQLAlchemyCompiler(Backend):
    def __init__(self, expr_dag):
        self._expr_dag = expr_dag
        self._expr_to_sqlalchemy = ExprDictionary()
        self._id_gen = itertools.count(1)

        self._sa_engine = None

    def _new_alias(self):
        return 't%s' % next(self._id_gen)

    def compile(self, expr, traversed=None):
        if traversed is None:
            traversed = set()

        for node in traverse_until_source(expr):
            if id(node) not in traversed:
                node.accept(self)
                traversed.add(id(node))

        sa_expr = self._expr_to_sqlalchemy[self._expr_dag.root]
        if is_source_collection(expr):
            sa_expr = select([sa_expr])
        return sa_expr

    def _add(self, expr, op):
        self._expr_to_sqlalchemy[expr] = op

    def _gen_select_columns(self, fields):
        sa_exprs = []
        for field in fields:
            if not isinstance(field, Scalar) or field._value is None:
                sa_exprs.append(self._expr_to_sqlalchemy[field].label(field.name))
            else:
                sa_exprs.append(
                    literal(self._expr_to_sqlalchemy[field]).label(field.name))
        return sa_exprs

    def visit_source_collection(self, expr):
        table = next(expr.data_source())

        if not isinstance(table, SATable):
            raise ValueError('Source data must be a sqlalchemy table')

        if table.bind and self._sa_engine is None:
            self._sa_engine = table.bind

        self._add(expr, table.alias(self._new_alias()))

    def visit_project_collection(self, expr):
        selects = select(self._gen_select_columns(expr._fields))\
            .select_from(self._expr_to_sqlalchemy[expr.input])

        if expr is not self._expr_dag.root:
            selects = selects.alias(self._new_alias())
        self._add(expr, selects)

    def visit_apply_collection(self, expr):
        raise NotImplementedError

    def visit_filter_collection(self, expr):
        input = self._expr_to_sqlalchemy[expr.input]
        predicate = self._expr_to_sqlalchemy[expr._predicate]
        filtered = input.select(predicate)

        if expr is not self._expr_dag.root:
            filtered = filtered.alias(self._new_alias())
        self._add(expr, filtered)

    def visit_slice_collection(self, expr):
        input = self._expr_to_sqlalchemy[expr.input]
        sliced = expr._indexes
        if sliced[2] is not None:
            raise NotImplementedError
        if sliced[0] is not None and sliced[0].value < 0:
            raise CompileError('start number must be greater than 0')
        if sliced[1] is not None and sliced[1].value <= 0:
            raise CompileError('end number must be greater than 0')

        kw = dict()
        if sliced[0] is not None and sliced[0].value > 0:
            kw['offset'] = sliced[0].value
        if sliced[1] is not None and sliced[1].value > 0:
            kw['limit'] = sliced[1].value
        input = input.select(**kw)

        if expr is not self._expr_dag.root:
            input = input.alias(self._new_alias())

        self._add(expr, input)

    def visit_element_op(self, expr):
        input = self._expr_to_sqlalchemy.get(expr.input)
        if isinstance(expr, element.IsNull):
            sa_expr = input.is_(None)
        elif isinstance(expr, element.NotNull):
            sa_expr = input.isnot(None)
        elif isinstance(expr, element.FillNa):
            sa_expr = case([(input.is_(None), expr.fill_value)], else_=input)
        elif isinstance(expr, (element.IsIn, element.NotIn)):
            op = input.in_ if isinstance(expr, element.IsIn) else input.notin_
            if expr._values is None:
                sa_expr = op([None])
            elif len(expr._values) == 1 and isinstance(expr._values[0], SequenceExpr):
                right = select([self._expr_to_sqlalchemy[expr._values[0]]])
                sa_expr = op(right)
            else:
                sa_expr = op(tuple(self._expr_to_sqlalchemy[it] for it in expr._values))
        elif isinstance(expr, element.Between):
            if not expr.inclusive:
                raise NotImplementedError
            sa_expr = input.between(
                self._expr_to_sqlalchemy[expr._left],
                self._expr_to_sqlalchemy[expr._right]
            )
        elif isinstance(expr, element.IfElse):
            sa_expr = case([(input, self._expr_to_sqlalchemy[expr._then])],
                           else_=self._expr_to_sqlalchemy[expr._else])
        elif isinstance(expr, element.Switch):
            conditions = [self._expr_to_sqlalchemy[cond] for cond in expr._conditions]
            thens = [self._expr_to_sqlalchemy[then] for then in expr._thens]
            sa_else = self._expr_to_sqlalchemy[expr._default] \
                if expr._default is not None else expr._default
            if expr._input is None:
                sa_expr = case(lzip(conditions, thens), else_=sa_else)
            else:
                sa_expr = case(dict(lzip(conditions, thens)),
                               value=input, else_=sa_else)
        else:
            raise NotImplementedError

        self._add(expr, sa_expr)

    def visit_binary_op(self, expr):
        if isinstance(expr, Power):
            op = func.pow
        elif isinstance(expr, FloorDivide):
            op = operator.div if six.PY2 else operator.truediv
        elif isinstance(expr, (Add, Substract)) and expr.dtype == df_types.datetime:
            if isinstance(expr, Add) and \
                    all(child.dtype == df_types.datetime for child in (expr.lhs, expr.rhs)):
                raise CompileError('Cannot add two datetimes')
            if isinstance(expr.rhs, DTScalar) or (isinstance(expr, Add) and expr.lhs, DTScalar):
                if isinstance(expr.rhs, DTScalar):
                    dt, scalar = expr.lhs, expr.rhs
                else:
                    dt, scalar = expr.rhs, expr.lhs
                val = scalar.value
                if isinstance(expr, Substract):
                    val = -val

                dt_type = type(scalar).__name__[:-6]
                sa_dt = self._expr_to_sqlalchemy[dt]
                try:
                    key = DATE_KEY_DIC[dt_type]
                except KeyError:
                    raise NotImplementedError
                if self._sa_engine and self._sa_engine.name == 'mysql':
                    if dt_type == 'MilliSecond':
                        val, dt_type = val * 1000, 'MicroSecond'
                    sa_expr = func.date_add(sa_dt, text('interval %d %s' % (val, dt_type.lower())))
                else:
                    sa_expr = sa_dt + timedelta(**{key: val})
                self._add(expr, sa_expr)
                return
            else:
                raise NotImplementedError
        elif isinstance(expr, Substract) and expr._lhs.dtype == df_types.datetime and \
                expr._rhs.dtype == df_types.datetime:
            sa_expr = self._expr_to_sqlalchemy[expr._lhs] - self._expr_to_sqlalchemy[expr._rhs]
            if self._sa_engine and self._sa_engine.name == 'mysql':
                sa_expr = func.abs(func.microsecond(sa_expr)
                                   .cast(types.df_type_to_sqlalchemy_type(expr.dtype))) / 1000
            else:
                sa_expr = func.abs(extract('MICROSECONDS', sa_expr)
                                   .cast(types.df_type_to_sqlalchemy_type(expr.dtype))) / 1000
            self._add(expr, sa_expr)
            return
        elif isinstance(expr, Mod):
            lhs, rhs = self._expr_to_sqlalchemy[expr._lhs], self._expr_to_sqlalchemy[expr._rhs]
            sa_expr = BINARY_OP[expr.node_name](lhs, rhs)
            if not is_constant_scalar(expr._rhs):
                sa_expr = case([(rhs > 0, func.abs(sa_expr))], else_=sa_expr)
            elif expr._rhs.value > 0:
                sa_expr = func.abs(sa_expr)
            self._add(expr, sa_expr)
            return
        else:
            op = BINARY_OP[expr.node_name]
        lhs, rhs = self._expr_to_sqlalchemy[expr._lhs], self._expr_to_sqlalchemy[expr._rhs]
        sa_expr = op(lhs, rhs)
        self._add(expr, sa_expr)

    def visit_unary_op(self, expr):
        if isinstance(expr, Abs):
            op = func.abs
        else:
            op = UNARY_OP[expr.node_name]
        self._add(expr, op(self._expr_to_sqlalchemy[expr._input]))

    def visit_math(self, expr):
        try:
            op = getattr(func, MATH_COMPILE_DIC[expr.node_name])
            sa_expr = op(self._expr_to_sqlalchemy[expr._input])
        except KeyError:
            if expr.node_name == 'Log':
                if expr._base is not None:
                    sa_expr = SALog('log', self._expr_to_sqlalchemy[expr._base],
                                    self._expr_to_sqlalchemy[expr._base],
                                    self._expr_to_sqlalchemy[expr._input])
                else:
                    sa_expr = SALog('log', None, self._expr_to_sqlalchemy[expr._input])
            elif expr.node_name == 'Log2':
                sa_expr = SALog('log', 2, 2, self._expr_to_sqlalchemy[expr._input])
                sa_expr = sa_expr.cast(types.df_type_to_sqlalchemy_type(expr.dtype))
            elif expr.node_name == 'Log10':
                sa_expr = SALog('log', 10, 10, self._expr_to_sqlalchemy[expr._input])
                sa_expr = sa_expr.cast(types.df_type_to_sqlalchemy_type(expr.dtype))
            elif expr.node_name == 'Trunc':
                input = self._expr_to_sqlalchemy[expr._input]
                decimals = 0 if expr._decimals is None else self._expr_to_sqlalchemy[expr._decimals]
                sa_expr = SATruncate('trunc', input, decimals)
            elif expr.node_name == 'Round':
                decimals = 0 if expr._decimals is None else self._expr_to_sqlalchemy[expr._decimals]
                sa_expr = func.round(self._expr_to_sqlalchemy[expr._input], decimals)
            else:
                raise NotImplementedError

        self._add(expr, sa_expr)

    def visit_string_op(self, expr):
        if isinstance(expr, strings.Capitalize):
            input = self._expr_to_sqlalchemy[expr._input]
            tp = types.df_type_to_sqlalchemy_type(expr.dtype)
            sa_expr = func.upper(func.substr(input, 1, 1)).cast(tp) + \
                      func.lower(func.substr(input, 2)).cast(tp)
        elif isinstance(expr, strings.Contains) and not expr.regex:
            sa_expr = self._expr_to_sqlalchemy[expr._input].contains(
                self._expr_to_sqlalchemy[expr._pat])
        elif isinstance(expr, strings.Endswith):
            sa_expr = self._expr_to_sqlalchemy[expr._input].endswith(
                self._expr_to_sqlalchemy[expr._pat])
        elif isinstance(expr, strings.Startswith):
            sa_expr = self._expr_to_sqlalchemy[expr._input].startswith(
                self._expr_to_sqlalchemy[expr._pat])
        elif isinstance(expr, strings.Replace) and not expr.regex:
            sa_expr = func.replace(self._expr_to_sqlalchemy[expr._input],
                                   self._expr_to_sqlalchemy[expr._pat],
                                   self._expr_to_sqlalchemy[expr._repl])
        elif isinstance(expr, strings.Get):
            sa_expr = func.substr(self._expr_to_sqlalchemy[expr._input],
                                  self._expr_to_sqlalchemy[expr._index] + 1, 1)
        elif isinstance(expr, strings.Len):
            sa_expr = func.length(self._expr_to_sqlalchemy[expr._input])
        elif isinstance(expr, (strings.Ljust, strings.Rjust, strings.Pad)):
            if isinstance(expr, strings.Pad):
                if expr.side == 'both':
                    raise NotImplementedError
                op = func.lpad if expr.side == 'left' else func.rpad
            else:
                op = func.lpad if isinstance(expr, strings.Ljust) else func.rpad
            sa_expr = op(self._expr_to_sqlalchemy[expr._input],
                         self._expr_to_sqlalchemy[expr._width],
                         self._expr_to_sqlalchemy[expr._fillchar])
        elif isinstance(expr, (strings.Lower, strings.Upper)):
            op = func.lower if isinstance(expr, strings.Lower) else func.upper
            sa_expr = op(self._expr_to_sqlalchemy[expr._input])
        elif isinstance(expr, (strings.Lstrip, strings.Rstrip, strings.Strip)):
            if expr._to_strip is None:
                raise NotImplementedError
            op = func.ltrim if isinstance(expr, strings.Lstrip) else (
                func.rtrim if isinstance(expr, strings.Rstrip) else func.btrim
            )
            sa_expr = op(self._expr_to_sqlalchemy[expr._input],
                         self._expr_to_sqlalchemy[expr._to_strip])
        elif isinstance(expr, strings.Repeat):
            sa_expr = func.repeat(self._expr_to_sqlalchemy[expr._input],
                                  self._expr_to_sqlalchemy[expr._repeats])
        elif isinstance(expr, strings.Slice):
            if expr.end is None and expr.step is None:
                sa_expr = func.substr(self._expr_to_sqlalchemy[expr._input],
                                      self._expr_to_sqlalchemy[expr._start] + 1)
            elif isinstance(expr.start, six.integer_types) and \
                    isinstance(expr.end, six.integer_types) and \
                    expr.step is None and expr.start > 0 and expr.end > 0:
                length = expr.end - expr.start
                sa_expr = func.substr(self._expr_to_sqlalchemy[expr._input],
                                      expr.start + 1, length)
            else:
                raise NotImplementedError
        elif isinstance(expr, strings.Title):
            sa_expr = func.initcap(self._expr_to_sqlalchemy[expr._input])
        else:
            raise NotImplementedError

        self._add(expr, sa_expr)

    def visit_datetime_op(self, expr):
        class_name = type(expr).__name__
        input = self._expr_to_sqlalchemy[expr._input]

        if class_name in DATE_PARTS_DIC:
            if self._sa_engine and self._sa_engine.name == 'mysql':
                if class_name == 'UnixTimestamp':
                    fun = func.unix_timestamp
                else:
                    fun = getattr(func, class_name.lower())
                sa_expr = fun(input).cast(types.df_type_to_sqlalchemy_type(expr.dtype))
            else:
                sa_expr = func.date_part(DATE_PARTS_DIC[class_name], input)\
                    .cast(types.df_type_to_sqlalchemy_type(expr.dtype))
        elif isinstance(expr, Date):
            if self._sa_engine and self._sa_engine.name == 'mysql':
                sa_expr = func.date(input).cast(types.df_type_to_sqlalchemy_type(expr.dtype))
            else:
                sa_expr = func.date_trunc('day', input)
        elif isinstance(expr, WeekDay):
            if self._sa_engine and self._sa_engine.name == 'mysql':
                sa_expr = (func.dayofweek(input).cast(types.df_type_to_sqlalchemy_type(expr.dtype)) + 5) % 7
            else:
                sa_expr = (func.date_part('dow', input).cast(types.df_type_to_sqlalchemy_type(expr.dtype)) + 6) % 7
        else:
            raise NotImplementedError

        self._add(expr, sa_expr)

    def visit_groupby(self, expr):
        bys, having, aggs, fields = tuple(expr.args[1:])
        if fields is None:
            fields = bys + aggs

        selects = select(self._gen_select_columns(fields))
        if len(fields) == 1 and isinstance(fields[0], (Count, GroupedCount)):
            selects = selects.select_from(self._expr_to_sqlalchemy[fields[0].input])
        grouped = selects.group_by(*self._gen_select_columns(bys))
        if having:
            grouped = grouped.having(self._expr_to_sqlalchemy[having])

        if expr is not self._expr_dag.root:
            grouped = grouped.alias(self._new_alias())

        self._add(expr, grouped)

    def visit_mutate(self, expr):
        bys, mutates, fields = tuple(expr.args[1:])
        if fields is None:
            fields = bys + mutates

        selects = select(self._gen_select_columns(fields))
        if expr is not self._expr_dag.root:
            selects = selects.alias(self._new_alias())

        self._add(expr, selects)

    def visit_sort_column(self, expr):
        if isinstance(expr.input, CollectionExpr):
            sa_expr = self._expr_to_sqlalchemy[expr.input].c[expr.source_name]
        else:
            sa_expr = self._expr_to_sqlalchemy[expr.input]
        if not expr._ascending:
            sa_expr = desc(sa_expr)

        self._add(expr, sa_expr)

    def visit_sort(self, expr):
        input = self._expr_to_sqlalchemy[expr.input]
        sa_expr = input.select(order_by=[self._expr_to_sqlalchemy[e]
                                         for e in expr._sorted_fields])
        if expr is not self._expr_dag.root:
            sa_expr = sa_expr.alias(self._new_alias())
        self._add(expr, sa_expr)

    def visit_distinct(self, expr):
        sa_expr = select(self._gen_select_columns(expr._unique_fields), distinct=True)

        if expr is not self._expr_dag.root:
            sa_expr = sa_expr.alias(self._new_alias())
        self._add(expr, sa_expr)

    def visit_column(self, expr):
        table = self._expr_to_sqlalchemy[expr.input]
        col = table.c[expr.source_name]

        if expr._source_data_type != expr._data_type:
            col = col.cast(types.df_type_to_sqlalchemy_type(expr._data_type))

        self._add(expr, col)

    def visit_reduction(self, expr):
        if getattr(expr, '_unique', False):
            raise NotImplementedError

        input = self._expr_to_sqlalchemy[expr.input]

        # TODO: MEDIAN does not support
        if isinstance(expr, (Max, GroupedMax)):
            f = func.max
        elif isinstance(expr, (Min, GroupedMin)):
            f = func.min
        elif isinstance(expr, (Count, GroupedCount)):
            f = func.count
        elif isinstance(expr, (Sum, GroupedSum)):
            f = func.sum
        elif isinstance(expr, (Var, GroupedVar)) and expr._ddof in (0, 1):
            f = func.var_pop if expr._ddof == 0 else func.var_samp
        elif isinstance(expr, (Std, GroupedStd)) and expr._ddof in (0, 1):
            f = func.stddev_pop if expr._ddof == 0 else func.stddev_samp
        elif isinstance(expr, (Mean, GroupedMean)):
            f = func.avg
        elif isinstance(expr, (NUnique, GroupedNUnique)):
            f = lambda *x: func.count(distinct(*x))
        elif isinstance(expr, (Cat, GroupedCat)):
            f = lambda x: func.array_to_string(func.array_agg(x),
                                               self._expr_to_sqlalchemy[expr._sep])
        else:
            raise NotImplementedError

        if isinstance(expr, (Count, GroupedCount)) and \
                isinstance(expr.input, CollectionExpr):
            reduced = f()
        elif isinstance(expr, (NUnique, GroupedNUnique)):
            if len(expr.inputs) > 1:
                raise NotImplementedError
            reduced = f(*(self._expr_to_sqlalchemy[i] for i in expr.inputs))
        else:
            reduced = f(input)
        self._add(expr, reduced)

    def visit_cum_window(self, expr):
        input = self._expr_to_sqlalchemy[expr._input]
        if expr._distinct.value is True:
            raise NotImplementedError
        try:
            func_name = WINDOW_COMPILE_DIC[expr.node_name]
        except KeyError:
            raise NotImplementedError
        f = getattr(func, func_name)
        partition_by = self._gen_select_columns(expr._partition_by) \
            if expr._partition_by else None
        order_by = self._gen_select_columns(expr._order_by) \
            if expr._order_by else None
        rows = (self._expr_to_sqlalchemy[expr._preceding] if expr._preceding else None,
                self._expr_to_sqlalchemy[expr._following] if expr._following else None)
        rows = None if all(r is None for r in rows) else rows

        sa_expr = f(input).over(partition_by=partition_by, order_by=order_by, rows=rows)
        self._add(expr, sa_expr)

    def visit_rank_window(self, expr):
        try:
            func_name = WINDOW_COMPILE_DIC[expr.node_name]
        except KeyError:
            raise NotImplementedError
        f = getattr(func, func_name)
        partition_by = self._gen_select_columns(expr._partition_by) \
            if expr._partition_by else None
        order_by = self._gen_select_columns(expr._order_by) \
            if expr._order_by else None

        sa_expr = f().over(partition_by=partition_by, order_by=order_by)
        if isinstance(expr, PercentRank):
            sa_expr = sa_expr.cast(types.df_type_to_sqlalchemy_type(expr.dtype))
        self._add(expr, sa_expr)

    def visit_shift_window(self, expr):
        input = self._expr_to_sqlalchemy[expr._input]
        try:
            func_name = WINDOW_COMPILE_DIC[expr.node_name]
        except KeyError:
            raise NotImplementedError
        f = getattr(func, func_name)
        partition_by = self._gen_select_columns(expr._partition_by) \
            if expr._partition_by else None
        order_by = self._gen_select_columns(expr._order_by) \
            if expr._order_by else None

        args = (input, self._expr_to_sqlalchemy[expr._offset])
        if expr._default:
            args += (literal(self._expr_to_sqlalchemy[expr._default]).cast(
                types.df_type_to_sqlalchemy_type(expr._input.dtype)),)

        sa_expr = f(*args).over(partition_by=partition_by, order_by=order_by)
        self._add(expr, sa_expr)

    def visit_scalar(self, expr):
        if expr._value is not None:
            if expr.dtype == df_types.string:
                val = utils.to_str(expr.value) \
                    if isinstance(expr.value, six.text_type) else expr.value
                self._add(expr, val)
                return
            else:
                self._add(expr, expr._value)
        else:
            self._add(expr, None)

    def visit_cast(self, expr):
        to_type = types.df_type_to_sqlalchemy_type(expr.dtype)
        self._add(expr, self._expr_to_sqlalchemy[expr.input].cast(to_type))

    def visit_join(self, expr):
        lhs, rhs = self._expr_to_sqlalchemy[expr._lhs], self._expr_to_sqlalchemy[expr._rhs]
        if isinstance(expr, RightJoin):
            lhs, rhs = rhs, lhs
        on = self._expr_to_sqlalchemy[expr._predicate]
        kw = dict()
        if isinstance(expr, OuterJoin):
            kw['full'] = True
        elif isinstance(expr, (LeftJoin, RightJoin)):
            kw['isouter'] = True
        joined = join(lhs, rhs, onclause=on, **kw)

        self._add(expr, joined)

    def visit_union(self, expr):
        lhs, rhs = self._expr_to_sqlalchemy[expr._lhs], self._expr_to_sqlalchemy[expr._rhs]
        if is_source_collection(expr._lhs):
            lhs = select([lhs])
        elif isinstance(lhs, Alias):
            lhs = lhs.element
        if is_source_collection(expr._rhs):
            rhs = select([rhs])
        elif isinstance(rhs, Alias):
            rhs = rhs.element
        method = union if expr._distinct else union_all
        unioned = method(lhs, rhs)

        if expr is not self._expr_dag.root:
            unioned = unioned.alias(self._new_alias())

        self._add(expr, unioned)