diff --git a/python/cudf/cudf/core/groupby/groupby.py b/python/cudf/cudf/core/groupby/groupby.py index 8d4bca3b5d91..3f9237e1e6d5 100644 --- a/python/cudf/cudf/core/groupby/groupby.py +++ b/python/cudf/cudf/core/groupby/groupby.py @@ -882,8 +882,26 @@ def rank(x): result = self.agg(rank) - # pandas always returns floats: - return result.astype(np.dtype(np.float64)) + # pandas always returns floats, staying within the value column's + # dtype family: numpy -> float64, masked (Int64/Float32/...) -> + # Float64, arrow -> double[pyarrow] + target = np.dtype(np.float64) + if result.ndim == 1: + source_dtype = ( + self.obj.dtype if self.obj.ndim == 1 else result.dtype + ) + return result.astype(get_dtype_of_same_kind(source_dtype, target)) + return result.astype( + { + label: get_dtype_of_same_kind( + self.obj._data[label].dtype + if self.obj.ndim == 2 and label in self.obj._data + else result_dtype, + target, + ) + for label, result_dtype in result._dtypes + } + ) @property def _groupby(self): diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index ec54cec49757..b661a90f5d49 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -1877,10 +1877,6 @@ def pytest_unconfigure(config): "tests/groupby/methods/test_nth.py::test_slice[arg2-expected_rows2]": "TODO: Add a reason for failure", "tests/groupby/methods/test_nth.py::test_slice[arg3-expected_rows3]": "TODO: Add a reason for failure", "tests/groupby/methods/test_quantile.py::test_groupby_quantile_nonmulti_levels_order": "tm.assert_equal compares MultiIndex.levels, which cudf returns as a list rather than the FrozenList pandas produces", - "tests/groupby/methods/test_rank.py::test_rank_avg_even_vals[True-int32]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_rank.py::test_rank_avg_even_vals[True-int64]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_rank.py::test_rank_avg_even_vals[True-uint32]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_rank.py::test_rank_avg_even_vals[True-uint64]": "TODO: Add a reason for failure", "tests/groupby/test_all_methods.py::test_not_c_contiguous_mask[all]": "assert not True", "tests/groupby/test_all_methods.py::test_not_c_contiguous_mask[any]": "assert not True", "tests/groupby/test_all_methods.py::test_not_c_contiguous_mask[bfill]": "assert not True", diff --git a/python/cudf/cudf/tests/groupby/test_rank.py b/python/cudf/cudf/tests/groupby/test_rank.py index b9337d2d68b3..a3e6586dd742 100644 --- a/python/cudf/cudf/tests/groupby/test_rank.py +++ b/python/cudf/cudf/tests/groupby/test_rank.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import numpy as np import pandas as pd @@ -58,3 +58,27 @@ def test_groupby_rank_fails(): ) with pytest.raises(NotImplementedError): gdf.groupby(["a"]).rank(method="min", axis=1) + + +@pytest.mark.parametrize( + "dtype, expected_dtype", + [ + ("int64", "float64"), + ("float32", "float64"), + ("Int64", "Float64"), + ("Float32", "Float64"), + ], +) +def test_rank_dtype_family(dtype, expected_dtype): + # pandas' groupby rank always returns floats within the value column's + # dtype family: numpy -> float64, masked -> Float64 + pdf = pd.DataFrame( + {"key": ["a"] * 4, "val": pd.array([1, 2, 2, 3], dtype=dtype)} + ) + gdf = cudf.from_pandas(pdf) + + expect = pdf.groupby("key").rank() + got = gdf.groupby("key").rank() + + assert str(got["val"].dtype) == expected_dtype + assert_groupby_results_equal(expect, got)