column_device_view.cuh

前往此文件的文档。

/*
 * 版权所有 (c) 2019-2025, NVIDIA CORPORATION.
 *
 * 根据 Apache License, Version 2.0 ("许可") 获得许可；
 * 您除非遵守许可条款，否则不得使用此文件。
 * 您可获取许可协议的副本，地址为：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则根据许可分发的软件
 * 按"现状"提供，不附带任何明示或暗示的保证或条件。
 * 请参阅许可协议以了解管理权限和
 * 限制的特定语言。
 */

#pragma once

#include <cudf/column/column_view.hpp>
#include <cudf/detail/offsets_iterator.cuh>
#include <cudf/detail/utilities/alignment.hpp>
#include <cudf/fixed_point/fixed_point.hpp>
#include <cudf/lists/list_view.hpp>
#include <cudf/strings/string_view.cuh>
#include <cudf/strings/strings_column_view.hpp>
#include <cudf/structs/struct_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/bit.hpp>
#include <cudf/utilities/default_stream.hpp>
#include <cudf/utilities/span.hpp>
#include <cudf/utilities/traits.hpp>
#include <cudf/utilities/type_dispatcher.hpp>

#include <rmm/cuda_stream_view.hpp>

#include <cuda/std/optional>
#include <cuda/std/type_traits>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/pair.h>

#include <algorithm>
#include <functional>
#include <type_traits>

namespace CUDF_EXPORT cudf {

struct nullate {
 struct YES : cuda::std::true_type {};
 struct NO : cuda::std::false_type {};
 struct DYNAMIC {
 DYNAMIC() = delete;
 constexpr explicit DYNAMIC(bool b) noexcept : value{b} {}
 CUDF_HOST_DEVICE constexpr operator bool() const noexcept { return value; }
 bool value;
 };
};

namespace detail {
/**
 * @brief 设备端列视图的基类。
 *
 * 此类提供了在设备上访问列元数据和数据的通用功能。
 * 它旨在用作 `column_device_view` 和`
 * `mutable_column_device_view` 的基类。
 */
class alignas(16) column_device_view_base {
 public
 column_device_view_base() = delete;
 ~column_device_view_base() = default;
 /** @brief 拷贝构造函数 */
 column_device_view_base(column_device_view_base const&) = default;
 /** @brief 移动构造函数 */
 column_device_view_base(column_device_view_base&&) = default;
 /** @brief 拷贝赋值运算符 */
 column_device_view_base& operator=(column_device_view_base const&) = default;
 /** @brief 移动赋值运算符 */
 operator=(column_device_view_base&&) = default;

 /**
 * @brief 返回指向列原始数据缓冲区的指针。
 *
 * 此指针在设备上有效。
 *
 * @tparam T 元素类型。应与列的数据类型布局兼容。
 * @return 指向原始数据缓冲区的指针。
 */
 template <typename T = void,
 CUDF_ENABLE_IF(std::is_same_v<T, void> or is_rep_layout_compatible<T>())>
 [[nodiscard]] CUDF_HOST_DEVICE T const* head() const noexcept
 {
 return static_cast<T const*>(_data);
 }

 /**
 * @brief 返回指向列数据缓冲区第一个元素的指针。
 *
 * 此指针在设备上有效。
 *
 * @tparam T 元素类型。应与列的数据类型布局兼容。
 * @return 指向第一个元素的指针。
 */
 template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
 [[nodiscard]] CUDF_HOST_DEVICE T const* data() const noexcept
 {
 return head<T>() + _offset;
 }

 /**
 * @brief 返回列中的元素数量。
 */
 [[nodiscard]] CUDF_HOST_DEVICE size_type size() const noexcept { return _size; }

 /**
 * @brief 返回列的数据类型。
 */
 [[nodiscard]] CUDF_HOST_DEVICE data_type type() const noexcept { return _type; }

 /**
 * @brief 如果列有关联的 null mask，则返回 true。
 *
 * 此函数被标记为 `CUDF_HOST_DEVICE`，因为它需要从内核和`
 * 设备函数中调用。
 */
 [[nodiscard]] CUDF_HOST_DEVICE bool nullable() const noexcept { return nullptr != _null_mask; }

