rolling.hpp

转到此文件的文档。

/*
 * Copyright (c) 2019-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.
 * You may obtain a copy of the License at
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#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 {
/**
 * @brief Strongly typed wrapper for bounded closed rolling windows.
 *
 */
struct bounded_closed {
 cudf::scalar const& delta_;

/**
 * @brief Strongly typed wrapper for bounded open rolling windows.
 *
 */
struct bounded_open {
 cudf::scalar const& delta_;

/**
 * @brief Strongly typed wrapper for unbounded rolling windows.
 *
 */
struct unbounded {
 constexpr cudf::scalar const* delta() const noexcept { return nullptr; }
};
/**
 * @brief Strongly typed wrapper for current_row rolling windows.
 *
 */
struct current_row {
 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 Applies a fixed-size rolling window function to the values in a 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,
 * providing default outputs for nulls.
 *
 */
std::unique_ptr<column> rolling_window(
 column_view const& input,
 size_type preceding_window,
 size_type following_window,
 size_type min_periods,
 rolling_aggregation const& agg,
 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,
 * providing default outputs for nulls.
 *
 */
std::unique_ptr<column> rolling_window(
 column_view const& input,
 column_view const& default_outputs,
 size_type preceding_window,
 size_type following_window,
 size_type min_periods,
 rolling_aggregation const& agg,
 rmm::cuda_stream_view stream = cudf::get_default_stream(),
 rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());

/**
 * @brief Abstraction for window boundary sizes.
 *
 */
struct window_bounds {
 public
 /**
 * @brief Construct bounded window boundary.
 *
 * @param value The boundary size, in number of rows.
 * @return A bounded window boundary.
 */
 static window_bounds get(size_type value) { return window_bounds(false, value); }

 /**
 * @brief Construct unbounded window boundary.
 *
 * @return An unbounded window boundary.
 */
 static window_bounds unbounded()
 {
 return window_bounds(true, std::numeric_limits<cudf::size_type>::max());
 }

 /**
 * @brief Checks if the window boundary is unbounded.
 *
 * @return `true` if boundary is unbounded, `false` otherwise.
 */
 [[nodiscard]] bool is_unbounded() const { return _is_unbounded; }

 /**
 * @brief Gets the row-boundary for this window_bounds.
 *
 * @note Undefined behavior if boundary is unbounded.
 *
 * @return Row boundary size.
 */
 [[nodiscard]] size_type value() const { return _value; }

 private
 explicit window_bounds(bool is_unbounded_, size_type value_ = 0)
 : _is_unbounded{is_unbounded_}, _value{value_}
 {
 }

 bool const _is_unbounded;
 size_type const _value;
};

/**
 * @brief Applies a grouping-aware, fixed-size rolling window function to the values in a column.
 *
 */
std::unique_ptr<column> grouped_rolling_window(
 table_view const& group_keys,
 column_view const& input,
 size_type preceding_window,
 size_type following_window,
 size_type min_periods,