diff --git a/python/cudf/cudf/core/groupby/groupby.py b/python/cudf/cudf/core/groupby/groupby.py index 8d4bca3b5d91..b1755ebbabf2 100644 --- a/python/cudf/cudf/core/groupby/groupby.py +++ b/python/cudf/cudf/core/groupby/groupby.py @@ -496,6 +496,25 @@ def _collect_series_key_column_names(obj, by) -> dict[int, Hashable]: return result +class GroupByNthSelector: + """Mirror of :class:`pandas.core.groupby.indexing.GroupByNthSelector`. + + ``GroupBy.nth`` supports both the call form ``gb.nth(n, dropna=...)`` + and the index form ``gb.nth[n]``. + """ + + def __init__(self, groupby_object: GroupBy) -> None: + self.groupby_object = groupby_object + + def __call__( + self, n, dropna: Literal["any", "all", None] = None + ) -> Series | DataFrame: + return self.groupby_object._nth(n, dropna) + + def __getitem__(self, n) -> Series | DataFrame: + return self.groupby_object._nth(n) + + class GroupBy(Serializable, Reducible, Scannable): obj: Series | DataFrame @@ -575,6 +594,10 @@ def __init__( # Must be done before ``nans_to_nulls`` which breaks identity. by_series_col_names = _collect_series_key_column_names(obj, by) + # Row-filter operations (``nth``) must return the original values, + # preserving the NaN-vs-null distinction that ``nans_to_nulls`` + # erases below. + self._obj_original = obj if get_option("mode.pandas_compatible"): obj = obj.nans_to_nulls() self.obj = obj @@ -1202,9 +1225,34 @@ def agg(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs): ): cast_dtype = orig_dtype elif agg not in {list, "collect"}: - create_dtype = get_dtype_of_same_kind( - orig_dtype, create_dtype - ) + if ( + isinstance(orig_dtype, np.dtype) + and orig_dtype.kind == "O" + and is_dtype_obj_string(create_dtype) + ): + # a string-producing aggregation (first/last/min/ + # max/nth) on an object-dtype column stays object, + # matching pandas. Scoped here rather than in + # get_dtype_of_same_kind: other callers (e.g. merge + # key coalescing) re-infer str for object inputs. + create_dtype = orig_dtype + elif ( + isinstance(orig_dtype, pd.DatetimeTZDtype) + and isinstance(create_dtype, np.dtype) + and create_dtype.kind == "M" + ): + # libcudf has no timezone notion: a DatetimeTZColumn + # feeds its stored UTC instants to libcudf and the + # result comes back as a tz-naive timestamp column. + # Reattach the original tz (the values are unchanged + # UTC instants, so this is lossless). + create_dtype = pd.DatetimeTZDtype( + np.datetime_data(create_dtype)[0], orig_dtype.tz + ) + else: + create_dtype = get_dtype_of_same_kind( + orig_dtype, create_dtype + ) result_col = ColumnBase.create(plc_result, create_dtype) if agg == "cumcount": @@ -1762,27 +1810,137 @@ def tail(self, n: int = 5, *, preserve_order: bool = True): n, take_head=False, preserve_order=preserve_order ) - @_performance_tracking - def nth(self, n, dropna: Literal["any", "all", None] = None): + @property + def nth(self): """ - Return the nth row from each group. + Take the nth row from each group if n is an int, otherwise a + subset of rows. + + Like pandas, supports both the call form ``gb.nth(n, dropna=...)`` + and the index form ``gb.nth[n]``. + + Parameters + ---------- + n : int, slice or list of ints and slices + A single nth value for the row, a slice with non-negative + step or a list of nth values and slices. Negative values + count from the end of each group. + dropna : {'any', 'all', None}, default None + Apply the specified dropna operation before counting which + row is the nth row. Only supported in the call form and not + currently implemented in cuDF (raises ``NotImplementedError``; + falls back to pandas under ``cudf.pandas``). + + Returns + ------- + Series or DataFrame + The nth row(s) of each group, keeping the original index and + row order (like a filter operation, the group keys are not + added as an index level). + + Examples + -------- + >>> import cudf + >>> df = cudf.DataFrame({"A": [1, 1, 2, 1, 2], + ... "B": [None, 2, 3, 4, 5]}) + >>> gb = df.groupby("A") + >>> gb.nth(0) + A B + 0 1 + 2 2 3 + >>> gb.nth(-1) + A B + 3 1 4 + 4 2 5 + >>> gb.nth[:2] + A B + 0 1 + 1 1 2 + 2 2 3 + 4 2 5 """ + return GroupByNthSelector(self) + + @_performance_tracking + def _nth(self, n, dropna: Literal["any", "all", None] = None): + """Positional row filter mirroring pandas' GroupBy.nth.""" if dropna is not None: raise NotImplementedError("dropna is not currently supported.") - self.obj["__groupbynth_order__"] = range(0, len(self.obj)) - # We perform another groupby here to have the grouping columns - # be a part of dataframe columns. - result = self.obj.groupby(self.grouping.keys).agg(lambda x: x.nth(n)) - sizes = self.size().reindex(result.index) - result = result[sizes > n] + # Normalize and validate ``n`` like pandas' + # GroupByIndexingMixin._make_mask_from_positional_indexer. + if isinstance(n, (int, np.integer)): + args: list = [int(n)] + elif isinstance(n, slice): + args = [n] + elif isinstance(n, (list, tuple, np.ndarray)): + args = list(n) + else: + raise TypeError( + f"Invalid index {type(n)}. " + "Must be integer, list-like, slice or a tuple of " + "integers and slices" + ) + for arg in args: + if isinstance(arg, slice): + if (arg.step or 1) < 0: + raise ValueError( + f"Invalid step {arg.step}. Must be non-negative" + ) + elif not isinstance(arg, (int, np.integer)): + raise TypeError( + f"Invalid index {type(n)}. " + "Must be integer, list-like, slice or a tuple of " + "integers and slices" + ) - result.index = self.obj.index.take( - result._data["__groupbynth_order__"] + # Per-row position within its group and group size, in group-major + # order (same construction as ``_head_tail``). + _, offsets, _, _ = self._grouped() + group_offsets = np.asarray(offsets, dtype=SIZE_TYPE_DTYPE) + size_per_group = np.diff(group_offsets) + sizes = np.repeat(size_per_group, size_per_group) + pos = np.arange(len(sizes), dtype=SIZE_TYPE_DTYPE) - np.repeat( + group_offsets[:-1], size_per_group ) - del result._data["__groupbynth_order__"] - del self.obj._data["__groupbynth_order__"] - return result + + mask = np.zeros(len(sizes), dtype=bool) + for arg in args: + if isinstance(arg, slice): + step = arg.step or 1 + if arg.start is None: + start = np.zeros_like(sizes) + elif arg.start >= 0: + start = np.full_like(sizes, arg.start) + else: + # ``slice.indices`` clamps a negative start at 0 and + # the step alignment begins at the clamped value + start = np.maximum(sizes + arg.start, 0) + submask = pos >= start + if step > 1: + submask &= (pos - start) % step == 0 + if arg.stop is not None: + if arg.stop >= 0: + submask &= pos < arg.stop + else: + submask &= pos < np.maximum(sizes + arg.stop, 0) + mask |= submask + elif arg >= 0: + mask |= pos == arg + else: + mask |= pos == sizes + arg + + # Map the selected group-major rows back to positions in the + # original object and gather from the *pre-nans_to_nulls* object: + # pandas' nth is a row filter, so values, dtypes, index and row + # order are those of the original rows. + to_take = as_column(np.nonzero(mask)[0].astype(SIZE_TYPE_DTYPE)) + _, _, (ordering,) = self._groups([self._range_column_from_obj]) + original_positions = ordering.take(to_take) + original_positions = original_positions.take( + original_positions.argsort() + ) + return self._obj_original.take(original_positions) @_performance_tracking def ngroup(self, ascending=True): diff --git a/python/cudf/cudf/pandas/_wrappers/pandas.py b/python/cudf/cudf/pandas/_wrappers/pandas.py index dc523cd1f4e1..341c65658e33 100644 --- a/python/cudf/cudf/pandas/_wrappers/pandas.py +++ b/python/cudf/cudf/pandas/_wrappers/pandas.py @@ -1167,6 +1167,17 @@ def Index__setattr__(self, name, value): ) +# ``GroupBy.nth`` is a property returning a selector object on both +# sides; registering the selector pair as an intermediate proxy gives +# ``gb.nth`` a proxied result whose ``__call__``/``__getitem__`` get the +# usual call-time fast/slow dispatch (with the slow side re-derived from +# the recorded ``getattr`` provenance on fallback). +GroupByNthSelector = make_intermediate_proxy_type( + "GroupByNthSelector", + cudf.core.groupby.groupby.GroupByNthSelector, + pd.core.groupby.indexing.GroupByNthSelector, +) + SeriesGroupBy = make_intermediate_proxy_type( "SeriesGroupBy", cudf.core.groupby.groupby.SeriesGroupBy, diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index ec54cec49757..ffdc3dd005d2 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -857,7 +857,6 @@ def pytest_unconfigure(config): "tests/extension/test_datetime.py::TestDatetimeArray::test_astype_str": "AssertionError: Series are different", "tests/extension/test_datetime.py::TestDatetimeArray::test_astype_string[string[pyarrow]]": "AssertionError: Series are different", "tests/extension/test_datetime.py::TestDatetimeArray::test_astype_string[string[python]]": "AssertionError: Series are different", - "tests/extension/test_datetime.py::TestDatetimeArray::test_groupby_agg_extension": "AssertionError: Attributes of DataFrame.iloc[:, 0] (column name='B') are different", "tests/extension/test_datetime.py::TestDatetimeArray::test_is_extension_array_dtype": "AssertionError", "tests/extension/test_datetime.py::TestDatetimeArray::test_reduce_frame[mean-False]": "AssertionError: Attributes of ExtensionArray are different", "tests/extension/test_datetime.py::TestDatetimeArray::test_reduce_frame[mean-True]": "AssertionError: Attributes of ExtensionArray are different", @@ -1822,7 +1821,6 @@ def pytest_unconfigure(config): "tests/groupby/aggregate/test_cython.py::test_cythonized_aggers[prod]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_cython.py::test_cythonized_aggers[sum]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_other.py::test_agg_over_numpy_arrays": "TODO: Add a reason for failure", - "tests/groupby/aggregate/test_other.py::test_agg_timezone_round_trip": "AssertionError: assert Timestamp('2016-01-01 20:00:00') == Timestamp('2016-01-01 12:00:00-0800', tz='US/Pacific')", "tests/groupby/methods/test_describe.py::test_describe_duplicate_columns": "AssertionError: DataFrame are different", "tests/groupby/methods/test_describe.py::test_describe_with_duplicate_output_column_names[False-keys0]": "AssertionError: DataFrame are different", "tests/groupby/methods/test_describe.py::test_describe_with_duplicate_output_column_names[False-keys1]": "AssertionError: DataFrame are different", @@ -1834,48 +1832,6 @@ def pytest_unconfigure(config): "tests/groupby/methods/test_groupby_shift_diff.py::test_group_diff_real_frame[int8]": "TODO: Add a reason for failure", "tests/groupby/methods/test_groupby_shift_diff.py::test_group_diff_real_series[int16]": "TODO: Add a reason for failure", "tests/groupby/methods/test_groupby_shift_diff.py::test_group_diff_real_series[int8]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_first_last_nth": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_first_last_nth_dtypes": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_first_last_nth_nan_dtype": "AssertionError: Attributes of Series are different", - "tests/groupby/methods/test_nth.py::test_first_last_tz[data0-expected_first0-expected_last0]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='time') are different", - "tests/groupby/methods/test_nth.py::test_first_last_with_None[first]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='value') are different", - "tests/groupby/methods/test_nth.py::test_first_last_with_None[last]": "AssertionError: Attributes of DataFrame.iloc[:, 1] (column name='value') are different", - "tests/groupby/methods/test_nth.py::test_groupby_last_first_nth_with_none[NoneType-first]": "AssertionError: Attributes of Series are different", - "tests/groupby/methods/test_nth.py::test_groupby_last_first_nth_with_none[NoneType-last]": "AssertionError: Attributes of Series are different", - "tests/groupby/methods/test_nth.py::test_groupby_last_first_nth_with_none[NoneType-nth]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_head_tail_dropna_false": "AssertionError: DataFrame Expected type , found ins...", - "tests/groupby/methods/test_nth.py::test_negative_step": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_np_ints": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth2": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth3": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth4": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth5": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[None-b]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[None-selection1]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[None-selection2]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[all-b]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[all-selection1]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[all-selection2]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[any-b]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[any-selection1]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_after_selection[any-selection2]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_column_order": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_indexed": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_multi_grouper": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_multi_index_as_expected": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_nan_in_grouper[all]": "NotImplementedError: dropna is not currently supported.", - "tests/groupby/methods/test_nth.py::test_nth_nan_in_grouper[any]": "NotImplementedError: dropna is not currently supported.", - "tests/groupby/methods/test_nth.py::test_nth_nan_in_grouper_series[None]": "KeyError: 'Label scalar is out of bounds'", - "tests/groupby/methods/test_nth.py::test_nth_nan_in_grouper_series[all]": "NotImplementedError: dropna is not currently supported.", - "tests/groupby/methods/test_nth.py::test_nth_nan_in_grouper_series[any]": "NotImplementedError: dropna is not currently supported.", - "tests/groupby/methods/test_nth.py::test_nth_with_na_object[NoneType--1]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_with_na_object[float0--1]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_nth_with_na_object[float1--1]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_slice[arg0-expected_rows0]": "TODO: Add a reason for failure", - "tests/groupby/methods/test_nth.py::test_slice[arg1-expected_rows1]": "TODO: Add a reason for failure", - "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", @@ -2046,13 +2002,11 @@ def pytest_unconfigure(config): "tests/groupby/test_groupby_subclass.py::test_groupby_resample_preserves_subclass[DataFrame]": "TODO: Add a reason for failure", "tests/groupby/test_grouping.py::TestGetGroup::test_get_group_grouped_by_tuple": "TODO: Add a reason for failure", "tests/groupby/test_grouping.py::TestGetGroup::test_get_group_grouped_by_tuple_with_lambda": "TODO: Add a reason for failure", - "tests/groupby/test_grouping.py::TestGetGroup::test_groupby_with_single_column": "TODO: Add a reason for failure", "tests/groupby/test_grouping.py::TestGrouping::test_groupby_apply_empty_with_group_keys_false": "AssertionError: DataFrame are different", "tests/groupby/test_grouping.py::TestGrouping::test_groupby_multiindex_partial_indexing_equivalence": "TODO: Add a reason for failure", "tests/groupby/test_grouping.py::TestGrouping::test_groupby_tuple_keys_handle_multiindex": "TypeError: unhashable type: 'list'", "tests/groupby/test_grouping.py::TestGrouping::test_multiindex_columns_empty_level": "TODO: Add a reason for failure", "tests/groupby/test_grouping.py::TestSelection::test_indices_grouped_by_tuple_with_lambda": "TODO: Add a reason for failure", - "tests/groupby/test_indexing.py::test_multiindex": "TODO: Add a reason for failure", "tests/groupby/test_missing.py::test_groupby_column_index_name_lost_fill_funcs[bfill]": "AssertionError: Index are different", "tests/groupby/test_missing.py::test_groupby_column_index_name_lost_fill_funcs[ffill]": "AssertionError: Index are different", "tests/groupby/test_missing.py::test_indices_with_missing": "TODO: Add a reason for failure", @@ -3374,7 +3328,6 @@ def pytest_unconfigure(config): "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_table_not_series": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_table_with_iterator_values": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_table_with_mixed_nested_tuples": "TODO: Add a reason for failure", - "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_tz_in_values": "AssertionError: Attributes of DataFrame.iloc[:, 0] (column name='2016-08-12 00:00:00-07:00') are different", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_with_categorical[False-False]": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_with_categorical[False-None]": "TODO: Add a reason for failure", "tests/reshape/test_pivot.py::TestPivotTable::test_pivot_with_categorical[False-True]": "TODO: Add a reason for failure", diff --git a/python/cudf/cudf/tests/groupby/test_nth.py b/python/cudf/cudf/tests/groupby/test_nth.py index 2c98a7ac7750..82ddedea95f1 100644 --- a/python/cudf/cudf/tests/groupby/test_nth.py +++ b/python/cudf/cudf/tests/groupby/test_nth.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import numpy as np import pandas as pd @@ -58,3 +58,49 @@ def test_groupby_consecutive_operations(): expected = pg.cumsum() assert_groupby_results_equal(actual, expected, check_dtype=False) + + +@pytest.mark.parametrize( + "arg", + [ + 0, + -1, + [0, 1], + slice(None, 2), + slice(1, None), + slice(None, None, 2), + slice(-2, None), + ], +) +def test_nth_selector_indexing(arg): + # GroupBy.nth mirrors pandas' GroupByNthSelector: it supports both the + # call form gb.nth(n) and the index form gb.nth[n], acting as a + # positional row filter that keeps the original index and row order. + pdf = pd.DataFrame( + { + "a": [1, 1, 1, 2, 2, 3], + "b": [10, 20, 30, 40, 50, 60], + }, + index=[5, 4, 3, 2, 1, 0], + ) + gdf = cudf.from_pandas(pdf) + + assert_groupby_results_equal( + pdf.groupby("a").nth[arg], + gdf.groupby("a").nth[arg], + ) + if not isinstance(arg, slice): + assert_groupby_results_equal( + pdf.groupby("a").nth(arg), + gdf.groupby("a").nth(arg), + ) + + +def test_nth_invalid_args(): + gb = cudf.DataFrame({"a": [1, 1, 2], "b": [1, 2, 3]}).groupby("a") + with pytest.raises(TypeError, match="Invalid index"): + gb.nth(3.14) + with pytest.raises(ValueError, match="Invalid step"): + gb.nth(slice(None, None, -1)) + with pytest.raises(NotImplementedError): + gb.nth(0, dropna="any")