Refactor `collect_set` to use `cudf::distinct` and `cudf::lists::distinct` #11228

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

rapids-bot merged 53 commits into rapidsai:branch-22.08 from ttnghia:refactor_collect_set

Jul 15, 2022

cpp/src/groupby/sort/aggregate.cpp

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -31,7 +31,7 @@
  
    #include <cudf/detail/unary.hpp>

    #include <cudf/dictionary/dictionary_column_view.hpp>

    #include <cudf/groupby.hpp>

    #include <cudf/lists/detail/drop_list_duplicates.hpp>

    #include <cudf/lists/detail/stream_compaction.hpp>

    #include <cudf/table/table.hpp>

    #include <cudf/table/table_view.hpp>

    #include <cudf/types.hpp>

    @@ -99,7 +99,7 @@ void aggregate_result_functor::operator()<aggregation::SUM>(aggregation const& a
  
        agg,

        detail::group_sum(

          get_grouped_values(), helper.num_groups(stream), helper.group_labels(stream), stream, mr));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::PRODUCT>(aggregation const& agg)

    @@ -111,7 +111,7 @@ void aggregate_result_functor::operator()<aggregation::PRODUCT>(aggregation cons
  
        agg,

        detail::group_product(

          get_grouped_values(), helper.num_groups(stream), helper.group_labels(stream), stream, mr));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::ARGMAX>(aggregation const& agg)

    @@ -126,7 +126,7 @@ void aggregate_result_functor::operator()<aggregation::ARGMAX>(aggregation const
  
                                            helper.key_sort_order(stream),

                                            stream,

                                            mr));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::ARGMIN>(aggregation const& agg)

    @@ -141,7 +141,7 @@ void aggregate_result_functor::operator()<aggregation::ARGMIN>(aggregation const
  
                                            helper.key_sort_order(stream),

                                            stream,

                                            mr));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::MIN>(aggregation const& agg)

    @@ -181,7 +181,7 @@ void aggregate_result_functor::operator()<aggregation::MIN>(aggregation const& a
  
      }();

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::MAX>(aggregation const& agg)

    @@ -221,7 +221,7 @@ void aggregate_result_functor::operator()<aggregation::MAX>(aggregation const& a
  
      }();

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::MEAN>(aggregation const& agg)

    @@ -248,7 +248,7 @@ void aggregate_result_functor::operator()<aggregation::MEAN>(aggregation const&
  
                                       stream,

                                       mr);

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::M2>(aggregation const& agg)

    @@ -263,7 +263,7 @@ void aggregate_result_functor::operator()<aggregation::M2>(aggregation const& ag
  
        values,

        agg,

        detail::group_m2(get_grouped_values(), mean_result, helper.group_labels(stream), stream, mr));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::VARIANCE>(aggregation const& agg)

    @@ -286,7 +286,7 @@ void aggregate_result_functor::operator()<aggregation::VARIANCE>(aggregation con
  
                                      stream,

                                      mr);

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::STD>(aggregation const& agg)

    @@ -300,7 +300,7 @@ void aggregate_result_functor::operator()<aggregation::STD>(aggregation const& a
  
      auto result = cudf::detail::unary_operation(var_result, unary_operator::SQRT, stream, mr);

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::QUANTILE>(aggregation const& agg)

    @@ -321,7 +321,7 @@ void aggregate_result_functor::operator()<aggregation::QUANTILE>(aggregation con
  
                                            stream,

                                            mr);

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::MEDIAN>(aggregation const& agg)

    @@ -341,7 +341,7 @@ void aggregate_result_functor::operator()<aggregation::MEDIAN>(aggregation const
  
                                            stream,

                                            mr);

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::NUNIQUE>(aggregation const& agg)

    @@ -358,7 +358,7 @@ void aggregate_result_functor::operator()<aggregation::NUNIQUE>(aggregation cons
  
                                          stream,

                                          mr);

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::NTH_ELEMENT>(aggregation const& agg)

    @@ -404,7 +404,7 @@ void aggregate_result_functor::operator()<aggregation::COLLECT_LIST>(aggregation
  
                                          stream,

                                          mr);

      cache.add_result(values, agg, std::move(result));

    };

    }

    template <>

    void aggregate_result_functor::operator()<aggregation::COLLECT_SET>(aggregation const& agg)

    @@ -426,9 +426,9 @@ void aggregate_result_functor::operator()<aggregation::COLLECT_SET>(aggregation
  
      cache.add_result(

        values,

        agg,

        lists::detail::drop_list_duplicates(

          lists_column_view(collect_result->view()), nulls_equal, nans_equal, stream, mr));

    };

        lists::detail::distinct(

          lists_column_view{collect_result->view()}, nulls_equal, nans_equal, stream, mr));

    }

    /**

     * @brief Perform merging for the lists that correspond to the same key value.

    @@ -455,7 +455,7 @@ void aggregate_result_functor::operator()<aggregation::MERGE_LISTS>(aggregation
  
        agg,

        detail::group_merge_lists(

          get_grouped_values(), helper.group_offsets(stream), helper.num_groups(stream), stream, mr));

    };

    }

    /**

     * @brief Perform merging for the lists corresponding to the same key value, then dropping duplicate

    @@ -473,13 +473,13 @@ void aggregate_result_functor::operator()<aggregation::MERGE_LISTS>(aggregation
  
     * column for this aggregation.

     *

     * Firstly, this aggregation performs `MERGE_LISTS` to concatenate the input lists (corresponding to

     * the same key) into intermediate lists, then it calls `lists::drop_list_duplicates` on them to

     * the same key) into intermediate lists, then it calls `lists::distinct` on them to

     * remove duplicate list entries. As such, the input (partial results) to this aggregation should be

     * generated by (distributed) `COLLECT_LIST` aggregations, not `COLLECT_SET`, to avoid unnecessarily

     * removing duplicate entries for the partial results.

     *

     * Since duplicate list entries will be removed, the parameters `null_equality` and `nan_equality`

     * are needed for calling to `lists::drop_list_duplicates`.

     * are needed for calling `lists::distinct`.

     */

    template <>

    void aggregate_result_functor::operator()<aggregation::MERGE_SETS>(aggregation const& agg)

    @@ -494,12 +494,12 @@ void aggregate_result_functor::operator()<aggregation::MERGE_SETS>(aggregation c
  
      auto const& merge_sets_agg = dynamic_cast<cudf::detail::merge_sets_aggregation const&>(agg);

      cache.add_result(values,

                       agg,

                       lists::detail::drop_list_duplicates(lists_column_view(merged_result->view()),

                                                           merge_sets_agg._nulls_equal,

                                                           merge_sets_agg._nans_equal,

                                                           stream,

                                                           mr));

    };

                       lists::detail::distinct(lists_column_view{merged_result->view()},

                                               merge_sets_agg._nulls_equal,

                                               merge_sets_agg._nans_equal,

                                               stream,

                                               mr));

    }

    /**

     * @brief Perform merging for the M2 values that correspond to the same key value.

    @@ -528,7 +528,7 @@ void aggregate_result_functor::operator()<aggregation::MERGE_M2>(aggregation con
  
        agg,

        detail::group_merge_m2(

          get_grouped_values(), helper.group_offsets(stream), helper.num_groups(stream), stream, mr));

    };

    }

    /**

     * @brief Creates column views with only valid elements in both input column views

    @@ -600,7 +600,7 @@ void aggregate_result_functor::operator()<aggregation::COVARIANCE>(aggregation c
  
                                                cov_agg._ddof,

                                                stream,

                                                mr));

    };

    }

    /**

     * @brief Perform correlation between two child columns of non-nullable struct column.

    @@ -710,7 +710,7 @@ void aggregate_result_functor::operator()<aggregation::TDIGEST>(aggregation cons
  
                         max_centroids,

                         stream,

                         mr));

    };

    }

    /**

     * @brief Generate a merged tdigest column from a grouped set of input tdigest columns.

    @@ -752,7 +752,7 @@ void aggregate_result_functor::operator()<aggregation::MERGE_TDIGEST>(aggregatio
  
                                                                  max_centroids,

                                                                  stream,

                                                                  mr));

    };

    }

    }  // namespace detail

cpp/src/reductions/collect_ops.cu

-Original file line number
+Diff line change
@@ Expand Up / @@ -18,34 +18,36 @@ @@
     #include <cudf/detail/copy_if.cuh>
     #include <cudf/detail/iterator.cuh>
     #include <cudf/detail/reduction_functions.hpp>
-    #include <cudf/lists/drop_list_duplicates.hpp>
-    #include <cudf/lists/lists_column_factories.hpp>
+    #include <cudf/detail/stream_compaction.hpp>
     #include <cudf/lists/lists_column_view.hpp>
     #include <cudf/scalar/scalar.hpp>
     #include <cudf/scalar/scalar_factories.hpp>
     namespace cudf {
     namespace reduction {
-    std::unique_ptr<scalar> drop_duplicates(list_scalar const& scalar,
-                                            null_equality nulls_equal,
-                                            nan_equality nans_equal,
-                                            rmm::cuda_stream_view stream,
-                                            rmm::mr::device_memory_resource* mr)
+    namespace {
+    /**
+     * @brief Check if we need to handle nulls in the input column.
+     *
+     * @param input The input column
+     * @param null_handling The null handling policy
+     * @return A boolean value indicating if we need to handle nulls
+     */
+    bool need_handle_nulls(column_view const& input, null_policy null_handling)
     {
-      auto list_wrapper   = lists::detail::make_lists_column_from_scalar(scalar, 1, stream, mr);
-      auto lcw            = lists_column_view(list_wrapper->view());
-      auto no_dup_wrapper = lists::drop_list_duplicates(lcw, nulls_equal, nans_equal, mr);
-      auto no_dup         = lists_column_view(no_dup_wrapper->view()).get_sliced_child(stream);
-      return make_list_scalar(no_dup, stream, mr);
+      return null_handling == null_policy::EXCLUDE && input.has_nulls();
     }
