Antalya 26.8: Backport of #117009 Speed up aggregation batches with non-null states

src/AggregateFunctions/Combinators/AggregateFunctionIf.h

@@ -162,6 +162,21 @@ class AggregateFunctionIf final : public IAggregateFunctionHelper<AggregateFunct
         nested_func->addBatch(row_begin, row_end, places, place_offset, columns, arena, num_arguments - 1);
     }
 
+    void addBatchWithNonNullPlaces(
+        size_t row_begin,
+        size_t row_end,
+        AggregateDataPtr * __restrict places,
+        size_t place_offset,
+        const IColumn ** columns,
+        Arena * arena,
+        ssize_t) const override
+    {
+        if (only_null_condition)
+            return;
+        nested_func->addBatchWithNonNullPlaces(
+            row_begin, row_end, places, place_offset, columns, arena, num_arguments - 1);
+    }
+
     void addBatchSinglePlace(
         size_t row_begin,
         size_t row_end,

src/AggregateFunctions/Combinators/AggregateFunctionOrFill.h

@@ -164,6 +164,33 @@ class AggregateFunctionOrFill final : public IAggregateFunctionHelper<AggregateF
         }
     }
 
+    void addBatchWithNonNullPlaces( /// NOLINT
+        size_t row_begin,
+        size_t row_end,
+        AggregateDataPtr * places,
+        size_t place_offset,
+        const IColumn ** columns,
+        Arena * arena,
+        ssize_t if_argument_pos = -1) const override
+    {
+        if (if_argument_pos >= 0)
+        {
+            const auto & flags = assert_cast<const ColumnUInt8 &>(*columns[if_argument_pos]).getData();
+            for (size_t i = row_begin; i < row_end; ++i)
+            {
+                if (flags[i])
+                    add(places[i] + place_offset, columns, i, arena);
+            }
+        }
+        else
+        {
+            nested_function->addBatchWithNonNullPlaces(
+                row_begin, row_end, places, place_offset, columns, arena, if_argument_pos);
+            for (size_t i = row_begin; i < row_end; ++i)
+                (places[i] + place_offset)[size_of_data] = 1;
+        }
+    }
+
     void addBatchSinglePlace( /// NOLINT
         size_t row_begin,
         size_t row_end,

src/AggregateFunctions/Combinators/AggregateFunctionTuple.cpp

@@ -185,7 +185,7 @@ void AggregateFunctionTuple::add(AggregateDataPtr __restrict place, const IColum
     }
 }
 
-template <bool has_null_map, typename GetPlace>
+template <bool has_null_map, bool all_places_are_non_null, typename GetPlace>
 void AggregateFunctionTuple::addBatchImpl(
     size_t row_begin,
     size_t row_end,
@@ -301,7 +301,9 @@ void AggregateFunctionTuple::addBatchImpl(
                 if (null_map[i])
                     continue;
             }
-            if (AggregateDataPtr place = get_place(i))
+            if constexpr (all_places_are_non_null)
+                add_row(get_place(i), i);
+            else if (AggregateDataPtr place = get_place(i))
                 add_row(place, i);
         }
     }
@@ -314,7 +316,9 @@ void AggregateFunctionTuple::addBatchImpl(
                 if (null_map[i])
                     continue;
             }
-            if (AggregateDataPtr place = get_place(i))
+            if constexpr (all_places_are_non_null)
+                add_row(get_place(i), i);
+            else if (AggregateDataPtr place = get_place(i))
                 add_row(place, i);
         }
     }