 /**
 * @brief 返回指向 null mask 缓冲区的指针，如果不可空则返回 `nullptr`。
 *
 * 此指针在设备上有效。
 */
 [[nodiscard]] CUDF_HOST_DEVICE bitmask_type const* null_mask() const noexcept
 {
 return _null_mask;
 }

 /**
 * @brief 返回列数据相对于缓冲区起点的偏移量。
 */
 [[nodiscard]] CUDF_HOST_DEVICE size_type offset() const noexcept { return _offset; }

 /**
 * @brief 检查给定索引处的元素是否有效（非 null）。
 *
 * 此函数被标记为 `__device__`，因为它仅适用于从设备`
 * 内核调用。
 *
 * @param element_index 要检查的元素的索引。
 * @return 如果元素有效则返回 true，否则返回 false。
 */
 [[nodiscard]] __device__ bool is_valid(size_type element_index) const noexcept
 {
 return not nullable() or is_valid_nocheck(element_index);
 }

 /**
 * @brief 检查给定索引处的元素是否有效（非 null），假设列可空。
 *
 * 此函数被标记为 `__device__`，因为它仅适用于从设备`
 * 内核调用。
 *
 * @param element_index 要检查的元素的索引。
 * @return 如果元素有效则返回 true，否则返回 false。
 */
 [[nodiscard]] __device__ bool is_valid_nocheck(size_type element_index) const noexcept
 {
 return bit_is_set(_null_mask, offset() + element_index);
 }

 /**
 * @brief 检查给定索引处的元素是否为 null。
 *
 * 此函数被标记为 `__device__`，因为它仅适用于从设备`
 * 内核调用。
 *
 * @param element_index 要检查的元素的索引。
 * @return 如果元素为 null 则返回 true，否则返回 false。
 */
 [[nodiscard]] __device__ bool is_null(size_type element_index) const noexcept
 {
 return not is_valid(element_index);
 }

 /**
 * @brief 检查给定索引处的元素是否为 null，假设列可空。
 *
 * 此函数被标记为 `__device__`，因为它仅适用于从设备`
 * 内核调用。
 *
 * @param element_index 要检查的元素的索引。
 * @return 如果元素为 null 则返回 true，否则返回 false。
 */
 [[nodiscard]] __device__ bool is_null_nocheck(size_type element_index) const noexcept
 {
 return not is_valid_nocheck(element_index);
 }

 /**
 * @brief 获取第 `word_index` 个位掩码字。
 *
 * @param word_index 要检索的位掩码字的索引。
 * @return 位掩码字。
 */
 [[nodiscard]] __device__ bitmask_type get_mask_word(size_type word_index) const noexcept
 {
 return null_mask()[word_index];
 }

 protected
 data_type _type{type_id::EMPTY};
 cudf::size_type _size{};
 void const* _data{};
 bitmask_type const* _null_mask{};
 size_type _offset{};

 /**
 * @brief 构造一个新的设备端列视图基类对象，使用原始指针和元数据。
 */
 CUDF_HOST_DEVICE column_device_view_base(data_type type,
 size_type size,
 void const* data,
 bitmask_type const* null_mask,
 size_type offset)
 : _type{type}, _size{size}, _data{data}, _null_mask{null_mask}, _offset{offset}
 {
 }

 template <typename C, typename T, typename = void>
 struct has_element_accessor_impl : cuda::std::false_type {};

 template <typename C, typename T>
 struct has_element_accessor_impl<
 C,
 T,
 void_t<decltype(cuda::std::declval<C>().template element<T>(cuda::std::declval<size_type>()))>>
 : cuda::std::true_type {};
};
// @cond
// 前向声明
template <typename T>
struct value_accessor;
template <typename T, typename Nullate>
struct optional_accessor;
template <typename T, bool has_nulls>
struct pair_accessor;
template <typename T, bool has_nulls>
struct pair_rep_accessor;
template <typename T>
struct mutable_value_accessor;
// @endcond
} // namespace detail

/**
 * @brief 设备上不可变列数据和元数据的视图。
 *
 * `column_device_view` 是用于设备内核的列数据的不可变视图。
 * 它包含指向列数据、null mask 和子列的设备指针。
 */
class alignas(16) column_device_view : public detail::column_device_view_base {
 public
 column_device_view() = delete;
 ~column_device_view() = default;
 /** @brief 拷贝构造函数 */
 column_device_view(column_device_view const&) = default;
 /** @brief 移动构造函数 */
 column_device_view(column_device_view&&) = default;
 /** @brief 拷贝赋值运算符 */
 column_device_view& operator=(column_device_view const&) = default;
 /** @brief 移动赋值运算符 */
 column_device_view& operator=(column_device_view&&) = default;

