diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index bdcd6c737641..a4d92d0fe6ea 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py @@ -846,6 +846,7 @@ def _build_window_order_index( ob_nulls_last: bool, value_col: plc.Column | None = None, value_desc: bool = False, + reverse: bool = False, stream: Stream, ) -> plc.Column: """Compute a stable row ordering for unary operations in a grouped context.""" @@ -866,7 +867,9 @@ def _build_window_order_index( if order_by_col is not None: cols.append(order_by_col.obj) orders.append( - plc.types.Order.DESCENDING if ob_desc else plc.types.Order.ASCENDING + plc.types.Order.DESCENDING + if ob_desc ^ reverse + else plc.types.Order.ASCENDING ) nulls.append( plc.types.NullOrder.AFTER @@ -876,7 +879,9 @@ def _build_window_order_index( # Use the row id to break ties cols.append(row_id) - orders.append(plc.types.Order.ASCENDING) + orders.append( + plc.types.Order.DESCENDING if reverse else plc.types.Order.ASCENDING + ) nulls.append(plc.types.NullOrder.AFTER) return plc.sorting.stable_sorted_order( @@ -914,10 +919,11 @@ def _grouped_window_scan_setup( ob_desc: bool, ob_nulls_last: bool, grouper: plc.groupby.GroupBy, + reverse: bool = False, stream: Stream, require_sorted_groups: bool = False, ) -> tuple[plc.Column | None, list[Column] | None, plc.groupby.GroupBy]: - if order_by_col is None and not require_sorted_groups: + if order_by_col is None and not (require_sorted_groups or reverse): # keep the original ordering return None, None, grouper order_index = self._build_window_order_index( @@ -926,6 +932,7 @@ def _grouped_window_scan_setup( order_by_col=order_by_col, ob_desc=ob_desc, ob_nulls_last=ob_nulls_last, + reverse=reverse, stream=stream, ) by_cols_for_scan = self._gather_columns(by_cols, order_index, stream=stream) @@ -1262,54 +1269,73 @@ def do_evaluate( # noqa: D102 ) if cum_named := unary_window_ops["cum_sum"]: - # A fill_null_with_strategy fill runs on the scan output, which is - # always in sorted-group order, so it needs a sorted grouper even when - # there is no order_by. - has_fill = any( - isinstance(ne.value, expr.UnaryFunction) - and ne.value.name == "fill_null_with_strategy" - for ne in cum_named - ) - order_index, cum_sum_by_cols_for_scan, local = ( - self._grouped_window_scan_setup( - by_cols, - row_id=row_id, - order_by_col=order_by_col - if self._order_by_expr is not None - else None, - ob_desc=self.options[2] - if self._order_by_expr is not None - else False, - ob_nulls_last=self.options[3] - if self._order_by_expr is not None - else False, - grouper=grouper, - stream=df.stream, - require_sorted_groups=has_fill, + cum_reverse = [] + for ne in cum_named: + v = ne.value + assert isinstance(v, expr.UnaryFunction) + if v.name == "fill_null_with_strategy": + cum_sum_expr = v.children[0] + assert isinstance(cum_sum_expr, expr.UnaryFunction) + cum_reverse.append(bool(cum_sum_expr.options[0])) + else: + cum_reverse.append(bool(v.options[0])) + for is_reverse in (False, True): + subset = [ + ne + for ne, rev in zip(cum_named, cum_reverse, strict=True) + if rev is is_reverse + ] + if not subset: + continue + # A fill_null_with_strategy fill runs on the scan output, which is + # always in sorted-group order, so it needs a sorted grouper even when + # there is no order_by. + has_fill = any( + isinstance(ne.value, expr.UnaryFunction) + and ne.value.name == "fill_null_with_strategy" + for ne in cum_named ) - ) - names, dtypes, tables = self._apply_unary_op( - CumSumOp( - named_exprs=cum_named, - order_index=order_index, - by_cols_for_scan=cum_sum_by_cols_for_scan, - local_grouper=local, - ), - df, - grouper, - ) - broadcasted_cols.extend( - self._reorder_to_input( - row_id, - by_cols, - df.num_rows, - tables, - names, - dtypes, - order_index=order_index, - stream=df.stream, + order_index, cum_sum_by_cols_for_scan, local = ( + self._grouped_window_scan_setup( + by_cols, + row_id=row_id, + order_by_col=order_by_col + if self._order_by_expr is not None + else None, + ob_desc=self.options[2] + if self._order_by_expr is not None + else False, + ob_nulls_last=self.options[3] + if self._order_by_expr is not None + else False, + grouper=grouper, + reverse=is_reverse, + stream=df.stream, + require_sorted_groups=has_fill, + ) + ) + names, dtypes, tables = self._apply_unary_op( + CumSumOp( + named_exprs=subset, + order_index=order_index, + by_cols_for_scan=cum_sum_by_cols_for_scan, + local_grouper=local, + ), + df, + grouper, + ) + broadcasted_cols.extend( + self._reorder_to_input( + row_id, + by_cols, + df.num_rows, + tables, + names, + dtypes, + order_index=order_index, + stream=df.stream, + ) ) - ) if shift_named := unary_window_ops["shift"]: order_index, shift_by_cols_for_scan, local = ( diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/unary.py b/python/cudf_polars/cudf_polars/dsl/expressions/unary.py index 6bd975788d3c..cc1fb2afb80f 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/unary.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/unary.py @@ -134,6 +134,7 @@ class UnaryFunction(Expr): "repeat_by", "replace", "replace_strict", + "reverse", "round", "round_sig_figs", "search_sorted", @@ -207,12 +208,6 @@ def __init__( if self.name not in UnaryFunction._supported_fns: raise NotImplementedError(f"Unary function {name=}") # pragma: no cover - if self.name in UnaryFunction._supported_cum_aggs: - (reverse,) = self.options - if reverse: - raise NotImplementedError( - "reverse=True is not supported for cumulative aggregations" - ) if self.name == "index_of" and plc.traits.is_nested(children[0].dtype.plc_type): raise NotImplementedError("index_of on nested types is not supported") if self.name == "fill_null_with_strategy" and self.options[1] not in {0, None}: @@ -1371,6 +1366,12 @@ def do_evaluate( plc.copying.shift(column.obj, offset, fill_scalar, stream=df.stream), dtype=self.dtype, ) + elif self.name == "reverse": + column = self.children[0].evaluate(df, context=context) + return Column( + plc.copying.reverse(column.obj, stream=df.stream), + dtype=self.dtype, + ) elif self.name == "reinterpret": column = self.children[0].evaluate(df, context=context) return column.astype(self.dtype, stream=df.stream) @@ -1621,6 +1622,8 @@ def do_evaluate( ) elif self.name in UnaryFunction._supported_cum_aggs: column = self.children[0].evaluate(df, context=context) + (reverse,) = self.options + # https://github.com/rapidsai/cudf/issues/23208 for a native reverse scan if self.name == "cum_count": # cum_count is the cumulative count of non-null values. counts = plc.unary.cast( @@ -1628,16 +1631,24 @@ def do_evaluate( self.dtype.plc_type, stream=df.stream, ) - return Column( - plc.reduce.scan( - counts, - plc.aggregation.sum(), - plc.reduce.ScanType.INCLUSIVE, - stream=df.stream, - ), - dtype=self.dtype, + if reverse: + # A reverse cumulative aggregation is a forward one over + # the reversed column, reversed back into place. + counts = plc.copying.reverse(counts, stream=df.stream) + result = plc.reduce.scan( + counts, + plc.aggregation.sum(), + plc.reduce.ScanType.INCLUSIVE, + stream=df.stream, ) + if reverse: + result = plc.copying.reverse(result, stream=df.stream) + return Column(result, dtype=self.dtype) plc_col = column.obj + if reverse: + # A reverse cumulative aggregation is a forward one over the + # reversed column, reversed back into place. + plc_col = plc.copying.reverse(plc_col, stream=df.stream) col_type = column.dtype.plc_type # cum_sum casts # Int8, UInt8, Int16, UInt16 -> Int64 for overflow prevention @@ -1678,12 +1689,12 @@ def do_evaluate( elif self.name == "cum_max": agg = plc.aggregation.max() - return Column( - plc.reduce.scan( - plc_col, agg, plc.reduce.ScanType.INCLUSIVE, stream=df.stream - ), - dtype=self.dtype, + result = plc.reduce.scan( + plc_col, agg, plc.reduce.ScanType.INCLUSIVE, stream=df.stream ) + if reverse: + result = plc.copying.reverse(result, stream=df.stream) + return Column(result, dtype=self.dtype) raise NotImplementedError( f"Unimplemented unary function {self.name=}" ) # pragma: no cover; init trips first diff --git a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py index 8e5914a4cfe5..23082789811e 100644 --- a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py +++ b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py @@ -470,6 +470,7 @@ def pytest_report_header(config: pytest.Config) -> str: "tests/unit/sql/test_window_functions.py::test_window_multiple_named_window": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", "tests/unit/functions/test_concat.py::test_concat_horizontal_lazy_strict_raises_shape_error_27415": "horizontal-concat strict height-mismatch raised inside an ExceptionGroup under the streaming engine", "tests/unit/io/test_io_plugin.py::test_defer_validate_true": "correct SchemaError raised but wrapped in an ExceptionGroup under the streaming engine", + "tests/unit/lazyframe/test_projections.py::test_merge_sorted_projection_pd": "https://github.com/rapidsai/cudf/issues/23055", "tests/unit/operations/test_slice.py::test_hconcat_tail_unequal_heights_strict_raises_27552": "horizontal-concat strict height-mismatch raised inside an ExceptionGroup under the streaming engine", } diff --git a/python/cudf_polars/tests/expressions/test_agg.py b/python/cudf_polars/tests/expressions/test_agg.py index b363619e192f..bd14e9ba0a03 100644 --- a/python/cudf_polars/tests/expressions/test_agg.py +++ b/python/cudf_polars/tests/expressions/test_agg.py @@ -195,12 +195,21 @@ def test_cum_count(engine: pl.GPUEngine, data): @pytest.mark.parametrize("cum_agg", sorted(expr.UnaryFunction._supported_cum_aggs)) -def test_cum_agg_reverse_unsupported(engine: pl.GPUEngine, cum_agg): - df = pl.LazyFrame({"a": [1, 2, 3]}) - expr = getattr(pl.col("a"), cum_agg)(reverse=True) - q = df.select(expr) - - assert_ir_translation_raises(q, engine, NotImplementedError) +@pytest.mark.parametrize( + "data,dtype", + [ + ([1, 2, 3, 4, 5], pl.Int32), + ([1, None, 3, None, 5], pl.Int32), + ([None, None, None], pl.Int32), + ([2, 3, 4], pl.Int8), + ([1.5, 2.0, 0.5, 4.0], pl.Float64), + ([], pl.Int32), + ], +) +def test_cum_agg_reverse(engine: pl.GPUEngine, cum_agg, data, dtype): + df = pl.LazyFrame({"a": pl.Series(data, dtype=dtype)}) + q = df.select(getattr(pl.col("a"), cum_agg)(reverse=True)) + assert_gpu_result_equal(q, engine=engine, check_exact=False) @pytest.mark.parametrize("q", [0.5, pl.lit(0.5)]) diff --git a/python/cudf_polars/tests/expressions/test_reverse.py b/python/cudf_polars/tests/expressions/test_reverse.py new file mode 100644 index 000000000000..462592edf190 --- /dev/null +++ b/python/cudf_polars/tests/expressions/test_reverse.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import pytest + +import polars as pl + +from cudf_polars.testing.asserts import assert_gpu_result_equal + + +@pytest.mark.parametrize( + "series", + [ + pl.Series([1, 2, 3, 4, 5, 6, 7]), + pl.Series([1, 2, None, 4, None, 6, 7]), + pl.Series([1.5, 2.5, float("nan"), -3.0, float("inf"), 0.0, -0.0]), + pl.Series(["a", "bb", None, "dddd", "e", "ff", "ggg"]), + pl.Series([1]), + pl.Series([None, None, None], dtype=pl.Int64), + pl.Series([], dtype=pl.Int64), + ], +) +def test_reverse(engine: pl.GPUEngine, series: pl.Series) -> None: + lf = pl.LazyFrame({"a": series}) + q = lf.select(pl.col("a").reverse()) + assert_gpu_result_equal(q, engine=engine) diff --git a/python/pylibcudf/pylibcudf/copying.pxd b/python/pylibcudf/pylibcudf/copying.pxd index 4143e8469945..350e3e85827a 100644 --- a/python/pylibcudf/pylibcudf/copying.pxd +++ b/python/pylibcudf/pylibcudf/copying.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libcpp cimport bool as cbool @@ -90,6 +90,10 @@ cpdef Column shift( DeviceMemoryResource mr=*, ) +cpdef ColumnOrTable reverse( + ColumnOrTable input, object stream = *, DeviceMemoryResource mr=* +) + cpdef list slice(ColumnOrTable input, list indices, object stream = *) cpdef list split(ColumnOrTable input, list splits, object stream = *) diff --git a/python/pylibcudf/pylibcudf/copying.pyi b/python/pylibcudf/pylibcudf/copying.pyi index bdff6cddad57..dc04f2e7324c 100644 --- a/python/pylibcudf/pylibcudf/copying.pyi +++ b/python/pylibcudf/pylibcudf/copying.pyi @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from enum import IntEnum @@ -72,6 +72,11 @@ def shift( stream: CudaStreamLike | None = None, mr: DeviceMemoryResource | None = None, ) -> Column: ... +def reverse( + input: ColumnOrTable, + stream: CudaStreamLike | None = None, + mr: DeviceMemoryResource | None = None, +) -> ColumnOrTable: ... def slice( input: ColumnOrTable, indices: list[int], diff --git a/python/pylibcudf/pylibcudf/copying.pyx b/python/pylibcudf/pylibcudf/copying.pyx index 9b09536d314d..0fe79953238d 100644 --- a/python/pylibcudf/pylibcudf/copying.pyx +++ b/python/pylibcudf/pylibcudf/copying.pyx @@ -59,6 +59,7 @@ __all__ = [ "empty_like", "gather", "get_element", + "reverse", "scatter", "shift", "slice", @@ -455,6 +456,47 @@ cpdef Column shift( return Column.from_libcudf(move(c_result), _stream, mr) +cpdef ColumnOrTable reverse( + ColumnOrTable input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None +): + """Reverse the rows of a column or table. + + For details, see :cpp:func:`reverse`. + + Parameters + ---------- + input : Union[Column, Table] + The column or table to reverse. + stream : Stream | None + CUDA stream on which to perform the operation. + mr : DeviceMemoryResource | None + Device memory resource used to allocate the returned result's memory. + + Returns + ------- + Union[Column, Table] + The reversed column or table. + """ + cdef unique_ptr[table] c_tbl_result + cdef unique_ptr[column] c_col_result + cdef Stream _stream = _get_stream(stream) + cdef cudaStream_t _cs = _stream.view().value() + cdef column_view c_input_column + cdef table_view c_input_table + + mr = _get_memory_resource(mr) + if ColumnOrTable is Column: + c_input_column = input.view() + with nogil: + c_col_result = cpp_copying.reverse(c_input_column, _cs, mr.get_mr()) + return Column.from_libcudf(move(c_col_result), _stream, mr) + else: + c_input_table = input.view() + with nogil: + c_tbl_result = cpp_copying.reverse(c_input_table, _cs, mr.get_mr()) + return Table.from_libcudf(move(c_tbl_result), _stream, mr) + + cpdef list slice( ColumnOrTable input, list indices: list[int], diff --git a/python/pylibcudf/pylibcudf/libcudf/copying.pxd b/python/pylibcudf/pylibcudf/libcudf/copying.pxd index 36c95fa777c6..f211dbb47b57 100644 --- a/python/pylibcudf/pylibcudf/libcudf/copying.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/copying.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libc.stdint cimport int32_t, int64_t, uint8_t from libcpp cimport bool @@ -43,6 +43,18 @@ cdef extern from "cudf/copying.hpp" namespace "cudf" nogil: device_async_resource_ref mr ) except +libcudf_exception_handler + cdef unique_ptr[table] reverse( + const table_view& source_table, + cudaStream_t stream, + device_async_resource_ref mr + ) except +libcudf_exception_handler + + cdef unique_ptr[column] reverse( + const column_view& source_column, + cudaStream_t stream, + device_async_resource_ref mr + ) except +libcudf_exception_handler + cdef unique_ptr[table] scatter ( const table_view& source_table, const column_view& scatter_map, diff --git a/python/pylibcudf/tests/test_copying.py b/python/pylibcudf/tests/test_copying.py index efe638b97e02..fe55ba82b8d5 100644 --- a/python/pylibcudf/tests/test_copying.py +++ b/python/pylibcudf/tests/test_copying.py @@ -455,6 +455,22 @@ def test_empty_like_table(source_table): assert rcol.type() == icol.type() +def test_reverse_column(target_column): + pa_target_column, plc_target_column = target_column + result = plc.copying.reverse(plc_target_column) + reversed_indices = pa.array(range(len(pa_target_column) - 1, -1, -1)) + expected = pa_target_column.take(reversed_indices) + assert_column_eq(expected, result) + + +def test_reverse_table(source_table): + pa_source_table, plc_source_table = source_table + result = plc.copying.reverse(plc_source_table) + reversed_indices = pa.array(range(pa_source_table.num_rows - 1, -1, -1)) + expected = pa_source_table.take(reversed_indices) + assert_table_eq(expected, result) + + @pytest.mark.parametrize("size", [None, 10]) def test_allocate_like(input_column, size): _, plc_input_column = input_column