Backport ClickHouse#87049 to 25.8: Fix compatibility of some aggregate function states with String argument

Author: robot-clickhouse

src/AggregateFunctions/AggregateFunctionGroupArray.cpp

@@ -456,15 +456,15 @@ struct GroupArrayNodeGeneral : public GroupArrayNodeBase<GroupArrayNodeGeneral>
     static Node * allocate(const IColumn & column, size_t row_num, Arena * arena)
     {
         const char * begin = arena->alignedAlloc(sizeof(Node), alignof(Node));
-        StringRef value = column.serializeValueIntoArena(row_num, *arena, begin);
+        StringRef value = column.serializeAggregationStateValueIntoArena(row_num, *arena, begin);
 
         Node * node = reinterpret_cast<Node *>(const_cast<char *>(begin));
         node->size = value.size;
 
         return node;
     }
 
-    void insertInto(IColumn & column) { std::ignore = column.deserializeAndInsertFromArena(data()); }
+    void insertInto(IColumn & column) { std::ignore = column.deserializeAndInsertAggregationStateValueFromArena(data()); }
 };
 
 template <typename Node, bool has_sampler>

src/AggregateFunctions/AggregateFunctionGroupArrayIntersect.cpp

@@ -235,7 +235,7 @@ class AggregateFunctionGroupArrayIntersectGeneric final
                 else
                 {
                     const char * begin = nullptr;
-                    StringRef serialized = data_column->serializeValueIntoArena(offset + i, *arena, begin);
+                    StringRef serialized = data_column->serializeAggregationStateValueIntoArena(offset + i, *arena, begin);
                     chassert(serialized.data != nullptr);
                     set.emplace(SerializedKeyHolder{serialized, *arena}, it, inserted);
                 }
@@ -255,7 +255,7 @@ class AggregateFunctionGroupArrayIntersectGeneric final
                 else
                 {
                     const char * begin = nullptr;
-                    StringRef serialized = data_column->serializeValueIntoArena(offset + i, *arena, begin);
+                    StringRef serialized = data_column->serializeAggregationStateValueIntoArena(offset + i, *arena, begin);
                     chassert(serialized.data != nullptr);
                     it = set.find(serialized);
 
@@ -344,7 +344,7 @@ class AggregateFunctionGroupArrayIntersectGeneric final
             if constexpr (is_plain_column)
                 data_to.insertData(elem.getValue().data, elem.getValue().size);
             else
-                std::ignore = data_to.deserializeAndInsertFromArena(elem.getValue().data);
+                std::ignore = data_to.deserializeAndInsertAggregationStateValueFromArena(elem.getValue().data);
         }
     }
 };

src/AggregateFunctions/AggregateFunctionTopK.cpp

@@ -310,7 +310,7 @@ class AggregateFunctionTopKGeneric final
         else
         {
             const char * begin = nullptr;
-            StringRef str_serialized = columns[0]->serializeValueIntoArena(row_num, *arena, begin);
+            StringRef str_serialized = columns[0]->serializeAggregationStateValueIntoArena(row_num, *arena, begin);
             if constexpr (is_weighted)
                 set.insert(str_serialized, columns[1]->getUInt(row_num));
             else

src/AggregateFunctions/Combinators/AggregateFunctionDistinct.h

@@ -116,7 +116,7 @@ struct AggregateFunctionDistinctMultipleGenericData : public AggregateFunctionDi
         StringRef value(begin, 0);
         for (size_t i = 0; i < columns_num; ++i)
         {
-            auto cur_ref = columns[i]->serializeValueIntoArena(row_num, *arena, begin);
+            auto cur_ref = columns[i]->serializeAggregationStateValueIntoArena(row_num, *arena, begin);
             value.data = cur_ref.data - value.size;
             value.size += cur_ref.size;
         }
@@ -140,7 +140,7 @@ struct AggregateFunctionDistinctMultipleGenericData : public AggregateFunctionDi
                 history.emplace(ArenaKeyHolder{value, *arena}, it, inserted);
                 const char * pos = it->getValue().data;
                 for (auto & column : argument_columns)
-                    pos = column->deserializeAndInsertFromArena(pos);
+                    pos = column->deserializeAndInsertAggregationStateValueFromArena(pos);
             }
         }
     }

src/AggregateFunctions/KeyHolderHelpers.h

@@ -17,7 +17,7 @@ static auto getKeyHolder(const IColumn & column, size_t row_num, Arena & arena)
     else
     {
         const char * begin = nullptr;
-        StringRef serialized = column.serializeValueIntoArena(row_num, arena, begin);
+        StringRef serialized = column.serializeAggregationStateValueIntoArena(row_num, arena, begin);
         assert(serialized.data != nullptr);
         return SerializedKeyHolder{serialized, arena};
     }
@@ -29,7 +29,7 @@ static void deserializeAndInsert(StringRef str, IColumn & data_to)
     if constexpr (is_plain_column)
         data_to.insertData(str.data, str.size);
     else
-        std::ignore = data_to.deserializeAndInsertFromArena(str.data);
+        std::ignore = data_to.deserializeAndInsertAggregationStateValueFromArena(str.data);
 }
 
 }

src/Columns/ColumnArray.cpp

@@ -243,6 +243,27 @@ StringRef ColumnArray::serializeValueIntoArena(size_t n, Arena & arena, char con
 }
 
 
