Add stream parameter to Set Operations (Public List APIs) (#14305)

This PR marks the conclusion of the List API series. The PR introduces the stream parameter to the Set operations (Public List Comparison and Intersection APIs). Comparison and Intersection (`set_operations.hpp`) ``` have_overlap intersect_distinct union_distinct difference_distinct ``` Reference [13744](#13744) Authors: - Suraj Aralihalli (https://github.com/SurajAralihalli) - Nghia Truong (https://github.com/ttnghia) Approvers: - Nghia Truong (https://github.com/ttnghia) - Mike Wilson (https://github.com/hyperbolic2346) URL: #14305
rapidsai · Oct 27, 2023 · f6099ca · f6099ca
1 parent d8f0790
commit f6099ca
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 7 deletions.
diff --git a/cpp/benchmarks/lists/set_operations.cpp b/cpp/benchmarks/lists/set_operations.cpp
@@ -54,6 +54,7 @@ void nvbench_set_op(nvbench::state& state, BenchFuncPtr bfunc)
           cudf::lists_column_view{*rhs},
           cudf::null_equality::EQUAL,
           cudf::nan_equality::ALL_EQUAL,
+          cudf::get_default_stream(),
           rmm::mr::get_current_device_resource());
   });
 }

diff --git a/cpp/include/cudf/lists/set_operations.hpp b/cpp/include/cudf/lists/set_operations.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.
@@ -53,13 +53,15 @@ namespace cudf::lists {
  *        to be `UNEQUAL` which means only non-null elements are checked for overlapping
  * @param nans_equal Flag to specify whether floating-point NaNs should be considered as equal
  * @param mr Device memory resource used to allocate the returned object
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @return A column of type BOOL containing the check results
  */
 std::unique_ptr<column> have_overlap(
   lists_column_view const& lhs,
   lists_column_view const& rhs,
   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());
 
 /**
@@ -87,6 +89,7 @@ std::unique_ptr<column> have_overlap(
  * @param rhs The input lists column for the other side
  * @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 A lists column containing the intersection results
  */
@@ -95,6 +98,7 @@ std::unique_ptr<column> intersect_distinct(
   lists_column_view const& rhs,
   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());
 
 /**
@@ -122,6 +126,7 @@ std::unique_ptr<column> intersect_distinct(
  * @param rhs The input lists column for the other side
  * @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 A lists column containing the union results
  */
@@ -130,6 +135,7 @@ std::unique_ptr<column> union_distinct(
   lists_column_view const& rhs,
   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());
 
 /**
@@ -157,6 +163,7 @@ std::unique_ptr<column> union_distinct(
  * @param rhs The input lists column of elements to exclude
  * @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 A lists column containing the difference results
  */
@@ -165,6 +172,7 @@ std::unique_ptr<column> difference_distinct(
   lists_column_view const& rhs,
   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/set_operations.cu b/cpp/src/lists/set_operations.cu
@@ -278,42 +278,44 @@ std::unique_ptr<column> have_overlap(lists_column_view const& lhs,
                                      lists_column_view const& rhs,
                                      null_equality nulls_equal,
                                      nan_equality nans_equal,
+                                     rmm::cuda_stream_view stream,
                                      rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::have_overlap(lhs, rhs, nulls_equal, nans_equal, cudf::get_default_stream(), mr);
+  return detail::have_overlap(lhs, rhs, nulls_equal, nans_equal, stream, mr);
 }
 
 std::unique_ptr<column> intersect_distinct(lists_column_view const& lhs,
                                            lists_column_view const& rhs,
                                            null_equality nulls_equal,
                                            nan_equality nans_equal,
+                                           rmm::cuda_stream_view stream,
                                            rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::intersect_distinct(
-    lhs, rhs, nulls_equal, nans_equal, cudf::get_default_stream(), mr);
+  return detail::intersect_distinct(lhs, rhs, nulls_equal, nans_equal, stream, mr);
 }
 
 std::unique_ptr<column> union_distinct(lists_column_view const& lhs,
                                        lists_column_view const& rhs,
                                        null_equality nulls_equal,
                                        nan_equality nans_equal,
+                                       rmm::cuda_stream_view stream,
                                        rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::union_distinct(lhs, rhs, nulls_equal, nans_equal, cudf::get_default_stream(), mr);
+  return detail::union_distinct(lhs, rhs, nulls_equal, nans_equal, stream, mr);
 }
 
 std::unique_ptr<column> difference_distinct(lists_column_view const& lhs,
                                             lists_column_view const& rhs,
                                             null_equality nulls_equal,
                                             nan_equality nans_equal,
+                                            rmm::cuda_stream_view stream,
                                             rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::difference_distinct(
-    lhs, rhs, nulls_equal, nans_equal, cudf::get_default_stream(), mr);
+  return detail::difference_distinct(lhs, rhs, nulls_equal, nans_equal, stream, mr);
 }
 
 }  // namespace cudf::lists
diff --git a/cpp/tests/streams/lists_test.cpp b/cpp/tests/streams/lists_test.cpp
@@ -25,6 +25,7 @@
 #include <cudf/lists/filling.hpp>
 #include <cudf/lists/gather.hpp>
 #include <cudf/lists/reverse.hpp>
+#include <cudf/lists/set_operations.hpp>
 #include <cudf/lists/sorting.hpp>
 #include <cudf/lists/stream_compaction.hpp>
 
@@ -166,3 +167,47 @@ TEST_F(ListTest, Distinct)
                         cudf::nan_equality::ALL_EQUAL,
                         cudf::test::get_default_stream());
 }
