column_device_view.cuh

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
 
#include <cudf/column/column_device_view_base.cuh>
#include <cudf/column/column_view.hpp>
#include <cudf/detail/utilities/alignment.hpp>
#include <cudf/lists/list_view.hpp>
#include <cudf/strings/strings_column_view.hpp>
#include <cudf/structs/struct_view.hpp>
#include <cudf/utilities/default_stream.hpp>
#include <cudf/utilities/span.hpp>
#include <cudf/utilities/type_dispatcher.hpp>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/pair.h>
 
#include <functional>
 
namespace CUDF_EXPORT cudf {
 
class alignas(16) column_device_view : public column_device_view_core {
 public:
  using base = column_device_view_core;  
 
  column_device_view()                          = delete;
  ~column_device_view()                         = default;
  column_device_view(column_device_view const&) = default;  
  column_device_view(column_device_view&&)      = default;  
  column_device_view& operator=(column_device_view const&) = default;
  column_device_view& operator=(column_device_view&&) = default;
 
  column_device_view(column_view column, void* h_ptr, void* d_ptr);
 
  [[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,
                              static_cast<column_device_view*>(d_children),
                              this->num_child_columns()};
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  [[nodiscard]] __device__ T element(size_type element_index) const noexcept
  {
    return base::element<T>(element_index);
  }
 
  template <typename T, CUDF_ENABLE_IF(cuda::std::is_same_v<T, string_view>)>
  [[nodiscard]] __device__ T element(size_type element_index) const noexcept
  {
    return base::element<T>(element_index);
  }
 
  template <typename T, CUDF_ENABLE_IF(cudf::is_fixed_point<T>())>
  [[nodiscard]] __device__ T element(size_type element_index) const noexcept
  {
    return base::element<T>(element_index);
  }
 
 private:
  struct index_element_fn {
    template <typename IndexType,
              CUDF_ENABLE_IF(is_index_type<IndexType>() and cuda::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 cuda::std::is_signed_v<IndexType>))>
    __device__ size_type operator()(Args&&... args)
    {
      CUDF_UNREACHABLE("dictionary indices must be a signed integral type");
    }
  };
 
 public:
  template <typename T, CUDF_ENABLE_IF(cuda::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 = child(0);
    return dictionary32{type_dispatcher(indices.type(), index_element_fn{}, indices, index)};
  }
 
  template <typename T>
  CUDF_HOST_DEVICE static constexpr bool has_element_accessor()
  {
    return has_element_accessor_impl<column_device_view, T>::value;
  }
 
  using count_it = thrust::counting_iterator<size_type>;
  template <typename T>
  using const_iterator = thrust::transform_iterator<detail::value_accessor<T>, count_it>;
 
  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}};
  }
 
  template <typename T, CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
  [[nodiscard]] const_iterator<T> end() const
  {
    return const_iterator<T>{count_it{size()}, detail::value_accessor<T>{*this}};
  }
 
  template <typename T, typename Nullate>
  using const_optional_iterator =
    thrust::transform_iterator<detail::optional_accessor<T, Nullate>, count_it>;
 
  template <typename T, bool has_nulls>
  using const_pair_iterator =
    thrust::transform_iterator<detail::pair_accessor<T, has_nulls>, count_it>;
 
  template <typename T, bool has_nulls>
  using const_pair_rep_iterator =
    thrust::transform_iterator<detail::pair_rep_accessor<T, has_nulls>, count_it>;
 
  template <typename T,
            typename Nullate,
            CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
  auto optional_begin(Nullate has_nulls) const
  {
    return const_optional_iterator<T, Nullate>{
      count_it{0}, detail::optional_accessor<T, Nullate>{*this, has_nulls}};
  }
 
  template <typename T,
            bool has_nulls,
            CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
  [[nodiscard]] const_pair_iterator<T, has_nulls> pair_begin() const
  {
    return const_pair_iterator<T, has_nulls>{count_it{0},
                                             detail::pair_accessor<T, has_nulls>{*this}};
  }
 
  template <typename T,
            bool has_nulls,
            CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
  [[nodiscard]] const_pair_rep_iterator<T, has_nulls> pair_rep_begin() const
  {
    return const_pair_rep_iterator<T, has_nulls>{count_it{0},
                                                 detail::pair_rep_accessor<T, has_nulls>{*this}};
  }
 
  template <typename T,
            typename Nullate,
            CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
  auto optional_end(Nullate has_nulls) const
  {
    return const_optional_iterator<T, Nullate>{
      count_it{size()}, detail::optional_accessor<T, Nullate>{*this, has_nulls}};
  }
 
  template <typename T,
            bool has_nulls,
            CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
  [[nodiscard]] const_pair_iterator<T, has_nulls> pair_end() const
  {
    return const_pair_iterator<T, has_nulls>{count_it{size()},
                                             detail::pair_accessor<T, has_nulls>{*this}};
  }
 
  template <typename T,
            bool has_nulls,
            CUDF_ENABLE_IF(column_device_view::has_element_accessor<T>())>
  [[nodiscard]] const_pair_rep_iterator<T, has_nulls> pair_rep_end() const
  {
    return const_pair_rep_iterator<T, has_nulls>{count_it{size()},
                                                 detail::pair_rep_accessor<T, has_nulls>{*this}};
  }
 
  static std::unique_ptr<column_device_view, std::function<void(column_device_view*)>> create(
    column_view source_view, rmm::cuda_stream_view stream = cudf::get_default_stream());
 
  void destroy();
 
  static std::size_t extent(column_view const& source_view);
 
  [[nodiscard]] __device__ column_device_view child(size_type child_index) const noexcept
  {
    return static_cast<column_device_view*>(d_children)[child_index];
  }
 
  [[nodiscard]] __device__ device_span<column_device_view const> children() const noexcept
  {
    return {static_cast<column_device_view*>(d_children), static_cast<std::size_t>(_num_children)};
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE size_type num_child_columns() const noexcept
  {
    return _num_children;
  }
 
 private:
  CUDF_HOST_DEVICE column_device_view(data_type type,
                                      size_type size,
                                      void const* data,
                                      bitmask_type const* null_mask,
                                      size_type offset,
                                      column_device_view* children,
                                      size_type num_children)
    : column_device_view_core{type, size, data, null_mask, offset, children, num_children}
  {
  }
 
  column_device_view(column_view source);
};
 
