contiguous_split.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <cudf/packed_types.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/memory_resource.hpp>
 
#include <cstdint>
#include <memory>
#include <span>
#include <vector>
 
namespace CUDF_EXPORT cudf {
 
std::vector<packed_table> contiguous_split(
  cudf::table_view const& input,
  std::vector<size_type> const& splits,
  cuda::stream_ref 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,
    cuda::stream_ref 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,
    cuda::stream_ref 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,
                    cuda::stream_ref stream           = cudf::get_default_stream(),
                    rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref());
 
std::size_t packed_size(
  cudf::table_view const& input,
  cuda::stream_ref stream                = cudf::get_default_stream(),
  rmm::device_async_resource_ref temp_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);
 
class packed_metadata_view {
 public:
  class column_view {
   public:
    [[nodiscard]] data_type type() const;
 
    [[nodiscard]] size_type num_rows() const;
 
    [[nodiscard]] size_type null_count() const;
 
    [[nodiscard]] size_type num_children() const;
 
    [[nodiscard]] column_view child(size_type i) const;
 
   private:
    friend class packed_metadata_view;
    data_type _type{type_id::EMPTY};
    size_type _size{};
    size_type _null_count{};
    size_type _num_children{};
    // Span from this entry to the end of the metadata buffer (needed for child traversal).
    std::span<std::uint8_t const> _buffer;
    explicit column_view(std::span<std::uint8_t const> buffer);
  };
 
  explicit packed_metadata_view(std::span<std::uint8_t const> buffer);
 
  [[nodiscard]] size_type num_columns() const;
 
  [[nodiscard]] size_type num_rows() const;
 
  [[nodiscard]] column_view column(size_type i) const;
 
 private:
  // Span from the first top-level column entry to the end of the metadata buffer.
  std::span<std::uint8_t const> _entries;
  size_type _num_columns{};
  // Table row count, read directly from the serialized table header.
  size_type _num_rows{};
};
 
}  // namespace CUDF_EXPORT cudf