Remove ._with_dtype_metadata in cudf classic groupby (rapidsai#21389)

Towards rapidsai#21229

The groupby results from pylibcudf needed to be preserved as `pylibcudf.Column`s before passed to another function that resolves the correct resulting type.

The added test to `conftest-patch.py` _appears_ to a separate cudf bug arising. `test_groupby_agg_extension` is using an indexing API to compute the expected result, but it doesn't seem to preserving the original string variant correctly.

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

Approvers:
  - GALI PREM SAGAR (https://github.com/galipremsagar)

URL: rapidsai#21389

Author: mroeschke

python/cudf/cudf/core/groupby/groupby.py

@@ -53,6 +53,7 @@
     CUDF_STRING_DTYPE,
     SIZE_TYPE_DTYPE,
     cudf_dtype_to_pa_type,
+    dtype_from_pylibcudf_column,
     get_dtype_of_same_kind,
 )
 from cudf.utils.performance_tracking import _performance_tracking
@@ -850,24 +851,21 @@ def _groups(
     def _aggregate(
         self, values: tuple[ColumnBase, ...], aggregations
     ) -> tuple[
-        list[list[ColumnBase]],
+        list[list[plc.Column]],
         list[ColumnBase],
         list[list[tuple[str, str]]],
     ]:
         included_aggregations = []
         column_included = []
         requests = []
-        # For any post-processing needed after pylibcudf aggregations
-        adjustments = []
-        result_columns: list[list[ColumnBase]] = []
+        result_columns: list[list[plc.Column]] = []
 
         for i, (col, aggs) in enumerate(
             zip(values, aggregations, strict=True)
         ):
             valid_aggregations = get_valid_aggregation(col.dtype)
             included_aggregations_i = []
             col_aggregations = []
-            adjustments_i = []
             for agg in aggs:
                 str_agg = str(agg)
                 if _is_unsupported_agg_for_type(col.dtype, str_agg):
@@ -881,12 +879,6 @@ def _aggregate(
                 ):
                     included_aggregations_i.append((agg, agg_obj.kind))
                     col_aggregations.append(agg_obj.plc_obj)
-                    if str_agg == "cumcount":
-                        # pandas 0-indexes cumulative count, see
-                        # https://github.com/rapidsai/cudf/issues/10237
-                        adjustments_i.append(lambda col: (col - 1))
-                    else:
-                        adjustments_i.append(lambda col: col)
             included_aggregations.append(included_aggregations_i)
             result_columns.append([])
             if col_aggregations:
@@ -896,7 +888,6 @@ def _aggregate(
                     )
                 )
                 column_included.append(i)
-                adjustments.append(adjustments_i)
 
         if not requests and any(len(v) > 0 for v in aggregations):
             raise pd.errors.DataError(
@@ -911,19 +902,15 @@ def _aggregate(
                     else plc_groupby.aggregate(requests)
                 )
 
-        for i, result, adjustments_i in zip(
-            column_included, results, adjustments, strict=True
-        ):
-            result_columns[i] = [
-                adj(ColumnBase.from_pylibcudf(col))
-                for col, adj in zip(
-                    result.columns(), adjustments_i, strict=True
-                )
-            ]
+        for i, result in zip(column_included, results, strict=True):
+            result_columns[i] = result.columns()
 
         return (
             result_columns,
-            [ColumnBase.from_pylibcudf(key) for key in keys.columns()],
+            [
+                ColumnBase.create(key, dtype_from_pylibcudf_column(key))
+                for key in keys.columns()
+            ],
             included_aggregations,
         )
 
@@ -1096,52 +1083,52 @@ def agg(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs):
             orig_dtypes,
             strict=True,
         ):
-            for agg_tuple, col in zip(aggs, cols, strict=True):
+            for agg_tuple, plc_result in zip(aggs, cols, strict=True):
                 agg, agg_kind = agg_tuple
                 agg_name = agg.__name__ if callable(agg) else agg
                 if multilevel:
                     key = (col_name, agg_name)
                 else:
                     key = col_name
+
+                create_dtype = dtype_from_pylibcudf_column(plc_result)
+                cast_dtype = None
                 if agg in {list, "collect"}:
                     # Collect wraps the original dtype in ListDtype (e.g., int -> list<int>)
-                    new_dtype = get_dtype_of_same_kind(
+                    create_dtype = get_dtype_of_same_kind(
                         orig_dtype, ListDtype(orig_dtype)
                     )
-                    col = ColumnBase.create(col.plc_column, new_dtype)
-
-                # Default: use column as-is
-                data[key] = col
-
                 # Override for specific aggregation types that need dtype adjustments
                 if agg_kind in {"COUNT", "SIZE", "ARGMIN", "ARGMAX"}:
-                    data[key] = col.astype(
-                        get_dtype_of_same_kind(orig_dtype, np.dtype(np.int64))
+                    cast_dtype = get_dtype_of_same_kind(
+                        orig_dtype, np.dtype(np.int64)
                     )
                 elif (
                     self.obj.empty
                     and (
                         isinstance(agg_name, str)
                         and agg_name in Reducible._SUPPORTED_REDUCTIONS
                     )
-                    and len(col) == 0
+                    and plc_result.size() == 0
                     and not isinstance(
-                        col.dtype,
+                        create_dtype,
                         (ListDtype, StructDtype, DecimalDtype),
                     )
                 ):
-                    data[key] = col.astype(orig_dtype)
+                    cast_dtype = orig_dtype
                 elif agg not in {list, "collect"}:
-                    # For non-collect aggregations, apply original dtype metadata
-                    if isinstance(orig_dtype, DecimalDtype):
-                        # `col` has a different precision than `orig_dtype`
-                        # hence we only preserve the kind of the dtype
-                        # and not the precision.
-                        data[key] = col._with_type_metadata(
-                            get_dtype_of_same_kind(orig_dtype, col.dtype)
-                        )
-                    else:
-                        data[key] = col._with_type_metadata(orig_dtype)
+                    create_dtype = get_dtype_of_same_kind(
+                        orig_dtype, create_dtype
+                    )
+
+                result_col = ColumnBase.create(plc_result, create_dtype)
+                if agg == "cumcount":
+                    # pandas 0-indexes cumulative count, see
+                    # https://github.com/rapidsai/cudf/issues/10237
+                    result_col = result_col - 1
+                if cast_dtype is not None:
+                    result_col = result_col.astype(cast_dtype)
+                data[key] = result_col
         data = ColumnAccessor(data, multiindex=multilevel)
         if not multilevel:
             data = data.rename_levels({np.nan: None}, level=0)

python/cudf/cudf/pandas/scripts/conftest-patch.py

@@ -2039,6 +2039,12 @@ def set_copy_on_write_option():
     "tests/extension/test_string.py::TestStringArray::test_container_shift[python-True-0-indices1-True]",
     "tests/extension/test_string.py::TestStringArray::test_container_shift[python-True-2-indices2-False]",
     "tests/extension/test_string.py::TestStringArray::test_container_shift[python-True-2-indices2-True]",
+    "tests/extension/test_string.py::TestStringArray::test_groupby_agg_extension[string=string[pyarrow]-True]",
+    "tests/extension/test_string.py::TestStringArray::test_groupby_agg_extension[string=string[pyarrow]-False]",
+    "tests/extension/test_string.py::TestStringArray::test_groupby_agg_extension[string=str[pyarrow]-True]",
+    "tests/extension/test_string.py::TestStringArray::test_groupby_agg_extension[string=str[pyarrow]-False]",
+    "tests/extension/test_string.py::TestStringArray::test_groupby_agg_extension[string=str[python]-True]",
+    "tests/extension/test_string.py::TestStringArray::test_groupby_agg_extension[string=str[python]-False]",
     "tests/extension/test_string.py::TestStringArray::test_grouping_grouper[pyarrow-False]",
     "tests/extension/test_string.py::TestStringArray::test_grouping_grouper[pyarrow-True]",
     "tests/extension/test_string.py::TestStringArray::test_grouping_grouper[pyarrow_numpy-False]",