+
+TEST_F(ListTest, DifferenceDistinct)
+{
+  cudf::test::lists_column_wrapper<int> list_col_a{{0, 1}, {2, 3, 7, 8}, {4, 5}};
+  cudf::test::lists_column_wrapper<int> list_col_b{{0, 1}, {1, 3, 6, 8}, {5}};
+  cudf::lists::difference_distinct(list_col_a,
+                                   list_col_b,
+                                   cudf::null_equality::EQUAL,
+                                   cudf::nan_equality::ALL_EQUAL,
+                                   cudf::test::get_default_stream());
+}
+
+TEST_F(ListTest, IntersectDistinct)
+{
+  cudf::test::lists_column_wrapper<int> list_col_a{{0, 1}, {2, 3, 7, 8}, {4, 5}};
+  cudf::test::lists_column_wrapper<int> list_col_b{{0, 1}, {1, 3, 6, 8}, {5}};
+  cudf::lists::intersect_distinct(list_col_a,
+                                  list_col_b,
+                                  cudf::null_equality::EQUAL,
+                                  cudf::nan_equality::ALL_EQUAL,
+                                  cudf::test::get_default_stream());
+}
+
+TEST_F(ListTest, UnionDistinct)
+{
+  cudf::test::lists_column_wrapper<int> list_col_a{{0, 1}, {2, 3, 7, 8}, {4, 5}};
+  cudf::test::lists_column_wrapper<int> list_col_b{{0, 1}, {1, 3, 6, 8}, {5}};
+  cudf::lists::union_distinct(list_col_a,
+                              list_col_b,
+                              cudf::null_equality::EQUAL,
+                              cudf::nan_equality::ALL_EQUAL,
+                              cudf::test::get_default_stream());
+}
+
+TEST_F(ListTest, HaveOverlap)
+{
+  cudf::test::lists_column_wrapper<int> list_col_a{{0, 1}, {2, 3, 7, 8}, {4, 5}};
+  cudf::test::lists_column_wrapper<int> list_col_b{{0, 1}, {1, 3, 6, 8}, {5}};
+  cudf::lists::have_overlap(list_col_a,
+                            list_col_b,
+                            cudf::null_equality::EQUAL,
+                            cudf::nan_equality::ALL_EQUAL,
+                            cudf::test::get_default_stream());
+}