diff --git a/cpp/include/cudf/lists/extract.hpp b/cpp/include/cudf/lists/extract.hpp index e92354134e8..14c0f59e17d 100644 --- a/cpp/include/cudf/lists/extract.hpp +++ b/cpp/include/cudf/lists/extract.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2022, NVIDIA CORPORATION. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -59,12 +59,14 @@ namespace lists { * * @param lists_column Column to extract elements from. * @param index The row within each sublist to retrieve. + * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned column's device memory. * @return Column of extracted elements. */ std::unique_ptr extract_list_element( lists_column_view const& lists_column, size_type index, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @@ -97,6 +99,7 @@ std::unique_ptr extract_list_element( * @param lists_column Column to extract elements from. * @param indices The column whose rows indicate the element index to be retrieved from each list * row. + * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned column's device memory. * @return Column of extracted elements. * @throws cudf::logic_error If the sizes of `lists_column` and `indices` do not match. @@ -104,6 +107,7 @@ std::unique_ptr extract_list_element( std::unique_ptr extract_list_element( lists_column_view const& lists_column, column_view const& indices, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @} */ // end of group diff --git a/cpp/include/cudf/lists/filling.hpp b/cpp/include/cudf/lists/filling.hpp index 059ed5ffd33..3730e16482d 100644 --- a/cpp/include/cudf/lists/filling.hpp +++ b/cpp/include/cudf/lists/filling.hpp @@ -17,7 +17,9 @@ #pragma once #include +#include +#include #include #include @@ -57,12 +59,14 @@ namespace cudf::lists { * * @param starts First values in the result sequences. * @param sizes Numbers of values in the result sequences. + * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned column's device memory. * @return The result column containing generated sequences. */ std::unique_ptr sequences( column_view const& starts, column_view const& sizes, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @@ -96,6 +100,7 @@ std::unique_ptr sequences( * @param starts First values in the result sequences. * @param steps Increment values for the result sequences. * @param sizes Numbers of values in the result sequences. + * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned column's device memory. * @return The result column containing generated sequences. */ @@ -103,6 +108,7 @@ std::unique_ptr sequences( column_view const& starts, column_view const& steps, column_view const& sizes, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @} */ // end of group diff --git a/cpp/include/cudf/lists/gather.hpp b/cpp/include/cudf/lists/gather.hpp index 38bed9ede43..5e6ab6816e6 100644 --- a/cpp/include/cudf/lists/gather.hpp +++ b/cpp/include/cudf/lists/gather.hpp @@ -65,6 +65,7 @@ namespace lists { * @param bounds_policy Can be `DONT_CHECK` or `NULLIFY`. Selects whether or not to nullify the * output list row's element, when the gather index falls outside the range `[-n, n)`, * where `n` is the number of elements in list row corresponding to the gather-map row. + * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource to allocate any returned objects * @return column with elements in list of rows gathered based on `gather_map_list` * @@ -73,6 +74,7 @@ std::unique_ptr segmented_gather( lists_column_view const& source_column, lists_column_view const& gather_map_list, out_of_bounds_policy bounds_policy = out_of_bounds_policy::DONT_CHECK, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @} */ // end of group diff --git a/cpp/include/cudf/lists/reverse.hpp b/cpp/include/cudf/lists/reverse.hpp index 226d417c53a..864cd796f72 100644 --- a/cpp/include/cudf/lists/reverse.hpp +++ b/cpp/include/cudf/lists/reverse.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022, NVIDIA CORPORATION. + * Copyright (c) 2022-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -42,11 +42,13 @@ namespace cudf::lists { * @endcode * * @param input Lists column for this operation + * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used to allocate the returned column's device memory * @return New lists column with reversed lists */ std::unique_ptr reverse( lists_column_view const& input, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @} */ // end of doxygen group diff --git a/cpp/include/cudf/lists/sorting.hpp b/cpp/include/cudf/lists/sorting.hpp index c203c452b0d..39a52c75a98 100644 --- a/cpp/include/cudf/lists/sorting.hpp +++ b/cpp/include/cudf/lists/sorting.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021-2022, NVIDIA CORPORATION. + * Copyright (c) 2021-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -46,6 +46,7 @@ namespace lists { * @param source_column View of the list column of numeric types to sort * @param column_order The desired sort order * @param null_precedence The desired order of null compared to other elements in the list + * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource to allocate any returned objects * @return list column with elements in each list sorted. * @@ -54,6 +55,7 @@ std::unique_ptr sort_lists( lists_column_view const& source_column, order column_order, null_order null_precedence, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @@ -66,6 +68,7 @@ std::unique_ptr stable_sort_lists( lists_column_view const& source_column, order column_order, null_order null_precedence, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @} */ // end of group diff --git a/cpp/include/cudf/lists/stream_compaction.hpp b/cpp/include/cudf/lists/stream_compaction.hpp index 5ddaa992184..3ac4f6861ec 100644 --- a/cpp/include/cudf/lists/stream_compaction.hpp +++ b/cpp/include/cudf/lists/stream_compaction.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2022, NVIDIA CORPORATION. + * Copyright (c) 2022-2023, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -54,12 +54,14 @@ namespace cudf::lists { * * @param input The input list column view to be filtered * @param boolean_mask A nullable list of bools column used to filter `input` elements + * @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 * @return List column of the same type as `input`, containing filtered list rows */ std::unique_ptr apply_boolean_mask( lists_column_view const& input, lists_column_view const& boolean_mask, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @@ -78,6 +80,7 @@ std::unique_ptr apply_boolean_mask( * @param input The input lists column * @param nulls_equal Flag to specify whether null elements should be considered as equal * @param nans_equal Flag to specify whether floating-point NaNs should be considered as equal + * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used to allocate the returned object * @return The resulting lists column containing lists without duplicates */ @@ -85,6 +88,7 @@ std::unique_ptr distinct( lists_column_view const& input, null_equality nulls_equal = null_equality::EQUAL, nan_equality nans_equal = nan_equality::ALL_EQUAL, + rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()); /** @} */ // end of group diff --git a/cpp/src/lists/copying/segmented_gather.cu b/cpp/src/lists/copying/segmented_gather.cu index 79d33e7c17d..855ceadf33f 100644 --- a/cpp/src/lists/copying/segmented_gather.cu +++ b/cpp/src/lists/copying/segmented_gather.cu @@ -116,11 +116,11 @@ std::unique_ptr segmented_gather(lists_column_view const& value_column, std::unique_ptr segmented_gather(lists_column_view const& source_column, lists_column_view const& gather_map_list, out_of_bounds_policy bounds_policy, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::segmented_gather( - source_column, gather_map_list, bounds_policy, cudf::get_default_stream(), mr); + return detail::segmented_gather(source_column, gather_map_list, bounds_policy, stream, mr); } } // namespace lists diff --git a/cpp/src/lists/extract.cu b/cpp/src/lists/extract.cu index 5d4a20d1cb8..365e9ef8255 100644 --- a/cpp/src/lists/extract.cu +++ b/cpp/src/lists/extract.cu @@ -196,10 +196,11 @@ std::unique_ptr extract_list_element(lists_column_view lists_column, */ std::unique_ptr extract_list_element(lists_column_view const& lists_column, size_type index, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::extract_list_element(lists_column, index, cudf::get_default_stream(), mr); + return detail::extract_list_element(lists_column, index, stream, mr); } /** @@ -209,12 +210,13 @@ std::unique_ptr extract_list_element(lists_column_view const& lists_colu */ std::unique_ptr extract_list_element(lists_column_view const& lists_column, column_view const& indices, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); CUDF_EXPECTS(indices.size() == lists_column.size(), "Index column must have as many elements as lists column."); - return detail::extract_list_element(lists_column, indices, cudf::get_default_stream(), mr); + return detail::extract_list_element(lists_column, indices, stream, mr); } } // namespace lists diff --git a/cpp/src/lists/reverse.cu b/cpp/src/lists/reverse.cu index a2af85b5dad..6c00f8b64b4 100644 --- a/cpp/src/lists/reverse.cu +++ b/cpp/src/lists/reverse.cu @@ -86,10 +86,12 @@ std::unique_ptr reverse(lists_column_view const& input, } // namespace detail -std::unique_ptr reverse(lists_column_view const& input, rmm::mr::device_memory_resource* mr) +std::unique_ptr reverse(lists_column_view const& input, + rmm::cuda_stream_view stream, + rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::reverse(input, cudf::get_default_stream(), mr); + return detail::reverse(input, stream, mr); } } // namespace cudf::lists diff --git a/cpp/src/lists/segmented_sort.cu b/cpp/src/lists/segmented_sort.cu index 49054ebb046..0b70773f4b2 100644 --- a/cpp/src/lists/segmented_sort.cu +++ b/cpp/src/lists/segmented_sort.cu @@ -119,20 +119,21 @@ std::unique_ptr stable_sort_lists(lists_column_view const& input, std::unique_ptr sort_lists(lists_column_view const& input, order column_order, null_order null_precedence, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::sort_lists(input, column_order, null_precedence, cudf::get_default_stream(), mr); + return detail::sort_lists(input, column_order, null_precedence, stream, mr); } std::unique_ptr stable_sort_lists(lists_column_view const& input, order column_order, null_order null_precedence, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::stable_sort_lists( - input, column_order, null_precedence, cudf::get_default_stream(), mr); + return detail::stable_sort_lists(input, column_order, null_precedence, stream, mr); } } // namespace lists diff --git a/cpp/src/lists/sequences.cu b/cpp/src/lists/sequences.cu index aaee5608cc3..f92ba782da7 100644 --- a/cpp/src/lists/sequences.cu +++ b/cpp/src/lists/sequences.cu @@ -208,19 +208,21 @@ std::unique_ptr sequences(column_view const& starts, std::unique_ptr sequences(column_view const& starts, column_view const& sizes, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::sequences(starts, sizes, cudf::get_default_stream(), mr); + return detail::sequences(starts, sizes, stream, mr); } std::unique_ptr sequences(column_view const& starts, column_view const& steps, column_view const& sizes, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::sequences(starts, steps, sizes, cudf::get_default_stream(), mr); + return detail::sequences(starts, steps, sizes, stream, mr); } } // namespace cudf::lists diff --git a/cpp/src/lists/stream_compaction/apply_boolean_mask.cu b/cpp/src/lists/stream_compaction/apply_boolean_mask.cu index ad43fbd5b00..ce972d89150 100644 --- a/cpp/src/lists/stream_compaction/apply_boolean_mask.cu +++ b/cpp/src/lists/stream_compaction/apply_boolean_mask.cu @@ -101,10 +101,11 @@ std::unique_ptr apply_boolean_mask(lists_column_view const& input, std::unique_ptr apply_boolean_mask(lists_column_view const& input, lists_column_view const& boolean_mask, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::apply_boolean_mask(input, boolean_mask, cudf::get_default_stream(), mr); + return detail::apply_boolean_mask(input, boolean_mask, stream, mr); } } // namespace cudf::lists diff --git a/cpp/src/lists/stream_compaction/distinct.cu b/cpp/src/lists/stream_compaction/distinct.cu index 48d8babb4fa..eb21787b3fa 100644 --- a/cpp/src/lists/stream_compaction/distinct.cu +++ b/cpp/src/lists/stream_compaction/distinct.cu @@ -76,10 +76,11 @@ std::unique_ptr distinct(lists_column_view const& input, std::unique_ptr distinct(lists_column_view const& input, null_equality nulls_equal, nan_equality nans_equal, + rmm::cuda_stream_view stream, rmm::mr::device_memory_resource* mr) { CUDF_FUNC_RANGE(); - return detail::distinct(input, nulls_equal, nans_equal, cudf::get_default_stream(), mr); + return detail::distinct(input, nulls_equal, nans_equal, stream, mr); } } // namespace cudf::lists diff --git a/cpp/tests/groupby/histogram_tests.cpp b/cpp/tests/groupby/histogram_tests.cpp index c5833f40cf2..612486d8e5c 100644 --- a/cpp/tests/groupby/histogram_tests.cpp +++ b/cpp/tests/groupby/histogram_tests.cpp @@ -67,6 +67,7 @@ auto groupby_histogram(cudf::column_view const& keys, auto sorted_histograms = cudf::lists::sort_lists(cudf::lists_column_view{*sorted_vals}, cudf::order::ASCENDING, cudf::null_order::BEFORE, + cudf::get_default_stream(), rmm::mr::get_current_device_resource()); return std::pair{std::move(sorted_keys), std::move(sorted_histograms)}; diff --git a/cpp/tests/streams/lists_test.cpp b/cpp/tests/streams/lists_test.cpp index e292b551d83..82a4cb8aa4a 100644 --- a/cpp/tests/streams/lists_test.cpp +++ b/cpp/tests/streams/lists_test.cpp @@ -21,6 +21,12 @@ #include #include #include +#include +#include +#include +#include +#include +#include class ListTest : public cudf::test::BaseFixture {}; @@ -85,3 +91,78 @@ TEST_F(ListTest, CountElements) cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7}, {4, 5}}; cudf::lists::count_elements(list_col, cudf::test::get_default_stream()); } + +TEST_F(ListTest, ExtractListElementFromIndex) +{ + cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7}, {4, 5}}; + cudf::lists::extract_list_element(list_col, -1, cudf::test::get_default_stream()); +} + +TEST_F(ListTest, ExtractListElementFromIndices) +{ + cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7}, {4, 5}}; + cudf::test::fixed_width_column_wrapper indices({-1, -2, -1}); + cudf::lists::extract_list_element(list_col, indices, cudf::test::get_default_stream()); +} + +TEST_F(ListTest, SegmentedGather) +{ + cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7, 8}, {4, 5}}; + cudf::test::lists_column_wrapper gather_map_list{{0}, {1, 2}, {1}}; + cudf::lists::segmented_gather(list_col, + gather_map_list, + cudf::out_of_bounds_policy::DONT_CHECK, + cudf::test::get_default_stream()); +} + +TEST_F(ListTest, Sequences) +{ + cudf::test::fixed_width_column_wrapper starts({0, 1, 2, 3, 4}); + cudf::test::fixed_width_column_wrapper sizes({0, 1, 2, 2, 1}); + cudf::lists::sequences(starts, sizes, cudf::test::get_default_stream()); +} + +TEST_F(ListTest, SequencesWithSteps) +{ + cudf::test::fixed_width_column_wrapper starts({0, 1, 2, 3, 4}); + cudf::test::fixed_width_column_wrapper steps({2, 1, 1, 1, -3}); + cudf::test::fixed_width_column_wrapper sizes({0, 1, 2, 2, 1}); + cudf::lists::sequences(starts, steps, sizes, cudf::test::get_default_stream()); +} + +TEST_F(ListTest, Reverse) +{ + cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7, 8}, {4, 5}}; + cudf::lists::reverse(list_col, cudf::test::get_default_stream()); +} + +TEST_F(ListTest, SortLists) +{ + cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7, 8}, {4, 5}}; + cudf::lists::sort_lists( + list_col, cudf::order::DESCENDING, cudf::null_order::AFTER, cudf::test::get_default_stream()); +} + +TEST_F(ListTest, StableSortLists) +{ + cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7, 8}, {4, 5}}; + cudf::lists::stable_sort_lists( + list_col, cudf::order::DESCENDING, cudf::null_order::AFTER, cudf::test::get_default_stream()); +} + +TEST_F(ListTest, ApplyBooleanMask) +{ + cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7, 8}, {4, 5}}; + cudf::test::lists_column_wrapper boolean_mask{{0, 1}, {1, 1, 1, 0}, {0, 1}}; + cudf::lists::apply_boolean_mask(list_col, boolean_mask, cudf::test::get_default_stream()); +} + +TEST_F(ListTest, Distinct) +{ + cudf::test::lists_column_wrapper list_col{{0, 1}, {2, 3, 7, 8}, {4, 5}}; + cudf::test::lists_column_wrapper boolean_mask{{0, 1}, {1, 1, 1, 0}, {0, 1}}; + cudf::lists::distinct(list_col, + cudf::null_equality::EQUAL, + cudf::nan_equality::ALL_EQUAL, + cudf::test::get_default_stream()); +}