Extend and simplify API for calculation of range-based rolling window offsets

cpp/CMakeLists.txt

@@ -643,10 +643,13 @@ add_library(
   src/reshape/interleave_columns.cu
   src/reshape/tile.cu
   src/rolling/detail/optimized_unbounded_window.cpp
+  src/rolling/detail/range_following.cu
+  src/rolling/detail/range_preceding.cu
   src/rolling/detail/rolling_collect_list.cu
   src/rolling/detail/rolling_fixed_window.cu
   src/rolling/detail/rolling_variable_window.cu
   src/rolling/grouped_rolling.cu
+  src/rolling/range_rolling.cu
   src/rolling/range_window_bounds.cpp
   src/rolling/rolling.cu
   src/round/round.cu

cpp/include/cudf/detail/rolling.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -29,6 +29,26 @@
 namespace CUDF_EXPORT cudf {
 namespace detail {
 
+namespace rolling {
+/**
+ * @brief Direction tag for a range-based rolling window.
+ */
+enum class direction : bool {
+  PRECEDING,  ///< A preceding window.
+  FOLLOWING,  ///< A following window.
+};
+/**
+ * @brief Wrapper for preprocessed group information from sorted group keys.
+ */
+struct preprocessed_group_info {
+  rmm::device_uvector<size_type> const& labels;   ///< Mapping from row index to group label
+  rmm::device_uvector<size_type> const& offsets;  ///< Mapping from group label to row offsets
+  rmm::device_uvector<size_type> const&
+    nulls_per_group;  ///< Mapping from group label to null count in the group
+};
+
+}  // namespace rolling
+
 /**
  * @copydoc std::unique_ptr<column> rolling_window(
  *            column_view const& input,
@@ -48,5 +68,67 @@ std::unique_ptr<column> rolling_window(column_view const& input,
                                        rmm::cuda_stream_view stream,
                                        rmm::device_async_resource_ref mr);
 
+/**
+ * @brief Make a column representing the window offsets for a range-based window
+ *
+ * @tparam Direction Is this a preceding window or a following one.
+ *
+ * @param group_keys Table defining grouping of the windows. May be empty. If
+ * non-empty, group keys must be sorted.
+ * @param orderby Column use to define window ranges. If @p group_keys is empty,
+ * must be sorted. If
+ * @p group_keys is non-empty, must be sorted within each group. As well as
+ * being sorted, must be sorted consistently with the @p order and @p null_order
+ * parameters.
+ * @param order The sort order of the @p orderby column.
+ * @param null_order The sort order of nulls in the @p orderby column.
+ * @param row_delta Pointer to scalar providing the delta for the window range.
+ * May be null, but only if the @p window_type is @p CURRENT_ROW or @p
+ * UNBOUNDED. Note that @p row_delta is always added to the current row value.
+ * @param window_type The type of window we are computing bounds for.
+ * @param stream CUDA stream used for device memory operations and kernel
+ * launches.
+ * @param mr Device memory resource used for allocations.
+ */
+template <rolling::direction Direction>
+[[nodiscard]] std::unique_ptr<column> make_range_window(
+  column_view const& orderby,
+  std::optional<rolling::preprocessed_group_info> const& grouping,
+  order order,
+  null_order null_order,
+  range_window_type window,
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref mr);
+
+/**
+ * @copydoc `make_range_window`
+ *
+ * This behaves as `make_range_window`, but template-specialised to `PRECEDING` windows. This way we
+ * can compile it in a separate translation unit from `make_following_range_window`.
+ */
+[[nodiscard]] std::unique_ptr<column> make_preceding_range_window(
+  column_view const& orderby,
+  std::optional<rolling::preprocessed_group_info> const& grouping,
+  order order,
+  null_order null_order,
+  range_window_type window,
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref mr);
+
+/**
+ * @copydoc `make_range_window`
+ *
+ * This behaves as `make_range_window`, but template-specialised to `FOLLOWING` windows. This way we
+ * can compile it in a separate translation unit from `make_preceding_range_window`.
+ */
+[[nodiscard]] std::unique_ptr<column> make_following_range_window(
+  column_view const& orderby,
+  std::optional<rolling::preprocessed_group_info> const& grouping,
+  order order,
+  null_order null_order,
+  range_window_type window,
+  rmm::cuda_stream_view stream,
+  rmm::device_async_resource_ref mr);
+
 }  // namespace detail
 }  // namespace CUDF_EXPORT cudf

cpp/include/cudf/rolling.hpp

@@ -19,10 +19,15 @@
 #include <cudf/aggregation.hpp>
 #include <cudf/rolling/range_window_bounds.hpp>
 #include <cudf/types.hpp>
+#include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/export.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 
+#include <rmm/resource_ref.hpp>
+
 #include <memory>
+#include <optional>
+#include <variant>
 
 namespace CUDF_EXPORT cudf {
 /**
@@ -31,6 +36,98 @@ namespace CUDF_EXPORT cudf {
  * @file
  */
 
+/**
+ * @brief Strongly typed wrapper for bounded closed rolling windows.
+ *
+ * The endpoints of this window are included.
+ */
+struct bounded_closed {
+  cudf::scalar const& delta;  ///< Delta from the current row in the window. Must be valid,
+                              ///< behaviour is undefined if not.
+};
+
+/**
+ * @brief Strongly typed wrapper for bounded closed rolling windows.
+ *
+ * The endpoints of this window are excluded.
+ */
+struct bounded_open {
+  cudf::scalar const& delta;  ///< Delta from the current row in the window. Must be valid,
+                              ///< behaviour is undefined if not.
+};
+
+/**
+ * @brief Strongly typed wrapper for unbounded rolling windows.
+ *
+ * This window runs to the begin/end of the current row's group.
+ */
+struct unbounded {};
+/**
+ * @brief Strongly typed wrapper for current_row rolling windows.
+ *
+ * This window contains all rows that are equal to the current row.
+ */
+struct current_row {};
+
+/**
+ * @brief The type of the range-based rolling window endpoint.
+ */
+using range_window_type = std::variant<unbounded, current_row, bounded_closed, bounded_open>;
+
+/**
+ * @brief Constructs preceding and following columns given window range specifications.
+ *
+ * @param group_keys Possibly empty table of sorted keys defining groups.
+ * @param orderby Column defining window ranges. Must be sorted. If `group_keys` is non-empty, must
+ * be sorted groupwise.
+ * @param order Sort order of the `orderby` column.
+ * @param null_order Null sort order in the sorted `orderby` column.
+ * @param preceding Type of the preceding window.
+ * @param following Type of the following window.
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @param mr Device memory resource used to allocate the returned column's device memory
+ * @return pair of preceding and following columns that define the window bounds for each row,
+ * suitable for passing to `rolling_window`.
+ */
+std::pair<std::unique_ptr<column>, std::unique_ptr<column>> make_range_windows(
+  table_view const& group_keys,
+  column_view const& orderby,
+  order order,
+  null_order null_order,
+  range_window_type preceding,
+  range_window_type following,
+  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+
+/**
+ * @brief Constructs preceding and following columns given window range specifications.
+ *
+ * @param group_keys Possibly empty table of sorted keys defining groups.
+ * @param orderby Column defining window ranges. Must be sorted. If `group_keys` is non-empty, must
+ * be sorted groupwise.
+ * @param order Sort order of the `orderby` column.
+ * @param preceding Type of the preceding window.
+ * @param following Type of the following window.
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @param mr Device memory resource used to allocate the returned column's device memory
+ *
+ * @note Using `make_range_window_bounds(table_view const&, column_view const&, order, null_order,
+ * window_type, window_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)` is preferred,
+ * since this function requires the launch of an additional kernel to deduce the null order of the
+ * orderby column.
+ *
+ * @return pair of preceding and following columns that define the window bounds for each row,
+ * suitable for passing to `rolling_window`.
+ */
+std::pair<std::unique_ptr<column>, std::unique_ptr<column>> make_range_windows(
+  table_view const& group_keys,
+  column_view const& orderby,
+  order order,
+  range_window_type preceding,
+  range_window_type following,
+  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
+
 /**
  * @brief  Applies a fixed-size rolling window function to the values in a column.
  *
@@ -577,6 +674,33 @@ std::unique_ptr<column> grouped_range_rolling_window(
   rmm::cuda_stream_view stream      = cudf::get_default_stream(),
   rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
+/**
+ * @brief Apply a grouping-aware range-based rolling window function to a sequence of columns.
+ *
+ * @param group_keys Possibly empty table of sorted keys defining groups.
+ * @param orderby Column defining window ranges. Must be sorted. If `group_keys` is non-empty, must
+ * be sorted groupwise.
+ * @param order Sort order of the `orderby` column.
+ * @param null_order Null sort order in the sorted `orderby` column.
+ * @param preceding Type of the preceding window.
+ * @param following Type of the following window.
+ * @param min_periods Minimum number of observations in the window required to have a value.
+ * @param requests Vector of pairs of columns and aggregation requests.
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @param mr Device memory resource used to allocate the returned column's device memory
+ * @return A table of results, one column per input request.
+ */
+std::unique_ptr<table> grouped_range_rolling_window(
+  table_view const& group_keys,
+  column_view const& orderby,
+  order order,
+  null_order null_order,
+  range_window_type preceding,
+  range_window_type following,
+  size_type min_periods,
+  std::vector<std::pair<column_view const&, rolling_aggregation const&>> requests,
+  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
+  rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 /**
  * @brief  Applies a variable-size rolling window function to the values in a column.
  *