grouped_rolling_test.cpp

/*
 * Copyright (c) 2019-2020, 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.
 */

#include <tests/utilities/base_fixture.hpp>
#include <tests/utilities/column_utilities.hpp>
#include <tests/utilities/column_wrapper.hpp>
#include <tests/utilities/cudf_gtest.hpp>
#include <tests/utilities/type_lists.hpp>

#include <cudf/aggregation.hpp>
#include <cudf/detail/aggregation/aggregation.hpp>
#include <cudf/rolling.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/utilities/bit.hpp>
#include <src/rolling/rolling_detail.hpp>

#include <thrust/iterator/constant_iterator.h>

#include <algorithm>
#include <vector>

using cudf::bitmask_type;

const std::string cuda_func{
  R"***(
    template <typename OutType, typename InType>
    __device__ void CUDA_GENERIC_AGGREGATOR(OutType *ret, InType *in_col, cudf::size_type start,
                                            cudf::size_type count) {
      OutType val = 0;
      for (cudf::size_type i = 0; i < count; i++) {
        val += in_col[start + i];
      }
      *ret = val;
    }
  )***"};

const std::string ptx_func{
  R"***(
  //
  // Generated by NVIDIA NVVM Compiler
  //
  // Compiler Build ID: CL-24817639
  // Cuda compilation tools, release 10.0, V10.0.130
  // Based on LLVM 3.4svn
  //

  .version 6.3
  .target sm_70
  .address_size 64

  // .globl	_ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE
  .common .global .align 8 .u64 _ZN08NumbaEnv8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE;

  .visible .func  (.param .b32 func_retval0) _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE(
  .param .b64 _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_0,
  .param .b64 _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_1,
  .param .b64 _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_2,
  .param .b64 _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_3,
  .param .b64 _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_4,
  .param .b64 _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_5,
  .param .b64 _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_6,
  .param .b64 _ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_7
  )
  {
  .reg .pred 	%p<3>;
  .reg .b32 	%r<6>;
  .reg .b64 	%rd<18>;


  ld.param.u64 	%rd6, [_ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_0];
  ld.param.u64 	%rd7, [_ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_5];
  ld.param.u64 	%rd8, [_ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_6];
  ld.param.u64 	%rd9, [_ZN8__main__7add$241E5ArrayIiLi1E1A7mutable7alignedE_paam_7];
  mov.u64 	%rd15, 0;
  mov.u64 	%rd16, %rd15;

  BB0_1:
  mov.u64 	%rd2, %rd16;
  mov.u32 	%r5, 0;
  setp.ge.s64	%p1, %rd15, %rd8;
  mov.u64 	%rd17, %rd15;
  @%p1 bra 	BB0_3;

  mul.lo.s64 	%rd12, %rd15, %rd9;
  add.s64 	%rd13, %rd12, %rd7;
  ld.u32 	%r5, [%rd13];
  add.s64 	%rd17, %rd15, 1;

  BB0_3:
  cvt.s64.s32	%rd14, %r5;
  add.s64 	%rd16, %rd14, %rd2;
  setp.lt.s64	%p2, %rd15, %rd8;
  mov.u64 	%rd15, %rd17;
  @%p2 bra 	BB0_1;

  st.u64 	[%rd6], %rd2;
  mov.u32 	%r4, 0;
  st.param.b32	[func_retval0+0], %r4;
  ret;
  }
  )***"};

using cudf::size_type;
using cudf::test::fixed_width_column_wrapper;

