lists_column_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
 *
 *     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_device_view.cuh>
#include <cudf/lists/lists_column_view.hpp>
#include <cudf/types.hpp>
 
#include <cuda_runtime.h>
 
namespace cudf::detail {
 
class lists_column_device_view : private column_device_view {
 public:
  lists_column_device_view()                                = delete;
  ~lists_column_device_view()                               = default;
  lists_column_device_view(lists_column_device_view const&) = default;  
  lists_column_device_view(lists_column_device_view&&)      = default;  
  lists_column_device_view& operator=(lists_column_device_view const&) = default;
  lists_column_device_view& operator=(lists_column_device_view&&) = default;
 
  CUDF_HOST_DEVICE lists_column_device_view(column_device_view const& underlying_)
    : column_device_view(underlying_)
  {
#ifdef __CUDA_ARCH__
    cudf_assert(underlying_.type().id() == type_id::LIST and
                "lists_column_device_view only supports lists");
#else
    CUDF_EXPECTS(underlying_.type().id() == type_id::LIST,
                 "lists_column_device_view only supports lists");
#endif
  }
 
  using column_device_view::is_null;
  using column_device_view::nullable;
  using column_device_view::offset;
  using column_device_view::size;
 
  [[nodiscard]] __device__ inline column_device_view offsets() const
  {
    return column_device_view::child(lists_column_view::offsets_column_index);
  }
 
  [[nodiscard]] __device__ inline size_type offset_at(size_type idx) const
  {
    return offsets().size() > 0 ? offsets().element<size_type>(offset() + idx) : 0;
  }
 
  [[nodiscard]] __device__ inline column_device_view child() const
  {
    return column_device_view::child(lists_column_view::child_column_index);
  }
 
  [[nodiscard]] __device__ inline column_device_view get_sliced_child() const
  {
    auto start = offset_at(0);
    auto end   = offset_at(size());
    return child().slice(start, end - start);
  }
};
 
}  // namespace cudf::detail