def getitem_str_impl(arr, slice_index, start, count):
rank = hpat.distributed_api.get_rank()
k = slice_index.stop
# get total characters for allocation
n_chars = np.uint64(0)
if k > start:
# if slice end is beyond the start of this subset we have to send our elements
my_end = min(count, k - start)
my_arr = arr[:my_end]
else:
my_arr = arr[:0]
# get the total number of chars in our array, then gather all arrays into one
# and compute total number of chars in all arrays
n_chars = num_total_chars(my_arr)
my_arr = hpat.distributed_api.gatherv(my_arr)
n_chars = hpat.distributed_api.dist_reduce(n_chars,
np.int32(reduce_op))
if rank != 0:
out_arr = pre_alloc_string_array(k, n_chars)
else:
out_arr = my_arr
# actual communication
hpat.distributed_api.bcast(out_arr)
return out_arr
def getitem_str_impl(arr, slice_index, start, count):
rank = hpat.distributed_api.get_rank()
k = slice_index.stop
# get total characters for allocation
n_chars = np.uint64(0)
if k > count:
my_end = min(count, max(k-start, 0))
my_arr = arr[:my_end]
my_arr = hpat.distributed_api.gatherv(my_arr)
n_chars = hpat.distributed_api.dist_reduce(
num_total_chars(my_arr), np.int32(reduce_op))
if rank == 0:
out_arr = my_arr
else:
if rank == 0:
my_arr = arr[:k]
n_chars = num_total_chars(my_arr)
out_arr = my_arr
n_chars = bcast_scalar(n_chars)
if rank != 0:
out_arr = pre_alloc_string_array(k, n_chars)
# actual communication
hpat.distributed_api.bcast(out_arr)
return out_arr
def prealloc_impl(arr):
rank = hpat.distributed_api.get_rank()
n_loc = bcast_scalar(len(arr))
n_all_char = bcast_scalar(np.int64(num_total_chars(arr)))
if rank != MPI_ROOT:
arr = pre_alloc_string_array(n_loc, n_all_char)
return arr
def finalize_str_impl(arr, shuffle_meta, is_contig):
hpat.distributed_api.alltoall(shuffle_meta.send_counts,
shuffle_meta.recv_counts, 1)
hpat.distributed_api.alltoall(shuffle_meta.send_counts_char,
shuffle_meta.recv_counts_char, 1)
shuffle_meta.n_out = shuffle_meta.recv_counts.sum()
n_all_chars = shuffle_meta.recv_counts_char.sum()
shuffle_meta.out_arr = pre_alloc_string_array(shuffle_meta.n_out,
n_all_chars)
shuffle_meta.send_disp = hpat.hiframes_join.calc_disp(
shuffle_meta.send_counts)
shuffle_meta.recv_disp = hpat.hiframes_join.calc_disp(
shuffle_meta.recv_counts)
shuffle_meta.send_disp_char = hpat.hiframes_join.calc_disp(
shuffle_meta.send_counts_char)
shuffle_meta.recv_disp_char = hpat.hiframes_join.calc_disp(
shuffle_meta.recv_counts_char)
if not is_contig:
shuffle_meta.n_send = shuffle_meta.send_counts.sum()
shuffle_meta.send_arr_lens = np.empty(shuffle_meta.n_send,
np.uint32)
# shuffle_meta.send_arr_lens = np.empty(len(arr), np.uint32)
# s_n_all_chars = num_total_chars(arr)
s_n_all_chars = shuffle_meta.send_counts_char.sum()
shuffle_meta.send_arr_chars_arr = np.empty(s_n_all_chars, np.uint8)
shuffle_meta.send_arr_chars = get_ctypes_ptr(
shuffle_meta.send_arr_chars_arr.ctypes)
def set_string_to_array(A):
# TODO: support unicode
num_total_chars = num_total_chars_set_string(A)
num_strs = len(A)
str_arr = pre_alloc_string_array(num_strs, num_total_chars)
populate_str_arr_from_set(A, str_arr)
return str_arr
def _str_replace_noregex_impl(str_arr, pat, val):
numba.parfor.init_prange()
n = len(str_arr)
n_total_chars = 0
str_list = hpat.str_ext.alloc_str_list(n)
for i in numba.parfor.internal_prange(n):
out_str = str_arr[i].replace(pat, val)
str_list[i] = out_str
n_total_chars += get_utf8_size(out_str)
numba.parfor.init_prange()
out_arr = pre_alloc_string_array(n, n_total_chars)
for i in numba.parfor.internal_prange(n):
_str = str_list[i]
out_arr[i] = _str
return hpat.hiframes.api.init_series(out_arr)
def ensure_capacity_str(arr, new_size, n_chars):
# new_size is right after write index
new_arr = arr
curr_len = len(arr)
curr_num_chars = num_total_chars(arr)
# TODO: corner case test
#print("new alloc", new_size, curr_len, getitem_str_offset(arr, new_size-1), n_chars, curr_num_chars)
if curr_len < new_size or getitem_str_offset(arr, new_size-1) + n_chars > curr_num_chars:
new_len = 2 * curr_len
new_num_chars = 2 * curr_num_chars
new_arr = pre_alloc_string_array(new_len, new_num_chars)
copy_str_arr_slice(new_arr, arr, new_size-1)
return new_arr
def gatherv_str_arr_impl(data):
rank = hpat.distributed_api.get_rank()
n_loc = len(data)
n_all_chars = num_total_chars(data)
# allocate send lens arrays
send_arr_lens = np.empty(n_loc, np.uint32) # XXX offset type is uint32
send_data_ptr = get_data_ptr(data)
for i in range(n_loc):
_str = data[i]
send_arr_lens[i] = len(_str)
recv_counts = gather_scalar(np.int32(n_loc))
recv_counts_char = gather_scalar(np.int32(n_all_chars))
n_total = recv_counts.sum()
n_total_char = recv_counts_char.sum()
# displacements
all_data = StringArray(['']) # dummy arrays on non-root PEs
displs = np.empty(0, np.int32)
displs_char = np.empty(0, np.int32)
if rank == MPI_ROOT:
all_data = pre_alloc_string_array(n_total, n_total_char)
displs = hpat.hiframes.join.calc_disp(recv_counts)
displs_char = hpat.hiframes.join.calc_disp(recv_counts_char)
offset_ptr = get_offset_ptr(all_data)
data_ptr = get_data_ptr(all_data)
c_gatherv(
send_arr_lens.ctypes,
np.int32(n_loc),
offset_ptr,
recv_counts.ctypes,
displs.ctypes,
int32_typ_enum)
c_gatherv(
send_data_ptr,
np.int32(n_all_chars),
data_ptr,
recv_counts_char.ctypes,
displs_char.ctypes,
char_typ_enum)
convert_len_arr_to_offset(offset_ptr, n_total)
return all_data
def _str_replace_regex_impl(str_arr, pat, val):
numba.parfor.init_prange()
e = hpat.str_ext.compile_regex(unicode_to_std_str(pat))
val = unicode_to_std_str(val)
n = len(str_arr)
n_total_chars = 0
str_list = hpat.str_ext.alloc_str_list(n)
for i in numba.parfor.internal_prange(n):
# TODO: support unicode
in_str = unicode_to_std_str(str_arr[i])
out_str = std_str_to_unicode(
hpat.str_ext.str_replace_regex(in_str, e, val))
str_list[i] = out_str
n_total_chars += len(out_str)
numba.parfor.init_prange()
out_arr = pre_alloc_string_array(n, n_total_chars)
for i in numba.parfor.internal_prange(n):
_str = str_list[i]
out_arr[i] = _str
return hpat.hiframes.api.init_series(out_arr)
def trim_arr_str(arr, size):
# print("trim size", size, arr[size-1], getitem_str_offset(arr, size))
new_arr = pre_alloc_string_array(size, np.int64(getitem_str_offset(arr, size)))
copy_str_arr_slice(new_arr, arr, size)
return new_arr
def empty_like_type_str_arr(n, arr):
# average character heuristic
avg_chars = 20 # heuristic
if len(arr) != 0:
avg_chars = num_total_chars(arr) // len(arr)
return pre_alloc_string_array(n, n * avg_chars)
def set_string_to_array(str_set):
num_total_chars = num_total_chars_set_string(str_set)
num_strs = len(str_set)
str_arr = pre_alloc_string_array(num_strs, num_total_chars)
populate_str_arr_from_set(str_set, str_arr)
return str_arr