def _add_offset_to_slice(self, slice_var, offset_var, out_nodes, scope,
loc):
if isinstance(slice_var, slice):
f_text = """def f(offset):
return slice({} + offset, {} + offset)
""".format(slice_var.start, slice_var.stop)
loc = {}
exec_(f_text, {}, loc)
f = loc['f']
args = [offset_var]
arg_typs = (types.intp,)
else:
def f(old_slice, offset):
return slice(old_slice.start + offset, old_slice.stop + offset)
args = [slice_var, offset_var]
slice_type = self.typemap[slice_var.name]
arg_typs = (slice_type, types.intp,)
_globals = self.func_ir.func_id.func.__globals__
f_ir = compile_to_numba_ir(f, _globals, self.typingctx, arg_typs,
self.typemap, self.calltypes)
_, block = f_ir.blocks.popitem()
replace_arg_nodes(block, args)
new_index = block.body[-2].value.value
out_nodes.extend(block.body[:-2]) # ignore return nodes
return new_index
def _get_stencil_start_ind(self, start_length, gen_nodes, scope, loc):
if isinstance(start_length, int):
return abs(min(start_length, 0))
def get_start_ind(s_length):
return abs(min(s_length, 0))
f_ir = compile_to_numba_ir(get_start_ind, {}, self.typingctx,
(types.intp,), self.typemap, self.calltypes)
assert len(f_ir.blocks) == 1
block = f_ir.blocks.popitem()[1]
replace_arg_nodes(block, [start_length])
gen_nodes += block.body[:-2]
ret_var = block.body[-2].value.value
return ret_var
def sort_distributed_run(sort_node, array_dists, typemap, calltypes, typingctx,
targetctx, dist_pass):
parallel = True
in_vars = list(sort_node.df_in_vars.values())
out_vars = list(sort_node.df_out_vars.values())
for v in sort_node.key_arrs + sort_node.out_key_arrs + in_vars + out_vars:
if (array_dists[v.name] != distributed.Distribution.OneD
and array_dists[v.name] != distributed.Distribution.OneD_Var):
parallel = False
loc = sort_node.loc
scope = sort_node.key_arrs[0].scope
# copy arrays when not inplace
nodes = []
key_arrs = sort_node.key_arrs
if not sort_node.inplace:
new_keys = []
for v in key_arrs:
new_key = _copy_array_nodes(v, nodes, typingctx, typemap,
calltypes)
new_keys.append(new_key)
key_arrs = new_keys
new_in_vars = []
for v in in_vars:
v_cp = _copy_array_nodes(v, nodes, typingctx, typemap, calltypes)
new_in_vars.append(v_cp)
in_vars = new_in_vars
key_name_args = ', '.join("key" + str(i) for i in range(len(key_arrs)))
col_name_args = ', '.join(["c" + str(i) for i in range(len(in_vars))])
# TODO: use *args
func_text = "def f({}, {}):\n".format(key_name_args, col_name_args)
func_text += " key_arrs = ({},)\n".format(key_name_args)
func_text += " data = ({}{})\n".format(
col_name_args, "," if len(in_vars) == 1 else
"") # single value needs comma to become tuple
func_text += " hpat.hiframes.sort.local_sort(key_arrs, data, {})\n".format(
sort_node.ascending)
func_text += " return key_arrs, data\n"
loc_vars = {}
exec(func_text, {}, loc_vars)
sort_impl = loc_vars['f']
key_typ = types.Tuple([typemap[v.name] for v in key_arrs])
data_tup_typ = types.Tuple([typemap[v.name] for v in in_vars])
f_block = compile_to_numba_ir(
sort_impl, {
'hpat': hpat,
'to_string_list': to_string_list,
'cp_str_list_to_array': cp_str_list_to_array
}, typingctx, tuple(list(key_typ.types) + list(data_tup_typ.types)),
typemap, calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, key_arrs + in_vars)
nodes += f_block.body[:-2]
ret_var = nodes[-1].target
# get key tup
key_arrs_tup_var = ir.Var(scope, mk_unique_var('key_data'), loc)
typemap[key_arrs_tup_var.name] = key_typ
gen_getitem(key_arrs_tup_var, ret_var, 0, calltypes, nodes)
# get data tup
data_tup_var = ir.Var(scope, mk_unique_var('sort_data'), loc)
typemap[data_tup_var.name] = data_tup_typ
gen_getitem(data_tup_var, ret_var, 1, calltypes, nodes)
if not parallel:
for i, var in enumerate(sort_node.out_key_arrs):
gen_getitem(var, key_arrs_tup_var, i, calltypes, nodes)
for i, var in enumerate(out_vars):
gen_getitem(var, data_tup_var, i, calltypes, nodes)
return nodes
ascending = sort_node.ascending
# parallel case
def par_sort_impl(key_arrs, data):
out_key, out_data = parallel_sort(key_arrs, data, ascending)
# TODO: use k-way merge instead of sort
# sort output
hpat.hiframes.sort.local_sort(out_key, out_data)
return out_key, out_data
f_block = compile_to_numba_ir(
par_sort_impl, {
'hpat': hpat,
'parallel_sort': parallel_sort,
'to_string_list': to_string_list,
'cp_str_list_to_array': cp_str_list_to_array,
'ascending': ascending
}, typingctx, (key_typ, data_tup_typ), typemap,
calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [key_arrs_tup_var, data_tup_var])
nodes += f_block.body[:-2]
ret_var = nodes[-1].target
# get output key
key_tup = ir.Var(scope, mk_unique_var('sort_keys'), loc)
typemap[key_tup.name] = key_typ
gen_getitem(key_tup, ret_var, 0, calltypes, nodes)
# get data tup
data_tup = ir.Var(scope, mk_unique_var('sort_data'), loc)
typemap[data_tup.name] = data_tup_typ
gen_getitem(data_tup, ret_var, 1, calltypes, nodes)
for i, var in enumerate(sort_node.out_key_arrs):
gen_getitem(var, key_tup, i, calltypes, nodes)
for i, var in enumerate(out_vars):
gen_getitem(var, data_tup, i, calltypes, nodes)
# TODO: handle 1D balance for inplace case
return nodes
def gen_parquet_read(self, file_name, lhs):
scope = file_name.scope
loc = file_name.loc
table_types = None
# lhs is temporary and will possibly be assigned to user variable
assert lhs.name.startswith('$')
if lhs.name in self.reverse_copies and self.reverse_copies[
lhs.name] in self.locals:
table_types = self.locals[self.reverse_copies[lhs.name]]
self.locals.pop(self.reverse_copies[lhs.name])
convert_types = {}
# user-specified type conversion
if lhs.name in self.reverse_copies and (self.reverse_copies[lhs.name] +
':convert') in self.locals:
convert_types = self.locals[self.reverse_copies[lhs.name] +
':convert']
self.locals.pop(self.reverse_copies[lhs.name] + ':convert')
if table_types is None:
fname_def = guard(get_definition, self.func_ir, file_name)
if (not isinstance(fname_def, (ir.Const, ir.Global, ir.FreeVar))
or not isinstance(fname_def.value, str)):
raise ValueError("Parquet schema not available")
file_name_str = fname_def.value
col_names, col_types = parquet_file_schema(file_name_str)
# remove Pandas index if exists
# TODO: handle index properly when indices are supported
_rm_pd_index(col_names, col_types)
else:
col_names = list(table_types.keys())
col_types = list(table_types.values())
out_nodes = []
# get arrow readers once
def init_arrow_readers(fname):
arrow_readers = get_arrow_readers(unicode_to_char_ptr(fname))
f_block = compile_to_numba_ir(
init_arrow_readers, {
'get_arrow_readers': _get_arrow_readers,
'unicode_to_char_ptr': unicode_to_char_ptr,
}).blocks.popitem()[1]
replace_arg_nodes(f_block, [file_name])
out_nodes += f_block.body[:-3]
arrow_readers_var = out_nodes[-1].target
col_arrs = []
for i, cname in enumerate(col_names):
# get column type from schema
c_type = col_types[i]
if cname in convert_types:
c_type = convert_types[cname].dtype
# create a variable for column and assign type
varname = mk_unique_var(cname)
#self.locals[varname] = c_type
cvar = ir.Var(scope, varname, loc)
col_arrs.append(cvar)
out_nodes += get_column_read_nodes(c_type, cvar, arrow_readers_var,
i)
# delete arrow readers
def cleanup_arrow_readers(readers):
s = del_arrow_readers(readers)
f_block = compile_to_numba_ir(cleanup_arrow_readers, {
'del_arrow_readers': _del_arrow_readers,
}).blocks.popitem()[1]
replace_arg_nodes(f_block, [arrow_readers_var])
out_nodes += f_block.body[:-3]
return col_names, col_arrs, out_nodes
def join_distributed_run(join_node, array_dists, typemap, calltypes, typingctx, targetctx):
parallel = True
for v in (list(join_node.left_vars.values())
+ list(join_node.right_vars.values())
+ list(join_node.df_out_vars.values())):
if (array_dists[v.name] != distributed.Distribution.OneD
and array_dists[v.name] != distributed.Distribution.OneD_Var):
parallel = False
# TODO: rebalance if output distributions are 1D instead of 1D_Var
loc = join_node.loc
# get column variables
left_key_var = join_node.left_vars[join_node.left_key]
right_key_var = join_node.right_vars[join_node.right_key]
left_other_col_vars = [v for (n, v) in sorted(join_node.left_vars.items())
if n != join_node.left_key]
right_other_col_vars = [v for (n, v) in sorted(join_node.right_vars.items())
if n != join_node.right_key]
# get column types
left_other_col_typ = [typemap[v.name] for v in left_other_col_vars]
right_other_col_typ = [typemap[v.name] for v in right_other_col_vars]
arg_typs = tuple([typemap[left_key_var.name], typemap[right_key_var.name]]
+ left_other_col_typ + right_other_col_typ)
# arg names of non-key columns
left_other_names = ["t1_c" + str(i)
for i in range(len(left_other_col_vars))]
right_other_names = ["t2_c" + str(i)
for i in range(len(right_other_col_vars))]
# all arg names
left_arg_names = ['t1_key'] + left_other_names
right_arg_names = ['t2_key'] + right_other_names
func_text = "def f(t1_key, t2_key,{}{}{}):\n".format(
",".join(left_other_names),
("," if len(left_other_names) != 0 else ""),
",".join(right_other_names))
func_text += " data_left = ({}{})\n".format(",".join(left_other_names),
"," if len(left_other_names) != 0 else "")
func_text += " data_right = ({}{})\n".format(",".join(right_other_names),
"," if len(right_other_names) != 0 else "")
if parallel:
func_text += " t1_key, data_left = parallel_join(t1_key, data_left)\n"
#func_text += " print(t2_key, data_right)\n"
func_text += " t2_key, data_right = parallel_join(t2_key, data_right)\n"
#func_text += " print(t2_key, data_right)\n"
local_left_data = "t1_key" + (", " if len(left_other_names) != 0 else "") + ",".join(["data_left[{}]".format(i) for i in range(len(left_other_names))])
local_right_data = "t2_key" + (", " if len(right_other_names) != 0 else "") + ",".join(["data_right[{}]".format(i) for i in range(len(right_other_names))])
else:
local_left_data = ",".join(left_arg_names)
local_right_data = ",".join(right_arg_names)
# local sort
func_text += " local_sort_f1(t1_key, data_left)\n"
func_text += " local_sort_f2(t2_key, data_right)\n"
# align output variables for local merge
# add keys first (TODO: remove dead keys)
merge_out = [join_node.df_out_vars[join_node.left_key]]
merge_out.append(join_node.df_out_vars[join_node.right_key])
merge_out += [join_node.df_out_vars[n] for (n, v) in sorted(join_node.left_vars.items())
if n != join_node.left_key]
merge_out += [join_node.df_out_vars[n] for (n, v) in sorted(join_node.right_vars.items())
if n != join_node.right_key]
out_names = ["t3_c" + str(i) for i in range(len(merge_out))]
func_text += " out_t1_key, out_t2_key, out_data_left, out_data_right = hpat.hiframes_join.local_merge_new(t1_key, t2_key, data_left, data_right)\n"
for i in range(len(left_other_names)):
func_text += " left_{} = out_data_left[{}]\n".format(i, i)
for i in range(len(right_other_names)):
func_text += " right_{} = out_data_right[{}]\n".format(i, i)
func_text += " {} = out_t1_key\n".format(out_names[0])
func_text += " {} = out_t2_key\n".format(out_names[1])
for i in range(len(left_other_names)):
func_text += " {} = left_{}\n".format(out_names[i+2], i)
for i in range(len(right_other_names)):
func_text += " {} = right_{}\n".format(out_names[i+2+len(left_other_names)], i)
# func_text += " {} = hpat.hiframes_join.local_merge({}, {}, {})\n".format(
# ",".join(out_names), len(left_arg_names),
# local_left_data, local_right_data)
loc_vars = {}
exec(func_text, {}, loc_vars)
join_impl = loc_vars['f']
left_data_tup_typ = types.Tuple([typemap[v.name] for v in left_other_col_vars])
_local_sort_f1 = hpat.hiframes_sort.get_local_sort_func(typemap[left_key_var.name], left_data_tup_typ)
right_data_tup_typ = types.Tuple([typemap[v.name] for v in right_other_col_vars])
_local_sort_f2 = hpat.hiframes_sort.get_local_sort_func(typemap[right_key_var.name], right_data_tup_typ)
f_block = compile_to_numba_ir(join_impl,
{'hpat': hpat, 'np': np,
'to_string_list': to_string_list,
'cp_str_list_to_array': cp_str_list_to_array,
'local_sort_f1': _local_sort_f1,
'local_sort_f2': _local_sort_f2,
'parallel_join': parallel_join},
typingctx, arg_typs,
typemap, calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [left_key_var, right_key_var]
+ left_other_col_vars + right_other_col_vars)
nodes = f_block.body[:-3]
for i in range(len(merge_out)):
nodes[-len(merge_out) + i].target = merge_out[i]
return nodes
def join_distributed_run(join_node, array_dists, typemap, calltypes,
typingctx):
parallel = True
for v in (list(join_node.left_vars.values()) +
list(join_node.right_vars.values()) +
list(join_node.df_out_vars.values())):
if (array_dists[v.name] != distributed.Distribution.OneD
and array_dists[v.name] != distributed.Distribution.OneD_Var):
parallel = False
if (typemap[v.name] != types.Array(types.intp, 1, 'C')
and typemap[v.name] != types.Array(types.float64, 1, 'C')):
raise ValueError(
"Only int64 and float64 columns are currently supported in join"
)
# TODO: rebalance if output distributions are 1D instead of 1D_Var
loc = join_node.loc
# get column variables
left_key_var = join_node.left_vars[join_node.left_key]
right_key_var = join_node.right_vars[join_node.right_key]
if (typemap[left_key_var.name] != types.Array(types.intp, 1, 'C')
or typemap[right_key_var.name] != types.Array(types.intp, 1, 'C')):
raise ValueError("Only int64 keys are currently supported in join")
left_other_col_vars = [
v for (n, v) in sorted(join_node.left_vars.items())
if n != join_node.left_key
]
right_other_col_vars = [
v for (n, v) in sorted(join_node.right_vars.items())
if n != join_node.right_key
]
# get column types
left_other_col_typ = [typemap[v.name] for v in left_other_col_vars]
right_other_col_typ = [typemap[v.name] for v in right_other_col_vars]
arg_typs = tuple(
[typemap[left_key_var.name], typemap[right_key_var.name]] +
left_other_col_typ + right_other_col_typ)
# arg names of non-key columns
left_other_names = [
"t1_c" + str(i) for i in range(len(left_other_col_vars))
]
right_other_names = [
"t2_c" + str(i) for i in range(len(right_other_col_vars))
]
# all arg names
left_arg_names = ['t1_key'] + left_other_names
right_arg_names = ['t2_key'] + right_other_names
func_text = "def f(t1_key, t2_key,{}{}{}):\n".format(
",".join(left_other_names),
("," if len(left_other_names) != 0 else ""),
",".join(right_other_names))
if parallel:
# get send/recv counts
func_text += " (t1_send_counts, t1_recv_counts, t1_send_disp,\n"
func_text += " t1_recv_disp, t1_recv_size) = hpat.hiframes_join.get_sendrecv_counts(t1_key)\n"
func_text += " (t2_send_counts, t2_recv_counts, t2_send_disp,\n"
func_text += " t2_recv_disp, t2_recv_size) = hpat.hiframes_join.get_sendrecv_counts(t2_key)\n"
#func_text += " hpat.cprint(t1_recv_size, t2_recv_size)\n"
# prepare for shuffle
# allocate send/recv buffers
for a in left_arg_names:
func_text += " send_{} = np.empty_like({})\n".format(a, a)
func_text += " recv_{} = np.empty(t1_recv_size, {}.dtype)\n".format(
a, a)
for a in right_arg_names:
func_text += " send_{} = np.empty_like({})\n".format(a, a)
func_text += " recv_{} = np.empty(t2_recv_size, {}.dtype)\n".format(
a, a)
left_send_names = ",".join(["send_" + v for v in left_arg_names])
left_recv_names = ",".join(["recv_" + v for v in left_arg_names])
func_text += (
" hpat.hiframes_join.shuffle_data(t1_send_counts," +
" t1_recv_counts, t1_send_disp, t1_recv_disp, {},{},{})\n".format(
",".join(left_arg_names), left_send_names, left_recv_names))
right_send_names = ",".join(["send_" + v for v in right_arg_names])
right_recv_names = ",".join(["recv_" + v for v in right_arg_names])
func_text += (
" hpat.hiframes_join.shuffle_data(t2_send_counts," +
" t2_recv_counts, t2_send_disp, t2_recv_disp, {},{},{})\n".format(
",".join(right_arg_names), right_send_names, right_recv_names))
local_left_data = left_recv_names
local_right_data = right_recv_names
else:
local_left_data = ",".join(left_arg_names)
local_right_data = ",".join(right_arg_names)
func_text += " hpat.hiframes_join.sort({})\n".format(local_left_data)
func_text += " hpat.hiframes_join.sort({})\n".format(local_right_data)
# align output variables for local merge
# add keys first (TODO: remove dead keys)
merge_out = [join_node.df_out_vars[join_node.left_key]]
merge_out += [
join_node.df_out_vars[n]
for (n, v) in sorted(join_node.left_vars.items())
if n != join_node.left_key
]
merge_out.append(join_node.df_out_vars[join_node.right_key])
merge_out += [
join_node.df_out_vars[n]
for (n, v) in sorted(join_node.right_vars.items())
if n != join_node.right_key
]
out_names = ["t3_c" + str(i) for i in range(len(merge_out))]
func_text += " {} = hpat.hiframes_join.local_merge({}, {}, {})\n".format(
",".join(out_names), len(left_arg_names), local_left_data,
local_right_data)
# TODO: delete buffers
#delete_buffers((t1_send_counts, t1_recv_counts, t1_send_disp, t1_recv_disp))
#delete_buffers((t2_send_counts, t2_recv_counts, t2_send_disp, t2_recv_disp))
loc_vars = {}
exec(func_text, {}, loc_vars)
f = loc_vars['f']
f_block = compile_to_numba_ir(f, {
'hpat': hpat,
'np': np
}, typingctx, arg_typs, typemap, calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [left_key_var, right_key_var] +
left_other_col_vars + right_other_col_vars)
nodes = f_block.body[:-3]
for i in range(len(merge_out)):
nodes[-len(merge_out) + i].target = merge_out[i]
return nodes
def sort_distributed_run(sort_node, array_dists, typemap, calltypes, typingctx,
targetctx):
parallel = True
data_vars = list(sort_node.df_vars.values())
for v in [sort_node.key_arr] + data_vars:
if (array_dists[v.name] != distributed.Distribution.OneD
and array_dists[v.name] != distributed.Distribution.OneD_Var):
parallel = False
key_arr = sort_node.key_arr
col_name_args = ', '.join(["c" + str(i) for i in range(len(data_vars))])
# TODO: use *args
func_text = "def f(key_arr, {}):\n".format(col_name_args)
func_text += " data = ({}{})\n".format(
col_name_args, "," if len(data_vars) == 1 else
"") # single value needs comma to become tuple
func_text += " local_sort_f(key_arr, data)\n"
loc_vars = {}
exec(func_text, {}, loc_vars)
sort_impl = loc_vars['f']
key_typ = typemap[key_arr.name]
data_tup_typ = types.Tuple(
[typemap[v.name] for v in sort_node.df_vars.values()])
_local_sort_f = get_local_sort_func(key_typ, data_tup_typ)
f_block = compile_to_numba_ir(
sort_impl, {
'hpat': hpat,
'local_sort_f': _local_sort_f,
'to_string_list': to_string_list,
'cp_str_list_to_array': cp_str_list_to_array
}, typingctx, tuple([key_typ] + list(data_tup_typ.types)), typemap,
calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [sort_node.key_arr] + data_vars)
nodes = f_block.body[:-3]
if not parallel:
return nodes
# parallel case
# TODO: refactor with previous call, use *args?
# get data variable tuple
func_text = "def f({}):\n".format(col_name_args)
func_text += " data = ({}{})\n".format(
col_name_args, "," if len(data_vars) == 1 else
"") # single value needs comma to become tuple
loc_vars = {}
exec(func_text, {}, loc_vars)
tup_impl = loc_vars['f']
f_block = compile_to_numba_ir(tup_impl, {}, typingctx,
list(data_tup_typ.types), typemap,
calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, data_vars)
nodes += f_block.body[:-3]
data_tup_var = nodes[-1].target
def par_sort_impl(key_arr, data):
out, out_data = parallel_sort(key_arr, data)
# TODO: use k-way merge instead of sort
# sort output
local_sort_f(out, out_data)
res_data = out_data
res = out
f_block = compile_to_numba_ir(
par_sort_impl, {
'hpat': hpat,
'parallel_sort': parallel_sort,
'to_string_list': to_string_list,
'cp_str_list_to_array': cp_str_list_to_array,
'local_sort_f': _local_sort_f
}, typingctx, (key_typ, data_tup_typ), typemap,
calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [sort_node.key_arr, data_tup_var])
nodes += f_block.body[:-3]
# set vars since new arrays are created after communication
data_tup = nodes[-2].target
# key
nodes.append(
ir.Assign(nodes[-1].target, sort_node.key_arr, sort_node.key_arr.loc))
for i, var in enumerate(data_vars):
getitem = ir.Expr.static_getitem(data_tup, i, None, var.loc)
calltypes[getitem] = None
nodes.append(ir.Assign(getitem, var, var.loc))
return nodes
def join_distributed_run(join_node, array_dists, typemap, calltypes, typingctx,
targetctx, dist_pass):
left_parallel, right_parallel = _get_table_parallel_flags(
join_node, array_dists)
method = 'hash'
# method = 'sort'
# TODO: rebalance if output distributions are 1D instead of 1D_Var
loc = join_node.loc
n_keys = len(join_node.left_keys)
# get column variables
left_key_vars = tuple(join_node.left_vars[c] for c in join_node.left_keys)
right_key_vars = tuple(join_node.right_vars[c]
for c in join_node.right_keys)
left_other_col_vars = tuple(v
for (n,
v) in sorted(join_node.left_vars.items())
if n not in join_node.left_keys)
right_other_col_vars = tuple(
v for (n, v) in sorted(join_node.right_vars.items())
if n not in join_node.right_keys)
# get column types
arg_vars = (left_key_vars + right_key_vars + left_other_col_vars +
right_other_col_vars)
arg_typs = tuple(typemap[v.name] for v in arg_vars)
scope = arg_vars[0].scope
# arg names of non-key columns
left_other_names = tuple("t1_c" + str(i)
for i in range(len(left_other_col_vars)))
right_other_names = tuple("t2_c" + str(i)
for i in range(len(right_other_col_vars)))
left_key_names = tuple("t1_key" + str(i) for i in range(n_keys))
right_key_names = tuple("t2_key" + str(i) for i in range(n_keys))
func_text = "def f({}, {},{}{}{}):\n".format(
",".join(left_key_names), ",".join(right_key_names),
",".join(left_other_names),
("," if len(left_other_names) != 0 else ""),
",".join(right_other_names))
func_text += " t1_keys = ({},)\n".format(",".join(left_key_names))
func_text += " t2_keys = ({},)\n".format(",".join(right_key_names))
func_text += " data_left = ({}{})\n".format(
",".join(left_other_names), "," if len(left_other_names) != 0 else "")
func_text += " data_right = ({}{})\n".format(
",".join(right_other_names),
"," if len(right_other_names) != 0 else "")
if join_node.how == 'asof':
if left_parallel or right_parallel:
assert left_parallel and right_parallel
# only the right key needs to be aligned
func_text += " t2_keys, data_right = parallel_asof_comm(t1_keys, t2_keys, data_right)\n"
else:
if left_parallel:
func_text += " t1_keys, data_left = parallel_join(t1_keys, data_left)\n"
if right_parallel:
func_text += " t2_keys, data_right = parallel_join(t2_keys, data_right)\n"
#func_text += " print(t2_key, data_right)\n"
if method == 'sort' and join_node.how != 'asof':
# asof key is already sorted, TODO: add error checking
# local sort
func_text += " hpat.hiframes.sort.local_sort(t1_keys, data_left)\n"
func_text += " hpat.hiframes.sort.local_sort(t2_keys, data_right)\n"
# align output variables for local merge
# add keys first (TODO: remove dead keys)
out_l_key_vars = tuple(join_node.df_out_vars[c]
for c in join_node.left_keys)
out_r_key_vars = tuple(join_node.df_out_vars[c]
for c in join_node.right_keys)
# create dummy variable if right key is not actually returned
# using the same output left key causes errors for asof case
if join_node.left_keys == join_node.right_keys:
out_r_key_vars = tuple(
ir.Var(scope, mk_unique_var('dummy_k'), loc)
for _ in range(n_keys))
for v, w in zip(out_r_key_vars, out_l_key_vars):
typemap[v.name] = typemap[w.name]
merge_out = out_l_key_vars + out_r_key_vars
merge_out += tuple(join_node.df_out_vars[n]
for (n, v) in sorted(join_node.left_vars.items())
if n not in join_node.left_keys)
merge_out += tuple(join_node.df_out_vars[n]
for (n, v) in sorted(join_node.right_vars.items())
if n not in join_node.right_keys)
out_names = ["t3_c" + str(i) for i in range(len(merge_out))]
if join_node.how == 'asof':
func_text += (
" out_t1_keys, out_t2_keys, out_data_left, out_data_right"
" = hpat.hiframes.join.local_merge_asof(t1_keys, t2_keys, data_left, data_right)\n"
)
elif method == 'sort':
func_text += (
" out_t1_keys, out_t2_keys, out_data_left, out_data_right"
" = hpat.hiframes.join.local_merge_new(t1_keys, t2_keys, data_left, data_right, {}, {})\n"
.format(join_node.how in ('left', 'outer'),
join_node.how == 'outer'))
else:
assert method == 'hash'
func_text += (
" out_t1_keys, out_t2_keys, out_data_left, out_data_right"
" = hpat.hiframes.join.local_hash_join(t1_keys, t2_keys, data_left, data_right, {}, {})\n"
.format(join_node.how in ('left', 'outer'),
join_node.how == 'outer'))
for i in range(len(left_other_names)):
func_text += " left_{} = out_data_left[{}]\n".format(i, i)
for i in range(len(right_other_names)):
func_text += " right_{} = out_data_right[{}]\n".format(i, i)
for i in range(n_keys):
func_text += " t1_keys_{} = out_t1_keys[{}]\n".format(i, i)
for i in range(n_keys):
func_text += " t2_keys_{} = out_t2_keys[{}]\n".format(i, i)
for i in range(n_keys):
func_text += " {} = t1_keys_{}\n".format(out_names[i], i)
for i in range(n_keys):
func_text += " {} = t2_keys_{}\n".format(out_names[n_keys + i], i)
for i in range(len(left_other_names)):
func_text += " {} = left_{}\n".format(out_names[i + 2 * n_keys], i)
for i in range(len(right_other_names)):
func_text += " {} = right_{}\n".format(
out_names[i + 2 * n_keys + len(left_other_names)], i)
loc_vars = {}
exec(func_text, {}, loc_vars)
join_impl = loc_vars['f']
# print(func_text)
glbs = {
'hpat': hpat,
'np': np,
'to_string_list': to_string_list,
'cp_str_list_to_array': cp_str_list_to_array,
'parallel_join': parallel_join,
'parallel_asof_comm': parallel_asof_comm
}
f_block = compile_to_numba_ir(join_impl, glbs, typingctx, arg_typs,
typemap, calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, arg_vars)
nodes = f_block.body[:-3]
for i in range(len(merge_out)):
nodes[-len(merge_out) + i].target = merge_out[i]
return nodes
def _gen_rolling_call(self, args, col_var, win_size, center, func,
out_var):
loc = col_var.loc
scope = col_var.scope
if func == 'apply':
if len(args) != 1:
raise ValueError("One argument expected for rolling apply")
kernel_func = guard(get_definition, self.func_ir, args[0])
elif func in ['sum', 'mean', 'min', 'max', 'std', 'var']:
if len(args) != 0:
raise ValueError(
"No argument expected for rolling {}".format(func))
g_pack = "np"
if func in ['std', 'var']:
g_pack = "hpat.hiframes_api"
if isinstance(win_size, int) and win_size < LARGE_WIN_SIZE:
# unroll if size is less than 5
kernel_args = ','.join(
['a[{}]'.format(-i) for i in range(win_size)])
kernel_expr = '{}.{}(np.array([{}]))'.format(
g_pack, func, kernel_args)
if func == 'sum': # simplify sum
kernel_expr = '+'.join(
['a[{}]'.format(-i) for i in range(win_size)])
else:
kernel_expr = '{}.{}(a[(-w+1):1])'.format(g_pack, func)
func_text = 'def g(a, w):\n return {}\n'.format(kernel_expr)
loc_vars = {}
exec(func_text, {}, loc_vars)
kernel_func = loc_vars['g']
init_nodes = []
if isinstance(win_size, int):
win_size_var = ir.Var(scope, mk_unique_var("win_size"), loc)
init_nodes.append(
ir.Assign(ir.Const(win_size, loc), win_size_var, loc))
win_size = win_size_var
index_offsets, win_tuple, option_nodes = self._gen_rolling_init(
win_size, func, center)
init_nodes += option_nodes
other_args = [win_size]
if func == 'apply':
other_args = None
options = {'neighborhood': win_tuple}
fir_globals = self.func_ir.func_id.func.__globals__
stencil_nodes = gen_stencil_call(col_var, out_var, kernel_func,
index_offsets, fir_globals,
other_args, options)
def f(A, w):
A[:w - 1] = np.nan
f_block = compile_to_numba_ir(f, {'np': np}).blocks.popitem()[1]
replace_arg_nodes(f_block, [out_var, win_size])
setitem_nodes = f_block.body[:-3] # remove none return
if center:
def f1(A, w):
A[:w // 2] = np.nan
def f2(A, w):
A[-(w // 2):] = np.nan
f_block = compile_to_numba_ir(f1, {'np': np}).blocks.popitem()[1]
replace_arg_nodes(f_block, [out_var, win_size])
setitem_nodes1 = f_block.body[:-3] # remove none return
f_block = compile_to_numba_ir(f2, {'np': np}).blocks.popitem()[1]
replace_arg_nodes(f_block, [out_var, win_size])
setitem_nodes2 = f_block.body[:-3] # remove none return
setitem_nodes = setitem_nodes1 + setitem_nodes2
return init_nodes + stencil_nodes + setitem_nodes
def join_distributed_run(join_node, array_dists, typemap, calltypes, typingctx,
targetctx):
parallel = True
for v in (list(join_node.left_vars.values()) +
list(join_node.right_vars.values()) +
list(join_node.df_out_vars.values())):
if (array_dists[v.name] != distributed.Distribution.OneD
and array_dists[v.name] != distributed.Distribution.OneD_Var):
parallel = False
# TODO: rebalance if output distributions are 1D instead of 1D_Var
loc = join_node.loc
# get column variables
left_key_var = join_node.left_vars[join_node.left_key]
right_key_var = join_node.right_vars[join_node.right_key]
left_other_col_vars = [
v for (n, v) in sorted(join_node.left_vars.items())
if n != join_node.left_key
]
right_other_col_vars = [
v for (n, v) in sorted(join_node.right_vars.items())
if n != join_node.right_key
]
# get column types
left_other_col_typ = [typemap[v.name] for v in left_other_col_vars]
right_other_col_typ = [typemap[v.name] for v in right_other_col_vars]
arg_typs = tuple(
[typemap[left_key_var.name], typemap[right_key_var.name]] +
left_other_col_typ + right_other_col_typ)
# arg names of non-key columns
left_other_names = [
"t1_c" + str(i) for i in range(len(left_other_col_vars))
]
right_other_names = [
"t2_c" + str(i) for i in range(len(right_other_col_vars))
]
func_text = "def f(t1_key, t2_key,{}{}{}):\n".format(
",".join(left_other_names),
("," if len(left_other_names) != 0 else ""),
",".join(right_other_names))
func_text += " data_left = ({}{})\n".format(
",".join(left_other_names), "," if len(left_other_names) != 0 else "")
func_text += " data_right = ({}{})\n".format(
",".join(right_other_names),
"," if len(right_other_names) != 0 else "")
if parallel:
if join_node.how == 'asof':
# only the right key needs to be aligned
func_text += " t2_key, data_right = parallel_asof_comm(t1_key, t2_key, data_right)\n"
else:
func_text += " t1_key, data_left = parallel_join(t1_key, data_left)\n"
func_text += " t2_key, data_right = parallel_join(t2_key, data_right)\n"
#func_text += " print(t2_key, data_right)\n"
if join_node.how != 'asof':
# asof key is already sorted, TODO: add error checking
# local sort
func_text += " local_sort_f1(t1_key, data_left)\n"
func_text += " local_sort_f2(t2_key, data_right)\n"
# align output variables for local merge
# add keys first (TODO: remove dead keys)
out_l_key_var = join_node.df_out_vars[join_node.left_key]
out_r_key_var = join_node.df_out_vars[join_node.right_key]
# create dummy variable if right key is not actually returned
# using the same output left key causes errors for asof case
if join_node.left_key == join_node.right_key:
out_r_key_var = ir.Var(out_l_key_var.scope, mk_unique_var('dummy_k'),
loc)
typemap[out_r_key_var.name] = typemap[out_l_key_var.name]
merge_out = [out_l_key_var, out_r_key_var]
merge_out += [
join_node.df_out_vars[n]
for (n, v) in sorted(join_node.left_vars.items())
if n != join_node.left_key
]
merge_out += [
join_node.df_out_vars[n]
for (n, v) in sorted(join_node.right_vars.items())
if n != join_node.right_key
]
out_names = ["t3_c" + str(i) for i in range(len(merge_out))]
if join_node.how == 'asof':
func_text += (
" out_t1_key, out_t2_key, out_data_left, out_data_right"
" = hpat.hiframes_join.local_merge_asof(t1_key, t2_key, data_left, data_right)\n"
)
else:
func_text += (
" out_t1_key, out_t2_key, out_data_left, out_data_right"
" = hpat.hiframes_join.local_merge_new(t1_key, t2_key, data_left, data_right)\n"
)
for i in range(len(left_other_names)):
func_text += " left_{} = out_data_left[{}]\n".format(i, i)
for i in range(len(right_other_names)):
func_text += " right_{} = out_data_right[{}]\n".format(i, i)
func_text += " {} = out_t1_key\n".format(out_names[0])
func_text += " {} = out_t2_key\n".format(out_names[1])
for i in range(len(left_other_names)):
func_text += " {} = left_{}\n".format(out_names[i + 2], i)
for i in range(len(right_other_names)):
func_text += " {} = right_{}\n".format(
out_names[i + 2 + len(left_other_names)], i)
# func_text += " {} = hpat.hiframes_join.local_merge({}, {}, {})\n".format(
# ",".join(out_names), len(left_arg_names),
# local_left_data, local_right_data)
loc_vars = {}
exec(func_text, {}, loc_vars)
join_impl = loc_vars['f']
# print(func_text)
left_data_tup_typ = types.Tuple(
[typemap[v.name] for v in left_other_col_vars])
_local_sort_f1 = hpat.hiframes_sort.get_local_sort_func(
typemap[left_key_var.name], left_data_tup_typ)
right_data_tup_typ = types.Tuple(
[typemap[v.name] for v in right_other_col_vars])
_local_sort_f2 = hpat.hiframes_sort.get_local_sort_func(
typemap[right_key_var.name], right_data_tup_typ)
f_block = compile_to_numba_ir(
join_impl, {
'hpat': hpat,
'np': np,
'to_string_list': to_string_list,
'cp_str_list_to_array': cp_str_list_to_array,
'local_sort_f1': _local_sort_f1,
'local_sort_f2': _local_sort_f2,
'parallel_join': parallel_join,
'parallel_asof_comm': parallel_asof_comm
}, typingctx, arg_typs, typemap, calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [left_key_var, right_key_var] +
left_other_col_vars + right_other_col_vars)
nodes = f_block.body[:-3]
for i in range(len(merge_out)):
nodes[-len(merge_out) + i].target = merge_out[i]
return nodes
def _handle_df_col_calls(self, lhs_name, rhs, assign):
if guard(find_callname, self.func_ir,
rhs) == ('count', 'hpat.hiframes_api'):
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(_column_count_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx, (self.typemap[in_arr.name], ), self.typemap,
self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-4].target = assign.target
return f_blocks
if guard(find_callname, self.func_ir,
rhs) == ('fillna', 'hpat.hiframes_api'):
out_arr = rhs.args[0]
in_arr = rhs.args[1]
val = rhs.args[2]
f_blocks = compile_to_numba_ir(
_column_fillna_impl, {
'numba': numba,
'np': np
}, self.typingctx,
(self.typemap[out_arr.name], self.typemap[in_arr.name],
self.typemap[val.name]), self.typemap, self.calltypes).blocks
first_block = min(f_blocks.keys())
replace_arg_nodes(f_blocks[first_block], [out_arr, in_arr, val])
return f_blocks
if guard(find_callname, self.func_ir,
rhs) == ('column_sum', 'hpat.hiframes_api'):
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(_column_sum_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx, (self.typemap[in_arr.name], ), self.typemap,
self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-4].target = assign.target
return f_blocks
if guard(find_callname, self.func_ir,
rhs) == ('mean', 'hpat.hiframes_api'):
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(_column_mean_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx, (self.typemap[in_arr.name], ), self.typemap,
self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-4].target = assign.target
return f_blocks
if guard(find_callname, self.func_ir,
rhs) == ('var', 'hpat.hiframes_api'):
in_arr = rhs.args[0]
f_blocks = compile_to_numba_ir(_column_var_impl, {
'numba': numba,
'np': np,
'hpat': hpat
}, self.typingctx, (self.typemap[in_arr.name], ), self.typemap,
self.calltypes).blocks
topo_order = find_topo_order(f_blocks)
first_block = topo_order[0]
last_block = topo_order[-1]
replace_arg_nodes(f_blocks[first_block], [in_arr])
# assign results to lhs output
f_blocks[last_block].body[-4].target = assign.target
return f_blocks
return
def _handle_concat(self, lhs, rhs):
if guard(find_callname, self.func_ir, rhs) == ('concat', 'pandas'):
if len(rhs.args) != 1 or len(rhs.kws) != 0:
raise ValueError(
"only a list/tuple argument is supported in concat")
df_list = guard(get_definition, self.func_ir, rhs.args[0])
assert isinstance(df_list, ir.Expr) and df_list.op == 'build_list'
nodes = []
done_cols = {}
i = 0
for df in df_list.items:
for (c, v) in self.df_vars[df.name].items():
if c in done_cols:
continue
# arguments to the generated function
args = [v]
# names of arguments to the generated function
arg_names = ['_hpat_c' + str(i)]
# arguments to the concatenate function
conc_arg_names = ['_hpat_c' + str(i)]
allocs = ""
i += 1
for other_df in df_list.items:
if other_df.name == df.name:
continue
if c in self.df_vars[other_df.name]:
args.append(self.df_vars[other_df.name][c])
arg_names.append('_hpat_c' + str(i))
conc_arg_names.append('_hpat_c' + str(i))
i += 1
else:
# use a df column for length
len_arg = list(
self.df_vars[other_df.name].values())[0]
len_name = '_hpat_len' + str(i)
args.append(len_arg)
arg_names.append(len_name)
i += 1
out_name = '_hpat_out' + str(i)
conc_arg_names.append(out_name)
i += 1
# TODO: fix type
allocs += " {} = np.full(len({}), np.nan)\n".format(
out_name, len_name)
func_text = "def f({}):\n".format(",".join(arg_names))
func_text += allocs
func_text += " s = np.concatenate(({}))\n".format(
",".join(conc_arg_names))
loc_vars = {}
exec(func_text, {}, loc_vars)
f = loc_vars['f']
f_block = compile_to_numba_ir(f, {
'hpat': hpat,
'np': np
}).blocks.popitem()[1]
replace_arg_nodes(f_block, args)
nodes += f_block.body[:-3]
done_cols[c] = nodes[-1].target
self.df_vars[lhs.name] = done_cols
self._update_df_cols()
return nodes
def csv_distributed_run(csv_node, array_dists, typemap, calltypes, typingctx,
targetctx, dist_pass):
parallel = True
if hpat.config.config_transport_mpi:
for v in csv_node.out_vars:
if (array_dists[v.name] != distributed.Distribution.OneD and
array_dists[v.name] != distributed.Distribution.OneD_Var):
parallel = False
else:
parallel = False
n_cols = len(csv_node.out_vars)
# TODO: rebalance if output distributions are 1D instead of 1D_Var
# get column variables
arg_names = ", ".join("arr" + str(i) for i in range(n_cols))
func_text = "def csv_impl(fname):\n"
func_text += " ({},) = _csv_reader_py(fname)\n".format(arg_names)
# print(func_text)
loc_vars = {}
exec(func_text, {}, loc_vars)
csv_impl = loc_vars['csv_impl']
csv_reader_py = _gen_csv_reader_py(csv_node.df_colnames,
csv_node.out_types, csv_node.usecols,
csv_node.sep, typingctx, targetctx,
parallel, csv_node.skiprows)
f_block = compile_to_numba_ir(csv_impl, {
'_csv_reader_py': csv_reader_py
}, typingctx, (string_type, ), typemap, calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [csv_node.file_name])
nodes = f_block.body[:-3]
for i in range(len(csv_node.out_vars)):
nodes[-len(csv_node.out_vars) + i].target = csv_node.out_vars[i]
# get global array sizes by calling allreduce on chunk lens
# TODO: get global size from C
for arr in csv_node.out_vars:
def f(A):
return hpat.distributed_api.dist_reduce(len(A), np.int32(_op))
f_block = compile_to_numba_ir(
f, {
'hpat': hpat,
'np': np,
'_op': hpat.distributed_api.Reduce_Type.Sum.value
}, typingctx, (typemap[arr.name], ), typemap,
calltypes).blocks.popitem()[1]
replace_arg_nodes(f_block, [arr])
nodes += f_block.body[:-2]
size_var = nodes[-1].target
dist_pass._array_sizes[arr.name] = [size_var]
out, start_var, end_var = dist_pass._gen_1D_div(
size_var, arr.scope, csv_node.loc, "$alloc", "get_node_portion",
hpat.distributed_api.get_node_portion)
dist_pass._array_starts[arr.name] = [start_var]
dist_pass._array_counts[arr.name] = [end_var]
nodes += out
return nodes
