column_device_view.cuh

Go to the documentation of this file.

/*
 * Copyright (c) 2019-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/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 <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/pair.h>
 
#include <algorithm>
 
namespace CUDF_EXPORT cudf {
 
struct nullate {
  struct YES : std::bool_constant<true> {};
  struct NO : std::bool_constant<false> {};
  struct DYNAMIC {
    DYNAMIC() = delete;
    constexpr explicit DYNAMIC(bool b) noexcept : value{b} {}
    constexpr operator bool() const noexcept { return value; }
    bool value;  
  };
};
 
namespace detail {
class alignas(16) column_device_view_base {
 public:
  column_device_view_base()                               = delete;
  ~column_device_view_base()                              = default;
  column_device_view_base(column_device_view_base const&) = default;  
  column_device_view_base(column_device_view_base&&)      = default;  
  column_device_view_base& operator=(column_device_view_base const&) = default;
  column_device_view_base& operator=(column_device_view_base&&) = default;
 
  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);
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  [[nodiscard]] CUDF_HOST_DEVICE T const* data() const noexcept
  {
    return head<T>() + _offset;
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE size_type size() const noexcept { return _size; }
 
  [[nodiscard]] CUDF_HOST_DEVICE data_type type() const noexcept { return _type; }
 
  [[nodiscard]] CUDF_HOST_DEVICE bool nullable() const noexcept { return nullptr != _null_mask; }
 
  [[nodiscard]] CUDF_HOST_DEVICE bitmask_type const* null_mask() const noexcept
  {
    return _null_mask;
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE size_type offset() const noexcept { return _offset; }
 
  [[nodiscard]] __device__ bool is_valid(size_type element_index) const noexcept
  {
    return not nullable() or is_valid_nocheck(element_index);
  }
 
  [[nodiscard]] __device__ bool is_valid_nocheck(size_type element_index) const noexcept
  {
    return bit_is_set(_null_mask, offset() + element_index);
  }
 
  [[nodiscard]] __device__ bool is_null(size_type element_index) const noexcept
  {
    return not is_valid(element_index);
  }
 
  [[nodiscard]] __device__ bool is_null_nocheck(size_type element_index) const noexcept
  {
    return not is_valid_nocheck(element_index);
  }
 
  [[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{};               
 
  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 : std::false_type {};
 
  template <typename C, typename T>
  struct has_element_accessor_impl<
    C,
    T,
    void_t<decltype(std::declval<C>().template element<T>(std::declval<size_type>()))>>
    : std::true_type {};
};
// @cond
// Forward declaration
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
 
class alignas(16) column_device_view : public detail::column_device_view_base {
 public:
  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,
                              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 data<T>()[element_index];
  }
 
  template <typename T, CUDF_ENABLE_IF(std::is_same_v<T, string_view>)>
  __device__ [[nodiscard]] 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:
  struct index_element_fn {
    template <typename IndexType,
              CUDF_ENABLE_IF(is_index_type<IndexType>() and std::is_unsigned_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_unsigned_v<IndexType>))>
    __device__ size_type operator()(Args&&... args)
    {
      CUDF_UNREACHABLE("dictionary indices must be an unsigned integral type");
    }
  };
 
 public:
  template <typename T, CUDF_ENABLE_IF(std::is_same_v<T, dictionary32>)>
  __device__ [[nodiscard]] 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)};
  }
 
  template <typename T, CUDF_ENABLE_IF(cudf::is_fixed_point<T>())>
  __device__ [[nodiscard]] 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}};
  }
 
  template <typename T>
  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 d_children[child_index];
  }
 
  [[nodiscard]] __device__ device_span<column_device_view const> children() const noexcept
  {
    return {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_base(type, size, data, null_mask, offset),
      d_children(children),
      _num_children(num_children)
  {
  }
 
 protected:
  column_device_view* d_children{};  
  size_type _num_children{};         
 
  column_device_view(column_view source);
};
 
class alignas(16) mutable_column_device_view : public detail::column_device_view_base {
 public:
  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 = void,
            CUDF_ENABLE_IF(std::is_same_v<T, void> or is_rep_layout_compatible<T>())>
  CUDF_HOST_DEVICE T* head() const noexcept
  {
    return const_cast<T*>(detail::column_device_view_base::head<T>());
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  CUDF_HOST_DEVICE T* data() const noexcept
  {
    return const_cast<T*>(detail::column_device_view_base::data<T>());
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  __device__ [[nodiscard]] T& element(size_type element_index) const noexcept
  {
    return data<T>()[element_index];
  }
 
  template <typename T>
  static constexpr bool has_element_accessor()
  {
    return has_element_accessor_impl<mutable_column_device_view, T>::value;
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE bitmask_type* null_mask() const noexcept
  {
    return const_cast<bitmask_type*>(detail::column_device_view_base::null_mask());
  }
 
  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 d_children[child_index];
  }
 
#ifdef __CUDACC__  // because set_bit in bit.hpp is wrapped with __CUDACC__
  __device__ void set_valid(size_type element_index) const noexcept
  {
    return set_bit(null_mask(), element_index);
  }
 
  __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>
  __device__ [[nodiscard]] 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>
  __device__ [[nodiscard]] 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