def visit_element_op()

in odps/df/backends/pd/compiler.py [0:0]

• 63 lines of code
• 31 McCabe index (conditional complexity)

    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)