Add Categorical

AlexanderKalistratov · May 26, 2020 · Apr 21, 2020 · May 13, 2020 · Apr 21, 2020 · Apr 23, 2020
commit a84d422d5985d9938a09050e1824fe50942b88de
diff --git a/sdc/datatypes/categorical/boxing.py b/sdc/datatypes/categorical/boxing.py
@@ -25,9 +25,16 @@
 # *****************************************************************************
 
 from numba.extending import box, unbox, NativeValue
+from numba.targets import boxing
+from numba.targets.imputils import lower_constant
+from numba.targets import arrayobj
+from numba import types
 
 from . import pandas_support
-from .types import CategoricalDtypeType
+from .types import (
+    CategoricalDtypeType,
+    Categorical,
+)
 
 
 @box(CategoricalDtypeType)
@@ -39,3 +46,45 @@ def box_CategoricalDtype(typ, val, c):
 @unbox(CategoricalDtypeType)
 def unbox_CategoricalDtype(typ, val, c):
     return NativeValue(c.context.get_dummy_value())
+
+
+@box(Categorical)
+def box_Categorical(typ, val, c):
+    pandas_module_name = c.context.insert_const_string(c.builder.module, "pandas")
+    pandas_module = c.pyapi.import_module_noblock(pandas_module_name)
+
+    constructor = c.pyapi.object_getattr_string(pandas_module, "Categorical")
+
+    empty_list = c.pyapi.list_new(c.context.get_constant(types.intp, 0))
+    args = c.pyapi.tuple_pack([empty_list])
+    categorical = c.pyapi.call(constructor, args)
+
+    dtype = box_CategoricalDtype(typ.pd_dtype, val, c)
+    c.pyapi.object_setattr_string(categorical, "_dtype", dtype)
+
+    codes = boxing.box_array(typ.codes, val, c)
+    c.pyapi.object_setattr_string(categorical, "_codes", codes)
+
+    c.pyapi.decref(codes)
+    c.pyapi.decref(dtype)
+    c.pyapi.decref(args)
+    c.pyapi.decref(empty_list)
+    c.pyapi.decref(constructor)
+    c.pyapi.decref(pandas_module)
+    return categorical
+
+
+@unbox(Categorical)
+def unbox_Categorical(typ, val, c):
+    codes = c.pyapi.object_getattr_string(val, "codes")
+    native_value = boxing.unbox_array(typ.codes, codes, c)
+    c.pyapi.decref(codes)
+    return native_value
+
+
+@lower_constant(Categorical)
+def constant_Categorical(context, builder, ty, pyval):
+    """
+    Create a constant Categorical.
+    """
+    return arrayobj.constant_array(context, builder, ty.codes, pyval.codes)
diff --git a/sdc/datatypes/categorical/models.py b/sdc/datatypes/categorical/models.py
@@ -27,7 +27,11 @@
 from numba.extending import models
 from numba.extending import register_model
 
-from .types import CategoricalDtypeType
+from .types import (
+    CategoricalDtypeType,
+    Categorical,
+)
 
 
 register_model(CategoricalDtypeType)(models.OpaqueModel)
+register_model(Categorical)(models.ArrayModel)
diff --git a/sdc/datatypes/categorical/pdimpl.py b/sdc/datatypes/categorical/pdimpl.py
@@ -31,8 +31,13 @@
 from numba.extending import type_callable
 from numba.extending import lower_builtin
 from numba import types
+from numba import typeof
+from numba import objmode
 
-from .types import CategoricalDtypeType
+from .types import (
+    CategoricalDtypeType,
+    Categorical,
+)
 
 
 # Possible alternative implementations:
@@ -120,3 +125,97 @@ def codegen(context, builder, signature, args):
 #     # All CategoricalDtype objects are dummy values in LLVM.
 #     # They only exist in the type level.
 #     return context.get_dummy_value()