+    }  // namespace
     std::unique_ptr<scalar> collect_list(column_view const& col,
                                          null_policy null_handling,
                                          rmm::cuda_stream_view stream,
                                          rmm::mr::device_memory_resource* mr)
     {
-      if (null_handling == null_policy::EXCLUDE && col.has_nulls()) {
+      if (need_handle_nulls(col, null_handling)) {
         auto d_view             = column_device_view::create(col, stream);
         auto filter             = detail::validity_accessor(*d_view);
         auto null_purged_table  = detail::copy_if(table_view{{col}}, filter, stream, mr);
@@ Expand All / @@ -72,9 +74,27 @@ std::unique_ptr<scalar> collect_set(column_view const& col, @@
                                         rmm::cuda_stream_view stream,
                                         rmm::mr::device_memory_resource* mr)
     {
-      auto scalar = collect_list(col, null_handling, stream, mr);
-      auto ls     = dynamic_cast<list_scalar*>(scalar.get());
-      return drop_duplicates(*ls, nulls_equal, nans_equal, stream, mr);
+      // `input_as_collect_list` is the result of the input column that has been processed to obey
+      // the given null handling behavior.
+      [[maybe_unused]] auto const [input_as_collect_list, unused_scalar] = [&] {
+        if (need_handle_nulls(col, null_handling)) {
+          // Only call `collect_list` when we need to handle nulls.
+          auto scalar = collect_list(col, null_handling, stream, mr);
+          return std::pair(static_cast<list_scalar*>(scalar.get())->view(), std::move(scalar));
+        }
+        return std::pair(col, std::unique_ptr<scalar>(nullptr));
+      }();
+      auto distinct_table = detail::distinct(table_view{{input_as_collect_list}},
+                                             std::vector<size_type>{0},
+                                             duplicate_keep_option::KEEP_ANY,
+                                             nulls_equal,
+                                             nans_equal,
+                                             stream,
+                                             mr);
+      return std::make_unique<list_scalar>(std::move(distinct_table->get_column(0)), true, stream, mr);
     }
     std::unique_ptr<scalar> merge_sets(lists_column_view const& col,
@@ Expand All @@
                                        rmm::cuda_stream_view stream,
                                        rmm::mr::device_memory_resource* mr)
     {
-      auto flatten_col = col.get_sliced_child(stream);
-      auto scalar      = std::make_unique<list_scalar>(flatten_col, true, stream, mr);
-      return drop_duplicates(*scalar, nulls_equal, nans_equal, stream, mr);
+      auto flatten_col    = col.get_sliced_child(stream);
+      auto distinct_table = detail::distinct(table_view{{flatten_col}},
+                                             std::vector<size_type>{0},
+                                             duplicate_keep_option::KEEP_ANY,
+                                             nulls_equal,
+                                             nans_equal,
+                                             stream,
+                                             mr);
+      return std::make_unique<list_scalar>(std::move(distinct_table->get_column(0)), true, stream, mr);
     }
     }  // namespace reduction
@@ Expand Down @@

cpp/src/rolling/detail/rolling.cuh

            
                      Original file line number
                      Diff line number
                      Diff line change
                  
    @@ -38,7 +38,7 @@
  
    #include <cudf/detail/utilities/device_operators.cuh>

    #include <cudf/dictionary/dictionary_column_view.hpp>

    #include <cudf/dictionary/dictionary_factories.hpp>

    #include <cudf/lists/detail/drop_list_duplicates.hpp>

    #include <cudf/lists/detail/stream_compaction.hpp>

    #include <cudf/types.hpp>

    #include <cudf/utilities/bit.hpp>

    #include <cudf/utilities/error.hpp>

    @@ -928,8 +928,8 @@ class rolling_aggregation_postprocessor final : public cudf::detail::aggregation
  
                                                         stream,

                                                         rmm::mr::get_current_device_resource());

        result = lists::detail::drop_list_duplicates(

          lists_column_view(collected_list->view()), agg._nulls_equal, agg._nans_equal, stream, mr);

        result = lists::detail::distinct(

          lists_column_view{collected_list->view()}, agg._nulls_equal, agg._nans_equal, stream, mr);

      }

      // perform the element-wise square root operation on result of VARIANCE

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor `collect_set` to use `cudf::distinct` and `cudf::lists::distinct` #11228

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Refactor collect_set to use cudf::distinct and cudf::lists::distinct #11228

Uh oh!

Refactor collect_set to use cudf::distinct and cudf::lists::distinct #11228

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Refactor `collect_set` to use `cudf::distinct` and `cudf::lists::distinct` #11228

Refactor `collect_set` to use `cudf::distinct` and `cudf::lists::distinct` #11228