Struct binary comparison op functionality for spark rapids

cpp/include/cudf/binaryop.hpp

@@ -214,5 +214,47 @@ cudf::data_type binary_operation_fixed_point_output_type(binary_operator op,
                                                          cudf::data_type const& lhs,
                                                          cudf::data_type const& rhs);
 
+namespace binops {
+
+/**
+ * @brief Computes output valid mask for op between a column and a scalar
+ *
+ * @param col     Column to compute the valid mask from
+ * @param s       Scalar to compute the valid mask from
+ * @param stream  CUDA stream used for device memory operations and kernel launches
+ * @param mr      Device memory resource used to allocate the returned valid mask
+ * @return        Computed validity mask
+ */
+rmm::device_buffer scalar_col_valid_mask_and(
+  column_view const& col,
+  scalar const& s,
+  rmm::cuda_stream_view stream        = cudf::default_stream_value,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+namespace compiled {
+namespace detail {
+
+/**
+ * @brief struct binary operation using `NaN` aware sorting physical element comparators
+ *
+ * @param out mutable view of output column
+ * @param lhs view of left operand column
+ * @param rhs view of right operand column
+ * @param is_lhs_scalar true if @p lhs is a single element column representing a scalar
+ * @param is_rhs_scalar true if @p rhs is a single element column representing a scalar
+ * @param op binary operator identifier
+ * @param stream CUDA stream used for device memory operations
+ */
+void apply_sorting_struct_binary_op(mutable_column_view& out,
+                                    column_view const& lhs,
+                                    column_view const& rhs,
+                                    bool is_lhs_scalar,
+                                    bool is_rhs_scalar,
+                                    binary_operator op,
+                                    rmm::cuda_stream_view stream = cudf::default_stream_value);
+}  // namespace detail
+}  // namespace compiled
+}  // namespace binops
+
 /** @} */  // end of group
 }  // namespace cudf

cpp/src/binaryop/compiled/binary_ops.cu

@@ -16,9 +16,11 @@
 
 #include "binary_ops.hpp"
 #include "operation.cuh"
+#include "struct_binary_ops.cuh"
 
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/column/column_factories.hpp>
+#include <cudf/detail/structs/utilities.hpp>
 #include <cudf/scalar/scalar_device_view.cuh>
 #include <cudf/strings/detail/utilities.cuh>
 