+
+
+# TODO: use dtype too
+def _reconstruct_Categorical(values):
+    values_list = [v.literal_value for v in values]
+    return pd.Categorical(values=values_list)
+
+
+@overload(pd.Categorical)
+def _Categorical(values, categories=None, ordered=None, dtype=None, fastpath=False):
+    """
+    Implementation of constructor for pandas Categorical via objmode.
+    """
+    # TODO: support other parameters (only values now)
+
+    ty = typeof(_reconstruct_Categorical(values))
+    tyname = ty.name
+    setattr(types, tyname, ty)
+
+    from textwrap import dedent
+    text = dedent(f"""
+    def impl(values, categories=None, ordered=None, dtype=None, fastpath=False):
+        with objmode(categorical='vars()["{tyname}"]'):
+            categorical = pd.Categorical(values, categories, ordered, dtype, fastpath)
+        return categorical
+    """)
+    globals, locals = {'objmode': objmode, 'pd': pd}, {}
+    exec(text, globals, locals)
+    impl = locals['impl']
+    return impl
+
+
+# @type_callable(pd.Categorical)
+# def type_Categorical_constructor(context):
+#     """
+#     Similar to @infer_global(np.array).
+#     """
+#     def typer(values, categories=None, ordered=None, dtype=None, fastpath=False):
+#         # from numba.typing import npydecl
+#         # codes = npydecl.NpArray(context).generic()(values)
+#         categorical = _reconstruct_Categorical(values)
+#         return typeof(categorical)
+
+#     return typer
+
+
+# @lower_builtin(pd.Categorical, types.Any)
+# # @lower_builtin(np.Categorical, types.Any, types.DTypeSpec)
+# def pd_Categorical(context, builder, sig, args):
+#     """
+#     Similar to @lower_builtin(np.array, ...).
+#     """
+#     from numba.targets import arrayobj
+#     codes = sig.return_type.codes
+#     return arrayobj.np_array(context, builder, sig.replace(return_type=codes), args)
+
+
+# via intrinsic
+# @overload(pd.Categorical)
+# def _Categorical(values, categories=None, ordered=None, dtype=None, fastpath=False):
+#     """
+#     Implementation of constructor for pandas Categorical.
+#     """
+#     def impl(values, categories=None, ordered=None, dtype=None, fastpath=False):
+#         return _Categorical_intrinsic(values, categories, ordered, dtype, fastpath)
+#     return impl
+
+
+# @intrinsic
+# def _Categorical_intrinsic(typingctx, values, categories, ordered, dtype, fastpath):
+#     """
+#     Creates Categorical object.
+#     """
+#     if isinstance(values, types.Tuple):
+#         values_list = [v.literal_value for v in values]
+#         categorical = pd.Categorical(values=values_list)
+#         return_type = typeof(categorical)
+
+#         def codegen(context, builder, signature, args):
+#             [values] = args
+#             # TODO: can not recall similar function
+#             native_value = boxing.unbox_array(typ.codes, codes, c)
+#             return native_value
+
+#     sig = return_type(values, categories, ordered, dtype, fastpath)
+#     return sig, codegen
+
+#     # return_type = Categorical(dtype=CategoricalDtypeType(), codes=types.Array(types.int8, 1, 'C'))
+#     # sig = return_type(values)
+
+#     # def codegen(context, builder, signature, args):
+#     #     return context.get_dummy_value()
+
+#     # return sig, codegen
diff --git a/sdc/datatypes/categorical/rewrites.py b/sdc/datatypes/categorical/rewrites.py
@@ -29,3 +29,4 @@
 
 
 register_tuplify(pd.CategoricalDtype, 'categories')
+register_tuplify(pd.Categorical, 'values')
diff --git a/sdc/datatypes/categorical/typeof.py b/sdc/datatypes/categorical/typeof.py
@@ -31,10 +31,23 @@
 import pandas as pd
 
 from numba.extending import typeof_impl
+from numba import numpy_support
+from numba import typeof
 
 from . import pandas_support
+from .types import Categorical
 
 
 @typeof_impl.register(pd.CategoricalDtype)
 def _typeof_CategoricalDtype(val, c):
     return pandas_support.from_dtype(val)
