host_udf.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2024, 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
 *
 *     http://www.apache.org/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/types.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/span.hpp>
#include <cudf/utilities/traits.hpp>
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/resource_ref.hpp>
 
#include <optional>
#include <unordered_map>
#include <unordered_set>
#include <variant>
 
namespace CUDF_EXPORT cudf {
struct host_udf_base {
  host_udf_base()          = default;
  virtual ~host_udf_base() = default;
 
  enum class groupby_data_attribute : int32_t {
    INPUT_VALUES,    
    GROUPED_VALUES,  
    SORTED_GROUPED_VALUES,  
    NUM_GROUPS,             
    GROUP_OFFSETS,          
    GROUP_LABELS            
  };
 
  struct data_attribute {
    using value_type = std::variant<groupby_data_attribute, std::unique_ptr<aggregation>>;
    value_type value;  
 
    data_attribute()                 = default;  
    data_attribute(data_attribute&&) = default;  
 
    template <typename T, CUDF_ENABLE_IF(std::is_same_v<T, groupby_data_attribute>)>
    data_attribute(T value_) : value{value_}
    {
    }
 
    template <typename T,
              CUDF_ENABLE_IF(std::is_same_v<T, aggregation> ||
                             std::is_same_v<T, groupby_aggregation>)>
    data_attribute(std::unique_ptr<T> value_) : value{std::move(value_)}
    {
      CUDF_EXPECTS(std::get<std::unique_ptr<aggregation>>(value) != nullptr,
                   "Invalid aggregation request.");
      if constexpr (std::is_same_v<T, aggregation>) {
        CUDF_EXPECTS(
          dynamic_cast<groupby_aggregation*>(std::get<std::unique_ptr<T>>(value).get()) != nullptr,
          "Requesting results from other aggregations is only supported in groupby "
          "aggregations.");
      }
    }
 
    data_attribute(data_attribute const& other);
 
    struct hash {
      std::size_t operator()(data_attribute const& attr) const;
    };  // struct hash
 
    struct equal_to {
      bool operator()(data_attribute const& lhs, data_attribute const& rhs) const;
    };  // struct equal_to
  };    // struct data_attribute
 
  using data_attribute_set_t =
    std::unordered_set<data_attribute, data_attribute::hash, data_attribute::equal_to>;
 
  [[nodiscard]] virtual data_attribute_set_t get_required_data() const { return {}; }
 
  using input_data_t = std::variant<column_view, size_type, device_span<size_type const>>;
 
  using input_map_t = std::
    unordered_map<data_attribute, input_data_t, data_attribute::hash, data_attribute::equal_to>;
 
  using output_t = std::variant<std::unique_ptr<column>>;
 
  [[nodiscard]] virtual output_t get_empty_output(std::optional<data_type> output_dtype,
                                                  rmm::cuda_stream_view stream,
                                                  rmm::device_async_resource_ref mr) const = 0;
 
  [[nodiscard]] virtual output_t operator()(input_map_t const& input,
                                            rmm::cuda_stream_view stream,
                                            rmm::device_async_resource_ref mr) const = 0;
 
  [[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;
};
  // end of group
}  // namespace CUDF_EXPORT cudf