[cudf] Adding memory_resources support for column wrappers

cpp/include/cudf_test/memory_resource_utilities.hpp

@@ -5,6 +5,7 @@
 
 #pragma once
 
+#include <cudf_test/column_wrapper.hpp>
 #include <cudf_test/cudf_gtest.hpp>
 #include <cudf_test/default_stream.hpp>
 
@@ -17,6 +18,7 @@
 #include <cuda/memory_resource>
 #include <cuda/stream_ref>
 
+#include <concepts>
 #include <cstddef>
 #include <functional>
 #include <utility>
@@ -154,6 +156,24 @@ class memory_resource_test_harness {
   rmm::mr::callback_memory_resource _failing_mr;
 };
 
+/**
+ * @brief Callable that accepts a statistics resource and returns a column wrapper.
+ */
+template <typename Factory>
+concept column_wrapper_statistics_resource_factory =
+  requires(Factory& factory, rmm::mr::statistics_resource_adaptor& mr) {
+    { std::invoke(factory, mr) } -> std::derived_from<detail::column_wrapper>;
+  };
+
+/**
+ * @brief Callable that accepts `cudf::memory_resources` and returns a column wrapper.
+ */
+template <typename Factory>
+concept column_wrapper_memory_resources_factory =
+  requires(Factory& factory, cudf::memory_resources mr) {
+    { std::invoke(factory, mr) } -> std::derived_from<detail::column_wrapper>;
+  };
+
 /**
  * @brief Verify that an owning result uses one explicitly supplied output resource.
  *
@@ -169,7 +189,7 @@ class memory_resource_test_harness {
  * @param output_expectation Expected relationship between live and total output bytes
  * @param stream Stream to synchronize before inspecting allocation counters
  */
-template <typename Factory>
+template <column_wrapper_statistics_resource_factory Factory>
 void expect_output_uses_resource(
   Factory&& factory,
   output_allocation_expectation output_expectation = output_allocation_expectation::EXACT,
@@ -207,7 +227,7 @@ void expect_output_uses_resource(
  * @param expectations Expected output and temporary allocation behavior
  * @param stream Stream to synchronize before inspecting allocation counters
  */
-template <typename Factory>
+template <column_wrapper_memory_resources_factory Factory>
 void expect_output_uses_distinct_resources(
   Factory&& factory,
   memory_resource_expectations expectations = {},

cpp/include/cudf_test/timestamp_utilities.cuh

@@ -1,11 +1,12 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
 #include <cudf_test/column_wrapper.hpp>
+#include <cudf_test/default_stream.hpp>
 
 #include <cudf/detail/iterator.cuh>
 #include <cudf/utilities/export.hpp>
@@ -32,11 +33,16 @@ using time_point_ms =
  * @param count The number of timestamps to create
  * @param start The first timestamp as a cuda::std::chrono::time_point
  * @param stop The last timestamp as a cuda::std::chrono::time_point
+ * @param stream CUDA stream used for device memory operations
+ * @param mr Memory resources used to allocate the returned column
  */
 template <typename T, bool nullable = false>
-inline cudf::test::fixed_width_column_wrapper<T, int64_t> generate_timestamps(int32_t count,
-                                                                              time_point_ms start,
-                                                                              time_point_ms stop)
+inline cudf::test::fixed_width_column_wrapper<T, int64_t> generate_timestamps(
+  int32_t count,
+  time_point_ms start,
+  time_point_ms stop,
+  rmm::cuda_stream_view stream = cudf::test::get_default_stream(),
+  cudf::memory_resources mr    = cudf::get_current_device_resource_ref())
 {
   using Rep        = typename T::rep;
   using Period     = typename T::period;
@@ -57,10 +63,10 @@ inline cudf::test::fixed_width_column_wrapper<T, int64_t> generate_timestamps(in
   if (nullable) {
     auto mask =
       cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 2 == 0; });
-    return cudf::test::fixed_width_column_wrapper<T, int64_t>(iter, iter + count, mask);
+    return cudf::test::fixed_width_column_wrapper<T, int64_t>(iter, iter + count, mask, stream, mr);
   } else {
     // This needs to be in an else to quash `statement_not_reachable` warnings
-    return cudf::test::fixed_width_column_wrapper<T, int64_t>(iter, iter + count);
+    return cudf::test::fixed_width_column_wrapper<T, int64_t>(iter, iter + count, stream, mr);
   }
 }
 

cpp/tests/utilities/column_utilities.cu

@@ -489,8 +489,9 @@ std::string stringify_column_differences(cudf::device_span<int const> difference
     buffer << depth_str << "differences:" << std::endl;
 
     auto source_table = cudf::table_view({lhs, rhs});
-    auto diff_column =
-      fixed_width_column_wrapper<int32_t>(h_differences.begin(), h_differences.end());
+    // Intermediate gather indices — allocate on temporary, not output.
+    auto diff_column = fixed_width_column_wrapper<int32_t>(
+      h_differences.begin(), h_differences.end(), stream, mr.get_temporary_mr());
     auto diff_table = cudf::gather(source_table,
                                    diff_column,
                                    cudf::out_of_bounds_policy::DONT_CHECK,