@@ -329,10 +333,23 @@ void AggregateFunctionTuple::addBatch( /// NOLINT
     Arena * arena,
     ssize_t if_argument_pos) const
 {
-    addBatchImpl<false>(row_begin, row_end, columns, nullptr, if_argument_pos, arena,
+    addBatchImpl<false, false>(row_begin, row_end, columns, nullptr, if_argument_pos, arena,
         [&](size_t i) { return places[i] ? places[i] + place_offset : nullptr; });
 }
 
+void AggregateFunctionTuple::addBatchWithNonNullPlaces( /// NOLINT
+    size_t row_begin,
+    size_t row_end,
+    AggregateDataPtr * places,
+    size_t place_offset,
+    const IColumn ** columns,
+    Arena * arena,
+    ssize_t if_argument_pos) const
+{
+    addBatchImpl<false, true>(row_begin, row_end, columns, nullptr, if_argument_pos, arena,
+        [&](size_t i) { return places[i] + place_offset; });
+}
+
 void AggregateFunctionTuple::addBatchSinglePlace( /// NOLINT
     size_t row_begin,
     size_t row_end,
@@ -341,7 +358,7 @@ void AggregateFunctionTuple::addBatchSinglePlace( /// NOLINT
     Arena * arena,
     ssize_t if_argument_pos) const
 {
-    addBatchImpl<false>(row_begin, row_end, columns, nullptr, if_argument_pos, arena,
+    addBatchImpl<false, true>(row_begin, row_end, columns, nullptr, if_argument_pos, arena,
         [&](size_t) { return place; });
 }
 
@@ -356,7 +373,7 @@ void AggregateFunctionTuple::addBatchSinglePlaceNotNull( /// NOLINT
 {
     /// Reached from `AggregateFunctionNullUnary` for `Nullable(Tuple(...))` inputs: rows whose tuple
     /// is NULL are skipped via the null map.
-    addBatchImpl<true>(row_begin, row_end, columns, null_map, if_argument_pos, arena,
+    addBatchImpl<true, true>(row_begin, row_end, columns, null_map, if_argument_pos, arena,
         [&](size_t) { return place; });
 }
 

src/AggregateFunctions/Combinators/AggregateFunctionTuple.h

@@ -89,6 +89,14 @@ class AggregateFunctionTuple final : public IAggregateFunctionHelper<AggregateFu
         const IColumn ** columns,
         Arena * arena,
         ssize_t if_argument_pos = -1) const override;
+    void addBatchWithNonNullPlaces( /// NOLINT
+        size_t row_begin,
+        size_t row_end,
+        AggregateDataPtr * places,
+        size_t place_offset,
+        const IColumn ** columns,
+        Arena * arena,
+        ssize_t if_argument_pos = -1) const override;
     void addBatchSinglePlace( /// NOLINT
         size_t row_begin,
         size_t row_end,
@@ -214,7 +222,7 @@ class AggregateFunctionTuple final : public IAggregateFunctionHelper<AggregateFu
     /// Shared implementation of the batch add overrides. Hoists the per-element column pointers, so
     /// no per-row unwrapping work remains in the row loop.
     /// `get_place` returns the aggregation state for a row, or nullptr when the row has none.
-    template <bool has_null_map, typename GetPlace>
+    template <bool has_null_map, bool all_places_are_non_null, typename GetPlace>
     void addBatchImpl(
         size_t row_begin,
         size_t row_end,

src/AggregateFunctions/IAggregateFunction.h

@@ -340,6 +340,21 @@ class IAggregateFunction : public std::enable_shared_from_this<IAggregateFunctio
         Arena * arena,
         ssize_t if_argument_pos = -1) const = 0;
 
+    /** A version of `addBatch` for callers that guarantee that every entry in `places` is non-null.
+      * Implementations that don't benefit from this guarantee can use the default implementation.
+      */
+    virtual void addBatchWithNonNullPlaces( /// NOLINT
+        size_t row_begin,
+        size_t row_end,
+        AggregateDataPtr * places,
+        size_t place_offset,
+        const IColumn ** columns,
+        Arena * arena,
+        ssize_t if_argument_pos = -1) const
+    {
+        addBatch(row_begin, row_end, places, place_offset, columns, arena, if_argument_pos);
+    }
+
     /// The version of "addBatch", that handle sparse columns as arguments.
     virtual void addBatchSparse(
         size_t row_begin,
@@ -571,6 +586,31 @@ class IAggregateFunctionHelper : public IAggregateFunction
         }
     }
 
+    void addBatchWithNonNullPlaces( /// NOLINT
+        size_t row_begin,
+        size_t row_end,
+        AggregateDataPtr * places,
+        size_t place_offset,
+        const IColumn ** columns,
+        Arena * arena,
+        ssize_t if_argument_pos = -1) const override
+    {
+        if (if_argument_pos >= 0)
+        {
+            const auto & flags = assert_cast<const ColumnUInt8 &>(*columns[if_argument_pos]).getData();
+            for (size_t i = row_begin; i < row_end; ++i)
+            {
+                if (flags[i])
+                    static_cast<const Derived *>(this)->add(places[i] + place_offset, columns, i, arena);
+            }
+        }
+        else
+        {
+            for (size_t i = row_begin; i < row_end; ++i)
+                static_cast<const Derived *>(this)->add(places[i] + place_offset, columns, i, arena);
+        }
+    }
+
     void serializeBatch(const PaddedPODArray<AggregateDataPtr> & data, size_t start, size_t size, WriteBuffer & buf, std::optional<size_t> version) const final // NOLINT
     {
         for (size_t i = start; i < size; ++i)

src/AggregateFunctions/TimeSeries/AggregateFunctionTimeseriesBase.h

@@ -199,6 +199,22 @@ class AggregateFunctionTimeseriesBase :
         }
     }
 
+    void addBatchWithNonNullPlaces(
+        size_t row_begin,
+        size_t row_end,
+        AggregateDataPtr * places,
+        size_t place_offset,
+        const IColumn ** columns,
+        Arena * arena,
+        ssize_t if_argument_pos) const override
+    {
+        if (array_arguments)
+            Base::addBatchWithNonNullPlaces(
+                row_begin, row_end, places, place_offset, columns, arena, if_argument_pos);
+        else
+            addBatch(row_begin, row_end, places, place_offset, columns, arena, if_argument_pos);
+    }
+
     void addBatchSinglePlace(
         size_t row_begin,
         size_t row_end,

src/Interpreters/AdaptiveAggregationKernels.cpp

@@ -667,6 +667,7 @@ void NO_INLINE Aggregator::executeFrozenImpl(
             /*key_start=*/row_begin,
             /*has_only_one_value_since_last_reset=*/false,
             /*all_keys_are_const=*/false,
+            /*all_places_are_non_null=*/false,
             /*use_compiled_functions=*/false);
 }
 
@@ -1447,6 +1448,7 @@ void NO_INLINE Aggregator::drainAdaptiveBucketImpl(
         /*key_start=*/slice_begin,
         /*has_only_one_value_since_last_reset=*/false,
         /*all_keys_are_const=*/false,
+        /*all_places_are_non_null=*/true,
         use_compiled_functions);
 }
 

src/Interpreters/Aggregator.cpp

@@ -1761,6 +1761,8 @@ void NO_INLINE Aggregator::executeImplBatch(
     state.resetCache();
 
     [[maybe_unused]] std::vector<DestroyedState> destroyed_states;
+    /// Assign at the branch tails so `no_more_keys` is not live across either loop.
+    bool all_places_are_non_null = false;
 
     /// For all rows.
     if (!no_more_keys)
@@ -1877,6 +1879,8 @@ void NO_INLINE Aggregator::executeImplBatch(
                 }
             }
         }
+
+        all_places_are_non_null = !top_k;
     }
     else
     {
@@ -1891,6 +1895,8 @@ void NO_INLINE Aggregator::executeImplBatch(
                 aggregate_data = overflow_row;
             places[i] = aggregate_data;
         }
+
+        all_places_are_non_null = false;
     }
 
     if constexpr (top_k)