@@ -44,15 +46,23 @@ namespace {
  */
 struct scalar_as_column_view {
   using return_type = typename std::pair<column_view, std::unique_ptr<column>>;
-  template <typename T, std::enable_if_t<(is_fixed_width<T>())>* = nullptr>
+  template <typename T, CUDF_ENABLE_IF(is_fixed_width<T>())>
   return_type operator()(scalar const& s, rmm::cuda_stream_view, rmm::mr::device_memory_resource*)
   {
     auto& h_scalar_type_view = static_cast<cudf::scalar_type_t<T>&>(const_cast<scalar&>(s));
     auto col_v =
       column_view(s.type(), 1, h_scalar_type_view.data(), (bitmask_type const*)s.validity_data());
     return std::pair{col_v, std::unique_ptr<column>(nullptr)};
   }
-  template <typename T, std::enable_if_t<(!is_fixed_width<T>())>* = nullptr>
+  template <typename T, CUDF_ENABLE_IF(is_struct<T>())>
+  return_type operator()(scalar const& s,
+                         rmm::cuda_stream_view stream,
+                         rmm::mr::device_memory_resource* mr)
+  {
+    auto col = make_column_from_scalar(s, 1, stream, mr);
+    return std::pair{col->view(), std::move(col)};
+  }
+  template <typename T, CUDF_ENABLE_IF(!is_fixed_width<T>() and !is_struct<T>())>
   return_type operator()(scalar const&, rmm::cuda_stream_view, rmm::mr::device_memory_resource*)
   {
     CUDF_FAIL("Unsupported type");
@@ -375,6 +385,79 @@ void binary_operation(mutable_column_view& out,
   auto [rhsv, aux] = scalar_to_column_view(rhs, stream);
   operator_dispatcher(out, lhs, rhsv, false, true, op, stream);
 }
+
+namespace detail {
+void apply_sorting_struct_binary_op(mutable_column_view& out,
+                                    column_view const& lhs,
+                                    column_view const& rhs,
+                                    bool is_lhs_scalar,
+                                    bool is_rhs_scalar,
+                                    binary_operator op,
+                                    rmm::cuda_stream_view stream)
+{
+  CUDF_EXPECTS(lhs.type().id() == type_id::STRUCT && rhs.type().id() == type_id::STRUCT,
+               "Both columns must be struct columns");
+  CUDF_EXPECTS(!cudf::structs::detail::is_or_has_nested_lists(lhs) and
+                 !cudf::structs::detail::is_or_has_nested_lists(rhs),
+               "Lists not supported");
+  // Struct child column type and structure mismatches are caught within the two_table_comparator
+  switch (op) {
+    case binary_operator::EQUAL: [[fallthrough]];
+    case binary_operator::NOT_EQUAL:
+      detail::apply_struct_equality_op(
+        out,
+        lhs,
+        rhs,
+        is_lhs_scalar,
+        is_rhs_scalar,
+        op,
+        cudf::experimental::row::equality::nan_equal_physical_equality_comparator{},
+        stream);
+      break;
+    case binary_operator::LESS:
+      detail::apply_struct_binary_op<ops::Less>(
+        out,
+        lhs,
+        rhs,
+        is_lhs_scalar,
+        is_rhs_scalar,
+        cudf::experimental::row::lexicographic::sorting_physical_element_comparator{},
+        stream);
+      break;
+    case binary_operator::GREATER:
+      detail::apply_struct_binary_op<ops::Greater>(
+        out,
+        lhs,
+        rhs,
+        is_lhs_scalar,
+        is_rhs_scalar,
+        cudf::experimental::row::lexicographic::sorting_physical_element_comparator{},
+        stream);
+      break;
+    case binary_operator::LESS_EQUAL:
+      detail::apply_struct_binary_op<ops::LessEqual>(
+        out,
+        lhs,
+        rhs,
+        is_lhs_scalar,
+        is_rhs_scalar,
+        cudf::experimental::row::lexicographic::sorting_physical_element_comparator{},
+        stream);
+      break;
+    case binary_operator::GREATER_EQUAL:
+      detail::apply_struct_binary_op<ops::GreaterEqual>(
+        out,
+        lhs,
+        rhs,
+        is_lhs_scalar,
+        is_rhs_scalar,
+        cudf::experimental::row::lexicographic::sorting_physical_element_comparator{},
+        stream);
+      break;
+    default: CUDF_FAIL("Unsupported operator for structs");
+  }
+}
+}  // namespace detail
 }  // namespace compiled
 }  // namespace binops
 }  // namespace cudf

cpp/src/binaryop/compiled/binary_ops.cuh

@@ -272,10 +272,7 @@ void for_each(rmm::cuda_stream_view stream, cudf::size_type size, Functor f)
   const int grid_size = util::div_rounding_up_safe(size, 2 * block_size);
   for_each_kernel<<<grid_size, block_size, 0, stream.value()>>>(size, std::forward<Functor&&>(f));
 }
