host_udf.hpp

前往此文件的文档。

/*
 * Copyright (c) 2024-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/column/column_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/span.hpp>

#include <rmm/cuda_stream_view.hpp>
#include <rmm/resource_ref.hpp>

#include <functional>
#include <optional>

namespace CUDF_EXPORT cudf {
class host_udf_base {
 // Declare constructor private to prevent the users from deriving from this class.
 private
 host_udf_base() = default;

 // Only allow deriving from the structs below.
 friend struct reduce_host_udf;
 friend struct segmented_reduce_host_udf;
 friend struct groupby_host_udf;

 public
 virtual ~host_udf_base() = default;

 [[nodiscard]] virtual std::size_t do_hash() const
 {
 return std::hash<int>{}(static_cast<int>(aggregation::Kind::HOST_UDF));
 }

 [[nodiscard]] virtual bool is_equal(host_udf_base const& other) const = 0;

 [[nodiscard]] virtual std::unique_ptr<host_udf_base> clone() const = 0;
};

struct reduce_host_udf : host_udf_base {
 [[nodiscard]] virtual std::unique_ptr<scalar> operator()(
 column_view const& input,
 data_type output_dtype,
 std::optional<std::reference_wrapper<scalar const>> init,
 rmm::cuda_stream_view stream,
 rmm::device_async_resource_ref mr) const = 0;
};

struct segmented_reduce_host_udf : host_udf_base {
 [[nodiscard]] virtual std::unique_ptr<column> operator()(
 column_view const& input,
 device_span<size_type const> offsets,
 data_type output_dtype,
 null_policy null_handling,
 std::optional<std::reference_wrapper<scalar const>> init,
 rmm::cuda_stream_view stream,
 rmm::device_async_resource_ref mr) const = 0;
};

// Forward declaration.
namespace groupby ::detail {
struct aggregate_result_functor;
}

struct groupby_host_udf : host_udf_base {
 [[nodiscard]] virtual std::unique_ptr<column> get_empty_output(
 rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const = 0;

 [[nodiscard]] virtual std::unique_ptr<column> operator()(
 rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const = 0;

 private
 // Allow the struct `aggregate_result_functor` to set its private callback variables.
 friend struct groupby::detail::aggregate_result_functor;

 std::function<column_view(void)> callback_input_values;

 std::function<column_view(void)> callback_grouped_values;

 std::function<column_view(void)> callback_sorted_grouped_values;

 std::function<size_type(void)> callback_num_groups;

 std::function<device_span<size_type const>(void)> callback_group_offsets;

 std::function<device_span<size_type const>(void)> callback_group_labels;

 std::function<column_view(std::unique_ptr<aggregation>)> callback_compute_aggregation;

 protected
 [[nodiscard]] column_view get_input_values() const
 {
 CUDF_EXPECTS(callback_input_values, "Uninitialized callback_input_values.");
 return callback_input_values();
 }

 [[nodiscard]] column_view get_grouped_values() const
 {
 CUDF_EXPECTS(callback_grouped_values, "Uninitialized callback_grouped_values.");
 return callback_grouped_values();
 }

 [[nodiscard]] column_view get_sorted_grouped_values() const
 {
 CUDF_EXPECTS(callback_sorted_grouped_values, "Uninitialized callback_sorted_grouped_values.");
 return callback_sorted_grouped_values();
 }

 [[nodiscard]] size_type get_num_groups() const
 {
 CUDF_EXPECTS(callback_num_groups, "Uninitialized callback_num_groups.");
 return callback_num_groups();
 }

 [[nodiscard]] device_span<size_type const> get_group_offsets() const
 {
 CUDF_EXPECTS(callback_group_offsets, "Uninitialized callback_group_offsets.");
 return callback_group_offsets();
 }

 [[nodiscard]] device_span<size_type const> get_group_labels() const
 {
 CUDF_EXPECTS(callback_group_labels, "Uninitialized callback_group_labels.");
 return callback_group_labels();
 }

 [[nodiscard]] column_view compute_aggregation(std::unique_ptr<aggregation> other_agg) const
 {
 CUDF_EXPECTS(callback_compute_aggregation, "Uninitialized callback for computing aggregation.");
 return callback_compute_aggregation(std::move(other_agg));
 }
};
 // end of group
} // namespace CUDF_EXPORT cudf