 /**
 * @brief 从主机端列视图和指向子列视图的设备指针构造一个新的设备端列视图对象。
 *
 * 此构造函数仅打算由工厂函数 `create` 调用。
 *
 * @param column 用于构造的主机端列视图。
 * @param h_ptr 指向支持此视图的设备内存的主机端指针。用于计算到数据缓冲区的偏移量，该偏移量是根据 `column_view` 及其 `data()` 成员的相对地址计算的。
 * @param d_ptr 指向支持此视图的设备内存的设备端指针。
 */
 column_device_view(column_view column, void* h_ptr, void* d_ptr);

 /**
 * @brief 返回元素子范围的新视图。
 *
 * 返回的视图与原始视图具有相同的类型、null mask 和子项。
 * 新视图的数据指针和 null mask 指针根据请求的子范围进行偏移。
 *
 * @param offset 子范围的起始索引。
 * @param size 子范围中的元素数量。
 * @return 指定子范围的新设备端列视图。
 */
 [[nodiscard]] CUDF_HOST_DEVICE column_device_view slice(size_type offset,
 size_type size) const noexcept
 {
 return column_device_view{this->type(),
 size,
 this->head(),
 this->null_mask(),
 this->offset() + offset,
 d_children,
 this->num_child_columns()};
 }

 /**
 * @brief 返回给定索引处的元素。
 *
 * 索引相对于此视图的起始位置。
 *
 * @tparam T 元素类型。必须与列的数据类型布局兼容。
 * @param element_index 要检索的元素的索引。
 * @return 给定索引处的元素。
 */
 template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
 [[nodiscard]] __device__ T element(size_type element_index) const noexcept
 {
 return data<T>()[element_index];
 }

 /**
 * @brief 对于字符串列，返回给定索引处的元素。
 *
 * 索引相对于此视图的起始位置。
 *
 * @tparam T 必须为 `string_view`。
 * @param element_index 要检索的元素的索引。
 * @return 给定索引处的字符串元素。
 */
 template <typename T, CUDF_ENABLE_IF(std::is_same_v<T, string_view>)>
 [[nodiscard]] __device__ T element(size_type element_index) const noexcept
 {
 size_type index = element_index + offset(); // account for this view's _offset
 char const* d_strings = static_cast<char const*>(_data);
 auto const offsets = d_children[strings_column_view::offsets_column_index];
 auto const itr = cudf::detail::input_offsetalator(offsets.head(), offsets.type());
 auto const offset = itr[index];
 return string_view{d_strings + offset, static_cast<cudf::size_type>(itr[index + 1] - offset)};
 }

 private
 /**
 * @brief 从索引类型列的 device_view 中提取元素的函数对象。
 */
 struct index_element_fn {
 template <typename IndexType,
 CUDF_ENABLE_IF(is_index_type<IndexType>() and std::is_signed_v<IndexType>)>
 __device__ size_type operator()(column_device_view const& indices, size_type index)
 {
 return static_cast<size_type>(indices.element<IndexType>(index));
 }

 template <typename IndexType,
 typename... Args,
 CUDF_ENABLE_IF(not(is_index_type<IndexType>() and std::is_signed_v<IndexType>))>
 __device__ size_type operator()(Args&&... args)
 {
 CUDF_UNREACHABLE("dictionary indices must be a signed integral type");
 }
 };

 public
 /**
 * @brief 对于字典列，返回给定索引处的元素。
 *
 * 索引相对于此视图的起始位置。
 *
 * @tparam T 必须为 `dictionary32`。
 * @param element_index 要检索的元素的索引。
 * @return 给定索引处的字典元素。
 */
 template <typename T, CUDF_ENABLE_IF(std::is_same_v<T, dictionary32>)>
 [[nodiscard]] __device__ T element(size_type element_index) const noexcept
 {
 size_type index = element_index + offset(); // account for this view's _offset
 auto const indices = d_children[0];
 return dictionary32{type_dispatcher(indices.type(), index_element_fn{}, indices, index)};
 }

