Expose streams in public search APIs (#14034)

Contributes to #925 Authors: - Vyas Ramasubramani (https://github.com/vyasr) Approvers: - David Wendt (https://github.com/davidwendt) - Nghia Truong (https://github.com/ttnghia) URL: #14034
rapidsai · Sep 6, 2023 · e81d79e · e81d79e
1 parent ea59dbf
commit e81d79e
Show file tree

Hide file tree

Showing 6 changed files with 90 additions and 10 deletions.
diff --git a/cpp/include/cudf/search.hpp b/cpp/include/cudf/search.hpp
@@ -63,6 +63,7 @@ namespace cudf {
  * @param needles Values for which to find the insert locations in the search space
  * @param column_order Vector of column sort order
  * @param null_precedence Vector of null_precedence enums needles
+ * @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 A non-nullable column of elements containing the insertion points
  */
@@ -71,6 +72,7 @@ std::unique_ptr<column> lower_bound(
   table_view const& needles,
   std::vector<order> const& column_order,
   std::vector<null_order> const& null_precedence,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -103,6 +105,7 @@ std::unique_ptr<column> lower_bound(
  * @param needles Values for which to find the insert locations in the search space
  * @param column_order Vector of column sort order
  * @param null_precedence Vector of null_precedence enums needles
+ * @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 A non-nullable column of elements containing the insertion points
  */
@@ -111,6 +114,7 @@ std::unique_ptr<column> upper_bound(
   table_view const& needles,
   std::vector<order> const& column_order,
   std::vector<null_order> const& null_precedence,
+  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
 /**
@@ -128,9 +132,12 @@ std::unique_ptr<column> upper_bound(
  *
  * @param haystack The column containing search space
  * @param needle A scalar value to check for existence in the search space
+ * @param stream CUDA stream used for device memory operations and kernel launches
  * @return true if the given `needle` value exists in the `haystack` column
  */
-bool contains(column_view const& haystack, scalar const& needle);
+bool contains(column_view const& haystack,
+              scalar const& needle,
+              rmm::cuda_stream_view stream = cudf::get_default_stream());
 
 /**
  * @brief Check if the given `needles` values exists in the `haystack` column.
@@ -149,12 +156,14 @@ bool contains(column_view const& haystack, scalar const& needle);
  *
  * @param haystack The column containing search space
  * @param needles A column of values to check for existence in the search space
+ * @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 A BOOL column indicating if each element in `needles` exists in the search space
  */
 std::unique_ptr<column> contains(
   column_view const& haystack,
   column_view const& needles,
+  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/search/contains_column.cu b/cpp/src/search/contains_column.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2019-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.
@@ -154,10 +154,11 @@ std::unique_ptr<column> contains(column_view const& haystack,
 
 std::unique_ptr<column> contains(column_view const& haystack,
                                  column_view const& needles,
+                                 rmm::cuda_stream_view stream,
                                  rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::contains(haystack, needles, cudf::get_default_stream(), mr);
+  return detail::contains(haystack, needles, stream, mr);
 }
 
 }  // namespace cudf
diff --git a/cpp/src/search/contains_scalar.cu b/cpp/src/search/contains_scalar.cu
@@ -160,10 +160,10 @@ bool contains(column_view const& haystack, scalar const& needle, rmm::cuda_strea
 
 }  // namespace detail
 
-bool contains(column_view const& haystack, scalar const& needle)
+bool contains(column_view const& haystack, scalar const& needle, rmm::cuda_stream_view stream)
 {
   CUDF_FUNC_RANGE();
-  return detail::contains(haystack, needle, cudf::get_default_stream());
+  return detail::contains(haystack, needle, stream);
 }
 
 }  // namespace cudf
diff --git a/cpp/src/search/search_ordered.cu b/cpp/src/search/search_ordered.cu
@@ -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.
@@ -144,22 +144,22 @@ std::unique_ptr<column> lower_bound(table_view const& haystack,
                                     table_view const& needles,
                                     std::vector<order> const& column_order,
                                     std::vector<null_order> const& null_precedence,
+                                    rmm::cuda_stream_view stream,
                                     rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::lower_bound(
-    haystack, needles, column_order, null_precedence, cudf::get_default_stream(), mr);
+  return detail::lower_bound(haystack, needles, column_order, null_precedence, stream, mr);
 }
 
 std::unique_ptr<column> upper_bound(table_view const& haystack,
                                     table_view const& needles,
                                     std::vector<order> const& column_order,
                                     std::vector<null_order> const& null_precedence,
+                                    rmm::cuda_stream_view stream,
                                     rmm::mr::device_memory_resource* mr)
 {
   CUDF_FUNC_RANGE();
-  return detail::upper_bound(
-    haystack, needles, column_order, null_precedence, cudf::get_default_stream(), mr);
+  return detail::upper_bound(haystack, needles, column_order, null_precedence, stream, mr);
 }
 
 }  // namespace cudf
diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
@@ -626,6 +626,7 @@ ConfigureTest(STREAM_GROUPBY_TEST streams/groupby_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_CONCATENATE_TEST streams/concatenate_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_FILLING_TEST streams/filling_test.cpp STREAM_MODE testing)
 ConfigureTest(STREAM_REPLACE_TEST streams/replace_test.cpp STREAM_MODE testing)
+ConfigureTest(STREAM_SEARCH_TEST streams/search_test.cpp STREAM_MODE testing)
 
 # ##################################################################################################
 # Install tests ####################################################################################

diff --git a/cpp/tests/streams/search_test.cpp b/cpp/tests/streams/search_test.cpp
@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) 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.
+ * 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 <cudf/column/column_view.hpp>
+#include <cudf/scalar/scalar.hpp>
+#include <cudf/search.hpp>
+
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/column_wrapper.hpp>
+#include <cudf_test/default_stream.hpp>
+
+class SearchTest : public cudf::test::BaseFixture {};
+
+TEST_F(SearchTest, LowerBound)
+{
+  cudf::test::fixed_width_column_wrapper<int32_t> column{10, 20, 30, 40, 50};
+  cudf::test::fixed_width_column_wrapper<int32_t> values{0, 7, 10, 11, 30, 32, 40, 47, 50, 90};
+  cudf::test::fixed_width_column_wrapper<cudf::size_type> expect{0, 0, 0, 1, 2, 3, 3, 4, 4, 5};
+
+  cudf::lower_bound({cudf::table_view{{column}}},
+                    {cudf::table_view{{values}}},
+                    {cudf::order::ASCENDING},
+                    {cudf::null_order::BEFORE},
+                    cudf::test::get_default_stream());
+}
+
+TEST_F(SearchTest, UpperBound)
+{
+  cudf::test::fixed_width_column_wrapper<int32_t> column{10, 20, 30, 40, 50};
+  cudf::test::fixed_width_column_wrapper<int32_t> values{0, 7, 10, 11, 30, 32, 40, 47, 50, 90};
+  cudf::test::fixed_width_column_wrapper<cudf::size_type> expect{0, 0, 0, 1, 2, 3, 3, 4, 4, 5};
+
+  cudf::upper_bound({cudf::table_view{{column}}},
+                    {cudf::table_view{{values}}},
+                    {cudf::order::ASCENDING},
+                    {cudf::null_order::BEFORE},
+                    cudf::test::get_default_stream());
+}
+
+TEST_F(SearchTest, ContainsScalar)
+{
+  cudf::test::fixed_width_column_wrapper<int32_t> column{0, 1, 17, 19, 23, 29, 71};
+  cudf::numeric_scalar<int32_t> scalar{23, true, cudf::test::get_default_stream()};
+
+  cudf::contains(column, scalar, cudf::test::get_default_stream());
+}
+
+TEST_F(SearchTest, ContainsColumn)
+{
+  cudf::test::fixed_width_column_wrapper<int32_t> haystack{0, 1, 17, 19, 23, 29, 71};
+  cudf::test::fixed_width_column_wrapper<int32_t> needles{17, 19, 45, 72};
+
+  cudf::test::fixed_width_column_wrapper<bool> expect{1, 1, 0, 0};
+
+  cudf::contains(haystack, needles, cudf::test::get_default_stream());
+}