Optimize agg with collation & Optimize mem utils

CMakeLists.txt

@@ -511,6 +511,11 @@ if (ARCH_AMD64)
     else()
         add_definitions(-DTIFLASH_COMPILER_VPCLMULQDQ_SUPPORT=0)
     endif()
+
+    check_cxx_compiler_flag("-mmovbe" TIFLASH_COMPILER_MOVBE_SUPPORT)
+    if (TIFLASH_COMPILER_MOVBE_SUPPORT)
+        set(COMPILER_MOVBE_FLAG "-mmovbe")
+    endif()
 else()
     add_definitions(-DTIFLASH_COMPILER_VPCLMULQDQ_SUPPORT=0)
 endif()

dbms/src/AggregateFunctions/AggregateFunctionMinMaxAny.h

@@ -200,21 +200,21 @@ struct SingleValueDataString
     {
         if (collator == nullptr)
             return a < b;
-        return collator->compare(a.data, a.size, b.data, b.size) < 0;
+        return collator->compareFastPath(a.data, a.size, b.data, b.size) < 0;
     }
 
     bool greater(const StringRef & a, const StringRef & b) const
     {
         if (collator == nullptr)
             return a > b;
-        return collator->compare(a.data, a.size, b.data, b.size) > 0;
+        return collator->compareFastPath(a.data, a.size, b.data, b.size) > 0;
     }
 
     bool equalTo(const StringRef & a, const StringRef & b) const
     {
         if (collator == nullptr)
             return a == b;
-        return collator->compare(a.data, a.size, b.data, b.size) == 0;
+        return collator->compareFastPath(a.data, a.size, b.data, b.size) == 0;
     }
 
 public:

dbms/src/AggregateFunctions/IAggregateFunction.h

@@ -439,7 +439,7 @@ struct AggregationCollatorsWrapper<true>
         if (likely(collators.size() > column_index))
         {
             if (collators[column_index] != nullptr)
-                return collators[column_index]->sortKey(in.data, in.size, sort_key_containers[column_index]);
+                return collators[column_index]->sortKeyFastPath(in.data, in.size, sort_key_containers[column_index]);
             return in;
         }
         else if (collators.empty())

dbms/src/Columns/ColumnString.cpp

@@ -18,6 +18,7 @@
 #include <Common/HashTable/Hash.h>
 #include <DataStreams/ColumnGathererStream.h>
 #include <Storages/Transaction/CollatorUtils.h>
+#include <common/memcpy.h>
 #include <fmt/core.h>
 
 