+StringRef ColumnArray::serializeAggregationStateValueIntoArena(size_t n, Arena & arena, char const *& begin) const
+{
+    size_t array_size = sizeAt(n);
+    size_t offset = offsetAt(n);
+
+    char * pos = arena.allocContinue(sizeof(array_size), begin);
+    memcpy(pos, &array_size, sizeof(array_size));
+
+    StringRef res(pos, sizeof(array_size));
+
+    for (size_t i = 0; i < array_size; ++i)
+    {
+        auto value_ref = getData().serializeAggregationStateValueIntoArena(offset + i, arena, begin);
+        res.data = value_ref.data - res.size;
+        res.size += value_ref.size;
+    }
+
+    return res;
+}
+
+
 char * ColumnArray::serializeValueIntoMemory(size_t n, char * memory) const
 {
     size_t array_size = sizeAt(n);
@@ -287,6 +308,18 @@ const char * ColumnArray::deserializeAndInsertFromArena(const char * pos)
     return pos;
 }
 
+const char * ColumnArray::deserializeAndInsertAggregationStateValueFromArena(const char * pos)
+{
+    size_t array_size = unalignedLoad<size_t>(pos);
+    pos += sizeof(array_size);
+
+    for (size_t i = 0; i < array_size; ++i)
+        pos = getData().deserializeAndInsertAggregationStateValueFromArena(pos);
+
+    getOffsets().push_back(getOffsets().back() + array_size);
+    return pos;
+}
+
 const char * ColumnArray::skipSerializedInArena(const char * pos) const
 {
     size_t array_size = unalignedLoad<size_t>(pos);

src/Columns/ColumnArray.h

@@ -79,9 +79,11 @@ class ColumnArray final : public COWHelper<IColumnHelper<ColumnArray>, ColumnArr
     bool isDefaultAt(size_t n) const override;
     void insertData(const char * pos, size_t length) override;
     StringRef serializeValueIntoArena(size_t n, Arena & arena, char const *& begin) const override;
+    StringRef serializeAggregationStateValueIntoArena(size_t n, Arena & arena, char const *& begin) const override;
     char * serializeValueIntoMemory(size_t, char * memory) const override;
     std::optional<size_t> getSerializedValueSize(size_t n) const override;
     const char * deserializeAndInsertFromArena(const char * pos) override;
+    const char * deserializeAndInsertAggregationStateValueFromArena(const char * pos) override;
     const char * skipSerializedInArena(const char * pos) const override;
     void updateHashWithValue(size_t n, SipHash & hash) const override;
     WeakHash32 getWeakHash32() const override;

src/Columns/ColumnLowCardinality.cpp

@@ -279,6 +279,11 @@ StringRef ColumnLowCardinality::serializeValueIntoArena(size_t n, Arena & arena,
     return getDictionary().serializeValueIntoArena(getIndexes().getUInt(n), arena, begin);
 }
 
+StringRef ColumnLowCardinality::serializeAggregationStateValueIntoArena(size_t n, Arena & arena, char const *& begin) const
+{
+    return getDictionary().serializeAggregationStateValueIntoArena(getIndexes().getUInt(n), arena, begin);
+}
+
 char * ColumnLowCardinality::serializeValueIntoMemory(size_t n, char * memory) const
 {
     return getDictionary().serializeValueIntoMemory(getIndexes().getUInt(n), memory);
@@ -312,6 +317,16 @@ const char * ColumnLowCardinality::deserializeAndInsertFromArena(const char * po
     return new_pos;
 }
 
+const char * ColumnLowCardinality::deserializeAndInsertAggregationStateValueFromArena(const char * pos)
+{
+    compactIfSharedDictionary();
+
+    const char * new_pos;
+    idx.insertPosition(getDictionary().uniqueDeserializeAndInsertAggregationStateValueFromArena(pos, new_pos));
+
+    return new_pos;
+}
+
 const char * ColumnLowCardinality::skipSerializedInArena(const char * pos) const
 {
     return getDictionary().skipSerializedInArena(pos);

src/Columns/ColumnLowCardinality.h

@@ -102,11 +102,13 @@ class ColumnLowCardinality final : public COWHelper<IColumnHelper<ColumnLowCardi
     void popBack(size_t n) override { idx.popBack(n); }
 
     StringRef serializeValueIntoArena(size_t n, Arena & arena, char const *& begin) const override;
+    StringRef serializeAggregationStateValueIntoArena(size_t n, Arena & arena, char const *& begin) const override;
     char * serializeValueIntoMemory(size_t n, char * memory) const override;
 
     void collectSerializedValueSizes(PaddedPODArray<UInt64> & sizes, const UInt8 * is_null) const override;
 
     const char * deserializeAndInsertFromArena(const char * pos) override;
+    const char * deserializeAndInsertAggregationStateValueFromArena(const char * pos) override;
 
     const char * skipSerializedInArena(const char * pos) const override;
 

src/Columns/ColumnMap.cpp

@@ -155,6 +155,11 @@ StringRef ColumnMap::serializeValueIntoArena(size_t n, Arena & arena, char const
     return nested->serializeValueIntoArena(n, arena, begin);
 }
 
+StringRef ColumnMap::serializeAggregationStateValueIntoArena(size_t n, Arena & arena, char const *& begin) const
+{
+    return nested->serializeAggregationStateValueIntoArena(n, arena, begin);
+}
+
 char * ColumnMap::serializeValueIntoMemory(size_t n, char * memory) const
 {
     return nested->serializeValueIntoMemory(n, memory);
@@ -170,6 +175,11 @@ const char * ColumnMap::deserializeAndInsertFromArena(const char * pos)
     return nested->deserializeAndInsertFromArena(pos);
 }
 
+const char * ColumnMap::deserializeAndInsertAggregationStateValueFromArena(const char * pos)
+{
+    return nested->deserializeAndInsertAggregationStateValueFromArena(pos);
+}
+
 const char * ColumnMap::skipSerializedInArena(const char * pos) const
 {
     return nested->skipSerializedInArena(pos);