#!/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
import json
import re
import time
import uuid
import warnings
from datetime import datetime

from ..core import Backend
from ...expr.expressions import *
from ...expr.arithmetic import Power
from ...expr.reduction import GroupedSequenceReduction, GroupedCount, Count, \
    GroupedCat, Cat, NUnique, GroupedNUnique, ToList, GroupedToList, Quantile, \
    GroupedQuantile
from ...expr.merge import JoinCollectionExpr
from ...expr.datetimes import DTScalar
from ...expr.collections import PivotCollectionExpr
from ...expr import arithmetic, element, composites
from ...utils import traverse_until_source
from ....dag import DAG
from ..errors import CompileError
from ..utils import refresh_dynamic
from . import types
from ... import types as df_types
from ....models import FileResource, TableResource, TableSchema
from .... import compat
from ....lib import cloudpickle
from ....lib.xnamedtuple import xnamedtuple
from ....compat import lzip, Version

try:
    import numpy as np
    import pandas as pd
except ImportError:
    pd = None
    np = None


if pd is not None:
    PD_APPLY_HAS_RESULT_TYPE = Version(pd.__version__) >= Version('0.23.0')
else:
    PD_APPLY_HAS_RESULT_TYPE = False


BINARY_OP_TO_PANDAS = {
    'Add': operator.add,
    'Substract': operator.sub,
    'Multiply': operator.mul,
    'Divide': operator.div if six.PY2 else operator.truediv,
    'Mod': operator.mod,
    'FloorDivide': operator.floordiv,
    'Power': operator.pow,
    '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_TO_PANDAS = {
    'Negate': operator.neg,
    'Invert': operator.inv,
    'Abs': operator.abs
}

if pd:
    SORT_CUM_WINDOW_OP_TO_PANDAS = {
        'CumSum': lambda s: s.expanding(min_periods=1).sum(),
        'CumMean': lambda s: s.expanding(min_periods=1).mean(),
        'CumMedian': lambda s: s.expanding(min_periods=1).median(),
        'CumStd': lambda s: s.expanding(min_periods=1).std(),
        'CumMin': lambda s: s.expanding(min_periods=1).min(),
        'CumMax': lambda s: s.expanding(min_periods=1).max(),
        'CumCount': lambda s: s.expanding(min_periods=1).count(),
    }

if np:
    CUM_WINDOW_OP_TO_PANDAS = {
        'CumSum': np.sum,
        'CumMean': np.mean,
        'CumMedian': np.median,
        'CumStd': np.std,
        'CumMin': np.min,
        'CumMax': np.max,
        'CumCount': lambda x: len(x),
    }

JOIN_DICT = {
    'INNER': 'inner',
    'LEFT OUTER': 'left',
    'RIGHT OUTER': 'right',
    'FULL OUTER': 'outer'
}


def _explode(obj):
    if obj and isinstance(obj, tuple):
        obj = obj[0]
    if obj is None:
        return
    if isinstance(obj, dict):
        for k, v in six.iteritems(obj):
            yield k, v
    else:
        for v in obj:
            yield v


def _pos_explode(obj):
    if obj and isinstance(obj, tuple):
        obj = obj[0]
    if obj is None:
        return
    for idx, v in enumerate(obj):
        yield idx, v


def _filter_none(col):
    import numpy as np
    if hasattr(col, 'dropna'):
        col = col.dropna()
    else:
        try:
            col = col[~np.isnan(col)]
        except TypeError:
            col = col[np.fromiter((v is not None for v in col), np.bool_)]
    return col


BUILTIN_FUNCS = {
    'EXPLODE': _explode,
    'POSEXPLODE': _pos_explode,
}


