column_wrapper.hpp

/*
 * 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_test/column_utilities.hpp>
#include <cudf_test/default_stream.hpp>
 
#include <cudf/column/column.hpp>
#include <cudf/column/column_factories.hpp>
#include <cudf/copying.hpp>
#include <cudf/detail/concatenate.hpp>
#include <cudf/detail/iterator.cuh>
#include <cudf/detail/utilities/vector_factories.hpp>
#include <cudf/dictionary/encode.hpp>
#include <cudf/fixed_point/fixed_point.hpp>
#include <cudf/lists/lists_column_view.hpp>
#include <cudf/null_mask.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/bit.hpp>
#include <cudf/utilities/default_stream.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/memory_resource.hpp>
#include <cudf/utilities/traits.hpp>
#include <cudf/utilities/type_dispatcher.hpp>
 
#include <rmm/device_buffer.hpp>
 
#include <thrust/copy.h>
#include <thrust/functional.h>
#include <thrust/host_vector.h>
#include <thrust/iterator/constant_iterator.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
 
#include <algorithm>
#include <iterator>
#include <memory>
#include <numeric>
 
namespace CUDF_EXPORT cudf {
namespace test {
namespace detail {
class column_wrapper {
 public:
  operator column_view() const { return wrapped->view(); }
 
  operator mutable_column_view() { return wrapped->mutable_view(); }
 
  std::unique_ptr<cudf::column> release() { return std::move(wrapped); }
 
 protected:
  std::unique_ptr<cudf::column> wrapped{};  
};
 
template <typename From, typename To>
struct fixed_width_type_converter {
  template <typename FromT                                      = From,
            typename ToT                                        = To,
            std::enable_if_t<std::is_same_v<FromT, ToT>, void>* = nullptr>
  constexpr ToT operator()(FromT element) const
  {
    return element;
  }
 
  template <
    typename FromT          = From,
    typename ToT            = To,
    std::enable_if_t<!std::is_same_v<FromT, ToT> && (cudf::is_convertible<FromT, ToT>::value ||
                                                     std::is_constructible_v<ToT, FromT>),
                     void>* = nullptr>
  constexpr ToT operator()(FromT element) const
  {
    return static_cast<ToT>(element);
  }
 
  template <
    typename FromT                                                                  = From,
    typename ToT                                                                    = To,
    std::enable_if_t<std::is_integral_v<FromT> && cudf::is_timestamp<ToT>(), void>* = nullptr>
  constexpr ToT operator()(FromT element) const
  {
    return ToT{typename ToT::duration{element}};
  }
};
 
template <typename ElementTo,
          typename ElementFrom,
          typename InputIterator,
          std::enable_if_t<not cudf::is_fixed_point<ElementTo>()>* = nullptr>
rmm::device_buffer make_elements(InputIterator begin, InputIterator end)
{
  static_assert(cudf::is_fixed_width<ElementTo>(), "Unexpected non-fixed width type.");
  auto transformer     = fixed_width_type_converter<ElementFrom, ElementTo>{};
  auto transform_begin = thrust::make_transform_iterator(begin, transformer);
  auto const size      = cudf::distance(begin, end);
  auto const elements  = thrust::host_vector<ElementTo>(transform_begin, transform_begin + size);
  return rmm::device_buffer{
    elements.data(), size * sizeof(ElementTo), cudf::test::get_default_stream()};
}
 
// The two signatures below are identical to the above overload apart from
// SFINAE so doxygen sees it as a duplicate.
 
template <typename ElementTo,
          typename ElementFrom,
          typename InputIterator,
          std::enable_if_t<not cudf::is_fixed_point<ElementFrom>() and
                           cudf::is_fixed_point<ElementTo>()>* = nullptr>
rmm::device_buffer make_elements(InputIterator begin, InputIterator end)
{
  using RepType        = typename ElementTo::rep;
  auto transformer     = fixed_width_type_converter<ElementFrom, RepType>{};
  auto transform_begin = thrust::make_transform_iterator(begin, transformer);
  auto const size      = cudf::distance(begin, end);
  auto const elements  = thrust::host_vector<RepType>(transform_begin, transform_begin + size);
  return rmm::device_buffer{
    elements.data(), size * sizeof(RepType), cudf::test::get_default_stream()};
}
 
template <typename ElementTo,
          typename ElementFrom,
          typename InputIterator,
          std::enable_if_t<cudf::is_fixed_point<ElementFrom>() and
                           cudf::is_fixed_point<ElementTo>()>* = nullptr>
rmm::device_buffer make_elements(InputIterator begin, InputIterator end)
{
  using namespace numeric;
  using RepType = typename ElementTo::rep;
 
  CUDF_EXPECTS(std::all_of(begin, end, [](ElementFrom v) { return v.scale() == 0; }),
               "Only zero-scale fixed-point values are supported");
 
  auto to_rep            = [](ElementTo fp) { return fp.value(); };
  auto transformer_begin = thrust::make_transform_iterator(begin, to_rep);
  auto const size        = cudf::distance(begin, end);
  auto const elements = thrust::host_vector<RepType>(transformer_begin, transformer_begin + size);
  return rmm::device_buffer{
    elements.data(), size * sizeof(RepType), cudf::test::get_default_stream()};
}
 
template <typename ValidityIterator>
std::pair<std::vector<bitmask_type>, cudf::size_type> make_null_mask_vector(ValidityIterator begin,
                                                                            ValidityIterator end)
{
  auto const size      = cudf::distance(begin, end);
  auto const num_words = cudf::bitmask_allocation_size_bytes(size) / sizeof(bitmask_type);
 
  auto null_mask  = std::vector<bitmask_type>(num_words, 0);
  auto null_count = cudf::size_type{0};
  for (auto i = 0; i < size; ++i) {
    if (*(begin + i)) {
      set_bit_unsafe(null_mask.data(), i);
    } else {
      ++null_count;
    }
  }
 
  return {std::move(null_mask), null_count};
}
 
template <typename ValidityIterator>
std::pair<rmm::device_buffer, cudf::size_type> make_null_mask(ValidityIterator begin,
                                                              ValidityIterator end)
{
  auto [null_mask, null_count] = make_null_mask_vector(begin, end);
  auto d_mask                  = rmm::device_buffer{null_mask.data(),
                                   cudf::bitmask_allocation_size_bytes(cudf::distance(begin, end)),
                                   cudf::test::get_default_stream()};
  return {std::move(d_mask), null_count};
}
 
template <typename StringsIterator, typename ValidityIterator>
auto make_chars_and_offsets(StringsIterator begin, StringsIterator end, ValidityIterator v)
{
  std::vector<char> chars{};
  std::vector<cudf::size_type> offsets(1, 0);
  for (auto str = begin; str < end; ++str) {
    std::string tmp = (*v++) ? std::string(*str) : std::string{};
    chars.insert(chars.end(), std::cbegin(tmp), std::cend(tmp));
    auto const last_offset = static_cast<std::size_t>(offsets.back());
    auto const next_offset = last_offset + tmp.length();
    CUDF_EXPECTS(
      next_offset < static_cast<std::size_t>(std::numeric_limits<cudf::size_type>::max()),
      "Cannot use strings_column_wrapper to build a large strings column");
    offsets.push_back(static_cast<cudf::size_type>(next_offset));
  }
  return std::pair(std::move(chars), std::move(offsets));
};
}  // namespace detail
 
template <typename ElementTo, typename SourceElementT = ElementTo>
class fixed_width_column_wrapper : public detail::column_wrapper {
 public:
  fixed_width_column_wrapper() : column_wrapper{}
  {
    std::vector<ElementTo> empty;
    wrapped.reset(
      new cudf::column{cudf::data_type{cudf::type_to_id<ElementTo>()},
                       0,
                       detail::make_elements<ElementTo, SourceElementT>(empty.begin(), empty.end()),
                       rmm::device_buffer{},
                       0});
  }
 
  template <typename InputIterator>
  fixed_width_column_wrapper(InputIterator begin, InputIterator end) : column_wrapper{}
  {
    auto const size = cudf::distance(begin, end);
    wrapped.reset(new cudf::column{cudf::data_type{cudf::type_to_id<ElementTo>()},
                                   size,
                                   detail::make_elements<ElementTo, SourceElementT>(begin, end),
                                   rmm::device_buffer{},
                                   0});
  }
 
  template <typename InputIterator, typename ValidityIterator>
  fixed_width_column_wrapper(InputIterator begin, InputIterator end, ValidityIterator v)
    : column_wrapper{}
  {
    auto const size              = cudf::distance(begin, end);
    auto [null_mask, null_count] = detail::make_null_mask(v, v + size);
    wrapped.reset(new cudf::column{cudf::data_type{cudf::type_to_id<ElementTo>()},
                                   size,
                                   detail::make_elements<ElementTo, SourceElementT>(begin, end),
                                   std::move(null_mask),
                                   null_count});
  }
 
  template <typename ElementFrom>
  fixed_width_column_wrapper(std::initializer_list<ElementFrom> elements)
    : fixed_width_column_wrapper(std::cbegin(elements), std::cend(elements))
  {
  }
 
  template <typename ElementFrom>
  fixed_width_column_wrapper(std::initializer_list<ElementFrom> elements,
                             std::initializer_list<bool> validity)
    : fixed_width_column_wrapper(std::cbegin(elements), std::cend(elements), std::cbegin(validity))
  {
  }
 
  template <typename ValidityIterator, typename ElementFrom>
  fixed_width_column_wrapper(std::initializer_list<ElementFrom> element_list, ValidityIterator v)
    : fixed_width_column_wrapper(std::cbegin(element_list), std::cend(element_list), v)
  {
  }
 
  template <typename InputIterator>
  fixed_width_column_wrapper(InputIterator begin,
                             InputIterator end,
                             std::initializer_list<bool> const& validity)
    : fixed_width_column_wrapper(begin, end, std::cbegin(validity))
  {
  }
 
  template <typename ElementFrom>
  fixed_width_column_wrapper(std::initializer_list<std::pair<ElementFrom, bool>> elements)
  {
    auto begin =
      thrust::make_transform_iterator(elements.begin(), [](auto const& e) { return e.first; });
    auto end = begin + elements.size();
    auto v =
      thrust::make_transform_iterator(elements.begin(), [](auto const& e) { return e.second; });
    wrapped = fixed_width_column_wrapper<ElementTo, ElementFrom>(begin, end, v).release();
  }
};
 
template <typename Rep>
class fixed_point_column_wrapper : public detail::column_wrapper {
 public:
  template <typename FixedPointRepIterator>
  fixed_point_column_wrapper(FixedPointRepIterator begin,
                             FixedPointRepIterator end,
                             numeric::scale_type scale)
    : column_wrapper{}
  {
    CUDF_EXPECTS(numeric::is_supported_representation_type<Rep>(), "not valid representation type");
 
    auto const size      = cudf::distance(begin, end);
    auto const elements  = thrust::host_vector<Rep>(begin, end);
    auto const id        = type_to_id<numeric::fixed_point<Rep, numeric::Radix::BASE_10>>();
    auto const data_type = cudf::data_type{id, static_cast<int32_t>(scale)};
 
    wrapped.reset(new cudf::column{
      data_type,
      size,
      rmm::device_buffer{elements.data(), size * sizeof(Rep), cudf::test::get_default_stream()},
      rmm::device_buffer{},
      0});
  }
 
  fixed_point_column_wrapper(std::initializer_list<Rep> values, numeric::scale_type scale)
    : fixed_point_column_wrapper(std::cbegin(values), std::cend(values), scale)
  {
  }
 
  template <typename FixedPointRepIterator, typename ValidityIterator>
  fixed_point_column_wrapper(FixedPointRepIterator begin,
                             FixedPointRepIterator end,
                             ValidityIterator v,
                             numeric::scale_type scale)
    : column_wrapper{}
  {
    CUDF_EXPECTS(numeric::is_supported_representation_type<Rep>(), "not valid representation type");
 
    auto const size              = cudf::distance(begin, end);
    auto const elements          = thrust::host_vector<Rep>(begin, end);
    auto const id                = type_to_id<numeric::fixed_point<Rep, numeric::Radix::BASE_10>>();
    auto const data_type         = cudf::data_type{id, static_cast<int32_t>(scale)};
    auto [null_mask, null_count] = detail::make_null_mask(v, v + size);
    wrapped.reset(new cudf::column{
      data_type,
      size,
      rmm::device_buffer{elements.data(), size * sizeof(Rep), cudf::test::get_default_stream()},
      std::move(null_mask),
      null_count});
  }
 
  fixed_point_column_wrapper(std::initializer_list<Rep> elements,
                             std::initializer_list<bool> validity,
                             numeric::scale_type scale)
    : fixed_point_column_wrapper(
        std::cbegin(elements), std::cend(elements), std::cbegin(validity), scale)
  {
  }
 
  template <typename ValidityIterator>
  fixed_point_column_wrapper(std::initializer_list<Rep> element_list,
                             ValidityIterator v,
                             numeric::scale_type scale)
    : fixed_point_column_wrapper(std::cbegin(element_list), std::cend(element_list), v, scale)
  {
  }
 
  template <typename FixedPointRepIterator>
  fixed_point_column_wrapper(FixedPointRepIterator begin,
                             FixedPointRepIterator end,
                             std::initializer_list<bool> const& validity,
                             numeric::scale_type scale)
    : fixed_point_column_wrapper(begin, end, std::cbegin(validity), scale)
  {
  }
};
 
class strings_column_wrapper : public detail::column_wrapper {
 public:
  strings_column_wrapper() : strings_column_wrapper(std::initializer_list<std::string>{}) {}
 
  template <typename StringsIterator>
  strings_column_wrapper(StringsIterator begin, StringsIterator end) : column_wrapper{}
  {
    size_type num_strings = std::distance(begin, end);
    if (num_strings == 0) {
      wrapped = cudf::make_empty_column(cudf::type_id::STRING);
      return;
    }
    auto all_valid        = thrust::make_constant_iterator(true);
    auto [chars, offsets] = detail::make_chars_and_offsets(begin, end, all_valid);
    auto d_chars          = cudf::detail::make_device_uvector_async(
      chars, cudf::test::get_default_stream(), cudf::get_current_device_resource_ref());
    auto d_offsets = std::make_unique<cudf::column>(
      cudf::detail::make_device_uvector_sync(
        offsets, cudf::test::get_default_stream(), cudf::get_current_device_resource_ref()),
      rmm::device_buffer{},
      0);
    wrapped =
      cudf::make_strings_column(num_strings, std::move(d_offsets), d_chars.release(), 0, {});
  }
 
  template <typename StringsIterator, typename ValidityIterator>
  strings_column_wrapper(StringsIterator begin, StringsIterator end, ValidityIterator v)
    : column_wrapper{}
  {
    size_type num_strings = std::distance(begin, end);
    if (num_strings == 0) {
      wrapped = cudf::make_empty_column(cudf::type_id::STRING);
      return;
    }
    auto [chars, offsets]        = detail::make_chars_and_offsets(begin, end, v);
    auto [null_mask, null_count] = detail::make_null_mask_vector(v, v + num_strings);
    auto d_chars                 = cudf::detail::make_device_uvector_async(
      chars, cudf::test::get_default_stream(), cudf::get_current_device_resource_ref());
    auto d_offsets = std::make_unique<cudf::column>(
      cudf::detail::make_device_uvector_async(
        offsets, cudf::test::get_default_stream(), cudf::get_current_device_resource_ref()),
      rmm::device_buffer{},
      0);
    auto d_bitmask = cudf::detail::make_device_uvector_sync(
      null_mask, cudf::test::get_default_stream(), cudf::get_current_device_resource_ref());
    wrapped = cudf::make_strings_column(
      num_strings, std::move(d_offsets), d_chars.release(), null_count, d_bitmask.release());
  }
 
  strings_column_wrapper(std::initializer_list<std::string> strings)
    : strings_column_wrapper(std::cbegin(strings), std::cend(strings))
  {
  }
 
  template <typename ValidityIterator>
  strings_column_wrapper(std::initializer_list<std::string> strings, ValidityIterator v)
    : strings_column_wrapper(std::cbegin(strings), std::cend(strings), v)
  {
  }
 
  strings_column_wrapper(std::initializer_list<std::string> strings,
                         std::initializer_list<bool> validity)
    : strings_column_wrapper(std::cbegin(strings), std::cend(strings), std::cbegin(validity))
  {
  }
 
  strings_column_wrapper(std::initializer_list<std::pair<std::string, bool>> strings)
  {
    auto begin =
      thrust::make_transform_iterator(strings.begin(), [](auto const& s) { return s.first; });
    auto end = begin + strings.size();
    auto v =
      thrust::make_transform_iterator(strings.begin(), [](auto const& s) { return s.second; });
    wrapped = strings_column_wrapper(begin, end, v).release();
  }
};
 
template <typename KeyElementTo, typename SourceElementT = KeyElementTo>
class dictionary_column_wrapper : public detail::column_wrapper {
 public:
  operator dictionary_column_view() const { return cudf::dictionary_column_view{wrapped->view()}; }
 
  dictionary_column_wrapper() : column_wrapper{}
  {
    wrapped = cudf::make_empty_column(cudf::type_id::DICTIONARY32);
  }
 
  template <typename InputIterator>
  dictionary_column_wrapper(InputIterator begin, InputIterator end) : column_wrapper{}
  {
    wrapped =
      cudf::dictionary::encode(fixed_width_column_wrapper<KeyElementTo, SourceElementT>(begin, end),
                               cudf::data_type{type_id::UINT32},
                               cudf::test::get_default_stream());
  }
 
  template <typename InputIterator, typename ValidityIterator>
  dictionary_column_wrapper(InputIterator begin, InputIterator end, ValidityIterator v)
    : column_wrapper{}
  {
    wrapped = cudf::dictionary::encode(
      fixed_width_column_wrapper<KeyElementTo, SourceElementT>(begin, end, v),
      cudf::data_type{type_id::UINT32},
      cudf::test::get_default_stream());
  }
 
  template <typename ElementFrom>
  dictionary_column_wrapper(std::initializer_list<ElementFrom> elements)
    : dictionary_column_wrapper(std::cbegin(elements), std::cend(elements))
  {
  }
 
  template <typename ElementFrom>
  dictionary_column_wrapper(std::initializer_list<ElementFrom> elements,
                            std::initializer_list<bool> validity)
    : dictionary_column_wrapper(std::cbegin(elements), std::cend(elements), std::cbegin(validity))
  {
  }
 
  template <typename ValidityIterator, typename ElementFrom>
  dictionary_column_wrapper(std::initializer_list<ElementFrom> element_list, ValidityIterator v)
    : dictionary_column_wrapper(std::cbegin(element_list), std::cend(element_list), v)
  {
  }
 
  template <typename InputIterator>
  dictionary_column_wrapper(InputIterator begin,
                            InputIterator end,
                            std::initializer_list<bool> const& validity)
    : dictionary_column_wrapper(begin, end, std::cbegin(validity))
  {
  }
};
 
template <>
class dictionary_column_wrapper<std::string> : public detail::column_wrapper {
 public:
  operator dictionary_column_view() const { return cudf::dictionary_column_view{wrapped->view()}; }
 
  [[nodiscard]] column_view keys() const
  {
    return cudf::dictionary_column_view{wrapped->view()}.keys();
  }
 
  [[nodiscard]] column_view indices() const
  {
    return cudf::dictionary_column_view{wrapped->view()}.indices();
  }
 
  dictionary_column_wrapper() : dictionary_column_wrapper(std::initializer_list<std::string>{}) {}
 
  template <typename StringsIterator>
  dictionary_column_wrapper(StringsIterator begin, StringsIterator end) : column_wrapper{}
  {
    wrapped = cudf::dictionary::encode(strings_column_wrapper(begin, end),
                                       cudf::data_type{type_id::UINT32},
                                       cudf::test::get_default_stream());
  }
 
  template <typename StringsIterator, typename ValidityIterator>
  dictionary_column_wrapper(StringsIterator begin, StringsIterator end, ValidityIterator v)
    : column_wrapper{}
  {
    wrapped = cudf::dictionary::encode(strings_column_wrapper(begin, end, v),
                                       cudf::data_type{type_id::UINT32},
                                       cudf::test::get_default_stream());
  }
 
  dictionary_column_wrapper(std::initializer_list<std::string> strings)
    : dictionary_column_wrapper(std::cbegin(strings), std::cend(strings))
  {
  }
 
  template <typename ValidityIterator>
  dictionary_column_wrapper(std::initializer_list<std::string> strings, ValidityIterator v)
    : dictionary_column_wrapper(std::cbegin(strings), std::cend(strings), v)
  {
  }
 
  dictionary_column_wrapper(std::initializer_list<std::string> strings,
                            std::initializer_list<bool> validity)
    : dictionary_column_wrapper(std::cbegin(strings), std::cend(strings), std::cbegin(validity))
  {
  }
};
 
template <typename T, typename SourceElementT = T>
class lists_column_wrapper : public detail::column_wrapper {
 public:
  operator lists_column_view() const { return cudf::lists_column_view{wrapped->view()}; }
 
  template <typename Element = T, std::enable_if_t<cudf::is_fixed_width<Element>()>* = nullptr>
  lists_column_wrapper(std::initializer_list<SourceElementT> elements) : column_wrapper{}
  {
    build_from_non_nested(
      cudf::test::fixed_width_column_wrapper<T, SourceElementT>(elements).release());
  }
 
  template <typename Element = T,
            typename InputIterator,
            std::enable_if_t<cudf::is_fixed_width<Element>()>* = nullptr>
  lists_column_wrapper(InputIterator begin, InputIterator end) : column_wrapper{}
  {
    build_from_non_nested(
      cudf::test::fixed_width_column_wrapper<T, SourceElementT>(begin, end).release());
  }
 
  template <typename Element = T,
            typename ValidityIterator,
            std::enable_if_t<cudf::is_fixed_width<Element>()>* = nullptr>
  lists_column_wrapper(std::initializer_list<SourceElementT> elements, ValidityIterator v)
    : column_wrapper{}
  {
    build_from_non_nested(
      cudf::test::fixed_width_column_wrapper<T, SourceElementT>(elements, v).release());
  }
 
  template <typename Element = T,
            typename InputIterator,
            typename ValidityIterator,
            std::enable_if_t<cudf::is_fixed_width<Element>()>* = nullptr>
  lists_column_wrapper(InputIterator begin, InputIterator end, ValidityIterator v)
    : column_wrapper{}
  {
    build_from_non_nested(
      cudf::test::fixed_width_column_wrapper<T, SourceElementT>(begin, end, v).release());
  }
 
  template <typename Element                                              = T,
            std::enable_if_t<std::is_same_v<Element, cudf::string_view>>* = nullptr>
  lists_column_wrapper(std::initializer_list<std::string> elements) : column_wrapper{}
  {
    build_from_non_nested(
      cudf::test::strings_column_wrapper(elements.begin(), elements.end()).release());
  }
 
  template <typename Element = T,
            typename ValidityIterator,
            std::enable_if_t<std::is_same_v<Element, cudf::string_view>>* = nullptr>
  lists_column_wrapper(std::initializer_list<std::string> elements, ValidityIterator v)
    : column_wrapper{}
  {
    build_from_non_nested(
      cudf::test::strings_column_wrapper(elements.begin(), elements.end(), v).release());
  }
 
  lists_column_wrapper(std::initializer_list<lists_column_wrapper<T, SourceElementT>> elements)
    : column_wrapper{}
  {
    std::vector<bool> valids;
    build_from_nested(elements, valids);
  }
 
  lists_column_wrapper() : column_wrapper{}
  {
    build_from_non_nested(make_empty_column(cudf::type_to_id<T>()));
  }
 
  template <typename ValidityIterator>
  lists_column_wrapper(std::initializer_list<lists_column_wrapper<T, SourceElementT>> elements,
                       ValidityIterator v)
    : column_wrapper{}
  {
    std::vector<bool> validity;
    std::transform(elements.begin(),
                   elements.end(),
                   v,
                   std::back_inserter(validity),
                   [](lists_column_wrapper const& l, bool valid) { return valid; });
    build_from_nested(elements, validity);
  }
 
  static lists_column_wrapper<T> make_one_empty_row_column(bool valid = true)
  {
    cudf::test::fixed_width_column_wrapper<cudf::size_type> offsets{0, 0};
    cudf::test::fixed_width_column_wrapper<int> values{};
    return lists_column_wrapper<T>(
      1,
      offsets.release(),
      values.release(),
      valid ? 0 : 1,
      valid ? rmm::device_buffer{} : cudf::create_null_mask(1, cudf::mask_state::ALL_NULL));
  }
 
 private:
  lists_column_wrapper(size_type num_rows,
                       std::unique_ptr<cudf::column>&& offsets,
                       std::unique_ptr<cudf::column>&& values,
                       size_type null_count,
                       rmm::device_buffer&& null_mask)
  {
    // construct the list column
    wrapped = make_lists_column(num_rows,
                                std::move(offsets),
                                std::move(values),
                                null_count,
                                std::move(null_mask),
                                cudf::test::get_default_stream());
  }
 
  void build_from_nested(std::initializer_list<lists_column_wrapper<T, SourceElementT>> elements,
                         std::vector<bool> const& v)
  {
    auto valids = cudf::detail::make_counting_transform_iterator(
      0, [&v](auto i) { return v.empty() ? true : v[i]; });
 
    // compute the expected hierarchy and depth
    auto const hierarchy_and_depth =
      std::accumulate(elements.begin(),
                      elements.end(),
                      std::pair<column_view, int32_t>{{}, -1},
                      [](auto acc, lists_column_wrapper const& lcw) {
                        return lcw.depth > acc.second ? std::pair(lcw.get_view(), lcw.depth) : acc;
                      });
    column_view expected_hierarchy = hierarchy_and_depth.first;
    int32_t const expected_depth   = hierarchy_and_depth.second;
 
    // preprocess columns so that every column_view in 'cols' is an equivalent hierarchy
    auto [cols, stubs] = preprocess_columns(elements, expected_hierarchy, expected_depth);
 
    // generate offsets
    size_type count = 0;
    std::vector<size_type> offsetv;
    std::transform(cols.cbegin(),
                   cols.cend(),
                   valids,
                   std::back_inserter(offsetv),
                   [&](cudf::column_view const& col, bool valid) {
                     // nulls are represented as a repeated offset
                     size_type ret = count;
                     if (valid) { count += col.size(); }
                     return ret;
                   });
    // add the final offset
    offsetv.push_back(count);
    auto offsets =
      cudf::test::fixed_width_column_wrapper<size_type>(offsetv.begin(), offsetv.end()).release();
 
    // concatenate them together, skipping children that are null.
    std::vector<column_view> children;
    thrust::copy_if(
      std::cbegin(cols), std::cend(cols), valids, std::back_inserter(children), thrust::identity{});
 
    auto data = children.empty() ? cudf::empty_like(expected_hierarchy)
                                 : cudf::concatenate(children,
                                                     cudf::test::get_default_stream(),
                                                     cudf::get_current_device_resource_ref());
 
    // increment depth
    depth = expected_depth + 1;
 
    auto [null_mask, null_count] = [&] {
      if (v.size() <= 0) return std::make_pair(rmm::device_buffer{}, cudf::size_type{0});
      return cudf::test::detail::make_null_mask(v.begin(), v.end());
    }();
 
    // construct the list column
    wrapped = make_lists_column(cols.size(),
                                std::move(offsets),
                                std::move(data),
                                null_count,
                                std::move(null_mask),
                                cudf::test::get_default_stream());
  }
 
  void build_from_non_nested(std::unique_ptr<column> c)
  {
    CUDF_EXPECTS(c->type().id() == type_id::EMPTY || !cudf::is_nested(c->type()),
                 "Unexpected type");
 
    std::vector<size_type> offsetv;
    if (c->size() > 0) {
      offsetv.push_back(0);
      offsetv.push_back(c->size());
    }
    auto offsets =
      cudf::test::fixed_width_column_wrapper<size_type>(offsetv.begin(), offsetv.end()).release();
 
    // construct the list column. mark this as a root
    root  = true;
    depth = 0;
 
    size_type num_elements = offsets->size() == 0 ? 0 : offsets->size() - 1;
    wrapped                = make_lists_column(num_elements,
                                std::move(offsets),
                                std::move(c),
                                0,
                                rmm::device_buffer{},
                                cudf::test::get_default_stream());
  }
 
  std::unique_ptr<column> normalize_column(column_view const& col,
                                           column_view const& expected_hierarchy)
  {
    // if are at the bottom of the short column, it must be empty
    if (col.type().id() != type_id::LIST) {
      CUDF_EXPECTS(col.is_empty(), "Encountered mismatched column!");
 
      auto remainder = empty_like(expected_hierarchy);
      return remainder;
    }
 
    lists_column_view lcv(col);
    return make_lists_column(
      col.size(),
      std::make_unique<column>(lcv.offsets()),
      normalize_column(lists_column_view(col).child(),
                       lists_column_view(expected_hierarchy).child()),
      col.null_count(),
      cudf::copy_bitmask(
        col, cudf::test::get_default_stream(), cudf::get_current_device_resource_ref()),
      cudf::test::get_default_stream());
  }
 
  std::pair<std::vector<column_view>, std::vector<std::unique_ptr<column>>> preprocess_columns(
    std::initializer_list<lists_column_wrapper<T, SourceElementT>> const& elements,
    column_view& expected_hierarchy,
    int expected_depth)
  {
    std::vector<std::unique_ptr<column>> stubs;
    std::vector<column_view> cols;
 
    // preprocess the incoming lists.
    // - unwrap any "root" lists
    // - handle incomplete hierarchies
    std::transform(elements.begin(),
                   elements.end(),
                   std::back_inserter(cols),
                   [&](lists_column_wrapper const& l) -> column_view {
                     // depth mismatch.  attempt to normalize the short column.
                     // this function will also catch if this is a legitimately broken
                     // set of input
                     if (l.depth < expected_depth) {
                       if (l.root) {
                         // this exception distinguishes between the following two cases:
                         //
                         // { {{{1, 2, 3}}}, {} }
                         // In this case, row 0 is a List<List<List<int>>>, whereas row 1 is
                         // just a List<> which is an apparent mismatch.  However, because row 1
                         // is empty we will allow that to semantically mean
                         // "a List<List<List<int>>> that's empty at the top level"
                         //
                         // { {{{1, 2, 3}}}, {4, 5, 6} }
                         // In this case, row 1 is a concrete List<int> with actual values.
                         // There is no way to rectify the differences so we will treat it as a
                         // true column mismatch.
                         CUDF_EXPECTS(l.wrapped->size() == 0, "Mismatch in column types!");
                         stubs.push_back(empty_like(expected_hierarchy));
                       } else {
                         stubs.push_back(normalize_column(l.get_view(), expected_hierarchy));
                       }
                       return *(stubs.back());
                     }
                     // the empty hierarchy case
                     return l.get_view();
                   });
 
    return {std::move(cols), std::move(stubs)};
  }
 
  [[nodiscard]] column_view get_view() const
  {
    return root ? lists_column_view(*wrapped).child() : *wrapped;
  }
 
  int depth = 0;
  bool root = false;
};
 
class structs_column_wrapper : public detail::column_wrapper {
 public:
  structs_column_wrapper(std::vector<std::unique_ptr<cudf::column>>&& child_columns,
                         std::vector<bool> const& validity = {})
  {
    init(std::move(child_columns), validity);
  }
 
  structs_column_wrapper(
    std::initializer_list<std::reference_wrapper<detail::column_wrapper>> child_column_wrappers,
    std::vector<bool> const& validity = {})
  {
    std::vector<std::unique_ptr<cudf::column>> child_columns;
    child_columns.reserve(child_column_wrappers.size());
    std::transform(child_column_wrappers.begin(),
                   child_column_wrappers.end(),
                   std::back_inserter(child_columns),
                   [&](auto const& column_wrapper) {
                     return std::make_unique<cudf::column>(column_wrapper.get(),
                                                           cudf::test::get_default_stream());
                   });
    init(std::move(child_columns), validity);
  }
 
  template <typename V>
  structs_column_wrapper(
    std::initializer_list<std::reference_wrapper<detail::column_wrapper>> child_column_wrappers,
    V validity_iter)
  {
    std::vector<std::unique_ptr<cudf::column>> child_columns;
    child_columns.reserve(child_column_wrappers.size());
    std::transform(child_column_wrappers.begin(),
                   child_column_wrappers.end(),
                   std::back_inserter(child_columns),
                   [&](auto const& column_wrapper) {
                     return std::make_unique<cudf::column>(column_wrapper.get(),
                                                           cudf::test::get_default_stream());
                   });
    init(std::move(child_columns), validity_iter);
  }
 
 private:
  void init(std::vector<std::unique_ptr<cudf::column>>&& child_columns,
            std::vector<bool> const& validity)
  {
    size_type num_rows = child_columns.empty() ? 0 : child_columns[0]->size();
 
    CUDF_EXPECTS(std::all_of(child_columns.begin(),
                             child_columns.end(),
                             [&](auto const& p_column) { return p_column->size() == num_rows; }),
                 "All struct member columns must have the same row count.");
 
    CUDF_EXPECTS(validity.size() <= 0 || static_cast<size_type>(validity.size()) == num_rows,
                 "Validity buffer must have as many elements as rows in the struct column.");
 
    auto [null_mask, null_count] = [&] {
      if (validity.size() <= 0) return std::make_pair(rmm::device_buffer{}, cudf::size_type{0});
      return cudf::test::detail::make_null_mask(validity.begin(), validity.end());
    }();
 
    wrapped = cudf::make_structs_column(num_rows,
                                        std::move(child_columns),
                                        null_count,
                                        std::move(null_mask),
                                        cudf::test::get_default_stream());
  }
 
  template <typename V>
  void init(std::vector<std::unique_ptr<cudf::column>>&& child_columns, V validity_iterator)
  {
    size_type const num_rows = child_columns.empty() ? 0 : child_columns[0]->size();
 
    CUDF_EXPECTS(std::all_of(child_columns.begin(),
                             child_columns.end(),
                             [&](auto const& p_column) { return p_column->size() == num_rows; }),
                 "All struct member columns must have the same row count.");
 
    std::vector<bool> validity(num_rows);
    std::copy(validity_iterator, validity_iterator + num_rows, validity.begin());
 
    init(std::move(child_columns), validity);
  }
};
 
}  // namespace test
}  // namespace CUDF_EXPORT cudf