Nested struct binop comparison

cpp/include/cudf/detail/structs/utilities.hpp

@@ -245,6 +245,20 @@ std::tuple<cudf::table_view, std::vector<rmm::device_buffer>> superimpose_parent
   table_view const& table,
   rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+/**
+ * @brief Checks if column or any of its children are struct columns with structs that are null.
+ *
+ * This function searchings for structs that are null and differentiates between them and structs
+ * containing null values. Struct nulls add a column to the table result of the flatten column
+ * utility. The existence of struct nulls necessitates the use of column_nullability::FORCE when
+ * flattening the column for comparison.
+ *
+ * @param col Column to check for structs containing nulls
+ * @return true If the column is or contains a struct column with struct nulls
+ * @return false If the column is not a struct column or does not contain struct nulls
+ */
+bool contains_struct_nulls(column_view const& col);
 }  // namespace detail
 }  // namespace structs
 }  // namespace cudf

cpp/src/binaryop/compiled/binary_ops.cu

@@ -19,13 +19,20 @@
 
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/column/column_factories.hpp>
+#include <cudf/detail/iterator.cuh>
+#include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/scalar/scalar_device_view.cuh>
 #include <cudf/strings/detail/utilities.cuh>
+#include <cudf/table/row_operators.cuh>
+#include <cudf/table/table_device_view.cuh>
+#include <cudf/table/table_view.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
+#include <cudf/detail/structs/utilities.hpp>
+
 namespace cudf {
 namespace binops {
 namespace compiled {
@@ -239,6 +246,111 @@ struct null_considering_binop {
   }
 };
 
+rmm::device_uvector<order> get_orders(binary_operator op,
+                                      uint32_t const num_columns,
+                                      rmm::cuda_stream_view stream)
+{
+  std::vector<order> op_modifier(
+    num_columns,
+    (op == binary_operator::LESS || op == binary_operator::GREATER_EQUAL) ? order::ASCENDING
+                                                                          : order::DESCENDING);
+  return cudf::detail::make_device_uvector_async(op_modifier, stream);
+}
+
+template <typename Comparator, typename OptionalOutIter>
+void struct_compare_tabulation(mutable_column_view& out,
+                               Comparator compare,
+                               binary_operator op,
+                               OptionalOutIter optional_iter,
+                               rmm::cuda_stream_view stream)
+{
+  (op == binary_operator::EQUAL || op == binary_operator::LESS || op == binary_operator::GREATER)
+    ? thrust::tabulate(rmm::exec_policy(stream),
+                       out.begin<bool>(),
+                       out.end<bool>(),
+                       [optional_iter, compare] __device__(size_type i) {
+                         return optional_iter[i].has_value() and compare(i, i);
+                       })
+    : thrust::tabulate(rmm::exec_policy(stream),
+                       out.begin<bool>(),
+                       out.end<bool>(),
+                       [optional_iter, compare] __device__(size_type i) {
+                         return optional_iter[i].has_value() and not compare(i, i);
+                       });
+}
+
+void struct_binary_operation(mutable_column_view& out,
+                             column_view const& lhs,
+                             column_view const& rhs,
+                             binary_operator op,
+                             rmm::cuda_stream_view stream)
+{
+  bool const has_struct_nulls =
+    structs::detail::contains_struct_nulls(lhs) || structs::detail::contains_struct_nulls(rhs);
+  auto const lhs_superimposed = structs::detail::superimpose_parent_nulls(lhs);
+  auto const rhs_superimposed = structs::detail::superimpose_parent_nulls(rhs);
+  auto const lhs_flattener    = structs::detail::flatten_nested_columns(
+    table_view{{std::get<0>(lhs_superimposed)}},
+    {},
+    {},
+    has_struct_nulls ? structs::detail::column_nullability::FORCE
+                        : structs::detail::column_nullability::MATCH_INCOMING);
+  auto const rhs_flattener = structs::detail::flatten_nested_columns(
+    table_view{{std::get<0>(rhs_superimposed)}},
+    {},
+    {},
+    has_struct_nulls ? structs::detail::column_nullability::FORCE
+                     : structs::detail::column_nullability::MATCH_INCOMING);
+
+  auto lhs_flat = lhs_flattener.flattened_columns();
+  auto rhs_flat = rhs_flattener.flattened_columns();
+
+  auto d_out     = column_device_view::create(out, stream);
+  auto d_lhs     = table_device_view::create(lhs_flat);
+  auto d_rhs     = table_device_view::create(rhs_flat);
+  bool has_nulls = has_nested_nulls(lhs_flat) || has_nested_nulls(rhs_flat);
+
+  if (op == binary_operator::EQUAL || op == binary_operator::NOT_EQUAL) {
+    if (has_nulls) {
+      auto out_iter = cudf::detail::make_optional_iterator<bool>(
+        *d_out, contains_nulls::DYNAMIC{}, out.has_nulls());
+      struct_compare_tabulation(
+        out, row_equality_comparator<true>{*d_lhs, *d_rhs, true}, op, out_iter, stream);
+    } else {
+      auto out_iter = cudf::detail::make_optional_iterator<bool>(
+        *d_out, contains_nulls::DYNAMIC{}, out.has_nulls());
+      struct_compare_tabulation(
+        out, row_equality_comparator<false>{*d_lhs, *d_rhs, true}, op, out_iter, stream);
+    }
+  } else if (op == binary_operator::LESS || op == binary_operator::LESS_EQUAL ||
+             op == binary_operator::GREATER || op == binary_operator::GREATER_EQUAL) {
+    if (has_nulls) {
+      auto out_iter = cudf::detail::make_optional_iterator<bool>(
+        *d_out, contains_nulls::DYNAMIC{}, out.has_nulls());
+      struct_compare_tabulation(
+        out,
+        row_lexicographic_comparator<true>{
+          *d_lhs, *d_rhs, get_orders(op, lhs_flat.num_columns(), stream).data()},
+        op,
+        out_iter,
+        stream);
+    } else {
+      auto out_iter = cudf::detail::make_optional_iterator<bool>(
+        *d_out, contains_nulls::DYNAMIC{}, out.has_nulls());
+      struct_compare_tabulation(
+        out,
+        row_lexicographic_comparator<false>{
+          *d_lhs, *d_rhs, get_orders(op, lhs_flat.num_columns(), stream).data()},
+        op,
+        out_iter,
+        stream);
+    }
+    //  } else if (op == binary_operator::NULL_EQUALS) {
+  } else {
+    CUDF_FAIL("Unsupported operator for these types");
+  }
+}
+
 }  // namespace
 
 std::unique_ptr<column> string_null_min_max(scalar const& lhs,
@@ -351,10 +463,16 @@ void binary_operation(mutable_column_view& out,
                       binary_operator op,
                       rmm::cuda_stream_view stream)
 {
-  auto lhsd = column_device_view::create(lhs, stream);
-  auto rhsd = column_device_view::create(rhs, stream);
-  auto outd = mutable_column_device_view::create(out, stream);
-  operator_dispatcher(*outd, *lhsd, *rhsd, false, false, op, stream);
+  if (lhs.type().id() == type_id::STRUCT && rhs.type().id() == type_id::STRUCT) {
+    CUDF_EXPECTS(struct_children_support_operation(out.type(), lhs, rhs, op),
+                 "Unsupported operator for these types");
+    struct_binary_operation(out, lhs, rhs, op, stream);
+  } else {
+    auto lhsd = column_device_view::create(lhs, stream);
+    auto rhsd = column_device_view::create(rhs, stream);
+    auto outd = mutable_column_device_view::create(out, stream);
+    operator_dispatcher(*outd, *lhsd, *rhsd, false, false, op, stream);
+  }
 }
 // scalar_vector
 void binary_operation(mutable_column_view& out,
@@ -363,10 +481,17 @@ void binary_operation(mutable_column_view& out,
                       binary_operator op,
                       rmm::cuda_stream_view stream)
 {
-  auto [lhsd, aux] = scalar_to_column_device_view(lhs, stream);
-  auto rhsd        = column_device_view::create(rhs, stream);
-  auto outd        = mutable_column_device_view::create(out, stream);
-  operator_dispatcher(*outd, *lhsd, *rhsd, true, false, op, stream);
+  if (lhs.type().id() == type_id::STRUCT && rhs.type().id() == type_id::STRUCT) {
+    auto lhs_col = make_column_from_scalar(lhs, rhs.size(), stream);
+    CUDF_EXPECTS(struct_children_support_operation(out.type(), lhs_col->view(), rhs, op),
+                 "Unsupported operator for these types");
+    struct_binary_operation(out, lhs_col->view(), rhs, op, stream);
+  } else {
+    auto [lhsd, aux] = scalar_to_column_device_view(lhs, stream);
+    auto rhsd        = column_device_view::create(rhs, stream);
+    auto outd        = mutable_column_device_view::create(out, stream);
+    operator_dispatcher(*outd, *lhsd, *rhsd, true, false, op, stream);
+  }
 }
 // vector_scalar
 void binary_operation(mutable_column_view& out,
@@ -375,10 +500,18 @@ void binary_operation(mutable_column_view& out,
                       binary_operator op,
                       rmm::cuda_stream_view stream)
 {
-  auto lhsd        = column_device_view::create(lhs, stream);
-  auto [rhsd, aux] = scalar_to_column_device_view(rhs, stream);
-  auto outd        = mutable_column_device_view::create(out, stream);
-  operator_dispatcher(*outd, *lhsd, *rhsd, false, true, op, stream);
+  if (lhs.type().id() == type_id::STRUCT && rhs.type().id() == type_id::STRUCT) {
+    auto rhs_col = make_column_from_scalar(rhs, lhs.size(), stream);
+    CUDF_EXPECTS(struct_children_support_operation(out.type(), lhs, rhs_col->view(), op),
+                 "Unsupported operator for these types");
+
+    struct_binary_operation(out, lhs, rhs_col->view(), op, stream);
+  } else {
+    auto lhsd        = column_device_view::create(lhs, stream);
+    auto [rhsd, aux] = scalar_to_column_device_view(rhs, stream);
+    auto outd        = mutable_column_device_view::create(out, stream);
+    operator_dispatcher(*outd, *lhsd, *rhsd, false, true, op, stream);
+  }
 }
 
 }  // namespace compiled

cpp/src/binaryop/compiled/binary_ops.hpp

@@ -167,11 +167,25 @@ std::optional<data_type> get_common_type(data_type out, data_type lhs, data_type
  * @param out output type of the binary operation
  * @param lhs first operand type of the binary operation
  * @param rhs second operand type of the binary operation
- * @param op binary operator enum.
+ * @param op binary operator enum
  * @return true if given binary operator supports given input and output types.
  */
 bool is_supported_operation(data_type out, data_type lhs, data_type rhs, binary_operator op);
 
+/**
+ * @brief Check if input binary operation is supported for the given input and output types.
+ *
+ * @param out output type of the binary operation
+ * @param lhs left column of the binary operation
+ * @param rhs right column of the binary operation
+ * @param op binary operator enum
+ * @return true if given binary operator supports given input and output types.
+ */
+bool struct_children_support_operation(data_type out,
+                                       column_view const& lhs,
+                                       column_view const& rhs,
+                                       binary_operator op);
+
 // Defined in individual .cu files.
 /**
  * @brief Deploys single type or double type dispatcher that runs binary operation on each element

cpp/src/binaryop/compiled/util.cpp

@@ -18,6 +18,8 @@
 
 #include <cudf/binaryop.hpp>
 #include <cudf/column/column_device_view.cuh>
+#include <cudf/column/column_view.hpp>
+#include <cudf/types.hpp>
 #include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
@@ -181,6 +183,30 @@ std::optional<data_type> get_common_type(data_type out, data_type lhs, data_type
 
 bool is_supported_operation(data_type out, data_type lhs, data_type rhs, binary_operator op)
 {
+  if (lhs.id() == type_id::STRUCT && rhs.id() == type_id::STRUCT) {
+    return op == binary_operator::EQUAL || op == binary_operator::NOT_EQUAL ||
+           op == binary_operator::LESS || op == binary_operator::LESS_EQUAL ||
+           op == binary_operator::GREATER || op == binary_operator::GREATER_EQUAL;
+  }
   return double_type_dispatcher(lhs, rhs, is_supported_operation_functor{}, out, op);
 }
+
+bool struct_children_support_operation(data_type out,
+                                       column_view const& lhs,
+                                       column_view const& rhs,
+                                       binary_operator op)
+{
+  if (lhs.type().id() == type_id::STRUCT && rhs.type().id() == type_id::STRUCT) {
+    return lhs.num_children() == rhs.num_children() &&
+           std::all_of(thrust::counting_iterator<size_type>(0),
+                       thrust::counting_iterator<size_type>(lhs.num_children()),
+                       [&](size_type i) {
+                         return struct_children_support_operation(
+                           out, lhs.child(i), rhs.child(i), op);
+                       });
+
+  } else {
+    return is_supported_operation(out, lhs.type(), rhs.type(), op);
+  }
+}
 }  // namespace cudf::binops::compiled

cpp/src/structs/utilities.cpp

@@ -437,6 +437,14 @@ std::tuple<cudf::table_view, std::vector<rmm::device_buffer>> superimpose_parent
   return {table_view{superimposed_columns}, std::move(superimposed_nullmasks)};
 }
 
+bool contains_struct_nulls(column_view const& col)
+{
+  return (col.type().id() == type_id::STRUCT && col.has_nulls()) ||
+         std::all_of(col.child_begin(), col.child_end(), [](auto const& child) {
+           return contains_struct_nulls(child);
+         });
+}
+
 }  // namespace detail
 }  // namespace structs
 }  // namespace cudf

cpp/tests/CMakeLists.txt

@@ -175,6 +175,7 @@ ConfigureTest(BINARY_TEST
     binaryop/binop-compiled-test.cpp
     binaryop/binop-compiled-fixed_point-test.cpp
     binaryop/binop-generic-ptx-test.cpp
+    binaryop/binop-struct-test.cpp
     )
 
 ###################################################################################################