template <typename T>
class GroupedRollingTest : public cudf::test::BaseFixture {
 protected:
  // input as column_wrapper
  void run_test_col(cudf::table_view const& keys,
                    cudf::column_view const& input,
                    std::vector<size_type> const& expected_grouping,
                    size_type const& preceding_window,
                    size_type const& following_window,
                    size_type min_periods,
                    std::unique_ptr<cudf::aggregation> const& op)
  {
    std::unique_ptr<cudf::column> output;

    // wrap windows
    EXPECT_NO_THROW(output = cudf::grouped_rolling_window(
                      keys, input, preceding_window, following_window, min_periods, op));

    auto reference = create_reference_output(
      op, input, expected_grouping, preceding_window, following_window, min_periods);

#ifndef NDEBUG
    std::cout << "input:\n";
    cudf::test::print(input, std::cout, ", ");
    std::cout << "\n";
    std::cout << "output:\n";
    cudf::test::print(*output, std::cout, ", ");
    std::cout << "\n";
    std::cout << "reference:\n";
    cudf::test::print(*reference, std::cout, ", ");
    std::cout << "\n";
    std::cout << "\n";
#endif

    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*output, *reference);
  }

  void run_test_col_agg(cudf::table_view const& keys,
                        cudf::column_view const& input,
                        std::vector<size_type> const& expected_grouping,
                        size_type preceding_window,
                        size_type following_window,
                        size_type min_periods)
  {
    // Skip grouping-tests on bool8 keys. sort_helper does not support this.
    if (keys.num_columns() > 0 && cudf::is_boolean(keys.column(0).type())) { return; }

    // test all supported aggregators
    run_test_col(keys,
                 input,
                 expected_grouping,
                 preceding_window,
                 following_window,
                 min_periods,
                 cudf::make_min_aggregation());
    run_test_col(keys,
                 input,
                 expected_grouping,
                 preceding_window,
                 following_window,
                 min_periods,
                 cudf::make_count_aggregation());
    run_test_col(keys,
                 input,
                 expected_grouping,
                 preceding_window,
                 following_window,
                 min_periods,
                 cudf::make_count_aggregation(cudf::null_policy::INCLUDE));
    run_test_col(keys,
                 input,
                 expected_grouping,
                 preceding_window,
                 following_window,
                 min_periods,
                 cudf::make_max_aggregation());

    if (!cudf::is_timestamp(input.type())) {
      run_test_col(keys,
                   input,
                   expected_grouping,
                   preceding_window,
                   following_window,
                   min_periods,
                   cudf::make_sum_aggregation());
      run_test_col(keys,
                   input,
                   expected_grouping,
                   preceding_window,
                   following_window,
                   min_periods,
                   cudf::make_mean_aggregation());
    }
    run_test_col(keys,
                 input,
                 expected_grouping,
                 preceding_window,
                 following_window,
                 min_periods,
                 cudf::make_row_number_aggregation());

    // >>> test UDFs <<<
    if (input.type() == cudf::data_type{cudf::type_id::INT32} && !input.has_nulls()) {
      auto cuda_udf_agg = cudf::make_udf_aggregation(
        cudf::udf_type::CUDA, cuda_func, cudf::data_type{cudf::type_id::INT64});
      run_test_col(keys,
                   input,
                   expected_grouping,
                   preceding_window,
                   following_window,
                   min_periods,
                   cuda_udf_agg);

      auto ptx_udf_agg = cudf::make_udf_aggregation(
        cudf::udf_type::PTX, ptx_func, cudf::data_type{cudf::type_id::INT64});
      run_test_col(keys,
                   input,
                   expected_grouping,
                   preceding_window,
                   following_window,
                   min_periods,
                   ptx_udf_agg);
    }
  }

 private:
  // use SFINAE to only instantiate for supported combinations

  // specialization for COUNT_VALID, COUNT_ALL
  template <bool include_nulls>
  std::unique_ptr<cudf::column> create_count_reference_output(
    cudf::column_view const& input,
    std::vector<size_type> const& group_offsets,
    size_type const& preceding_window,
    size_type const& following_window,
    size_type min_periods)
  {
    size_type num_rows = input.size();
    thrust::host_vector<cudf::size_type> ref_data(num_rows);
    thrust::host_vector<bool> ref_valid(num_rows);

    // input data and mask

    std::vector<bitmask_type> in_valid = cudf::test::bitmask_to_host(input);
    bitmask_type* valid_mask           = in_valid.data();

    for (size_type i = 0; i < num_rows; i++) {
      // load sizes
      min_periods = std::max(min_periods, 1);  // at least one observation is required

      // compute bounds
      auto group_end_index   = std::upper_bound(group_offsets.begin(), group_offsets.end(), i);
      auto group_start_index = group_end_index - 1;

      size_type start       = std::min(num_rows, std::max(0, i - preceding_window + 1));
      size_type end         = std::min(num_rows, std::max(0, i + following_window + 1));
      size_type start_index = std::max(*group_start_index, std::min(start, end));
      size_type end_index   = std::min(*group_end_index, std::max(start, end));

      // aggregate
      size_type count = 0;
      for (size_type j = start_index; j < end_index; j++) {
        if (include_nulls || !input.nullable() || cudf::bit_is_set(valid_mask, j)) count++;
      }

      ref_valid[i] = (count >= min_periods);
      if (ref_valid[i]) ref_data[i] = count;
    }

    fixed_width_column_wrapper<cudf::size_type> col(
      ref_data.begin(), ref_data.end(), ref_valid.begin());
    return col.release();
  }

  std::unique_ptr<cudf::column> create_row_number_reference_output(
    cudf::column_view const& input,
    std::vector<size_type> const& group_offsets,
    size_type const& preceding_window,
    size_type const& following_window,
    size_type min_periods)
  {
    size_type num_rows = input.size();
    thrust::host_vector<cudf::size_type> ref_data(num_rows);
    thrust::host_vector<bool> ref_valid(num_rows);

    // input data and mask

    std::vector<bitmask_type> in_valid = cudf::test::bitmask_to_host(input);
    bitmask_type* valid_mask           = in_valid.data();

    for (size_type i = 0; i < num_rows; i++) {
      // load sizes
      min_periods = std::max(min_periods, 1);  // at least one observation is required

      // compute bounds
      auto group_end_index   = std::upper_bound(group_offsets.begin(), group_offsets.end(), i);
      auto group_start_index = group_end_index - 1;

      size_type start       = std::min(num_rows, std::max(0, i - preceding_window + 1));
      size_type end         = std::min(num_rows, std::max(0, i + following_window + 1));
      size_type start_index = std::max(*group_start_index, std::min(start, end));
      size_type end_index   = std::min(*group_end_index, std::max(start, end));

      // aggregate
      size_type count{end_index - start_index};
      size_type row_number{i - start_index + 1};

      ref_valid[i] = (count >= min_periods);
      ref_data[i]  = row_number;
    }

    fixed_width_column_wrapper<cudf::size_type> col(
      ref_data.begin(), ref_data.end(), ref_valid.begin());
    return col.release();
  }

  template <typename agg_op,
            cudf::aggregation::Kind k,
            typename OutputType,
            bool is_mean,
            std::enable_if_t<cudf::detail::is_rolling_supported<T, agg_op, k>()>* = nullptr>
  std::unique_ptr<cudf::column> create_reference_output(cudf::column_view const& input,
                                                        std::vector<size_type> const& group_offsets,
                                                        size_type const& preceding_window,
                                                        size_type const& following_window,
                                                        size_type min_periods)
  {
    size_type num_rows = input.size();
    thrust::host_vector<OutputType> ref_data(num_rows);
    thrust::host_vector<bool> ref_valid(num_rows);

    // input data and mask
    thrust::host_vector<T> in_col;
    std::vector<bitmask_type> in_valid;
    std::tie(in_col, in_valid) = cudf::test::to_host<T>(input);
    bitmask_type* valid_mask   = in_valid.data();

    agg_op op;
    for (size_type i = 0; i < num_rows; i++) {
      OutputType val = agg_op::template identity<OutputType>();

      // load sizes
      min_periods = std::max(min_periods, 1);  // at least one observation is required

      // compute bounds
      auto group_end_index   = std::upper_bound(group_offsets.begin(), group_offsets.end(), i);
      auto group_start_index = group_end_index - 1;

      size_type start = std::min(
        num_rows, std::max(0, i - preceding_window + 1));  // Preceding window includes current row.
      size_type end         = std::min(num_rows, std::max(0, i + following_window + 1));
      size_type start_index = std::max(*group_start_index, std::min(start, end));
      size_type end_index   = std::min(*group_end_index, std::max(start, end));

      // aggregate
      size_type count = 0;
      for (size_type j = start_index; j < end_index; j++) {
        if (!input.nullable() || cudf::bit_is_set(valid_mask, j)) {
          val = op(static_cast<OutputType>(in_col[j]), val);
          count++;
        }
      }

      ref_valid[i] = (count >= min_periods);
      if (ref_valid[i]) {
        cudf::detail::rolling_store_output_functor<OutputType, is_mean>{}(ref_data[i], val, count);
      }
    }

    fixed_width_column_wrapper<OutputType> col(ref_data.begin(), ref_data.end(), ref_valid.begin());
    return col.release();
  }

  template <typename agg_op,
            cudf::aggregation::Kind k,
            typename OutputType,
            bool is_mean,
            std::enable_if_t<!cudf::detail::is_rolling_supported<T, agg_op, k>()>* = nullptr>
  std::unique_ptr<cudf::column> create_reference_output(cudf::column_view const& input,
                                                        std::vector<size_type> const& group_offsets,
                                                        size_type const& preceding_window_col,
                                                        size_type const& following_window_col,
                                                        size_type min_periods)
  {
    CUDF_FAIL("Unsupported combination of type and aggregation");
  }

  std::unique_ptr<cudf::column> create_reference_output(
    std::unique_ptr<cudf::aggregation> const& op,
    cudf::column_view const& input,
    std::vector<size_type> const& group_offsets,
    size_type const& preceding_window,
    size_type const& following_window,
    size_type min_periods)
  {
    // unroll aggregation types
    switch (op->kind) {
      case cudf::aggregation::SUM:
        return create_reference_output<cudf::DeviceSum,
                                       cudf::aggregation::SUM,
                                       cudf::detail::target_type_t<T, cudf::aggregation::SUM>,
                                       false>(
          input, group_offsets, preceding_window, following_window, min_periods);
      case cudf::aggregation::MIN:
        return create_reference_output<cudf::DeviceMin,
                                       cudf::aggregation::MIN,
                                       cudf::detail::target_type_t<T, cudf::aggregation::MIN>,
                                       false>(
          input, group_offsets, preceding_window, following_window, min_periods);
      case cudf::aggregation::MAX:
        return create_reference_output<cudf::DeviceMax,
                                       cudf::aggregation::MAX,
                                       cudf::detail::target_type_t<T, cudf::aggregation::MAX>,
                                       false>(
          input, group_offsets, preceding_window, following_window, min_periods);
      case cudf::aggregation::COUNT_VALID:
        return create_count_reference_output<false>(
          input, group_offsets, preceding_window, following_window, min_periods);
      case cudf::aggregation::COUNT_ALL:
        return create_count_reference_output<true>(
          input, group_offsets, preceding_window, following_window, min_periods);
      case cudf::aggregation::ROW_NUMBER:
        return create_row_number_reference_output(
          input, group_offsets, preceding_window, following_window, min_periods);
      case cudf::aggregation::MEAN:
        return create_reference_output<cudf::DeviceSum,
                                       cudf::aggregation::MEAN,
                                       cudf::detail::target_type_t<T, cudf::aggregation::MEAN>,
                                       true>(
          input, group_offsets, preceding_window, following_window, min_periods);
      // >>> UDFs <<<
      case cudf::aggregation::CUDA:
        return create_reference_output<cudf::DeviceSum,
                                       cudf::aggregation::SUM,
                                       cudf::detail::target_type_t<T, cudf::aggregation::SUM>,
                                       false>(
          input, group_offsets, preceding_window, following_window, min_periods);
      case cudf::aggregation::PTX:
        return create_reference_output<cudf::DeviceSum,
                                       cudf::aggregation::SUM,
                                       cudf::detail::target_type_t<T, cudf::aggregation::SUM>,
                                       false>(
          input, group_offsets, preceding_window, following_window, min_periods);
      default: return fixed_width_column_wrapper<T>({}).release();
    }
  }
};