 /**
 * @brief 对于 fixed_point 类型，返回给定索引处的元素。
 *
 * @tparam T 必须是 fixed_point 类型。
 * @param element_index 要检索的元素的索引。
 * @return 给定索引处的 fixed_point 元素。
 */
 template <typename T, CUDF_ENABLE_IF(cudf::is_fixed_point<T>())>
 [[nodiscard]] __device__ T element(size_type element_index) const noexcept
 {
 using namespace numeric;
 using rep = typename T::rep;
 auto const scale = scale_type{_type.scale()};
 return T{scaled_integer<rep>{data<rep>()[element_index], scale}};
 }

 /**
 * @brief 如果此类型受 `element()` 支持，则返回 true。
 *
 * @tparam T 要检查的类型。
 */
 template <typename T>
 CUDF_HOST_DEVICE static constexpr bool has_element_accessor()
 {
 return has_element_accessor_impl<column_device_view, T>::value;
 }

 /** @brief 列索引计数迭代器的类型别名 */
 using count_it = thrust::counting_iterator<size_type>;
 /**
 * @brief 列元素常量迭代器的类型别名。
 *
 * 解引用迭代器返回元素的副本。
 */
 template <typename T>
 using const_iterator = thrust::transform_iterator<detail::value_accessor<T>, count_it>;

 /**
 * @brief 返回指向第一个元素的迭代器。
 *
 * 列类型必须受 `element()` 支持。
 *
 * @tparam T 元素类型。
 * @return 指向第一个元素的迭代器。
 */
 template <typename T, CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
 [[nodiscard]] const_iterator<T> begin() const
 {
 return const_iterator<T>{count_it{0}, detail::value_accessor<T>{*this}};
 }

 /**
 * @brief 返回指向列末尾的迭代器。
 *
 * 列类型必须受 `element()` 支持。
 *
 * @tparam T 元素类型。
 * @return 指向列末尾的迭代器。
 */
 template <typename T, CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
 [[nodiscard]] const_iterator<T> end() const
 {
 __device__ void set_null(size_type element_index) const noexcept
 {
 return clear_bit(null_mask(), element_index);
 }

#endif

 __device__ void set_mask_word(size_type word_index, bitmask_type new_word) const noexcept
 {
 null_mask()[word_index] = new_word;
 }

 static std::size_t extent(mutable_column_view source_view);

 void destroy();

 private
 mutable_column_device_view* d_children{};
 size_type _num_children{};

 mutable_column_device_view(mutable_column_view source);
};

namespace detail {

#ifdef __CUDACC__ // because set_bit in bit.hpp is wrapped with __CUDACC__

__device__ inline bitmask_type get_mask_offset_word(bitmask_type const* __restrict__ source,
 size_type destination_word_index,
 size_type source_begin_bit,
 size_type source_end_bit)
{
 size_type source_word_index = destination_word_index + word_index(source_begin_bit);
 bitmask_type curr_word = source[source_word_index];
 bitmask_type next_word = 0;
 if (word_index(source_end_bit - 1) >
 word_index(source_begin_bit +
 destination_word_index * detail::size_in_bits<bitmask_type>())) {
 next_word = source[source_word_index + 1];
 }
 return __funnelshift_r(curr_word, next_word, source_begin_bit);
}

#endif

template <typename T>
struct value_accessor {
 column_device_view const col;

 value_accessor(column_device_view const& _col) : col{_col}
 {
 CUDF_EXPECTS(type_id_matches_device_storage_type<T>(col.type().id()), "the data type mismatch");
 }

 __device__ T operator()(cudf::size_type i) const { return col.element<T>(i); }
};

template <typename T, typename Nullate>
struct optional_accessor {
 column_device_view const col;

 optional_accessor(column_device_view const& _col, Nullate with_nulls)
 : col{_col}, has_nulls{with_nulls}
 {
 CUDF_EXPECTS(type_id_matches_device_storage_type<T>(col.type().id()), "the data type mismatch");
 if (with_nulls) { CUDF_EXPECTS(_col.nullable(), "Unexpected non-nullable column."); }
 }

