diff --git a/python/cudf/cudf/core/groupby/groupby.py b/python/cudf/cudf/core/groupby/groupby.py index 2519fda326a..2ed9bed5b49 100644 --- a/python/cudf/cudf/core/groupby/groupby.py +++ b/python/cudf/cudf/core/groupby/groupby.py @@ -1204,8 +1204,7 @@ def _jit_groupby_apply( {None: chunk_results}, index=group_names ) result.index.names = self.grouping.names - result = result.reset_index() - result[None] = result.pop(0) + return result @_cudf_nvtx_annotate @@ -1364,10 +1363,12 @@ def mult(df): ... lambda group: group['b'].max() - group['b'].min(), ... engine='jit' ... ) - a None - 0 1 1 - 1 2 1 - 2 3 1 + a + 1 1 + 2 1 + 3 1 + dtype: int64 + """ if self.obj.empty: diff --git a/python/cudf/cudf/tests/test_groupby.py b/python/cudf/cudf/tests/test_groupby.py index 7d22cb70803..b01b44da201 100644 --- a/python/cudf/cudf/tests/test_groupby.py +++ b/python/cudf/cudf/tests/test_groupby.py @@ -395,7 +395,7 @@ def groupby_jit_data(): def run_groupby_apply_jit_test(data, func, keys, *args): - expect_groupby_obj = data.to_pandas().groupby(keys, as_index=False) + expect_groupby_obj = data.to_pandas().groupby(keys) got_groupby_obj = data.groupby(keys) # compare cuDF jit to pandas @@ -475,7 +475,19 @@ def func(df): @pytest.mark.parametrize("dtype", ["float64"]) @pytest.mark.parametrize("func", ["idxmax", "idxmin"]) -@pytest.mark.parametrize("special_val", [np.nan, np.inf, -np.inf]) +@pytest.mark.parametrize( + "special_val", + [ + pytest.param( + np.nan, + marks=pytest.mark.xfail( + reason="https://github.com/rapidsai/cudf/issues/13832" + ), + ), + np.inf, + -np.inf, + ], +) def test_groupby_apply_jit_idx_reductions_special_vals( func, groupby_jit_data, dtype, special_val ): @@ -494,19 +506,10 @@ def func(df): groupby_jit_data["val1"] = special_val groupby_jit_data["val1"] = groupby_jit_data["val1"].astype(dtype) - expect = ( - groupby_jit_data.to_pandas() - .groupby("key1", as_index=False) - .apply(func) - ) - - grouped = groupby_jit_data.groupby("key1") - sorted = grouped._grouped()[3].to_pandas() - expect_vals = sorted["key1"].drop_duplicates().index - expect[None] = expect_vals + expect = groupby_jit_data.to_pandas().groupby("key1").apply(func) + got = groupby_jit_data.groupby("key1").apply(func, engine="jit") - got = grouped.apply(func, engine="jit") - assert_eq(expect, got) + assert_eq(expect, got, check_dtype=False) @pytest.mark.parametrize(