// // ------------- expected failures --------------------

class GroupedRollingErrorTest : public cudf::test::BaseFixture {
};

// negative sizes
TEST_F(GroupedRollingErrorTest, NegativeMinPeriods)
{
  // Construct agg column.
  const std::vector<size_type> col_data{0, 1, 2, 0, 4};
  const std::vector<bool> col_valid{1, 1, 1, 0, 1};
  fixed_width_column_wrapper<size_type> input{col_data.begin(), col_data.end(), col_valid.begin()};

  // Construct Grouping keys table-view.
  const auto N_ELEMENTS{col_data.size()};
  const std::vector<size_type> grouping_key_vec(N_ELEMENTS, 0);
  fixed_width_column_wrapper<size_type> grouping_keys_col(
    grouping_key_vec.begin(), grouping_key_vec.end(), col_valid.begin());
  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{grouping_keys_col}};

  EXPECT_THROW(
    cudf::grouped_rolling_window(grouping_keys, input, 2, 2, -2, cudf::make_sum_aggregation()),
    cudf::logic_error);
}

TEST_F(GroupedRollingErrorTest, EmptyInput)
{
  cudf::test::fixed_width_column_wrapper<int32_t> empty_col{};
  std::unique_ptr<cudf::column> output;
  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{}};
  EXPECT_NO_THROW(output = cudf::grouped_rolling_window(
                    grouping_keys, empty_col, 2, 0, 2, cudf::make_sum_aggregation()));
  EXPECT_EQ(output->size(), 0);
}