+
+
+@typeof_impl.register(pd.Categorical)
+def _typeof_Categorical(val, c):
+    try:
+        dtype = pandas_support.from_dtype(val.dtype)
+    except NotImplementedError:
+        raise ValueError("Unsupported Categorical dtype: %s" % (val.dtype,))
+    codes = typeof(val.codes)
+    return Categorical(dtype=dtype, codes=codes)
diff --git a/sdc/datatypes/categorical/types.py b/sdc/datatypes/categorical/types.py
@@ -30,6 +30,13 @@
 
 from numba import types
 
+import numpy as np
+
+
+__all__ = [
+    'CategoricalDtypeType',
+    'Categorical',
+]
 
 
 # TODO: consider renaming to CategoricalDtype b/c Categorical - not CategoricalType
@@ -47,3 +54,68 @@ def __init__(self, categories=None, ordered=None):
         name = 'CategoricalDtype(categories={}, ordered={})'.format(
             self.categories, self.ordered)
         super().__init__(name=name)
+
+    def __len__(self):
+        return len(self.categories) if self.categories else 0
+
+    @property
+    def dtype(self):
+        # TODO: take dtype from categories array
+        return types.int64
+
+    def int_type(self):
+        """
+        Return minimal int type to fit all categories.
+        """
+        dtype = types.int64
+        n_cats = len(self.categories)
+        if n_cats < np.iinfo(np.int8).max:
+            dtype = types.int8
+        elif n_cats < np.iinfo(np.int16).max:
+            dtype = types.int16
+        elif n_cats < np.iinfo(np.int32).max:
+            dtype = types.int32
+        return dtype
+
+
+# TODO: make ArrayCompatible. It will make reuse Array boxing, unboxing.
+class Categorical(types.Type):
+    """
+    Numba type for pandas Categorical.
+
+    Contains:
+        codes -> array-like
+        dtype -> CategoricalDtypeType
+    """
+    def __init__(self, dtype, codes=None):
+        assert(isinstance(dtype, CategoricalDtypeType))
+        self.pd_dtype = dtype
+        self.codes = codes or types.Array(self.pd_dtype.int_type(), ndim=1, layout='C')
+        # TODO: store dtype for categories values and use it for dtype
+        super().__init__(name=self.__repr__())
+
+    def __repr__(self):
+        def Array__repr__(array):
+            return "Array({}, {}, {})".format(
+                self.codes.dtype.__repr__(),
+                self.codes.ndim.__repr__(),
+                self.codes.layout.__repr__()
+            )
+
+        dtype = self.pd_dtype.__repr__()
+        codes = Array__repr__(self.codes)
+        return 'Categorical(dtype={}, codes={})'.format(dtype, codes)
+
+    @property
+    def categories(self):
+        return self.pd_dtype.categories
+
+    # Properties for model
+
+    @property
+    def ndim(self):
+        return self.codes.ndim
+
+    @property
+    def dtype(self):
+        return self.codes.dtype
diff --git a/sdc/hiframes/pd_categorical_ext.py b/sdc/hiframes/pd_categorical_ext.py
@@ -49,7 +49,7 @@ def __init__(self, dtype):
             dtype, 1, 'C', name='CategoricalArray({})'.format(dtype))
 
 
-@unbox(CategoricalArray)
+# @unbox(CategoricalArray)
 def unbox_categorical_array(typ, val, c):
     arr_obj = c.pyapi.object_getattr_string(val, "codes")
     # c.pyapi.print_object(arr_obj)
@@ -71,7 +71,7 @@ def get_categories_int_type(cat_dtype):
     return dtype
 
 
-@box(CategoricalArray)
+# @box(CategoricalArray)
 def box_categorical_array(typ, val, c):
     dtype = typ.dtype
     mod_name = c.context.insert_const_string(c.builder.module, "pandas")

diff --git a/sdc/tests/categorical/__init__.py b/sdc/tests/categorical/__init__.py
@@ -24,4 +24,5 @@
 # EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 # *****************************************************************************
 
+from . import test_categorical
 from . import test_categoricaldtype
Original file line number	Diff line number	Diff line change
Expand Up		@@ -29,3 +29,4 @@


		register_tuplify(pd.CategoricalDtype, 'categories')
		register_tuplify(pd.Categorical, 'values')