[cudf] Adding memory_resources support for column wrappers

cpp/include/cudf_test/base_fixture.hpp

@@ -1,17 +1,20 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
 #include <cudf_test/cudf_gtest.hpp>
+#include <cudf_test/default_stream.hpp>
 #include <cudf_test/file_utilities.hpp>
+#include <cudf_test/memory_resource_utilities.hpp>
 
 #include <cudf/utilities/export.hpp>
 #include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/traits.hpp>
 
+#include <rmm/cuda_stream_view.hpp>
 #include <rmm/resource_ref.hpp>
 
 #include <cuda/memory_resource>
@@ -39,6 +42,34 @@ class BaseFixture : public ::testing::Test {
   rmm::device_async_resource_ref mr() { return _mr; }
 };
 
+/**
+ * @brief Base fixture that instruments tests with a memory-resource harness.
+ *
+ * Each test instantiates a fresh harness. Tests should construct results with `resources()`.
+ * `TearDown` asserts that no output or temporary allocations remain live.
+ */
+struct BaseFixtureWithHarness : public BaseFixture {
+  /**
+   * @brief Assert that the harness has no live output or temporary allocations.
+   */
+  void TearDown() override { _harness.expect_no_live_allocations(stream()); }
+
+  /**
+   * @brief Return the default stream used by tests inheriting from this fixture.
+   * @return CUDA stream view
+   */
+  [[nodiscard]] rmm::cuda_stream_view stream() const { return cudf::test::get_default_stream(); }
+
+  /**
+   * @brief Return the harness output and temporary memory resources.
+   * @return Explicit output and temporary resources that do not consult the current resource
+   */
+  cudf::memory_resources resources() { return _harness.resources(); }
+
+ protected:
+  memory_resource_test_harness _harness{mr()};
+};
+
 /**
  * @brief Base test fixture that takes a parameter.
  *

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

@@ -6,6 +6,7 @@
 #pragma once
 
 #include <cudf_test/column_wrapper.hpp>
+#include <cudf_test/default_stream.hpp>
 
 #include <cudf/detail/iterator.cuh>
 #include <cudf/utilities/export.hpp>
@@ -31,11 +32,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;
@@ -56,10 +62,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

@@ -490,8 +490,8 @@ 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());
+    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,
@@ -542,6 +542,8 @@ struct column_comparator_impl {
     auto lhs_tview = table_view{{lhs}};
     auto rhs_tview = table_view{{rhs}};
 
+    // TODO: Pass `mr` once two_table_comparator / equality preprocessed_table::create accept
+    // memory_resources instead of allocating from the current device resource.
     auto const comparator =
       cudf::detail::row::equality::two_table_comparator{lhs_tview, rhs_tview, stream};
     auto const has_nulls = cudf::has_nulls(lhs_tview) or cudf::has_nulls(rhs_tview);
@@ -883,6 +885,9 @@ bool expect_columns_equal(cudf::column_view const& lhs,
                           cuda::stream_ref stream,
                           cudf::memory_resources mr)
 {
+  // TODO: equality row preprocessing (two_table_comparator / preprocessed_table::create) still
+  // allocates from the current device resource; pass `mr` through once that path accepts
+  // memory_resources so callers need not disable failing current-resource scopes.
   check_non_empty_nulls(lhs, rhs, stream);
   auto lhs_indices = generate_all_row_indices(lhs.size(), stream, mr);
   auto rhs_indices = generate_all_row_indices(rhs.size(), stream, mr);