Aggregator support batch serialize

dbms/src/AggregateFunctions/AggregateFunctionGroupUniqArray.h

@@ -182,18 +182,18 @@ class AggregateFunctionGroupUniqArrayGeneric
         {
             // We have to copy the keys to our arena.
             assert(arena != nullptr);
-            cur_set.emplace(ArenaKeyHolder{rhs_elem.getValue(), *arena}, it, inserted);
+            cur_set.emplace(ArenaKeyHolder{rhs_elem.getValue(), arena}, it, inserted);
         }
     }
 
     void insertResultInto(ConstAggregateDataPtr __restrict place, IColumn & to, Arena *) const override
     {
-        ColumnArray & arr_to = assert_cast<ColumnArray &>(to);
+        auto & arr_to = assert_cast<ColumnArray &>(to);
         ColumnArray::Offsets & offsets_to = arr_to.getOffsets();
         IColumn & data_to = arr_to.getData();
 
         auto & set = this->data(place).value;
-        offsets_to.push_back((offsets_to.size() == 0 ? 0 : offsets_to.back()) + set.size());
+        offsets_to.push_back((offsets_to.empty() ? 0 : offsets_to.back()) + set.size());
 
         for (auto & elem : set)
             deserializeAndInsert<is_plain_column>(elem.getValue(), data_to);

dbms/src/AggregateFunctions/KeyHolderHelpers.h

@@ -24,14 +24,14 @@ inline auto getKeyHolder(const IColumn & column, size_t row_num, Arena & arena)
 {
     if constexpr (is_plain_column)
     {
-        return ArenaKeyHolder{column.getDataAt(row_num), arena};
+        return ArenaKeyHolder{column.getDataAt(row_num), &arena};
     }
     else
     {
         const char * begin = nullptr;
         StringRef serialized = column.serializeValueIntoArena(row_num, arena, begin);
         assert(serialized.data != nullptr);
-        return SerializedKeyHolder{serialized, arena};
+        return SerializedKeyHolder{serialized, &arena};
     }
 }
 

dbms/src/Columns/ColumnAggregateFunction.h

@@ -165,10 +165,27 @@ class ColumnAggregateFunction final : public COWPtrHelper<IColumn, ColumnAggrega
 
     const char * deserializeAndInsertFromArena(const char * src_arena, const TiDB::TiDBCollatorPtr &) override;
 
+    void countSerializeByteSizeUnique(
+        PaddedPODArray<size_t> & /* byte_size */,
+        const TiDB::TiDBCollatorPtr & /* collator */) const override
+    {
+        throw Exception(
+            "Method countSerializeByteSizeUnique is not supported for " + getName(),
+            ErrorCodes::NOT_IMPLEMENTED);
+    }
     void countSerializeByteSize(PaddedPODArray<size_t> & /* byte_size */) const override
     {
         throw Exception("Method countSerializeByteSize is not supported for " + getName(), ErrorCodes::NOT_IMPLEMENTED);
     }
+    void countSerializeByteSizeUniqueForColumnArray(
+        PaddedPODArray<size_t> & /* byte_size */,
+        const IColumn::Offsets & /* offsets */,
+        const TiDB::TiDBCollatorPtr & /* collator */) const override
+    {
+        throw Exception(
+            "Method countSerializeByteSizeUniqueForColumnArray is not supported for " + getName(),
+            ErrorCodes::NOT_IMPLEMENTED);
+    }
     void countSerializeByteSizeForColumnArray(
         PaddedPODArray<size_t> & /* byte_size */,
         const IColumn::Offsets & /* offsets */) const override
@@ -178,6 +195,16 @@ class ColumnAggregateFunction final : public COWPtrHelper<IColumn, ColumnAggrega
             ErrorCodes::NOT_IMPLEMENTED);
     }
 
