Fix NaN handling in drop_list_duplicates

cpp/src/lists/drop_list_duplicates.cu

@@ -22,7 +22,6 @@
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/lists/detail/sorting.hpp>
 #include <cudf/lists/drop_list_duplicates.hpp>
-#include <cudf/table/row_operators.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
@@ -35,6 +34,90 @@ namespace lists {
 namespace detail {
 namespace {
 using offset_type = lists_column_view::offset_type;
+
+/**
+ * @brief Performs an equality comparison between two rows in a single column
+ *
+ * If the row elements are not of floating point types, this functor will return the
+ * same comparison result as `cudf::element_equality_comparator`.
+ * For floating point types, elements that hold NaN value will be considered as unequal.
+ *
+ * @tparam has_nulls Indicates the potential for null values in the column
+ */
+template <bool has_nulls = true>
+class element_comparator_fn {
+ public:
+  /**
+   * @brief Construct type-dispatched function object for comparing equality
+   * between two rows
+   *
+   * @param d_view The column containing the rows to compare
+   * @param nulls_are_equal Indicates if two null rows are treated as having the same value
+   */
+  __host__ __device__ element_comparator_fn(column_device_view d_view, bool nulls_are_equal)
+    : d_view{d_view}, nulls_are_equal{nulls_are_equal}
+  {
+  }
+
+  template <class T, std::enable_if_t<cudf::is_equality_comparable<T, T>()>* = nullptr>
+  __device__ bool operator()(size_type i, size_type j) const noexcept
+  {
+    if (has_nulls) {
+      bool const nullable = d_view.nullable();
+      bool const lhs_is_null{nullable and d_view.is_null(i)};
+      bool const rhs_is_null{nullable and d_view.is_null(j)};
+      if (lhs_is_null and rhs_is_null) {
+        return nulls_are_equal;
+      } else if (lhs_is_null != rhs_is_null) {
+        return false;
+      }
+    }
+
+    // For floating point types, if both element(i) and element(j) are NaNs then this comparison
+    // will return `false`. This is the desired behavior for `drop_list_duplicates`.
+    return d_view.element<T>(i) == d_view.element<T>(j);
+  }
+
+  template <typename Element,
+            std::enable_if_t<not cudf::is_equality_comparable<Element, Element>()>* = nullptr>
+  __device__ bool operator()(size_type, size_type) const
+  {
+    release_assert(false && "Attempted to compare elements of uncomparable types.");
+    return false;
+  }
+
+ private:
+  column_device_view d_view;
+  bool nulls_are_equal;
+};
+
+/**
+ * @brief Perform an equality comparison between two rows in a single column
+ *
+ * @tparam has_nulls Indicates the potential for null values in the column
+ */
+template <bool has_nulls>
+class element_comparator {
+ public:
+  element_comparator(column_device_view _d_view, null_equality _nulls_equal)
+    : d_view{_d_view}, nulls_equal{_nulls_equal}
+  {
+  }
+
+  __device__ bool operator()(size_type i, size_type j) const noexcept
+  {
+    return cudf::type_dispatcher(
+      d_view.type(),
+      element_comparator_fn<has_nulls>{d_view, nulls_equal == null_equality::EQUAL},
+      i,
+      j);
+  }
+
+ private:
+  column_device_view d_view;
+  null_equality nulls_equal;
+};
+
 /**
  * @brief Copy list entries and entry list offsets ignoring duplicates
  *
@@ -51,36 +134,46 @@ using offset_type = lists_column_view::offset_type;
  * @return A pair of columns, the first one contains unique list entries and the second one
  * contains their corresponding list offsets
  */
-template <bool has_nulls>
 std::vector<std::unique_ptr<column>> get_unique_entries_and_list_offsets(
   column_view const& all_lists_entries,
   column_view const& entries_list_offsets,
   null_equality nulls_equal,
   rmm::cuda_stream_view stream,
   rmm::mr::device_memory_resource* mr)
 {
-  // Create an intermediate table, since the comparator only work on tables
-  auto const device_input_table =
-    cudf::table_device_view::create(table_view{{all_lists_entries}}, stream);
-  auto const comp = row_equality_comparator<has_nulls>(
-    *device_input_table, *device_input_table, nulls_equal == null_equality::EQUAL);
-
   auto const num_entries = all_lists_entries.size();
+
   // Allocate memory to store the indices of the unique entries
   auto const unique_indices = cudf::make_numeric_column(
     entries_list_offsets.type(), num_entries, mask_state::UNALLOCATED, stream);
   auto const unique_indices_begin = unique_indices->mutable_view().begin<offset_type>();
 
-  auto const copy_end = thrust::unique_copy(
-    rmm::exec_policy(stream),
-    thrust::make_counting_iterator(0),
-    thrust::make_counting_iterator(num_entries),
-    unique_indices_begin,
-    [list_offsets = entries_list_offsets.begin<offset_type>(), comp] __device__(auto i, auto j) {
-      return list_offsets[i] == list_offsets[j] && comp(i, j);
-    });
+  offset_type* copy_end{nullptr};
+  auto const d_view_entries = column_device_view::create(all_lists_entries, stream);
+  auto const list_offsets   = entries_list_offsets.begin<offset_type>();
+  if (all_lists_entries.has_nulls()) {
+    element_comparator<true> const comp{*d_view_entries, nulls_equal};
+    copy_end = thrust::unique_copy(rmm::exec_policy(stream),
+                                   thrust::make_counting_iterator(0),
+                                   thrust::make_counting_iterator(num_entries),
+                                   unique_indices_begin,
+                                   [list_offsets, comp] __device__(auto i, auto j) {
+                                     return list_offsets[i] == list_offsets[j] && comp(i, j);
+                                   });
+  } else {
+    element_comparator<false> const comp{*d_view_entries, nulls_equal};
+    copy_end = thrust::unique_copy(rmm::exec_policy(stream),
+                                   thrust::make_counting_iterator(0),
+                                   thrust::make_counting_iterator(num_entries),
+                                   unique_indices_begin,
+                                   [list_offsets, comp] __device__(auto i, auto j) {
+                                     return list_offsets[i] == list_offsets[j] && comp(i, j);
+                                   });
+  }
 
   // Collect unique entries and entry list offsets
+  // The new null_count and bitmask of the unique entries will also be generated
+  // by the gather function
   auto const indices = cudf::detail::slice(
     unique_indices->view(), 0, thrust::distance(unique_indices_begin, copy_end));
   return cudf::detail::gather(table_view{{all_lists_entries, entries_list_offsets}},
@@ -255,12 +348,8 @@ std::unique_ptr<column> drop_list_duplicates(lists_column_view const& lists_colu
     detail::generate_entry_list_offsets(all_lists_entries.size(), lists_offsets->view(), stream);
 
   // Copy non-duplicated entries (along with their list offsets) to new arrays
-  auto unique_entries_and_list_offsets =
-    all_lists_entries.has_nulls()
-      ? detail::get_unique_entries_and_list_offsets<true>(
-          all_lists_entries, entries_list_offsets->view(), nulls_equal, stream, mr)
-      : detail::get_unique_entries_and_list_offsets<false>(
-          all_lists_entries, entries_list_offsets->view(), nulls_equal, stream, mr);
+  auto unique_entries_and_list_offsets = detail::get_unique_entries_and_list_offsets(
+    all_lists_entries, entries_list_offsets->view(), nulls_equal, stream, mr);
 
   // Generate offsets for the new lists column
   detail::generate_offsets(unique_entries_and_list_offsets.front()->size(),
@@ -269,6 +358,10 @@ std::unique_ptr<column> drop_list_duplicates(lists_column_view const& lists_colu
                            stream);
 
   // Construct a new lists column without duplicated entries
+  // Reuse the null_count and bitmask of the lists_column: those are the null information for
+  // the list elements (rows)
+  // For the entries of those lists (rows), their null_count and bitmask were generated separately
+  // during the step `get_unique_entries_and_list_offsets` above
   return make_lists_column(lists_column.size(),
                            std::move(lists_offsets),
                            std::move(unique_entries_and_list_offsets.front()),