Support nested structs in rank and dense rank

cpp/CMakeLists.txt

@@ -232,7 +232,6 @@ add_library(cudf
     src/groupby/sort/group_sum.cu
     src/groupby/sort/scan.cpp
     src/groupby/sort/group_count_scan.cu
-    src/groupby/sort/group_dense_rank_scan.cu
     src/groupby/sort/group_max_scan.cu
     src/groupby/sort/group_min_scan.cu
     src/groupby/sort/group_rank_scan.cu

cpp/src/groupby/sort/group_rank_scan.cu

@@ -16,96 +16,129 @@
 
 #include <cudf/column/column.hpp>
 #include <cudf/column/column_factories.hpp>
-#include <cudf/detail/aggregation/aggregation.cuh>
+#include <cudf/detail/utilities/device_operators.cuh>
 #include <cudf/table/row_operators.cuh>
-#include <cudf/types.hpp>
 #include <cudf/utilities/span.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
 
-#include <thrust/logical.h>
+#include <structs/utilities.hpp>
 
 namespace cudf {
 namespace groupby {
 namespace detail {
 namespace {
-template <bool has_nested_nulls>
-std::unique_ptr<column> generate_ranks(column_view const& order_by,
+/**
+ * @brief generate grouped row ranks or dense ranks using a row comparison then scan the results
+ *
+ * @tparam has_nulls if the order_by column has nulls
+ * @tparam value_resolver flag value resolver function with boolean first and row number arguments
+ * @tparam scan_operator scan function ran on the flag values
+ * @param order_by input column to generate ranks for
+ * @param group_labels ID of group that the corresponding value belongs to
+ * @param group_offsets group index offsets with group ID indices
+ * @param resolver flag value resolver
+ * @param scan_op scan operation ran on the flag results
+ * @param stream CUDA stream used for device memory operations and kernel launches
+ * @param mr Device memory resource used to allocate the returned column's device memory
+ * @return std::unique_ptr<column> rank values
+ */
+template <bool has_nulls, typename value_resolver, typename scan_operator>
+std::unique_ptr<column> rank_generator(column_view const& order_by,
                                        device_span<size_type const> group_labels,
                                        device_span<size_type const> group_offsets,
+                                       value_resolver resolver,
+                                       scan_operator scan_op,
                                        rmm::cuda_stream_view stream,
                                        rmm::mr::device_memory_resource* mr)
 {
-  auto const flat_order =
-    order_by.type().id() == type_id::STRUCT
-      ? table_view{std::vector<column_view>{order_by.child_begin(), order_by.child_end()}}
-      : table_view{{order_by}};
-  auto const d_flat_order = table_device_view::create(flat_order, stream);
-  row_equality_comparator<has_nested_nulls> comparator(*d_flat_order, *d_flat_order, true);
-  auto ranks = make_fixed_width_column(
-    data_type{type_to_id<size_type>()}, order_by.size(), mask_state::UNALLOCATED, stream, mr);
+  auto const superimposed = structs::detail::superimpose_parent_nulls(order_by, stream, mr);
+  table_view const order_table{{std::get<0>(superimposed)}};
+  auto const flattener = cudf::structs::detail::flatten_nested_columns(
+    order_table, {}, {}, structs::detail::column_nullability::MATCH_INCOMING);
+  auto const d_flat_order = table_device_view::create(std::get<0>(flattener), stream);
+  row_equality_comparator<has_nulls> comparator(*d_flat_order, *d_flat_order, true);
+  auto ranks         = make_fixed_width_column(data_type{type_to_id<size_type>()},
+                                       order_table.num_rows(),
+                                       mask_state::UNALLOCATED,
+                                       stream,
+                                       mr);
   auto mutable_ranks = ranks->mutable_view();
 
-  if (order_by.type().id() == type_id::STRUCT && order_by.has_nulls()) {
-    auto const d_col_order = column_device_view::create(order_by, stream);
-    thrust::tabulate(rmm::exec_policy(stream),
-                     mutable_ranks.begin<size_type>(),
-                     mutable_ranks.end<size_type>(),
-                     [comparator,
-                      d_col_order = *d_col_order,
-                      labels      = group_labels.data(),
-                      offsets     = group_offsets.data()] __device__(size_type row_index) {
-                       auto group_start = offsets[labels[row_index]];
-                       if (row_index == group_start) { return 1; }
-                       bool const lhs_is_null{d_col_order.is_null(row_index)};
-                       bool const rhs_is_null{d_col_order.is_null(row_index - 1)};
-                       if (lhs_is_null && rhs_is_null) {
-                         return 0;
-                       } else if (lhs_is_null != rhs_is_null) {
-                         return row_index - group_start + 1;
-                       }
-                       return comparator(row_index, row_index - 1) ? 0
-                                                                   : row_index - group_start + 1;
-                     });
-  } else {
-    thrust::tabulate(rmm::exec_policy(stream),
-                     mutable_ranks.begin<size_type>(),
-                     mutable_ranks.end<size_type>(),
-                     [comparator,
-                      labels  = group_labels.data(),
-                      offsets = group_offsets.data()] __device__(size_type row_index) {
-                       auto group_start = offsets[labels[row_index]];
-                       return row_index != group_start && comparator(row_index, row_index - 1)
-                                ? 0
-                                : row_index - group_start + 1;
-                     });
-  }
+  thrust::tabulate(
+    rmm::exec_policy(stream),
+    mutable_ranks.begin<size_type>(),
+    mutable_ranks.end<size_type>(),
+    [comparator, resolver, labels = group_labels.data(), offsets = group_offsets.data()] __device__(
+      size_type row_index) {
+      auto group_start = offsets[labels[row_index]];
+      return resolver(row_index == group_start || !comparator(row_index, row_index - 1),
+                      row_index - group_start);
+    });
 
   thrust::inclusive_scan_by_key(rmm::exec_policy(stream),
                                 group_labels.begin(),
                                 group_labels.end(),
                                 mutable_ranks.begin<size_type>(),
                                 mutable_ranks.begin<size_type>(),
                                 thrust::equal_to<size_type>{},
-                                DeviceMax{});
+                                scan_op);
 
   return ranks;
 }
 }  // namespace
+
 std::unique_ptr<column> rank_scan(column_view const& order_by,
                                   device_span<size_type const> group_labels,
                                   device_span<size_type const> group_offsets,
                                   rmm::cuda_stream_view stream,
                                   rmm::mr::device_memory_resource* mr)
 {
-  if ((order_by.type().id() == type_id::STRUCT &&
-       has_nested_nulls(
-         table_view{std::vector<column_view>{order_by.child_begin(), order_by.child_end()}})) ||
-      (order_by.type().id() != type_id::STRUCT && order_by.has_nulls())) {
-    return generate_ranks<true>(order_by, group_labels, group_offsets, stream, mr);
+  if (has_nested_nulls(table_view{{order_by}})) {
+    return rank_generator<true>(
+      order_by,
+      group_labels,
+      group_offsets,
+      [] __device__(bool equality, auto row_index) { return equality ? row_index + 1 : 0; },
+      DeviceMax{},
+      stream,
+      mr);
+  }
+  return rank_generator<false>(
+    order_by,
+    group_labels,
+    group_offsets,
+    [] __device__(bool equality, auto row_index) { return equality ? row_index + 1 : 0; },
+    DeviceMax{},
+    stream,
+    mr);
+}
+
+std::unique_ptr<column> dense_rank_scan(column_view const& order_by,
+                                        device_span<size_type const> group_labels,
+                                        device_span<size_type const> group_offsets,
+                                        rmm::cuda_stream_view stream,
+                                        rmm::mr::device_memory_resource* mr)
+{
+  if (has_nested_nulls(table_view{{order_by}})) {
+    return rank_generator<true>(
+      order_by,
+      group_labels,
+      group_offsets,
+      [] __device__(bool equality, auto row_index) { return equality; },
+      DeviceSum{},
+      stream,
+      mr);
   }
-  return generate_ranks<false>(order_by, group_labels, group_offsets, stream, mr);
+  return rank_generator<false>(
+    order_by,
+    group_labels,
+    group_offsets,
+    [] __device__(bool equality, auto row_index) { return equality; },
+    DeviceSum{},
+    stream,
+    mr);
 }
 
 }  // namespace detail

cpp/src/groupby/sort/scan.cpp

@@ -29,6 +29,8 @@
 
 #include <rmm/cuda_stream_view.hpp>
 
+#include <structs/utilities.hpp>
+
 #include <memory>
 
 namespace cudf {
@@ -114,12 +116,8 @@ void scan_result_functor::operator()<aggregation::RANK>(aggregation const& agg)
   CUDF_EXPECTS(helper.is_presorted(),
                "Rank aggregate in groupby scan requires the keys to be presorted");
   auto const order_by = get_grouped_values();
-  CUDF_EXPECTS(order_by.type().id() != type_id::LIST,
+  CUDF_EXPECTS(!cudf::structs::detail::is_or_has_nested_lists(order_by),
                "Unsupported list type in grouped rank scan.");
-  CUDF_EXPECTS(std::none_of(order_by.child_begin(),
-                            order_by.child_end(),
-                            [](auto const& col) { return is_nested(col.type()); }),
-               "Unsupported nested columns in grouped rank scan.");
 
   cache.add_result(
     col_idx,
@@ -135,12 +133,8 @@ void scan_result_functor::operator()<aggregation::DENSE_RANK>(aggregation const&
   CUDF_EXPECTS(helper.is_presorted(),
                "Dense rank aggregate in groupby scan requires the keys to be presorted");
   auto const order_by = get_grouped_values();
-  CUDF_EXPECTS(order_by.type().id() != type_id::LIST,
+  CUDF_EXPECTS(!cudf::structs::detail::is_or_has_nested_lists(order_by),
                "Unsupported list type in grouped dense_rank scan.");
-  CUDF_EXPECTS(std::none_of(order_by.child_begin(),
-                            order_by.child_end(),
-                            [](auto const& col) { return is_nested(col.type()); }),
-               "Unsupported nested columns in grouped dense_rank scan.");
 
   cache.add_result(
     col_idx,