From aa98bb9c45f538ea5dcd067ca9428819addbf48f Mon Sep 17 00:00:00 2001 From: "Todd A. Anderson" Date: Fri, 4 May 2018 14:56:11 -0700 Subject: [PATCH] int packed version of datetime.date --- hpat/pd_timestamp_ext.py | 227 +++++++++++++++++++++++++++------------ 1 file changed, 159 insertions(+), 68 deletions(-) diff --git a/hpat/pd_timestamp_ext.py b/hpat/pd_timestamp_ext.py index 6c149ce8a..4dc71a7c2 100644 --- a/hpat/pd_timestamp_ext.py +++ b/hpat/pd_timestamp_ext.py @@ -1,7 +1,8 @@ import numba from numba import types from numba.extending import (typeof_impl, type_callable, models, register_model, NativeValue, - make_attribute_wrapper, lower_builtin, box, unbox, lower_cast, overload_method) + make_attribute_wrapper, lower_builtin, box, unbox, lower_cast, + lower_getattr, infer_getattr, overload_method) from numba import cgutils from numba.targets.boxing import unbox_array from numba.typing.templates import infer_getattr, AttributeTemplate, bound_function @@ -82,79 +83,169 @@ def impl_ctor_pandas_dts(context, builder, sig, args): #--------------------------------------------------------------- -class DatetimeDateType(types.Type): - def __init__(self): - super(DatetimeDateType, self).__init__( - name='DatetimeDateType()') - -datetime_date_type = DatetimeDateType() - -@typeof_impl.register(datetime.date) -def typeof_pd_timestamp(val, c): - return datetime_date_type - - -@register_model(DatetimeDateType) -class DatetimeDateModel(models.StructModel): - def __init__(self, dmm, fe_type): - members = [ - ('year', types.int64), - ('month', types.int64), - ('day', types.int64), - ] - models.StructModel.__init__(self, dmm, fe_type, members) - -make_attribute_wrapper(DatetimeDateType, 'year', 'year') -make_attribute_wrapper(DatetimeDateType, 'month', 'month') -make_attribute_wrapper(DatetimeDateType, 'day', 'day') +intdatetime = True -@unbox(DatetimeDateType) -def unbox_datetime_date(typ, val, c): - year_obj = c.pyapi.object_getattr_string(val, "year") - month_obj = c.pyapi.object_getattr_string(val, "month") - day_obj = c.pyapi.object_getattr_string(val, "day") +if intdatetime: + class DatetimeDateType(types.Type): + def __init__(self): + super(DatetimeDateType, self).__init__( + name='DatetimeDateType()') + self.bitwidth = 64 - dt_date = cgutils.create_struct_proxy(typ)(c.context, c.builder) - dt_date.year = c.pyapi.long_as_longlong(year_obj) - dt_date.month = c.pyapi.long_as_longlong(month_obj) - dt_date.day = c.pyapi.long_as_longlong(day_obj) - - c.pyapi.decref(year_obj) - c.pyapi.decref(month_obj) - c.pyapi.decref(day_obj) + datetime_date_type = DatetimeDateType() - is_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred()) - return NativeValue(dt_date._getvalue(), is_error=is_error) - -@box(DatetimeDateType) -def box_datetime_date(typ, val, c): - dt_date = cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val) - year_obj = c.pyapi.long_from_longlong(dt_date.year) - month_obj = c.pyapi.long_from_longlong(dt_date.month) - day_obj = c.pyapi.long_from_longlong(dt_date.day) - dt_obj = c.pyapi.unserialize(c.pyapi.serialize_object(datetime.date)) - res = c.pyapi.call_function_objargs(dt_obj, (year_obj, month_obj, day_obj)) - c.pyapi.decref(year_obj) - c.pyapi.decref(month_obj) - c.pyapi.decref(day_obj) - return res + @typeof_impl.register(datetime.date) + def typeof_pd_timestamp(val, c): + return datetime_date_type -@type_callable(datetime.date) -def type_timestamp(context): - def typer(year, month, day): - # TODO: check types + register_model(DatetimeDateType)(models.IntegerModel) + + @infer_getattr + class DatetimeAttribute(AttributeTemplate): + key = DatetimeDateType + + def generic_resolve(self, typ, attr): + return types.int64 + + @lower_getattr(DatetimeDateType, 'year') + def datetime_get_year(context, builder, typ, val): + return builder.lshr(val, lir.Constant(lir.IntType(64), 32)) + + @lower_getattr(DatetimeDateType, 'month') + def datetime_get_year(context, builder, typ, val): + return builder.and_(builder.lshr(val, lir.Constant(lir.IntType(64), 16)), lir.Constant(lir.IntType(64), 0xFFFF)) + + @lower_getattr(DatetimeDateType, 'day') + def datetime_get_year(context, builder, typ, val): + return builder.and_(val, lir.Constant(lir.IntType(64), 0xFFFF)) + + @numba.njit(nopython=True) + def convert_datetime_date_array_to_native(x): + return np.array([(val.day + (val.month << 16) + (val.year << 32)) for val in x], dtype=np.int64) + + @unbox(DatetimeDateType) + def unbox_datetime_date(typ, val, c): + year_obj = c.pyapi.object_getattr_string(val, "year") + month_obj = c.pyapi.object_getattr_string(val, "month") + day_obj = c.pyapi.object_getattr_string(val, "day") + + yll = c.pyapi.long_as_longlong(year_obj) + mll = c.pyapi.long_as_longlong(month_obj) + dll = c.pyapi.long_as_longlong(day_obj) + + nopython_date = c.builder.add(dll, + c.builder.add(c.builder.shl(yll, lir.Constant(lir.IntType(64), 32)), + c.builder.shl(mll, lir.Constant(lir.IntType(64), 16)))) + + c.pyapi.decref(year_obj) + c.pyapi.decref(month_obj) + c.pyapi.decref(day_obj) + + is_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred()) + return NativeValue(nopython_date, is_error=is_error) + + @box(DatetimeDateType) + def box_datetime_date(typ, val, c): + year_obj = c.pyapi.long_from_longlong(c.builder.lshr(val, lir.Constant(lir.IntType(64), 32))) + month_obj = c.pyapi.long_from_longlong(c.builder.and_(c.builder.lshr(val, lir.Constant(lir.IntType(64), 16)), lir.Constant(lir.IntType(64), 0xFFFF))) + day_obj = c.pyapi.long_from_longlong(c.builder.and_(val, lir.Constant(lir.IntType(64), 0xFFFF))) + + dt_obj = c.pyapi.unserialize(c.pyapi.serialize_object(datetime.date)) + res = c.pyapi.call_function_objargs(dt_obj, (year_obj, month_obj, day_obj)) + c.pyapi.decref(year_obj) + c.pyapi.decref(month_obj) + c.pyapi.decref(day_obj) + return res + + @type_callable(datetime.date) + def type_timestamp(context): + def typer(year, month, day): + # TODO: check types + return datetime_date_type + return typer + + @lower_builtin(datetime.date, types.int64, types.int64, types.int64) + def impl_ctor_timestamp(context, builder, sig, args): + typ = sig.return_type + year, month, day = args + nopython_date = builder.add(day, + builder.add(builder.shl(year, lir.Constant(lir.IntType(64), 32)), + builder.shl(month, lir.Constant(lir.IntType(64), 16)))) + return nopython_date +else: + class DatetimeDateType(types.Type): + def __init__(self): + super(DatetimeDateType, self).__init__( + name='DatetimeDateType()') + + datetime_date_type = DatetimeDateType() + + @typeof_impl.register(datetime.date) + def typeof_pd_timestamp(val, c): return datetime_date_type - return typer -@lower_builtin(datetime.date, types.int64, types.int64, types.int64) -def impl_ctor_timestamp(context, builder, sig, args): - typ = sig.return_type - year, month, day = args - ts = cgutils.create_struct_proxy(typ)(context, builder) - ts.year = year - ts.month = month - ts.day = day - return ts._getvalue() + + @register_model(DatetimeDateType) + class DatetimeDateModel(models.StructModel): + def __init__(self, dmm, fe_type): + members = [ + ('year', types.int64), + ('month', types.int64), + ('day', types.int64), + ] + models.StructModel.__init__(self, dmm, fe_type, members) + + make_attribute_wrapper(DatetimeDateType, 'year', 'year') + make_attribute_wrapper(DatetimeDateType, 'month', 'month') + make_attribute_wrapper(DatetimeDateType, 'day', 'day') + + @unbox(DatetimeDateType) + def unbox_datetime_date(typ, val, c): + year_obj = c.pyapi.object_getattr_string(val, "year") + month_obj = c.pyapi.object_getattr_string(val, "month") + day_obj = c.pyapi.object_getattr_string(val, "day") + + dt_date = cgutils.create_struct_proxy(typ)(c.context, c.builder) + dt_date.year = c.pyapi.long_as_longlong(year_obj) + dt_date.month = c.pyapi.long_as_longlong(month_obj) + dt_date.day = c.pyapi.long_as_longlong(day_obj) + + c.pyapi.decref(year_obj) + c.pyapi.decref(month_obj) + c.pyapi.decref(day_obj) + + is_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred()) + return NativeValue(dt_date._getvalue(), is_error=is_error) + + @box(DatetimeDateType) + def box_datetime_date(typ, val, c): + dt_date = cgutils.create_struct_proxy(typ)(c.context, c.builder, value=val) + year_obj = c.pyapi.long_from_longlong(dt_date.year) + month_obj = c.pyapi.long_from_longlong(dt_date.month) + day_obj = c.pyapi.long_from_longlong(dt_date.day) + dt_obj = c.pyapi.unserialize(c.pyapi.serialize_object(datetime.date)) + res = c.pyapi.call_function_objargs(dt_obj, (year_obj, month_obj, day_obj)) + c.pyapi.decref(year_obj) + c.pyapi.decref(month_obj) + c.pyapi.decref(day_obj) + return res + + @type_callable(datetime.date) + def type_timestamp(context): + def typer(year, month, day): + # TODO: check types + return datetime_date_type + return typer + + @lower_builtin(datetime.date, types.int64, types.int64, types.int64) + def impl_ctor_timestamp(context, builder, sig, args): + typ = sig.return_type + year, month, day = args + ts = cgutils.create_struct_proxy(typ)(context, builder) + ts.year = year + ts.month = month + ts.day = day + return ts._getvalue() #------------------------------------------------------------------------