Cleaning up for loops with make_(counting_)transform_iterator (#6546

)

Doing some cleanup - using `make_counting_transform_iterator` and `make_transform_iterator` to clean up `for` loops and other code.

Making an argument that `make_counting_transform_iterator` should be moved out of `cudf::test::`.

Also, will clean up the `f`/`begin` and `op` naming conventions.

Depends on: #7306

Authors:
  - Conor Hoekstra (@codereport)

Approvers:
  - David (@davidwendt)
  - Jake Hemstad (@jrhemstad)
  - Mark Harris (@harrism)

URL: #6546

cpp/include/cudf/column/column_view.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

cpp/src/column/column.cu

@@ -32,6 +32,8 @@
 #include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
+#include <thrust/iterator/transform_iterator.h>
+
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 
@@ -221,10 +223,9 @@ struct create_column_from_view {
   template <typename ColumnType, std::enable_if_t<cudf::is_fixed_width<ColumnType>()> * = nullptr>
   std::unique_ptr<column> operator()()
   {
-    std::vector<std::unique_ptr<column>> children;
-    for (size_type i = 0; i < view.num_children(); ++i) {
-      children.emplace_back(std::make_unique<column>(view.child(i), stream, mr));
-    }
+    auto op       = [&](auto const &child) { return std::make_unique<column>(child, stream, mr); };
+    auto begin    = thrust::make_transform_iterator(view.child_begin(), op);
+    auto children = std::vector<std::unique_ptr<column>>(begin, begin + view.num_children());
 
     return std::make_unique<column>(
       view.type(),

cpp/src/column/column_device_view.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,6 +15,7 @@
  */
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/column/column_view.hpp>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
 
@@ -44,12 +45,10 @@ std::unique_ptr<ColumnDeviceView, std::function<void(ColumnDeviceView*)>>
 create_device_view_from_view(ColumnView const& source, rmm::cuda_stream_view stream)
 {
   size_type num_children = source.num_children();
-  // First calculate the size of memory needed to hold the
-  // child columns. This is done by calling extent()
-  // for each of the children.
-  auto get_extent = thrust::make_transform_iterator(
-    thrust::make_counting_iterator(0),
-    [&source](auto i) { return ColumnDeviceView::extent(source.child(i)); });
+  // First calculate the size of memory needed to hold the child columns. This is done by calling
+  // extent() for each of the children.
+  auto get_extent = cudf::detail::make_counting_transform_iterator(
+    0, [&source](auto i) { return ColumnDeviceView::extent(source.child(i)); });
 
   // pad the allocation for aligning the first pointer
   auto const descendant_storage_bytes = std::accumulate(

cpp/src/column/column_view.cpp

@@ -20,7 +20,10 @@
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/traits.hpp>
 
+#include <thrust/iterator/transform_iterator.h>
+
 #include <exception>
+#include <numeric>
 #include <vector>
 
 namespace cudf {
@@ -124,11 +127,9 @@ mutable_column_view::operator column_view() const
 
 size_type count_descendants(column_view parent)
 {
-  size_type count{parent.num_children()};
-  for (size_type i = 0; i < parent.num_children(); ++i) {
-    count += count_descendants(parent.child(i));
-  }
-  return count;
+  auto descendants = [](auto const& child) { return count_descendants(child); };
+  auto begin       = thrust::make_transform_iterator(parent.child_begin(), descendants);
+  return std::accumulate(begin, begin + parent.num_children(), size_type{parent.num_children()});
 }
 
 column_view logical_cast(column_view const& input, data_type type)

cpp/src/copying/copy.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,6 +26,8 @@
 
 #include <rmm/cuda_stream_view.hpp>
 
+#include <thrust/iterator/transform_iterator.h>
+
 #include <algorithm>
 
 namespace cudf {
@@ -58,11 +60,9 @@ std::unique_ptr<column> allocate_like(column_view const& input,
   CUDF_EXPECTS(is_fixed_width(input.type()), "Expects only fixed-width type column");
   mask_state allocate_mask = should_allocate_mask(mask_alloc, input.nullable());
 
-  std::vector<std::unique_ptr<column>> children{};
-  children.reserve(input.num_children());
-  for (size_type index = 0; index < input.num_children(); index++) {
-    children.emplace_back(allocate_like(input.child(index), size, mask_alloc, stream, mr));
-  }
+  auto op = [&](auto const& child) { return allocate_like(child, size, mask_alloc, stream, mr); };
+  auto begin = thrust::make_transform_iterator(input.child_begin(), op);
+  std::vector<std::unique_ptr<column>> children(begin, begin + input.num_children());
 
   return std::make_unique<column>(input.type(),
                                   size,

cpp/src/copying/slice.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 #include <cudf/column/column_view.hpp>
 #include <cudf/copying.hpp>
 #include <cudf/detail/copy.hpp>
+#include <cudf/detail/iterator.cuh>
 #include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/utilities/error.hpp>
@@ -33,30 +34,27 @@ std::vector<column_view> slice(column_view const& input,
 {
   CUDF_EXPECTS(indices.size() % 2 == 0, "indices size must be even");
 
-  if (indices.empty()) { return {}; }
+  if (indices.empty()) return {};
 
   auto null_counts = cudf::detail::segmented_count_unset_bits(input.null_mask(), indices, stream);
+  auto const children = std::vector<column_view>(input.child_begin(), input.child_end());
 
-  std::vector<column_view> children{};
-  for (size_type i = 0; i < input.num_children(); i++) { children.push_back(input.child(i)); }
-
-  std::vector<column_view> result{};
-  for (size_t i = 0; i < indices.size() / 2; i++) {
+  auto op = [&](auto i) {
     auto begin = indices[2 * i];
     auto end   = indices[2 * i + 1];
     CUDF_EXPECTS(begin >= 0, "Starting index cannot be negative.");
     CUDF_EXPECTS(end >= begin, "End index cannot be smaller than the starting index.");
     CUDF_EXPECTS(end <= input.size(), "Slice range out of bounds.");
-    result.emplace_back(input.type(),
-                        end - begin,
-                        input.head(),
-                        input.null_mask(),
-                        null_counts[i],
-                        input.offset() + begin,
-                        children);
-  }
-
-  return result;
+    return column_view{input.type(),
+                       end - begin,
+                       input.head(),
+                       input.null_mask(),
+                       null_counts[i],
+                       input.offset() + begin,
+                       children};
+  };
+  auto begin = cudf::detail::make_counting_transform_iterator(0, op);
+  return std::vector<column_view>{begin, begin + indices.size() / 2};
 }
 
 }  // namespace detail
@@ -75,15 +73,13 @@ std::vector<cudf::table_view> slice(cudf::table_view const& input,
   CUDF_EXPECTS(indices.size() % 2 == 0, "indices size must be even");
   if (indices.empty()) { return {}; }
 
-  // 2d arrangement of column_views that represent the outgoing table_views
-  // sliced_table[i][j]
+  // 2d arrangement of column_views that represent the outgoing table_views sliced_table[i][j]
   // where i is the i'th column of the j'th table_view
-  std::vector<std::vector<cudf::column_view>> sliced_table;
+  auto op = [&indices](auto const& c) { return cudf::slice(c, indices); };
+  auto f  = thrust::make_transform_iterator(input.begin(), op);
+
+  auto sliced_table = std::vector<std::vector<cudf::column_view>>(f, f + input.num_columns());
   sliced_table.reserve(indices.size() + 1);
-  std::transform(input.begin(),
-                 input.end(),
-                 std::back_inserter(sliced_table),
-                 [&indices](cudf::column_view const& c) { return cudf::slice(c, indices); });
 
   std::vector<cudf::table_view> result{};
   // distribute columns into outgoing table_views

cpp/src/interop/to_arrow.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -101,17 +101,15 @@ struct dispatch_to_arrow {
     rmm::cuda_stream_view stream)
   {
     std::vector<std::shared_ptr<arrow::Array>> child_arrays;
-    std::vector<size_type> child_indices(input_view.num_children());
-    std::iota(child_indices.begin(), child_indices.end(), 0);
-    std::transform(child_indices.begin(),
-                   child_indices.end(),
-                   metadata.begin(),
-                   std::back_inserter(child_arrays),
-                   [&input_view, &ar_mr, &stream](auto const& i, auto const& meta) {
-                     auto c = input_view.child(i);
-                     return type_dispatcher(
-                       c.type(), dispatch_to_arrow{}, c, c.type().id(), meta, ar_mr, stream);
-                   });
+    std::transform(
+      input_view.child_begin(),
+      input_view.child_end(),
+      metadata.begin(),
+      std::back_inserter(child_arrays),
+      [&ar_mr, &stream](auto const& child, auto const& meta) {
+        return type_dispatcher(
+          child.type(), dispatch_to_arrow{}, child, child.type().id(), meta, ar_mr, stream);
+      });
     return child_arrays;
   }