def register_unary_op(op):
"""
Register a lowering implementation for the
unary op `op`.
Because the lowering implementations compile the final
op separately using a lambda and compile_internal, `op`
needs to be tied to each lowering implementation using
a closure.
This function makes and lowers a closure for one op.
"""
to_lower_op = make_unary_op(op)
cuda_lower(op, MaskedType)(to_lower_op)
def register_const_op(op):
to_lower_op = make_const_op(op)
cuda_lower(op, MaskedType, types.Number)(to_lower_op)
cuda_lower(op, types.Number, MaskedType)(to_lower_op)
cuda_lower(op, MaskedType, types.Boolean)(to_lower_op)
cuda_lower(op, types.Boolean, MaskedType)(to_lower_op)
cuda_lower(op, MaskedType, types.NPDatetime)(to_lower_op)
cuda_lower(op, types.NPDatetime, MaskedType)(to_lower_op)
cuda_lower(op, MaskedType, types.NPTimedelta)(to_lower_op)
cuda_lower(op, types.NPTimedelta, MaskedType)(to_lower_op)
cuda_lower(op, MaskedType, types.Number)(to_lower_op)
cuda_lower(op, types.Number, MaskedType)(to_lower_op)
cuda_lower(op, MaskedType, types.Boolean)(to_lower_op)
cuda_lower(op, types.Boolean, MaskedType)(to_lower_op)
cuda_lower(op, MaskedType, types.NPDatetime)(to_lower_op)
cuda_lower(op, types.NPDatetime, MaskedType)(to_lower_op)
cuda_lower(op, MaskedType, types.NPTimedelta)(to_lower_op)
cuda_lower(op, types.NPTimedelta, MaskedType)(to_lower_op)
# register all lowering at init
for binary_op in arith_ops + bitwise_ops + comparison_ops:
register_arithmetic_op(binary_op)
register_const_op(binary_op)
# null op impl can be shared between all ops
cuda_lower(binary_op, MaskedType, NAType)(masked_scalar_null_op_impl)
cuda_lower(binary_op, NAType, MaskedType)(masked_scalar_null_op_impl)
# register all lowering at init
for unary_op in unary_ops:
register_unary_op(unary_op)
@cuda_lower(operator.is_, MaskedType, NAType)
@cuda_lower(operator.is_, NAType, MaskedType)
def masked_scalar_is_null_impl(context, builder, sig, args):
"""
Implement `MaskedType` is `NA`
"""
if isinstance(sig.args[1], NAType):
masked_type, na = sig.args
def register_const_op(op):
to_lower_op = make_const_op(op)
cuda_lower(op, MaskedType, types.Number)(to_lower_op)
cuda_lower(op, types.Number, MaskedType)(to_lower_op)