+    void serializeToPosUnique(
+        PaddedPODArray<char *> & /* pos */,
+        size_t /* start */,
+        size_t /* length */,
+        bool /* has_null */,
+        const TiDB::TiDBCollatorPtr & /* collator */,
+        String * /* sort_key_container */) const override
+    {
+        throw Exception("Method serializeToPosUnique is not supported for " + getName(), ErrorCodes::NOT_IMPLEMENTED);
+    }
     void serializeToPos(
         PaddedPODArray<char *> & /* pos */,
         size_t /* start */,
@@ -186,6 +213,19 @@ class ColumnAggregateFunction final : public COWPtrHelper<IColumn, ColumnAggrega
     {
         throw Exception("Method serializeToPos is not supported for " + getName(), ErrorCodes::NOT_IMPLEMENTED);
     }
+    void serializeToPosUniqueForColumnArray(
+        PaddedPODArray<char *> & /* pos */,
+        size_t /* start */,
+        size_t /* length */,
+        bool /* has_null */,
+        const IColumn::Offsets & /* offsets */,
+        const TiDB::TiDBCollatorPtr & /* collator */,
+        String * /* sort_key_container */) const override
+    {
+        throw Exception(
+            "Method serializeToPosUniqueForColumnArray is not supported for " + getName(),
+            ErrorCodes::NOT_IMPLEMENTED);
+    }
     void serializeToPosForColumnArray(
         PaddedPODArray<char *> & /* pos */,
         size_t /* start */,
@@ -198,14 +238,34 @@ class ColumnAggregateFunction final : public COWPtrHelper<IColumn, ColumnAggrega
             ErrorCodes::NOT_IMPLEMENTED);
     }
 
-    void deserializeAndInsertFromPos(PaddedPODArray<char *> & /* pos */, bool /* use_nt_align_buffer */) override
+    void deserializeAndInsertFromPosUnique(
+        PaddedPODArray<const char *> & /* pos */,
+        bool /* use_nt_align_buffer */,
+        const TiDB::TiDBCollatorPtr & /* collator */) override
+    {
+        throw Exception(
+            "Method deserializeAndInsertFromPosUnique is not supported for " + getName(),
+            ErrorCodes::NOT_IMPLEMENTED);
+    }
+    void deserializeAndInsertFromPos(PaddedPODArray<const char *> & /* pos */, bool /* use_nt_align_buffer */) override
     {
         throw Exception(
             "Method deserializeAndInsertFromPos is not supported for " + getName(),
             ErrorCodes::NOT_IMPLEMENTED);
     }