// incorrect type/aggregation combo: sum of timestamps
TEST_F(GroupedRollingErrorTest, SumTimestampNotSupported)
{
  constexpr size_type size{10};
  fixed_width_column_wrapper<cudf::timestamp_D, cudf::timestamp_D::rep> input_D(
    thrust::make_counting_iterator(0), thrust::make_counting_iterator(size));
  fixed_width_column_wrapper<cudf::timestamp_s, cudf::timestamp_s::rep> input_s(
    thrust::make_counting_iterator(0), thrust::make_counting_iterator(size));
  fixed_width_column_wrapper<cudf::timestamp_ms, cudf::timestamp_ms::rep> input_ms(
    thrust::make_counting_iterator(0), thrust::make_counting_iterator(size));
  fixed_width_column_wrapper<cudf::timestamp_us, cudf::timestamp_us::rep> input_us(
    thrust::make_counting_iterator(0), thrust::make_counting_iterator(size));
  fixed_width_column_wrapper<cudf::timestamp_ns, cudf::timestamp_ns::rep> input_ns(
    thrust::make_counting_iterator(0), thrust::make_counting_iterator(size));

  // Construct table-view of grouping keys.
  std::vector<size_type> grouping_keys_vec(size, 0);  // `size` elements, each == 0.
  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{
    fixed_width_column_wrapper<size_type>(grouping_keys_vec.begin(), grouping_keys_vec.end())}};

  EXPECT_THROW(
    cudf::grouped_rolling_window(grouping_keys, input_D, 2, 2, 0, cudf::make_sum_aggregation()),
    cudf::logic_error);
  EXPECT_THROW(
    cudf::grouped_rolling_window(grouping_keys, input_s, 2, 2, 0, cudf::make_sum_aggregation()),
    cudf::logic_error);
  EXPECT_THROW(
    cudf::grouped_rolling_window(grouping_keys, input_ms, 2, 2, 0, cudf::make_sum_aggregation()),
    cudf::logic_error);
  EXPECT_THROW(
    cudf::grouped_rolling_window(grouping_keys, input_us, 2, 2, 0, cudf::make_sum_aggregation()),
    cudf::logic_error);
  EXPECT_THROW(
    cudf::grouped_rolling_window(grouping_keys, input_ns, 2, 2, 0, cudf::make_sum_aggregation()),
    cudf::logic_error);
}

TYPED_TEST_CASE(GroupedRollingTest, cudf::test::FixedWidthTypesWithoutFixedPoint);

TYPED_TEST(GroupedRollingTest, SimplePartitionedStaticWindowsWithGroupKeys)
{
  const auto col_data = {0, 10, 20, 30, 40, 50, 60, 70, 80, 90};
  const size_type DATA_SIZE{static_cast<size_type>(col_data.size())};
  const std::vector<bool> col_mask(DATA_SIZE, true);
  fixed_width_column_wrapper<TypeParam, int32_t> input(
    col_data.begin(), col_data.end(), col_mask.begin());

  // 2 grouping keys, with effectively 3 groups of at most 4 rows each:
  //   1. key_0 {0, 0, 0, ...0}
  //   2. key_1 {0, 0, 0, 0, 1, 1, 1, 1, 2, 2}
  std::vector<int64_t> key_0_vec(DATA_SIZE, 0);
  std::vector<int64_t> key_1_vec;
  int i{0};
  std::generate_n(
    std::back_inserter(key_1_vec), DATA_SIZE, [&i]() { return i++ / 4; });  // Groups of 4.
  const fixed_width_column_wrapper<TypeParam, int64_t> key_0(key_0_vec.begin(), key_0_vec.end());
  const fixed_width_column_wrapper<TypeParam, int64_t> key_1(key_1_vec.begin(), key_1_vec.end());
  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{key_0, key_1}};

  size_type preceding_window = 2;
  size_type following_window = 1;
  std::vector<size_type> expected_group_offsets{0, 4, 8, DATA_SIZE};

  this->run_test_col_agg(
    grouping_keys, input, expected_group_offsets, preceding_window, following_window, 1);
}

TYPED_TEST(GroupedRollingTest, SimplePartitionedStaticWindowWithNoGroupKeys)
{
  const auto col_data =
    cudf::test::make_type_param_vector<TypeParam>({0, 10, 20, 30, 40, 50, 60, 70, 80, 90});
  const size_type DATA_SIZE{static_cast<size_type>(col_data.size())};
  const std::vector<bool> col_mask(DATA_SIZE, true);
  fixed_width_column_wrapper<TypeParam> input(col_data.begin(), col_data.end(), col_mask.begin());

  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{}};

  size_type preceding_window = 2;
  size_type following_window = 1;
  std::vector<size_type> expected_group_offsets{0, DATA_SIZE};

  this->run_test_col_agg(
    grouping_keys, input, expected_group_offsets, preceding_window, following_window, 1);
}

// all rows are invalid
TYPED_TEST(GroupedRollingTest, AllInvalid)
{
  const auto col_data =
    cudf::test::make_type_param_vector<TypeParam>({0, 10, 20, 30, 40, 50, 60, 70, 80, 90});
  const size_type DATA_SIZE{static_cast<size_type>(col_data.size())};
  const std::vector<bool> col_mask(DATA_SIZE, false);
  fixed_width_column_wrapper<TypeParam> input(col_data.begin(), col_data.end(), col_mask.begin());

  // 2 grouping keys, with effectively 3 groups of at most 4 rows each:
  //   1. key_0 {0, 0, 0, ...0}
  //   2. key_1 {0, 0, 0, 0, 1, 1, 1, 1, 2, 2}
  std::vector<int64_t> key_0_vec(DATA_SIZE, 0);
  std::vector<int64_t> key_1_vec;
  int i{0};
  std::generate_n(
    std::back_inserter(key_1_vec), DATA_SIZE, [&i]() { return i++ / 4; });  // Groups of 4.
  const fixed_width_column_wrapper<TypeParam, int64_t> key_0(key_0_vec.begin(), key_0_vec.end());
  const fixed_width_column_wrapper<TypeParam, int64_t> key_1(key_1_vec.begin(), key_1_vec.end());
  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{key_0, key_1}};

  size_type preceding_window = 2;
  size_type following_window = 1;
  std::vector<size_type> expected_group_offsets{0, 4, 8, DATA_SIZE};

  this->run_test_col_agg(
    grouping_keys, input, expected_group_offsets, preceding_window, following_window, 1);
}

