diff --git a/python/cudf/cudf/_lib/pylibcudf/libcudf/CMakeLists.txt b/python/cudf/cudf/_lib/pylibcudf/libcudf/CMakeLists.txt index b04e94f1546..b96d6590292 100644 --- a/python/cudf/cudf/_lib/pylibcudf/libcudf/CMakeLists.txt +++ b/python/cudf/cudf/_lib/pylibcudf/libcudf/CMakeLists.txt @@ -13,7 +13,7 @@ # ============================================================================= set(cython_sources aggregation.pyx binaryop.pyx copying.pyx expressions.pyx reduce.pyx replace.pyx - round.pyx stream_compaction.pyx types.pyx unary.pyx + rolling.pyx round.pyx stream_compaction.pyx types.pyx unary.pyx ) set(linked_libraries cudf::cudf) diff --git a/python/cudf/cudf/_lib/pylibcudf/libcudf/rolling.pxd b/python/cudf/cudf/_lib/pylibcudf/libcudf/rolling.pxd index d7844f99a73..435f6ffae4e 100644 --- a/python/cudf/cudf/_lib/pylibcudf/libcudf/rolling.pxd +++ b/python/cudf/cudf/_lib/pylibcudf/libcudf/rolling.pxd @@ -1,16 +1,26 @@ # Copyright (c) 2020-2024, NVIDIA CORPORATION. +from libc.stdint cimport int32_t +from libcpp cimport bool from libcpp.memory cimport unique_ptr +from libcpp.utility cimport pair from cudf._lib.types import cudf_to_np_types, np_to_cudf_types from cudf._lib.pylibcudf.libcudf.aggregation cimport rolling_aggregation from cudf._lib.pylibcudf.libcudf.column.column cimport column from cudf._lib.pylibcudf.libcudf.column.column_view cimport column_view +from cudf._lib.pylibcudf.libcudf.scalar.scalar cimport scalar from cudf._lib.pylibcudf.libcudf.types cimport size_type cdef extern from "cudf/rolling.hpp" namespace "cudf" nogil: + cpdef enum class window_type(int32_t): + LEFT_CLOSED + RIGHT_CLOSED + CLOSED + OPEN + cdef unique_ptr[column] rolling_window( column_view source, column_view preceding_window, @@ -24,3 +34,10 @@ cdef extern from "cudf/rolling.hpp" namespace "cudf" nogil: size_type following_window, size_type min_periods, rolling_aggregation& agg) except + + + cdef pair[unique_ptr[column], unique_ptr[column]] windows_from_offset( + column_view input, + scalar length, + scalar offset, + window_type, + bool) except + diff --git a/python/cudf/cudf/_lib/pylibcudf/libcudf/rolling.pyx b/python/cudf/cudf/_lib/pylibcudf/libcudf/rolling.pyx new file mode 100644 index 00000000000..e69de29bb2d diff --git a/python/cudf/cudf/_lib/pylibcudf/rolling.pxd b/python/cudf/cudf/_lib/pylibcudf/rolling.pxd index cdadee68d43..c9c026371e4 100644 --- a/python/cudf/cudf/_lib/pylibcudf/rolling.pxd +++ b/python/cudf/cudf/_lib/pylibcudf/rolling.pxd @@ -1,19 +1,31 @@ # Copyright (c) 2024, NVIDIA CORPORATION. +from libc.stdint cimport int32_t +from libcpp cimport bool +from cudf._lib.pylibcudf.libcudf.rolling cimport window_type from cudf._lib.pylibcudf.libcudf.types cimport size_type from .aggregation cimport Aggregation from .column cimport Column +from .scalar cimport Scalar -ctypedef fused WindowType: +ctypedef fused WindowArg: Column size_type cpdef Column rolling_window( Column source, - WindowType preceding_window, - WindowType following_window, + WindowArg preceding_window, + WindowArg following_window, size_type min_periods, Aggregation agg, ) + +cpdef tuple[Column, Column] windows_from_offset( + Column input, + Scalar length, + Scalar offset, + window_type window_type, + bool only_preceding, +) diff --git a/python/cudf/cudf/_lib/pylibcudf/rolling.pyx b/python/cudf/cudf/_lib/pylibcudf/rolling.pyx index 7aa7828a5dd..0982304bfff 100644 --- a/python/cudf/cudf/_lib/pylibcudf/rolling.pyx +++ b/python/cudf/cudf/_lib/pylibcudf/rolling.pyx @@ -1,22 +1,27 @@ # Copyright (c) 2024, NVIDIA CORPORATION. from cython.operator cimport dereference +from libcpp cimport bool from libcpp.memory cimport unique_ptr -from libcpp.utility cimport move +from libcpp.utility cimport move, pair from cudf._lib.pylibcudf.libcudf cimport rolling as cpp_rolling from cudf._lib.pylibcudf.libcudf.aggregation cimport rolling_aggregation from cudf._lib.pylibcudf.libcudf.column.column cimport column +from cudf._lib.pylibcudf.libcudf.rolling cimport window_type from cudf._lib.pylibcudf.libcudf.types cimport size_type from .aggregation cimport Aggregation from .column cimport Column +from .scalar cimport Scalar + +from cudf._lib.pylibcudf.libcudf.rolling import window_type as WindowType # no-cython-lint, isort: skip cpdef Column rolling_window( Column source, - WindowType preceding_window, - WindowType following_window, + WindowArg preceding_window, + WindowArg following_window, size_type min_periods, Aggregation agg, ): @@ -48,7 +53,7 @@ cpdef Column rolling_window( # TODO: Consider making all the conversion functions nogil functions that # reclaim the GIL internally for just the necessary scope like column.view() cdef const rolling_aggregation *c_agg = agg.view_underlying_as_rolling() - if WindowType is Column: + if WindowArg is Column: with nogil: result = move( cpp_rolling.rolling_window( @@ -71,3 +76,28 @@ cpdef Column rolling_window( ) ) return Column.from_libcudf(move(result)) + + +cpdef tuple[Column, Column] windows_from_offset( + Column input, + Scalar length, + Scalar offset, + window_type typ, + bool only_preceding, +): + cdef pair[unique_ptr[column], unique_ptr[column]] result + with nogil: + result = cpp_rolling.windows_from_offset( + input.view(), + dereference(length.c_obj), + dereference(offset.c_obj), + typ, + only_preceding, + ) + if only_preceding: + return (Column.from_libcudf(move(result.first)), None) + else: + return ( + Column.from_libcudf(move(result.first)), + Column.from_libcudf(move(result.second)), + )