groupby.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
 * 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 <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(
    host_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(
    host_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,
    host_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,
    host_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(
    host_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(
    host_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(
    host_span<scan_request const> requests,
    rmm::cuda_stream_view stream,
    rmm::device_async_resource_ref mr);
};
}  // namespace groupby
}  // namespace CUDF_EXPORT cudf