// window = following_window = 0
TYPED_TEST(GroupedRollingTest, ZeroWindow)
{
  const auto col_data = {0, 10, 20, 30, 40, 50, 60, 70, 80, 90};
  const size_type DATA_SIZE{static_cast<size_type>(col_data.size())};
  const std::vector<bool> col_mask(DATA_SIZE, true);
  fixed_width_column_wrapper<TypeParam, int32_t> input(
    col_data.begin(), col_data.end(), col_mask.begin());

  // 2 grouping keys, with effectively 3 groups of at most 4 rows each:
  //   1. key_0 {0, 0, 0, ...0}
  //   2. key_1 {0, 0, 0, 0, 1, 1, 1, 1, 2, 2}
  std::vector<int64_t> key_0_vec(DATA_SIZE, 0);
  std::vector<int64_t> key_1_vec;
  int i{0};
  std::generate_n(
    std::back_inserter(key_1_vec), DATA_SIZE, [&i]() { return i++ / 4; });  // Groups of 4.
  const fixed_width_column_wrapper<TypeParam, int64_t> key_0(key_0_vec.begin(), key_0_vec.end());
  const fixed_width_column_wrapper<TypeParam, int64_t> key_1(key_1_vec.begin(), key_1_vec.end());
  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{key_0, key_1}};

  size_type preceding_window = 0;
  size_type following_window = 0;
  std::vector<size_type> expected_group_offsets{0, 4, 8, DATA_SIZE};

  this->run_test_col_agg(
    grouping_keys, input, expected_group_offsets, preceding_window, following_window, 1);
}

// ------------- non-fixed-width types --------------------

using GroupedRollingTestStrings = GroupedRollingTest<cudf::string_view>;

TEST_F(GroupedRollingTestStrings, StringsUnsupportedOperators)
{
  cudf::test::strings_column_wrapper input{{"This", "is", "not", "a", "string", "type"},
                                           {1, 1, 1, 0, 1, 0}};

  const size_type DATA_SIZE{static_cast<cudf::column_view>(input).size()};
  const std::vector<size_type> key_col_vec(DATA_SIZE, 0);
  const cudf::table_view key_cols{std::vector<cudf::column_view>{
    fixed_width_column_wrapper<size_type>(key_col_vec.begin(), key_col_vec.end())}};

  EXPECT_THROW(cudf::grouped_rolling_window(key_cols, input, 2, 2, 0, cudf::make_sum_aggregation()),
               cudf::logic_error);
  EXPECT_THROW(
    cudf::grouped_rolling_window(key_cols, input, 2, 2, 0, cudf::make_mean_aggregation()),
    cudf::logic_error);
}

template <typename T>
class GroupedTimeRangeRollingTest : public cudf::test::BaseFixture {
 protected:
  // input as column_wrapper
  void run_test_col(cudf::table_view const& keys,
                    cudf::column_view const& timestamp_column,
                    cudf::order const& timestamp_order,
                    cudf::column_view const& input,
                    std::vector<size_type> const& expected_grouping,
                    size_type const& preceding_window_in_days,
                    size_type const& following_window_in_days,
                    size_type min_periods,
                    std::unique_ptr<cudf::aggregation> const& op)
  {
    std::unique_ptr<cudf::column> output;

    // wrap windows
    EXPECT_NO_THROW(output = cudf::grouped_time_range_rolling_window(keys,
                                                                     timestamp_column,
                                                                     timestamp_order,
                                                                     input,
                                                                     preceding_window_in_days,
                                                                     following_window_in_days,
                                                                     min_periods,
                                                                     op));

    auto reference = create_reference_output(op,
                                             timestamp_column,
                                             timestamp_order,
                                             input,
                                             expected_grouping,
                                             preceding_window_in_days,
                                             following_window_in_days,
                                             min_periods);

#ifndef NDEBUG
    std::cout << "input:\n";
    cudf::test::print(input, std::cout, ", ");
    std::cout << "\n";
    std::cout << "output:\n";
    cudf::test::print(*output, std::cout, ", ");
    std::cout << "\n";
    std::cout << "reference:\n";
    cudf::test::print(*reference, std::cout, ", ");
    std::cout << "\n";
    std::cout << "\n";
#endif

    CUDF_TEST_EXPECT_COLUMNS_EQUAL(*output, *reference);
  }

  void run_test_col_agg(cudf::table_view const& keys,
                        cudf::column_view const& timestamp_column,
                        cudf::order const& timestamp_order,
                        cudf::column_view const& input,
                        std::vector<size_type> const& expected_grouping,
                        size_type preceding_window_in_days,
                        size_type following_window_in_days,
                        size_type min_periods)
  {
    // Skip grouping-tests on bool8 keys. sort_helper does not support this.
    if (keys.num_columns() > 0 && cudf::is_boolean(keys.column(0).type())) { return; }

    // test all supported aggregators
    run_test_col(keys,
                 timestamp_column,
                 timestamp_order,
                 input,
                 expected_grouping,
                 preceding_window_in_days,
                 following_window_in_days,
                 min_periods,
                 cudf::make_min_aggregation());
    run_test_col(keys,
                 timestamp_column,
                 timestamp_order,
                 input,
                 expected_grouping,
                 preceding_window_in_days,
                 following_window_in_days,
                 min_periods,
                 cudf::make_count_aggregation());
    run_test_col(keys,
                 timestamp_column,
                 timestamp_order,
                 input,
                 expected_grouping,
                 preceding_window_in_days,
                 following_window_in_days,
                 min_periods,
                 cudf::make_count_aggregation(cudf::null_policy::INCLUDE));
    run_test_col(keys,
                 timestamp_column,
                 timestamp_order,
                 input,
                 expected_grouping,
                 preceding_window_in_days,
                 following_window_in_days,
                 min_periods,
                 cudf::make_max_aggregation());
    if (!cudf::is_timestamp(input.type())) {
      run_test_col(keys,
                   timestamp_column,
                   timestamp_order,
                   input,
                   expected_grouping,
                   preceding_window_in_days,
                   following_window_in_days,
                   min_periods,
                   cudf::make_sum_aggregation());
      run_test_col(keys,
                   timestamp_column,
                   timestamp_order,
                   input,
                   expected_grouping,
                   preceding_window_in_days,
                   following_window_in_days,
                   min_periods,
                   cudf::make_mean_aggregation());
    }
    run_test_col(keys,
                 timestamp_column,
                 timestamp_order,
                 input,
                 expected_grouping,
                 preceding_window_in_days,
                 following_window_in_days,
                 min_periods,
                 cudf::make_row_number_aggregation());
  }

