Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
cases.
"""
require(isinstance(index_var, ir.Var))
# case where the index is a const itself in outer function
var_const = guard(_get_const_two_irs, stencil_ir, func_ir, index_var)
if var_const is not None:
return var_const
# get index definition
index_def = ir_utils.get_definition(stencil_ir, index_var)
# match inner_var = unary(index_var)
var_const = guard(
_get_const_unary_expr, stencil_ir, func_ir, index_def)
if var_const is not None:
return var_const
# match inner_var = arg1 + arg2
var_const = guard(
_get_const_binary_expr, stencil_ir, func_ir, index_def)
if var_const is not None:
return var_const
raise GuardException
topo_order = find_topo_order(blocks)
for label in topo_order:
block = blocks[label]
new_body = []
equiv_set = self.array_analysis.get_equiv_set(label)
for instr in block.body:
parfor = None
if isinstance(instr, ir.Assign):
loc = instr.loc
lhs = instr.target
expr = instr.value
callname = guard(find_callname, self.func_ir, expr)
if (callname == ('reduce', 'builtins')
or callname == ('reduce', '_functools')):
# reduce function with generic function
parfor = guard(self._reduce_to_parfor, equiv_set, lhs,
expr.args, loc)
if parfor:
instr = parfor
new_body.append(instr)
block.body = new_body
return
from numba.stencil import StencilFunc
# could be make_function from list comprehension which is ok
func_def = guard(get_definition, self.state.func_ir, rhs.func)
if isinstance(func_def, ir.Expr) and func_def.op == 'make_function':
return [assign]
if isinstance(func_def, ir.Global) and isinstance(func_def.value, StencilFunc):
return [assign]
warnings.warn(
"function call couldn't be found for initial analysis")
return [assign]
else:
func_name, func_mod = fdef
string_methods_types = (SplitViewStringMethodsType, StringMethodsType)
if isinstance(func_mod, ir.Var) and isinstance(self.state.typemap[func_mod.name], string_methods_types):
f_def = guard(get_definition, self.state.func_ir, rhs.func)
str_def = guard(get_definition, self.state.func_ir, f_def.value)
if str_def is None: # TODO: check for errors
raise ValueError("invalid series.str")
series_var = str_def.value
# functions which are used from Numba directly by calling from StringMethodsType
# other functions (for example, 'capitalize' is not presented in Numba) goes to be replaced here
if func_name not in sdc.hiframes.pd_series_ext.str2str_methods_excluded:
return self._run_series_str_method(assign, assign.target, series_var, func_name, rhs)
# replace _get_type_max_value(arr.dtype) since parfors
# arr.dtype transformation produces invalid code for dt64
# TODO: min
if fdef == ('_get_type_max_value', 'sdc.hiframes.hiframes_typed'):
if self.state.typemap[rhs.args[0].name] == types.DType(types.NPDatetime('ns')):
def _run_series_str_replace(self, assign, lhs, arr, rhs, nodes):
regex = True
# TODO: refactor arg parsing
kws = dict(rhs.kws)
if 'regex' in kws:
regex = guard(find_const, self.state.func_ir, kws['regex'])
if regex is None:
raise ValueError(
"str.replace regex argument should be constant")
impl = (series_kernels._str_replace_regex_impl if regex
else series_kernels._str_replace_noregex_impl)
return self._replace_func(
impl,
[arr, rhs.args[0], rhs.args[1]], pre_nodes=nodes,
extra_globals={'unicode_to_std_str': unicode_to_std_str,
'std_str_to_unicode': std_str_to_unicode,
'pre_alloc_string_array': pre_alloc_string_array,
'get_utf8_size': get_utf8_size,
're': re}
)
f = loc_vars['f']
f_ir = compile_to_numba_ir(
f, {'numba': numba, 'np': np, 'Row': Row, 'sdc': sdc})
# fix definitions to enable finding sentinel
f_ir._definitions = build_definitions(f_ir.blocks)
topo_order = find_topo_order(f_ir.blocks)
# find sentinel function and replace with user func
for l in topo_order:
block = f_ir.blocks[l]
for i, stmt in enumerate(block.body):
if (isinstance(stmt, ir.Assign)
and isinstance(stmt.value, ir.Expr)
and stmt.value.op == 'call'):
fdef = guard(get_definition, f_ir, stmt.value.func)
if isinstance(fdef, ir.Global) and fdef.name == 'map_func':
update_globals(func, _globals)
inline_closure_call(f_ir, _globals, block, i, func)
# fix the global value to avoid typing errors
fdef.value = 1
fdef.name = 'A'
break
arg_typs = tuple(df_typ.data[df_typ.columns.index(c)] for c in used_cols)
f_typemap, f_return_type, f_calltypes = numba.typed_passes.type_inference_stage(
self.state.typingctx, f_ir, arg_typs, None)
self.state.typemap.update(f_typemap)
self.state.calltypes.update(f_calltypes)
nodes = []
col_vars = [self._get_dataframe_data(df_var, c, nodes) for c in used_cols]
loc = instr.loc
target = instr.target
index = instr.index if isinstance(instr, ir.SetItem) else instr.index_var
value = instr.value
target_typ = self.typemap[target.name]
index_typ = self.typemap[index.name]
value_typ = self.typemap[value.name]
if isinstance(target_typ, types.npytypes.Array):
if (isinstance(index_typ, types.npytypes.Array) and
isinstance(index_typ.dtype, types.Boolean) and
target_typ.ndim == index_typ.ndim):
if isinstance(value_typ, types.Number):
instr = self._setitem_to_parfor(equiv_set,
loc, target, index, value)
elif isinstance(value_typ, types.npytypes.Array):
val_def = guard(get_definition, self.func_ir,
value.name)
if (isinstance(val_def, ir.Expr) and
val_def.op == 'getitem' and
val_def.index.name == index.name):
instr = self._setitem_to_parfor(equiv_set,
loc, target, index, val_def.value)
else:
shape = equiv_set.get_shape(instr)
if shape != None:
instr = self._setitem_to_parfor(equiv_set,
loc, target, index, value, shape=shape)
new_body.append(instr)
block.body = new_body
def _get_dataframe_data(self, df_var, col_name, nodes):
# optimization: return data var directly if not ambiguous
# (no multiple init_dataframe calls for the same df_var with control
# flow)
# e.g. A = init_dataframe(A, None, 'A')
# XXX assuming init_dataframe is the only call to create a dataframe
# and dataframe._data is never overwritten
df_typ = self.state.typemap[df_var.name]
ind = df_typ.columns.index(col_name)
var_def = guard(get_definition, self.state.func_ir, df_var)
call_def = guard(find_callname, self.state.func_ir, var_def)
if call_def == ('init_dataframe', 'sdc.hiframes.pd_dataframe_ext'):
return var_def.args[ind]
loc = df_var.loc
ind_var = ir.Var(df_var.scope, mk_unique_var('col_ind'), loc)
self.state.typemap[ind_var.name] = types.IntegerLiteral(ind)
nodes.append(ir.Assign(ir.Const(ind, loc), ind_var, loc))
# XXX use get_series_data() for getting data instead of S._data
# to enable alias analysis
f_block = compile_to_numba_ir(
lambda df, c_ind: sdc.hiframes.pd_dataframe_ext.get_dataframe_data(
df, c_ind),
{'sdc': sdc},
self.state.typingctx,
(df_typ, self.state.typemap[ind_var.name]),
def _get_dataframe_index(self, df_var, nodes):
df_typ = self.state.typemap[df_var.name]
n_cols = len(df_typ.columns)
var_def = guard(get_definition, self.state.func_ir, df_var)
call_def = guard(find_callname, self.state.func_ir, var_def)
if call_def == ('init_dataframe', 'sdc.hiframes.pd_dataframe_ext'):
return var_def.args[n_cols]
# XXX use get_series_data() for getting data instead of S._data
# to enable alias analysis
f_block = compile_to_numba_ir(
lambda df: sdc.hiframes.pd_dataframe_ext.get_dataframe_index(df),
{'sdc': sdc},
self.state.typingctx,
(df_typ,),
self.state.typemap,
self.state.calltypes
).blocks.popitem()[1]
replace_arg_nodes(f_block, [df_var])
nodes += f_block.body[:-2]
return nodes[-1].target
def _get_dataframe_index(self, df_var, nodes):
df_typ = self.state.typemap[df_var.name]
n_cols = len(df_typ.columns)
var_def = guard(get_definition, self.state.func_ir, df_var)
call_def = guard(find_callname, self.state.func_ir, var_def)
if call_def == ('init_dataframe', 'sdc.hiframes.pd_dataframe_ext'):
return var_def.args[n_cols]
# XXX use get_series_data() for getting data instead of S._data
# to enable alias analysis
f_block = compile_to_numba_ir(
lambda df: sdc.hiframes.pd_dataframe_ext.get_dataframe_index(df),
{'sdc': sdc},
self.state.typingctx,
(df_typ,),
self.state.typemap,
self.state.calltypes
).blocks.popitem()[1]
replace_arg_nodes(f_block, [df_var])
nodes += f_block.body[:-2]
return nodes[-1].target
def _run_setitem(self, inst):
target_typ = self.state.typemap[inst.target.name]
nodes = []
index_var = (inst.index_var if isinstance(inst, ir.StaticSetItem)
else inst.index)
index_typ = self.state.typemap[index_var.name]
if self._is_df_iat_var(inst.target):
df_var = guard(get_definition, self.state.func_ir, inst.target).value
df_typ = self.state.typemap[df_var.name]
val = inst.value
# df.iat[n,1] = 3
if isinstance(index_typ, types.Tuple) and len(index_typ) == 2:
ind_def = guard(get_definition, self.state.func_ir, index_var)
col_ind = guard(find_const, self.state.func_ir, ind_def.items[1])
col_name = df_typ.columns[col_ind]
in_arr = self._get_dataframe_data(df_var, col_name, nodes)
row_ind = ind_def.items[0]
def _impl(A, row_ind, val):
A[row_ind] = val
return self._replace_func(_impl,
[in_arr, row_ind, val], pre_nodes=nodes)
return [inst]