Revert "Add primitive row dispatch support for semi/anti join and cudf::contains"

cpp/CMakeLists.txt

@@ -681,9 +681,6 @@ add_library(
   src/search/contains_column.cu
   src/search/contains_scalar.cu
   src/search/contains_table.cu
-  src/search/contains_table_impl.cu
-  src/search/contains_table_impl_primitive.cu
-  src/search/contains_table_impl_nested.cu
   src/search/search_ordered.cu
   src/sort/is_sorted.cu
   src/sort/rank.cu

cpp/include/cudf/table/primitive_row_operators.cuh

@@ -143,19 +143,6 @@ class row_equality_comparator {
                                                           rhs_row_index);
   }
 
-  /**
-   * @brief Compares the specified rows for equality.
-   *
-   * @param lhs_index The index of the first row to compare (in the lhs table)
-   * @param rhs_index The index of the second row to compare (in the rhs table)
-   * @return Boolean indicating if both rows are equal
-   */
-  __device__ bool operator()(cudf::experimental::row::lhs_index_type lhs_index,
-                             cudf::experimental::row::rhs_index_type rhs_index) const
-  {
-    return (*this)(static_cast<size_type>(lhs_index), static_cast<size_type>(rhs_index));
-  }
-
  private:
   cudf::nullate::DYNAMIC _has_nulls;
   table_device_view _lhs;

cpp/src/search/contains_table.cu

@@ -14,24 +14,175 @@
  * limitations under the License.
  */
 
-#include "contains_table_impl.cuh"
+#include "join/join_common_utils.cuh"
 
+#include <cudf/detail/cuco_helpers.hpp>
+#include <cudf/detail/null_mask.hpp>
 #include <cudf/detail/search.hpp>
+#include <cudf/hashing/detail/helper_functions.cuh>
 #include <cudf/table/experimental/row_operators.cuh>
-#include <cudf/table/primitive_row_operators.cuh>
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
+#include <cudf/utilities/memory_resource.hpp>
 #include <cudf/utilities/type_checks.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <cuco/static_set.cuh>