 private:
  // use SFINAE to only instantiate for supported combinations

  // specialization for COUNT_VALID, COUNT_ALL
  template <bool include_nulls>
  std::unique_ptr<cudf::column> create_count_reference_output(
    cudf::column_view const& timestamp_column,
    cudf::order const& timestamp_order,
    cudf::column_view const& input,
    std::vector<size_type> const& group_offsets,
    size_type const& preceding_window_in_days,
    size_type const& following_window_in_days,
    size_type min_periods)
  {
    assert(timestamp_column.type().id() == cudf::type_id::TIMESTAMP_DAYS);  // Testing with DAYS.

    auto timestamp_vec = cudf::test::to_host<int32_t>(timestamp_column).first;

    size_type num_rows = input.size();
    thrust::host_vector<cudf::size_type> ref_data(num_rows);
    thrust::host_vector<bool> ref_valid(num_rows);

    // input data and mask
    std::vector<bitmask_type> in_valid = cudf::test::bitmask_to_host(input);
    bitmask_type* valid_mask           = in_valid.data();

    for (size_type i = 0; i < num_rows; i++) {
      // load sizes
      min_periods = std::max(min_periods, 1);  // at least one observation is required

      // compute bounds
      auto group_end_index   = std::upper_bound(group_offsets.begin(), group_offsets.end(), i);
      auto group_start_index = group_end_index - 1;

      size_type start_index = i;
      size_type end_index   = i;

      if (timestamp_order == cudf::order::ASCENDING) {
        while ((start_index - 1) >= *group_start_index &&
               timestamp_vec[start_index - 1] >= (timestamp_vec[i] - preceding_window_in_days)) {
          --start_index;
        }

        while ((end_index + 1) < *group_end_index &&
               timestamp_vec[end_index + 1] <= (timestamp_vec[i] + following_window_in_days)) {
          ++end_index;
        }
        ++end_index;  // One past the last.
      } else {
        while ((start_index - 1) >= *group_start_index &&
               timestamp_vec[start_index - 1] <= (timestamp_vec[i] + preceding_window_in_days)) {
          --start_index;
        }

        while ((end_index + 1) < *group_end_index &&
               timestamp_vec[end_index + 1] >= (timestamp_vec[i] - following_window_in_days)) {
          ++end_index;
        }
        ++end_index;  // One past the last.
      }

      // aggregate
      size_type count = 0;
      for (size_type j = start_index; j < end_index; j++) {
        if (include_nulls || !input.nullable() || cudf::bit_is_set(valid_mask, j)) count++;
      }

      ref_valid[i] = (count >= min_periods);
      if (ref_valid[i]) ref_data[i] = count;
    }

    fixed_width_column_wrapper<cudf::size_type> col(
      ref_data.begin(), ref_data.end(), ref_valid.begin());
    return col.release();
  }

  std::unique_ptr<cudf::column> create_row_number_reference_output(
    cudf::column_view const& timestamp_column,
    cudf::order const& timestamp_order,
    cudf::column_view const& input,
    std::vector<size_type> const& group_offsets,
    size_type const& preceding_window_in_days,
    size_type const& following_window_in_days,
    size_type min_periods)
  {
    assert(timestamp_column.type().id() == cudf::type_id::TIMESTAMP_DAYS);  // Testing with DAYS.

    auto timestamp_vec = cudf::test::to_host<int32_t>(timestamp_column).first;

    size_type num_rows = input.size();
    thrust::host_vector<cudf::size_type> ref_data(num_rows);
    thrust::host_vector<bool> ref_valid(num_rows);

    // input data and mask

    std::vector<bitmask_type> in_valid = cudf::test::bitmask_to_host(input);
    bitmask_type* valid_mask           = in_valid.data();

    for (size_type i = 0; i < num_rows; i++) {
      // load sizes
      min_periods = std::max(min_periods, 1);  // at least one observation is required

      // compute bounds
      auto group_end_index   = std::upper_bound(group_offsets.begin(), group_offsets.end(), i);
      auto group_start_index = group_end_index - 1;

      size_type start_index = i;
      size_type end_index   = i;

      if (timestamp_order == cudf::order::ASCENDING) {
        while ((start_index - 1) >= *group_start_index &&
               timestamp_vec[start_index - 1] >= (timestamp_vec[i] - preceding_window_in_days)) {
          --start_index;
        }

        while ((end_index + 1) < *group_end_index &&
               timestamp_vec[end_index + 1] <= (timestamp_vec[i] + following_window_in_days)) {
          ++end_index;
        }
        ++end_index;  // One past the last.
      } else {
        while ((start_index - 1) >= *group_start_index &&
               timestamp_vec[start_index - 1] <= (timestamp_vec[i] + preceding_window_in_days)) {
          --start_index;
        }

        while ((end_index + 1) < *group_end_index &&
               timestamp_vec[end_index + 1] >= (timestamp_vec[i] - following_window_in_days)) {
          ++end_index;
        }
        ++end_index;  // One past the last.
      }

      // aggregate
      size_type count{end_index - start_index};
      size_type row_number{i - start_index + 1};

      ref_valid[i] = (count >= min_periods);
      ref_data[i]  = row_number;
    }

    fixed_width_column_wrapper<cudf::size_type> col(
      ref_data.begin(), ref_data.end(), ref_valid.begin());
    return col.release();
  }