@@ -94,7 +95,7 @@ void ColumnString::insertRangeFrom(const IColumn & src, size_t start, size_t len
 
     size_t old_chars_size = chars.size();
     chars.resize(old_chars_size + nested_length);
-    memcpy(&chars[old_chars_size], &src_concrete.chars[nested_offset], nested_length);
+    inline_memcpy(&chars[old_chars_size], &src_concrete.chars[nested_offset], nested_length);
 
     if (start == 0 && offsets.empty())
     {

dbms/src/Columns/ColumnString.h

@@ -211,10 +211,10 @@ class ColumnString final : public COWPtrHelper<IColumn, ColumnString>
 
         StringRef res;
 
-        if (collator != nullptr)
+        if (likely(collator != nullptr))
         {
             // Skip last zero byte.
-            auto sort_key = collator->sortKey(reinterpret_cast<const char *>(src), string_size - 1, sort_key_container);
+            auto sort_key = collator->sortKeyFastPath(reinterpret_cast<const char *>(src), string_size - 1, sort_key_container);
             string_size = sort_key.size;
             src = sort_key.data;
         }
@@ -244,10 +244,10 @@ class ColumnString final : public COWPtrHelper<IColumn, ColumnString>
     {
         size_t string_size = sizeAt(n);
         size_t offset = offsetAt(n);
-        if (collator != nullptr)
+        if (likely(collator != nullptr))
         {
             // Skip last zero byte.
-            auto sort_key = collator->sortKey(reinterpret_cast<const char *>(&chars[offset]), string_size - 1, sort_key_container);
+            auto sort_key = collator->sortKeyFastPath(reinterpret_cast<const char *>(&chars[offset]), string_size - 1, sort_key_container);
             string_size = sort_key.size;
             hash.update(reinterpret_cast<const char *>(&string_size), sizeof(string_size));
             hash.update(sort_key.data, sort_key.size);

dbms/src/Columns/ColumnsCommon.cpp

@@ -12,13 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#if __SSE2__
-#include <emmintrin.h>
-#endif
-
 #include <Columns/ColumnsCommon.h>
 #include <Columns/IColumn.h>
-
+#include <common/memcpy.h>
 
 namespace DB
 {
@@ -146,7 +142,7 @@ struct ResultOffsetsBuilder
     {
         const auto offsets_size_old = res_offsets.size();
         res_offsets.resize(offsets_size_old + SIMD_BYTES);
-        memcpy(&res_offsets[offsets_size_old], src_offsets_pos, SIMD_BYTES * sizeof(IColumn::Offset));
+        inline_memcpy(&res_offsets[offsets_size_old], src_offsets_pos, SIMD_BYTES * sizeof(IColumn::Offset));
 
         if (!first)
         {
@@ -194,7 +190,7 @@ void filterArraysImplGeneric(
 {
     const size_t size = src_offsets.size();
     if (size != filt.size())
-        throw Exception("Size of filter doesn't match size of column.", ErrorCodes::SIZES_OF_COLUMNS_DOESNT_MATCH);
+        throw Exception(fmt::format("size of filter {} doesn't match size of column {}", filt.size(), size), ErrorCodes::SIZES_OF_COLUMNS_DOESNT_MATCH);
 
     ResultOffsetsBuilder result_offsets_builder(res_offsets);
 
@@ -223,7 +219,7 @@ void filterArraysImplGeneric(
 
         const auto elems_size_old = res_elems.size();
         res_elems.resize(elems_size_old + size);
-        memcpy(&res_elems[elems_size_old], &src_elems[offset], size * sizeof(T));
+        inline_memcpy(&res_elems[elems_size_old], &src_elems[offset], size * sizeof(T));
     };
 
 #if __SSE2__
@@ -233,7 +229,7 @@ void filterArraysImplGeneric(
 
     while (filt_pos < filt_end_aligned)
     {
-        const auto mask = _mm_movemask_epi8(_mm_cmpgt_epi8(
+        uint32_t mask = _mm_movemask_epi8(_mm_cmpgt_epi8(
             _mm_loadu_si128(reinterpret_cast<const __m128i *>(filt_pos)),
             zero_vec));
 
@@ -254,13 +250,16 @@ void filterArraysImplGeneric(
             /// copy elements for SIMD_BYTES arrays at once
             const auto elems_size_old = res_elems.size();
             res_elems.resize(elems_size_old + chunk_size);
-            memcpy(&res_elems[elems_size_old], &src_elems[chunk_offset], chunk_size * sizeof(T));
+            inline_memcpy(&res_elems[elems_size_old], &src_elems[chunk_offset], chunk_size * sizeof(T));
         }
         else
         {
-            for (size_t i = 0; i < SIMD_BYTES; ++i)
-                if (filt_pos[i])
-                    copy_array(offsets_pos + i);
+            while (mask)
+            {
+                size_t index = __builtin_ctz(mask);
+                copy_array(offsets_pos + index);
+                mask = mask & (mask - 1);
+            }
         }
 
         filt_pos += SIMD_BYTES;

dbms/src/Common/ColumnsHashing.h

@@ -25,6 +25,7 @@
 #include <Core/Defines.h>
 #include <Interpreters/AggregationCommon.h>
 #include <Storages/Transaction/Collator.h>
+#include <common/memcpy.h>
 #include <common/unaligned.h>
 
 #include <memory>
@@ -72,7 +73,7 @@ struct HashMethodOneNumber
     using Base::getHash; /// (const Data & data, size_t row, Arena & pool) -> size_t
 
     /// Is used for default implementation in HashMethodBase.
-    FieldType getKeyHolder(size_t row, Arena *, std::vector<String> &) const
+    ALWAYS_INLINE inline FieldType getKeyHolder(size_t row, Arena *, std::vector<String> &) const
     {
         if constexpr (std::is_same_v<FieldType, Int256>)
             return vec[row];
@@ -99,7 +100,7 @@ struct HashMethodString
     HashMethodString(const ColumnRawPtrs & key_columns, const Sizes & /*key_sizes*/, const TiDB::TiDBCollators & collators)
     {
         const IColumn & column = *key_columns[0];
-        const ColumnString & column_string = assert_cast<const ColumnString &>(column);
+        const auto & column_string = assert_cast<const ColumnString &>(column);
         offsets = column_string.getOffsets().data();
         chars = column_string.getChars().data();
         if (!collators.empty())
@@ -110,15 +111,15 @@ struct HashMethodString
         }
     }
 
-    auto getKeyHolder(ssize_t row, [[maybe_unused]] Arena * pool, std::vector<String> & sort_key_containers) const
+    ALWAYS_INLINE inline auto getKeyHolder(ssize_t row, [[maybe_unused]] Arena * pool, std::vector<String> & sort_key_containers) const
     {
         auto last_offset = row == 0 ? 0 : offsets[row - 1];
         StringRef key(chars + last_offset, offsets[row] - last_offset - 1);
 
         if constexpr (place_string_to_arena)
         {
-            if (collator)
-                key = collator->sortKey(key.data, key.size, sort_key_containers[0]);
+            if (likely(collator))
+                key = collator->sortKeyFastPath(key.data, key.size, sort_key_containers[0]);
             return ArenaKeyHolder{key, *pool};
         }
         else
@@ -127,6 +128,98 @@ struct HashMethodString
         }
     }
 
+protected:
+    friend class columns_hashing_impl::HashMethodBase<Self, Value, Mapped, use_cache>;
+};
+
+/// For the case when there is multi string key.
+template <typename Value, typename Mapped>
+struct HashMethodMultiString
+    : public columns_hashing_impl::HashMethodBase<HashMethodMultiString<Value, Mapped>, Value, Mapped, false>
+{
+    using Self = HashMethodMultiString<Value, Mapped>;
+    using Base = columns_hashing_impl::HashMethodBase<Self, Value, Mapped, false>;
+
+    std::vector<const IColumn::Offset *> offsets;
+    std::vector<const UInt8 *> chars;
+    TiDB::TiDBCollators collators;
+    bool all_collators_padding_bin = false;
+
+    HashMethodMultiString(const ColumnRawPtrs & key_columns, const Sizes &, const TiDB::TiDBCollators & collators_)
+        : collators(collators_)
+    {
+        size_t num = key_columns.size();
+        offsets.resize(num);
+        chars.resize(num);
+
+        for (size_t i = 0; i < num; ++i)
+        {
+            const IColumn & column = *key_columns[i];
+            const auto & column_string = assert_cast<const ColumnString &>(column);
+            offsets[i] = column_string.getOffsets().data();
+            chars[i] = column_string.getChars().data();
+        }
+        if (!collators.empty())
+        {
+            all_collators_padding_bin = std::all_of(collators.begin(), collators.end(), [](auto & x) {
+                return x->isPaddingBinary();
+            });
+        }
+    }
+
+    template <typename F>
+    ALWAYS_INLINE inline SerializedKeyHolder genSerializedKeyHolder(ssize_t row, Arena * pool, F && fn_handle_key) const
+    {
+        auto num = offsets.size();
+
+        static_assert(std::is_same_v<size_t, decltype(reinterpret_cast<const StringRef *>(0)->size)>);
+
+        const char * begin = nullptr;
+        size_t sum_size = 0;
+
+        for (size_t key_index = 0; key_index < num; ++key_index)
+        {
+            auto last_offset = row == 0 ? 0 : offsets[key_index][row - 1];
+            StringRef key(chars[key_index] + last_offset, offsets[key_index][row] - last_offset - 1);
+
+            key = fn_handle_key(key_index, key);
+
+            char * pos = pool->allocContinue(key.size + sizeof(key.size), begin);
+            {
+                memcpy(pos, &key.size, sizeof(key.size));
+                inline_memcpy(pos + sizeof(key.size), key.data, key.size);
+            }
+
+            sum_size += key.size + sizeof(key.size);
+        }
+        return SerializedKeyHolder{{begin, sum_size}, *pool};
+    }
+
+    ALWAYS_INLINE inline auto getKeyHolder(ssize_t row, Arena * pool, std::vector<String> & sort_key_containers) const
+    {
+        if (likely(all_collators_padding_bin))
+        {
+            return genSerializedKeyHolder(row, pool, [](size_t, StringRef key) {
+                return DB::BinCollatorSortKey<true>(key.data, key.size);
+            });
+        }
+
+        if (unlikely(collators.empty()))
+        {
+            return genSerializedKeyHolder(row, pool, [](size_t, StringRef key) {
+                return key;
+            });
+        }
+        else
+        {
+            return genSerializedKeyHolder(row, pool, [&](size_t key_index, StringRef key) {
+                if (collators[key_index])
+                    return collators[key_index]->sortKey(key.data, key.size, sort_key_containers[key_index]);
+                return key;
+            });
+        }
+    }
+
 protected:
     friend class columns_hashing_impl::HashMethodBase<Self, Value, Mapped, false>;
 };
@@ -147,20 +240,20 @@ struct HashMethodFixedString
     HashMethodFixedString(const ColumnRawPtrs & key_columns, const Sizes & /*key_sizes*/, const TiDB::TiDBCollators & collators)
     {
         const IColumn & column = *key_columns[0];
-        const ColumnFixedString & column_string = assert_cast<const ColumnFixedString &>(column);
+        const auto & column_string = assert_cast<const ColumnFixedString &>(column);
         n = column_string.getN();
         chars = &column_string.getChars();
         if (!collators.empty())
             collator = collators[0];
     }
 
-    auto getKeyHolder(size_t row, [[maybe_unused]] Arena * pool, std::vector<String> & sort_key_containers) const
+    ALWAYS_INLINE inline auto getKeyHolder(size_t row, [[maybe_unused]] Arena * pool, std::vector<String> & sort_key_containers) const
     {
         StringRef key(&(*chars)[row * n], n);
 
         if (collator)
         {
-            key = collator->sortKey(key.data, key.size, sort_key_containers[0]);
+            key = collator->sortKeyFastPath(key.data, key.size, sort_key_containers[0]);
         }
 
         if constexpr (place_string_to_arena)
@@ -280,7 +373,7 @@ struct HashMethodKeysFixed
 #endif
     }
 
-    ALWAYS_INLINE Key getKeyHolder(size_t row, Arena *, std::vector<String> &) const
+    ALWAYS_INLINE inline Key getKeyHolder(size_t row, Arena *, std::vector<String> &) const
     {
         if constexpr (has_nullable_keys)
         {
@@ -357,7 +450,7 @@ struct HashMethodSerialized
         , collators(collators_)
     {}
 
-    ALWAYS_INLINE SerializedKeyHolder getKeyHolder(size_t row, Arena * pool, std::vector<String> & sort_key_containers) const
+    ALWAYS_INLINE inline SerializedKeyHolder getKeyHolder(size_t row, Arena * pool, std::vector<String> & sort_key_containers) const
     {
         return SerializedKeyHolder{
             serializeKeysToPoolContiguous(row, keys_size, key_columns, collators, sort_key_containers, *pool),
@@ -385,7 +478,7 @@ struct HashMethodHashed
         , collators(collators_)
     {}
 
-    ALWAYS_INLINE Key getKeyHolder(size_t row, Arena *, std::vector<String> & sort_key_containers) const
+    ALWAYS_INLINE inline Key getKeyHolder(size_t row, Arena *, std::vector<String> & sort_key_containers) const
     {
         return hash128(row, key_columns.size(), key_columns, collators, sort_key_containers);
     }

dbms/src/Common/ColumnsHashingImpl.h

@@ -128,21 +128,21 @@ class HashMethodBase
     using Cache = LastElementCache<Value, consecutive_keys_optimization>;
 
     template <typename Data>
-    ALWAYS_INLINE EmplaceResult emplaceKey(Data & data, size_t row, Arena & pool, std::vector<String> & sort_key_containers)
+    ALWAYS_INLINE inline EmplaceResult emplaceKey(Data & data, size_t row, Arena & pool, std::vector<String> & sort_key_containers)
     {
         auto key_holder = static_cast<Derived &>(*this).getKeyHolder(row, &pool, sort_key_containers);
         return emplaceImpl(key_holder, data);
     }
 
     template <typename Data>
-    ALWAYS_INLINE FindResult findKey(Data & data, size_t row, Arena & pool, std::vector<String> & sort_key_containers)
+    ALWAYS_INLINE inline FindResult findKey(Data & data, size_t row, Arena & pool, std::vector<String> & sort_key_containers)
     {
         auto key_holder = static_cast<Derived &>(*this).getKeyHolder(row, &pool, sort_key_containers);
         return findKeyImpl(keyHolderGetKey(key_holder), data);
     }
 
     template <typename Data>
-    ALWAYS_INLINE size_t getHash(const Data & data, size_t row, Arena & pool, std::vector<String> & sort_key_containers)
+    ALWAYS_INLINE inline size_t getHash(const Data & data, size_t row, Arena & pool, std::vector<String> & sort_key_containers)
     {
         auto key_holder = static_cast<Derived &>(*this).getKeyHolder(row, &pool, sort_key_containers);
         return data.hash(keyHolderGetKey(key_holder));
@@ -167,7 +167,7 @@ class HashMethodBase
     }
 
     template <typename Data, typename KeyHolder>
-    ALWAYS_INLINE EmplaceResult emplaceImpl(KeyHolder & key_holder, Data & data)
+    ALWAYS_INLINE inline EmplaceResult emplaceImpl(KeyHolder & key_holder, Data & data)
     {
         if constexpr (Cache::consecutive_keys_optimization)
         {
@@ -220,7 +220,7 @@ class HashMethodBase
     }
 
     template <typename Data, typename Key>
-    ALWAYS_INLINE FindResult findKeyImpl(Key key, Data & data)
+    ALWAYS_INLINE inline FindResult findKeyImpl(Key key, Data & data)
     {
         if constexpr (Cache::consecutive_keys_optimization)
         {