 __device__ inline cuda::std::optional<T> operator()(cudf::size_type i) const
 {
 if (has_nulls) {
 return (col.is_valid_nocheck(i)) ? cuda::std::optional<T>{col.element<T>(i)}
 : cuda::std::optional<T>{cuda::std::nullopt};
 }
 return cuda::std::optional<T>{col.element<T>(i)};
 }

 Nullate has_nulls{};
};

template <typename T, bool has_nulls = false>
struct pair_accessor {
 column_device_view const col;

 pair_accessor(column_device_view const& _col) : col{_col}
 {
 CUDF_EXPECTS(type_id_matches_device_storage_type<T>(col.type().id()), "the data type mismatch");
 if (has_nulls) { CUDF_EXPECTS(_col.nullable(), "Unexpected non-nullable column."); }
 }

 __device__ inline thrust::pair<T, bool> operator()(cudf::size_type i) const
 {
 return {col.element<T>(i), (has_nulls ? col.is_valid_nocheck(i) : true)};
 }
};

template <typename T, bool has_nulls = false>
struct pair_rep_accessor {
 column_device_view const col;

 using rep_type = device_storage_type_t<T>;

 pair_rep_accessor(column_device_view const& _col) : col{_col}
 {
 CUDF_EXPECTS(type_id_matches_device_storage_type<T>(col.type().id()), "the data type mismatch");
 if (has_nulls) { CUDF_EXPECTS(_col.nullable(), "Unexpected non-nullable column."); }
 }

 __device__ inline thrust::pair<rep_type, bool> operator()(cudf::size_type i) const
 {
 return {get_rep<T>(i), (has_nulls ? col.is_valid_nocheck(i) : true)};
 }

 private
 template <typename R, std::enable_if_t<std::is_same_v<R, rep_type>, void>* = nullptr>
 [[nodiscard]] __device__ inline auto get_rep(cudf::size_type i) const
 {
 return col.element<R>(i);
 }

 template <typename R, std::enable_if_t<not std::is_same_v<R, rep_type>, void>* = nullptr>
 [[nodiscard]] __device__ inline auto get_rep(cudf::size_type i) const
 {
 return col.element<R>(i).value();
 }
};

template <typename T>
struct mutable_value_accessor {
 mutable_column_device_view col;

 mutable_value_accessor(mutable_column_device_view& _col) : col{_col}
 {
 CUDF_EXPECTS(type_id_matches_device_storage_type<T>(col.type().id()), "the data type mismatch");
 }

 __device__ T& operator()(cudf::size_type i) { return col.element<T>(i); }
};

template <typename ColumnDeviceView, typename ColumnViewIterator>
ColumnDeviceView* child_columns_to_device_array(ColumnViewIterator child_begin,
 ColumnViewIterator child_end,
 void* h_ptr,
 void* d_ptr)
{
 ColumnDeviceView* d_children = detail::align_ptr_for_type<ColumnDeviceView>(d_ptr);
 auto num_children = std::distance(child_begin, child_end);
 if (num_children > 0) {
 // The beginning of the memory must be the fixed-sized ColumnDeviceView
 // struct objects in order for d_children to be used as an array.
 auto h_column = detail::align_ptr_for_type<ColumnDeviceView>(h_ptr);
 auto d_column = d_children;

 // Any child data is assigned past the end of this array: h_end and d_end.
 auto h_end = reinterpret_cast<int8_t*>(h_column + num_children);
 auto d_end = reinterpret_cast<int8_t*>(d_column + num_children);
 std::for_each(child_begin, child_end, [&](auto const& col) {
 // inplace-new each child into host memory
 new (h_column) ColumnDeviceView(col, h_end, d_end);
 h_column++; // advance to next child
 // update the pointers for holding this child column's child data
 auto col_child_data_size = ColumnDeviceView::extent(col) - sizeof(ColumnDeviceView);
 h_end += col_child_data_size;
 d_end += col_child_data_size;
 });
 }
 return d_children;
}

} // namespace detail
} // namespace CUDF_EXPORT cudf