From 8707c0ac9178f794606b85155ddc0d02909d6e95 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Mon, 10 Aug 2026 13:03:54 -0500 Subject: [PATCH] Migrate Nanoarrow test helpers to memory_resources --- cpp/include/cudf_test/nanoarrow_utils.hpp | 78 ++++++++++++++++---- cpp/tests/interop/from_arrow_host_test.cpp | 8 +- cpp/tests/interop/from_arrow_stream_test.cpp | 39 +++++++++- cpp/tests/interop/from_arrow_test.cpp | 25 ++++--- cpp/tests/interop/to_arrow_device_test.cpp | 77 +++++++++++-------- 5 files changed, 164 insertions(+), 63 deletions(-) diff --git a/cpp/include/cudf_test/nanoarrow_utils.hpp b/cpp/include/cudf_test/nanoarrow_utils.hpp index d323d10ba39b..881dbf44d0ac 100644 --- a/cpp/include/cudf_test/nanoarrow_utils.hpp +++ b/cpp/include/cudf_test/nanoarrow_utils.hpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -100,8 +101,10 @@ std::enable_if_t() and !std::is_same_v, 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 -std::enable_if_t, void> populate_from_col(ArrowArray* arr, - cudf::column_view view) +std::enable_if_t, 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, void> populate_from_col(ArrowArray* ar ArrowArrayValidityBitmap(arr)->buffer.data = const_cast(reinterpret_cast(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(bitmask.first->data()); NANOARROW_THROW_NOT_OK(ArrowBufferSetAllocator( ArrowArrayBuffer(arr, 1), @@ -131,7 +134,9 @@ std::enable_if_t, void> populate_from_col(ArrowArray* ar // of the device buffers. template std::enable_if_t, 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, void> populate_from_col( ArrowArrayBuffer(arr, 2)->size_bytes = sview.chars_size(cudf::get_default_stream()); ArrowArrayBuffer(arr, 2)->data = const_cast(view.data()); } else { - auto zero = cudf::detail::device_scalar(0, cudf::get_default_stream()); + auto zero = + cudf::detail::device_scalar(0, cudf::get_default_stream(), mr.get_output_mr()); uint8_t const* ptr = reinterpret_cast(zero.data()); nanoarrow::BufferInitWrapped(ArrowArrayBuffer(arr, 1), std::move(zero), ptr, 4); } } template -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(dview.indices().data()); - populate_from_col(arr->dictionary, dview.keys()); + if constexpr (std::is_same_v or std::is_same_v) { + populate_from_col(arr->dictionary, dview.keys(), mr); + } else { + static_cast(mr); + populate_from_col(arr->dictionary, dview.keys()); + } } using vector_of_columns = std::vector>; +/** + * @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, 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 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 get_cudf_table( + cudf::memory_resources mr = cudf::get_current_device_resource_ref()); template struct nanoarrow_storage_type {}; @@ -388,11 +416,27 @@ nanoarrow::UniqueArray get_nanoarrow_list_array(std::initializer_list data, return get_nanoarrow_list_array(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, 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, 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 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, 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()); diff --git a/cpp/tests/interop/from_arrow_host_test.cpp b/cpp/tests/interop/from_arrow_host_test.cpp index 3cb451165f79..7715ee4bb096 100644 --- a/cpp/tests/interop/from_arrow_host_test.cpp +++ b/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, 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(test_data.int64_data, test_data.validity); auto string_array = get_nanoarrow_array(test_data.string_data, test_data.validity); cudf::dictionary_column_view view(table->get_column(2).view()); - auto keys = cudf::test::to_host(view.keys()).first; - auto indices = cudf::test::to_host(view.indices()).first; + auto keys = cudf::test::to_host(view.keys(), mr).first; + auto indices = cudf::test::to_host(view.indices(), mr).first; auto dict_array = get_nanoarrow_dict_array(std::vector(keys.begin(), keys.end()), std::vector(indices.begin(), indices.end()), test_data.validity); diff --git a/cpp/tests/interop/from_arrow_stream_test.cpp b/cpp/tests/interop/from_arrow_stream_test.cpp index 74f9fc1df31d..d3335d8ba25c 100644 --- a/cpp/tests/interop/from_arrow_stream_test.cpp +++ b/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 #include +#include + struct FromArrowStreamTest : public cudf::test::BaseFixture {}; void makeStreamFromArrays(std::vector arrays, @@ -29,14 +31,16 @@ void makeStreamFromArrays(std::vector arrays, } std::tuple, 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> tables; // The schema is unique across all tables. nanoarrow::UniqueSchema schema; std::vector 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); diff --git a/cpp/tests/interop/from_arrow_test.cpp b/cpp/tests/interop/from_arrow_test.cpp index 3ee6c378f558..5a4d2e4bf395 100644 --- a/cpp/tests/interop/from_arrow_test.cpp +++ b/cpp/tests/interop/from_arrow_test.cpp @@ -27,23 +27,27 @@ #include -std::unique_ptr get_cudf_table() +std::unique_ptr get_cudf_table(cudf::memory_resources mr) { + auto const temporary_mr = mr.get_temporary_mr(); std::vector> columns; columns.emplace_back(cudf::test::fixed_width_column_wrapper( - {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({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({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({1, 2, 5, 7}); - auto indices = cudf::test::fixed_width_column_wrapper({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({1, 2, 5, 7}, temporary_mr); + auto indices = + cudf::test::fixed_width_column_wrapper({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( - {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 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({{1, 2}, {3, 4}, {}, {6}, {7, 8, // 9}}).release()); diff --git a/cpp/tests/interop/to_arrow_device_test.cpp b/cpp/tests/interop/to_arrow_device_test.cpp index 7c4e5ea043f0..a73efaab3ff9 100644 --- a/cpp/tests/interop/to_arrow_device_test.cpp +++ b/cpp/tests/interop/to_arrow_device_test.cpp @@ -19,35 +19,42 @@ #include std::tuple, 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) { + auto const temporary_mr = mr.get_temporary_mr(); generated_test_data test_data(length); std::vector> columns; - columns.emplace_back(cudf::test::fixed_width_column_wrapper(test_data.int64_data.begin(), - test_data.int64_data.end(), - test_data.validity.begin()) - .release()); - columns.emplace_back(cudf::test::strings_column_wrapper(test_data.string_data.begin(), - test_data.string_data.end(), - test_data.validity.begin()) - .release()); - auto col4 = cudf::test::fixed_width_column_wrapper( - test_data.int64_data.begin(), test_data.int64_data.end(), test_data.validity.begin()); - columns.emplace_back(cudf::dictionary::encode(col4)); - columns.emplace_back(cudf::test::fixed_width_column_wrapper(test_data.bool_data.begin(), - test_data.bool_data.end(), - test_data.bool_validity.begin()) - .release()); + columns.emplace_back( + cudf::test::fixed_width_column_wrapper( + test_data.int64_data.begin(), test_data.int64_data.end(), test_data.validity.begin(), mr) + .release()); + columns.emplace_back( + cudf::test::strings_column_wrapper( + test_data.string_data.begin(), test_data.string_data.end(), test_data.validity.begin(), mr) + .release()); + auto col4 = cudf::test::fixed_width_column_wrapper(test_data.int64_data.begin(), + test_data.int64_data.end(), + test_data.validity.begin(), + temporary_mr); + columns.emplace_back(cudf::dictionary::encode( + col4, cudf::data_type{cudf::type_id::INT32}, cudf::get_default_stream(), mr.get_output_mr())); + columns.emplace_back( + cudf::test::fixed_width_column_wrapper( + test_data.bool_data.begin(), test_data.bool_data.end(), test_data.bool_validity.begin(), mr) + .release()); auto list_child_column = cudf::test::fixed_width_column_wrapper(test_data.list_int64_data.begin(), test_data.list_int64_data.end(), - test_data.list_int64_data_validity.begin()); + test_data.list_int64_data_validity.begin(), + mr); auto list_offsets_column = cudf::test::fixed_width_column_wrapper( - test_data.list_offsets.begin(), test_data.list_offsets.end()); - auto [list_mask, list_nulls] = cudf::bools_to_mask(cudf::test::fixed_width_column_wrapper( - test_data.list_validity.begin(), test_data.list_validity.end())); + test_data.list_offsets.begin(), test_data.list_offsets.end(), mr); + auto list_validity = cudf::test::fixed_width_column_wrapper( + test_data.list_validity.begin(), test_data.list_validity.end(), temporary_mr); + auto [list_mask, list_nulls] = + cudf::bools_to_mask(list_validity, cudf::get_default_stream(), mr.get_output_mr()); columns.emplace_back(cudf::make_lists_column(length, list_offsets_column.release(), list_child_column.release(), @@ -55,19 +62,25 @@ get_nanoarrow_cudf_table(cudf::size_type length) std::move(*list_mask))); auto int_column = cudf::test::fixed_width_column_wrapper( - test_data.int64_data.begin(), test_data.int64_data.end(), test_data.validity.begin()) + test_data.int64_data.begin(), test_data.int64_data.end(), test_data.validity.begin(), mr) .release(); auto str_column = cudf::test::strings_column_wrapper( - test_data.string_data.begin(), test_data.string_data.end(), test_data.validity.begin()) + test_data.string_data.begin(), test_data.string_data.end(), test_data.validity.begin(), mr) .release(); vector_of_columns cols; cols.push_back(std::move(int_column)); cols.push_back(std::move(str_column)); - auto [null_mask, null_count] = cudf::bools_to_mask(cudf::test::fixed_width_column_wrapper( - test_data.bool_data_validity.begin(), test_data.bool_data_validity.end())); - columns.emplace_back( - cudf::make_structs_column(length, std::move(cols), null_count, std::move(*null_mask))); + auto struct_validity = cudf::test::fixed_width_column_wrapper( + test_data.bool_data_validity.begin(), test_data.bool_data_validity.end(), temporary_mr); + auto [null_mask, null_count] = + cudf::bools_to_mask(struct_validity, cudf::get_default_stream(), mr.get_output_mr()); + columns.emplace_back(cudf::make_structs_column(length, + std::move(cols), + null_count, + std::move(*null_mask), + cudf::get_default_stream(), + mr.get_output_mr())); nanoarrow::UniqueSchema schema; ArrowSchemaInit(schema.get()); @@ -157,27 +170,27 @@ get_nanoarrow_cudf_table(cudf::size_type length) } std::tuple, nanoarrow::UniqueSchema, nanoarrow::UniqueArray> -get_nanoarrow_tables(cudf::size_type length) +get_nanoarrow_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); nanoarrow::UniqueArray arrow; NANOARROW_THROW_NOT_OK(ArrowArrayInitFromSchema(arrow.get(), schema.get(), nullptr)); arrow->length = length; populate_from_col(arrow->children[0], table->get_column(0).view()); - populate_from_col(arrow->children[1], table->get_column(1).view()); + populate_from_col(arrow->children[1], table->get_column(1).view(), mr); populate_dict_from_col( - arrow->children[2], cudf::dictionary_column_view(table->get_column(2).view())); + arrow->children[2], cudf::dictionary_column_view(table->get_column(2).view()), mr); - populate_from_col(arrow->children[3], table->get_column(3).view()); + populate_from_col(arrow->children[3], table->get_column(3).view(), mr); cudf::lists_column_view list_view{table->get_column(4).view()}; populate_list_from_col(arrow->children[4], list_view); populate_from_col(arrow->children[4]->children[0], list_view.child()); cudf::structs_column_view struct_view{table->get_column(5).view()}; populate_from_col(arrow->children[5]->children[0], struct_view.child(0)); - populate_from_col(arrow->children[5]->children[1], struct_view.child(1)); + populate_from_col(arrow->children[5]->children[1], struct_view.child(1), mr); arrow->children[5]->length = struct_view.size(); arrow->children[5]->null_count = struct_view.null_count(); NANOARROW_THROW_NOT_OK(