  template <typename agg_op,
            cudf::aggregation::Kind k,
            typename OutputType,
            bool is_mean,
            std::enable_if_t<cudf::detail::is_rolling_supported<T, agg_op, k>()>* = nullptr>
  std::unique_ptr<cudf::column> create_reference_output(cudf::column_view const& timestamp_column,
                                                        cudf::order const& timestamp_order,
                                                        cudf::column_view const& input,
                                                        std::vector<size_type> const& group_offsets,
                                                        size_type const& preceding_window_in_days,
                                                        size_type const& following_window_in_days,
                                                        size_type min_periods)
  {
    assert(timestamp_column.type().id() == cudf::type_id::TIMESTAMP_DAYS);  // Testing with DAYS.

    auto timestamp_vec = cudf::test::to_host<int32_t>(timestamp_column).first;

    size_type num_rows = input.size();
    thrust::host_vector<OutputType> ref_data(num_rows);
    thrust::host_vector<bool> ref_valid(num_rows);

    // input data and mask
    thrust::host_vector<T> in_col;
    std::vector<bitmask_type> in_valid;
    std::tie(in_col, in_valid) = cudf::test::to_host<T>(input);
    bitmask_type* valid_mask   = in_valid.data();

    agg_op op;
    for (size_type i = 0; i < num_rows; i++) {
      OutputType val = agg_op::template identity<OutputType>();

      // load sizes
      min_periods = std::max(min_periods, 1);  // at least one observation is required

      // compute bounds
      auto group_end_index   = std::upper_bound(group_offsets.begin(), group_offsets.end(), i);
      auto group_start_index = group_end_index - 1;

      size_type start_index = i;
      size_type end_index   = i;

      if (timestamp_order == cudf::order::ASCENDING) {
        while ((start_index - 1) >= *group_start_index &&
               timestamp_vec[start_index - 1] >= (timestamp_vec[i] - preceding_window_in_days)) {
          --start_index;
        }

        while ((end_index + 1) < *group_end_index &&
               timestamp_vec[end_index + 1] <= (timestamp_vec[i] + following_window_in_days)) {
          ++end_index;
        }
        ++end_index;  // One past the last.
      } else {
        while ((start_index - 1) >= *group_start_index &&
               timestamp_vec[start_index - 1] <= (timestamp_vec[i] + preceding_window_in_days)) {
          --start_index;
        }

        while ((end_index + 1) < *group_end_index &&
               timestamp_vec[end_index + 1] >= (timestamp_vec[i] - following_window_in_days)) {
          ++end_index;
        }
        ++end_index;
      }

      // aggregate
      size_type count = 0;
      for (size_type j = start_index; j < end_index; j++) {
        if (!input.nullable() || cudf::bit_is_set(valid_mask, j)) {
          val = op(static_cast<OutputType>(in_col[j]), val);
          count++;
        }
      }

      ref_valid[i] = (count >= min_periods);
      if (ref_valid[i]) {
        cudf::detail::rolling_store_output_functor<OutputType, is_mean>{}(ref_data[i], val, count);
      }
    }

    fixed_width_column_wrapper<OutputType> col(ref_data.begin(), ref_data.end(), ref_valid.begin());
    return col.release();
  }

  template <typename agg_op,
            cudf::aggregation::Kind k,
            typename OutputType,
            bool is_mean,
            std::enable_if_t<!cudf::detail::is_rolling_supported<T, agg_op, k>()>* = nullptr>
  std::unique_ptr<cudf::column> create_reference_output(cudf::column_view const& timestamp_column,
                                                        cudf::order const& timestamp_order,
                                                        cudf::column_view const& input,
                                                        std::vector<size_type> const& group_offsets,
                                                        size_type const& preceding_window_col,
                                                        size_type const& following_window_col,
                                                        size_type min_periods)
  {
    CUDF_FAIL("Unsupported combination of type and aggregation");
  }

  std::unique_ptr<cudf::column> create_reference_output(
    std::unique_ptr<cudf::aggregation> const& op,
    cudf::column_view const& timestamp_column,
    cudf::order const& timestamp_order,
    cudf::column_view const& input,
    std::vector<size_type> const& group_offsets,
    size_type const& preceding_window,
    size_type const& following_window,
    size_type min_periods)
  {
    // unroll aggregation types
    switch (op->kind) {
      case cudf::aggregation::SUM:
        return create_reference_output<cudf::DeviceSum,
                                       cudf::aggregation::SUM,
                                       cudf::detail::target_type_t<T, cudf::aggregation::SUM>,
                                       false>(timestamp_column,
                                              timestamp_order,
                                              input,
                                              group_offsets,
                                              preceding_window,
                                              following_window,
                                              min_periods);
      case cudf::aggregation::MIN:
        return create_reference_output<cudf::DeviceMin,
                                       cudf::aggregation::MIN,
                                       cudf::detail::target_type_t<T, cudf::aggregation::MIN>,
                                       false>(timestamp_column,
                                              timestamp_order,
                                              input,
                                              group_offsets,
                                              preceding_window,
                                              following_window,
                                              min_periods);
      case cudf::aggregation::MAX:
        return create_reference_output<cudf::DeviceMax,
                                       cudf::aggregation::MAX,
                                       cudf::detail::target_type_t<T, cudf::aggregation::MAX>,
                                       false>(timestamp_column,
                                              timestamp_order,
                                              input,
                                              group_offsets,
                                              preceding_window,
                                              following_window,
                                              min_periods);
      case cudf::aggregation::COUNT_VALID:
        return create_count_reference_output<false>(timestamp_column,
                                                    timestamp_order,
                                                    input,
                                                    group_offsets,
                                                    preceding_window,
                                                    following_window,
                                                    min_periods);
      case cudf::aggregation::COUNT_ALL:
        return create_count_reference_output<true>(timestamp_column,
                                                   timestamp_order,
                                                   input,
                                                   group_offsets,
                                                   preceding_window,
                                                   following_window,
                                                   min_periods);
      case cudf::aggregation::ROW_NUMBER:
        return create_row_number_reference_output(timestamp_column,
                                                  timestamp_order,
                                                  input,
                                                  group_offsets,
                                                  preceding_window,
                                                  following_window,
                                                  min_periods);
      case cudf::aggregation::MEAN:
        return create_reference_output<cudf::DeviceSum,
                                       cudf::aggregation::MEAN,
                                       cudf::detail::target_type_t<T, cudf::aggregation::MEAN>,
                                       true>(timestamp_column,
                                             timestamp_order,
                                             input,
                                             group_offsets,
                                             preceding_window,
                                             following_window,
                                             min_periods);
      default: return fixed_width_column_wrapper<T>({}).release();
    }
  }
};

