def _if_else()

in torcharrow/velox_rt/numerical_column_cpu.py [0:0]

• 27 lines of code
• 11 McCabe index (conditional complexity)

    def _if_else(self, then_, else_):
        """Vectorized if-then-else"""
        self._prototype_support_warning("_if_else")

        if not dt.is_boolean(self.dtype):
            raise TypeError("condition must be a boolean vector")
        if not isinstance(then_, IColumn):
            then_ = ta.Column(then_)
        if not isinstance(else_, IColumn):
            else_ = ta.Column(else_)
        lub = dt.common_dtype(then_.dtype, else_.dtype)

        if lub is None or dt.is_void(lub):
            raise TypeError(
                "then and else branches must have compatible types, got {then_.dtype} and {else_.dtype}, respectively"
            )
        if isinstance(then_, NumericalColumnCpu) and isinstance(
            else_, NumericalColumnCpu
        ):
            col = velox.Column(get_velox_type(lub))
            for i in range(len(self)):
                if self._getmask(i):
                    col.append_null()
                else:
                    col.append(
                        then_._getdata(i) if self._getdata(i) else else_._getdata(i)
                    )
            return ColumnFromVelox._from_velox(self.device, lub, col, True)

        else:
            # refer back to default handling...
            return INumericalColumn._if_else(self, then_, else_)