groupby.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2019-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/column/column_view.hpp>
#include <cudf/replace.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/span.hpp>
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <memory>
#include <utility>
#include <vector>
 
namespace cudf {
namespace groupby {
namespace detail {
namespace sort {
class 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::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
  std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> aggregate(
    host_span<aggregation_request const> requests,
    rmm::cuda_stream_view stream,
    rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
  std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> scan(
    host_span<scan_request const> requests,
    rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
  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::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
  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::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
  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::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
 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