class PandasCompiler(Backend):
    """
    PandasCompiler will compile an Expr into a DAG
    in which each node is a pair of <expr, function>.
    """

    def __init__(self, expr_dag):
        self._dag = DAG()
        self._expr_to_dag_node = dict()
        self._expr_dag = expr_dag
        self._callbacks = list()

    def compile(self, expr):
        try:
            return self._compile(expr)
        finally:
            self._cleanup()

    def _cleanup(self):
        for callback in self._callbacks:
            callback()
        self._callbacks = list()

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

        root = self._retrieve_until_find_root(expr)

        if root is not None and id(root) not in traversed:
            self._compile_join_node(root, traversed)
            traversed.add(id(root))

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

        return self._dag

    def _compile_join_node(self, expr, traversed):
        nodes = []

        self._compile(expr._lhs, traversed)
        nodes.append(expr._lhs)
        self._compile(expr._rhs, traversed)
        nodes.append(expr._rhs)
        for node in expr._predicate:
            nodes.append(node._lhs)
            self._compile(node._lhs, traversed)
            nodes.append(node._rhs)
            self._compile(node._rhs, traversed)

        expr.accept(self)
        for node in nodes:
            self._dag.add_edge(self._expr_to_dag_node[node], self._expr_to_dag_node[expr])

        cached_args = expr.args

        def cb():
            for arg_name, arg in zip(expr._args, cached_args):
                setattr(expr, arg_name, arg)
        self._callbacks.append(cb)

        for arg_name in expr._args:
            setattr(expr, arg_name, None)

    @classmethod
    def _retrieve_until_find_root(cls, expr):
        for node in traverse_until_source(expr, top_down=True, unique=True):
            if isinstance(node, JoinCollectionExpr):
                return node

    def _add_node(self, expr, handle):
        children = expr.children()

        node = (expr, handle)
        self._dag.add_node(node)
        self._expr_to_dag_node[expr] = node

        # the dependencies do not exist in self._expr_to_dag_node
        predecessors = [self._expr_to_dag_node[child] for child in children
                        if child in self._expr_to_dag_node]
        [self._dag.add_edge(p, node) for p in predecessors]

    @staticmethod
    def _attempt_pickle_unpickle(*args):
        try:
            cloudpickle.loads(cloudpickle.dumps(args))
        except Exception as ex:  # pragma: no cover
            warnings.warn(
                "Failed serializing user-defined function %r. Might get errors "
                "when running remotely. Error message: %s" % (args, ex),
                RuntimeWarning,
            )

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

        if not isinstance(df, pd.DataFrame):
            raise ValueError('Expr data must be a pandas DataFrame.')

        def handle(_):
            copy_data = len(self._expr_dag.descendants(expr)) > 0
            if not copy_data and list(df.columns) == expr.schema.names:
                # no need to copy when there is no descendants for the dataframe
                return df
            else:
                # make a copy to avoid modify
                return df.rename(columns=dict(zip(df.columns, expr.schema.names)))

        self._add_node(expr, handle)

    @classmethod
    def _get_children_vals(cls, kw, expr=None, children=None):
        children = children or expr.children()
        return [kw.get(child) for child in children]

    @classmethod
    def _merge_values(cls, exprs, kw):
        fields = [kw.get(expr) for expr in exprs]
        size = max(len(f) for f, e in zip(fields, exprs) if isinstance(e, SequenceExpr))
        fields = [pd.Series([f] * size) if isinstance(e, Scalar) else f
                  for f, e in zip(fields, exprs)]

        return pd.concat(fields, axis=1, keys=[e.name for e in exprs])

    def visit_project_collection(self, expr):
        def handle(kw):
            children = expr.children()

            fields = self._get_children_vals(kw, children=children)[1:]
            names = expr.schema.names

            if isinstance(expr, Summary):
                size = 1
            else:
                size = max(len(f) for f, e in zip(fields, expr._fields)
                           if isinstance(e, SequenceExpr))
            for i in range(len(fields)):
                if not isinstance(fields[i], pd.Series):
                    fields[i] = pd.Series([fields[i]] * size)

            return pd.concat(fields, axis=1, keys=names)

        self._add_node(expr, handle)

    def visit_filter_partition_collection(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)
            df, predicate = children_vals[0:1]
            return df[predicate][expr.schema.names]

        self._add_node(expr, handle)

    def visit_filter_collection(self, expr):
        def handle(kw):
            df, predicate = tuple(self._get_children_vals(kw, expr))
            return df[predicate]

        self._add_node(expr, handle)

    def visit_slice_collection(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)

            df = children_vals[0]
            start, end, step = expr.start, expr.stop, expr.step

            return df[start: end: step]

        self._add_node(expr, handle)

    def visit_element_op(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)
            input, args = children_vals[0], children_vals[1:]

            if isinstance(expr.input, Scalar):
                input = pd.Series([input])

            def run():
                if isinstance(expr, element.IsNull):
                    return input.isnull()
                elif isinstance(expr, element.NotNull):
                    return input.notnull()
                elif isinstance(expr, element.IsNa):
                    return input.isna()
                elif isinstance(expr, element.NotNa):
                    return input.notna()
                elif isinstance(expr, element.FillNa):
                    return input.fillna(args[0])
                elif isinstance(expr, element.IsIn):
                    if isinstance(expr._values[0], SequenceExpr):
                        return input.isin(list(args[0]))
                    else:
                        return input.isin(args)
                elif isinstance(expr, element.NotIn):
                    if isinstance(expr._values[0], SequenceExpr):
                        return ~input.isin(list(args[0]))
                    else:
                        return ~input.isin(args)
                elif isinstance(expr, element.IfElse):
                    return pd.Series(np.where(input, args[0], args[1]), name=expr.name, index=input.index)
                elif isinstance(expr, element.Switch):
                    case = None if expr.case is None else kw.get(expr.case)
                    default = None if expr.default is None else kw.get(expr.default)

                    conditions = [kw.get(it) for it in expr.conditions]
                    thens = [kw.get(it) for it in expr.thens]
                    if case is not None:
                        conditions = [case == condition for condition in conditions]
                        condition_exprs = [expr.case == cond for cond in expr.conditions]
                    else:
                        condition_exprs = expr.conditions

                    size = max(len(val) for e, val in zip(condition_exprs + expr.thens, conditions + thens)
                               if isinstance(e, SequenceExpr))

                    curr = pd.Series([None] * size)
                    for condition, then in zip(conditions, thens):
                        curr = curr.where(-condition, then)
                    if default is not None:
                        return curr.fillna(default)
                    return curr
                elif isinstance(expr, element.Between):
                    return input.between(*args)
                elif isinstance(expr, element.Cut):
                    bins = [bin.value for bin in expr.bins]
                    if expr.include_under:
                        bins.insert(0, -float('inf'))
                    if expr.include_over:
                        bins.append(float('inf'))
                    labels = [l.value for l in expr.labels]
                    return pd.cut(input, bins, right=expr.right, labels=labels,
                                  include_lowest=expr.include_lowest)

            if isinstance(expr.input, Scalar):
                return run()[0]
            else:
                return run()

        self._add_node(expr, handle)

    def visit_binary_op(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)

            if expr.lhs.dtype == df_types.datetime and expr.rhs.dtype == df_types.datetime:
                return ((pd.to_datetime(children_vals[0]) - pd.to_datetime(children_vals[1])) /
                        np.timedelta64(1, 'ms')).astype(np.int64)

            op = BINARY_OP_TO_PANDAS[expr.node_name]
            if isinstance(expr, Power) and isinstance(expr.dtype, df_types.Integer):
                return op(*children_vals).astype(types.df_type_to_np_type(expr.dtype))
            return op(*children_vals)

        self._add_node(expr, handle)

    def visit_unary_op(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)
            op = UNARY_OP_TO_PANDAS[expr.node_name]
            return op(*children_vals)

        self._add_node(expr, handle)

    def visit_math(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)
            if isinstance(expr, math.Log) and expr._base is not None:
                base = expr._base.value
                return np.log(children_vals[0]) / np.log(base)
            elif isinstance(expr, math.Trunc):
                decimals = expr._decimals.value
                order = 10 ** decimals
                return np.trunc(children_vals[0] * order) / order
            else:
                op = getattr(np, expr.node_name.lower())
                return op(*children_vals)

        self._add_node(expr, handle)

    def visit_string_op(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)

            input = children_vals[0]
            if isinstance(expr.input, Scalar):
                input = pd.Series([input])

            assert len(expr._args) == len(expr.args)
            kv = dict((name.lstrip('_'), self._get(arg, kw))
                      for name, arg in zip(expr._args[1:], expr.args[1:]))
            op = expr.node_name.lower()

            if op == 'get':
                res = getattr(getattr(input, 'str'), op)(children_vals[1])
            elif op == 'strptime':
                res = input.map(lambda x: datetime.strptime(x, children_vals[1]))
            elif op == 'extract':
                def extract(x, pat, flags, group):
                    regex = re.compile(pat, flags=flags)
                    m = regex.match(x)
                    if m:
                        return m.group(group)
                df = self._merge_values([expr.input, expr._pat, expr._flags, expr._group], kw)
                return pd.Series([extract(*r[1]) for r in df.iterrows()])
            elif op == 'split':
                return input.apply(lambda v: v.split(kv['pat'], kv['n']) if v is not None else None)
            elif op == 'stringtodict':
                def _parse_dict(x):
                    return dict(it.split(kv['kv_delim'], 1) for it in x.split(kv['item_delim']))
                return input.apply(lambda v: _parse_dict(v) if v is not None else None)
            else:
                if op == 'slice':
                    kv['stop'] = kv.pop('end', None)
                elif op == 'replace':
                    assert 'regex' in kv
                    if kv['regex']:
                        kv.pop('regex')
                    else:
                        kv['pat'] = re.escape(kv['pat'])
                        kv.pop('regex')
                res = getattr(getattr(input, 'str'), op)(**kv)
            if isinstance(expr.input, Scalar):
                return res[0]
            else:
                return res

        self._add_node(expr, handle)

    def visit_datetime_op(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)

            input = children_vals[0]
            if isinstance(expr.input, Scalar):
                input = pd.Series([input])

            assert len(children_vals) == len(expr.args)
            kv = dict(zip([arg.lstrip('_') for arg in expr._args[1:]],
                          children_vals[1:]))
            op = expr.node_name.lower()

            res = getattr(getattr(input, 'dt'), op)
            if not isinstance(res, pd.Series):
                res = res(**kv)

            if isinstance(expr.input, Scalar):
                return res[0]
            else:
                return res

        self._add_node(expr, handle)

    def visit_groupby(self, expr):
        def handle(kw):
            fields_exprs = expr._fields or expr._by + expr._aggregations

            fields = [[kw.get(field), ] if isinstance(field, Scalar) else kw.get(field)
                      for field in fields_exprs]
            length = max(len(it) for it in fields)
            for i in range(len(fields)):
                bys = self._get_compiled_bys(kw, expr._by, length)
                if isinstance(fields_exprs[i], SequenceExpr):
                    is_reduction = False
                    for n in itertools.chain(*(fields_exprs[i].all_path(expr.input))):
                        if isinstance(n, GroupedSequenceReduction):
                            is_reduction = True
                            break
                    if not is_reduction:
                        fields[i] = fields[i].groupby(bys).first()
                elif len(fields[i]) == 1:
                    fields[i] = pd.Series(fields[i] * length,
                                          name=fields_exprs[i].name).groupby(bys).first()

            df = pd.concat(fields, axis=1)
            if expr._having is not None:
                having = kw.get(expr._having)
                if all(not isinstance(e, GroupedSequenceReduction)
                       for e in itertools.chain(*expr._having.all_path(expr.input))):
                    # the having comes from the by fields, we need to do Series.groupby explicitly.
                    bys = self._get_compiled_bys(kw, expr._by, len(having))
                    having = having.groupby(bys).first()
                df = df[having]
            return pd.DataFrame(
                df.values, columns=[f.name for f in fields_exprs])[expr.schema.names]

        self._add_node(expr, handle)

    def visit_mutate(self, expr):
        def handle(kw):
            bys = self._get_compiled_bys(kw, expr._by, len(kw.get(expr.input)))
            bys = pd.concat(bys)
            bys.sort_values(inplace=True)

            wins = [kw.get(f) for f in expr._window_fields]
            return pd.DataFrame(pd.concat([bys] + wins, axis=1).values,
                                columns=expr.schema.names)

        self._add_node(expr, handle)

    def visit_value_counts(self, expr):
        def handle(kw):
            by = kw.get(expr._by)
            sort = kw.get(expr._sort)
            ascending = kw.get(expr._ascending)
            dropna = kw.get(expr._dropna)

            df = by.value_counts(sort=sort, ascending=ascending, dropna=dropna).to_frame()
            df.reset_index(inplace=True)
            return pd.DataFrame(df.values, columns=expr.schema.names)

        self._add_node(expr, handle)

    def visit_sort(self, expr):
        def handle(kw):
            input = kw.get(expr.input)
            names = expr.schema.names

            sorted_columns = OrderedDict()
            for field in expr._sorted_fields:
                name = str(uuid.uuid4())
                sorted_columns[name] = kw.get(field)
            input = input.assign(**sorted_columns)

            return input.sort_values(list(six.iterkeys(sorted_columns)),
                                     ascending=expr._ascending)[names]

        self._add_node(expr, handle)

    def visit_sort_column(self, expr):
        def handle(kw):
            input = kw.get(expr.input)
            if isinstance(expr.input, CollectionExpr):
                return input[expr._source_name]
            else:
                return input

        self._add_node(expr, handle)

    def visit_distinct(self, expr):
        def handle(kw):
            children_vals = self._get_children_vals(kw, expr)
            fields = children_vals[1:]

            ret = pd.concat(fields, axis=1, keys=expr.schema.names).drop_duplicates()
            ret.reset_index(drop=True, inplace=True)
            return ret

        self._add_node(expr, handle)

    def _get(self, item, kw):
        if item is None:
            return
        if isinstance(item, (list, tuple, set)):
            return type(item)(kw.get(it) for it in item)
        return kw.get(item)

    def visit_sample(self, expr):
        def handle(kw):
            input = self._get(expr.input, kw)
            parts = self._get(expr._parts, kw)
            i = self._get(expr._i, kw)
            n = self._get(expr._n, kw)
            frac = self._get(expr._frac, kw)
            replace = self._get(expr._replace, kw)
            weights = self._get(expr._weights, kw)
            strata = self._get(expr._strata, kw)
            random_state = self._get(expr._random_state, kw)
            if expr._sampled_fields:
                collection = pd.DataFrame(
                    pd.concat([kw.get(e) for e in expr._sampled_fields], axis=1).values,
                    columns=[str(uuid.uuid4()) for _ in expr._sampled_fields])
            else:
                collection = input

            if parts is not None and frac is None:
                frac = 1 / float(parts)
            if i is not None and (len(i) != 1 or i[0] > 0):
                raise NotImplementedError
            if not strata:
                sampled = collection.sample(n=n, frac=frac, replace=replace, weights=weights,
                                            random_state=random_state)
            else:
                frames = []
                frac = json.loads(frac) if expr._frac else dict()
                n = json.loads(n) if expr._n else dict()
                for val in itertools.chain(six.iterkeys(frac), six.iterkeys(n)):
                    v_frac = frac.get(val)
                    v_n = n.get(val)
                    filtered = collection[collection[strata].astype(str) == val]
                    sampled = filtered.sample(n=v_n, frac=v_frac, replace=replace, random_state=random_state)
                    frames.append(sampled)
                if frames:
                    sampled = pd.concat(frames)
                else:
                    sampled = pd.DataFrame(columns=collection.columns)
            if expr._sampled_fields:
                return pd.concat([input, sampled], axis=1, join='inner')[
                    [n for n in input.columns.tolist()]]
            return sampled

        self._add_node(expr, handle)

    def _get_names(self, x, force_list=False):
        if x is None:
            return x
        res = [it.name for it in x]
        if not force_list and len(res) == 1:
            return res[0]
        return res

    def _get_pivot_handler(self, expr):
        def handle(kw):
            df = self._merge_values(expr._group + expr._columns + expr._values, kw)

            pivoted = df.pivot(index=self._get_names(expr._group),
                               columns=self._get_names(expr._columns))
            columns = pivoted.columns.levels
            pivoted.reset_index(inplace=True)

            names = self._get_names(expr._group, True)
            tps = [g.dtype for g in expr._group]
            if len(columns[0]) == 1:
                tp = expr._values[0].dtype
                for name in columns[1]:
                    names.append(name)
                    tps.append(tp)
            else:
                for value_name, value_col in zip(columns[0], expr._values):
                    for name in columns[1]:
                        names.append('{0}_{1}'.format(name, value_name))
                        tps.append(value_col.dtype)
            expr._schema = TableSchema.from_lists(names, tps)

            res = pd.DataFrame(pivoted.values, columns=names)
            to_sub = CollectionExpr(_source_data=res, _schema=expr._schema)
            self._expr_dag.substitute(expr, to_sub)

            # trigger refresh of dynamic operations
            def func(expr):
                for c in traverse_until_source(expr, unique=True):
                    if c not in self._expr_to_dag_node:
                        c.accept(self)
            refresh_dynamic(to_sub, self._expr_dag, func=func)

            return to_sub, res

        return handle

    def _get_pivot_table_handler(self, expr):
        from ...expr.query import ExprVisitor

        class WrappedNumpyFunction(object):
            def __init__(self, fun):
                self._fun = fun

            def __call__(self, *args, **kwargs):
                return self._fun(*args, **kwargs)

        class AggFuncVisitor(ExprVisitor):
            def __init__(self, np_object, env):
                super(AggFuncVisitor, self).__init__(env)
                self.np_object = np_object

            def get_named_object(self, obj_name):
                if obj_name == 'count':
                    return WrappedNumpyFunction(np.size)
                elif obj_name == 'nunique':
                    return WrappedNumpyFunction(lambda x: np.size(np.unique(x)))
                elif obj_name == 'quantile':
                    return WrappedNumpyFunction(lambda x, prob: np.percentile(x, prob * 100))
                else:
                    return WrappedNumpyFunction(getattr(np, obj_name))

            def visit_Call(self, node):
                func = self.visit(node.func)
                args = [self.visit(n) for n in node.args]
                if isinstance(func, WrappedNumpyFunction):
                    args = [self.np_object] + args
                kwargs = OrderedDict([(kw.arg, self.visit(kw.value)) for kw in node.keywords])
                return func(*args, **kwargs)

        def get_real_aggfunc(aggfunc):
            if isinstance(aggfunc, six.string_types):
                if aggfunc == 'count':
                    return getattr(np, 'size')
                if aggfunc == 'nunique':
                    return lambda x: np.size(np.unique(x))
                if hasattr(np, aggfunc):
                    return getattr(np, aggfunc)

                def agg_eval(x):
                    visitor = AggFuncVisitor(x, {})
                    return visitor.eval(aggfunc, rewrite=False)

                return agg_eval

            if inspect.isclass(aggfunc):
                aggfunc = aggfunc()

                def func(x):
                    buffer = aggfunc.buffer()
                    for it in x:
                        aggfunc(buffer, it)
                    return aggfunc.getvalue(buffer)

                return func
            return aggfunc

        def handle(kw):
            columns = expr._columns if expr._columns else []
            df = self._merge_values(expr._group + columns + expr._values, kw)
            pivoted = df.pivot_table(index=self._get_names(expr._group),
                                     columns=self._get_names(expr._columns),
                                     values=self._get_names(expr._values),
                                     aggfunc=[get_real_aggfunc(f) for f in expr._agg_func],
                                     fill_value=expr.fill_value)
            levels = pivoted.columns.levels if isinstance(pivoted.columns, pd.MultiIndex) \
                else [pivoted.columns]
            pivoted.reset_index(inplace=True)

            names = self._get_names(expr._group, True)
            tps = [g.dtype for g in expr._group]
            columns_values = levels[-1] if expr._columns else [None, ]
            for agg_func_name in expr._agg_func_names:
                for value_col in expr._values:
                    for col in columns_values:
                        base = '{0}_'.format(col) if col is not None else ''
                        name = '{0}{1}_{2}'.format(base, value_col.name, agg_func_name)
                        names.append(name)
                        tps.append(value_col.dtype)
            if expr._columns:
                expr._schema = TableSchema.from_lists(names, tps)

            res = pd.DataFrame(pivoted.values, columns=names)
            to_sub = CollectionExpr(_source_data=res, _schema=expr._schema)
            self._expr_dag.substitute(expr, to_sub)

            # trigger refresh of dynamic operations
            def func(expr):
                for c in traverse_until_source(expr, unique=True):
                    if c not in self._expr_to_dag_node:
                        c.accept(self)

            refresh_dynamic(to_sub, self._expr_dag, func=func)

            return to_sub, res

        return handle

    def visit_pivot(self, expr):
        if isinstance(expr, PivotCollectionExpr):
            handle = self._get_pivot_handler(expr)
        else:
            handle = self._get_pivot_table_handler(expr)
        self._add_node(expr, handle)

    def _get_compiled_bys(self, kw, by_exprs, length):
        bys = [[kw.get(by), ] if isinstance(by, Scalar) else kw.get(by)
               for by in by_exprs]
        if any(isinstance(e, SequenceExpr) for e in by_exprs):
            size = max(len(by) for by, e in zip(bys, by_exprs)
                       if isinstance(e, SequenceExpr))
        else:
            size = length
        return [(by * size if len(by) == 1 else by) for by in bys]

    def _compile_grouped_reduction(self, kw, expr):
        if isinstance(expr, GroupedCount) and isinstance(expr._input, CollectionExpr):
            df = kw.get(expr.input)

            bys = [[kw.get(by), ] if isinstance(by, Scalar) else kw.get(by)
                   for by in expr._by]
            if any(isinstance(e, SequenceExpr) for e in expr._by):
                size = max(len(by) for by, e in zip(bys, expr._by)
                           if isinstance(e, SequenceExpr))
            else:
                size = len(df)
            bys = [(by * size if len(by) == 1 else by) for by in bys]
            return df.groupby(bys).size()

        if isinstance(expr, GroupedNUnique):
            input_df = pd.concat([kw.get(ip) for ip in expr.inputs], axis=1)
            bys = self._get_compiled_bys(kw, expr._by, len(input_df))
            return input_df.groupby(bys).apply(lambda x: pd.Series([len(x.drop_duplicates())]))[0]

        series = kw.get(expr.input) if isinstance(expr.input, SequenceExpr) \
            else pd.Series([kw.get(expr.input)], name=expr.input.name)

        bys = self._get_compiled_bys(kw, expr._by, len(series))
        if isinstance(expr.input, Scalar):
            series = pd.Series(series.repeat(len(bys[0])).values, index=bys[0].index)

        if isinstance(expr, GroupedCat):
            return series.groupby(bys).apply(lambda x: kw.get(expr._sep).join(x))

        if isinstance(expr, GroupedToList):
            if expr._unique:
                return series.groupby(bys).apply(lambda x: list(set(x)))
            else:
                return series.groupby(bys).apply(list)

        kv = dict()
        if hasattr(expr, '_ddof'):
            kv['ddof'] = expr._ddof
        op = expr.node_name.lower()
        op = 'size' if op == 'count' else op

        return getattr(series.groupby(bys), op)(**kv)

    def visit_reduction(self, expr):
        def handle(kw):
            if isinstance(expr, GroupedSequenceReduction):
                return self._compile_grouped_reduction(kw, expr)

            children_vals = self._get_children_vals(kw, expr)

            kv = dict()
            if hasattr(expr, '_ddof'):
                kv['ddof'] = expr._ddof

            op = expr.node_name.lower()
            op = 'size' if op == 'count' else op

            if isinstance(expr, NUnique):
                inputs = children_vals[:len(expr.inputs)]
                if len(expr.inputs) == 1:
                    inputs[0] = _filter_none(inputs[0])
                return len(pd.concat(inputs, axis=1).drop_duplicates())

            input = children_vals[0]
            if getattr(expr, '_unique', False):
                input = input.unique()

            if isinstance(expr, Count):
                if isinstance(expr.input, CollectionExpr):
                    return len(input)
                elif isinstance(expr.input, SequenceExpr):
                    return len(_filter_none(input))

            input = _filter_none(input)
            if isinstance(expr, (Cat, GroupedCat)):
                kv['sep'] = expr._sep.value if isinstance(expr._sep, Scalar) else expr._sep
                kv['na_rep'] = expr._na_rep.value \
                    if isinstance(expr._na_rep, Scalar) else expr._na_rep
                return getattr(getattr(input, 'str'), 'cat')(**kv)
            elif isinstance(expr, (ToList, GroupedToList)):
                return list(input)
            elif isinstance(expr, (Quantile, GroupedQuantile)):
                if isinstance(expr._prob, (list, set)):
                    return [np.percentile(input, p * 100) for p in expr._prob]
                else:
                    return np.percentile(input, expr._prob * 100)
            return getattr(input, op)(**kv)

        self._add_node(expr, handle)

    def visit_user_defined_aggregator(self, expr):
        def handle(kw):
            resources = self._get_resources(expr, kw)

            input = self._merge_values(expr._inputs, kw)

            func = expr._aggregator
            args = expr._func_args
            kwargs = expr._func_kwargs or dict()
            self._attempt_pickle_unpickle(func, args, kwargs)

            if resources:
                if not args and not kwargs:
                    agg = func(resources)
                else:
                    kwargs['resources'] = resources
                    agg = func(*args, **kwargs)
            else:
                agg = func(*args, **kwargs)

            if isinstance(expr, GroupedSequenceReduction):
                bys = [[kw.get(by), ] if isinstance(by, Scalar) else kw.get(by)
                       for by in expr._by]
            else:
                bys = [[1, ]]
            if expr._by and any(isinstance(e, SequenceExpr) for e in expr._by):
                size = max(len(by) for by, e in zip(bys, expr._by)
                           if isinstance(e, SequenceExpr))
            else:
                size = len(input)
            bys = [(by * size if len(by) == 1 else by) for by in bys]

            def iterrows(x):
                if getattr(expr, '_unique', False):
                    vset = set()
                    for it in x.iterrows():
                        if bytes(it[1].values.data) not in vset:
                            yield it
                        vset.add(bytes(it[1].values.data))
                else:
                    for it in x.iterrows():
                        yield it

            def f(x):
                buffer = agg.buffer()
                for it in iterrows(x):
                    agg(buffer, *it[1])
                ret = agg.getvalue(buffer)
                np_type = types.df_type_to_np_type(expr.dtype)
                return np.array([ret,], dtype=np_type)[0]

            res = input.groupby(bys).apply(f)
            if isinstance(expr, Scalar):
                return res.iloc[0]
            return res

        self._add_node(expr, handle)

    def visit_column(self, expr):
        def handle(kw):
            chidren_vals = self._get_children_vals(kw, expr)
            # FIXME: consider the name which is unicode
            return chidren_vals[0][expr._source_name]

        self._add_node(expr, handle)

    def _get_resources(self, expr, kw):
        if not expr._resources:
            return

        res = []
        collection_idx = 0
        for resource in expr._resources:
            if isinstance(resource, FileResource):
                res.append(resource.open())
            elif isinstance(resource, TableResource):
                def gen():
                    table = resource.get_source_table()
                    named_args = xnamedtuple('NamedArgs', table.table_schema.names)
                    partition = resource.get_source_table_partition()
                    with table.open_reader(partition=partition) as reader:
                        for r in reader:
                            yield named_args(*r.values)
                res.append(gen())
            else:
                resource = expr._collection_resources[collection_idx]
                collection_idx += 1

                df = kw.get(resource)

                def gen():
                    named_args = xnamedtuple('NamedArgs', resource.schema.names)
                    for r in df.iterrows():
                        yield named_args(*r[1])
                res.append(gen())

        return res

    def visit_function(self, expr):
        def handle(kw):
            resources = self._get_resources(expr, kw)
            if not expr._multiple:
                input = self._get_children_vals(kw, expr)[0]

                if isinstance(expr.inputs[0], Scalar):
                    input = pd.Series([input])

                func = expr._func
                args = expr._func_args
                kwargs = expr._func_kwargs
                self._attempt_pickle_unpickle(func, args, kwargs)

                if args is not None and len(args) > 0:
                    raise NotImplementedError
                if kwargs is not None and len(kwargs) > 0:
                    raise NotImplementedError

                if inspect.isclass(func):
                    if resources:
                        func = func(resources)
                    else:
                        func = func()
                else:
                    if resources:
                        func = func(resources)

                res = input.map(func)
                if isinstance(expr.inputs[0], Scalar):
                    return res[0]
                return res
            else:
                input = self._merge_values(expr.inputs, kw)

                def func(s):
                    names = [f.name for f in expr.inputs]
                    t = xnamedtuple('NamedArgs', names)
                    row = t(*s.tolist())
                    if not inspect.isfunction(expr._func):
                        if resources:
                            f = expr._func(resources)
                        else:
                            f = expr._func()
                    else:
                        if resources:
                            f = expr._func(resources)
                        else:
                            f = expr._func

                    res = f(row, *expr._func_args, **expr._func_kwargs)
                    if not inspect.isgeneratorfunction(f):
                        return res
                    return next(res)

                if PD_APPLY_HAS_RESULT_TYPE:
                    return input.apply(func, axis=1, result_type='reduce',
                                    args=expr._func_args, **expr._func_kwargs)
                else:
                    return input.apply(func, axis=1, reduce=True,
                                    args=expr._func_args, **expr._func_kwargs)

        self._add_node(expr, handle)

    def visit_reshuffle(self, expr):
        def handle(kw):
            if expr._sort_fields is not None:
                input = kw.get(expr._input)
                names = []
                for sort in expr._sort_fields:
                    name = str(uuid.uuid4())
                    input[name] = kw.get(sort)
                    names.append(name)
                input = input.sort_values(
                    names, ascending=[f._ascending for f in expr._sort_fields])
                return input[expr.schema.names]
            return kw.get(expr._input)

        self._add_node(expr, handle)

    def _check_output_types(self, pd_df, expect_df_types):
        for field, expect_df_type in zip(pd_df.columns, expect_df_types):
            arr = pd_df[field].values
            try:
                df_type = types.np_type_to_df_type(pd_df[field].dtype, arr=arr)
            except TypeError:
                # all element is None
                continue
            if not expect_df_type.can_implicit_cast(df_type):
                raise TypeError('Field(%s) has wrong type, expect %s, got %s' % (
                    field, expect_df_type, df_type
                ))
        return pd_df

    def visit_apply_collection(self, expr):
        def conv(l):
            if isinstance(l, tuple):
                l = list(l)
            elif not isinstance(l, list):
                l = [l, ]
            return l

        def handle(kw):
            resources = self._get_resources(expr, kw)

            input = self._merge_values(expr.fields, kw)

            names = [f.name for f in expr.fields]
            t = xnamedtuple('NamedArgs', names)

            expr._func_args = expr._func_args or ()
            expr._func_kwargs = expr._func_kwargs or {}

            func = expr._func
            self._attempt_pickle_unpickle(func, expr._func_args, expr._func_kwargs)

            if isinstance(func, six.string_types) and func.upper() in BUILTIN_FUNCS:
                func = BUILTIN_FUNCS[func.upper()]
            if inspect.isfunction(func):
                if resources:
                    func = func(resources)

                is_generator_function = inspect.isgeneratorfunction(func)
                close_func = None
                is_close_generator_function = False
            elif hasattr(func, '__call__'):
                if resources:
                    func = func(resources)
                else:
                    func = func()
                is_generator_function = inspect.isgeneratorfunction(func.__call__)
                close_func = getattr(func, 'close', None)
                is_close_generator_function = inspect.isgeneratorfunction(close_func)
            else:
                raise NotImplementedError

            rows = []
            indices = []
            idx = 0
            for s in input.iterrows():
                row = t(*s[1])
                res = func(row, *expr._func_args, **expr._func_kwargs)
                expand_num = 0
                if is_generator_function:
                    for l in res:
                        rows.append(conv(l))
                        expand_num += 1
                else:
                    if res:
                        rows.append(conv(res))
                        expand_num += 1
                if expand_num == 0 and expr._keep_nulls:
                    rows.append([None] * len(names))
                    expand_num += 1
                indices.extend([s[0]] * expand_num)
                idx = max(idx, s[0] + 1)
            if close_func:
                expand_num = 0
                if is_close_generator_function:
                    for l in close_func(*expr._func_args, **expr._func_kwargs):
                        rows.append(conv(l))
                        expand_num += 1
                else:
                    rows.append(close_func(*expr._func_args, **expr._func_kwargs))
                    expand_num += 1
                indices.extend([idx] * expand_num)
            if expr._lateral_view:
                out_df = pd.DataFrame(rows, columns=expr.schema.names,
                                      index=pd.Int64Index(indices))
            else:
                out_df = pd.DataFrame(rows, columns=expr.schema.names)
            return self._check_output_types(out_df, expr.schema.types)

        self._add_node(expr, handle)

    def visit_lateral_view(self, expr):
        def handle(kw):
            lv_sources = dict()
            for lv in expr.lateral_views:
                for col_name in lv.schema.names:
                    lv_sources[col_name] = lv

            children = expr.children()

            fields = self._get_children_vals(kw, children=children)[1:len(expr._fields) + 1]
            names = expr.schema.names

            idx = reduce(operator.and_, (set(f.index.tolist()) for f, e in zip(fields, expr._fields)
                                         if isinstance(e, SequenceExpr)))
            idx = pd.Int64Index(sorted(idx))
            result = pd.DataFrame(index=idx)
            lv_visited = set()
            for i in range(len(fields)):
                f = fields[i]
                if names[i] in lv_sources:
                    lv_src = lv_sources[names[i]]
                    if lv_src in lv_visited:
                        continue
                    lv_visited.add(lv_src)
                    f = kw[lv_src]
                elif not isinstance(f, pd.Series):
                    f = pd.Series([f] * len(idx), index=idx, name=names[i])
                result = result.join(f)

            return result

        self._add_node(expr, handle)

    def visit_composite_op(self, expr):
        def handle(kw):
            def _zip_args(fields):
                zip_args = []
                seq_index = None
                for it in fields:
                    if isinstance(it, SequenceExpr):
                        zip_args.append(kw[it])
                        seq_index = kw[it].index
                    else:
                        zip_args.append(itertools.repeat(kw[it]))
                return seq_index, zip_args

            children_vals = self._get_children_vals(kw, expr)
            _input = children_vals[0]
            if isinstance(expr, composites.ListDictLength):
                return _input.apply(lambda v: len(v) if v is not None else None)
            elif isinstance(expr, composites.ListDictGetItem):
                def _get_list_item(l, x):
                    try:
                        return l[x] if l is not None else None
                    except IndexError:
                        return None

                _value = children_vals[1]
                if isinstance(expr.input.dtype, df_types.List):
                    item_fun = _get_list_item
                else:
                    item_fun = lambda s, k: s.get(k) if s is not None else None
                if isinstance(expr, Scalar):
                    return item_fun(_input, _value)
                else:
                    if isinstance(expr.input, Scalar):
                        return _value.apply(lambda v: item_fun(_input, v))
                    if isinstance(expr._key, Scalar):
                        return _input.apply(lambda v: item_fun(v, _value))
                    seq_values = [item_fun(k, v) for k, v in compat.izip(_input, _value)]
                    return pd.Series(seq_values, index=_input.index, name=expr.name)
            elif isinstance(expr, composites.ListContains):
                _value = children_vals[1]
                contains_fun = lambda s, k: k in s if s is not None else None
                if isinstance(expr, Scalar):
                    return contains_fun(_input, _value)
                else:
                    if isinstance(expr.input, Scalar):
                        return _value.apply(lambda v: contains_fun(_input, v))
                    if isinstance(expr._value, Scalar):
                        return _input.apply(lambda v: contains_fun(v, _value))
                    seq_values = [contains_fun(k, v) for k, v in compat.izip(_input, _value)]
                    return pd.Series(seq_values, index=_input.index, name=expr.name)
            elif isinstance(expr, composites.ListSort):
                return _input.apply(lambda l: sorted(l) if l is not None else None)
            elif isinstance(expr, composites.DictKeys):
                return _input.apply(lambda d: list(six.iterkeys(d)) if d is not None else None)
            elif isinstance(expr, composites.DictValues):
                return _input.apply(lambda d: list(six.itervalues(d)) if d is not None else None)
            elif isinstance(expr, composites.ListBuilder):
                if isinstance(expr, Scalar):
                    return [kw[v] for v in expr._values]
                else:
                    seq_index, zip_args = _zip_args(expr._values)
                    seq_values = []
                    for r in compat.izip(*zip_args):
                        seq_values.append(list(r))
                    return pd.Series(seq_values, index=seq_index, name=expr.name)
            elif isinstance(expr, composites.DictBuilder):
                if isinstance(expr, Scalar):
                    return OrderedDict((kw[k], kw[v]) for k, v in compat.izip(expr._keys, expr._values))
                else:
                    seq_index, zip_args = _zip_args(expr._keys + expr._values)
                    seq_values = []
                    dict_len = len(expr._values)
                    for r in zip(*zip_args):
                        seq_values.append(OrderedDict((k, v) for k, v in compat.izip(r[:dict_len], r[dict_len:])))
                    return pd.Series(seq_values, index=seq_index, name=expr.name)
            else:
                raise NotImplementedError

        self._add_node(expr, handle)

    def visit_sequence(self, expr):
        raise NotImplementedError

    def visit_cum_window(self, expr):
        if expr.preceding is not None or expr.following is not None:
            raise NotImplementedError

        def handle(kw):
            input = kw.get(expr.input)
            bys = self._get_compiled_bys(kw, expr.partition_by, len(input))
            grouped = input.groupby(bys)
            if expr.order_by:
                sort = [kw.get(e) for e in expr.order_by]
                ascendings = [e._ascending for e in expr.order_by]
                for s in sort:
                    sort_name = str(uuid.uuid4())
                    s.name = sort_name
            else:
                sort = None
                ascendings = None

            def f(x):
                if sort:
                    df = pd.concat([x] + sort, join='inner', axis=1)
                    df.sort_values([s.name for s in sort], ascending=ascendings, inplace=True)
                    series = df[x.name]

                    if expr.node_name in SORT_CUM_WINDOW_OP_TO_PANDAS:
                        return SORT_CUM_WINDOW_OP_TO_PANDAS[expr.node_name](series)
                    elif expr.node_name == 'NthValue':
                        values = [None] * len(series)
                        if expr._skip_nulls:
                            new_series = _filter_none(series)
                        else:
                            new_series = series
                        if expr._nth < len(new_series):
                            values[expr._nth:] = [new_series.iloc[expr._nth]] * (len(series) - expr._nth)
                        return pd.Series(values, index=series.index)
                    else:
                        raise NotImplementedError
                else:
                    if expr.distinct:
                        new_x = x.drop_duplicates()
                    else:
                        new_x = x
                    if expr.node_name in CUM_WINDOW_OP_TO_PANDAS:
                        val = CUM_WINDOW_OP_TO_PANDAS[expr.node_name](new_x)
                    elif expr.node_name == 'NthValue':
                        if expr._skip_nulls:
                            new_series = _filter_none(x)
                        else:
                            new_series = x
                        if expr._nth < len(new_series):
                            val = new_series.iloc[expr._nth]
                        else:
                            val = None
                    else:
                        raise NotImplementedError
                    return pd.Series([val] * len(x), index=x.index)

            res = grouped.apply(f)
            if sort:
                for _ in bys:
                    res = res.reset_index(level=0, drop=True)
            return res

        self._add_node(expr, handle)

    def visit_rank_window(self, expr):
        def handle(kw):
            input = kw.get(expr.input)
            sort = [kw.get(e) * (1 if e._ascending else -1)
                    for e in expr.order_by]
            bys = self._get_compiled_bys(kw, expr.partition_by, len(input))

            sort_names = [str(uuid.uuid4()) for _ in sort]
            by_names = [str(uuid.uuid4()) for _ in bys]
            input_names = [input.name] if isinstance(input, pd.Series) else input.columns.tolist()
            df = pd.DataFrame(pd.concat([input] + sort + [pd.Series(b) for b in bys], axis=1).values,
                              columns=input_names + sort_names + by_names,
                              index=input.index)
            df.sort_values(sort_names, inplace=True)
            grouped = df.groupby(by_names)

            try:
                pd_fast_zip = pd._libs.lib.fast_zip
            except AttributeError:
                pd_fast_zip = pd.lib.fast_zip

            def f(x):
                s_df = pd.Series(pd_fast_zip([x[s].values for s in sort_names]), index=x.index)
                if expr.node_name == 'Rank':
                    return s_df.rank(method='min')
                elif expr.node_name == 'DenseRank':
                    return s_df.rank(method='dense')
                elif expr.node_name == 'RowNumber':
                    return pd.Series(compat.lrange(1, len(s_df) + 1), index=s_df.index)
                elif expr.node_name == 'PercentRank':
                    if len(s_df) == 1:
                        return pd.Series([0.0, ], index=s_df.index)
                    return (s_df.rank(method='min') - 1) / (len(s_df) - 1)
                elif expr.node_name == 'CumeDist':
                    return pd.Series([v * 1.0 / len(s_df) for v in compat.irange(1, len(s_df) + 1)],
                                     index=s_df.index)
                elif expr.node_name == 'QCut':
                    if len(s_df) <= 1:
                        return pd.Series([0] * len(s_df), index=s_df.index, dtype=np.int64)
                    return pd.Series(pd.qcut(compat.irange(1, len(s_df) + 1), expr._bins, labels=False),
                                     index=s_df.index, dtype=np.int64)
                else:
                    raise NotImplementedError

            res = grouped.apply(f)
            if isinstance(res, pd.DataFrame):
                res = res.iloc[0]
            else:
                for _ in bys:
                    res = res.reset_index(level=0, drop=True)
            return res

        self._add_node(expr, handle)

    def visit_shift_window(self, expr):
        def handle(kw):
            input = kw.get(expr.input)

            bys = self._get_compiled_bys(kw, expr.partition_by, len(input))
            grouped = input.groupby(bys)
            if expr.order_by:
                sort = [kw.get(e) for e in expr.order_by]
                ascendings = [e._ascending for e in expr.order_by]
                for s in sort:
                    sort_name = str(uuid.uuid4())
                    s.name = sort_name
            else:
                sort = None
                ascendings = None

            if expr.node_name == 'Lag':
                shift = kw.get(expr.offset)
            else:
                assert expr.node_name == 'Lead'
                shift = -kw.get(expr.offset)
            default = kw.get(expr.default)

            def f(x):
                if sort:
                    df = pd.concat([x] + sort, join='inner', axis=1)
                    df.sort_values([s.name for s in sort], ascending=ascendings, inplace=True)
                    series = df[x.name]
                else:
                    series = x

                res = series.shift(shift)
                if default is not None:
                    return res.fillna(default)
                return res

            res = grouped.apply(f)
            if sort:
                for _ in bys:
                    res = res.reset_index(level=0, drop=True)
            return res

        self._add_node(expr, handle)

    def visit_scalar(self, expr):
        def handle(_):
            if isinstance(expr, DTScalar):
                arg_name = type(expr).__name__.lower()[:-6] + 's'
                value = expr.value
                if arg_name == 'milliseconds':
                    arg_name = 'microseconds'
                    value *= 1000
                return pd.DateOffset(**{arg_name: value})
            if expr.value is not None:
                return expr.value

            return None

        self._add_node(expr, handle)

    def visit_cast(self, expr):
        def handle(kw):
            dtype = types.df_type_to_np_type(expr.dtype)
            input = self._get_children_vals(kw, expr)[0]
            if isinstance(expr._input, Scalar):
                return pd.Series([input]).astype(dtype)[0]
            return input.astype(dtype)

        self._add_node(expr, handle)

    @classmethod
    def _find_all_equalizations(cls, predicate, lhs, rhs):
        return [eq for eq in traverse_until_source(predicate, top_down=True, unique=True)
                if isinstance(eq, arithmetic.Equal) and
                eq.is_ancestor(lhs) and eq.is_ancestor(rhs)]

    def visit_join(self, expr):
        def handle(kw):
            left = kw.get(expr._lhs)
            right = kw.get(expr._rhs)

            eqs = expr._predicate

            left_ons = []
            right_ons = []
            on_same_names = set()
            for eq in eqs:
                if isinstance(eq._lhs, Column) and isinstance(eq._rhs, Column) and \
                        eq._lhs.source_name == eq._rhs.source_name:
                    left_ons.append(eq._lhs.source_name)
                    right_ons.append(eq._rhs.source_name)
                    on_same_names.add(eq._lhs.source_name)
                    continue

                left_name = str(uuid.uuid4())
                left[left_name] = kw.get(eq._lhs)
                left_ons.append(left_name)

                right_name = str(uuid.uuid4())
                right[right_name] = kw.get(eq._rhs)
                right_ons.append(right_name)

            for idx, collection in enumerate([left, right]):
                collection_expr = (expr._lhs, expr._rhs)[idx]
                for field_name in collection_expr.schema.names:
                    if field_name in expr._renamed_columns and field_name in on_same_names:
                        new_name = expr._renamed_columns[field_name][idx]
                        collection[new_name] = collection[field_name]

            merged = left.merge(right, how=JOIN_DICT[expr._how], left_on=left_ons,
                                right_on=right_ons,
                                suffixes=(expr._left_suffix, expr._right_suffix))
            cols = []
            for name in expr.schema.names:
                if name in merged:
                    cols.append(merged[name])
                else:
                    cols.append(merged[expr._column_origins[name][1]])

            return pd.concat(cols, axis=1, keys=expr.schema.names)

        # Just add node, shouldn't add edge here
        node = (expr, handle)
        self._dag.add_node(node)
        self._expr_to_dag_node[expr] = node

    def visit_extract_kv(self, expr):
        def handle(kw):
            from ... import types
            _input = kw.get(expr._input)
            columns = [getattr(_input, c.name) for c in expr._columns]
            kv_delim = kw.get(expr._kv_delimiter)
            item_delim = kw.get(expr._item_delimiter)
            default = kw.get(expr._default)

            kv_slot_map = dict()
            app_col_names = []

            def validate_kv(v):
                parts = v.split(kv_delim)
                if len(parts) != 2:
                    raise ValueError('Malformed KV pair: %s' % v)
                return parts[0]

            for col in columns:
                kv_slot_map[col.name] = dict()

                keys = col.apply(lambda s: [validate_kv(kv) for kv in s.split(item_delim)])
                for k in sorted(compat.reduce(lambda a, b: set(a) | set(b), keys, set())):
                    app_col_names.append('%s_%s' % (col.name, k))
                    kv_slot_map[col.name][k] = len(app_col_names) - 1

            type_adapter = None
            if isinstance(expr._column_type, types.Float):
                type_adapter = float
            elif isinstance(expr._column_type, types.Integer):
                type_adapter = int

            append_grid = [[default] * len(app_col_names) for _ in compat.irange(len(_input))]
            for col in columns:
                series = getattr(_input, col.name)
                for idx, v in enumerate(series):
                    for kv_item in v.split(item_delim):
                        k, v = kv_item.split(kv_delim)
                        if type_adapter:
                            v = type_adapter(v)
                        append_grid[idx][kv_slot_map[col.name][k]] = v

            intact_names = [c.name for c in expr._intact]
            intact_types = [c.dtype for c in expr._intact]
            intact_df = _input[intact_names]
            append_df = pd.DataFrame(append_grid, columns=app_col_names)
            expr._schema = TableSchema.from_lists(
                intact_names + app_col_names,
                intact_types + [expr._column_type] * len(app_col_names),
            )

            res = pd.concat([intact_df, append_df], axis=1)
            to_sub = CollectionExpr(_source_data=res, _schema=expr._schema)
            self._expr_dag.substitute(expr, to_sub)

            # trigger refresh of dynamic operations
            def func(expr):
                for c in traverse_until_source(expr, unique=True):
                    if c not in self._expr_to_dag_node:
                        c.accept(self)

            refresh_dynamic(to_sub, self._expr_dag, func=func)

            return to_sub, res

        self._add_node(expr, handle)

    def visit_union(self, expr):
        if expr._distinct:
            raise CompileError("Distinct union is not supported here.")

        def handle(kw):
            left = kw.get(expr._lhs)
            right = kw.get(expr._rhs)

            merged = pd.concat([left, right])
            return merged[expr.schema.names]

        self._add_node(expr, handle)

    def visit_concat(self, expr):
        def handle(kw):
            left = kw.get(expr._lhs)
            right = kw.get(expr._rhs)

            merged = pd.concat([left, right], axis=1)
            return merged[expr.schema.names]

        self._add_node(expr, handle)

    def visit_append_id(self, expr):
        def handle(kw):
            _input = kw.get(expr._input)
            id_col = kw.get(expr._id_col)

            id_seq = pd.DataFrame(compat.lrange(len(_input)), columns=[id_col])
            return pd.concat([id_seq, _input], axis=1)

        self._add_node(expr, handle)

    def visit_split(self, expr):
        def handle(kw):
            _input = kw.get(expr._input)
            frac = kw.get(expr._frac)
            seed = kw.get(expr._seed) if expr._seed else None
            split_id = kw.get(expr._split_id)

            if seed is not None:
                np.random.seed(seed)

            cols = list(_input.columns)
            factor_col = 'rand_factor_%d' % int(time.time())
            factor_df = pd.DataFrame(np.random.rand(len(_input)), columns=[factor_col])
            concated_df = pd.concat([factor_df, _input], axis=1)

            if split_id == 0:
                return concated_df[concated_df[factor_col] <= frac][cols]
            else:
                return concated_df[concated_df[factor_col] > frac][cols]

        self._add_node(expr, handle)