+#include <cuda/functional>
+#include <thrust/iterator/counting_iterator.h>
+
+#include <type_traits>
 
 namespace cudf::detail {
 
 namespace {
 
+using cudf::experimental::row::lhs_index_type;
+using cudf::experimental::row::rhs_index_type;
+
+/**
+ * @brief An hasher adapter wrapping both haystack hasher and needles hasher
+ */
+template <typename HaystackHasher, typename NeedleHasher>
+struct hasher_adapter {
+  hasher_adapter(HaystackHasher const& haystack_hasher, NeedleHasher const& needle_hasher)
+    : _haystack_hasher{haystack_hasher}, _needle_hasher{needle_hasher}
+  {
+  }
+
+  __device__ constexpr auto operator()(lhs_index_type idx) const noexcept
+  {
+    return _needle_hasher(static_cast<size_type>(idx));
+  }
+
+  __device__ constexpr auto operator()(rhs_index_type idx) const noexcept
+  {
+    return _haystack_hasher(static_cast<size_type>(idx));
+  }
+
+ private:
+  HaystackHasher const _haystack_hasher;
+  NeedleHasher const _needle_hasher;
+};
+
+/**
+ * @brief An comparator adapter wrapping both self comparator and two table comparator
+ */
+template <typename SelfEqual, typename TwoTableEqual>
+struct comparator_adapter {
+  comparator_adapter(SelfEqual const& self_equal, TwoTableEqual const& two_table_equal)
+    : _self_equal{self_equal}, _two_table_equal{two_table_equal}
+  {
+  }
+
+  __device__ constexpr auto operator()(rhs_index_type lhs_index,
+                                       rhs_index_type rhs_index) const noexcept
+  {
+    auto const lhs = static_cast<size_type>(lhs_index);
+    auto const rhs = static_cast<size_type>(rhs_index);
+
+    return _self_equal(lhs, rhs);
+  }
+
+  __device__ constexpr auto operator()(lhs_index_type lhs_index,
+                                       rhs_index_type rhs_index) const noexcept
+  {
+    return _two_table_equal(lhs_index, rhs_index);
+  }
+
+ private:
+  SelfEqual const _self_equal;
+  TwoTableEqual const _two_table_equal;
+};
+
+/**
+ * @brief Build a row bitmask for the input table.
+ *
+ * The output bitmask will have invalid bits corresponding to the input rows having nulls (at
+ * any nested level) and vice versa.
+ *
+ * @param input The input table
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @return A pair of pointer to the output bitmask and the buffer containing the bitmask
+ */
+std::pair<rmm::device_buffer, bitmask_type const*> build_row_bitmask(table_view const& input,
+                                                                     rmm::cuda_stream_view stream)
+{
+  auto const nullable_columns = get_nullable_columns(input);
+  CUDF_EXPECTS(nullable_columns.size() > 0,
+               "The input table has nulls thus it should have nullable columns.");
+
+  // If there are more than one nullable column, we compute `bitmask_and` of their null masks.
+  // Otherwise, we have only one nullable column and can use its null mask directly.
+  if (nullable_columns.size() > 1) {
+    auto row_bitmask =
+      cudf::detail::bitmask_and(
+        table_view{nullable_columns}, stream, cudf::get_current_device_resource_ref())
+        .first;
+    auto const row_bitmask_ptr = static_cast<bitmask_type const*>(row_bitmask.data());
+    return std::pair(std::move(row_bitmask), row_bitmask_ptr);
+  }
+
+  return std::pair(rmm::device_buffer{0, stream}, nullable_columns.front().null_mask());
+}
+
+/**
+ * @brief Invokes the given `func` with desired comparators based on the specified `compare_nans`
+ * parameter
+ *
+ * @tparam HasNested Flag indicating whether there are nested columns in haystack or needles
+ * @tparam Hasher Type of device hash function
+ * @tparam Func Type of the helper function doing `contains` check
+ *
+ * @param compare_nulls Control whether nulls should be compared as equal or not
+ * @param compare_nans Control whether floating-point NaNs values should be compared as equal or not
+ * @param haystack_has_nulls Flag indicating whether haystack has nulls or not
+ * @param has_any_nulls Flag indicating whether there are nested nulls is either haystack or needles
+ * @param self_equal Self table comparator
+ * @param two_table_equal Two table comparator
+ * @param d_hasher Device hash functor
+ * @param func The input functor to invoke
+ */
+template <bool HasNested, typename Hasher, typename Func>
+void dispatch_nan_comparator(
+  null_equality compare_nulls,
+  nan_equality compare_nans,
+  bool haystack_has_nulls,
+  bool has_any_nulls,
+  cudf::experimental::row::equality::self_comparator self_equal,
+  cudf::experimental::row::equality::two_table_comparator two_table_equal,
+  Hasher const& d_hasher,
+  Func&& func)
+{
+  // Distinguish probing scheme CG sizes between nested and flat types for better performance
+  auto const probing_scheme = [&]() {
+    if constexpr (HasNested) {
+      return cuco::linear_probing<4, Hasher>{d_hasher};
+    } else {
+      return cuco::linear_probing<1, Hasher>{d_hasher};
+    }
+  }();
+
+  if (compare_nans == nan_equality::ALL_EQUAL) {
+    using nan_equal_comparator =
+      cudf::experimental::row::equality::nan_equal_physical_equality_comparator;
+    auto const d_self_equal = self_equal.equal_to<HasNested>(
+      nullate::DYNAMIC{haystack_has_nulls}, compare_nulls, nan_equal_comparator{});
+    auto const d_two_table_equal = two_table_equal.equal_to<HasNested>(
+      nullate::DYNAMIC{has_any_nulls}, compare_nulls, nan_equal_comparator{});
+    func(d_self_equal, d_two_table_equal, probing_scheme);
+  } else {
+    using nan_unequal_comparator = cudf::experimental::row::equality::physical_equality_comparator;
+    auto const d_self_equal      = self_equal.equal_to<HasNested>(
+      nullate::DYNAMIC{haystack_has_nulls}, compare_nulls, nan_unequal_comparator{});
+    auto const d_two_table_equal = two_table_equal.equal_to<HasNested>(
+      nullate::DYNAMIC{has_any_nulls}, compare_nulls, nan_unequal_comparator{});
+    func(d_self_equal, d_two_table_equal, probing_scheme);
+  }
+}
+
 }  // namespace
 
 rmm::device_uvector<bool> contains(table_view const& haystack,
@@ -52,69 +203,94 @@ rmm::device_uvector<bool> contains(table_view const& haystack,
   auto const preprocessed_haystack =
     cudf::experimental::row::equality::preprocessed_table::create(haystack, stream);
 
+  auto const haystack_hasher   = cudf::experimental::row::hash::row_hasher(preprocessed_haystack);
+  auto const d_haystack_hasher = haystack_hasher.device_hasher(nullate::DYNAMIC{has_any_nulls});
+  auto const needle_hasher     = cudf::experimental::row::hash::row_hasher(preprocessed_needles);
+  auto const d_needle_hasher   = needle_hasher.device_hasher(nullate::DYNAMIC{has_any_nulls});
+  auto const d_hasher          = hasher_adapter{d_haystack_hasher, d_needle_hasher};
+
+  auto const self_equal = cudf::experimental::row::equality::self_comparator(preprocessed_haystack);
+  auto const two_table_equal = cudf::experimental::row::equality::two_table_comparator(
+    preprocessed_needles, preprocessed_haystack);
+
   // The output vector.
   auto contained = rmm::device_uvector<bool>(needles.num_rows(), stream, mr);
 
-  if (cudf::is_primitive_row_op_compatible(haystack)) {
-    auto const d_haystack_hasher =
-      cudf::row::primitive::row_hasher{nullate::DYNAMIC{has_any_nulls}, preprocessed_haystack};
-    auto const d_needle_hasher =
-      cudf::row::primitive::row_hasher{nullate::DYNAMIC{has_any_nulls}, preprocessed_needles};
-    auto const d_hasher     = hasher_adapter{d_haystack_hasher, d_needle_hasher};
-    auto const d_self_equal = cudf::row::primitive::row_equality_comparator{
-      nullate::DYNAMIC{has_any_nulls}, preprocessed_haystack, preprocessed_haystack, compare_nulls};
-    auto const d_two_table_equal = cudf::row::primitive::row_equality_comparator{
-      nullate::DYNAMIC{has_any_nulls}, preprocessed_needles, preprocessed_haystack, compare_nulls};
-    auto const d_equal = comparator_adapter{d_self_equal, d_two_table_equal};
-    perform_contains(haystack,
-                     needles,
-                     haystack_has_nulls,
-                     needles_has_nulls,
-                     compare_nulls,
-                     d_equal,
-                     cuco::linear_probing<1, decltype(d_hasher)>{d_hasher},
-                     contained,
-                     stream);
+  auto const haystack_iter = cudf::detail::make_counting_transform_iterator(
+    size_type{0}, cuda::proclaim_return_type<rhs_index_type>([] __device__(auto idx) {
+      return rhs_index_type{idx};
+    }));
+  auto const needles_iter = cudf::detail::make_counting_transform_iterator(
+    size_type{0}, cuda::proclaim_return_type<lhs_index_type>([] __device__(auto idx) {
+      return lhs_index_type{idx};
+    }));
+
+  auto const helper_func =
+    [&](auto const& d_self_equal, auto const& d_two_table_equal, auto const& probing_scheme) {
+      auto const d_equal = comparator_adapter{d_self_equal, d_two_table_equal};
+
+      auto set = cuco::static_set{
+        cuco::extent{compute_hash_table_size(haystack.num_rows())},
+        cuco::empty_key{rhs_index_type{-1}},
+        d_equal,
+        probing_scheme,
+        {},
+        {},
+        cudf::detail::cuco_allocator<char>{rmm::mr::polymorphic_allocator<char>{}, stream},
+        stream.value()};
+
+      if (haystack_has_nulls && compare_nulls == null_equality::UNEQUAL) {
+        auto const bitmask_buffer_and_ptr = build_row_bitmask(haystack, stream);
+        auto const row_bitmask_ptr        = bitmask_buffer_and_ptr.second;
+
+        // If the haystack table has nulls but they are compared unequal, don't insert them.
+        // Otherwise, it was known to cause performance issue:
+        // - https://github.com/rapidsai/cudf/pull/6943
+        // - https://github.com/rapidsai/cudf/pull/8277
+        set.insert_if_async(haystack_iter,
+                            haystack_iter + haystack.num_rows(),
+                            thrust::counting_iterator<size_type>(0),  // stencil
+                            row_is_valid{row_bitmask_ptr},
+                            stream.value());
+      } else {
+        set.insert_async(haystack_iter, haystack_iter + haystack.num_rows(), stream.value());
+      }
+
+      if (needles_has_nulls && compare_nulls == null_equality::UNEQUAL) {
+        auto const bitmask_buffer_and_ptr = build_row_bitmask(needles, stream);
+        auto const row_bitmask_ptr        = bitmask_buffer_and_ptr.second;
+        set.contains_if_async(needles_iter,
+                              needles_iter + needles.num_rows(),
+                              thrust::counting_iterator<size_type>(0),  // stencil
+                              row_is_valid{row_bitmask_ptr},
+                              contained.begin(),
+                              stream.value());
+      } else {
+        set.contains_async(
+          needles_iter, needles_iter + needles.num_rows(), contained.begin(), stream.value());
+      }
+    };
+
+  if (cudf::detail::has_nested_columns(haystack)) {
+    dispatch_nan_comparator<true>(compare_nulls,
+                                  compare_nans,
+                                  haystack_has_nulls,
+                                  has_any_nulls,
+                                  self_equal,
+                                  two_table_equal,
+                                  d_hasher,
+                                  helper_func);
   } else {
-    auto const haystack_hasher   = cudf::experimental::row::hash::row_hasher(preprocessed_haystack);
-    auto const d_haystack_hasher = haystack_hasher.device_hasher(nullate::DYNAMIC{has_any_nulls});
-    auto const needle_hasher     = cudf::experimental::row::hash::row_hasher(preprocessed_needles);
-    auto const d_needle_hasher   = needle_hasher.device_hasher(nullate::DYNAMIC{has_any_nulls});
-    auto const d_hasher          = hasher_adapter{d_haystack_hasher, d_needle_hasher};
-
-    auto const self_equal =
-      cudf::experimental::row::equality::self_comparator(preprocessed_haystack);
-    auto const two_table_equal = cudf::experimental::row::equality::two_table_comparator(
-      preprocessed_needles, preprocessed_haystack);
-
-    if (cudf::detail::has_nested_columns(haystack)) {
-      dispatch_nan_comparator<true>(haystack,
-                                    needles,
-                                    compare_nulls,
-                                    compare_nans,
-                                    haystack_has_nulls,
-                                    needles_has_nulls,
-                                    has_any_nulls,
-                                    self_equal,
-                                    two_table_equal,
-                                    d_hasher,
-                                    contained,
-                                    stream);
-    } else {
-      dispatch_nan_comparator<false>(haystack,
-                                     needles,
-                                     compare_nulls,
-                                     compare_nans,
-                                     haystack_has_nulls,
-                                     needles_has_nulls,
-                                     has_any_nulls,
-                                     self_equal,
-                                     two_table_equal,
-                                     d_hasher,
-                                     contained,
-                                     stream);
-    }
+    dispatch_nan_comparator<false>(compare_nulls,
+                                   compare_nans,
+                                   haystack_has_nulls,
+                                   has_any_nulls,
+                                   self_equal,
+                                   two_table_equal,
+                                   d_hasher,
+                                   helper_func);
   }
+
   return contained;
 }