host_udf.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#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