Add decrease interface for aggregation

dbms/src/AggregateFunctions/AggregateFunctionArray.h

@@ -67,13 +67,25 @@ class AggregateFunctionArray final : public IAggregateFunctionHelper<AggregateFu
     size_t alignOfData() const override { return nested_func->alignOfData(); }
 
     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);
+    }
+
+    template <bool is_add>
+    void addOrDecrease(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena * arena) const
     {
         const IColumn * nested[num_arguments];
 
         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 +94,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))
@@ -92,9 +104,16 @@ class AggregateFunctionArray final : public IAggregateFunctionHelper<AggregateFu
         }
 
         for (size_t i = begin; i < end; ++i)
-            nested_func->add(place, nested, i, arena);
+            if constexpr (is_add)
+                nested_func->add(place, nested, i, arena);
+            else
+                nested_func->decrease(place, nested, i, arena);
     }
 
+    void reset(AggregateDataPtr __restrict place) const override { nested_func->reset(place); }
+
+    void prepareWindow(AggregateDataPtr __restrict place) const override { nested_func->prepareWindow(place); }
+
     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)
@@ -67,6 +73,8 @@ class AggregateFunctionAvg final
             return std::make_shared<DataTypeFloat64>();
     }
 
+    void prepareWindow(AggregateDataPtr __restrict) const override {}
+
     void add(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena *) const override
     {
         if constexpr (IsDecimal<T>)
@@ -78,6 +86,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;
+    }
+
+    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,15 @@ 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 prepareWindow(AggregateDataPtr __restrict) const override {}
+
     void addBatchSinglePlace(
         size_t start_offset,
         size_t batch_size,
@@ -173,6 +184,15 @@ 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 prepareWindow(AggregateDataPtr __restrict) const override {}
+
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena *) const override
     {
         data(place).count += data(rhs).count;
@@ -234,6 +254,19 @@ 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;
+    }
+
+    void reset(AggregateDataPtr __restrict place) const override { data(place).reset(); }
+
+    void prepareWindow(AggregateDataPtr __restrict) const override {}
+
     void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs, Arena *) const override
     {
         data(place).count += data(rhs).count;

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,10 @@ class AggregateFunctionForEach final
         size_t old_size = state.dynamic_array_size;
         if (old_size < new_size)
         {
-            state.array_of_aggregate_datas = arena.realloc(
+            if unlikely (arena == nullptr)
+                throw Exception("Get nullptr 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);
@@ -149,13 +152,25 @@ class AggregateFunctionForEach final
     bool hasTrivialDestructor() const override { return nested_func->hasTrivialDestructor(); }
 
     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);
+    }
+
+    template <bool is_add>
+    void addOrDecrease(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena * arena) const
     {
         const IColumn * nested[num_arguments];
 
         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 +179,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,20 +188,45 @@ 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)
         {
-            nested_func->add(nested_state, nested, i, arena);
+            if constexpr (is_add)
+                nested_func->add(nested_state, nested, i, arena);
+            else
+                nested_func->decrease(nested_state, nested, i, arena);
+            nested_state += nested_size_of_data;
+        }
+    }
+
+    void reset(AggregateDataPtr __restrict place) const override
+    {
+        AggregateFunctionForEachData & state = ensureAggregateData(place, 0, nullptr);
+        char * nested_state = state.array_of_aggregate_datas;
+        for (size_t i = 0; i < state.dynamic_array_size; i++)
+        {
+            nested_func->reset(nested_state);
+            nested_state += nested_size_of_data;
+        }
+    }
+
+    void prepareWindow(AggregateDataPtr __restrict place) const override
+    {
+        AggregateFunctionForEachData & state = ensureAggregateData(place, 0, nullptr);
+        char * nested_state = state.array_of_aggregate_datas;
+        for (size_t i = 0; i < state.dynamic_array_size; i++)
+        {
+            nested_func->prepareWindow(nested_state);
             nested_state += nested_size_of_data;
         }
     }
 
     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 +260,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 +274,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/AggregateFunctionIf.h

@@ -76,6 +76,17 @@ class AggregateFunctionIf final : public IAggregateFunctionHelper<AggregateFunct
             nested_func->add(place, columns, row_num, arena);
     }
 
+    void decrease(AggregateDataPtr __restrict place, const IColumn ** columns, size_t row_num, Arena * arena)
+        const override
+    {
+        if (static_cast<const ColumnUInt8 &>(*columns[num_arguments - 1]).getData()[row_num])
+            nested_func->decrease(place, columns, row_num, arena);
+    }
+
+    void reset(AggregateDataPtr __restrict place) const override { nested_func->reset(place); }
+
+    void prepareWindow(AggregateDataPtr __restrict place) const override { nested_func->prepareWindow(place); }
+
     void addBatch(
         size_t start_offset,
         size_t batch_size,

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