list_device_view.cuh

/*
 * Copyright (c) 2020-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
 *
 * 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/detail/iterator.cuh>
#include <cudf/lists/lists_column_device_view.cuh>
#include <cudf/types.hpp>
#include <cudf/utilities/type_dispatcher.hpp>

#include <cuda_runtime.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <thrust/pair.h>

namespace CUDF_EXPORT cudf {

class list_device_view {
 using lists_column_device_view = cudf::detail::lists_column_device_view;

 public
 list_device_view() = default;

 __device__ inline list_device_view(lists_column_device_view const& lists_column,
 size_type const& row_index)
 : lists_column(lists_column), _row_index(row_index)
 {
 column_device_view const& offsets = lists_column.offsets();
 cudf_assert(row_index >= 0 && row_index < lists_column.size() && row_index < offsets.size() &&
 "row_index 超出范围");

 begin_offset = offsets.element<size_type>(row_index + lists_column.offset());
 cudf_assert(begin_offset >= 0 && begin_offset <= lists_column.child().size() &&
 "begin_offset 超出范围。");
 _size = offsets.element<size_type>(row_index + 1 + lists_column.offset()) - begin_offset;
 }

 ~list_device_view() = default;

 [[nodiscard]] __device__ inline size_type element_offset(size_type idx) const
 {
 cudf_assert(idx >= 0 && idx < size() && "idx 超出范围");
 return begin_offset + idx;
 }

 template <typename T>
 __device__ inline T element(size_type idx) const
 {
 return lists_column.child().element<T>(element_offset(idx));
 }

 [[nodiscard]] __device__ inline bool is_null(size_type idx) const
 {
 cudf_assert(idx >= 0 && idx < size() && "索引超出范围。");
 auto element_offset = begin_offset + idx;
 return lists_column.child().is_null(element_offset);
 }

 [[nodiscard]] __device__ inline bool is_null() const { return lists_column.is_null(_row_index); }

 [[nodiscard]] __device__ inline size_type size() const { return _size; }

 [[nodiscard]] __device__ inline size_type row_index() const { return _row_index; }

 [[nodiscard]] __device__ inline lists_column_device_view const& get_column() const
 {
 return lists_column;
 }

 template <typename T>
 struct pair_accessor;

 template <typename T>
 struct pair_rep_accessor;

 template <typename T>
 using const_pair_iterator =
 thrust::transform_iterator<pair_accessor<T>, thrust::counting_iterator<cudf::size_type>>;

 template <typename T>
 using const_pair_rep_iterator =
 thrust::transform_iterator<pair_rep_accessor<T>, thrust::counting_iterator<cudf::size_type>>;

 template <typename T>
 [[nodiscard]] __device__ inline const_pair_iterator<T> pair_begin() const
 {
 return const_pair_iterator<T>{thrust::counting_iterator<size_type>(0), pair_accessor<T>{*this}};
 }

 template <typename T>
 [[nodiscard]] __device__ inline const_pair_iterator<T> pair_end() const
 {
 return const_pair_iterator<T>{thrust::counting_iterator<size_type>(size()),
 pair_accessor<T>{*this}};
 }

 template <typename T>
 [[nodiscard]] __device__ inline const_pair_rep_iterator<T> pair_rep_begin() const
 {
 return const_pair_rep_iterator<T>{thrust::counting_iterator<size_type>(0),
 pair_rep_accessor<T>{*this}};
 }

 template <typename T>
 [[nodiscard]] __device__ inline const_pair_rep_iterator<T> pair_rep_end() const
 {
 return const_pair_rep_iterator<T>{thrust::counting_iterator<size_type>(size()),
 pair_rep_accessor<T>{*this}};
 }

 private
 lists_column_device_view const& lists_column;
 size_type _row_index{}; // 在 Lists 列向量中的行索引。
 size_type _size{}; // *此*列表行中的元素数量。

 size_type begin_offset; // 此列表在 list_column_device_view 中开始的偏移量。

 template <typename T>
 struct pair_accessor {
 list_device_view const& list;

 explicit CUDF_HOST_DEVICE inline pair_accessor(list_device_view const& _list) : list{_list} {}

 __device__ inline thrust::pair<T, bool> operator()(cudf::size_type i) const
 {
 return {list.element<T>(i), !list.is_null(i)};
 }
 };

 template <typename T>
 struct pair_rep_accessor {
 list_device_view const& list;

 using rep_type = device_storage_type_t<T>;

 explicit CUDF_HOST_DEVICE inline pair_rep_accessor(list_device_view const& _list) : list{_list}
 {
 }

 __device__ inline thrust::pair<rep_type, bool> operator()(cudf::size_type i) const
 {
 return {get_rep<T>(i), !list.is_null(i)};
 }

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

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

struct list_size_functor {
 detail::lists_column_device_view const d_column;
 CUDF_HOST_DEVICE inline list_size_functor(detail::lists_column_device_view const& d_col)
 : d_column(d_col)
 {
 }
 __device__ inline size_type operator()(size_type idx)
 {
 if (d_column.is_null(idx)) return size_type{0};
 return d_column.offset_at(idx + 1) - d_column.offset_at(idx);
 }
};

CUDF_HOST_DEVICE auto inline make_list_size_iterator(detail::lists_column_device_view const& c)
{
 return detail::make_counting_transform_iterator(0, list_size_functor{c});
}

} // 命名空间 CUDF_EXPORT cudf