Add new ColumnBase.create API (rapidsai#21187)

The new API avoids the design inconsistencies inherent to the current `from_pylibcudf`->`_with_type_metadata` API. The old approach is highly error-prone, and missing `_with_type_metadata` calls have been responsible for a large number of bugs over the lifetime of the project. With the various recently completed refactorings, a ColumnBase subclass is strictly defined by an underlying pylibcudf Column and a dtype, which broke the (incorrect, but convenient) assumption of 1-1 correspondence between pylibcudf data types and cudf ColumnBase types. Since the same type of pylibcudf data could be underlying multiple cudf ColumnBase subclasses, and since the dtype is the primary way to indicate this, forcing both to be provided on construction is more correct and more robust.

I'll be opening subsequent PRs to roll this change out to more of the package, this PR just includes a few examples to clarify the new API's usage.

Authors:
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Matthew Roeschke (https://github.com/mroeschke)

URL: rapidsai#21187

Author: vyasr

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

@@ -200,10 +200,9 @@ def _binaryop(self, other: ColumnBinaryOperand, op: str) -> ColumnBase:
                     )
             # We'll compare self's decategorized values later for non-CategoricalColumn
         else:
-            codes = column.as_column(
+            other = column.as_column(
                 self._encode(other), length=len(self), dtype=self.codes.dtype
-            )
-            other = codes._with_type_metadata(self.dtype)
+            )._with_type_metadata(self.dtype)
         equality_ops = {"__eq__", "__ne__", "NULL_EQUALS", "NULL_NOT_EQUALS"}
         if not self.ordered and op not in equality_ops:
             raise TypeError(

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

@@ -431,13 +431,7 @@ def set_mask(self, mask: Buffer | None) -> Self:
         new_plc_column = self.plc_column.with_mask(new_mask, new_null_count)
         return cast(
             "Self",
-            (
-                type(self)
-                .from_pylibcudf(
-                    new_plc_column,
-                )
-                ._with_type_metadata(self.dtype)
-            ),
+            ColumnBase.create(new_plc_column, self.dtype),
         )
 
     @property
@@ -617,12 +611,129 @@ def _wrap_buffers(col: plc.Column) -> plc.Column:
             validate=False,
         )
 
+    @staticmethod
+    def create(col: plc.Column, dtype: DtypeObj) -> ColumnBase:
+        """
+        Create a Column from a pylibcudf.Column with an explicit cudf dtype.
+
+        This is the primary factory for ColumnBase construction. It always requires
+        an explicit dtype to ensure type safety. If you need to infer the dtype from
+        the pylibcudf Column, use dtype_from_pylibcudf_column() first:
+
+            dtype = dtype_from_pylibcudf_column(plc_col)
+            col = ColumnBase.create(plc_col, dtype)
+        """
+        # Wrap buffers recursively
+        wrapped = ColumnBase._wrap_buffers(col)
+
+        # Dispatch to the appropriate subclass based on dtype
+        target_cls = ColumnBase._dispatch_subclass_from_dtype(dtype)
+
+        # Validate dtype compatibility with the column structure using the
+        # target subclass's _validate_args method (includes recursive validation)
+        wrapped, dtype = target_cls._validate_args(wrapped, dtype)
+
+        # Construct the instance using the subclass's _from_preprocessed method
+        # Skip validation since we already validated above
+        return target_cls._from_preprocessed(
+            plc_column=wrapped,
+            dtype=dtype,
+            validate=False,
+        )
+
+    @staticmethod
+    def _dispatch_subclass_from_dtype(dtype: DtypeObj) -> type[ColumnBase]:
+        """
+        Dispatch to the appropriate ColumnBase subclass based on dtype.
+
+        This function determines which ColumnBase subclass should be used
+        to construct a column with the given dtype.
+        """
+        # Special pandas extension types
+        if isinstance(dtype, pd.DatetimeTZDtype):
+            return cudf.core.column.datetime.DatetimeTZColumn
+        if isinstance(dtype, CategoricalDtype):
+            return cudf.core.column.CategoricalColumn
+
+        # Temporal types (by kind)
+        if dtype.kind == "M":
+            return cudf.core.column.DatetimeColumn
+        if dtype.kind == "m":
+            return cudf.core.column.TimeDeltaColumn
+
+        # String types
+        if (
+            dtype == CUDF_STRING_DTYPE
+            or (hasattr(dtype, "kind") and dtype.kind == "U")
+            or isinstance(dtype, pd.StringDtype)
+            or (isinstance(dtype, pd.ArrowDtype) and dtype.kind == "U")
+        ):
+            return cudf.core.column.StringColumn
+
+        # cuDF custom types
+        if isinstance(dtype, ListDtype):
+            return cudf.core.column.ListColumn
+        if isinstance(dtype, IntervalDtype):
+            return cudf.core.column.IntervalColumn
+        if isinstance(dtype, StructDtype):
+            return cudf.core.column.StructColumn
+
+        # Decimal types
+        if isinstance(dtype, cudf.Decimal128Dtype):
+            return cudf.core.column.Decimal128Column
+        if isinstance(dtype, cudf.Decimal64Dtype):
+            return cudf.core.column.Decimal64Column
+        if isinstance(dtype, cudf.Decimal32Dtype):
+            return cudf.core.column.Decimal32Column
+
+        # Numerical types
+        if dtype.kind in "iufb":
+            return cudf.core.column.NumericalColumn
+
+        raise TypeError(f"Unrecognized dtype: {dtype}")
+
+    @staticmethod
+    def _validate_dtype_recursively(col: plc.Column, dtype: DtypeObj) -> None:
+        """
+        Validate dtype compatibility by dispatching to the appropriate ColumnBase
+        subclass's _validate_args method.
+
+        This method is used for recursive validation in nested types (List, Struct,
+        Interval). It dispatches to the correct ColumnBase subclass based on dtype
+        and calls its _validate_args method, which may recursively call this method
+        for nested children.
+
+        Parameters
+        ----------
+        col : plc.Column
+            The pylibcudf Column to validate.
+        dtype : DtypeObj
+            The cudf dtype to validate against.
+
+        Raises
+        ------
+        ValueError
+            If the dtype is incompatible with the Column.
+        """
+        # Skip validation for empty columns (INT8 with all nulls). These are created
+        # by _wrap_buffers() from EMPTY columns and may have inaccurate dtype metadata.
+        # For example, an empty list [] has element_type=object but child is INT8.
+        if (
+            col.type().id() == plc.TypeId.INT8
+            and col.null_count() == col.size()
+        ):
+            return
+
+        # Dispatch to the appropriate subclass and use its _validate_args
+        target_cls = ColumnBase._dispatch_subclass_from_dtype(dtype)
+        target_cls._validate_args(col, dtype)
+
     @staticmethod
     def from_pylibcudf(col: plc.Column) -> ColumnBase:
         """Create a Column from a pylibcudf.Column.
 
         This function will generate a Column pointing to the provided pylibcudf
-        Column.  It will directly access the data and mask buffers of the
+        Column. It will directly access the data and mask buffers of the
         pylibcudf Column, so the newly created object is not tied to the
         lifetime of the original pylibcudf.Column.
 
@@ -636,51 +747,17 @@ def from_pylibcudf(col: plc.Column) -> ColumnBase:
         pylibcudf.Column
             A new pylibcudf.Column referencing the same data.
         """
+        # Wrap buffers first so that dtypes are compatible with dtype_from_pylibcudf_column
         wrapped = ColumnBase._wrap_buffers(col)
-
         dtype = dtype_from_pylibcudf_column(wrapped)
-
-        cls: type[ColumnBase]
-        if isinstance(dtype, pd.DatetimeTZDtype):
-            cls = cudf.core.column.datetime.DatetimeTZColumn
-        elif dtype.kind == "M":
-            cls = cudf.core.column.DatetimeColumn
-        elif dtype.kind == "m":
-            cls = cudf.core.column.TimeDeltaColumn
-        elif (
-            dtype == CUDF_STRING_DTYPE
-            or dtype.kind == "U"
-            or isinstance(dtype, pd.StringDtype)
-            or (isinstance(dtype, pd.ArrowDtype) and dtype.kind == "U")
-        ):
-            cls = cudf.core.column.StringColumn
-        elif isinstance(dtype, ListDtype):
-            cls = cudf.core.column.ListColumn
-        elif isinstance(dtype, IntervalDtype):
-            cls = cudf.core.column.IntervalColumn
-        elif isinstance(dtype, StructDtype):
-            cls = cudf.core.column.StructColumn
-        elif isinstance(dtype, cudf.Decimal64Dtype):
-            cls = cudf.core.column.Decimal64Column
-        elif isinstance(dtype, cudf.Decimal32Dtype):
-            cls = cudf.core.column.Decimal32Column
-        elif isinstance(dtype, cudf.Decimal128Dtype):
-            cls = cudf.core.column.Decimal128Column
-        elif dtype.kind in "iufb":
-            cls = cudf.core.column.NumericalColumn
-        else:
-            raise TypeError(f"Unrecognized dtype: {dtype}")
-
-        return cls._from_preprocessed(
-            plc_column=wrapped,
-            dtype=dtype,
-        )
+        return ColumnBase.create(wrapped, dtype)
 
     @classmethod
     def _from_preprocessed(
         cls,
         plc_column: plc.Column,
         dtype: DtypeObj,
+        validate: bool = True,
     ) -> Self:
         # TODO: This function bypassess some of the buffer copying/wrapping that would
         # be done in from_pylibcudf, so it is only ever safe to call this in situations
@@ -689,7 +766,8 @@ def _from_preprocessed(
         # in from_pylibcudf, but for now it is necessary for the various
         # _with_type_metadata calls.
         self = cls.__new__(cls)
-        plc_column, dtype = self._validate_args(plc_column, dtype)
+        if validate:
+            plc_column, dtype = self._validate_args(plc_column, dtype)
         self.plc_column = plc_column
         self._dtype = dtype
         self._distinct_count = {}
@@ -921,9 +999,9 @@ def dropna(self) -> Self:
         if self.has_nulls():
             return cast(
                 "Self",
-                ColumnBase.from_pylibcudf(
-                    stream_compaction.drop_nulls([self])[0]
-                )._with_type_metadata(self.dtype),
+                ColumnBase.create(
+                    stream_compaction.drop_nulls([self])[0], self.dtype
+                ),
             )
         else:
             return self.copy()
@@ -1120,6 +1198,11 @@ def copy(self, deep: bool = True) -> Self:
         plc_col = self.plc_column
         if deep:
             plc_col = plc_col.copy()
+        # For nested types (e.g., list<list<int>>), self.dtype may not accurately
+        # reflect the actual plc_column structure. Some operations (like groupby
+        # collect on a list column) create nested structures but don't update the
+        # stored dtype to reflect the new nesting level. Using _with_type_metadata()
+        # is more permissive and handles these cases.
         return cast(
             "Self",
             (
@@ -1360,7 +1443,8 @@ def _scatter_by_column(
         else:
             return cast(
                 "Self",
-                ColumnBase.from_pylibcudf(
+                type(self)
+                .from_pylibcudf(
                     copying.scatter(
                         cast("list[plc.Scalar]", [value])
                         if isinstance(value, plc.Scalar)
@@ -1369,7 +1453,8 @@ def _scatter_by_column(
                         [self],
                         bounds_check=bounds_check,
                     )[0]
-                )._with_type_metadata(self.dtype),
+                )
+                ._with_type_metadata(self.dtype),
             )
 
     def _check_scatter_key_length(
@@ -1468,10 +1553,9 @@ def fillna(
                 input_col.plc_column,
                 plc_replace,
             )
-            result = type(self).from_pylibcudf(plc_column)
         return cast(
             "Self",
-            result._with_type_metadata(self.dtype),
+            ColumnBase.create(plc_column, self.dtype),
         )
 
     def is_valid(self) -> ColumnBase:
@@ -1778,11 +1862,7 @@ def sort_values(
             )
             return cast(
                 "Self",
-                (
-                    type(self)
-                    .from_pylibcudf(plc_table.columns()[0])
-                    ._with_type_metadata(self.dtype)
-                ),
+                ColumnBase.create(plc_table.columns()[0], self.dtype),
             )
 
     def distinct_count(self, dropna: bool = True) -> int:
@@ -1902,9 +1982,9 @@ def apply_boolean_mask(self, mask: ColumnBase) -> ColumnBase:
         if mask.dtype.kind != "b":
             raise ValueError("boolean_mask is not boolean type.")
 
-        return ColumnBase.from_pylibcudf(
-            stream_compaction.apply_boolean_mask([self], mask)[0]
-        )._with_type_metadata(self.dtype)
+        return ColumnBase.create(
+            stream_compaction.apply_boolean_mask([self], mask)[0], self.dtype
+        )
 
     def argsort(
         self,
@@ -2029,9 +2109,10 @@ def unique(self) -> Self:
         else:
             return cast(
                 "Self",
-                ColumnBase.from_pylibcudf(
-                    stream_compaction.drop_duplicates([self], keep="first")[0]
-                )._with_type_metadata(self.dtype),
+                ColumnBase.create(
+                    stream_compaction.drop_duplicates([self], keep="first")[0],
+                    self.dtype,
+                ),
             )
 
     @staticmethod
@@ -2174,7 +2255,7 @@ def deserialize(cls, header: dict, frames: list) -> ColumnBase:
         assert len(frames) == 0, (
             f"{len(frames)} frame(s) remaining after deserialization"
         )
-        return cls.from_pylibcudf(plc_column)._with_type_metadata(dtype)
+        return ColumnBase.create(plc_column, dtype)
 
     def unary_operator(self, unaryop: str) -> ColumnBase:
         raise TypeError(
@@ -2369,9 +2450,7 @@ def split_by_offsets(
             for col in cols:
                 yield cast(
                     "Self",
-                    type(self)
-                    .from_pylibcudf(col)
-                    ._with_type_metadata(self.dtype),
+                    ColumnBase.create(col, self.dtype),
                 )
 
     def one_hot_encode(self, categories: ColumnBase) -> Generator[ColumnBase]:
@@ -2386,20 +2465,23 @@ def one_hot_encode(self, categories: ColumnBase) -> Generator[ColumnBase]:
                 type(self).from_pylibcudf(col) for col in plc_table.columns()
             )
 
+    # TODO: Currently this method is only used once, in ExponentialMovingWindow. That
+    # suggests a potential refactoring opportunity to make EWM play better with the rest
+    # of our aggregation/reduction framework.
     def scan(self, scan_op: str, inclusive: bool, **kwargs: Any) -> Self:
         with self.access(mode="read", scope="internal"):
-            return cast(
-                "Self",
-                type(self).from_pylibcudf(
-                    plc.reduce.scan(
-                        self.plc_column,
-                        aggregation.make_aggregation(scan_op, kwargs).plc_obj,
-                        plc.reduce.ScanType.INCLUSIVE
-                        if inclusive
-                        else plc.reduce.ScanType.EXCLUSIVE,
-                    )
-                ),
+            plc_result = plc.reduce.scan(
+                self.plc_column,
+                aggregation.make_aggregation(scan_op, kwargs).plc_obj,
+                plc.reduce.ScanType.INCLUSIVE
+                if inclusive
+                else plc.reduce.ScanType.EXCLUSIVE,
             )
+        return cast("Self", ColumnBase.create(plc_result, self.dtype))
+
+    def _scan(self, op: str) -> ColumnBase:
+        """Default cumulative scan implementation for DataFrame.cum* methods."""
+        return self.scan(op.replace("cum", ""), inclusive=True)
 
     def reduce(self, reduction_op: str, **kwargs: Any) -> ScalarLike:
         col_dtype = self._reduction_result_dtype(reduction_op)
@@ -3428,8 +3510,9 @@ def concat_columns(objs: Sequence[ColumnBase]) -> ColumnBase:
     with access_columns(  # type: ignore[assignment]
         *objs_with_len, mode="read", scope="internal"
     ) as objs_with_len:
-        return ColumnBase.from_pylibcudf(
+        return ColumnBase.create(
             plc.concatenate.concatenate(
                 [col.plc_column for col in objs_with_len]
-            )
-        )._with_type_metadata(objs_with_len[0].dtype)
+            ),
+            objs_with_len[0].dtype,
+        )