Support VARIANCE and STD aggregation in rolling op

cpp/include/cudf/aggregation.hpp

@@ -117,6 +117,7 @@ class rolling_aggregation : public virtual aggregation {
 
  protected:
   rolling_aggregation() {}
+  rolling_aggregation(aggregation::Kind a) : aggregation{a} {}
 };
 
 enum class udf_type : bool { CUDA, PTX };

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

@@ -314,7 +314,7 @@ class m2_aggregation : public aggregation {
 /**
  * @brief Derived class for specifying a standard deviation/variance aggregation
  */
-class std_var_aggregation : public aggregation {
+class std_var_aggregation : public rolling_aggregation {
  public:
   size_type _ddof;  ///< Delta degrees of freedom
 
@@ -328,7 +328,7 @@ class std_var_aggregation : public aggregation {
   size_t do_hash() const override { return this->aggregation::do_hash() ^ hash_impl(); }
 
  protected:
-  std_var_aggregation(aggregation::Kind k, size_type ddof) : aggregation(k), _ddof{ddof}
+  std_var_aggregation(aggregation::Kind k, size_type ddof) : rolling_aggregation(k), _ddof{ddof}
   {
     CUDF_EXPECTS(k == aggregation::STD or k == aggregation::VARIANCE,
                  "std_var_aggregation can accept only STD, VARIANCE");
@@ -342,7 +342,10 @@ class std_var_aggregation : public aggregation {
  */
 class var_aggregation final : public std_var_aggregation {
  public:
-  var_aggregation(size_type ddof) : std_var_aggregation{aggregation::VARIANCE, ddof} {}
+  var_aggregation(size_type ddof)
+    : aggregation(aggregation::VARIANCE), std_var_aggregation{aggregation::VARIANCE, ddof}
+  {
+  }
 
   std::unique_ptr<aggregation> clone() const override
   {
@@ -361,7 +364,10 @@ class var_aggregation final : public std_var_aggregation {
  */
 class std_aggregation final : public std_var_aggregation {
  public:
-  std_aggregation(size_type ddof) : std_var_aggregation{aggregation::STD, ddof} {}
+  std_aggregation(size_type ddof)
+    : aggregation(aggregation::STD), std_var_aggregation{aggregation::STD, ddof}
+  {
+  }
 
   std::unique_ptr<aggregation> clone() const override
   {

cpp/src/aggregation/aggregation.cpp

@@ -428,6 +428,8 @@ std::unique_ptr<Base> make_variance_aggregation(size_type ddof)
   return std::make_unique<detail::var_aggregation>(ddof);
 }
 template std::unique_ptr<aggregation> make_variance_aggregation<aggregation>(size_type ddof);
+template std::unique_ptr<rolling_aggregation> make_variance_aggregation<rolling_aggregation>(
+  size_type ddof);
 
 /// Factory to create a STD aggregation
 template <typename Base>
@@ -436,6 +438,8 @@ std::unique_ptr<Base> make_std_aggregation(size_type ddof)
   return std::make_unique<detail::std_aggregation>(ddof);
 }
 template std::unique_ptr<aggregation> make_std_aggregation<aggregation>(size_type ddof);
+template std::unique_ptr<rolling_aggregation> make_std_aggregation<rolling_aggregation>(
+  size_type ddof);
 
 /// Factory to create a MEDIAN aggregation
 template <typename Base>

cpp/src/rolling/rolling_detail.cuh

@@ -33,6 +33,7 @@
 #include <cudf/detail/gather.hpp>
 #include <cudf/detail/groupby/sort_helper.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
+#include <cudf/detail/unary.hpp>
 #include <cudf/detail/utilities/cuda.cuh>
 #include <cudf/detail/utilities/device_operators.cuh>
 #include <cudf/detail/valid_if.cuh>
@@ -54,12 +55,14 @@
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_scalar.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <thrust/binary_search.h>
 #include <thrust/detail/execution_policy.h>
 #include <thrust/execution_policy.h>
 #include <thrust/find.h>
 #include <thrust/iterator/counting_iterator.h>
+#include <thrust/pair.h>
 #include <thrust/transform.h>
 
 #include <memory>
@@ -279,6 +282,86 @@ struct DeviceRollingCountAll {
   }
 };
 
+/**
+ * @brief Operator for applying a VAR rolling aggregation on a single window.
+ */
+template <typename InputType>
+struct DeviceRollingVariance {
+  size_type min_periods;
+  size_type ddof;
+
+  // what operations do we support
+  template <typename T = InputType, aggregation::Kind O = aggregation::VARIANCE>
+  static constexpr bool is_supported()
+  {
+    return is_fixed_width<InputType>() and not is_chrono<InputType>();
+  }
+
+  DeviceRollingVariance(size_type _min_periods, size_type _ddof)
+    : min_periods(_min_periods), ddof{_ddof}
+  {
+  }
+
+  template <typename OutputType, bool has_nulls>
+  bool __device__ operator()(column_device_view const& input,
+                             column_device_view const&,
+                             mutable_column_device_view& output,
+                             size_type start_index,
+                             size_type end_index,
+                             size_type current_index)
+  {
+    using DeviceInputType = device_storage_type_t<InputType>;
+
+    cudf::size_type count{0};  // valid counts in the window
+
+    // Count valid observations in window
+    if (has_nulls) {
+      count = thrust::count_if(thrust::seq,
+                               thrust::make_counting_iterator(start_index),
+                               thrust::make_counting_iterator(end_index),
+                               [&input](auto i) { return input.is_valid_nocheck(i); });
+    } else {
+      count = end_index - start_index;
+    }
+
+    // Variance/Std is non-negative, thus ddof should be strictly less than valid counts.
+    // Variance/Std of a lone value is statistically meaningless
+    bool output_is_valid = (count >= min_periods) and not(count == 1) and not(count <= ddof);
+
+    if (output_is_valid) {
+      // Welford algorithm, a numerically stable, single pass algorithm
+      // See https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
+      OutputType m, m2;
+      size_type running_count;
+
+      thrust::tie(running_count, m, m2) =
+        thrust::reduce(thrust::seq,
+                       thrust::make_counting_iterator(start_index),
+                       thrust::make_counting_iterator(end_index),
+                       thrust::make_tuple(size_type{0}, OutputType{0}, OutputType{0}),
+                       [&](auto acc, auto i) {
+                         if (has_nulls and input.is_null_nocheck(i)) { return acc; }
+
+                         OutputType m_acc, m2_acc, tmp1, tmp2;
+                         size_type r_count_acc;
+                         thrust::tie(r_count_acc, m_acc, m2_acc) = acc;
+                         OutputType x = static_cast<OutputType>(input.element<DeviceInputType>(i));
+
+                         r_count_acc++;
+                         tmp1 = x - m_acc;
+                         m_acc += tmp1 / r_count_acc;
+                         tmp2 = x - m_acc;
+                         m2_acc += tmp1 * tmp2;
+                         return thrust::make_tuple(r_count_acc, m_acc, m2_acc);
+                       });
+
+      output.element<OutputType>(current_index) = m2 / (count - ddof);
+    }
+
+    return output_is_valid;
+  }
+};
+
 /**
  * @brief Operator for applying a ROW_NUMBER rolling aggregation on a single window.
  */
@@ -506,6 +589,11 @@ struct corresponding_rolling_operator<InputType, aggregation::Kind::LEAD> {
   using type = DeviceRollingLead<InputType>;
 };
 
+template <typename InputType>
+struct corresponding_rolling_operator<InputType, aggregation::Kind::VARIANCE> {
+  using type = DeviceRollingVariance<InputType>;
+};
+
 template <typename InputType>
 struct corresponding_rolling_operator<InputType, aggregation::Kind::LAG> {
   using type = DeviceRollingLag<InputType>;
@@ -527,15 +615,24 @@ struct create_rolling_operator<
   InputType,
   op,
   std::enable_if_t<corresponding_rolling_operator<InputType, op>::type::is_supported()>> {
-  template <
-    typename T                                                                     = InputType,
-    aggregation::Kind O                                                            = op,
-    std::enable_if_t<O != aggregation::Kind::LEAD && O != aggregation::Kind::LAG>* = nullptr>
+  template <typename T                                          = InputType,
+            aggregation::Kind O                                 = op,
+            std::enable_if_t<O != aggregation::Kind::LEAD && O != aggregation::Kind::LAG &&
+                             O != aggregation::Kind::VARIANCE>* = nullptr>
   auto operator()(size_type min_periods, rolling_aggregation const&)
   {
     return typename corresponding_rolling_operator<InputType, op>::type(min_periods);
   }
 
+  template <typename T                                          = InputType,
+            aggregation::Kind O                                 = op,
+            std::enable_if_t<O == aggregation::Kind::VARIANCE>* = nullptr>
+  auto operator()(size_type min_periods, rolling_aggregation const& agg)
+  {
+    return DeviceRollingVariance<InputType>{
+      min_periods, dynamic_cast<cudf::detail::var_aggregation const&>(agg)._ddof};
+  }
+
   template <typename T                                      = InputType,
             aggregation::Kind O                             = op,
             std::enable_if_t<O == aggregation::Kind::LEAD>* = nullptr>
@@ -632,6 +729,16 @@ class rolling_aggregation_preprocessor final : public cudf::detail::simple_aggre
     return {};
   }
 
+  // STD aggregations depends on VARIANCE aggregation. Each element is applied
+  // with sqaured-root in the finalize() step.
+  std::vector<std::unique_ptr<aggregation>> visit(data_type,
+                                                  cudf::detail::std_aggregation const& agg) override
+  {
+    std::vector<std::unique_ptr<aggregation>> aggs;
+    aggs.push_back(make_variance_aggregation(agg._ddof));
+    return aggs;
+  }
+
   // LEAD and LAG have custom behaviors for non fixed-width types.
   std::vector<std::unique_ptr<aggregation>> visit(
     data_type col_type, cudf::detail::lead_lag_aggregation const& agg) override
@@ -750,6 +857,12 @@ class rolling_aggregation_postprocessor final : public cudf::detail::aggregation
       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
+  void visit(cudf::detail::std_aggregation const& agg) override
+  {
+    result = detail::unary_operation(intermediate->view(), unary_operator::SQRT, stream, mr);
+  }
+
   std::unique_ptr<column> get_result()
   {
     CUDF_EXPECTS(result != nullptr,