Add Covariance, Pearson correlation for sort groupby (libcudf)

cpp/CMakeLists.txt

@@ -219,6 +219,7 @@ add_library(cudf
     src/groupby/sort/group_argmax.cu
     src/groupby/sort/group_argmin.cu
     src/groupby/sort/group_collect.cu
+    src/groupby/sort/group_correlation.cu
     src/groupby/sort/group_count.cu
     src/groupby/sort/group_m2.cu
     src/groupby/sort/group_max.cu

cpp/include/cudf/aggregation.hpp

@@ -87,7 +87,9 @@ class aggregation {
     CUDA,            ///< CUDA UDF based reduction
     MERGE_LISTS,     ///< merge multiple lists values into one list
     MERGE_SETS,      ///< merge multiple lists values into one list then drop duplicate entries
-    MERGE_M2,        ///< merge partial values of M2 aggregation
+    MERGE_M2,        ///< merge partial values of M2 aggregation,
+    COVARIANCE,      ///< covariance between two sets of elements
+    CORRELATION,     ///< correlation between two sets of elements
     TDIGEST,         ///< create a tdigest from a set of input values
     MERGE_TDIGEST    ///< create a tdigest by merging multiple tdigests together
   };
@@ -146,6 +148,7 @@ class groupby_scan_aggregation : public virtual aggregation {
 };
 
 enum class udf_type : bool { CUDA, PTX };
+enum class correlation_type : int32_t { PEARSON, KENDALL, SPEARMAN };
 
 /// Factory to create a SUM aggregation
 template <typename Base = aggregation>
@@ -495,6 +498,26 @@ std::unique_ptr<Base> make_merge_sets_aggregation(null_equality nulls_equal = nu
 template <typename Base = aggregation>
 std::unique_ptr<Base> make_merge_m2_aggregation();
 
+/**
+ * @brief Factory to create a COVARIANCE aggregation
+ *
+ * Compute covariance between two columns.
+ * The input columns are child columns of a non-nullable struct columns.
+ */
+template <typename Base = aggregation>
+std::unique_ptr<Base> make_covariance_aggregation();
+
+/**
+ * @brief Factory to create a CORRELATION aggregation
+ *
+ * Compute correlation coefficient between two columns.
+ * The input columns are child columns of a non-nullable struct columns.
+ *
+ * @param[in] type: correlation_type
+ */
+template <typename Base = aggregation>
+std::unique_ptr<Base> make_correlation_aggregation(correlation_type type);
+
 /**
  * @brief Factory to create a TDIGEST aggregation
  *

cpp/include/cudf/detail/aggregation/aggregation.hpp

@@ -91,6 +91,10 @@ class simple_aggregations_collector {  // Declares the interface for the simple
                                                           class merge_sets_aggregation const& agg);
   virtual std::vector<std::unique_ptr<aggregation>> visit(data_type col_type,
                                                           class merge_m2_aggregation const& agg);
+  virtual std::vector<std::unique_ptr<aggregation>> visit(data_type col_type,
+                                                          class covariance_aggregation const& agg);
+  virtual std::vector<std::unique_ptr<aggregation>> visit(data_type col_type,
+                                                          class correlation_aggregation const& agg);
   virtual std::vector<std::unique_ptr<aggregation>> visit(data_type col_type,
                                                           class tdigest_aggregation const& agg);
   virtual std::vector<std::unique_ptr<aggregation>> visit(
@@ -129,6 +133,8 @@ class aggregation_finalizer {  // Declares the interface for the finalizer
   virtual void visit(class merge_lists_aggregation const& agg);
   virtual void visit(class merge_sets_aggregation const& agg);
   virtual void visit(class merge_m2_aggregation const& agg);
+  virtual void visit(class covariance_aggregation const& agg);
+  virtual void visit(class correlation_aggregation const& agg);
   virtual void visit(class tdigest_aggregation const& agg);
   virtual void visit(class merge_tdigest_aggregation const& agg);
 };
@@ -890,6 +896,57 @@ class merge_m2_aggregation final : public groupby_aggregation {
   void finalize(aggregation_finalizer& finalizer) const override { finalizer.visit(*this); }
 };
 
+/**
+ * @brief Derived aggregation class for specifying COVARIANCE aggregation
+ */
+class covariance_aggregation final : public groupby_aggregation {
+ public:
+  explicit covariance_aggregation() : aggregation{COVARIANCE} {}
+
+  std::unique_ptr<aggregation> clone() const override
+  {
+    return std::make_unique<covariance_aggregation>(*this);
+  }
+  std::vector<std::unique_ptr<aggregation>> get_simple_aggregations(
+    data_type col_type, simple_aggregations_collector& collector) const override
+  {
+    return collector.visit(col_type, *this);
+  }
+  void finalize(aggregation_finalizer& finalizer) const override { finalizer.visit(*this); }
+};
+
+/**
+ * @brief Derived aggregation class for specifying CORRELATION aggregation
+ */
+class correlation_aggregation final : public groupby_aggregation {
+ public:
+  explicit correlation_aggregation(correlation_type type) : aggregation{CORRELATION}, _type{type} {}
+  correlation_type _type;
+
+  bool is_equal(aggregation const& _other) const override
+  {
+    if (!this->aggregation::is_equal(_other)) { return false; }
+    auto const& other = dynamic_cast<correlation_aggregation const&>(_other);
+    return (_type == other._type);
+  }
+
+  size_t do_hash() const override { return this->aggregation::do_hash() ^ hash_impl(); }
+
+  std::unique_ptr<aggregation> clone() const override
+  {
+    return std::make_unique<correlation_aggregation>(*this);
+  }
+  std::vector<std::unique_ptr<aggregation>> get_simple_aggregations(
+    data_type col_type, simple_aggregations_collector& collector) const override
+  {
+    return collector.visit(col_type, *this);
+  }
+  void finalize(aggregation_finalizer& finalizer) const override { finalizer.visit(*this); }
+
+ protected:
+  size_t hash_impl() const { return std::hash<int>{}(static_cast<int>(_type)); }
+};
+
 /**
  * @brief Derived aggregation class for specifying TDIGEST aggregation
  */
@@ -1174,6 +1231,18 @@ struct target_type_impl<SourceType, aggregation::MERGE_M2> {
   using type = struct_view;
 };
 
+// Always use double for COVARIANCE
+template <typename SourceType>
+struct target_type_impl<SourceType, aggregation::COVARIANCE> {
+  using type = double;
+};
+
+// Always use double for CORRELATION
+template <typename SourceType>
+struct target_type_impl<SourceType, aggregation::CORRELATION> {
+  using type = double;
+};
+
 // Always use numeric types for TDIGEST
 template <typename Source>
 struct target_type_impl<Source,
@@ -1296,6 +1365,10 @@ CUDA_HOST_DEVICE_CALLABLE decltype(auto) aggregation_dispatcher(aggregation::Kin
       return f.template operator()<aggregation::MERGE_SETS>(std::forward<Ts>(args)...);
     case aggregation::MERGE_M2:
       return f.template operator()<aggregation::MERGE_M2>(std::forward<Ts>(args)...);
+    case aggregation::COVARIANCE:
+      return f.template operator()<aggregation::COVARIANCE>(std::forward<Ts>(args)...);
+    case aggregation::CORRELATION:
+      return f.template operator()<aggregation::CORRELATION>(std::forward<Ts>(args)...);
     case aggregation::TDIGEST:
       return f.template operator()<aggregation::TDIGEST>(std::forward<Ts>(args)...);
     case aggregation::MERGE_TDIGEST:

cpp/include/cudf/detail/aggregation/result_cache.hpp

@@ -18,20 +18,26 @@
 
 #include <cudf/column/column.hpp>
 #include <cudf/detail/aggregation/aggregation.hpp>
+#include <cudf/detail/hashing.hpp>
+#include <cudf/types.hpp>
 
 #include <unordered_map>
 
 namespace cudf {
 namespace detail {
-struct aggregation_equality {
-  bool operator()(aggregation const& lhs, aggregation const& rhs) const
+struct pair_column_aggregation_equal_to {
+  bool operator()(std::pair<column_view, aggregation const&> const& lhs,
+                  std::pair<column_view, aggregation const&> const& rhs) const
   {
-    return lhs.is_equal(rhs);
+    return is_shallow_equivalent(lhs.first, rhs.first) and lhs.second.is_equal(rhs.second);
   }
 };
 
-struct aggregation_hash {
-  size_t operator()(aggregation const& key) const noexcept { return key.do_hash(); }
+struct pair_column_aggregation_hash {
+  size_t operator()(std::pair<column_view, aggregation const&> const& key) const noexcept
+  {
+    return hash_combine(shallow_hash(key.first), key.second.do_hash());
+  }
 };
 
 class result_cache {
@@ -43,19 +49,19 @@ class result_cache {
 
   result_cache(size_t num_columns) : _cache(num_columns) {}
 
-  bool has_result(size_t col_idx, aggregation const& agg) const;
+  bool has_result(column_view const& input, aggregation const& agg) const;
 
-  void add_result(size_t col_idx, aggregation const& agg, std::unique_ptr<column>&& col);
+  void add_result(column_view const& input, aggregation const& agg, std::unique_ptr<column>&& col);
 
-  column_view get_result(size_t col_idx, aggregation const& agg) const;
+  column_view get_result(column_view const& input, aggregation const& agg) const;
 
-  std::unique_ptr<column> release_result(size_t col_idx, aggregation const& agg);
+  std::unique_ptr<column> release_result(column_view const& input, aggregation const& agg);
 
  private:
-  std::vector<std::unordered_map<std::reference_wrapper<aggregation const>,
-                                 std::pair<std::unique_ptr<aggregation>, std::unique_ptr<column>>,
-                                 aggregation_hash,
-                                 aggregation_equality>>
+  std::unordered_map<std::pair<column_view, std::reference_wrapper<aggregation const>>,
+                     std::pair<std::unique_ptr<aggregation>, std::unique_ptr<column>>,
+                     pair_column_aggregation_hash,
+                     pair_column_aggregation_equal_to>
     _cache;
 };
 

cpp/src/aggregation/aggregation.cpp

@@ -202,6 +202,16 @@ std::vector<std::unique_ptr<aggregation>> simple_aggregations_collector::visit(
   return visit(col_type, static_cast<aggregation const&>(agg));
 }
 
+std::vector<std::unique_ptr<aggregation>> simple_aggregations_collector::visit(
+  data_type col_type, covariance_aggregation const& agg)
+{
+  return visit(col_type, static_cast<aggregation const&>(agg));
+}
+std::vector<std::unique_ptr<aggregation>> simple_aggregations_collector::visit(
+  data_type col_type, correlation_aggregation const& agg)
+{
+  return visit(col_type, static_cast<aggregation const&>(agg));
+}
 std::vector<std::unique_ptr<aggregation>> simple_aggregations_collector::visit(
   data_type col_type, tdigest_aggregation const& agg)
 {
@@ -358,6 +368,16 @@ void aggregation_finalizer::visit(merge_m2_aggregation const& agg)
   visit(static_cast<aggregation const&>(agg));
 }
 
+void aggregation_finalizer::visit(covariance_aggregation const& agg)
+{
+  visit(static_cast<aggregation const&>(agg));
+}
+
+void aggregation_finalizer::visit(correlation_aggregation const& agg)
+{
+  visit(static_cast<aggregation const&>(agg));
+}
+
 void aggregation_finalizer::visit(tdigest_aggregation const& agg)
 {
   visit(static_cast<aggregation const&>(agg));
@@ -690,6 +710,26 @@ std::unique_ptr<Base> make_merge_m2_aggregation()
 template std::unique_ptr<aggregation> make_merge_m2_aggregation<aggregation>();
 template std::unique_ptr<groupby_aggregation> make_merge_m2_aggregation<groupby_aggregation>();
 
+/// Factory to create a COVARIANCE aggregation
+template <typename Base>
+std::unique_ptr<Base> make_covariance_aggregation()
+{
+  return std::make_unique<detail::covariance_aggregation>();
+}
+template std::unique_ptr<aggregation> make_covariance_aggregation<aggregation>();
+template std::unique_ptr<groupby_aggregation> make_covariance_aggregation<groupby_aggregation>();
+
+/// Factory to create a CORRELATION aggregation
+template <typename Base>
+std::unique_ptr<Base> make_correlation_aggregation(correlation_type type)
+{
+  return std::make_unique<detail::correlation_aggregation>(type);
+}
+template std::unique_ptr<aggregation> make_correlation_aggregation<aggregation>(
+  correlation_type type);
+template std::unique_ptr<groupby_aggregation> make_correlation_aggregation<groupby_aggregation>(
+  correlation_type type);
+
 template <typename Base>
 std::unique_ptr<Base> make_tdigest_aggregation(int max_centroids)
 {

cpp/src/aggregation/result_cache.cpp

@@ -19,39 +19,36 @@
 namespace cudf {
 namespace detail {
 
-bool result_cache::has_result(size_t col_idx, aggregation const& agg) const
+bool result_cache::has_result(column_view const& input, aggregation const& agg) const
 {
-  if (col_idx > _cache.size()) return false;
-
-  auto result_it = _cache[col_idx].find(agg);
-
-  return (result_it != _cache[col_idx].end());
+  return _cache.count({input, agg});
 }
 
-void result_cache::add_result(size_t col_idx, aggregation const& agg, std::unique_ptr<column>&& col)
+void result_cache::add_result(column_view const& input,
+                              aggregation const& agg,
+                              std::unique_ptr<column>&& col)
 {
   // We can't guarantee that agg will outlive the cache, so we need to take ownership of a copy.
   // To allow lookup by reference, make the key a reference and keep the owner in the value pair.
-  auto owned_agg  = agg.clone();
-  auto const& key = *owned_agg;
-  auto value      = std::make_pair(std::move(owned_agg), std::move(col));
-  _cache[col_idx].emplace(key, std::move(value));
+  auto owned_agg       = agg.clone();
+  auto const& key      = *owned_agg;
+  auto value           = std::make_pair(std::move(owned_agg), std::move(col));
+  _cache[{input, key}] = std::move(value);
 }
 
-column_view result_cache::get_result(size_t col_idx, aggregation const& agg) const
+column_view result_cache::get_result(column_view const& input, aggregation const& agg) const
 {
-  CUDF_EXPECTS(has_result(col_idx, agg), "Result does not exist in cache");
-
-  auto result_it = _cache[col_idx].find(agg);
+  auto result_it = _cache.find({input, agg});
+  CUDF_EXPECTS(result_it != _cache.end(), "Result does not exist in cache");
   return result_it->second.second->view();
 }
 
-std::unique_ptr<column> result_cache::release_result(size_t col_idx, aggregation const& agg)
+std::unique_ptr<column> result_cache::release_result(column_view const& input,
+                                                     aggregation const& agg)
 {
-  CUDF_EXPECTS(has_result(col_idx, agg), "Result does not exist in cache");
-
-  auto result_it = _cache[col_idx].extract(agg);
-  return std::move(result_it.mapped().second);
+  auto node = _cache.extract({input, agg});
+  CUDF_EXPECTS(not node.empty(), "Result does not exist in cache");
+  return std::move(node.mapped().second);
 }
 
 }  // namespace detail

cpp/src/groupby/common/utils.hpp

@@ -19,6 +19,8 @@
 #include <cudf/detail/aggregation/result_cache.hpp>
 #include <cudf/detail/groupby.hpp>
 #include <cudf/utilities/span.hpp>
+
+#include <memory>
 #include <vector>
 
 namespace cudf {
@@ -30,10 +32,24 @@ inline std::vector<aggregation_result> extract_results(host_span<RequestType con
                                                        cudf::detail::result_cache& cache)
 {
   std::vector<aggregation_result> results(requests.size());
-
+  std::unordered_map<std::pair<column_view, std::reference_wrapper<aggregation const>>,
+                     column_view,
+                     cudf::detail::pair_column_aggregation_hash,
+                     cudf::detail::pair_column_aggregation_equal_to>
+    repeated_result;
   for (size_t i = 0; i < requests.size(); i++) {
     for (auto&& agg : requests[i].aggregations) {
-      results[i].results.emplace_back(cache.release_result(i, *agg));
+      if (cache.has_result(requests[i].values, *agg)) {
+        results[i].results.emplace_back(cache.release_result(requests[i].values, *agg));
+        repeated_result[{requests[i].values, *agg}] = results[i].results.back()->view();
+      } else {
+        auto it = repeated_result.find({requests[i].values, *agg});
+        if (it != repeated_result.end()) {
+          results[i].results.emplace_back(std::make_unique<column>(it->second));
+        } else {
+          CUDF_FAIL("Cannot extract result from the cache");
+        }
+      }
     }
   }
   return results;