column_factories.hpp

/*
 * Copyright (c) 2023-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 <cuspatial/types.hpp>
 
#include <cudf/column/column_factories.hpp>
#include <cudf/filling.hpp>
#include <cudf/scalar/scalar_factories.hpp>
#include <cudf_test/column_wrapper.hpp>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <initializer_list>
#include <memory>
 
namespace cuspatial {
namespace test {
 
using namespace cudf;
using namespace cudf::test;
 
std::unique_ptr<column> coords_offsets(size_type num_points, rmm::cuda_stream_view stream)
{
  auto zero = make_fixed_width_scalar<size_type>(0, stream);
  auto two  = make_fixed_width_scalar<size_type>(2, stream);
 
  return sequence(num_points + 1, *zero, *two);
}
 
template <typename T>
std::unique_ptr<column> make_non_nullable_lists_column(std::unique_ptr<column> offset,
                                                       std::unique_ptr<column> child)
{
  auto size = offset->size() - 1;
  return make_lists_column(size, std::move(offset), std::move(child), 0, {});
}
 
template <typename T>
std::unique_ptr<column> make_non_nullable_lists_column(std::initializer_list<size_type> offset,
                                                       std::unique_ptr<column> child)
{
  auto d_offset = fixed_width_column_wrapper<size_type>(offset).release();
  return make_non_nullable_lists_column<T>(std::move(d_offset), std::move(child));
}
 
template <typename T>
std::unique_ptr<column> make_non_nullable_lists_column(std::unique_ptr<column> offset,
                                                       std::initializer_list<T> child)
{
  auto d_child = fixed_width_column_wrapper<T>(child).release();
  return make_non_nullable_lists_column<T>(std::move(offset), std::move(d_child));
}
 
template <typename T>
std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_point_column(
  std::initializer_list<T>&& point_coords, rmm::cuda_stream_view stream)
{
  auto num_points = point_coords.size() / 2;
 
  return {collection_type_id::SINGLE,
          make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), point_coords)};
}
 
template <typename T>
std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_point_column(
  std::initializer_list<cudf::size_type>&& multipoint_offsets,
  std::initializer_list<T> point_coords,
  rmm::cuda_stream_view stream)
{
  auto num_points = point_coords.size() / 2;
 
  return {collection_type_id::MULTI,
          make_non_nullable_lists_column<T>(
            multipoint_offsets,
            make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), point_coords))};
}
 
template <typename T>
std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_linestring_column(
  std::initializer_list<cudf::size_type>&& linestring_offsets,
  std::initializer_list<T>&& linestring_coords,
  rmm::cuda_stream_view stream)
{
  auto num_points = linestring_coords.size() / 2;
 
  return {
    collection_type_id::SINGLE,
    make_non_nullable_lists_column<T>(
      linestring_offsets,
      make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), linestring_coords))};
}
 
template <typename T>
std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_linestring_column(
  std::initializer_list<cudf::size_type>&& multilinestring_offsets,
  std::initializer_list<cudf::size_type>&& linestring_offsets,
  std::initializer_list<T> linestring_coords,
  rmm::cuda_stream_view stream)
{
  auto num_points = linestring_coords.size() / 2;
  return {
    collection_type_id::MULTI,
    make_non_nullable_lists_column<T>(
      multilinestring_offsets,
      make_non_nullable_lists_column<T>(
        linestring_offsets,
        make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), linestring_coords)))};
}
 
template <typename T>
std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_polygon_column(
  std::initializer_list<cudf::size_type>&& polygon_offsets,
  std::initializer_list<cudf::size_type>&& ring_offsets,
  std::initializer_list<T> polygon_coords,
  rmm::cuda_stream_view stream)
{
  auto num_points = polygon_coords.size() / 2;
  return {
    collection_type_id::SINGLE,
    make_non_nullable_lists_column<T>(
      polygon_offsets,
      make_non_nullable_lists_column<T>(
        ring_offsets,
        make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), polygon_coords)))};
}
 
template <typename T>
std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_polygon_column(
  std::initializer_list<cudf::size_type>&& multipolygon_offsets,
  std::initializer_list<cudf::size_type>&& polygon_offsets,
  std::initializer_list<cudf::size_type>&& ring_offsets,
  std::initializer_list<T> polygon_coords,
  rmm::cuda_stream_view stream)
{
  auto num_points = polygon_coords.size() / 2;
  return {
    collection_type_id::MULTI,
    make_non_nullable_lists_column<T>(
      multipolygon_offsets,
      make_non_nullable_lists_column<T>(
        polygon_offsets,
        make_non_nullable_lists_column<T>(
          ring_offsets,
          make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), polygon_coords))))};
}
 
}  // namespace test
}  // namespace cuspatial