TYPED_TEST_CASE(GroupedTimeRangeRollingTest, cudf::test::FixedWidthTypesWithoutFixedPoint);

TYPED_TEST(GroupedTimeRangeRollingTest,
           SimplePartitionedStaticWindowsWithGroupKeysAndTimeRangesAscending)
{
  const size_type DATA_SIZE{static_cast<size_type>(18)};
  const std::vector<int> col_data(DATA_SIZE, 1);
  const std::vector<bool> col_mask(DATA_SIZE, true);
  fixed_width_column_wrapper<TypeParam, int> input(
    col_data.begin(), col_data.end(), col_mask.begin());

  // 2 grouping keys, with effectively 3 groups of at most 6 rows each:
  //   1. key_0 {0, 0, 0, ...0}
  //   2. key_1 {0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2}
  std::vector<int64_t> key_0_vec(DATA_SIZE, 0);
  std::vector<int64_t> key_1_vec;
  int i{0};
  std::generate_n(
    std::back_inserter(key_1_vec), DATA_SIZE, [&i]() { return i++ / 6; });  // Groups of 6.
  const fixed_width_column_wrapper<TypeParam, int64_t> key_0(key_0_vec.begin(), key_0_vec.end());
  const fixed_width_column_wrapper<TypeParam, int64_t> key_1(key_1_vec.begin(), key_1_vec.end());
  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{key_0, key_1}};

  size_type preceding_window_in_days = 1;
  size_type following_window_in_days = 1;
  std::vector<size_type> expected_group_offsets{0, 6, 12, DATA_SIZE};

  // Timestamp column.
  std::vector<int32_t> timestamp_days_vec{0, 2, 3, 4, 5, 7, 0, 0, 1, 2, 3, 3, 0, 1, 2, 3, 3, 3};
  fixed_width_column_wrapper<cudf::timestamp_D, cudf::timestamp_D::rep> timestamp_days_ascending(
    timestamp_days_vec.begin(), timestamp_days_vec.end());

  this->run_test_col_agg(grouping_keys,
                         timestamp_days_ascending,
                         cudf::order::ASCENDING,
                         input,
                         expected_group_offsets,
                         preceding_window_in_days,
                         following_window_in_days,
                         1);
}

TYPED_TEST(GroupedTimeRangeRollingTest,
           SimplePartitionedStaticWindowsWithGroupKeysAndTimeRangesDescending)
{
  const size_type DATA_SIZE{static_cast<size_type>(18)};
  const std::vector<int> col_data(DATA_SIZE, 1);
  const std::vector<bool> col_mask(DATA_SIZE, true);
  fixed_width_column_wrapper<TypeParam, int> input(
    col_data.begin(), col_data.end(), col_mask.begin());

  // 2 grouping keys, with effectively 3 groups of at most 6 rows each:
  //   1. key_0 {0, 0, 0, ...0}
  //   2. key_1 {0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2}
  std::vector<int64_t> key_0_vec(DATA_SIZE, 0);
  std::vector<int64_t> key_1_vec;
  int i{0};
  std::generate_n(
    std::back_inserter(key_1_vec), DATA_SIZE, [&i]() { return i++ / 6; });  // Groups of 6.
  const fixed_width_column_wrapper<TypeParam, int64_t> key_0(key_0_vec.begin(), key_0_vec.end());
  const fixed_width_column_wrapper<TypeParam, int64_t> key_1(key_1_vec.begin(), key_1_vec.end());
  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{key_0, key_1}};

  size_type preceding_window_in_days = 1;
  size_type following_window_in_days = 2;
  std::vector<size_type> expected_group_offsets{0, 6, 12, DATA_SIZE};

  // Timestamp column.
  std::vector<int32_t> timestamp_days_vec{0, 2, 3, 4, 5, 7, 0, 0, 1, 2, 3, 3, 0, 1, 2, 3, 3, 3};
  fixed_width_column_wrapper<cudf::timestamp_D, cudf::timestamp_D::rep> timestamp_days_descending(
    timestamp_days_vec.rbegin(), timestamp_days_vec.rend());
  this->run_test_col_agg(grouping_keys,
                         timestamp_days_descending,
                         cudf::order::DESCENDING,
                         input,
                         expected_group_offsets,
                         preceding_window_in_days,
                         following_window_in_days,
                         1);
}

TYPED_TEST(GroupedTimeRangeRollingTest, SimplePartitionedStaticWindowsWithNoGroupingKeys)
{
  const size_type DATA_SIZE{static_cast<size_type>(6)};
  const std::vector<int> col_data(DATA_SIZE, 1);
  const std::vector<bool> col_mask(DATA_SIZE, true);
  fixed_width_column_wrapper<TypeParam, int> input(
    col_data.begin(), col_data.end(), col_mask.begin());

  const cudf::table_view grouping_keys{std::vector<cudf::column_view>{}};

  size_type preceding_window_in_days = 1;
  size_type following_window_in_days = 1;
  std::vector<size_type> expected_group_offsets{0, DATA_SIZE};

  // Timestamp column.
  std::vector<int32_t> timestamp_days_vec{0, 2, 3, 4, 5, 7};
  fixed_width_column_wrapper<cudf::timestamp_D, cudf::timestamp_D::rep> timestamp_days_ascending(
    timestamp_days_vec.begin(), timestamp_days_vec.end());

  this->run_test_col_agg(grouping_keys,
                         timestamp_days_ascending,
                         cudf::order::ASCENDING,
                         input,
                         expected_group_offsets,
                         preceding_window_in_days,
                         following_window_in_days,
                         1);
}

CUDF_TEST_PROGRAM_MAIN()