def masked_scalar_const_op_impl(context, builder, sig, args):
return_type = sig.return_type
result = cgutils.create_struct_proxy(return_type)(context, builder)
result.valid = context.get_constant(types.boolean, 0)
if isinstance(sig.args[0], MaskedType):
masked_type, const_type = sig.args
masked_value, const_value = args
indata = cgutils.create_struct_proxy(masked_type)(
context, builder, value=masked_value)
nb_sig = nb_signature(return_type.value_type,
masked_type.value_type, const_type)
compile_args = (indata.value, const_value)
else:
const_type, masked_type = sig.args
const_value, masked_value = args
indata = cgutils.create_struct_proxy(masked_type)(
context, builder, value=masked_value)
nb_sig = nb_signature(return_type.value_type, const_type,
masked_type.value_type)
compile_args = (const_value, indata.value)
with builder.if_then(indata.valid):
result.value = context.compile_internal(builder,
lambda x, y: op(x, y),
nb_sig, compile_args)
result.valid = context.get_constant(types.boolean, 1)
return result._getvalue()
def generic(self, args, kws):
if isinstance(args[0], MaskedType):
# MaskedType(dtype, valid) -> MaskedType(dtype, valid)
return nb_signature(args[0], args[0])
elif isinstance(args[0], SUPPORTED_NUMBA_TYPES):
# scalar_type -> MaskedType(scalar_type, True)
return_type = MaskedType(args[0])
return nb_signature(return_type, args[0])
def generic(self, args, kws):
"""
Typing for `Masked` + `NA`
Handles situations like `x + cudf.NA`
"""
if isinstance(args[0], MaskedType) and isinstance(args[1], NAType):
# In the case of op(Masked, NA), the result has the same
# dtype as the original regardless of what it is
return nb_signature(args[0], args[0], na_type,)
elif isinstance(args[0], NAType) and isinstance(args[1], MaskedType):
return nb_signature(args[1], na_type, args[1])
class MaskedConstructor(ConcreteTemplate):
key = classes.Masked
cases = [
nb_signature(MaskedType(t), t, types.boolean)
for t in (types.integer_domain | types.real_domain)
]
def masked_scalar_unary_op_impl(context, builder, sig, args):
"""
Implement <op> `MaskedType`
"""
# MaskedType(...)
masked_type_1 = sig.args[0]
# MaskedType(...)
masked_return_type = sig.return_type
m1 = cgutils.create_struct_proxy(masked_type_1)(context,
builder,
value=args[0])
# we will return an output struct
result = cgutils.create_struct_proxy(masked_return_type)(context,
builder)
# compute output validity
result.valid = m1.valid
with builder.if_then(m1.valid):
# Let numba handle generating the extra IR needed to perform
# operations on mixed types, by compiling the final core op between
# the two primitive values as a separate function and calling it
result.value = context.compile_internal(
builder,
lambda x: op(x),
nb_signature(
masked_return_type.value_type,
masked_type_1.value_type,
),
(m1.value, ),
)
return result._getvalue()
class MaskedConstructor(ConcreteTemplate):
key = api.Masked
units = ["ns", "ms", "us", "s"]
datetime_cases = {types.NPDatetime(u) for u in units}
timedelta_cases = {types.NPTimedelta(u) for u in units}
cases = [
nb_signature(MaskedType(t), t, types.boolean)
for t in (
types.integer_domain
| types.real_domain
| datetime_cases
| timedelta_cases
| {types.boolean}
)
]
def generic(self, args, kws):
"""
Typing for `Masked` <op> `Masked`
Numba expects a valid numba type to be returned if typing is successful
else `None` signifies the error state (this pattern is commonly used
in Numba)
"""
if isinstance(args[0], MaskedType) and isinstance(args[1], MaskedType):
# In the case of op(Masked, Masked), the return type is a Masked
# such that Masked.value is the primitive type that would have
# been resolved if we were just operating on the
# `value_type`s.
return_type = self.context.resolve_function_type(
self.key, (args[0].value_type, args[1].value_type), kws
).return_type
return nb_signature(MaskedType(return_type), args[0], args[1])
def generic(self, args, kws):
"""
Typing for `Masked` <op> a scalar (and vice-versa).
handles situations like `x + 1`
"""
# In the case of op(Masked, scalar), we resolve the type between
# the Masked value_type and the scalar's type directly
if isinstance(args[0], MaskedType) and isinstance(
args[1], types.Number):
to_resolve_types = (args[0].value_type, args[1])
elif isinstance(args[0], types.Number) and isinstance(
args[1], MaskedType):
to_resolve_types = (args[1].value_type, args[0])
return_type = self.context.resolve_function_type(
self.key, to_resolve_types, kws).return_type
return nb_signature(
MaskedType(return_type),
args[0],
args[1],
)
def masked_scalar_op_impl(context, builder, sig, args):
"""
Implement `MaskedType` <op> `MaskedType`
"""
# MaskedType(...), MaskedType(...)
masked_type_1, masked_type_2 = sig.args
# MaskedType(...)
masked_return_type = sig.return_type
# Let there be two actual LLVM structs backing the two inputs
# https://mapping-high-level-constructs-to-llvm-ir.readthedocs.io/en/latest/basic-constructs/structures.html
m1 = cgutils.create_struct_proxy(masked_type_1)(context,
builder,
value=args[0])
m2 = cgutils.create_struct_proxy(masked_type_2)(context,
builder,
value=args[1])
# we will return an output struct
result = cgutils.create_struct_proxy(masked_return_type)(context,
builder)
# compute output validity
valid = builder.and_(m1.valid, m2.valid)
result.valid = valid
with builder.if_then(valid):
# Let numba handle generating the extra IR needed to perform
# operations on mixed types, by compiling the final core op between
# the two primitive values as a separate function and calling it
result.value = context.compile_internal(
builder,
lambda x, y: op(x, y),
nb_signature(
masked_return_type.value_type,
masked_type_1.value_type,
masked_type_2.value_type,
),
(m1.value, m2.value),
)
return result._getvalue()
def generic(self, args, kws):
if isinstance(args[0], MaskedType):
return nb_signature(types.boolean, MaskedType(types.boolean))
def generic(self, args, kws):
if isinstance(args[0], MaskedType) and isinstance(args[1], NAType):
return nb_signature(types.boolean, args[0], na_type)
elif isinstance(args[1], MaskedType) and isinstance(args[0], NAType):
return nb_signature(types.boolean, na_type, args[1])
def generic(self, args, kws):
if len(args) == 1 and isinstance(args[0], MaskedType):
return_type = self.context.resolve_function_type(
self.key, (args[0].value_type,), kws
).return_type
return nb_signature(MaskedType(return_type), args[0])