Nested struct binop comparison

cpp/include/cudf/table/experimental/row_operators.cuh

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <cudf/binaryop.hpp>
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/detail/hashing.hpp>
 #include <cudf/detail/iterator.cuh>
@@ -45,6 +46,7 @@
 #include <utility>
 
 namespace cudf {
+
 namespace experimental {
 
 /**
@@ -68,28 +70,31 @@ struct dispatch_void_if_nested {
 };
 
 namespace row {
-
 namespace lexicographic {
 
 /**
- * @brief Computes whether one row is lexicographically *less* than another row.
+ * @brief Computes the lexicographic comparison between 2 rows.
  *
  * Lexicographic ordering is determined by:
  * - Two rows are compared element by element.
  * - The first mismatching element defines which row is lexicographically less
  * or greater than the other.
+ * - If the rows are compared without mismatched elements, the rows are equivalent
+ *
  *
  * Lexicographic ordering is exactly equivalent to doing an alphabetical sort of
  * two words, for example, `aac` would be *less* than (or precede) `abb`. The
  * second letter in both words is the first non-equal letter, and `a < b`, thus
  * `aac < abb`.
  *
  * @tparam Nullate A cudf::nullate type describing whether to check for nulls.
+ * @tparam NanConfig default configuration nans are equal, if set to true triggers specialized IEEE
+ * 754 compliant nan handling
  */
-template <typename Nullate>
+template <typename Nullate, bool NanConfig = false>
 class device_row_comparator {
-  friend class self_comparator;
-
+  // friend class self_comparator;
+ public:  // needs to be removed, pending strict typing for indices
   /**
    * @brief Construct a function object for performing a lexicographic
    * comparison between the rows of two tables.
@@ -139,13 +144,20 @@ class device_row_comparator {
      * @param null_precedence Indicates how null values are ordered with other values
      * @param depth The depth of the column if part of a nested column @see
      * preprocessed_table::depths
+     * @param nan_result Specifies what value should be returned if either element is `nan`
      */
     __device__ element_comparator(Nullate check_nulls,
                                   column_device_view lhs,
                                   column_device_view rhs,
                                   null_order null_precedence = null_order::BEFORE,
-                                  int depth                  = 0)
-      : _lhs{lhs}, _rhs{rhs}, _nulls{check_nulls}, _null_precedence{null_precedence}, _depth{depth}
+                                  int depth                  = 0,
+                                  weak_ordering nan_result   = weak_ordering::EQUIVALENT)
+      : _lhs{lhs},
+        _rhs{rhs},
+        _check_nulls{check_nulls},
+        _null_precedence{null_precedence},
+        _depth{depth},
+        _nan_result{nan_result}
     {
     }
 
@@ -162,7 +174,7 @@ class device_row_comparator {
     __device__ cuda::std::pair<weak_ordering, int> operator()(
       size_type const lhs_element_index, size_type const rhs_element_index) const noexcept
     {
-      if (_nulls) {
+      if (_check_nulls) {
         bool const lhs_is_null{_lhs.is_null(lhs_element_index)};
         bool const rhs_is_null{_rhs.is_null(rhs_element_index)};
 
@@ -171,8 +183,12 @@ class device_row_comparator {
         }
       }
 
-      return cuda::std::pair(relational_compare(_lhs.element<Element>(lhs_element_index),
-                                                _rhs.element<Element>(rhs_element_index)),
+      return cuda::std::pair(NanConfig
+                               ? relational_compare(_lhs.element<Element>(lhs_element_index),
+                                                    _rhs.element<Element>(rhs_element_index),
+                                                    _nan_result)
+                               : relational_compare(_lhs.element<Element>(lhs_element_index),
+                                                    _rhs.element<Element>(rhs_element_index)),
                              std::numeric_limits<int>::max());
     }
 
@@ -211,29 +227,32 @@ class device_row_comparator {
         ++depth;
       }
 
-      auto const comparator = element_comparator{_nulls, lcol, rcol, _null_precedence, depth};
+      auto const comparator =
+        element_comparator{_check_nulls, lcol, rcol, _null_precedence, depth, _nan_result};
       return cudf::type_dispatcher<dispatch_void_if_nested>(
         lcol.type(), comparator, lhs_element_index, rhs_element_index);
     }
 
    private:
     column_device_view const _lhs;
     column_device_view const _rhs;
-    Nullate const _nulls;
+    Nullate const _check_nulls;
     null_order const _null_precedence;
     int const _depth;
+    weak_ordering _nan_result;
   };
 
  public:
   /**
    * @brief Checks whether the row at `lhs_index` in the `lhs` table compares
-   * lexicographically less than the row at `rhs_index` in the `rhs` table.
+   * lexicographically less, greater, or equivalent to the row at `rhs_index` in the `rhs` table.
    *
    * @param lhs_index The index of row in the `lhs` table to examine
    * @param rhs_index The index of the row in the `rhs` table to examine
-   * @return `true` if row from the `lhs` table compares less than row in the `rhs` table
+   * @return weak ordering comparison of the row in the `lhs` table relative to the row in the `rhs`
+   * table
    */
-  __device__ bool operator()(size_type const lhs_index, size_type const rhs_index) const noexcept
+  __device__ weak_ordering operator()(size_type lhs_index, size_type rhs_index) const noexcept
   {
     int last_null_depth = std::numeric_limits<int>::max();
     for (size_type i = 0; i < _lhs.num_columns(); ++i) {
@@ -247,17 +266,31 @@ class device_row_comparator {
         _null_precedence.has_value() ? (*_null_precedence)[i] : null_order::BEFORE;
 
       auto const comparator =
-        element_comparator{_check_nulls, _lhs.column(i), _rhs.column(i), null_precedence, depth};
+        NanConfig ? element_comparator{_check_nulls,
+                                       _lhs.column(i),
+                                       _rhs.column(i),
+                                       null_precedence,
+                                       depth,
+                                       ascending ? weak_ordering::GREATER : weak_ordering::LESS}
+
+                  : element_comparator{_check_nulls,
+                                       _lhs.column(i),
+                                       _rhs.column(i),
+                                       null_precedence,
+                                       depth,
+                                       weak_ordering::EQUIVALENT};
 
       weak_ordering state;
       cuda::std::tie(state, last_null_depth) =
         cudf::type_dispatcher(_lhs.column(i).type(), comparator, lhs_index, rhs_index);
 
       if (state == weak_ordering::EQUIVALENT) { continue; }
 
-      return state == (ascending ? weak_ordering::LESS : weak_ordering::GREATER);
+      return ascending
+               ? state
+               : (state == weak_ordering::GREATER ? weak_ordering::LESS : weak_ordering::GREATER);
     }
-    return false;
+    return weak_ordering::EQUIVALENT;
   }
 
  private:
@@ -269,6 +302,42 @@ class device_row_comparator {
   std::optional<device_span<null_order const>> const _null_precedence;
 };  // class device_row_comparator
 
+/**
+ * @brief Wraps and interprets the result of templated Comparator that returns a weak_ordering.
+ * Returns true if the weak_ordering matches any of the templated values.
+ *
+ * Note that this should never be used with only `weak_ordering::EQUIVALENT`.
+ * An equality comparator should be used instead for optimal performance.
+ *
+ * @tparam Comparator generic comparator that returns a weak_ordering.
+ * @tparam values weak_ordering parameter pack of orderings to interpret as true
+ */
+template <typename Comparator, weak_ordering... values>
+struct weak_ordering_comparator_impl {
+  __device__ bool operator()(size_type const& lhs, size_type const& rhs)
+  {
+    weak_ordering const result = comparator(lhs, rhs);
+    return ((result == values) || ...);
+  }
+  Comparator comparator;
+};
+
+/**
+ * @brief Wraps and interprets the result of device_row_comparator, true if the result is
+ * weak_ordering::LESS meaning one row is lexicographically *less* than another row.
+ *
+ * @tparam Nullate A cudf::nullate type describing whether to check for nulls.
+ */
+template <typename Nullate, bool NanConfig = false>
+using less_comparator =
+  weak_ordering_comparator_impl<device_row_comparator<Nullate, NanConfig>, weak_ordering::LESS>;
+
+template <typename Nullate, bool NanConfig = false>
+using less_equivalent_comparator =
+  weak_ordering_comparator_impl<device_row_comparator<Nullate, NanConfig>,
+                                weak_ordering::LESS,
+                                weak_ordering::EQUIVALENT>;
+
 struct preprocessed_table {
   using table_device_view_owner =
     std::invoke_result_t<decltype(table_device_view::create), table_view, rmm::cuda_stream_view>;
@@ -416,11 +485,11 @@ class self_comparator {
    *
    * @tparam Nullate A cudf::nullate type describing whether to check for nulls.
    */
-  template <typename Nullate>
-  device_row_comparator<Nullate> device_comparator(Nullate nullate = {}) const
+  template <typename Nullate, bool NanConfig = false>
+  less_comparator<Nullate, NanConfig> device_comparator(Nullate nullate = {}) const
   {
-    return device_row_comparator(
-      nullate, *d_t, *d_t, d_t->depths(), d_t->column_order(), d_t->null_precedence());
+    return less_comparator<Nullate, NanConfig>{device_row_comparator<Nullate, NanConfig>(
+      nullate, *d_t, *d_t, d_t->depths(), d_t->column_order(), d_t->null_precedence())};
   }
 
  private:

cpp/include/cudf/table/row_operators.cuh

@@ -97,6 +97,26 @@ __device__ weak_ordering relational_compare(Element lhs, Element rhs)
   return detail::compare_elements(lhs, rhs);
 }
 
+/**
+ * @brief A specialization for floating-point `Element` type relational comparison
+ * to derive the order of the elements with respect to `lhs`. Returns specified weak_ordering if
+ * either value is `nan`, enabling IEEE 754 compliant comparison.
+ *
+ * This specialization allows `nan` values to be evaluated as not equal to any other value, while
+ * also not evaluating as greater or less than
+ *
+ * @param lhs first element
+ * @param rhs second element
+ * @param nan_result specifies what value should be returned if either element is `nan`
+ * @return Indicates the relationship between the elements in
+ * the `lhs` and `rhs` columns.
+ */
+template <typename Element, std::enable_if_t<std::is_floating_point<Element>::value>* = nullptr>
+__device__ weak_ordering relational_compare(Element lhs, Element rhs, weak_ordering nan_result)
+{
+  return isnan(lhs) or isnan(rhs) ? nan_result : detail::compare_elements(lhs, rhs);
+}
+
 /**
  * @brief Compare the nulls according to null order.
  *
@@ -123,11 +143,14 @@ inline __device__ auto null_compare(bool lhs_is_null, bool rhs_is_null, null_ord
  *
  * @param[in] lhs first element
  * @param[in] rhs second element
+ * @param nan_result ignored for non-floating point operation
  * @return Indicates the relationship between the elements in
  * the `lhs` and `rhs` columns.
  */
 template <typename Element, std::enable_if_t<not std::is_floating_point_v<Element>>* = nullptr>
-__device__ weak_ordering relational_compare(Element lhs, Element rhs)
+__device__ weak_ordering relational_compare(Element lhs,
+                                            Element rhs,
+                                            weak_ordering const nan_result = weak_ordering::GREATER)
 {
   return detail::compare_elements(lhs, rhs);
 }
@@ -138,12 +161,15 @@ __device__ weak_ordering relational_compare(Element lhs, Element rhs)
  *
  * @param lhs first element
  * @param rhs second element
+ * @param nan_result specifies what value should be returned if either element is `nan`
  * @return `true` if `lhs` == `rhs` else `false`.
  */
 template <typename Element, std::enable_if_t<std::is_floating_point_v<Element>>* = nullptr>
-__device__ bool equality_compare(Element lhs, Element rhs)
+__device__ bool equality_compare(Element lhs,
+                                 Element rhs,
+                                 nan_equality const nan_result = nan_equality::ALL_EQUAL)
 {
-  if (isnan(lhs) and isnan(rhs)) { return true; }
+  if (isnan(lhs) and isnan(rhs)) { return nan_result == nan_equality::ALL_EQUAL; }
   return lhs == rhs;
 }
 
@@ -153,10 +179,13 @@ __device__ bool equality_compare(Element lhs, Element rhs)
  *
  * @param lhs first element
  * @param rhs second element
+ * @param nan_result ignored for non-floating point operation
  * @return `true` if `lhs` == `rhs` else `false`.
  */
 template <typename Element, std::enable_if_t<not std::is_floating_point_v<Element>>* = nullptr>
-__device__ bool equality_compare(Element const lhs, Element const rhs)
+__device__ bool equality_compare(Element const lhs,
+                                 Element const rhs,
+                                 nan_equality const nan_result = nan_equality::ALL_EQUAL)
 {
   return lhs == rhs;
 }
@@ -179,13 +208,19 @@ class element_equality_comparator {
    * @param lhs The column containing the first element
    * @param rhs The column containing the second element (may be the same as lhs)
    * @param nulls_are_equal Indicates if two null elements are treated as equivalent
+   * @param nan_result specifies what value should be returned if either element is `nan`
    */
   __host__ __device__
   element_equality_comparator(Nullate has_nulls,
                               column_device_view lhs,
                               column_device_view rhs,
-                              null_equality nulls_are_equal = null_equality::EQUAL)
-    : lhs{lhs}, rhs{rhs}, nulls{has_nulls}, nulls_are_equal{nulls_are_equal}
+                              null_equality nulls_are_equal = null_equality::EQUAL,
+                              nan_equality nans_are_equal   = nan_equality::ALL_EQUAL)
+    : lhs{lhs},
+      rhs{rhs},
+      nulls{has_nulls},
+      nulls_are_equal{nulls_are_equal},
+      nans_are_equal{nans_are_equal}
   {
   }
 
@@ -212,7 +247,8 @@ class element_equality_comparator {
     }
 
     return equality_compare(lhs.element<Element>(lhs_element_index),
-                            rhs.element<Element>(rhs_element_index));
+                            rhs.element<Element>(rhs_element_index),
+                            nans_are_equal);
   }
 
   template <typename Element,
@@ -227,6 +263,7 @@ class element_equality_comparator {
   column_device_view rhs;
   Nullate nulls;
   null_equality nulls_are_equal;
+  nan_equality nans_are_equal;
 };
 
 template <typename Nullate>
@@ -235,19 +272,25 @@ class row_equality_comparator {
   row_equality_comparator(Nullate has_nulls,
                           table_device_view lhs,
                           table_device_view rhs,
-                          null_equality nulls_are_equal = null_equality::EQUAL)
-    : lhs{lhs}, rhs{rhs}, nulls{has_nulls}, nulls_are_equal{nulls_are_equal}
+                          null_equality nulls_are_equal = null_equality::EQUAL,
+                          nan_equality nans_are_equal   = nan_equality::ALL_EQUAL)
+    : lhs{lhs},
+      rhs{rhs},
+      nulls{has_nulls},
+      nulls_are_equal{nulls_are_equal},
+      nans_are_equal{nans_are_equal}
   {
     CUDF_EXPECTS(lhs.num_columns() == rhs.num_columns(), "Mismatched number of columns.");
   }
 
   __device__ bool operator()(size_type lhs_row_index, size_type rhs_row_index) const noexcept
   {
     auto equal_elements = [=](column_device_view l, column_device_view r) {
-      return cudf::type_dispatcher(l.type(),
-                                   element_equality_comparator{nulls, l, r, nulls_are_equal},
-                                   lhs_row_index,
-                                   rhs_row_index);
+      return cudf::type_dispatcher(
+        l.type(),
+        element_equality_comparator{nulls, l, r, nulls_are_equal, nans_are_equal},
+        lhs_row_index,
+        rhs_row_index);
     };
 
     return thrust::equal(thrust::seq, lhs.begin(), lhs.end(), rhs.begin(), equal_elements);
@@ -258,6 +301,7 @@ class row_equality_comparator {
   table_device_view rhs;
   Nullate nulls;
   null_equality nulls_are_equal;
+  nan_equality nans_are_equal;
 };
 
 /**