Loading...
Searching...
No Matches
column_factories.hpp
1/*
2 * Copyright (c) 2023-2024, NVIDIA CORPORATION.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#pragma once
18
19#include <cuspatial/types.hpp>
20
21#include <cudf/column/column_factories.hpp>
22#include <cudf/filling.hpp>
23#include <cudf/scalar/scalar_factories.hpp>
24#include <cudf_test/column_wrapper.hpp>
25
26#include <rmm/cuda_stream_view.hpp>
27
28#include <initializer_list>
29#include <memory>
30
31namespace cuspatial {
32namespace test {
33
34using namespace cudf;
35using namespace cudf::test;
36
37std::unique_ptr<column> coords_offsets(size_type num_points, rmm::cuda_stream_view stream)
38{
39 auto zero = make_fixed_width_scalar<size_type>(0, stream);
40 auto two = make_fixed_width_scalar<size_type>(2, stream);
41
42 return sequence(num_points + 1, *zero, *two);
43}
44
45template <typename T>
46std::unique_ptr<column> make_non_nullable_lists_column(std::unique_ptr<column> offset,
47 std::unique_ptr<column> child)
48{
49 auto size = offset->size() - 1;
50 return make_lists_column(size, std::move(offset), std::move(child), 0, {});
51}
52
53template <typename T>
54std::unique_ptr<column> make_non_nullable_lists_column(std::initializer_list<size_type> offset,
55 std::unique_ptr<column> child)
56{
57 auto d_offset = fixed_width_column_wrapper<size_type>(offset).release();
58 return make_non_nullable_lists_column<T>(std::move(d_offset), std::move(child));
59}
60
61template <typename T>
62std::unique_ptr<column> make_non_nullable_lists_column(std::unique_ptr<column> offset,
63 std::initializer_list<T> child)
64{
65 auto d_child = fixed_width_column_wrapper<T>(child).release();
66 return make_non_nullable_lists_column<T>(std::move(offset), std::move(d_child));
67}
68
86template <typename T>
87std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_point_column(
88 std::initializer_list<T>&& point_coords, rmm::cuda_stream_view stream)
89{
90 auto num_points = point_coords.size() / 2;
91
92 return {collection_type_id::SINGLE,
93 make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), point_coords)};
94}
95
114template <typename T>
115std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_point_column(
116 std::initializer_list<cudf::size_type>&& multipoint_offsets,
117 std::initializer_list<T> point_coords,
118 rmm::cuda_stream_view stream)
119{
120 auto num_points = point_coords.size() / 2;
121
122 return {collection_type_id::MULTI,
123 make_non_nullable_lists_column<T>(
124 multipoint_offsets,
125 make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), point_coords))};
126}
127
146template <typename T>
147std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_linestring_column(
148 std::initializer_list<cudf::size_type>&& linestring_offsets,
149 std::initializer_list<T>&& linestring_coords,
150 rmm::cuda_stream_view stream)
151{
152 auto num_points = linestring_coords.size() / 2;
153
154 return {
155 collection_type_id::SINGLE,
156 make_non_nullable_lists_column<T>(
157 linestring_offsets,
158 make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), linestring_coords))};
159}
160
184template <typename T>
185std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_linestring_column(
186 std::initializer_list<cudf::size_type>&& multilinestring_offsets,
187 std::initializer_list<cudf::size_type>&& linestring_offsets,
188 std::initializer_list<T> linestring_coords,
189 rmm::cuda_stream_view stream)
190{
191 auto num_points = linestring_coords.size() / 2;
192 return {
193 collection_type_id::MULTI,
194 make_non_nullable_lists_column<T>(
195 multilinestring_offsets,
196 make_non_nullable_lists_column<T>(
197 linestring_offsets,
198 make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), linestring_coords)))};
199}
200
224template <typename T>
225std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_polygon_column(
226 std::initializer_list<cudf::size_type>&& polygon_offsets,
227 std::initializer_list<cudf::size_type>&& ring_offsets,
228 std::initializer_list<T> polygon_coords,
229 rmm::cuda_stream_view stream)
230{
231 auto num_points = polygon_coords.size() / 2;
232 return {
233 collection_type_id::SINGLE,
234 make_non_nullable_lists_column<T>(
235 polygon_offsets,
236 make_non_nullable_lists_column<T>(
237 ring_offsets,
238 make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), polygon_coords)))};
239}
240
266template <typename T>
267std::pair<collection_type_id, std::unique_ptr<cudf::column>> make_polygon_column(
268 std::initializer_list<cudf::size_type>&& multipolygon_offsets,
269 std::initializer_list<cudf::size_type>&& polygon_offsets,
270 std::initializer_list<cudf::size_type>&& ring_offsets,
271 std::initializer_list<T> polygon_coords,
272 rmm::cuda_stream_view stream)
273{
274 auto num_points = polygon_coords.size() / 2;
275 return {
276 collection_type_id::MULTI,
277 make_non_nullable_lists_column<T>(
278 multipolygon_offsets,
279 make_non_nullable_lists_column<T>(
280 polygon_offsets,
281 make_non_nullable_lists_column<T>(
282 ring_offsets,
283 make_non_nullable_lists_column<T>(coords_offsets(num_points, stream), polygon_coords))))};
284}
285
286} // namespace test
287} // namespace cuspatial