def _switch()

in odps/df/expr/element.py [0:0]

• 33 lines of code
• 24 McCabe index (conditional complexity)

def _switch(expr, *args, **kw):
    """
    Similar to the case-when in SQL. Refer to the example below

    :param expr:
    :param args:
    :param kw:
    :return: sequence or scalar

    :Example:

    >>> # if df.id == 3 then df.name
    >>> # elif df.id == df.fid.abs() then df.name + 'test'
    >>> # default: 'test'
    >>> df.id.switch(3, df.name, df.fid.abs(), df.name + 'test', default='test')
    """
    default = _scalar(kw.get('default'))

    if len(args) <= 0:
        raise errors.ExpressionError('Switch must accept more than one condition')

    if all(isinstance(arg, tuple) and len(arg) == 2 for arg in args):
        conditions, thens = [list(tp) for tp in zip(*args)]
    else:
        conditions = [arg for i, arg in enumerate(args) if i % 2 == 0]
        thens = [arg for i, arg in enumerate(args) if i % 2 == 1]

    if len(conditions) == len(thens):
        conditions, thens = _scalar(conditions), _scalar(thens)
    else:
        raise errors.ExpressionError('Switch should be called by case and then pairs')

    if isinstance(expr, (Scalar, SequenceExpr)):
        case = expr
    else:
        case = None
        if not all(hasattr(it, 'dtype') and it.dtype == types.boolean for it in conditions):
            raise errors.ExpressionError('Switch must be called by all boolean conditions')

    res = thens if default is None else thens + [default, ]
    output_type = utils.highest_precedence_data_type(*(it.dtype for it in res))

    is_seq = isinstance(expr, SequenceExpr) or \
        any(isinstance(it, SequenceExpr) for it in conditions) or \
        any(isinstance(it, SequenceExpr) for it in res)
    if case is not None:
        is_seq = is_seq or isinstance(case, SequenceExpr)

    kwargs = dict()
    if is_seq:
        kwargs['_data_type'] = output_type
    else:
        kwargs['_value_type'] = output_type
    return Switch(_input=expr, _case=case, _conditions=conditions,
                  _thens=thens, _default=default, **kwargs)