+
+    void deserializeAndInsertFromPosUniqueForColumnArray(
+        PaddedPODArray<const char *> & /* pos */,
+        const IColumn::Offsets & /* array_offsets */,
+        bool /* use_nt_align_buffer */,
+        const TiDB::TiDBCollatorPtr & /* collator */) override
+    {
+        throw Exception(
+            "Method deserializeAndInsertFromPosUniqueForColumnArray is not supported for " + getName(),
+            ErrorCodes::NOT_IMPLEMENTED);
+    }
     void deserializeAndInsertFromPosForColumnArray(
-        PaddedPODArray<char *> & /* pos */,
+        PaddedPODArray<const char *> & /* pos */,
         const IColumn::Offsets & /* array_offsets */,
         bool /* use_nt_align_buffer */) override
     {

dbms/src/Columns/ColumnArray.cpp

@@ -218,7 +218,21 @@ const char * ColumnArray::deserializeAndInsertFromArena(const char * pos, const
     return pos;
 }
 
+void ColumnArray::countSerializeByteSizeUnique(
+    PaddedPODArray<size_t> & byte_size,
+    const TiDB::TiDBCollatorPtr & collator) const
+{
+    countSerializeByteSizeImpl<true>(byte_size, collator);
+}
+
 void ColumnArray::countSerializeByteSize(PaddedPODArray<size_t> & byte_size) const
+{
+    countSerializeByteSizeImpl<false>(byte_size, nullptr);
+}
+
+template <bool ensure_unique>
+void ColumnArray::countSerializeByteSizeImpl(PaddedPODArray<size_t> & byte_size, const TiDB::TiDBCollatorPtr & collator)
+    const
 {
     RUNTIME_CHECK_MSG(byte_size.size() == size(), "size of byte_size({}) != column size({})", byte_size.size(), size());
 
@@ -237,19 +251,41 @@ void ColumnArray::countSerializeByteSize(PaddedPODArray<size_t> & byte_size) con
     for (size_t i = 0; i < size; ++i)
         byte_size[i] += sizeof(UInt32);
 
-    getData().countSerializeByteSizeForColumnArray(byte_size, getOffsets());
+    if constexpr (ensure_unique)
+        getData().countSerializeByteSizeUniqueForColumnArray(byte_size, getOffsets(), collator);
+    else
+        getData().countSerializeByteSizeForColumnArray(byte_size, getOffsets());
+}
+
+void ColumnArray::serializeToPosUnique(
+    PaddedPODArray<char *> & pos,
+    size_t start,
+    size_t length,
+    bool has_null,
+    const TiDB::TiDBCollatorPtr & collator,
+    String * sort_key_container) const
+{
+    if (has_null)
+        serializeToPosImpl<true, true>(pos, start, length, collator, sort_key_container);
+    else
+        serializeToPosImpl<false, true>(pos, start, length, collator, sort_key_container);
 }
 
 void ColumnArray::serializeToPos(PaddedPODArray<char *> & pos, size_t start, size_t length, bool has_null) const
 {
     if (has_null)
-        serializeToPosImpl<true>(pos, start, length);
+        serializeToPosImpl<true, false>(pos, start, length, nullptr, nullptr);
     else
-        serializeToPosImpl<false>(pos, start, length);
+        serializeToPosImpl<false, false>(pos, start, length, nullptr, nullptr);
 }
 
-template <bool has_null>
-void ColumnArray::serializeToPosImpl(PaddedPODArray<char *> & pos, size_t start, size_t length) const
+template <bool has_null, bool ensure_unique>
+void ColumnArray::serializeToPosImpl(
+    PaddedPODArray<char *> & pos,
+    size_t start,
+    size_t length,
+    const TiDB::TiDBCollatorPtr & collator,
+    String * sort_key_container) const
 {
     RUNTIME_CHECK_MSG(length <= pos.size(), "length({}) > size of pos({})", length, pos.size());
     RUNTIME_CHECK_MSG(start + length <= size(), "start({}) + length({}) > size of column({})", start, length, size());
@@ -267,10 +303,37 @@ void ColumnArray::serializeToPosImpl(PaddedPODArray<char *> & pos, size_t start,
         pos[i] += sizeof(UInt32);
     }
 
-    getData().serializeToPosForColumnArray(pos, start, length, has_null, getOffsets());
+    if constexpr (ensure_unique)
+        getData().serializeToPosUniqueForColumnArray(
+            pos,
+            start,
+            length,
+            has_null,
+            getOffsets(),
+            collator,
+            sort_key_container);
+    else
+        getData().serializeToPosForColumnArray(pos, start, length, has_null, getOffsets());
 }
 
-void ColumnArray::deserializeAndInsertFromPos(PaddedPODArray<char *> & pos, bool use_nt_align_buffer)
+void ColumnArray::deserializeAndInsertFromPosUnique(
+    PaddedPODArray<const char *> & pos,
+    bool use_nt_align_buffer,
+    const TiDB::TiDBCollatorPtr & collator)
+{
+    deserializeAndInsertFromPosImpl<true>(pos, use_nt_align_buffer, collator);
+}
+
+void ColumnArray::deserializeAndInsertFromPos(PaddedPODArray<const char *> & pos, bool use_nt_align_buffer)
+{
+    deserializeAndInsertFromPosImpl<false>(pos, use_nt_align_buffer, nullptr);
+}
+
+template <bool ensure_unique>
+void ColumnArray::deserializeAndInsertFromPosImpl(
+    PaddedPODArray<const char *> & pos,
+    bool use_nt_align_buffer,
+    const TiDB::TiDBCollatorPtr & collator)
 {
     auto & offsets = getOffsets();
     size_t prev_size = offsets.size();
@@ -285,7 +348,10 @@ void ColumnArray::deserializeAndInsertFromPos(PaddedPODArray<char *> & pos, bool
         pos[i] += sizeof(UInt32);
     }
 
-    getData().deserializeAndInsertFromPosForColumnArray(pos, offsets, use_nt_align_buffer);
+    if constexpr (ensure_unique)
+        getData().deserializeAndInsertFromPosUniqueForColumnArray(pos, offsets, use_nt_align_buffer, collator);
+    else
+        getData().deserializeAndInsertFromPosForColumnArray(pos, offsets, use_nt_align_buffer);
 }
 
 void ColumnArray::flushNTAlignBuffer()

dbms/src/Columns/ColumnArray.h

@@ -44,6 +44,23 @@ class ColumnArray final : public COWPtrHelper<IColumn, ColumnArray>
 
     ColumnArray(const ColumnArray &) = default;
 
+    template <bool ensure_unique>
+    void countSerializeByteSizeImpl(PaddedPODArray<size_t> & byte_size, const TiDB::TiDBCollatorPtr & collator) const;
+
+    template <bool has_null, bool ensure_unique>
+    void serializeToPosImpl(
+        PaddedPODArray<char *> & pos,
+        size_t start,
+        size_t length,
+        const TiDB::TiDBCollatorPtr & collator,
+        String * sort_key_container) const;
+
+    template <bool ensure_unique>
+    void deserializeAndInsertFromPosImpl(
+        PaddedPODArray<const char *> & pos,
+        bool use_nt_align_buffer,
+        const TiDB::TiDBCollatorPtr & collator);
+
 public:
     /** Create immutable column using immutable arguments. This arguments may be shared with other columns.
       * Use IColumn::mutate in order to make mutable column and mutate shared nested columns.
@@ -82,7 +99,19 @@ class ColumnArray final : public COWPtrHelper<IColumn, ColumnArray>
         String &) const override;
     const char * deserializeAndInsertFromArena(const char * pos, const TiDB::TiDBCollatorPtr &) override;
 
+    void countSerializeByteSizeUnique(PaddedPODArray<size_t> & byte_size, const TiDB::TiDBCollatorPtr & collator)
+        const override;
     void countSerializeByteSize(PaddedPODArray<size_t> & byte_size) const override;
+
+    void countSerializeByteSizeUniqueForColumnArray(
+        PaddedPODArray<size_t> & /* byte_size */,
+        const IColumn::Offsets & /* array_offsets */,
+        const TiDB::TiDBCollatorPtr & /* collator */) const override
+    {
+        throw Exception(
+            "Method countSerializeByteSizeUniqueForColumnArray is not supported for " + getName(),
+            ErrorCodes::NOT_IMPLEMENTED);
+    }
     void countSerializeByteSizeForColumnArray(
         PaddedPODArray<size_t> & /* byte_size */,
         const IColumn::Offsets & /* array_offsets */) const override
