contiguous_split.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2023-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/table/table.hpp>
#include <cudf/types.hpp>
 
#include <rmm/resource_ref.hpp>
 
#include <memory>
#include <vector>
 
namespace cudf {
 
struct packed_columns {
  packed_columns()
    : metadata(std::make_unique<std::vector<uint8_t>>()),
      gpu_data(std::make_unique<rmm::device_buffer>())
  {
  }
 
  packed_columns(std::unique_ptr<std::vector<uint8_t>>&& md,
                 std::unique_ptr<rmm::device_buffer>&& gd)
    : metadata(std::move(md)), gpu_data(std::move(gd))
  {
  }
 
  std::unique_ptr<std::vector<uint8_t>> metadata;  
  std::unique_ptr<rmm::device_buffer> gpu_data;    
};
 
struct packed_table {
  cudf::table_view table;  
  packed_columns data;     
};
 
std::vector<packed_table> contiguous_split(
  cudf::table_view const& input,
  std::vector<size_type> const& splits,
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
namespace detail {
struct contiguous_split_state;
};
 
class chunked_pack {
 public:
  explicit chunked_pack(
    cudf::table_view const& input,
    std::size_t user_buffer_size,
    rmm::device_async_resource_ref temp_mr = rmm::mr::get_current_device_resource());
 
  ~chunked_pack();
 
  [[nodiscard]] std::size_t get_total_contiguous_size() const;
 
  [[nodiscard]] bool has_next() const;
 
  [[nodiscard]] std::size_t next(cudf::device_span<uint8_t> const& user_buffer);
 
  [[nodiscard]] std::unique_ptr<std::vector<uint8_t>> build_metadata() const;
 
  [[nodiscard]] static std::unique_ptr<chunked_pack> create(
    cudf::table_view const& input,
    std::size_t user_buffer_size,
    rmm::device_async_resource_ref temp_mr = rmm::mr::get_current_device_resource());
 
 private:
  // internal state of contiguous split
  std::unique_ptr<detail::contiguous_split_state> state;
};
 
packed_columns pack(cudf::table_view const& input,
                    rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
std::vector<uint8_t> pack_metadata(table_view const& table,
                                   uint8_t const* contiguous_buffer,
                                   size_t buffer_size);
 
table_view unpack(packed_columns const& input);
 
table_view unpack(uint8_t const* metadata, uint8_t const* gpu_data);
 
}  // namespace cudf