class alignas(16) mutable_column_device_view : public mutable_column_device_view_core {
 public:
  using base = mutable_column_device_view_core;  
 
  mutable_column_device_view()                                  = delete;
  ~mutable_column_device_view()                                 = default;
  mutable_column_device_view(mutable_column_device_view const&) = default;  
  mutable_column_device_view(mutable_column_device_view&&)      = default;  
  mutable_column_device_view& operator=(mutable_column_device_view const&) = default;
  mutable_column_device_view& operator=(mutable_column_device_view&&) = default;
 
  mutable_column_device_view(mutable_column_view column, void* h_ptr, void* d_ptr);
 
  static std::unique_ptr<mutable_column_device_view,
                         std::function<void(mutable_column_device_view*)>>
  create(mutable_column_view source_view,
         rmm::cuda_stream_view stream = cudf::get_default_stream());
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  [[nodiscard]] __device__ T& element(size_type element_index) const noexcept
  {
    return base::element<T>(element_index);
  }
 
  template <typename T>
  CUDF_HOST_DEVICE static constexpr bool has_element_accessor()
  {
    return has_element_accessor_impl<mutable_column_device_view, T>::value;
  }
 
  using count_it = thrust::counting_iterator<size_type>;
  template <typename T>
  using iterator = thrust::transform_iterator<detail::mutable_value_accessor<T>, count_it>;
 
  template <typename T, CUDF_ENABLE_IF(mutable_column_device_view::has_element_accessor<T>())>
  iterator<T> begin()
  {
    return iterator<T>{count_it{0}, detail::mutable_value_accessor<T>{*this}};
  }
 
  template <typename T, CUDF_ENABLE_IF(mutable_column_device_view::has_element_accessor<T>())>
  iterator<T> end()
  {
    return iterator<T>{count_it{size()}, detail::mutable_value_accessor<T>{*this}};
  }
 
  [[nodiscard]] __device__ mutable_column_device_view child(size_type child_index) const noexcept
  {
    return static_cast<mutable_column_device_view*>(d_children)[child_index];
  }
 
  static std::size_t extent(mutable_column_view source_view);
 
  void destroy();
 
 private:
  mutable_column_device_view(mutable_column_view source);
};
 
static_assert(sizeof(column_device_view) == sizeof(column_device_view_core),
              "column_device_view and raw_column_device_view must be bitwise-compatible");
 
static_assert(
  sizeof(mutable_column_device_view) == sizeof(mutable_column_device_view_core),
  "mutable_column_device_view and raw_mutable_column_device_view must be bitwise-compatible");
 
namespace detail {
 
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>
  [[nodiscard]] __device__ inline auto get_rep(cudf::size_type i) const
    requires(std::is_same_v<R, rep_type>)
  {
    return col.element<R>(i);
  }
 
  template <typename R>
  [[nodiscard]] __device__ inline auto get_rep(cudf::size_type i) const
    requires(not std::is_same_v<R, rep_type>)
  {
    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