Fully support nested types in cudf::contains

cpp/include/cudf/detail/search.hpp

@@ -105,12 +105,29 @@ rmm::device_uvector<bool> contains(
  * This function is designed for nested types only. It can also work with non-nested types
  * but with lower performance due to the complexity of the implementation.
  *
- * @param haystack The column containing search space.
- * @param needle A scalar value to check for existence in the search space.
- * @return true if the given `needle` value exists in the `haystack` column.
+ * @param haystack The column containing search space
+ * @param needle A scalar value to check for existence in the search space
+ * @return true if the given `needle` value exists in the `haystack` column
  */
 bool contains_nested_element(column_view const& haystack,
                              column_view const& needle,
                              rmm::cuda_stream_view stream);
 
+/**
+ * @brief Check if each row of the `needles` column exists in the `haystack` column.
+ *
+ * This function is designed for nested types only. It can also work with non-nested types
+ * but with lower performance due to the complexity of the implementation.
+ *
+ * @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> multi_contains_nested_elements(column_view const& haystack,
+                                                       column_view const& needles,
+                                                       rmm::cuda_stream_view stream,
+                                                       rmm::mr::device_memory_resource* mr);
+
 }  // namespace cudf::detail

cpp/src/search/contains_column.cu

@@ -17,6 +17,7 @@
 #include <hash/unordered_multiset.cuh>
 
 #include <cudf/column/column_factories.hpp>
+#include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/search.hpp>
 #include <cudf/dictionary/detail/search.hpp>
@@ -60,12 +61,16 @@ __device__ auto inline get_scalar_value(ScalarDView d_scalar)
 
 struct contains_scalar_dispatch {
   template <typename Type>
-  bool operator()(column_view const& haystack,
-                  scalar const& needle,
-                  rmm::cuda_stream_view stream) const
+  std::enable_if_t<!is_nested<Type>(), bool> operator()(column_view const& haystack,
+                                                        scalar const& needle,
+                                                        rmm::cuda_stream_view stream) const
   {
     CUDF_EXPECTS(haystack.type() == needle.type(), "scalar and column types must match");
 
+    // In case the input scalar is invalid, it should be handled at the caller site before
+    // dispatching to this function.
+    // (Handling such case is very simple: just check if the input haystack column has nulls).
+
     using DType           = device_storage_type_t<Type>;
     auto const d_haystack = column_device_view::create(haystack, stream);
     auto const d_needle =
@@ -90,28 +95,24 @@ struct contains_scalar_dispatch {
                }) > 0;
     }
   }
-};
 
-template <>
-bool contains_scalar_dispatch::operator()<cudf::list_view>(column_view const&,
-                                                           scalar const&,
-                                                           rmm::cuda_stream_view) const
-{
-  CUDF_FAIL("list_view type not supported yet");
-}
+  template <typename Type>
+  std::enable_if_t<is_nested<Type>(), bool> operator()(column_view const& haystack,
+                                                       scalar const& needle,
+                                                       rmm::cuda_stream_view stream) const
+  {
+    CUDF_EXPECTS(haystack.type() == needle.type(), "scalar and column types must match");
+    // Haystack and needle structure compatibility will be checked by the table comparator
+    // constructor during call to `contains_nested_element`.
 
-template <>
-bool contains_scalar_dispatch::operator()<cudf::struct_view>(column_view const& haystack,
-                                                             scalar const& needle,
-                                                             rmm::cuda_stream_view stream) const
-{
-  CUDF_EXPECTS(haystack.type() == needle.type(), "scalar and column types must match");
-  // Haystack and needle structure compatibility will be checked by the table comparator
-  // constructor during call to `contains_nested_element`.
+    // In case the input scalar is invalid, it should be handled at the caller site before
+    // dispatching to this function.
+    // (Handling such case is very simple: just check if the input haystack column has nulls).
 
-  auto const needle_as_col = make_column_from_scalar(needle, 1, stream);
-  return contains_nested_element(haystack, needle_as_col->view(), stream);
-}
+    auto const needle_as_col = make_column_from_scalar(needle, 1, stream);
+    return contains_nested_element(haystack, needle_as_col->view(), stream);
+  }
+};
 
 template <>
 bool contains_scalar_dispatch::operator()<cudf::dictionary32>(column_view const& haystack,
@@ -130,15 +131,32 @@ bool contains_scalar_dispatch::operator()<cudf::dictionary32>(column_view const&
 }
 
 struct multi_contains_dispatch {
-  template <typename Type>
+  template <typename ElementType, typename Haystack>
+  struct contains_fn {
+    bool __device__ operator()(size_type const idx) const
+    {
+      if (needles_have_nulls && needles.is_null_nocheck(idx)) {
+        // `true` or `false`: doesn't matter, this will be masked out as a null element.
+        return true;
+      }
+
+      return haystack.contains(needles.template element<ElementType>(idx));
+    }
+
+    Haystack const haystack;
+    column_device_view const needles;
+    bool const needles_have_nulls;
+  };
+
+  template <typename Type, CUDF_ENABLE_IF(!is_nested<Type>())>
   std::unique_ptr<column> operator()(column_view const& haystack,
                                      column_view const& needles,
                                      rmm::cuda_stream_view stream,
                                      rmm::mr::device_memory_resource* mr) const
   {
     auto result = make_numeric_column(data_type{type_to_id<bool>()},
                                       needles.size(),
-                                      copy_bitmask(needles),
+                                      copy_bitmask(needles, stream, mr),
                                       needles.null_count(),
                                       stream,
                                       mr);
@@ -152,53 +170,27 @@ struct multi_contains_dispatch {
     }
 
     auto const haystack_set    = cudf::detail::unordered_multiset<Type>::create(haystack, stream);
+    auto const haystack_set_dv = haystack_set.to_device();
     auto const needles_cdv_ptr = column_device_view::create(needles, stream);
-    auto const needles_it      = thrust::make_counting_iterator<size_type>(0);
-
-    if (needles.has_nulls()) {
-      thrust::transform(rmm::exec_policy(stream),
-                        needles_it,
-                        needles_it + needles.size(),
-                        out_begin,
-                        [haystack = haystack_set.to_device(),
-                         needles  = *needles_cdv_ptr] __device__(size_type const idx) {
-                          return needles.is_null_nocheck(idx) ||
-                                 haystack.contains(needles.template element<Type>(idx));
-                        });
-    } else {
-      thrust::transform(rmm::exec_policy(stream),
-                        needles_it,
-                        needles_it + needles.size(),
-                        out_begin,
-                        [haystack = haystack_set.to_device(),
-                         needles  = *needles_cdv_ptr] __device__(size_type const index) {
-                          return haystack.contains(needles.template element<Type>(index));
-                        });
-    }
 
+    thrust::transform(rmm::exec_policy(stream),
+                      thrust::make_counting_iterator<size_type>(0),
+                      thrust::make_counting_iterator<size_type>(needles.size()),
+                      out_begin,
+                      contains_fn<Type, decltype(haystack_set_dv)>{
+                        haystack_set_dv, *needles_cdv_ptr, needles.has_nulls()});
     return result;
   }
-};
 
-template <>
-std::unique_ptr<column> multi_contains_dispatch::operator()<list_view>(
-  column_view const&,
-  column_view const&,
-  rmm::cuda_stream_view,
-  rmm::mr::device_memory_resource*) const
-{
-  CUDF_FAIL("list_view type not supported");
-}
-
-template <>
-std::unique_ptr<column> multi_contains_dispatch::operator()<struct_view>(
-  column_view const&,
-  column_view const&,
-  rmm::cuda_stream_view,
-  rmm::mr::device_memory_resource*) const
-{
-  CUDF_FAIL("struct_view type not supported");
-}
+  template <typename Type, CUDF_ENABLE_IF(is_nested<Type>())>
+  std::unique_ptr<column> operator()(column_view const& haystack,
+                                     column_view const& needles,
+                                     rmm::cuda_stream_view stream,
+                                     rmm::mr::device_memory_resource* mr) const
+  {
+    return multi_contains_nested_elements(haystack, needles, stream, mr);
+  }
+};
 
 template <>
 std::unique_ptr<column> multi_contains_dispatch::operator()<dictionary32>(

cpp/src/search/contains_nested.cu

@@ -15,6 +15,8 @@
  */
 
 #include <cudf/detail/iterator.cuh>
+#include <cudf/detail/null_mask.hpp>
+#include <cudf/detail/search.hpp>
 #include <cudf/table/experimental/row_operators.cuh>
 #include <cudf/table/table_view.hpp>
 
@@ -63,4 +65,24 @@ bool contains_nested_element(column_view const& haystack,
            }) > 0;
 }
 
