Fully support nested types in cudf::contains

cpp/include/cudf/detail/search.hpp

@@ -73,12 +73,29 @@ std::unique_ptr<column> contains(column_view const& haystack,
  * 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/include/cudf/table/experimental/row_operators.cuh

@@ -133,6 +133,25 @@ using lhs_iterator = strong_index_iterator<lhs_index_type>;
  */
 using rhs_iterator = strong_index_iterator<rhs_index_type>;
 
+/**
+ * @brief Convert an index `idx` iterating in reverse order in the range `[-1, -size-1)` into the
+ *        valid index range `[0, size)`.
+ *
+ * This conversion is needed in situations when we want to use two strong index types but are
+ * limited to use only one. One example of such limitation is using row indices of both tables as
+ * keys in a hash table. To overcome such limitation, we can use non-negative indices to
+ * denote one strong index type and negative indices (offset by one) to denote the other. After
+ * recognizing the corresponding type, we need to convert the negative indices into their original
+ * values.
+ *
+ * @param idx The negative index iterating in reverse order in the range `[-1, -size-1)`
+ * @return The converted index iterating in forward order in the range `[0, size)`
+ */
+[[nodiscard]] __device__ auto constexpr remap_negative_index(size_type const idx)
+{
+  return -(idx + 1);
+}
+
 namespace lexicographic {
 
 /**
@@ -413,7 +432,9 @@ struct preprocessed_table {
     : _t(std::move(table)),
       _column_order(std::move(column_order)),
       _null_precedence(std::move(null_precedence)),
-      _depths(std::move(depths)){};
+      _depths(std::move(depths))
+  {
+  }
 
   /**
    * @brief Implicit conversion operator to a `table_device_view` of the preprocessed table.
@@ -963,6 +984,50 @@ struct strong_index_comparator_adapter {
   Comparator const comparator;
 };
 
+/**
+ * @brief The adapter functor for calling a row comparator with strong index types from just one
+ *        type of integer index.
+ *
+ * In order to call the underlying row comparator with strong index types, the input integer
+ * indices to the functor must be given such that one index must be non-negative and the other must
+ * be negative:
+ *  - The non-negative index will be converted to `lhs_index_type`.
+ *  - The negative index will be converted into a non-negative value using the
+ *    @ref `remap_negative_index` function and then converted to `rhs_index_type`.
+ *
+ * @tparam Comparator A class of device row comparator with strong index types
+ */
+template <typename Comparator>
+struct negative_index_comparator_adapter {
+  negative_index_comparator_adapter(Comparator&& comparator) : comparator(std::move(comparator)) {}
+  /**
+   * @brief Call the underlying row comparator with strong index types from row indices of
+   *        `size_type` type.
+   *
+   * From two row indices `i` and `j`, the non-negative index is converted to `lhs_index_type`
+   * while the negative index is firstly converted into a non-negative value using the
+   * @ref `remap_negative_index` function then further converted to `rhs_index_type`. The underlying
+   * row comparator will be called using these strong type indices.
+   *
+   * Note that there is not any validity check to be performed in this function. Caller is
+   * responsible to make sure that the indices `i` and `j` having exactly one non-negative and one
+   * negative values. Otherwise, the output is undefined.
+   *
+   * @param i Index of a row in one table to compare
+   * @param j Index of a row in the other table to compare
+   * @return The row equality comparison result
+   */
+  __device__ auto operator()(size_type const i, size_type const j) const noexcept
+  {
+    auto const lhs_idx = static_cast<lhs_index_type>(i >= 0 ? i : j);
+    auto const rhs_idx = static_cast<rhs_index_type>(remap_negative_index(i < 0 ? i : j));
+    return comparator(lhs_idx, rhs_idx);
+  }
+
+ private:
+  Comparator const comparator;
+};
+
 /**
  * @brief An owning object that can be used to equality compare rows of two different tables.
  *
@@ -1182,6 +1247,39 @@ class device_row_hasher {
   uint32_t const _seed;
 };
 
+/**
+ * @brief The adapter functor for calling a row hasher from row index given in a range of negative
+ *        values.
+ *
+ * In order to call the underlying hasher, the input negative index will be converted back to a
+ * non-negative value using the @ref `remap_negative_index` utility function.
+ *
+ * @tparam Hasher A class of device row hasher
+ */
+template <typename Hasher>
+struct negative_index_hasher_adapter {
+  negative_index_hasher_adapter(Hasher&& hasher_) : hasher(std::move(hasher_)) {}
+
+  /**
+   * Given a negative row index `idx`, this functor converts the index into a (valid) non-negative
+   * value using the @ref `remap_negative_index` function before calling to the underlying row
+   * hasher.
+   *
+   * Note that there is not any validity check for the input value `idx`, assuming that it is always
+   * negative. Otherwise, the output is undefined.
+   *
+   * @param idx A row index that is given in a range of negative values
+   * @return The resulting hash value
+   */
+  __device__ auto operator()(size_type const idx) const noexcept
+  {
+    return hasher(remap_negative_index(idx));
+  }
+
+ private:
+  Hasher const hasher;
+};
+
 // Inject row::equality::preprocessed_table into the row::hash namespace
 // As a result, row::equality::preprocessed_table and row::hash::preprocessed table are the same
 // type and are interchangeable.

cpp/src/search/contains.cu

@@ -40,13 +40,17 @@ namespace detail {
 namespace {
 
 struct contains_scalar_dispatch {
-  template <typename Type>
+  template <typename Type, CUDF_ENABLE_IF(!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>;
     using ScalarType      = cudf::scalar_type_t<Type>;
     auto const d_haystack = column_device_view::create(haystack, stream);
@@ -66,28 +70,24 @@ struct contains_scalar_dispatch {
       return thrust::find(rmm::exec_policy(stream), begin, end, val) != end;
     }
   }
-};
 
-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, CUDF_ENABLE_IF(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,
@@ -106,7 +106,24 @@ 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,
@@ -128,53 +145,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>(