Migrate Nanoarrow test helpers to memory_resources

cpp/include/cudf_test/nanoarrow_utils.hpp

@@ -13,13 +13,19 @@
 #include <cudf/strings/strings_column_view.hpp>
 #include <cudf/transform.hpp>
 #include <cudf/types.hpp>
+#include <cudf/utilities/default_stream.hpp>
 #include <cudf/utilities/error.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/traits.hpp>
 #include <cudf/wrappers/durations.hpp>
 
+#include <cuda/stream>
+
 #include <nanoarrow/nanoarrow.hpp>
 #include <nanoarrow/nanoarrow_device.h>
 
+#include <concepts>
+
 struct generated_test_data {
   generated_test_data(cudf::size_type length)
     : int64_data(length),
@@ -144,8 +150,8 @@ static ArrowBufferAllocator noop_alloc = (struct ArrowBufferAllocator){
 // populate an ArrowArray with pointers to the raw device buffers of a cudf::column_view
 // and use the no-op alloc so that the ArrowArray doesn't presume ownership of the data
 template <typename T>
-std::enable_if_t<cudf::is_fixed_width<T>() and !std::is_same_v<T, bool>, void> populate_from_col(
-  ArrowArray* arr, cudf::column_view view)
+void populate_from_col(ArrowArray* arr, cudf::column_view view)
+  requires(cudf::is_fixed_width<T>() && !cudf::is_boolean<T>())
 {
   arr->length     = view.size();
   arr->null_count = view.null_count();
@@ -163,8 +169,11 @@ 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)
+void populate_from_col(ArrowArray* arr,
+                       cudf::column_view view,
+                       cuda::stream_ref stream   = cudf::get_default_stream(),
+                       cudf::memory_resources mr = cudf::get_current_device_resource_ref())
+  requires(cudf::is_boolean<T>())
 {
   arr->length     = view.size();
   arr->null_count = view.null_count();
@@ -175,7 +184,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, stream, mr.get_output_mr());
   auto ptr     = reinterpret_cast<uint8_t*>(bitmask.first->data());
   NANOARROW_THROW_NOT_OK(ArrowBufferSetAllocator(
     ArrowArrayBuffer(arr, 1),
@@ -193,8 +202,11 @@ std::enable_if_t<std::is_same_v<T, bool>, void> populate_from_col(ArrowArray* ar
 // using no-op allocator so the ArrowArray knows it doesn't have ownership
 // 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)
+void populate_from_col(ArrowArray* arr,
+                       cudf::column_view view,
+                       cuda::stream_ref stream   = cudf::get_default_stream(),
+                       cudf::memory_resources mr = cudf::get_current_device_resource_ref())
+  requires(std::same_as<T, cudf::string_view>)
 {
   arr->length     = view.size();
   arr->null_count = view.null_count();
@@ -211,17 +223,20 @@ std::enable_if_t<std::is_same_v<T, cudf::string_view>, void> populate_from_col(
     ArrowArrayBuffer(arr, 1)->size_bytes = sizeof(int32_t) * sview.offsets().size();
     ArrowArrayBuffer(arr, 1)->data       = const_cast<uint8_t*>(sview.offsets().data<uint8_t>());
     NANOARROW_THROW_NOT_OK(ArrowBufferSetAllocator(ArrowArrayBuffer(arr, 2), noop_alloc));
-    ArrowArrayBuffer(arr, 2)->size_bytes = sview.chars_size(cudf::get_default_stream());
+    ArrowArrayBuffer(arr, 2)->size_bytes = sview.chars_size(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, 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,
+                            cuda::stream_ref stream   = cudf::get_default_stream(),
+                            cudf::memory_resources mr = cudf::get_current_device_resource_ref())
 {
   arr->length     = dview.size();
   arr->null_count = dview.null_count();
@@ -235,17 +250,40 @@ 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 (cudf::is_boolean<KEY_TYPE>() or std::same_as<KEY_TYPE, cudf::string_view>) {
+    populate_from_col<KEY_TYPE>(arr->dictionary, dview.keys(), stream, mr);
+  } else {
+    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 stream CUDA stream used for device memory operations and kernel launches
+ * @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,
+                     cuda::stream_ref stream   = cudf::get_default_stream(),
+                     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 stream CUDA stream used for device memory operations and kernel launches
+ * @param mr Memory resources used for returned table allocations and helper temporaries
+ * @return Generated cuDF table
+ */
+std::unique_ptr<cudf::table> get_cudf_table(
+  cuda::stream_ref stream   = cudf::get_default_stream(),
+  cudf::memory_resources mr = cudf::get_current_device_resource_ref());
 
 template <typename T>
 struct nanoarrow_storage_type {};
@@ -292,8 +330,9 @@ struct nanoarrow_decimal_type<__int128_t> {
 };
 
 template <typename T>
-std::enable_if_t<cudf::is_fixed_width<T>() and !std::is_same_v<T, bool>, nanoarrow::UniqueArray>
-get_nanoarrow_array(std::vector<T> const& data, std::vector<uint8_t> const& mask = {})
+nanoarrow::UniqueArray get_nanoarrow_array(std::vector<T> const& data,
+                                           std::vector<uint8_t> const& mask = {})
+  requires(cudf::is_fixed_width<T>() && !cudf::is_boolean<T>())
 {
   nanoarrow::UniqueArray tmp;
   NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(tmp.get(), nanoarrow_storage_type<T>::type));
@@ -322,8 +361,9 @@ get_nanoarrow_array(std::vector<T> const& data, std::vector<uint8_t> const& mask
 }
 
 template <typename T>
-std::enable_if_t<std::is_same_v<T, bool>, nanoarrow::UniqueArray> get_nanoarrow_array(
-  std::vector<bool> const& data, std::vector<bool> const& mask = {})
+nanoarrow::UniqueArray get_nanoarrow_array(std::vector<bool> const& data,
+                                           std::vector<bool> const& mask = {})
+  requires(cudf::is_boolean<T>())
 {
   nanoarrow::UniqueArray tmp;
   NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(tmp.get(), NANOARROW_TYPE_BOOL));
@@ -368,8 +408,9 @@ nanoarrow::UniqueArray get_nanoarrow_array(std::initializer_list<T> elements,
 }
 
 template <typename T>
-std::enable_if_t<std::is_same_v<T, cudf::string_view>, nanoarrow::UniqueArray> get_nanoarrow_array(
-  std::vector<std::string> const& data, std::vector<uint8_t> const& mask = {})
+nanoarrow::UniqueArray get_nanoarrow_array(std::vector<std::string> const& data,
+                                           std::vector<uint8_t> const& mask = {})
+  requires(std::same_as<T, cudf::string_view>)
 {
   nanoarrow::UniqueArray tmp;
   NANOARROW_THROW_NOT_OK(ArrowArrayInitFromType(tmp.get(), NANOARROW_TYPE_STRING));
@@ -451,20 +492,37 @@ 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 stream CUDA stream used for device memory operations and kernel launches
+ * @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,
+                         cuda::stream_ref stream   = cudf::get_default_stream(),
+                         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 stream CUDA stream used for device memory operations and kernel launches
+ * @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,
+                          cuda::stream_ref stream   = cudf::get_default_stream(),
+                          cudf::memory_resources mr = cudf::get_current_device_resource_ref());
 
 void slice_host_nanoarrow(ArrowArray* arr, int64_t start, int64_t end);
 
 template <typename T>
-std::enable_if_t<std::disjunction_v<std::is_same<T, int32_t>,
-                                    std::is_same<T, int64_t>,
-                                    std::is_same<T, __int128_t>>,
-                 std::size_t>
-get_decimal_precision()
+std::size_t get_decimal_precision()
+  requires(std::same_as<T, int32_t> || std::same_as<T, int64_t> || std::same_as<T, __int128_t>)
 {
   return std::numeric_limits<T>::digits10;
 }
@@ -505,5 +563,15 @@ 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 stream CUDA stream used for device memory operations and kernel launches
+ * @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,
+                     cuda::stream_ref stream   = cudf::get_default_stream(),
+                     cudf::memory_resources mr = cudf::get_current_device_resource_ref());

cpp/tests/interop/from_arrow_host_test.cpp

@@ -117,16 +117,18 @@ struct direct_arrow_c_producer {
 
 // 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,
+                          cuda::stream_ref stream,
+                          cudf::memory_resources mr)
 {
-  auto [table, schema, test_data] = get_nanoarrow_cudf_table(length);
+  auto [table, schema, test_data] = get_nanoarrow_cudf_table(length, stream, 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(), stream, mr).first;
+  auto indices    = cudf::test::to_host<uint32_t>(view.indices(), stream, 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

@@ -11,17 +11,61 @@
 
 #include <cudf/concatenate.hpp>
 #include <cudf/copying.hpp>
+#include <cudf/dictionary/dictionary_factories.hpp>
 #include <cudf/interop.hpp>
 #include <cudf/lists/lists_column_view.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/utilities/type_checks.hpp>
 
+#include <rmm/mr/statistics_resource_adaptor.hpp>
+
 #include <numeric>
 #include <vector>
 
 struct FromArrowStreamTest : public cudf::test::BaseFixture {};
 
+// Defined here rather than in from_arrow_test.cpp, which historically held this helper. That file
+// still depends on Arrow C++ and was dropped from the INTEROP_TEST target when the Arrow C++
+// dependency was removed from the C++ tests, so a definition placed there never reaches the link.
+std::unique_ptr<cudf::table> get_cudf_table(cuda::stream_ref stream, 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}, stream, mr)
+                         .release());
+  columns.emplace_back(
+    cudf::test::fixed_width_column_wrapper<int64_t>({1, 2, 3, 4, 5}, stream, mr).release());
+  columns.emplace_back(
+    cudf::test::strings_column_wrapper(
+      {"fff", "aaa", "", "fff", "ccc"}, {true, true, true, false, true}, stream, mr)
+      .release());
+
+  auto keys = cudf::test::fixed_width_column_wrapper<int32_t>({1, 2, 5, 7}, stream, temporary_mr);
+  auto indices = cudf::test::fixed_width_column_wrapper<int32_t>(
+    {0, 1, 2, 1, 3}, {1, 0, 1, 1, 1}, stream, temporary_mr);
+  columns.emplace_back(cudf::make_dictionary_column(keys, indices, 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}, stream, mr)
+      .release());
+  columns.emplace_back(cudf::test::strings_column_wrapper(
+                         {
+                           "",
+                           "abc",
+                           "def",
+                           "1",
+                           "2",
+                         },
+                         {0, 1, 1, 1, 1},
+                         stream,
+                         mr)
+                         .release());
+  return std::make_unique<cudf::table>(std::move(columns));
+}
+
 void makeStreamFromArrays(std::vector<nanoarrow::UniqueArray> arrays,
                           nanoarrow::UniqueSchema schema,
                           ArrowArrayStream* out)
@@ -35,14 +79,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, cuda::stream_ref stream, 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, stream, temporary_resources);
     tables.push_back(std::move(tbl));
     arrays.push_back(std::move(arr));
     if (i == 0) { sch.move(schema.get()); }
@@ -51,11 +97,11 @@ 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, stream, mr.get_output_mr());
 
-  ArrowArrayStream stream;
-  makeStreamFromArrays(std::move(arrays), std::move(schema), &stream);
-  return std::make_tuple(std::move(expected), std::move(schema), stream);
+  ArrowArrayStream arrow_stream;
+  makeStreamFromArrays(std::move(arrays), std::move(schema), &arrow_stream);
+  return std::make_tuple(std::move(expected), std::move(schema), arrow_stream);
 }
 
 std::tuple<std::unique_ptr<cudf::column>, nanoarrow::UniqueSchema, ArrowArrayStream>
@@ -93,6 +139,35 @@ 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 stream       = cudf::get_default_stream();
+
+  {
+    auto direct_table = get_cudf_table(stream, resources);
+    auto [generated_table, generated_schema, test_data] =
+      get_nanoarrow_cudf_table(3, stream, resources);
+    auto [device_table, device_schema, device_array] = get_nanoarrow_tables(0, stream, resources);
+    auto [host_table, host_schema, host_array] = get_nanoarrow_host_tables(3, stream, resources);
+    auto [stream_table, stream_schema, arrow_stream] = get_nanoarrow_stream(2, stream, resources);
+
+    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 (arrow_stream.release != nullptr) { arrow_stream.release(&arrow_stream); }
+  }
+
+  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);