def _normalize_agg_arg()

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

• 21 lines of code
• 13 McCabe index (conditional complexity)

    def _normalize_agg_arg(self, arg):
        res = []  # triple name, field, op
        if isinstance(arg, str):
            # normalize
            arg = [arg]
        if isinstance(arg, list):
            for op in arg:
                for f in self._item_fields:
                    res.append((f"{f.name}.{op}", f, op))
        elif isinstance(arg, dict):
            for n, ops in arg.items():
                fields = [f for f in self._item_fields if f.name == n]
                if len(fields) == 0:
                    raise ValueError(f"column ({n}) does not exist")
                # TODO handle duplicate columns, if ever...
                assert len(fields) == 1
                if isinstance(ops, str):
                    ops = [ops]
                for op in ops:
                    res.append((f"{n}.{op}", fields[0], op))
        else:
            raise TypeError(f"unexpected arg type ({type(arg).__name__})")
        return res