+std::unique_ptr<column> multi_contains_nested_elements(column_view const& haystack,
+                                                       column_view const& needles,
+                                                       rmm::cuda_stream_view stream,
+                                                       rmm::mr::device_memory_resource* mr)
+{
+  auto result_v = contains(table_view{{haystack}},
+                           table_view{{needles}},
+                           null_equality::EQUAL,
+                           nan_equality::ALL_EQUAL,
+                           stream,
+                           mr);
+
+  // todo: https://github.com/rapidsai/cudf/pull/11356
+  auto result =
+    std::make_unique<column>(data_type{type_to_id<bool>()}, needles.size(), result_v.release());
+  result->set_null_mask(copy_bitmask(needles, stream, mr), needles.null_count());
+
+  return result;
+}
+
 }  // namespace cudf::detail

cpp/tests/CMakeLists.txt

@@ -357,8 +357,8 @@ ConfigureTest(
 # ##################################################################################################
 # * search test -----------------------------------------------------------------------------------
 ConfigureTest(
-  SEARCH_TEST search/search_dictionary_test.cpp search/search_struct_test.cpp
-  search/search_test.cpp
+  SEARCH_TEST search/search_dictionary_test.cpp search/search_list_test.cpp
+  search/search_struct_test.cpp search/search_test.cpp
 )
 
 # ##################################################################################################