groupby.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <cudf/aggregation.hpp>
#include <cudf/column/column_view.hpp>
#include <cudf/replace.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/memory_resource.hpp>
#include <cudf/utilities/span.hpp>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <memory>
#include <span>
#include <utility>
#include <vector>
 
namespace CUDF_EXPORT cudf {
namespace groupby {
namespace detail {
namespace sort {
struct sort_groupby_helper;
 
}  // namespace sort
}  // namespace detail
 
struct aggregation_request {
  column_view values;                                              
  std::vector<std::unique_ptr<groupby_aggregation>> aggregations;  
};
 
struct scan_request {
  column_view values;  
  std::vector<std::unique_ptr<groupby_scan_aggregation>> aggregations;  
};
 
struct aggregation_result {
  std::vector<std::unique_ptr<column>> results{};
};
 
class groupby {
 public:
  groupby() = delete;
  ~groupby();
  groupby(groupby const&)            = delete;
  groupby(groupby&&)                 = delete;
  groupby& operator=(groupby const&) = delete;
  groupby& operator=(groupby&&)      = delete;
 
  explicit groupby(table_view const& keys,
                   null_policy null_handling                      = null_policy::EXCLUDE,
                   sorted keys_are_sorted                         = sorted::NO,
                   std::vector<order> const& column_order         = {},
                   std::vector<null_order> const& null_precedence = {});
 
  std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> aggregate(
    std::span<aggregation_request const> requests,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
  std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> scan(
    std::span<scan_request const> requests,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
  std::pair<std::unique_ptr<table>, std::unique_ptr<table>> shift(
    table_view const& values,
    std::span<size_type const> offsets,
    std::vector<std::reference_wrapper<scalar const>> const& fill_values,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
  struct groups {
    std::unique_ptr<table> keys;     
    std::vector<size_type> offsets;  
    std::unique_ptr<table> values;   
  };
 
  groups get_groups(cudf::table_view values           = {},
                    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
                    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
  std::pair<std::unique_ptr<table>, std::unique_ptr<table>> replace_nulls(
    table_view const& values,
    std::span<cudf::replace_policy const> replace_policies,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
 private:
  table_view _keys;                                      
  null_policy _include_null_keys{null_policy::EXCLUDE};  
  sorted _keys_are_sorted{sorted::NO};                   
  std::vector<order> _column_order{};                    
  std::vector<null_order> _null_precedence{};            
  std::unique_ptr<detail::sort::sort_groupby_helper>
    _helper;  
 
  detail::sort::sort_groupby_helper& helper();
 
  std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> dispatch_aggregation(
    std::span<aggregation_request const> requests,
    rmm::cuda_stream_view stream,
    rmm::device_async_resource_ref mr);
 
  // Sort-based groupby
  std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> sort_aggregate(
    std::span<aggregation_request const> requests,
    rmm::cuda_stream_view stream,
    rmm::device_async_resource_ref mr);
 
  std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> sort_scan(
    std::span<scan_request const> requests,
    rmm::cuda_stream_view stream,
    rmm::device_async_resource_ref mr);
};
 
struct streaming_aggregation_request {
  size_type column_index;                            
  std::unique_ptr<groupby_aggregation> aggregation;  
};
 
class streaming_groupby {
 public:
  streaming_groupby() = delete;
  ~streaming_groupby();
  streaming_groupby(streaming_groupby const&)            = delete;
  streaming_groupby& operator=(streaming_groupby const&) = delete;
 
  streaming_groupby(streaming_groupby&&) noexcept;
 
  streaming_groupby& operator=(streaming_groupby&&) noexcept;
 
  explicit streaming_groupby(host_span<size_type const> key_indices,
                             host_span<streaming_aggregation_request const> requests,
                             size_type max_distinct_keys,
                             null_policy null_handling = null_policy::EXCLUDE);
 
  void aggregate(table_view const& data, rmm::cuda_stream_view stream = cudf::get_default_stream());
 
  void merge(streaming_groupby const& other,
             rmm::cuda_stream_view stream = cudf::get_default_stream());
 
  [[nodiscard]] std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> finalize(
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const;
 
  [[nodiscard]] size_type distinct_keys() const noexcept;
 
 private:
  struct impl;
  std::unique_ptr<impl> _impl;
 
  void do_aggregate(table_view const& data, rmm::cuda_stream_view stream);
  void do_merge(streaming_groupby const& other, rmm::cuda_stream_view stream);
  [[nodiscard]] std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> do_finalize(
    rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const;
};
 
[[nodiscard]] bool is_streaming_groupby_supported(data_type values_type, aggregation::Kind kind);
 
}  // namespace groupby
}  // namespace CUDF_EXPORT cudf