Migrate Nanoarrow test helpers to memory_resources

cpp/include/cudf_test/nanoarrow_utils.hpp

@@ -14,6 +14,7 @@
 #include <cudf/transform.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/traits.hpp>
 #include <cudf/wrappers/durations.hpp>
 
@@ -100,8 +101,10 @@ std::enable_if_t<cudf::is_fixed_width<T>() and !std::is_same_v<T, bool>, void> p
 // still represent boolean arrays differently, we have to use bools_to_mask
 // and give the ArrowArray object ownership of the device data.
 template <typename T>
-std::enable_if_t<std::is_same_v<T, bool>, void> populate_from_col(ArrowArray* arr,
-                                                                  cudf::column_view view)
+std::enable_if_t<std::is_same_v<T, bool>, void> populate_from_col(
+  ArrowArray* arr,
+  cudf::column_view view,
+  cudf::memory_resources mr = cudf::get_current_device_resource_ref())
 {
   arr->length     = view.size();
   arr->null_count = view.null_count();
@@ -112,7 +115,7 @@ std::enable_if_t<std::is_same_v<T, bool>, void> populate_from_col(ArrowArray* ar
   ArrowArrayValidityBitmap(arr)->buffer.data =
     const_cast<uint8_t*>(reinterpret_cast<uint8_t const*>(view.null_mask()));
 
-  auto bitmask = cudf::bools_to_mask(view);
+  auto bitmask = cudf::bools_to_mask(view, cudf::get_default_stream(), mr.get_output_mr());
   auto ptr     = reinterpret_cast<uint8_t*>(bitmask.first->data());
   NANOARROW_THROW_NOT_OK(ArrowBufferSetAllocator(
     ArrowArrayBuffer(arr, 1),
@@ -131,7 +134,9 @@ std::enable_if_t<std::is_same_v<T, bool>, void> populate_from_col(ArrowArray* ar
 // of the device buffers.
 template <typename T>
 std::enable_if_t<std::is_same_v<T, cudf::string_view>, void> populate_from_col(
-  ArrowArray* arr, cudf::column_view view)
+  ArrowArray* arr,
+  cudf::column_view view,
+  cudf::memory_resources mr = cudf::get_current_device_resource_ref())
 {
   arr->length     = view.size();
   arr->null_count = view.null_count();
@@ -151,14 +156,17 @@ std::enable_if_t<std::is_same_v<T, cudf::string_view>, void> populate_from_col(
     ArrowArrayBuffer(arr, 2)->size_bytes = sview.chars_size(cudf::get_default_stream());
     ArrowArrayBuffer(arr, 2)->data       = const_cast<uint8_t*>(view.data<uint8_t>());
   } else {
-    auto zero          = cudf::detail::device_scalar<int32_t>(0, cudf::get_default_stream());
+    auto zero =
+      cudf::detail::device_scalar<int32_t>(0, cudf::get_default_stream(), mr.get_output_mr());
     uint8_t const* ptr = reinterpret_cast<uint8_t*>(zero.data());
     nanoarrow::BufferInitWrapped(ArrowArrayBuffer(arr, 1), std::move(zero), ptr, 4);
   }
 }
 
 template <typename KEY_TYPE, typename IND_TYPE>
-void populate_dict_from_col(ArrowArray* arr, cudf::dictionary_column_view dview)
+void populate_dict_from_col(ArrowArray* arr,
+                            cudf::dictionary_column_view dview,
+                            cudf::memory_resources mr = cudf::get_current_device_resource_ref())
 {
   arr->length     = dview.size();
   arr->null_count = dview.null_count();
@@ -172,17 +180,37 @@ void populate_dict_from_col(ArrowArray* arr, cudf::dictionary_column_view dview)
   ArrowArrayBuffer(arr, 1)->size_bytes = sizeof(IND_TYPE) * dview.indices().size();
   ArrowArrayBuffer(arr, 1)->data       = const_cast<uint8_t*>(dview.indices().data<uint8_t>());
 
-  populate_from_col<KEY_TYPE>(arr->dictionary, dview.keys());
+  if constexpr (std::is_same_v<KEY_TYPE, bool> or std::is_same_v<KEY_TYPE, cudf::string_view>) {
+    populate_from_col<KEY_TYPE>(arr->dictionary, dview.keys(), mr);
+  } else {
+    static_cast<void>(mr);
+    populate_from_col<KEY_TYPE>(arr->dictionary, dview.keys());
+  }
 }
 
 using vector_of_columns = std::vector<std::unique_ptr<cudf::column>>;
 
+/**
+ * @brief Create equivalent cuDF and device-backed nanoarrow tables.
+ *
+ * @param length Number of rows to generate
+ * @param mr Memory resources used for returned device allocations and helper temporaries
+ * @return cuDF table, Arrow schema, and Arrow array
+ */
 std::tuple<std::unique_ptr<cudf::table>, nanoarrow::UniqueSchema, nanoarrow::UniqueArray>
-get_nanoarrow_tables(cudf::size_type length = 10000);
+get_nanoarrow_tables(cudf::size_type length    = 10000,
+                     cudf::memory_resources mr = cudf::get_current_device_resource_ref());
 
 void populate_list_from_col(ArrowArray* arr, cudf::lists_column_view view);
 
-std::unique_ptr<cudf::table> get_cudf_table();
+/**
+ * @brief Create the standard cuDF table used by Arrow interop tests.
+ *
+ * @param mr Memory resources used for returned table allocations and helper temporaries
+ * @return Generated cuDF table
+ */
+std::unique_ptr<cudf::table> get_cudf_table(
+  cudf::memory_resources mr = cudf::get_current_device_resource_ref());
 
 template <typename T>
 struct nanoarrow_storage_type {};
@@ -388,11 +416,27 @@ nanoarrow::UniqueArray get_nanoarrow_list_array(std::initializer_list<T> data,
   return get_nanoarrow_list_array<T>(data_vector, offset, data_mask, list_mask);
 }
 
+/**
+ * @brief Create a cuDF table, matching Arrow schema, and source host data.
+ *
+ * @param length Number of rows to generate
+ * @param mr Memory resources used for returned table allocations and helper temporaries
+ * @return cuDF table, Arrow schema, and generated host data
+ */
 std::tuple<std::unique_ptr<cudf::table>, nanoarrow::UniqueSchema, generated_test_data>
-get_nanoarrow_cudf_table(cudf::size_type length);
-
+get_nanoarrow_cudf_table(cudf::size_type length,
+                         cudf::memory_resources mr = cudf::get_current_device_resource_ref());
+
+/**
+ * @brief Create equivalent cuDF and host-backed nanoarrow tables.
+ *
+ * @param length Number of rows to generate
+ * @param mr Memory resources used for returned table allocations and helper temporaries
+ * @return cuDF table, Arrow schema, and Arrow array
+ */
 std::tuple<std::unique_ptr<cudf::table>, nanoarrow::UniqueSchema, nanoarrow::UniqueArray>
-get_nanoarrow_host_tables(cudf::size_type length);
+get_nanoarrow_host_tables(cudf::size_type length,
+                          cudf::memory_resources mr = cudf::get_current_device_resource_ref());
 
 void slice_host_nanoarrow(ArrowArray* arr, int64_t start, int64_t end);
 
@@ -442,5 +486,13 @@ void makeStreamFromArrays(std::vector<nanoarrow::UniqueArray> arrays,
                           nanoarrow::UniqueSchema schema,
                           ArrowArrayStream* out);
 
+/**
+ * @brief Create a cuDF table and equivalent nanoarrow stream.
+ *
+ * @param num_copies Number of record batches in the stream
+ * @param mr Memory resources used for returned table allocations and helper temporaries
+ * @return Concatenated cuDF table, Arrow schema, and Arrow stream
+ */
 std::tuple<std::unique_ptr<cudf::table>, nanoarrow::UniqueSchema, ArrowArrayStream>
-get_nanoarrow_stream(int num_copies);
+get_nanoarrow_stream(int num_copies,
+                     cudf::memory_resources mr = cudf::get_current_device_resource_ref());

cpp/tests/interop/from_arrow_host_test.cpp

@@ -26,16 +26,16 @@
 
 // create a cudf::table and equivalent arrow table with host memory
 std::tuple<std::unique_ptr<cudf::table>, nanoarrow::UniqueSchema, nanoarrow::UniqueArray>
-get_nanoarrow_host_tables(cudf::size_type length)
+get_nanoarrow_host_tables(cudf::size_type length, cudf::memory_resources mr)
 {
-  auto [table, schema, test_data] = get_nanoarrow_cudf_table(length);
+  auto [table, schema, test_data] = get_nanoarrow_cudf_table(length, mr);
 
   auto int64_array = get_nanoarrow_array<int64_t>(test_data.int64_data, test_data.validity);
   auto string_array =
     get_nanoarrow_array<cudf::string_view>(test_data.string_data, test_data.validity);
   cudf::dictionary_column_view view(table->get_column(2).view());
-  auto keys       = cudf::test::to_host<int64_t>(view.keys()).first;
-  auto indices    = cudf::test::to_host<uint32_t>(view.indices()).first;
+  auto keys       = cudf::test::to_host<int64_t>(view.keys(), mr).first;
+  auto indices    = cudf::test::to_host<uint32_t>(view.indices(), mr).first;
   auto dict_array = get_nanoarrow_dict_array(std::vector<int64_t>(keys.begin(), keys.end()),
                                              std::vector<int32_t>(indices.begin(), indices.end()),
                                              test_data.validity);

cpp/tests/interop/from_arrow_stream_test.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -14,6 +14,8 @@
 #include <cudf/table/table_view.hpp>
 #include <cudf/utilities/type_checks.hpp>
 
+#include <rmm/mr/statistics_resource_adaptor.hpp>
+
 struct FromArrowStreamTest : public cudf::test::BaseFixture {};
 
 void makeStreamFromArrays(std::vector<nanoarrow::UniqueArray> arrays,
@@ -29,14 +31,16 @@ void makeStreamFromArrays(std::vector<nanoarrow::UniqueArray> arrays,
 }
 
 std::tuple<std::unique_ptr<cudf::table>, nanoarrow::UniqueSchema, ArrowArrayStream>
-get_nanoarrow_stream(int num_copies)
+get_nanoarrow_stream(int num_copies, cudf::memory_resources mr)
 {
+  auto const temporary_mr        = mr.get_temporary_mr();
+  auto const temporary_resources = cudf::memory_resources{temporary_mr, temporary_mr};
   std::vector<std::unique_ptr<cudf::table>> tables;
   // The schema is unique across all tables.
   nanoarrow::UniqueSchema schema;
   std::vector<nanoarrow::UniqueArray> arrays;
   for (auto i = 0; i < num_copies; ++i) {
-    auto [tbl, sch, arr] = get_nanoarrow_host_tables(3);
+    auto [tbl, sch, arr] = get_nanoarrow_host_tables(3, temporary_resources);
     tables.push_back(std::move(tbl));
     arrays.push_back(std::move(arr));
     if (i == 0) { sch.move(schema.get()); }
@@ -45,7 +49,7 @@ get_nanoarrow_stream(int num_copies)
   for (auto const& table : tables) {
     table_views.push_back(table->view());
   }
-  auto expected = cudf::concatenate(table_views);
+  auto expected = cudf::concatenate(table_views, cudf::get_default_stream(), mr.get_output_mr());
 
   ArrowArrayStream stream;
   makeStreamFromArrays(std::move(arrays), std::move(schema), &stream);
@@ -87,6 +91,33 @@ TEST_F(FromArrowStreamTest, BasicTest)
   CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), result->view());
 }
 
+TEST_F(FromArrowStreamTest, TestUtilityMemoryResourceControl)
+{
+  auto upstream     = this->mr();
+  auto output_mr    = rmm::mr::statistics_resource_adaptor(upstream);
+  auto temporary_mr = rmm::mr::statistics_resource_adaptor(upstream);
+  auto resources    = cudf::memory_resources{output_mr, temporary_mr};
+
+  {
+    auto direct_table                                   = get_cudf_table(resources);
+    auto [generated_table, generated_schema, test_data] = get_nanoarrow_cudf_table(3, resources);
+    auto [device_table, device_schema, device_array]    = get_nanoarrow_tables(0, resources);
+    auto [host_table, host_schema, host_array]          = get_nanoarrow_host_tables(3, resources);
+    auto [stream_table, stream_schema, stream]          = get_nanoarrow_stream(2, resources);
+
+    cudf::get_default_stream().synchronize();
+    EXPECT_GT(output_mr.get_bytes_counter().value, 0);
+    EXPECT_EQ(temporary_mr.get_bytes_counter().value, 0);
+    EXPECT_GT(temporary_mr.get_bytes_counter().total, 0);
+
+    if (stream.release != nullptr) { stream.release(&stream); }
+  }
+
+  cudf::get_default_stream().synchronize();
+  EXPECT_EQ(output_mr.get_bytes_counter().value, 0);
+  EXPECT_EQ(temporary_mr.get_bytes_counter().value, 0);
+}
+
 TEST_F(FromArrowStreamTest, EmptyTest)
 {
   auto [tbl, sch, arr] = get_nanoarrow_host_tables(0);

cpp/tests/interop/from_arrow_test.cpp

@@ -27,23 +27,27 @@
 
 #include <arrow/c/bridge.h>
 
-std::unique_ptr<cudf::table> get_cudf_table()
+std::unique_ptr<cudf::table> get_cudf_table(cudf::memory_resources mr)
 {
+  auto const temporary_mr = mr.get_temporary_mr();
   std::vector<std::unique_ptr<cudf::column>> columns;
   columns.emplace_back(cudf::test::fixed_width_column_wrapper<int32_t>(
-                         {1, 2, 5, 2, 7}, {true, false, true, true, true})
+                         {1, 2, 5, 2, 7}, {true, false, true, true, true}, mr)
                          .release());
-  columns.emplace_back(cudf::test::fixed_width_column_wrapper<int64_t>({1, 2, 3, 4, 5}).release());
-  columns.emplace_back(cudf::test::strings_column_wrapper({"fff", "aaa", "", "fff", "ccc"},
-                                                          {true, true, true, false, true})
+  columns.emplace_back(
+    cudf::test::fixed_width_column_wrapper<int64_t>({1, 2, 3, 4, 5}, mr).release());
+  columns.emplace_back(cudf::test::strings_column_wrapper(
+                         {"fff", "aaa", "", "fff", "ccc"}, {true, true, true, false, true}, mr)
                          .release());
 
-  auto keys    = cudf::test::fixed_width_column_wrapper<int32_t>({1, 2, 5, 7});
-  auto indices = cudf::test::fixed_width_column_wrapper<int32_t>({0, 1, 2, 1, 3}, {1, 0, 1, 1, 1});
-  columns.emplace_back(cudf::make_dictionary_column(keys, indices));
+  auto keys = cudf::test::fixed_width_column_wrapper<int32_t>({1, 2, 5, 7}, temporary_mr);
+  auto indices =
+    cudf::test::fixed_width_column_wrapper<int32_t>({0, 1, 2, 1, 3}, {1, 0, 1, 1, 1}, temporary_mr);
+  columns.emplace_back(
+    cudf::make_dictionary_column(keys, indices, cudf::get_default_stream(), mr.get_output_mr()));
 
   columns.emplace_back(cudf::test::fixed_width_column_wrapper<bool>(
-                         {true, false, true, false, true}, {true, false, true, true, false})
+                         {true, false, true, false, true}, {true, false, true, true, false}, mr)
                          .release());
   columns.emplace_back(cudf::test::strings_column_wrapper(
                          {
@@ -53,7 +57,8 @@ std::unique_ptr<cudf::table> get_cudf_table()
                            "1",
                            "2",
                          },
-                         {0, 1, 1, 1, 1})
+                         {0, 1, 1, 1, 1},
+                         mr)
                          .release());
   // columns.emplace_back(cudf::test::lists_column_wrapper<int>({{1, 2}, {3, 4}, {}, {6}, {7, 8,
   // 9}}).release());