From ec315d562916018bad9557e96e9105c062c91928 Mon Sep 17 00:00:00 2001 From: Robert Maynard Date: Mon, 17 Jun 2024 15:56:43 -0400 Subject: [PATCH] cudf::merge public API now support passing a user stream --- cpp/include/cudf/detail/merge.hpp | 1 + cpp/include/cudf/merge.hpp | 3 +- cpp/src/merge/merge.cu | 4 +- cpp/tests/CMakeLists.txt | 1 + cpp/tests/streams/merge_test.cpp | 119 ++++++++++++++++++++++++++++++ 5 files changed, 125 insertions(+), 3 deletions(-) create mode 100644 cpp/tests/streams/merge_test.cpp diff --git a/cpp/include/cudf/detail/merge.hpp b/cpp/include/cudf/detail/merge.hpp index 837eda0d7b5..56ac0554403 100644 --- a/cpp/include/cudf/detail/merge.hpp +++ b/cpp/include/cudf/detail/merge.hpp @@ -46,6 +46,7 @@ using index_vector = rmm::device_uvector; * std::vector const& key_cols, * std::vector const& column_order, * std::vector const& null_precedence, + * rmm::cuda_stream_view stream, * rmm::device_async_resource_ref mr) * * @param stream CUDA stream used for device memory operations and kernel launches diff --git a/cpp/include/cudf/merge.hpp b/cpp/include/cudf/merge.hpp index 29aa3ffe934..301e56c19b8 100644 --- a/cpp/include/cudf/merge.hpp +++ b/cpp/include/cudf/merge.hpp @@ -97,6 +97,7 @@ namespace cudf { * @param[in] column_order Sort order types of columns indexed by key_cols * @param[in] null_precedence Array indicating the order of nulls with respect * to non-nulls for the indexing columns (key_cols) + * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used to allocate the returned table's device memory * * @returns A table containing sorted data from all input tables @@ -106,7 +107,7 @@ std::unique_ptr merge( std::vector const& key_cols, std::vector const& column_order, std::vector const& null_precedence = {}, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource()); - /** @} */ // end of group } // namespace cudf diff --git a/cpp/src/merge/merge.cu b/cpp/src/merge/merge.cu index 630cf328579..7ecaa0fba56 100644 --- a/cpp/src/merge/merge.cu +++ b/cpp/src/merge/merge.cu @@ -694,11 +694,11 @@ std::unique_ptr merge(std::vector const& tables_to_merg std::vector const& key_cols, std::vector const& column_order, std::vector const& null_precedence, + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); - return detail::merge( - tables_to_merge, key_cols, column_order, null_precedence, cudf::get_default_stream(), mr); + return detail::merge(tables_to_merge, key_cols, column_order, null_precedence, stream, mr); } } // namespace cudf diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index eef09954647..c04f64a61bc 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -691,6 +691,7 @@ ConfigureTest(STREAM_INTEROP_TEST streams/interop_test.cpp STREAM_MODE testing) ConfigureTest(STREAM_JSONIO_TEST streams/io/json_test.cpp STREAM_MODE testing) ConfigureTest(STREAM_LABELING_BINS_TEST streams/labeling_bins_test.cpp STREAM_MODE testing) ConfigureTest(STREAM_LISTS_TEST streams/lists_test.cpp STREAM_MODE testing) +ConfigureTest(STREAM_MERGE_TEST streams/merge_test.cpp STREAM_MODE testing) ConfigureTest(STREAM_NULL_MASK_TEST streams/null_mask_test.cpp STREAM_MODE testing) ConfigureTest(STREAM_ORCIO_TEST streams/io/orc_test.cpp STREAM_MODE testing) ConfigureTest(STREAM_PARQUETIO_TEST streams/io/parquet_test.cpp STREAM_MODE testing) diff --git a/cpp/tests/streams/merge_test.cpp b/cpp/tests/streams/merge_test.cpp new file mode 100644 index 00000000000..9c5daf733ae --- /dev/null +++ b/cpp/tests/streams/merge_test.cpp @@ -0,0 +1,119 @@ +/* + * Copyright (c) 2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include + +template +class MergeTest_ : public cudf::test::BaseFixture {}; + +TYPED_TEST_SUITE(MergeTest_, cudf::test::FixedWidthTypes); + +TYPED_TEST(MergeTest_, MergeIsZeroWhenShouldNotBeZero) +{ + using columnFactoryT = cudf::test::fixed_width_column_wrapper; + + columnFactoryT leftColWrap1({1, 2, 3, 4, 5}); + cudf::test::fixed_width_column_wrapper rightColWrap1{}; + + std::vector key_cols{0}; + std::vector column_order; + column_order.push_back(cudf::order::ASCENDING); + std::vector null_precedence(column_order.size(), cudf::null_order::AFTER); + + cudf::table_view left_view{{leftColWrap1}}; + cudf::table_view right_view{{rightColWrap1}}; + cudf::table_view expected{{leftColWrap1}}; + + auto result = cudf::merge({left_view, right_view}, key_cols, column_order, null_precedence); + + int expected_len = 5; + ASSERT_EQ(result->num_rows(), expected_len); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result->view()); +} + +TYPED_TEST(MergeTest_, SingleTableInput) +{ + cudf::size_type inputRows = 40; + + auto sequence = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + cudf::test::fixed_width_column_wrapper + colWrap1(sequence, sequence + inputRows); + + std::vector key_cols{0}; + std::vector column_order{cudf::order::ASCENDING}; + std::vector null_precedence{}; + + cudf::table_view left_view{{colWrap1}}; + + std::unique_ptr p_outputTable; + CUDF_EXPECT_NO_THROW(p_outputTable = + cudf::merge({left_view}, key_cols, column_order, null_precedence)); + + auto input_column_view{left_view.column(0)}; + auto output_column_view{p_outputTable->view().column(0)}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(input_column_view, output_column_view); +} + +class MergeTest : public cudf::test::BaseFixture {}; + +TEST_F(MergeTest, KeysWithNulls) +{ + cudf::size_type nrows = 13200; // Ensures that thrust::merge uses more than one tile/block + auto data_iter = thrust::make_counting_iterator(0); + auto valids1 = + cudf::detail::make_counting_transform_iterator(0, [](auto row) { return row % 10 != 0; }); + cudf::test::fixed_width_column_wrapper data1(data_iter, data_iter + nrows, valids1); + auto valids2 = + cudf::detail::make_counting_transform_iterator(0, [](auto row) { return row % 15 != 0; }); + cudf::test::fixed_width_column_wrapper data2(data_iter, data_iter + nrows, valids2); + auto all_data = cudf::concatenate(std::vector{{data1, data2}}); + + std::vector column_orders{cudf::order::ASCENDING, cudf::order::DESCENDING}; + std::vector null_precedences{cudf::null_order::AFTER, cudf::null_order::BEFORE}; + + for (auto co : column_orders) + for (auto np : null_precedences) { + std::vector column_order{co}; + std::vector null_precedence{np}; + auto sorted1 = + cudf::sort(cudf::table_view({data1}), column_order, null_precedence)->release(); + auto col1 = sorted1.front()->view(); + auto sorted2 = + cudf::sort(cudf::table_view({data2}), column_order, null_precedence)->release(); + auto col2 = sorted2.front()->view(); + + auto result = cudf::merge( + {cudf::table_view({col1}), cudf::table_view({col2})}, {0}, column_order, null_precedence); + auto sorted_all = + cudf::sort(cudf::table_view({all_data->view()}), column_order, null_precedence); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(sorted_all->view().column(0), result->view().column(0)); + } +} + +CUDF_TEST_PROGRAM_MAIN()