Fix Groupby.get_group (#14728)

Closes #14727 Authors: - Richard (Rick) Zamora (https://github.com/rjzamora) - Charles Blackmon-Luca (https://github.com/charlesbluca) Approvers: - Vyas Ramasubramani (https://github.com/vyasr) - Charles Blackmon-Luca (https://github.com/charlesbluca) URL: #14728
rapidsai · Jan 12, 2024 · 7ca988f · 7ca988f
1 parent 5c78b7e
commit 7ca988f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/python/cudf/cudf/core/groupby/groupby.py b/python/cudf/cudf/core/groupby/groupby.py
@@ -377,7 +377,7 @@ def get_group(self, name, obj=None):
         if obj is None:
             obj = self.obj
 
-        return obj.loc[self.groups[name]]
+        return obj.loc[self.groups[name].drop_duplicates()]
 
     @_cudf_nvtx_annotate
     def size(self):

diff --git a/python/cudf/cudf/tests/groupby/test_indexing.py b/python/cudf/cudf/tests/groupby/test_indexing.py
@@ -1 +1,12 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
+import cudf
+from cudf.testing._utils import assert_eq
+
+
+def test_rank_return_type_compatible_mode():
+    # in compatible mode, rank() always returns floats
+    df = cudf.DataFrame({"a": range(10), "b": [0] * 10}, index=[0] * 10)
+    pdf = df.to_pandas()
+    expect = pdf.groupby("b").get_group(0)
+    result = df.groupby("b").get_group(0)
+    assert_eq(expect, result)