Cleaning up for loops with make_(counting_)transform_iterator

CHANGELOG.md

@@ -13,6 +13,7 @@
 - PR #6485 Add File IO to cuIO benchmarks
 - PR #6504 Update Java bindings version to 0.17-SNAPSHOT
 - PR #6527 Refactor DeviceColumnViewAccess to avoid JNI returning an array
+- RP #6546 Cleaning up `for` loops with `make_counting_transform_iterator`
 
 ## Bug Fixes
 - PR #6506 Fix DateTime type value truncation while writing to csv

cpp/src/column/column_view.cpp

@@ -20,6 +20,8 @@
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/traits.hpp>
 
+#include <cudf_test/column_wrapper.hpp>
+
 #include <exception>
 #include <vector>
 
@@ -124,11 +126,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 op    = [&](auto i) { return count_descendants(parent.child(i)); };
+  auto begin = cudf::test::make_counting_transform_iterator(0, op);
+  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

@@ -25,6 +25,7 @@
 #include <cudf/utilities/traits.hpp>
 
 #include <algorithm>
+#include "cudf_test/column_wrapper.hpp"
 
 namespace cudf {
 namespace detail {
@@ -56,11 +57,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, mr, stream));
-  }
+  auto op    = [&](auto i) { return allocate_like(input.child(i), size, mask_alloc, mr, stream); };
+  auto begin = cudf::test::make_counting_transform_iterator(0, op);
+  auto children = std::vector<std::unique_ptr<column>>(begin, begin + input.num_children());
 
   return std::make_unique<column>(input.type(),
                                   size,

cpp/src/copying/slice.cpp

@@ -22,6 +22,8 @@
 #include <cudf/utilities/error.hpp>
 
 #include <algorithm>
+#include "cudf_test/column_wrapper.hpp"
+#include "thrust/iterator/transform_iterator.h"
 
 namespace cudf {
 namespace detail {
@@ -31,30 +33,29 @@ 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);
 
-  std::vector<column_view> children{};
-  for (size_type i = 0; i < input.num_children(); i++) { children.push_back(input.child(i)); }
+  auto f = cudf::test::make_counting_transform_iterator(0, [&](auto i) { return input.child(i); });
+  auto const children = std::vector<column_view>(f, f + input.num_children());
 
-  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::test::make_counting_transform_iterator(0, op);
+  return {begin, begin + indices.size() / 2};
 }
 
 }  // namespace detail
@@ -73,15 +74,12 @@ 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;
-  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); });
+  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 + indices.size() + 1);
 
   std::vector<cudf::table_view> result{};
   // distribute columns into outgoing table_views

cpp/src/interop/to_arrow.cpp

@@ -28,6 +28,7 @@
 #include <cudf/utilities/type_dispatcher.hpp>
 
 #include <rmm/mr/device/per_device_resource.hpp>
+#include "cudf_test/column_wrapper.hpp"
 
 namespace cudf {
 namespace detail {
@@ -97,18 +98,12 @@ struct dispatch_to_arrow {
                                                                arrow::MemoryPool* ar_mr,
                                                                cudaStream_t 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(),
-                   std::back_inserter(child_arrays),
-                   [&input_view, &ar_mr, &stream](auto const& i) {
-                     auto c = input_view.child(i);
-                     return type_dispatcher(
-                       c.type(), dispatch_to_arrow{}, c, c.type().id(), ar_mr, stream);
-                   });
-    return child_arrays;
+    auto op = [&input_view, &ar_mr, &stream](auto i) -> std::shared_ptr<arrow::Array> {
+      auto c = input_view.child(i);
+      return type_dispatcher(c.type(), dispatch_to_arrow{}, c, c.type().id(), ar_mr, stream);
+    };
+    auto begin = cudf::test::make_counting_transform_iterator(0, op);
+    return {begin, begin + input_view.num_children()};
   }
 
   template <typename T>