Add decrease interface for aggregation

dbms/src/AggregateFunctions/AggregateFunctionArray.h

@@ -73,7 +73,7 @@ class AggregateFunctionArray final : public IAggregateFunctionHelper<AggregateFu
         for (size_t i = 0; i < num_arguments; ++i)
             nested[i] = &static_cast<const ColumnArray &>(*columns[i]).getData();
 
-        const ColumnArray & first_array_column = static_cast<const ColumnArray &>(*columns[0]);
+        const auto & first_array_column = static_cast<const ColumnArray &>(*columns[0]);
         const IColumn::Offsets & offsets = first_array_column.getOffsets();
 
         size_t begin = row_num == 0 ? 0 : offsets[row_num - 1];
@@ -82,7 +82,7 @@ class AggregateFunctionArray final : public IAggregateFunctionHelper<AggregateFu
         /// Sanity check. NOTE We can implement specialization for a case with single argument, if the check will hurt performance.
         for (size_t i = 1; i < num_arguments; ++i)
         {
-            const ColumnArray & ith_column = static_cast<const ColumnArray &>(*columns[i]);
+            const auto & ith_column = static_cast<const ColumnArray &>(*columns[i]);
             const IColumn::Offsets & ith_offsets = ith_column.getOffsets();
 
             if (ith_offsets[row_num] != end || (row_num != 0 && ith_offsets[row_num - 1] != begin))
@@ -95,6 +95,7 @@ class AggregateFunctionArray final : public IAggregateFunctionHelper<AggregateFu
             nested_func->add(place, nested, i, arena);
     }
 
+
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena * arena) const override
     {
         nested_func->merge(place, rhs, arena);

dbms/src/AggregateFunctions/AggregateFunctionAvg.h

@@ -29,6 +29,12 @@ struct AggregateFunctionAvgData
     T sum;
     UInt64 count;
 
+    void reset()
+    {
+        sum = T(0);
+        count = 0;
+    }
+
     AggregateFunctionAvgData()
         : sum(0)
         , count(0)
@@ -78,6 +84,19 @@ class AggregateFunctionAvg final
         ++this->data(place).count;
     }
 
+    void decrease(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena *) const override
+    {
+        if constexpr (IsDecimal<T>)
+            this->data(place).sum -= static_cast<const ColumnDecimal<T> &>(*columns[0]).getData()[row_num];
+        else
+            this->data(place).sum -= static_cast<const ColumnVector<T> &>(*columns[0]).getData()[row_num];
+
+        --this->data(place).count;
+        assert(this->data(place).count >= 0);
+    }
+
+    void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }
+
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena *) const override
     {
         this->data(place).sum += this->data(rhs).sum;

dbms/src/AggregateFunctions/AggregateFunctionCount.h

@@ -29,6 +29,8 @@ namespace DB
 struct AggregateFunctionCountData
 {
     UInt64 count = 0;
+
+    inline void reset() noexcept { count = 0; }
 };
 
 namespace ErrorCodes
@@ -52,6 +54,13 @@ class AggregateFunctionCount final
         ++data(place).count;
     }
 
+    void decrease(AggregateDataPtr __restrict place, const IColumn **, size_t, Arena *) const override
+    {
+        --data(place).count;
+    }
+
+    void reset(AggregateDataPtr __restrict place) const override { data(place).reset(); }
+
     void addBatchSinglePlace(
         size_t start_offset,
         size_t batch_size,
@@ -173,6 +182,13 @@ class AggregateFunctionCountNotNullUnary final
         data(place).count += !static_cast<const ColumnNullable &>(*columns[0]).isNullAt(row_num);
     }
 
+    void decrease(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena *) const override
+    {
+        data(place).count -= !static_cast<const ColumnNullable &>(*columns[0]).isNullAt(row_num);
+    }
+
+    void reset(AggregateDataPtr __restrict place) const override { data(place).reset(); }
+
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena *) const override
     {
         data(place).count += data(rhs).count;
@@ -234,6 +250,18 @@ class AggregateFunctionCountNotNullVariadic final
         ++data(place).count;
     }
 