-namespace detail {
-template <class T, class... Ts>
-inline constexpr bool is_any_v = std::disjunction<std::is_same<T, Ts>...>::value;
-}
+
 template <class BinaryOperator>
 void apply_binary_op(mutable_column_view& out,
                      column_view const& lhs,

cpp/src/binaryop/compiled/struct_binary_ops.cuh

@@ -0,0 +1,127 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "binary_ops.hpp"
+#include "operation.cuh"
+
+#include <cudf/binaryop.hpp>
+#include <cudf/column/column_device_view.cuh>
+#include <cudf/column/column_view.hpp>
+#include <cudf/detail/iterator.cuh>
+#include <cudf/table/experimental/row_operators.cuh>
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/exec_policy.hpp>
+
+namespace cudf::binops::compiled::detail {
+template <class T, class... Ts>
+inline constexpr bool is_any_v = std::disjunction<std::is_same<T, Ts>...>::value;
+
+template <class BinaryOperator,
+          typename PhysicalElementComparator =
+            cudf::experimental::row::lexicographic::sorting_physical_element_comparator>
+void apply_struct_binary_op(mutable_column_view& out,
+                            column_view const& lhs,
+                            column_view const& rhs,
+                            bool is_lhs_scalar,
+                            bool is_rhs_scalar,
+                            PhysicalElementComparator c  = {},
+                            rmm::cuda_stream_view stream = cudf::default_stream_value)
+{
+  auto compare_orders   = std::vector<order>(lhs.size(),
+                                           is_any_v<BinaryOperator, ops::Greater, ops::GreaterEqual>
+                                               ? order::DESCENDING
+                                               : order::ASCENDING);
+  auto tlhs             = table_view{{lhs}};
+  auto trhs             = table_view{{rhs}};
+  auto table_comparator = cudf::experimental::row::lexicographic::two_table_comparator{
+    tlhs, trhs, compare_orders, {}, stream};
+  auto outd = column_device_view::create(out, stream);
+  auto optional_iter =
+    cudf::detail::make_optional_iterator<bool>(*outd, nullate::DYNAMIC{out.has_nulls()});
+
+  if (is_any_v<BinaryOperator, ops::LessEqual, ops::GreaterEqual>) {
+    auto device_comparator = table_comparator.less_equivalent(
+      nullate::DYNAMIC{nullate::DYNAMIC{has_nested_nulls(tlhs) || has_nested_nulls(trhs)}}, c);
+    thrust::tabulate(
+      rmm::exec_policy(stream),
+      out.begin<bool>(),
+      out.end<bool>(),
+      [optional_iter, is_lhs_scalar, is_rhs_scalar, device_comparator] __device__(size_type i) {
+        return optional_iter[i].has_value() &&
+               device_comparator(cudf::experimental::row::lhs_index_type{is_lhs_scalar ? 0 : i},
+                                 cudf::experimental::row::rhs_index_type{is_rhs_scalar ? 0 : i});
+      });
+
+  } else {
+    auto device_comparator = table_comparator.less(
+      nullate::DYNAMIC{nullate::DYNAMIC{has_nested_nulls(tlhs) || has_nested_nulls(trhs)}}, c);
+    thrust::tabulate(
+      rmm::exec_policy(stream),
+      out.begin<bool>(),
+      out.end<bool>(),
+      [optional_iter, is_lhs_scalar, is_rhs_scalar, device_comparator] __device__(size_type i) {
+        return optional_iter[i].has_value() &&
+               device_comparator(cudf::experimental::row::lhs_index_type{is_lhs_scalar ? 0 : i},
+                                 cudf::experimental::row::rhs_index_type{is_rhs_scalar ? 0 : i});
+      });
+  }
+  return;
+}
+
+template <typename PhysicalEqualityComparator =
+            cudf::experimental::row::equality::physical_equality_comparator>
+void apply_struct_equality_op(mutable_column_view& out,
+                              column_view const& lhs,
+                              column_view const& rhs,
+                              bool is_lhs_scalar,
+                              bool is_rhs_scalar,
+                              binary_operator op,
+                              PhysicalEqualityComparator c = {},
+                              rmm::cuda_stream_view stream = cudf::default_stream_value)
+{
+  CUDF_EXPECTS(op == binary_operator::EQUAL || op == binary_operator::NOT_EQUAL,
+               "Unsupported operator for these types");
+
+  auto tlhs = table_view{{lhs}};
+  auto trhs = table_view{{rhs}};
+  auto table_comparator =
+    cudf::experimental::row::equality::two_table_comparator{tlhs, trhs, stream};
+  auto device_comparator = table_comparator.equal_to(
+    nullate::DYNAMIC{nullate::DYNAMIC{has_nested_nulls(tlhs) || has_nested_nulls(trhs)}},
+    null_equality::EQUAL,
+    c);
+
+  auto outd = column_device_view::create(out, stream);
+  auto optional_iter =
+    cudf::detail::make_optional_iterator<bool>(*outd, nullate::DYNAMIC{out.has_nulls()});
+  thrust::tabulate(rmm::exec_policy(stream),
+                   out.begin<bool>(),
+                   out.end<bool>(),
+                   [optional_iter,
+                    is_lhs_scalar,
+                    is_rhs_scalar,
+                    preserve_output = (op != binary_operator::NOT_EQUAL),
+                    device_comparator] __device__(size_type i) {
+                     auto lhs = cudf::experimental::row::lhs_index_type{is_lhs_scalar ? 0 : i};
+                     auto rhs = cudf::experimental::row::rhs_index_type{is_rhs_scalar ? 0 : i};
+                     return optional_iter[i].has_value() and
+                            (device_comparator(lhs, rhs) == preserve_output);
+                   });
+}
+}  // namespace cudf::binops::compiled::detail

java/src/main/native/include/maps_column_view.hpp

@@ -141,7 +141,7 @@ class maps_column_view {
    */
 
   std::unique_ptr<column>
-  contains(column_view const &key, rmm::cuda_stream_view stream = rmm::cuda_stream_default,
+  contains(column_view const &key, rmm::cuda_stream_view stream = cudf::default_stream_value,
            rmm::mr::device_memory_resource *mr = rmm::mr::get_current_device_resource()) const;
 
 private:

java/src/main/native/src/ColumnViewJni.cpp

@@ -24,6 +24,7 @@
 #include <cudf/column/column_factories.hpp>
 #include <cudf/concatenate.hpp>
 #include <cudf/datetime.hpp>
+#include <cudf/detail/null_mask.hpp>
 #include <cudf/filling.hpp>
 #include <cudf/hashing.hpp>
 #include <cudf/lists/contains.hpp>
@@ -1289,9 +1290,18 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_binaryOpVV(JNIEnv *env, j
     cudf::jni::auto_set_device(env);
     auto lhs = reinterpret_cast<cudf::column_view *>(lhs_view);
     auto rhs = reinterpret_cast<cudf::column_view *>(rhs_view);
-
     cudf::data_type n_data_type = cudf::jni::make_data_type(out_dtype, scale);
     cudf::binary_operator op = static_cast<cudf::binary_operator>(int_op);
+
+    if (lhs->type().id() == cudf::type_id::STRUCT) {
+      auto [new_mask, null_count] = cudf::bitmask_and(cudf::table_view{{*lhs, *rhs}});
+      auto out = make_fixed_width_column(n_data_type, lhs->size(), std::move(new_mask), null_count);
+      auto out_view = out->mutable_view();
+      cudf::binops::compiled::detail::apply_sorting_struct_binary_op(out_view, *lhs, *rhs, false,
+                                                                     false, op);
+      return release_as_jlong(out);
+    }
+
     return release_as_jlong(cudf::binary_operation(*lhs, *rhs, op, n_data_type));
   }
   CATCH_STD(env, 0);
@@ -1320,8 +1330,19 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_binaryOpVS(JNIEnv *env, j
     auto lhs = reinterpret_cast<cudf::column_view *>(lhs_view);
     cudf::scalar *rhs = reinterpret_cast<cudf::scalar *>(rhs_ptr);
     cudf::data_type n_data_type = cudf::jni::make_data_type(out_dtype, scale);
-
     cudf::binary_operator op = static_cast<cudf::binary_operator>(int_op);
+
+    if (lhs->type().id() == cudf::type_id::STRUCT) {
+      auto new_mask = cudf::binops::scalar_col_valid_mask_and(*lhs, *rhs);
+      auto out = make_fixed_width_column(n_data_type, lhs->size(), std::move(new_mask),
+                                         cudf::UNKNOWN_NULL_COUNT);
+      auto rhsv = cudf::make_column_from_scalar(*rhs, 1);
+      auto out_view = out->mutable_view();
+      cudf::binops::compiled::detail::apply_sorting_struct_binary_op(out_view, *lhs, rhsv->view(),
+                                                                     false, true, op);
+      return release_as_jlong(out);
+    }
+
     return release_as_jlong(cudf::binary_operation(*lhs, *rhs, op, n_data_type));
   }
   CATCH_STD(env, 0);