contiguous_split.hpp

转到此文件的文档。

/*
 * Copyright (c) 2023-2025, 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
 *
 * https://apache.ac.cn/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 <cudf/utilities/export.hpp>
#include <cudf/utilities/memory_resource.hpp>

#include <memory>
#include <vector>

namespace CUDF_EXPORT 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::cuda_stream_view stream = cudf::get_default_stream(),
 rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());

namespace detail {

struct contiguous_split_state;
} // namespace detail

class chunked_pack {
 public
 explicit chunked_pack(
 cudf::table_view const& input,
 std::size_t user_buffer_size,
 rmm::cuda_stream_view stream = cudf::get_default_stream(),
 rmm::device_async_resource_ref temp_mr = cudf::get_current_device_resource_ref());

 ~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::cuda_stream_view stream = cudf::get_default_stream(),
 rmm::device_async_resource_ref temp_mr = cudf::get_current_device_resource_ref());

 private
 // internal state of contiguous split
 std::unique_ptr<detail::contiguous_split_state> state;
};

packed_columns pack(cudf::table_view const& input,
 rmm::cuda_stream_view stream = cudf::get_default_stream(),
 rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());

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_EXPORT cudf