Add stream param to stream compaction APIs

cpp/include/cudf/detail/stream_compaction.hpp

@@ -29,9 +29,7 @@ namespace cudf {
 namespace detail {
 /**
  * @copydoc cudf::drop_nulls(table_view const&, std::vector<size_type> const&,
- * cudf::size_type, rmm::device_async_resource_ref)
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ * cudf::size_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)
  */
 std::unique_ptr<table> drop_nulls(table_view const& input,
  std::vector<size_type> const& keys,
@@ -41,9 +39,7 @@ std::unique_ptr<table> drop_nulls(table_view const& input,
 
 /**
  * @copydoc cudf::drop_nans(table_view const&, std::vector<size_type> const&,
- * cudf::size_type, rmm::device_async_resource_ref)
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ * cudf::size_type, rmm::cuda_stream_view, rmm::device_async_resource_ref)
  */
 std::unique_ptr<table> drop_nans(table_view const& input,
  std::vector<size_type> const& keys,
@@ -53,8 +49,6 @@ std::unique_ptr<table> drop_nans(table_view const& input,
 
 /**
  * @copydoc cudf::apply_boolean_mask
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
  */
 std::unique_ptr<table> apply_boolean_mask(table_view const& input,
  column_view const& boolean_mask,
@@ -63,8 +57,6 @@ std::unique_ptr<table> apply_boolean_mask(table_view const& input,
 
 /**
  * @copydoc cudf::unique
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
  */
 std::unique_ptr<table> unique(table_view const& input,
  std::vector<size_type> const& keys,
@@ -75,8 +67,6 @@ std::unique_ptr<table> unique(table_view const& input,
 
 /**
  * @copydoc cudf::distinct
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
  */
 std::unique_ptr<table> distinct(table_view const& input,
  std::vector<size_type> const& keys,
@@ -110,38 +100,30 @@ rmm::device_uvector<size_type> distinct_indices(table_view const& input,
  rmm::device_async_resource_ref mr);
 
 /**
- * @copydoc cudf::unique_count(column_view const&, null_policy, nan_policy)
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ * @copydoc cudf::unique_count(column_view const&, null_policy, nan_policy, rmm::cuda_stream_view)
  */
 cudf::size_type unique_count(column_view const& input,
  null_policy null_handling,
  nan_policy nan_handling,
  rmm::cuda_stream_view stream);
 
 /**
- * @copydoc cudf::unique_count(table_view const&, null_equality)
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ * @copydoc cudf::unique_count(table_view const&, null_equality, rmm::cuda_stream_view)
  */
 cudf::size_type unique_count(table_view const& input,
  null_equality nulls_equal,
  rmm::cuda_stream_view stream);
 
 /**
- * @copydoc cudf::distinct_count(column_view const&, null_policy, nan_policy)
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ * @copydoc cudf::distinct_count(column_view const&, null_policy, nan_policy, rmm::cuda_stream_view)
  */
 cudf::size_type distinct_count(column_view const& input,
  null_policy null_handling,
  nan_policy nan_handling,
  rmm::cuda_stream_view stream);
 
 /**
- * @copydoc cudf::distinct_count(table_view const&, null_equality)
- *
- * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ * @copydoc cudf::distinct_count(table_view const&, null_equality, rmm::cuda_stream_view)
  */
 cudf::size_type distinct_count(table_view const& input,
  null_equality nulls_equal,

cpp/include/cudf/lists/detail/stream_compaction.hpp

@@ -24,20 +24,15 @@
 namespace cudf::lists::detail {
 
 /**
- * @copydoc cudf::lists::apply_boolean_mask(lists_column_view const&, lists_column_view const&,
- * rmm::device_async_resource_ref)
- *
- * @param stream CUDA stream used for device memory operations and kernel launches
+ * @copydoc cudf::lists::apply_boolean_mask
  */
 std::unique_ptr<column> apply_boolean_mask(lists_column_view const& input,
  lists_column_view const& boolean_mask,
  rmm::cuda_stream_view stream,
  rmm::device_async_resource_ref mr);
 
 /**
- * @copydoc cudf::list::distinct
- *
- * @param stream CUDA stream used for device memory operations and kernel launches.
+ * @copydoc cudf::lists::distinct
  */
 std::unique_ptr<column> distinct(lists_column_view const& input,
  null_equality nulls_equal,

cpp/include/cudf/stream_compaction.hpp

@@ -66,6 +66,7 @@ namespace cudf {
  * @param[in] keys vector of indices representing key columns from `input`
  * @param[in] keep_threshold The minimum number of non-null fields in a row
  * required to keep the row.
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  * @param[in] mr Device memory resource used to allocate the returned table's device memory
  * @return Table containing all rows of the `input` with at least @p
  * keep_threshold non-null fields in @p keys.
@@ -74,6 +75,7 @@ std::unique_ptr<table> drop_nulls(
  table_view const& input,
  std::vector<size_type> const& keys,
  cudf::size_type keep_threshold,
+ rmm::cuda_stream_view stream = cudf::get_default_stream(),
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -98,13 +100,15 @@ std::unique_ptr<table> drop_nulls(
  *
  * @param[in] input The input `table_view` to filter
  * @param[in] keys vector of indices representing key columns from `input`
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  * @param[in] mr Device memory resource used to allocate the returned table's device memory
  * @return Table containing all rows of the `input` without nulls in the columns
  * of @p keys.
  */
 std::unique_ptr<table> drop_nulls(
  table_view const& input,
  std::vector<size_type> const& keys,
+ rmm::cuda_stream_view stream = cudf::get_default_stream(),
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -140,6 +144,7 @@ std::unique_ptr<table> drop_nulls(
  * @param[in] keys vector of indices representing key columns from `input`
  * @param[in] keep_threshold The minimum number of non-NAN elements in a row
  * required to keep the row.
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  * @param[in] mr Device memory resource used to allocate the returned table's device memory
  * @return Table containing all rows of the `input` with at least @p
  * keep_threshold non-NAN elements in @p keys.
@@ -148,6 +153,7 @@ std::unique_ptr<table> drop_nans(
  table_view const& input,
  std::vector<size_type> const& keys,
  cudf::size_type keep_threshold,
+ rmm::cuda_stream_view stream = cudf::get_default_stream(),
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -173,13 +179,15 @@ std::unique_ptr<table> drop_nans(
  *
  * @param[in] input The input `table_view` to filter
  * @param[in] keys vector of indices representing key columns from `input`
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  * @param[in] mr Device memory resource used to allocate the returned table's device memory
  * @return Table containing all rows of the `input` without NANs in the columns
  * of @p keys.
  */
 std::unique_ptr<table> drop_nans(
  table_view const& input,
  std::vector<size_type> const& keys,
+ rmm::cuda_stream_view stream = cudf::get_default_stream(),
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -199,13 +207,15 @@ std::unique_ptr<table> drop_nans(
  * @param[in] input The input table_view to filter
  * @param[in] boolean_mask A nullable column_view of type type_id::BOOL8 used
  * as a mask to filter the `input`.
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  * @param[in] mr Device memory resource used to allocate the returned table's device memory
  * @return Table containing copy of all rows of @p input passing
  * the filter defined by @p boolean_mask.
  */
 std::unique_ptr<table> apply_boolean_mask(
  table_view const& input,
  column_view const& boolean_mask,
+ rmm::cuda_stream_view stream = cudf::get_default_stream(),
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -240,6 +250,7 @@ enum class duplicate_keep_option {
  * @param[in] keep keep any, first, last, or none of the found duplicates
  * @param[in] nulls_equal flag to denote nulls are equal if null_equality::EQUAL, nulls are not
  * equal if null_equality::UNEQUAL
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  * @param[in] mr Device memory resource used to allocate the returned table's device
  * memory
  *
@@ -250,6 +261,7 @@ std::unique_ptr<table> unique(
  std::vector<size_type> const& keys,
  duplicate_keep_option keep,
  null_equality nulls_equal = null_equality::EQUAL,
+ rmm::cuda_stream_view stream = cudf::get_default_stream(),
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -268,6 +280,7 @@ std::unique_ptr<table> unique(
  * @param keep Copy any, first, last, or none of the found duplicates
  * @param nulls_equal Flag to specify whether null elements should be considered as equal
  * @param nans_equal Flag to specify whether NaN elements 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 table
  * @return Table with distinct rows in an unspecified order
  */
@@ -277,6 +290,7 @@ std::unique_ptr<table> distinct(
  duplicate_keep_option keep = duplicate_keep_option::KEEP_ANY,
  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::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -345,24 +359,28 @@ std::unique_ptr<table> stable_distinct(
  * @param[in] input The column_view whose consecutive groups of equivalent rows will be counted
  * @param[in] null_handling flag to include or ignore `null` while counting
  * @param[in] nan_handling flag to consider `NaN==null` or not
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  *
  * @return number of consecutive groups of equivalent rows in the column
  */
 cudf::size_type unique_count(column_view const& input,
  null_policy null_handling,
- nan_policy nan_handling);
+ nan_policy nan_handling,
+ rmm::cuda_stream_view stream = cudf::get_default_stream());
 
 /**
  * @brief Count the number of consecutive groups of equivalent rows in a table.
  *
  * @param[in] input Table whose consecutive groups of equivalent rows will be counted
  * @param[in] nulls_equal flag to denote if null elements should be considered equal
  * nulls are not equal if null_equality::UNEQUAL.
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  *
  * @return number of consecutive groups of equivalent rows in the column
  */
 cudf::size_type unique_count(table_view const& input,
- null_equality nulls_equal = null_equality::EQUAL);
+ null_equality nulls_equal = null_equality::EQUAL,
+ rmm::cuda_stream_view stream = cudf::get_default_stream());
 
 /**
  * @brief Count the distinct elements in the column_view.
@@ -381,24 +399,28 @@ cudf::size_type unique_count(table_view const& input,
  * @param[in] input The column_view whose distinct elements will be counted
  * @param[in] null_handling flag to include or ignore `null` while counting
  * @param[in] nan_handling flag to consider `NaN==null` or not
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  *
  * @return number of distinct rows in the table
  */
 cudf::size_type distinct_count(column_view const& input,
  null_policy null_handling,
- nan_policy nan_handling);
+ nan_policy nan_handling,
+ rmm::cuda_stream_view stream = cudf::get_default_stream());
 
 /**
  * @brief Count the distinct rows in a table.
  *
  * @param[in] input Table whose distinct rows will be counted
  * @param[in] nulls_equal flag to denote if null elements should be considered equal.
  * nulls are not equal if null_equality::UNEQUAL.
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches
  *
  * @return number of distinct rows in the table
  */
 cudf::size_type distinct_count(table_view const& input,
- null_equality nulls_equal = null_equality::EQUAL);
+ null_equality nulls_equal = null_equality::EQUAL,
+ rmm::cuda_stream_view stream = cudf::get_default_stream());
 
 /** @} */
 } // namespace cudf

cpp/src/stream_compaction/apply_boolean_mask.cu

@@ -91,9 +91,10 @@ std::unique_ptr<table> apply_boolean_mask(table_view const& input,
  */
 std::unique_ptr<table> apply_boolean_mask(table_view const& input,
  column_view const& boolean_mask,
+ rmm::cuda_stream_view stream,
  rmm::device_async_resource_ref 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

cpp/src/stream_compaction/distinct.cu

@@ -149,11 +149,11 @@ std::unique_ptr<table> distinct(table_view const& input,
  duplicate_keep_option keep,
  null_equality nulls_equal,
  nan_equality nans_equal,
+ rmm::cuda_stream_view stream,
  rmm::device_async_resource_ref mr)
 {
  CUDF_FUNC_RANGE();
- return detail::distinct(
- input, keys, keep, nulls_equal, nans_equal, cudf::get_default_stream(), mr);
+ return detail::distinct(input, keys, keep, nulls_equal, nans_equal, stream, mr);
 }
 
 std::unique_ptr<column> distinct_indices(table_view const& input,

cpp/src/stream_compaction/distinct_count.cu

@@ -217,15 +217,18 @@ cudf::size_type distinct_count(column_view const& input,
 
 cudf::size_type distinct_count(column_view const& input,
  null_policy null_handling,
- nan_policy nan_handling)
+ nan_policy nan_handling,
+ rmm::cuda_stream_view stream)
 {
  CUDF_FUNC_RANGE();
- return detail::distinct_count(input, null_handling, nan_handling, cudf::get_default_stream());
+ return detail::distinct_count(input, null_handling, nan_handling, stream);
 }
 
-cudf::size_type distinct_count(table_view const& input, null_equality nulls_equal)
+cudf::size_type distinct_count(table_view const& input,
+ null_equality nulls_equal,
+ rmm::cuda_stream_view stream)
 {
  CUDF_FUNC_RANGE();
- return detail::distinct_count(input, nulls_equal, cudf::get_default_stream());
+ return detail::distinct_count(input, nulls_equal, stream);
 }
 } // namespace cudf

cpp/src/stream_compaction/drop_nans.cu

@@ -117,20 +117,22 @@ std::unique_ptr<table> drop_nans(table_view const& input,
 std::unique_ptr<table> drop_nans(table_view const& input,
  std::vector<size_type> const& keys,
  cudf::size_type keep_threshold,
+ rmm::cuda_stream_view stream,
  rmm::device_async_resource_ref mr)
 {
  CUDF_FUNC_RANGE();
- return detail::drop_nans(input, keys, keep_threshold, cudf::get_default_stream(), mr);
+ return detail::drop_nans(input, keys, keep_threshold, stream, mr);
 }
 /*
  * Filters a table to remove nan elements.
  */
 std::unique_ptr<table> drop_nans(table_view const& input,
  std::vector<size_type> const& keys,
+ rmm::cuda_stream_view stream,
  rmm::device_async_resource_ref mr)
 {
  CUDF_FUNC_RANGE();
- return detail::drop_nans(input, keys, keys.size(), cudf::get_default_stream(), mr);
+ return detail::drop_nans(input, keys, keys.size(), stream, mr);
 }
 
 } // namespace cudf

cpp/src/stream_compaction/drop_nulls.cu

@@ -90,20 +90,22 @@ std::unique_ptr<table> drop_nulls(table_view const& input,
 std::unique_ptr<table> drop_nulls(table_view const& input,
  std::vector<size_type> const& keys,
  cudf::size_type keep_threshold,
+ rmm::cuda_stream_view stream,
  rmm::device_async_resource_ref mr)
 {
  CUDF_FUNC_RANGE();
- return detail::drop_nulls(input, keys, keep_threshold, cudf::get_default_stream(), mr);
+ return detail::drop_nulls(input, keys, keep_threshold, stream, mr);
 }
 /*
  * Filters a table to remove null elements.
  */
 std::unique_ptr<table> drop_nulls(table_view const& input,
  std::vector<size_type> const& keys,
+ rmm::cuda_stream_view stream,
  rmm::device_async_resource_ref mr)
 {
  CUDF_FUNC_RANGE();
- return detail::drop_nulls(input, keys, keys.size(), cudf::get_default_stream(), mr);
+ return detail::drop_nulls(input, keys, keys.size(), stream, mr);
 }
 
 } // namespace cudf