def execute_udaf_node_groupby(op, *args, **kwargs):
# construct a generator that yields the next group of data for
# every argument excluding the first (pandas performs the
# iteration for the first argument) for each argument that is a
# SeriesGroupBy.
#
# If the argument is not a SeriesGroupBy then keep repeating it
# until all groups are exhausted.
context = kwargs.pop('context', None)
assert context is not None, 'context is None'
iters = (
(data for _, data in arg)
if isinstance(arg, SeriesGroupBy)
else itertools.repeat(arg) for arg in args[1:]
)
funcsig = signature(func)
def aggregator(first, *rest, **kwargs):
# map(next, *rest) gets the inputs for the next group
# TODO: might be inefficient to do this on every call
args, kwargs = arguments_from_signature(
funcsig, first, *map(next, rest), **kwargs
)
return func(*args, **kwargs)
result = context.agg(args[0], aggregator, *iters, **kwargs)
return result
def wrapper(func):
num_params = sum(
param.kind in {param.POSITIONAL_OR_KEYWORD, param.POSITIONAL_ONLY}
for param in signature(func).parameters.values()
if param.default is empty)
num_declared = len(input_type)
if num_params != num_declared:
raise TypeError('Function {!r} has {:d} parameters, '
'input_type has {:d}. These must match'.format(
func.__name__,
num_params,
num_declared,
))
return func
def execute_udf_node_groupby(op, *args, **kwargs):
groupers = [
grouper for grouper in (getattr(arg, 'grouper', None)
for arg in args)
if grouper is not None
]
# all grouping keys must be identical
assert all(groupers[0] == grouper for grouper in groupers[1:])
# we're performing a scalar operation on grouped column, so
# perform the operation directly on the underlying Series and
# regroup after it's finished
arguments = [getattr(arg, 'obj', arg) for arg in args]
groupings = groupers[0].groupings
args, kwargs = arguments_from_signature(
signature(func), *arguments, **kwargs)
return func(*args, **kwargs).groupby(groupings)
def execute_udaf_node(op, *args, **kwargs):
args, kwargs = arguments_from_signature(
signature(func), *args, **kwargs
)
return func(*args, **kwargs)
