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

cpp/CMakeLists.txt

@@ -642,6 +642,7 @@ add_library(
   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,39 @@
 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
+
+/**
+ * @brief Compute the number of nulls in each group.
+ *
+ * @param orderby Column with null mask.
+ * @param offsets Offset array defining the (sorted) groups.
+ * @param stream CUDA stream used for kernel launches
+ * @return device_uvector containing the null count per group.
+ */
+[[nodiscard]] rmm::device_uvector<cudf::size_type> nulls_per_group(
+  column_view const& orderby,
+  rmm::device_uvector<size_type> const& offsets,
+  rmm::cuda_stream_view stream);
+
 /**
  * @copydoc std::unique_ptr<column> rolling_window(
  *            column_view const& input,
@@ -48,5 +81,33 @@ 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
+ *
+ * @param orderby Column use to define window ranges. If @p grouping is empty,
+ * must be sorted. If
+ * @p grouping 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 grouping Optional preprocessed grouping information.
+ * @param order The sort order of the @p orderby column.
+ * @param null_order The sort order of nulls in the @p orderby column.
+ * @param window Descriptor specifying the window type.
+ * @param stream CUDA stream used for device memory operations and kernel
+ * launches.
+ * @param mr Device memory resource used for allocations.
+ * @return Column representing the window offsets as requested, suitable for passing to
+ * `rolling_window`.
+ */
+[[nodiscard]] std::unique_ptr<column> make_range_window(
+  column_view const& orderby,
+  std::optional<rolling::preprocessed_group_info> const& grouping,
+  rolling::direction direction,
+  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,92 @@ 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 Return pointer to the row delta scalar.
+   * @return pointer to scalar, not null.
+   */
+  cudf::scalar const* delta() const noexcept { return &delta_; }
+};
+
+/**
+ * @brief Strongly typed wrapper for bounded open 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 Return pointer to the row delta scalar.
+   * @return pointer to scalar, not null.
+   */
+  cudf::scalar const* delta() const noexcept { return &delta_; }
+};
+
+/**
+ * @brief Strongly typed wrapper for unbounded rolling windows.
+ *
+ * This window runs to the begin/end of the current row's group.
+ */
+struct unbounded {
+  /**
+   * @brief Return a null row delta
+   * @return nullptr
+   */
+  constexpr cudf::scalar const* delta() const noexcept { return nullptr; }
+};
+/**
+ * @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 Return a null row delta
+   * @return nullptr
+   */
+  constexpr cudf::scalar const* delta() const noexcept { return nullptr; }
+};
+
+/**
+ * @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. Apples groupwise if
+ * `group_keys` is non-empty.
+ * @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  Applies a fixed-size rolling window function to the values in a column.
  *
@@ -443,6 +534,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.
  *