@@ -1913,6 +1919,7 @@ void NO_INLINE Aggregator::executeImplBatch(
             key_start,
             has_only_one_value,
             all_keys_are_const,
+            all_places_are_non_null,
             use_jit);
 }
 
@@ -1925,6 +1932,7 @@ void Aggregator::executeAggregateInstructions(
     size_t key_start,
     bool has_only_one_value_since_last_reset,
     bool all_keys_are_const,
+    bool all_places_are_non_null,
     bool use_compiled_functions [[maybe_unused]]) const
 {
 #if USE_EMBEDDED_COMPILER
@@ -1977,7 +1985,7 @@ void Aggregator::executeAggregateInstructions(
         }
         else
         {
-            addBatch(row_begin, row_end, inst, places, aggregates_pool);
+            addBatch(row_begin, row_end, inst, places, aggregates_pool, all_places_are_non_null);
         }
     }
 
@@ -2041,7 +2049,8 @@ void Aggregator::addBatch(
     size_t row_begin, size_t row_end,
     const AggregateFunctionInstruction * inst,
     AggregateDataPtr * places,
-    Arena * arena)
+    Arena * arena,
+    bool all_places_are_non_null)
 {
     if (inst->offsets)
         inst->batch_that->addBatchArray(
@@ -2056,6 +2065,12 @@ void Aggregator::addBatch(
             inst->state_offset,
             inst->batch_arguments,
             arena);
+    else if (all_places_are_non_null)
+        inst->batch_that->addBatchWithNonNullPlaces(
+            row_begin, row_end, places,
+            inst->state_offset,
+            inst->batch_arguments,
+            arena);
     else
         inst->batch_that->addBatch(
             row_begin, row_end, places,

src/Interpreters/Aggregator.h

@@ -928,6 +928,7 @@ class Aggregator final
         size_t key_start,
         bool has_only_one_value_since_last_reset,
         bool all_keys_are_const,
+        bool all_places_are_non_null,
         bool use_compiled_functions) const;
 
     /// For case when there are no keys (all aggregate into one row).
@@ -1314,7 +1315,8 @@ class Aggregator final
         size_t row_begin, size_t row_end,
         const AggregateFunctionInstruction * inst,
         AggregateDataPtr * places,
-        Arena * arena);
+        Arena * arena,
+        bool all_places_are_non_null);
 
     static void addBatchSinglePlace(
         size_t row_begin, size_t row_end,