def execute_until_in_scope(
expr,
scope: Scope,
timecontext: Optional[TimeContext] = None,
aggcontext=None,
clients=None,
post_execute_=None,
**kwargs,
) -> Scope:
"""Execute until our op is in `scope`.
Parameters
----------
expr : ibis.expr.types.Expr
scope : Scope
timecontext : Optional[TimeContext]
aggcontext : Optional[AggregationContext]
clients : List[ibis.backends.base.BaseBackend]
kwargs : Mapping
"""
# these should never be None
assert aggcontext is not None, 'aggcontext is None'
assert clients is not None, 'clients is None'
assert post_execute_ is not None, 'post_execute_ is None'
# base case: our op has been computed (or is a leaf data node), so
# return the corresponding value
op = expr.op()
if scope.get_value(op, timecontext) is not None:
return scope
if isinstance(op, ops.Literal):
# special case literals to avoid the overhead of dispatching
# execute_node
return Scope(
{
op:
execute_literal(
op, op.value, expr.type(), aggcontext=aggcontext, **kwargs)
},
timecontext,
)
# figure out what arguments we're able to compute on based on the
# expressions inputs. things like expressions, None, and scalar types are
# computable whereas ``list``s are not
computable_args = [arg for arg in op.inputs if is_computable_input(arg)]
# pre_executed_states is a list of states with same the length of
# computable_args, these states are passed to each arg
if timecontext:
arg_timecontexts = compute_time_context(
op,
num_args=len(computable_args),
timecontext=timecontext,
clients=clients,
scope=scope,
)
else:
arg_timecontexts = [None] * len(computable_args)
pre_executed_scope = pre_execute(
op,
*clients,
scope=scope,
timecontext=timecontext,
aggcontext=aggcontext,
**kwargs,
)
new_scope = scope.merge_scope(pre_executed_scope)
# Short circuit: if pre_execute puts op in scope, then we don't need to
# execute its computable_args
if new_scope.get_value(op, timecontext) is not None:
return new_scope
# recursively compute each node's arguments until we've changed type.
# compute_time_context should return with a list with the same length
# as computable_args, the two lists will be zipping together for
# further execution
if len(arg_timecontexts) != len(computable_args):
raise com.IbisError(
'arg_timecontexts differ with computable_arg in length '
f'for type:\n{type(op).__name__}.')
scopes = [
execute_until_in_scope(
arg,
new_scope,
timecontext=timecontext,
aggcontext=aggcontext,
post_execute_=post_execute_,
clients=clients,
**kwargs,
) if hasattr(arg, 'op') else Scope({arg: arg}, timecontext)
for (arg, timecontext) in zip(computable_args, arg_timecontexts)
]
# if we're unable to find data then raise an exception
if not scopes and computable_args:
raise com.UnboundExpressionError(
f'Unable to find data for expression:\n{repr(expr)}')
# there should be exactly one dictionary per computable argument
assert len(computable_args) == len(scopes)
new_scope = new_scope.merge_scopes(scopes)
# pass our computed arguments to this node's execute_node implementation
data = [
new_scope.get_value(arg.op(), timecontext)
if hasattr(arg, 'op') else arg
for (arg, timecontext) in zip(computable_args, arg_timecontexts)
]
result = execute_node(
op,
*data,
scope=scope,
timecontext=timecontext,
aggcontext=aggcontext,
clients=clients,
**kwargs,
)
computed = post_execute_(op, result, timecontext=timecontext)
return Scope({op: computed}, timecontext)
def execute_log_series_gb_series_gb(op, left, right, **kwargs):
result = execute_node(op, left.obj, right.obj, **kwargs)
return result.groupby(left.grouper.groupings)
def execute_decimal_log_with_np_integer_base(op, data, base, **kwargs):
return execute_node(op, data, int(base), **kwargs)
def execute_unary_op_series_gb(op, operand, **kwargs):
result = execute_node(op, operand.obj, **kwargs)
return result.groupby(operand.grouper.groupings)
def execute_decimal_log_with_real_base(op, data, base, **kwargs):
return execute_node(op, data, decimal.Decimal(base), **kwargs)