+    void decrease(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena *) const override
+    {
+        for (size_t i = 0; i < number_of_arguments; ++i)
+            if (is_nullable[i] && static_cast<const ColumnNullable &>(*columns[i]).isNullAt(row_num))
+                return;
+
+        --data(place).count;
+        assert(data(place).count >= 0);
+    }
+
+    void reset(AggregateDataPtr __restrict place) const override { data(place).reset(); }
+
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena *) const override
     {
         data(place).count += data(rhs).count;

dbms/src/AggregateFunctions/AggregateFunctionFactory.cpp

@@ -107,6 +107,26 @@ AggregateFunctionPtr AggregateFunctionFactory::get(
     return res;
 }
 
+AggregateFunctionPtr AggregateFunctionFactory::getForWindow(
+    const Context & context,
+    const String & name,
+    const DataTypes & argument_types,
+    const Array & parameters,
+    int recursion_level) const
+{
+    AggregateFunctionCombinatorPtr combinator
+        = AggregateFunctionCombinatorFactory::instance().tryFindSuffix("NullForWindow");
+    if (!combinator)
+        throw Exception(
+            "Logical error: cannot find aggregate function combinator to apply a function to Nullable for window "
+            "arguments.",
+            ErrorCodes::LOGICAL_ERROR);
+
+    DataTypes nested_types = combinator->transformArguments(argument_types);
+    AggregateFunctionPtr nested_function = getImpl(context, name, nested_types, parameters, recursion_level);
+    return combinator->transformAggregateFunction(nested_function, argument_types, parameters);
+}
+
 AggregateFunctionPtr AggregateFunctionFactory::getImpl(
     const Context & context,
     const String & name,

dbms/src/AggregateFunctions/AggregateFunctionFactory.h

@@ -67,6 +67,14 @@ class AggregateFunctionFactory final : public ext::Singleton<AggregateFunctionFa
         int recursion_level = 0,
         bool empty_input_as_null = false) const;
 
+    /// Throws an exception if not found.
+    AggregateFunctionPtr getForWindow(
+        const Context & context,
+        const String & name,
+        const DataTypes & argument_types,
+        const Array & parameters = {},
+        int recursion_level = 0) const;
+
     /// Returns nullptr if not found.
     AggregateFunctionPtr tryGet(
         const Context & context,

dbms/src/AggregateFunctions/AggregateFunctionForEach.h

@@ -66,7 +66,7 @@ class AggregateFunctionForEach final
     AggregateFunctionForEachData & ensureAggregateData(
         AggregateDataPtr __restrict place,
         size_t new_size,
-        Arena & arena) const
+        Arena * arena) const
     {
         AggregateFunctionForEachData & state = data(place);
 
@@ -75,7 +75,9 @@ class AggregateFunctionForEach final
         size_t old_size = state.dynamic_array_size;
         if (old_size < new_size)
         {
-            state.array_of_aggregate_datas = arena.realloc(
+            RUNTIME_CHECK_MSG(arena, "got null arena ptr in ensureAggregateData");
+
+            state.array_of_aggregate_datas = arena->realloc(
                 state.array_of_aggregate_datas,
                 old_size * nested_size_of_data,
                 new_size * nested_size_of_data);
@@ -155,7 +157,7 @@ class AggregateFunctionForEach final
         for (size_t i = 0; i < num_arguments; ++i)
             nested[i] = &static_cast<const ColumnArray &>(*columns[i]).getData();
 
-        const ColumnArray & first_array_column = static_cast<const ColumnArray &>(*columns[0]);
+        const auto & first_array_column = static_cast<const ColumnArray &>(*columns[0]);
         const IColumn::Offsets & offsets = first_array_column.getOffsets();
 
         size_t begin = row_num == 0 ? 0 : offsets[row_num - 1];
@@ -164,7 +166,7 @@ class AggregateFunctionForEach final
         /// Sanity check. NOTE We can implement specialization for a case with single argument, if the check will hurt performance.
         for (size_t i = 1; i < num_arguments; ++i)
         {
-            const ColumnArray & ith_column = static_cast<const ColumnArray &>(*columns[i]);
+            const auto & ith_column = static_cast<const ColumnArray &>(*columns[i]);
             const IColumn::Offsets & ith_offsets = ith_column.getOffsets();
 
             if (ith_offsets[row_num] != end || (row_num != 0 && ith_offsets[row_num - 1] != begin))
@@ -173,7 +175,7 @@ class AggregateFunctionForEach final
                     ErrorCodes::SIZES_OF_ARRAYS_DOESNT_MATCH);
         }
 
-        AggregateFunctionForEachData & state = ensureAggregateData(place, end - begin, *arena);
+        AggregateFunctionForEachData & state = ensureAggregateData(place, end - begin, arena);
 
         char * nested_state = state.array_of_aggregate_datas;
         for (size_t i = begin; i < end; ++i)
@@ -186,7 +188,7 @@ class AggregateFunctionForEach final
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena * arena) const override
     {
         const AggregateFunctionForEachData & rhs_state = data(rhs);
-        AggregateFunctionForEachData & state = ensureAggregateData(place, rhs_state.dynamic_array_size, *arena);
+        AggregateFunctionForEachData & state = ensureAggregateData(place, rhs_state.dynamic_array_size, arena);
 
         const char * rhs_nested_state = rhs_state.array_of_aggregate_datas;
         char * nested_state = state.array_of_aggregate_datas;
@@ -220,7 +222,7 @@ class AggregateFunctionForEach final
         size_t new_size = 0;
         readBinary(new_size, buf);
 
-        ensureAggregateData(place, new_size, *arena);
+        ensureAggregateData(place, new_size, arena);
 
         char * nested_state = state.array_of_aggregate_datas;
         for (size_t i = 0; i < new_size; ++i)
@@ -234,7 +236,7 @@ class AggregateFunctionForEach final
     {
         const AggregateFunctionForEachData & state = data(place);
 
-        ColumnArray & arr_to = static_cast<ColumnArray &>(to);
+        auto & arr_to = static_cast<ColumnArray &>(to);
         ColumnArray::Offsets & offsets_to = arr_to.getOffsets();
         IColumn & elems_to = arr_to.getData();
 

dbms/src/AggregateFunctions/AggregateFunctionGroupConcat.h

@@ -116,32 +116,36 @@ class AggregateFunctionGroupConcat final
 
     DataTypePtr getReturnType() const override { return result_is_nullable ? makeNullable(ret_type) : ret_type; }
 
-    /// reject nulls before add() of nested agg
-    void add(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena * arena) const override
+    template <bool is_add>
+    void addOrDecrease(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena * arena) const
     {
         if constexpr (only_one_column)
         {
             if (is_nullable[0])
             {
-                const ColumnNullable * column = static_cast<const ColumnNullable *>(columns[0]);
+                const auto * column = static_cast<const ColumnNullable *>(columns[0]);
                 if (!column->isNullAt(row_num))
                 {
                     this->setFlag(place);
                     const IColumn * nested_column = &column->getNestedColumn();
-                    this->nested_function->add(this->nestedPlace(place), &nested_column, row_num, arena);
+
+                    if constexpr (is_add)
+                        this->nested_function->add(this->nestedPlace(place), &nested_column, row_num, arena);
+                    else
+                        this->nested_function->decrease(this->nestedPlace(place), &nested_column, row_num, arena);
                 }
                 return;
             }
         }
         else
         {
             /// remove the row with null, except for sort columns
-            const ColumnTuple & tuple = static_cast<const ColumnTuple &>(*columns[0]);
+            const auto & tuple = static_cast<const ColumnTuple &>(*columns[0]);
             for (size_t i = 0; i < number_of_concat_items; ++i)
             {
                 if (is_nullable[i])
                 {
-                    const ColumnNullable & nullable_col = static_cast<const ColumnNullable &>(tuple.getColumn(i));
+                    const auto & nullable_col = static_cast<const ColumnNullable &>(tuple.getColumn(i));
                     if (nullable_col.isNullAt(row_num))
                     {
                         /// If at least one column has a null value in the current row,
@@ -152,7 +156,21 @@ class AggregateFunctionGroupConcat final
             }
         }
         this->setFlag(place);
-        this->nested_function->add(this->nestedPlace(place), columns, row_num, arena);
+        if constexpr (is_add)
+            this->nested_function->add(this->nestedPlace(place), columns, row_num, arena);
+        else
+            this->nested_function->decrease(this->nestedPlace(place), columns, row_num, arena);
+    }
+
+    void add(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena * arena) const override
+    {
+        addOrDecrease<true>(place, columns, row_num, arena);
+    }
+
+    void decrease(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena * arena)
+        const override
+    {
+        addOrDecrease<false>(place, columns, row_num, arena);
     }
 
     void insertResultInto(ConstAggregateDataPtr __restrict place, IColumn & to, Arena * arena) const override

dbms/src/AggregateFunctions/AggregateFunctionMaxIntersections.h

@@ -103,10 +103,10 @@ class AggregateFunctionIntersectionsMax final
         PointType right = static_cast<const ColumnVector<PointType> &>(*columns[1]).getData()[row_num];
 
         if (!isNaN(left))
-            this->data(place).value.push_back(std::make_pair(left, Int64(1)), arena);
+            this->data(place).value.push_back(std::make_pair(left, static_cast<Int64>(1)), arena);
 
         if (!isNaN(right))
-            this->data(place).value.push_back(std::make_pair(right, Int64(-1)), arena);
+            this->data(place).value.push_back(std::make_pair(right, static_cast<Int64>(-1)), arena);
     }
 
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena * arena) const override

dbms/src/AggregateFunctions/AggregateFunctionMerge.h

@@ -37,7 +37,7 @@ class AggregateFunctionMerge final : public IAggregateFunctionHelper<AggregateFu
     AggregateFunctionMerge(const AggregateFunctionPtr & nested_, const IDataType & argument)
         : nested_func(nested_)
     {
-        const DataTypeAggregateFunction * data_type = typeid_cast<const DataTypeAggregateFunction *>(&argument);
+        const auto * data_type = typeid_cast<const DataTypeAggregateFunction *>(&argument);
 
         if (!data_type || data_type->getFunctionName() != nested_func->getName())
             throw Exception(