@@ -92,10 +121,28 @@ class ColumnArray final : public COWPtrHelper<IColumn, ColumnArray>
             ErrorCodes::NOT_IMPLEMENTED);
     }
 
+    void serializeToPosUnique(
+        PaddedPODArray<char *> & pos,
+        size_t start,
+        size_t length,
+        bool has_null,
+        const TiDB::TiDBCollatorPtr & collator,
+        String * sort_key_container) const override;
     void serializeToPos(PaddedPODArray<char *> & pos, size_t start, size_t length, bool has_null) const override;
-    template <bool has_null>
-    void serializeToPosImpl(PaddedPODArray<char *> & pos, size_t start, size_t length) const;
 
+    void serializeToPosUniqueForColumnArray(
+        PaddedPODArray<char *> & /* pos */,
+        size_t /* start */,
+        size_t /* length */,
+        bool /* has_null */,
+        const IColumn::Offsets & /* array_offsets */,
+        const TiDB::TiDBCollatorPtr & /* collator */,
+        String * /* sort_key_container */) const override
+    {
+        throw Exception(
+            "Method serializeToPosUniqueForColumnArray is not supported for " + getName(),
+            ErrorCodes::NOT_IMPLEMENTED);
+    }
     void serializeToPosForColumnArray(
         PaddedPODArray<char *> & /* pos */,
         size_t /* start */,
@@ -108,9 +155,24 @@ class ColumnArray final : public COWPtrHelper<IColumn, ColumnArray>
             ErrorCodes::NOT_IMPLEMENTED);
     }
 
-    void deserializeAndInsertFromPos(PaddedPODArray<char *> & pos, bool use_nt_align_buffer) override;
+    void deserializeAndInsertFromPosUnique(
+        PaddedPODArray<const char *> & pos,
+        bool use_nt_align_buffer,
+        const TiDB::TiDBCollatorPtr & collator) override;
+    void deserializeAndInsertFromPos(PaddedPODArray<const char *> & pos, bool use_nt_align_buffer) override;
+
+    void deserializeAndInsertFromPosUniqueForColumnArray(
+        PaddedPODArray<const char *> & /* pos */,
+        const IColumn::Offsets & /* array_offsets */,
+        bool /* use_nt_align_buffer */,
+        const TiDB::TiDBCollatorPtr & /* collator */) override
+    {
+        throw Exception(
+            "Method deserializeAndInsertFromPosUniqueForColumnArray is not supported for " + getName(),
+            ErrorCodes::NOT_IMPLEMENTED);
+    }
     void deserializeAndInsertFromPosForColumnArray(
-        PaddedPODArray<char *> & /* pos */,
+        PaddedPODArray<const char *> & /* pos */,
         const IColumn::Offsets & /* array_offsets */,
         bool /* use_nt_align_buffer */) override
     {