def map_operator(name, inplace_name, op):
op_func = getattr(operator, name)
@lower(op_func, types.VarArg(types.Any))
def binop_impl(context, builder, sig, args):
impl = context.get_function(op, sig)
return impl(builder, args)
if inplace_name:
op_func = getattr(operator, inplace_name)
@lower(op_func, types.VarArg(types.Any))
def binop_inplace_impl(context, builder, sig, args):
first = sig.args[0]
if first.mutable:
impl = context.get_function(op + '=', sig)
else:
impl = context.get_function(op, sig)
return impl(builder, args)
def _implement_method(self, registry, attr):
@registry.lower((types.ClassInstanceType, attr),
types.ClassInstanceType, types.VarArg(types.Any))
def imp(context, builder, sig, args):
instance_type = sig.args[0]
method = instance_type.jitmethods[attr]
disp_type = types.Dispatcher(method)
call = context.get_function(disp_type, sig)
out = call(builder, args)
return imputils.impl_ret_new_ref(context, builder, sig.return_type,
out)
def implement_method(self, registry, instance_type, attr):
# TODO: we should be able to refactor this so that the registration
# can be at the top-level instead of at per instance type.
@registry.register
@imputils.implement((instance_type, attr), types.VarArg(types.Any))
def imp(context, builder, sig, args):
method = instance_type.jitmethods[attr]
call = self.targetctx.get_function(types.Dispatcher(method), sig)
out = call(builder, args)
return imputils.impl_ret_new_ref(context, builder, sig.return_type,
out)
def _getsetitem_gen(getset):
_dunder_meth = "__%s__" % getset
op = getattr(operator, getset)
@templates.infer_global(op)
class GetSetItem(templates.AbstractTemplate):
def generic(self, args, kws):
instance = args[0]
if isinstance(instance, types.ClassInstanceType) and \
_dunder_meth in instance.jitmethods:
meth = instance.jitmethods[_dunder_meth]
disp_type = types.Dispatcher(meth)
sig = disp_type.get_call_type(self.context, args, kws)
return sig
# lower both {g,s}etitem and __{g,s}etitem__ to catch the calls
# from python and numba
imputils.lower_builtin((types.ClassInstanceType, _dunder_meth),
types.ClassInstanceType,
types.VarArg(types.Any))(get_imp())
imputils.lower_builtin(op, types.ClassInstanceType,
types.VarArg(types.Any))(get_imp())
def map_operator(name, inplace_name, op):
op_func = getattr(operator, name)
reverse_args = (op == 'in')
@lower(op_func, types.VarArg(types.Any))
def binop_impl(context, builder, sig, args):
if reverse_args:
args = args[::-1]
sig = typing.signature(sig.return_type, *sig.args[::-1])
impl = context.get_function(op, sig)
return impl(builder, args)
if inplace_name:
op_func = getattr(operator, inplace_name)
@lower(op_func, types.VarArg(types.Any))
def binop_inplace_impl(context, builder, sig, args):
first = sig.args[0]
if first.mutable:
impl = context.get_function(op + '=', sig)
else:
impl = context.get_function(op, sig)
return impl(builder, args)
def _implement_method(self, registry, attr):
# create a separate instance of imp method to avoid closure clashing
def get_imp():
def imp(context, builder, sig, args):
instance_type = sig.args[0]
method = instance_type.jitmethods[attr]
disp_type = types.Dispatcher(method)
call = context.get_function(disp_type, sig)
out = call(builder, args)
_add_linking_libs(context, call)
return imputils.impl_ret_new_ref(context, builder,
sig.return_type, out)
return imp
def _getsetitem_gen(getset):
_dunder_meth = "__%s__" % getset
op = getattr(operator, getset)
@templates.infer_global(op)
class GetSetItem(templates.AbstractTemplate):
def generic(self, args, kws):
instance = args[0]
if isinstance(instance, types.ClassInstanceType) and \
_dunder_meth in instance.jitmethods:
meth = instance.jitmethods[_dunder_meth]
disp_type = types.Dispatcher(meth)
sig = disp_type.get_call_type(self.context, args, kws)
return sig
# lower both {g,s}etitem and __{g,s}etitem__ to catch the calls
# from python and numba
imputils.lower_builtin((types.ClassInstanceType, _dunder_meth),
types.ClassInstanceType,
types.VarArg(types.Any))(get_imp())
imputils.lower_builtin(op, types.ClassInstanceType,
types.VarArg(types.Any))(get_imp())
dunder_stripped = attr.strip('_')
if dunder_stripped in ("getitem", "setitem"):
_getsetitem_gen(dunder_stripped)
else:
registry.lower(
(types.ClassInstanceType, attr), types.ClassInstanceType,
types.VarArg(types.Any))(get_imp())
@lower("print_item", types.CharSeq)
def print_charseq(context, builder, sig, args):
[tx] = sig.args
[x] = args
py = context.get_python_api(builder)
xp = cgutils.alloca_once(builder, x.type)
builder.store(x, xp)
byteptr = builder.bitcast(xp, Type.pointer(Type.int(8)))
size = context.get_constant(types.intp, tx.count)
cstr = py.bytes_from_string_and_size(byteptr, size)
py.print_object(cstr)
py.decref(cstr)
res = context.get_dummy_value()
return impl_ret_untracked(context, builder, sig.return_type, res)
@lower(print, types.VarArg(types.Any))
def print_varargs(context, builder, sig, args):
py = context.get_python_api(builder)
for i, (argtype, argval) in enumerate(zip(sig.args, args)):
signature = typing.signature(types.none, argtype)
imp = context.get_function("print_item", signature)
imp(builder, [argval])
if i < len(args) - 1:
py.print_string(' ')
py.print_string('\n')
res = context.get_dummy_value()
return impl_ret_untracked(context, builder, sig.return_type, res)
def get_defaults(context):
"""
Get the default values for a slice's members:
(start for positive step, start for negative step,
stop for positive step, stop for negative step, step)
"""
maxint = (1 << (context.address_size - 1)) - 1
return (0, maxint, maxint, -maxint - 1, 1)
#---------------------------------------------------------------------------
# The slice structure
@lower_builtin(slice, types.VarArg(types.Any))
def slice_constructor_impl(context, builder, sig, args):
default_start_pos, default_start_neg, default_stop_pos, default_stop_neg, default_step = \
[context.get_constant(types.intp, x) for x in get_defaults(context)]
slice_args = [None] * 3
# Fetch non-None arguments
if len(args) == 1 and sig.args[0] is not types.none:
slice_args[1] = args[0]
else:
for i, (ty, val) in enumerate(zip(sig.args, args)):
if ty is not types.none:
slice_args[i] = val
# Fill omitted arguments
typ_to_format = {
types.int32: 'd',
types.uint32: 'u',
types.int64: 'lld',
types.uint64: 'llu',
types.float32: 'f',
types.float64: 'lf',
}
ll.add_symbol('print_str', hstr_ext.print_str)
ll.add_symbol('print_char', hstr_ext.print_char)
@lower_builtin(cprint, types.VarArg(types.Any))
def cprint_lower(context, builder, sig, args):
from sdc.str_ext import string_type, char_type
for i, val in enumerate(args):
typ = sig.args[i]
if typ == string_type:
fnty = lir.FunctionType(lir.VoidType(),
[lir.IntType(8).as_pointer()])
fn = builder.module.get_or_insert_function(fnty, name="print_str")
builder.call(fn, [val])
cgutils.printf(builder, " ")
continue
if typ == char_type:
fnty = lir.FunctionType(lir.VoidType(), [lir.IntType(8)])
fn = builder.module.get_or_insert_function(fnty, name="print_char")
#-------------------------------------------------------------------------------
# builtin `zip` implementation
@struct_factory(types.ZipType)
def make_zip_cls(zip_type):
"""
Return the Structure representation of the given *zip_type* (an
instance of types.ZipType).
"""
return cgutils.create_struct_proxy(zip_type)
@builtin
@implement(zip, types.VarArg(types.Any))
def make_zip_object(context, builder, sig, args):
zip_type = sig.return_type
assert len(args) == len(zip_type.source_types)
zipcls = make_zip_cls(zip_type)
zipobj = zipcls(context, builder)
for i, (arg, srcty) in enumerate(zip(args, sig.args)):
zipobj[i] = call_getiter(context, builder, srcty, arg)
return zipobj._getvalue()
@builtin
@register
@implement(types.print_item_type, types.Kind(types.CharSeq))
def print_charseq(context, builder, sig, args):
[x] = args
py = context.get_python_api(builder)
xp = cgutils.alloca_once(builder, x.type)
builder.store(x, xp)
byteptr = builder.bitcast(xp, Type.pointer(Type.int(8)))
size = context.get_constant(types.intp, x.type.elements[0].count)
cstr = py.bytes_from_string_and_size(byteptr, size)
py.print_object(cstr)
py.decref(cstr)
return context.get_dummy_value()
@register
@implement(types.print_type, types.VarArg(types.Any))
def print_varargs(context, builder, sig, args):
py = context.get_python_api(builder)
for i, (argtype, argval) in enumerate(zip(sig.args, args)):
signature = typing.signature(types.none, argtype)
imp = context.get_function(types.print_item_type, signature)
imp(builder, [argval])
if i == len(args) - 1:
py.print_string('\n')
else:
py.print_string(' ')
return context.get_dummy_value()
#---------------------------------------------------------------------------
# The slice structure
def make_slice(context, builder, typ, value=None):
"""
Create a slice structure, optionally initialized from the given LLVM
*value*.
"""
cls = cgutils.create_struct_proxy(typ)
return cls(context, builder, value=value)
@builtin
@implement(types.slice_type, types.VarArg(types.Any))
def slice_constructor_impl(context, builder, sig, args):
slice_args = []
for ty, val, default in zip_longest(sig.args, args, get_defaults(context)):
if ty in (types.none, None):
# Omitted or None
slice_args.append(context.get_constant(types.intp, default))
else:
slice_args.append(val)
start, stop, step = slice_args
ty = sig.return_type
sli = make_slice(context, builder, sig.return_type)
sli.start = start
sli.stop = stop
sli.step = step
builder.store(ncount, enum.count)
srcres = call_iternext(context, builder, enumty.source_type, enum.iter)
is_valid = srcres.is_valid()
result.set_valid(is_valid)
with builder.if_then(is_valid):
srcval = srcres.yielded_value()
result.yield_(context.make_tuple(builder, enumty.yield_type,
[count, srcval]))
#-------------------------------------------------------------------------------
# builtin `zip` implementation
@lower_builtin(zip, types.VarArg(types.Any))
def make_zip_object(context, builder, sig, args):
zip_type = sig.return_type
assert len(args) == len(zip_type.source_types)
zipobj = context.make_helper(builder, zip_type)
for i, (arg, srcty) in enumerate(zip(args, sig.args)):
zipobj[i] = call_getiter(context, builder, srcty, arg)
res = zipobj._getvalue()
return impl_ret_new_ref(context, builder, sig.return_type, res)
@lower_builtin('iternext', types.ZipType)
@iternext_impl
def iternext_numpy_nditer(context, builder, sig, args, result):
[nditerty] = sig.args
[nditer] = args
nditercls = make_array_ndenumerate_cls(nditerty)
nditer = nditercls(context, builder, value=nditer)
arrty = nditerty.array_type
arrcls = context.make_array(arrty)
arr = arrcls(context, builder, value=builder.load(nditer.array))
nditer.iternext_specific(context, builder, arrty, arr, result)
@builtin
@implement(numpy.ndindex, types.VarArg(types.Kind(types.Integer)))
def make_array_ndindex(context, builder, sig, args):
"""ndindex(*shape)"""
shape = [
context.cast(builder, arg, argty, types.intp)
for argty, arg in zip(sig.args, args)
]
nditercls = make_ndindex_cls(types.NumpyNdIndexType(len(shape)))
nditer = nditercls(context, builder)
nditer.init_specific(context, builder, shape)
return nditer._getvalue()
@builtin
if not isinstance(on, types.StringLiteral):
raise ValueError(
"'on' argument to rolling() should be constant string")
on = on.literal_value
selection = list(df.columns)
if on is not None:
selection.remove(on)
out_typ = RollingType(df, on, tuple(selection), False)
return signature(out_typ, *args)
# dummy lowering to avoid overload errors, remove after overload inline PR
# is merged
@lower_builtin(rolling_dummy, types.VarArg(types.Any))
def lower_rolling_dummy(context, builder, sig, args):
return context.get_constant_null(sig.return_type)
@infer
class GetItemDataFrameRolling2(AbstractTemplate):
key = 'static_getitem'
def generic(self, args, kws):
rolling, idx = args
# df.rolling('A')['B', 'C']
if isinstance(rolling, RollingType):
if isinstance(idx, tuple):
assert all(isinstance(c, str) for c in idx)
selection = idx
builder = llvmir.IRBuilder(dtor_fn.append_basic_block())
alloc_fe_type = instance_type.get_data_type()
alloc_type = context.get_value_type(alloc_fe_type)
ptr = builder.bitcast(dtor_fn.args[0], alloc_type.as_pointer())
data = context.make_helper(builder, alloc_fe_type, ref=ptr)
context.nrt.decref(builder, alloc_fe_type, data._getvalue())
builder.ret_void()
return dtor_fn
@ClassBuilder.class_impl_registry.lower(types.ClassType, types.VarArg(types.Any))
def ctor_impl(context, builder, sig, args):
"""
Generic constructor (__new__) for jitclasses.
"""
# Allocate the instance
inst_typ = sig.return_type
alloc_type = context.get_data_type(inst_typ.get_data_type())
alloc_size = context.get_abi_sizeof(alloc_type)
meminfo = context.nrt.meminfo_alloc_dtor(
builder,
context.get_constant(types.uintp, alloc_size),
imp_dtor(context, builder.module, inst_typ),
)
data_pointer = context.nrt.meminfo_data(builder, meminfo)
alloc_fe_type = instance_type.get_data_type()
alloc_type = context.get_value_type(alloc_fe_type)
data_struct = cgutils.create_struct_proxy(alloc_fe_type)
ptr = builder.bitcast(dtor_fn.args[0], alloc_type.as_pointer())
data = data_struct(context, builder, ref=ptr)
context.nrt_decref(builder, alloc_fe_type, data._getvalue())
builder.ret_void()
return dtor_fn
@imputils.implement(types.ClassType, types.VarArg(types.Any))
def ctor_impl(context, builder, sig, args):
# Allocate the instance
inst_typ = sig.return_type
alloc_type = context.get_data_type(inst_typ.get_data_type())
alloc_size = context.get_abi_sizeof(alloc_type)
meminfo = context.nrt_meminfo_alloc_dtor(
builder,
context.get_constant(types.uintp, alloc_size),
imp_dtor(context, builder.module, inst_typ),
)
data_pointer = context.nrt_meminfo_data(builder, meminfo)
data_pointer = builder.bitcast(data_pointer, alloc_type.as_pointer())
# Nullify all data
alloc_fe_type = instance_type.get_data_type()
alloc_type = context.get_value_type(alloc_fe_type)
ptr = builder.bitcast(dtor_fn.args[0], alloc_type.as_pointer())
data = context.make_helper(builder, alloc_fe_type, ref=ptr)
context.nrt_decref(builder, alloc_fe_type, data._getvalue())
builder.ret_void()
return dtor_fn
@ClassBuilder.class_impl_registry.lower(types.ClassType,
types.VarArg(types.Any))
def ctor_impl(context, builder, sig, args):
"""
Generic constructor (__new__) for jitclasses.
"""
# Allocate the instance
inst_typ = sig.return_type
alloc_type = context.get_data_type(inst_typ.get_data_type())
alloc_size = context.get_abi_sizeof(alloc_type)
meminfo = context.nrt_meminfo_alloc_dtor(
builder,
context.get_constant(types.uintp, alloc_size),
imp_dtor(context, builder.module, inst_typ),
)
data_pointer = context.nrt_meminfo_data(builder, meminfo)
def box_df(names, arrs):
return pd.DataFrame()
@infer_global(box_df)
class BoxDfTyper(AbstractTemplate):
def generic(self, args, kws):
assert not kws
assert len(args) % 2 == 0, "name and column pairs expected"
col_names = [a.literal_value for a in args[:len(args)//2]]
col_types = [a.dtype for a in args[len(args)//2:]]
df_typ = PandasDataFrameType(col_names, col_types)
return signature(df_typ, *args)
BoxDfTyper.support_literals = True
@lower_builtin(box_df, types.Literal, types.VarArg(types.Any))
def lower_box_df(context, builder, sig, args):
assert len(sig.args) % 2 == 0, "name and column pairs expected"
n_cols = len(sig.args)//2
col_names = [a.literal_value for a in sig.args[:n_cols]]
col_arrs = [a for a in args[n_cols:]]
arr_typs = [a for a in sig.args[n_cols:]]
pyapi = context.get_python_api(builder)
env_manager = context.get_env_manager(builder)
c = numba.pythonapi._BoxContext(context, builder, pyapi, env_manager)
gil_state = pyapi.gil_ensure() # acquire GIL
mod_name = context.insert_const_string(c.builder.module, "pandas")
class_obj = pyapi.import_module_noblock(mod_name)
res = pyapi.call_method(class_obj, "DataFrame", ())
name="get_join_sendrecv_counts")
total_size = builder.call(fn, call_args)
items = [
builder.load(send_counts),
builder.load(recv_counts),
builder.load(send_disp),
builder.load(recv_disp), total_size
]
out_tuple_typ = types.Tuple([
c_buffer_type, c_buffer_type, c_buffer_type, c_buffer_type, types.intp
])
return context.make_tuple(builder, out_tuple_typ, items)
@lower_builtin(shuffle_data, c_buffer_type, c_buffer_type, c_buffer_type,
c_buffer_type, types.VarArg(types.Any))
def lower_shuffle(context, builder, sig, args):
# assuming there are 4 buffer arguments, column vars, send arrs, recv arrs
assert (len(args) - 4) % 3 == 0
num_cols = (len(args) - 4) // 3
send_counts, recv_counts, send_disp, recv_disp = args[:4]
col_names = ["c" + str(i) for i in range(num_cols)]
send_names = ["send_c" + str(i) for i in range(num_cols)]
# create send buffer building function
func_text = "def f(send_disp, {}, {}):\n".format(",".join(col_names),
",".join(send_names))
func_text += " n_pes = hpat.distributed_api.get_size()\n"
func_text += " tmp_count = np.zeros(n_pes, dtype=np.int64)\n"
func_text += " for i in range(len(c0)):\n"
func_text += " node_id = c0[i] % n_pes\n"
func_text += " ind = send_disp[node_id] + tmp_count[node_id]\n"
if isinstance(as_index, sdc.utilities.utils.BooleanLiteral):
as_index = as_index.literal_value
else:
# XXX as_index type is just bool when value not passed. Therefore,
# we assume the default True value.
# TODO: more robust fix or just check
as_index = True
out_typ = DataFrameGroupByType(df, keys, tuple(selection),
as_index, False)
return signature(out_typ, *args)
# dummy lowering to avoid overload errors, remove after overload inline PR
# is merged
@lower_builtin(groupby_dummy, types.VarArg(types.Any))
def lower_groupby_dummy(context, builder, sig, args):
return context.get_constant_null(sig.return_type)
@infer
class GetItemDataFrameGroupBy2(AbstractTemplate):
key = 'static_getitem'
def generic(self, args, kws):
grpby, idx = args
# df.groupby('A')['B', 'C']
if isinstance(grpby, DataFrameGroupByType):
if isinstance(idx, tuple):
assert all(isinstance(c, str) for c in idx)
selection = idx
elif isinstance(idx, str):
lambda x: None,
levels,
)
else:
_names = name
levels_and_codes_pairs = sdc_tuple_map_elementwise(
_multi_index_create_levels_and_codes, levels, codes, _names)
_levels, _codes = sdc_tuple_unzip(levels_and_codes_pairs)
return init_multi_index(_levels, _codes)
return pd_multi_index_ctor_impl
@lower_builtin(MultiIndexTypeRef, types.VarArg(types.Any))
def multi_index_typeref_call_impl(context, builder, sig, args):
# FIXME_Numba#7111: this uses low-level API as a workaround for numba issue
# TO-DO: remove and use @overload(MultiIndexTypeRef), once issue is fixed
# and now we do the following:
# (1) lookup function type for the actual ctor (sdc_pandas_multi_index_ctor)
# (2) get compiled implementation for provided args (hardcodes 0 as selected overload template,
# i.e. we rely on the fact that sdc_pandas_multi_index_ctor was overloaded only once)
# (3) get the function descriptor from compiled result and emit the call to it
call_sig = context.typing_context._resolve_user_function_type(
sdc_pandas_multi_index_ctor, sig.args, {})
fnty = context.typing_context._lookup_global(sdc_pandas_multi_index_ctor)
disp = fnty.templates[0](context.typing_context)._get_impl(
call_sig.args, {})
cres = disp[0].get_compile_result(call_sig)
