Permanently back cuDF column by a pylibcudf.Column

python/cudf/cudf/core/column/categorical.py

@@ -11,6 +11,8 @@
 import pyarrow as pa
 from typing_extensions import Self
 
+import pylibcudf as plc
+
 import cudf
 from cudf.api.types import is_scalar
 from cudf.core.column import column
@@ -29,12 +31,11 @@
 if TYPE_CHECKING:
     from collections.abc import Mapping, MutableSequence, Sequence
 
-    from pylibcudf import Scalar as plc_Scalar
-
     from cudf._typing import (
         ColumnBinaryOperand,
         ColumnLike,
         Dtype,
+        DtypeObj,
         ScalarLike,
     )
     from cudf.core.buffer import Buffer
@@ -52,18 +53,6 @@
 _DEFAULT_CATEGORICAL_VALUE = np.int8(-1)
 
 
-def validate_categorical_children(children) -> None:
-    if not (
-        len(children) == 1
-        and isinstance(children[0], cudf.core.column.numerical.NumericalColumn)
-        and children[0].dtype.kind in "iu"
-    ):
-        # TODO: Enforce unsigned integer?
-        raise ValueError(
-            "Must specify exactly one child NumericalColumn of integers for representing the codes."
-        )
-
-
 class CategoricalColumn(column.ColumnBase):
     """
     Implements operations for Columns of Categorical type
@@ -94,35 +83,59 @@ class CategoricalColumn(column.ColumnBase):
         "__gt__",
         "__ge__",
     }
+    # TODO: See if we can narrow these integer types
+    _VALID_PLC_TYPES = {
+        plc.TypeId.INT8,
+        plc.TypeId.INT16,
+        plc.TypeId.INT32,
+        plc.TypeId.INT64,
+        plc.TypeId.UINT8,
+        plc.TypeId.UINT16,
+        plc.TypeId.UINT32,
+        plc.TypeId.UINT64,
+    }
 
     def __init__(
         self,
-        data: None,
+        plc_column: plc.Column,
         size: int,
         dtype: CategoricalDtype,
-        mask: Buffer | None,
         offset: int,
         null_count: int,
-        children: tuple[NumericalColumn],
-    ):
-        if data is not None:
-            raise ValueError(f"{data=} must be None")
-        validate_categorical_children(children)
+        exposed: bool,
+    ) -> None:
         if not isinstance(dtype, CategoricalDtype):
             raise ValueError(
                 f"{dtype=} must be cudf.CategoricalDtype instance."
             )
         super().__init__(
-            data=data,
+            plc_column=plc_column,
             size=size,
             dtype=dtype,
-            mask=mask,
             offset=offset,
             null_count=null_count,
-            children=children,
+            exposed=exposed,
         )
         self._codes = self.children[0].set_mask(self.mask)
 
+    def _get_data_buffer_from_pylibcudf_column(
+        self, plc_column: plc.Column, exposed: bool
+    ) -> None:
+        """
+        This column considers the plc_column (i.e. codes) as children
+        """
+        return None
+
+    def _get_children_from_pylibcudf_column(
+        self, plc_column: plc.Column, dtype: DtypeObj, exposed: bool
+    ) -> tuple[ColumnBase]:
+        """
+        This column considers the plc_column (i.e. codes) as children
+        """
+        return (
+            type(self).from_pylibcudf(plc_column, data_ptr_exposed=exposed),
+        )
+
     @property
     def base_size(self) -> int:
         return int(
@@ -136,15 +149,6 @@ def __contains__(self, item: ScalarLike) -> bool:
             return False
         return self._encode(item) in self.codes
 
-    def set_base_data(self, value):
-        if value is not None:
-            raise RuntimeError(
-                "CategoricalColumns do not use data attribute of Column, use "
-                "`set_base_children` instead"
-            )
-        else:
-            super().set_base_data(value)
-
     def _process_values_for_isin(
         self, values: Sequence
     ) -> tuple[ColumnBase, ColumnBase]:
@@ -157,7 +161,6 @@ def set_base_mask(self, value: Buffer | None) -> None:
 
     def set_base_children(self, value: tuple[NumericalColumn]) -> None:  # type: ignore[override]
         super().set_base_children(value)
-        validate_categorical_children(value)
         self._codes = value[0].set_mask(self.mask)
 
     @property
@@ -216,7 +219,7 @@ def __setitem__(self, key, value):
 
     def _fill(
         self,
-        fill_value: plc_Scalar,
+        fill_value: plc.Scalar,
         begin: int,
         end: int,
         inplace: bool = False,
@@ -373,7 +376,7 @@ def unique(self) -> Self:
 
     def _cast_self_and_other_for_where(
         self, other: ScalarLike | ColumnBase, inplace: bool
-    ) -> tuple[ColumnBase, plc_Scalar | ColumnBase]:
+    ) -> tuple[ColumnBase, plc.Scalar | ColumnBase]:
         if is_scalar(other):
             try:
                 other = self._encode(other)
@@ -568,7 +571,7 @@ def notnull(self) -> ColumnBase:
 
     def _validate_fillna_value(
         self, fill_value: ScalarLike | ColumnLike
-    ) -> plc_Scalar | ColumnBase:
+    ) -> plc.Scalar | ColumnBase:
         """Align fill_value for .fillna based on column type."""
         if is_scalar(fill_value):
             if fill_value != _DEFAULT_CATEGORICAL_VALUE:
@@ -710,13 +713,12 @@ def _concat(
     def _with_type_metadata(self: Self, dtype: Dtype) -> Self:
         if isinstance(dtype, CategoricalDtype):
             return type(self)(
-                data=self.data,  # type: ignore[arg-type]
+                plc_column=self.plc_column,
                 size=self.size,
                 dtype=dtype,
-                mask=self.base_mask,
                 offset=self.offset,
                 null_count=self.null_count,
-                children=self.base_children,  # type: ignore[arg-type]
+                exposed=False,
             )
 
         return self