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