From 8eb6ffcd49074e896a656bb0a7b489246aa6d28c Mon Sep 17 00:00:00 2001 From: ChangRui-Ryan Date: Tue, 14 Jul 2026 19:43:37 +0800 Subject: [PATCH 1/3] executor: Improve HashAgg string collation key caching --- dbms/src/Common/ColumnsHashing.h | 118 ++++++++- dbms/src/Common/ColumnsHashingImpl.h | 3 + .../tests/bench_aggregation_hash_map.cpp | 237 ++++++++++++++++++ .../tests/gtest_aggregation_executor.cpp | 185 ++++++++++++++ dbms/src/Interpreters/Aggregator.cpp | 4 + dbms/src/Interpreters/Aggregator.h | 3 + 6 files changed, 540 insertions(+), 10 deletions(-) diff --git a/dbms/src/Common/ColumnsHashing.h b/dbms/src/Common/ColumnsHashing.h index 3fa4db7e838..d74fb3dc72f 100644 --- a/dbms/src/Common/ColumnsHashing.h +++ b/dbms/src/Common/ColumnsHashing.h @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -38,6 +39,9 @@ extern const int LOGICAL_ERROR; namespace ColumnsHashing { +static constexpr size_t min_rows_to_use_collation_sort_key_cache = 1024; +static constexpr size_t max_collation_sort_key_cache_distinct_ratio = 32; + /// For the case when there is one numeric key. /// UInt8/16/32/64 for any type with corresponding bit width. template @@ -95,6 +99,67 @@ class KeyStringBatchHandlerBase size_t processed_row_idx = 0; std::vector sort_key_containers{}; std::vector batch_rows{}; + bool collation_sort_key_cache_active = false; + bool collation_sort_key_cache_fallback = false; + size_t collation_sort_key_cache_distinct_limit = 0; + HashMap raw_to_sort_key_index; + std::vector cached_sort_keys{}; + + template + StringRef getSortKeyWithCache( + const DerivedCollator & collator, + const StringRef & raw_key, + std::string & sort_key_container) + { + if (!collation_sort_key_cache_active) + return collator.sortKey(raw_key.data, raw_key.size, sort_key_container); + + if (auto it = raw_to_sort_key_index.find(raw_key); it != nullptr) + { + const auto & sort_key = cached_sort_keys[it->getMapped()]; + return {sort_key.data(), sort_key.size()}; + } + + if (raw_to_sort_key_index.size() >= collation_sort_key_cache_distinct_limit) + { + collation_sort_key_cache_active = false; + collation_sort_key_cache_fallback = true; + return collator.sortKey(raw_key.data, raw_key.size, sort_key_container); + } + + const auto sort_key_index = cached_sort_keys.size(); + cached_sort_keys.emplace_back(); + auto sort_key = collator.sortKey(raw_key.data, raw_key.size, cached_sort_keys.back()); + + typename decltype(raw_to_sort_key_index)::LookupResult it; + bool inserted = false; + raw_to_sort_key_index.emplace(raw_key, it, inserted); + RUNTIME_CHECK(inserted); + it->getMapped() = sort_key_index; + + return sort_key; + } + + template + void prepareNextBatchWithCollator( + const UInt8 * chars, + const IColumn::Offsets & offsets, + size_t cur_batch_size, + const DerivedCollator & collator) + { + for (size_t i = 0; i < cur_batch_size; ++i) + { + const auto row = processed_row_idx + i; + const auto last_offset = offsets[row - 1]; + // Remove last zero byte. + StringRef key(chars + last_offset, offsets[row] - last_offset - 1); + if constexpr (use_cache) + key = getSortKeyWithCache(collator, key, sort_key_containers[i]); + else + key = collator.sortKey(key.data, key.size, sort_key_containers[i]); + batch_rows[i] = key; + } + } template void prepareNextBatchType( @@ -108,23 +173,48 @@ class KeyStringBatchHandlerBase batch_rows.resize(cur_batch_size); - const auto * derived_collator = static_cast(collator); - for (size_t i = 0; i < cur_batch_size; ++i) + if constexpr (has_collator) { - const auto row = processed_row_idx + i; - const auto last_offset = offsets[row - 1]; - // Remove last zero byte. - StringRef key(chars + last_offset, offsets[row] - last_offset - 1); - if constexpr (has_collator) - key = derived_collator->sortKey(key.data, key.size, sort_key_containers[i]); - - batch_rows[i] = key; + const auto & derived_collator = *static_cast(collator); + if (collation_sort_key_cache_active) + prepareNextBatchWithCollator( + chars, + offsets, + cur_batch_size, + derived_collator); + else + prepareNextBatchWithCollator( + chars, + offsets, + cur_batch_size, + derived_collator); + } + else + { + for (size_t i = 0; i < cur_batch_size; ++i) + { + const auto row = processed_row_idx + i; + const auto last_offset = offsets[row - 1]; + // Remove last zero byte. + batch_rows[i] = StringRef(chars + last_offset, offsets[row] - last_offset - 1); + } } processed_row_idx += cur_batch_size; } protected: bool inited() const { return !sort_key_containers.empty(); } + bool hasCollationSortKeyCacheFallback() const { return collation_sort_key_cache_fallback; } + + void initCollationSortKeyCache(size_t rows) + { + if (rows < min_rows_to_use_collation_sort_key_cache) + return; + + collation_sort_key_cache_active = true; + collation_sort_key_cache_distinct_limit = rows / max_collation_sort_key_cache_distinct_ratio; + cached_sort_keys.reserve(collation_sort_key_cache_distinct_limit + 1); + } void init(size_t start_row, size_t max_batch_size) { @@ -207,6 +297,14 @@ struct HashMethodString collator = collators[0]; } + void initCollationSortKeyCache(size_t rows) + { + if (collator != nullptr && collator->isCI()) + BatchHandlerBase::initCollationSortKeyCache(rows); + } + + bool hasCollationSortKeyCacheFallback() const { return BatchHandlerBase::hasCollationSortKeyCacheFallback(); } + void initBatchHandler(size_t start_row, size_t max_batch_size) { assert(!BatchHandlerBase::inited()); diff --git a/dbms/src/Common/ColumnsHashingImpl.h b/dbms/src/Common/ColumnsHashingImpl.h index 4f6c56be6d4..d242c40eac9 100644 --- a/dbms/src/Common/ColumnsHashingImpl.h +++ b/dbms/src/Common/ColumnsHashingImpl.h @@ -129,6 +129,9 @@ class HashMethodBase using Cache = LastElementCache; using Derived = TDerived; + void initCollationSortKeyCache(size_t) {} + bool hasCollationSortKeyCacheFallback() const { return false; } + template ALWAYS_INLINE inline EmplaceResult emplaceKey( Data & data, diff --git a/dbms/src/Flash/tests/bench_aggregation_hash_map.cpp b/dbms/src/Flash/tests/bench_aggregation_hash_map.cpp index 46c9752ded6..e42327841c8 100644 --- a/dbms/src/Flash/tests/bench_aggregation_hash_map.cpp +++ b/dbms/src/Flash/tests/bench_aggregation_hash_map.cpp @@ -12,8 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include +#include #include #include #include @@ -27,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -339,5 +342,239 @@ try CATCH BENCHMARK_REGISTER_F(BenchProbeAggHashMap, basic); +class BenchStringCollationKeyCache : public ::benchmark::Fixture +{ +public: + void SetUp(const ::benchmark::State &) override + try + { + if (!AggregateFunctionFactory::instance().isAggregateFunctionName("count")) + DB::registerAggregateFunctions(); + + context = TiFlashTestEnv::getContext(); + params = buildParams(); + } + CATCH + + void TearDown(benchmark::State &) override { context.reset(); } + +protected: + static constexpr size_t rows_per_block = DEFAULT_BLOCK_SIZE; + static constexpr size_t block_num = 128; + static constexpr size_t total_rows = rows_per_block * block_num; + static constexpr size_t low_ndv = 64; + static constexpr size_t high_ndv = rows_per_block; + static constexpr size_t cache_distinct_limit + = rows_per_block / ColumnsHashing::max_collation_sort_key_cache_distinct_ratio; + static constexpr size_t ndv_transition_block = block_num / 2; + static constexpr const char * key_column_name = "key"; + static constexpr const char * count_column_name = "count(key)"; + + static std::vector generateKeys(size_t distinct_keys) + { + std::vector keys; + keys.reserve(distinct_keys); + for (size_t i = 0; i < distinct_keys; ++i) + { + if (i % 2 == 0) + keys.push_back(fmt::format("customer-region-{:05}-AlphaBetaGamma", i)); + else + keys.push_back(fmt::format("CUSTOMER-REGION-{:05}-alphabetagamma", i)); + } + return keys; + } + + static BlockPtr generateBlock( + const std::shared_ptr & data_type_string, + const std::vector & keys, + size_t block_index) + { + auto col_key = data_type_string->createColumn(); + col_key->reserve(rows_per_block); + for (size_t row = 0; row < rows_per_block; ++row) + { + const auto & key = keys[(block_index * rows_per_block + row) % keys.size()]; + col_key->insertData(key.data(), key.size()); + } + + ColumnsWithTypeAndName cols{ + ColumnWithTypeAndName(std::move(col_key), data_type_string, String(key_column_name)), + }; + return std::make_shared(cols); + } + + static std::vector generateData(size_t distinct_keys) + { + auto data_type_string = std::make_shared(); + auto keys = generateKeys(distinct_keys); + std::vector blocks; + blocks.reserve(block_num); + for (size_t block_index = 0; block_index < block_num; ++block_index) + blocks.push_back(generateBlock(data_type_string, keys, block_index)); + return blocks; + } + + static std::vector generateDataWithNDVTransition( + size_t first_distinct_keys, + size_t second_distinct_keys, + size_t transition_block) + { + RUNTIME_CHECK(transition_block < block_num); + auto data_type_string = std::make_shared(); + auto first_keys = generateKeys(first_distinct_keys); + auto second_keys = generateKeys(second_distinct_keys); + + std::vector blocks; + blocks.reserve(block_num); + for (size_t block_index = 0; block_index < block_num; ++block_index) + { + const auto & keys = block_index < transition_block ? first_keys : second_keys; + blocks.push_back(generateBlock(data_type_string, keys, block_index)); + } + return blocks; + } + + std::unique_ptr buildParams() + { + auto data_type_string = std::make_shared(); + Block src_header(NamesAndTypes{{String(key_column_name), data_type_string}}); + + AggregateDescription count_desc; + count_desc.function + = AggregateFunctionFactory::instance().get(*context, "count", {data_type_string}, {}, 0, false); + count_desc.arguments = {0}; + count_desc.argument_names = {String(key_column_name)}; + count_desc.column_name = String(count_column_name); + AggregateDescriptions aggregate_desc{std::move(count_desc)}; + + const auto & settings = context->getSettingsRef(); + SpillConfig spill_config( + context->getTemporaryPath(), + "BenchStringCollationKeyCache", + settings.max_cached_data_bytes_in_spiller, + settings.max_spilled_rows_per_file, + settings.max_spilled_bytes_per_file, + context->getFileProvider(), + /*for_all_constant_max_streams=*/1, + /*for_all_constant_block_size=*/rows_per_block); + TiDB::TiDBCollators collators{TiDB::ITiDBCollator::getCollator(TiDB::ITiDBCollator::UTF8MB4_GENERAL_CI)}; + + return std::make_unique( + src_header, + ColumnNumbers{0}, + KeyRefAggFuncMap{}, + AggFuncRefKeyMap{}, + aggregate_desc, + /*group_by_two_level_threshold=*/0, + /*group_by_two_level_threshold_bytes=*/0, + /*max_bytes_before_external_group_by=*/0, + /*empty_result_for_aggregation_by_empty_set=*/false, + spill_config, + /*max_block_size=*/rows_per_block, + /*use_magic_hash=*/false, + collators); + } + + static size_t runAggregation(const Aggregator::Params & params, const std::vector & blocks) + { + std::shared_ptr aggregator; + auto data_variants = std::make_shared(); + RegisterOperatorSpillContext register_operator_spill_context; + aggregator = std::make_shared( + params, + "BenchStringCollationKeyCache", + /*concurrency=*/1, + register_operator_spill_context, + /*is_auto_pass_through=*/false, + params.use_magic_hash); + data_variants->aggregator = aggregator.get(); + + Aggregator::AggProcessInfo agg_process_info(aggregator.get()); + for (const auto & block : blocks) + { + agg_process_info.resetBlock(*block); + do + { + aggregator->executeOnBlock(agg_process_info, *data_variants, 1); + } while (!agg_process_info.allBlockDataHandled()); + } + + std::vector variants{data_variants}; + auto merging_buckets = aggregator->mergeAndConvertToBlocks(variants, /*final=*/true, /*max_threads=*/1); + + size_t output_rows = 0; + for (;;) + { + auto block = merging_buckets->getData(0); + if (!block) + break; + output_rows += block.rows(); + } + return output_rows; + } + + void run(benchmark::State & state, const std::vector & blocks, const Aggregator::Params & params) + { + size_t output_rows = 0; + for (const auto & _ : state) + output_rows = runAggregation(params, blocks); + + benchmark::DoNotOptimize(output_rows); + state.SetItemsProcessed(static_cast(state.iterations() * total_rows)); + state.counters["rows_per_second"] + = benchmark::Counter(static_cast(state.iterations() * total_rows), benchmark::Counter::kIsRate); + state.counters["output_rows"] = static_cast(output_rows); + } + + ContextPtr context; + std::unique_ptr params; +}; + +BENCHMARK_DEFINE_F(BenchStringCollationKeyCache, low_ndv)(benchmark::State & state) +try +{ + auto blocks = generateData(low_ndv); + run(state, blocks, *params); +} +CATCH + +BENCHMARK_DEFINE_F(BenchStringCollationKeyCache, ndv_at_cache_limit)(benchmark::State & state) +try +{ + auto blocks = generateData(cache_distinct_limit); + run(state, blocks, *params); +} +CATCH + +BENCHMARK_DEFINE_F(BenchStringCollationKeyCache, ndv_over_cache_limit)(benchmark::State & state) +try +{ + auto blocks = generateData(cache_distinct_limit + 1); + run(state, blocks, *params); +} +CATCH + +BENCHMARK_DEFINE_F(BenchStringCollationKeyCache, high_ndv)(benchmark::State & state) +try +{ + auto blocks = generateData(high_ndv); + run(state, blocks, *params); +} +CATCH + +BENCHMARK_DEFINE_F(BenchStringCollationKeyCache, low_then_high_ndv)(benchmark::State & state) +try +{ + auto blocks = generateDataWithNDVTransition(low_ndv, high_ndv, ndv_transition_block); + run(state, blocks, *params); +} +CATCH + +BENCHMARK_REGISTER_F(BenchStringCollationKeyCache, low_ndv); +BENCHMARK_REGISTER_F(BenchStringCollationKeyCache, ndv_at_cache_limit); +BENCHMARK_REGISTER_F(BenchStringCollationKeyCache, ndv_over_cache_limit); +BENCHMARK_REGISTER_F(BenchStringCollationKeyCache, high_ndv); +BENCHMARK_REGISTER_F(BenchStringCollationKeyCache, low_then_high_ndv); + } // namespace tests } // namespace DB diff --git a/dbms/src/Flash/tests/gtest_aggregation_executor.cpp b/dbms/src/Flash/tests/gtest_aggregation_executor.cpp index b7d1af096b6..88cfe9dcc6b 100644 --- a/dbms/src/Flash/tests/gtest_aggregation_executor.cpp +++ b/dbms/src/Flash/tests/gtest_aggregation_executor.cpp @@ -13,6 +13,8 @@ // limitations under the License. #include +#include +#include #include #include #include @@ -1256,6 +1258,189 @@ try } CATCH +TEST_F(AggExecutorTestRunner, StringCollationKeyCacheGeneralCISemantics) +try +{ + const String db_name = "test_db"; + constexpr size_t rows = 1024; + DB::MockColumnInfoVec table_column_infos{ + {"key", TiDB::TP::TypeString, false, Poco::Dynamic::Var("utf8mb4_general_ci")}}; + + context.setCollation(TiDB::ITiDBCollator::UTF8MB4_GENERAL_CI); + context.context->setSetting("group_by_collation_sensitive", Field(static_cast(1))); + context.context->setSetting("max_block_size", Field(static_cast(rows))); + context.context->setSetting("group_by_two_level_threshold_bytes", Field(static_cast(0))); + + const String tbl_name = "string_collation_key_cache_general_ci_semantics"; + std::vector keys(rows); + for (size_t i = 0; i < rows; ++i) + { + switch (i % 4) + { + case 0: + keys[i] = "AAA"; + break; + case 1: + keys[i] = "aaa"; + break; + case 2: + keys[i] = "BBB"; + break; + default: + keys[i] = "bbb"; + break; + } + } + context.addMockTable({db_name, tbl_name}, table_column_infos, {toVec("key", keys)}); + + auto request + = buildDAGRequest({db_name, tbl_name}, {Count(lit(Field(static_cast(1))))}, {col("key")}, {"count(1)"}); + executeAndAssertColumnsEqual(request, {toVec("count(1)", ColumnWithUInt64{rows / 2, rows / 2})}); +} +CATCH + +TEST_F(AggExecutorTestRunner, StringCollationKeyCacheGeneralCILowAndHighNDV) +try +{ + const String db_name = "test_db"; + constexpr size_t rows = 1024; + DB::MockColumnInfoVec table_column_infos{ + {"key", TiDB::TP::TypeString, false, Poco::Dynamic::Var("utf8mb4_general_ci")}}; + + context.setCollation(TiDB::ITiDBCollator::UTF8MB4_GENERAL_CI); + context.context->setSetting("group_by_collation_sensitive", Field(static_cast(1))); + context.context->setSetting("max_block_size", Field(static_cast(rows))); + context.context->setSetting("group_by_two_level_threshold_bytes", Field(static_cast(0))); + + { + const String tbl_name = "string_collation_key_cache_low_ndv"; + std::vector keys(rows); + for (size_t i = 0; i < rows; ++i) + { + switch (i % 4) + { + case 0: + keys[i] = "AAA"; + break; + case 1: + keys[i] = "BBB"; + break; + case 2: + keys[i] = "CCC"; + break; + default: + keys[i] = "DDD"; + break; + } + } + context.addMockTable({db_name, tbl_name}, table_column_infos, {toVec("key", keys)}); + + auto request = buildDAGRequest( + {db_name, tbl_name}, + {Count(lit(Field(static_cast(1))))}, + {col("key")}, + {"count(1)"}); + executeAndAssertColumnsEqual( + request, + {toVec("count(1)", ColumnWithUInt64{rows / 4, rows / 4, rows / 4, rows / 4})}); + } + + { + const String tbl_name = "string_collation_key_cache_high_ndv"; + std::vector keys(rows); + ColumnWithUInt64 counts(rows, 1); + for (size_t i = 0; i < rows; ++i) + keys[i] = fmt::format("key_{:04}", i); + context.addMockTable({db_name, tbl_name}, table_column_infos, {toVec("key", keys)}); + + auto request = buildDAGRequest( + {db_name, tbl_name}, + {Count(lit(Field(static_cast(1))))}, + {col("key")}, + {"count(1)"}); + executeAndAssertColumnsEqual(request, {toVec("count(1)", counts)}); + } +} +CATCH + +TEST_F(AggExecutorTestRunner, StringCollationKeyCacheWorkerLevelOnceFallback) +try +{ + constexpr size_t rows = 1024; + const String key_column_name = "key"; + auto data_type_string = std::make_shared(); + Block src_header(NamesAndTypes{{key_column_name, data_type_string}}); + + context.context->setSetting("group_by_two_level_threshold", Field(static_cast(0))); + context.context->setSetting("group_by_two_level_threshold_bytes", Field(static_cast(0))); + context.context->setSetting("max_bytes_before_external_group_by", Field(static_cast(0))); + + const auto & settings = context.context->getSettingsRef(); + SpillConfig spill_config( + context.context->getTemporaryPath(), + "StringCollationKeyCacheWorkerLevelOnceFallback", + settings.max_cached_data_bytes_in_spiller, + settings.max_spilled_rows_per_file, + settings.max_spilled_bytes_per_file, + context.context->getFileProvider(), + settings.max_threads, + settings.max_block_size); + std::unordered_map collators{ + {key_column_name, TiDB::ITiDBCollator::getCollator(TiDB::ITiDBCollator::UTF8MB4_GENERAL_CI)}}; + auto params = AggregationInterpreterHelper::buildParams( + *context.context, + src_header, + /*before_agg_streams_size=*/2, + /*agg_streams_size=*/2, + {key_column_name}, + {}, + {}, + collators, + {}, + /*is_final_agg=*/true, + spill_config); + + RegisterOperatorSpillContext register_operator_spill_context; + auto aggregator = std::make_shared( + *params, + "StringCollationKeyCacheWorkerLevelOnceFallback", + /*concurrency=*/2, + register_operator_spill_context, + /*is_auto_pass_through=*/false, + params->use_magic_hash); + + auto execute_block = [&](AggregatedDataVariants & data_variants, size_t distinct_keys, size_t thread_num) { + auto column = data_type_string->createColumn(); + column->reserve(rows); + for (size_t row = 0; row < rows; ++row) + { + const auto key = fmt::format("key_{:04}", row % distinct_keys); + column->insertData(key.data(), key.size()); + } + Block block(ColumnsWithTypeAndName{ + ColumnWithTypeAndName(std::move(column), data_type_string, key_column_name), + }); + Aggregator::AggProcessInfo agg_process_info(aggregator.get()); + agg_process_info.resetBlock(block); + do + { + ASSERT_TRUE(aggregator->executeOnBlock(agg_process_info, data_variants, thread_num)); + } while (!agg_process_info.allBlockDataHandled()); + }; + + AggregatedDataVariants high_ndv_worker; + execute_block(high_ndv_worker, rows, /*thread_num=*/0); + EXPECT_TRUE(high_ndv_worker.string_collation_key_cache_fallback); + + execute_block(high_ndv_worker, /*distinct_keys=*/4, /*thread_num=*/0); + EXPECT_TRUE(high_ndv_worker.string_collation_key_cache_fallback); + + AggregatedDataVariants low_ndv_worker; + execute_block(low_ndv_worker, /*distinct_keys=*/4, /*thread_num=*/1); + EXPECT_FALSE(low_ndv_worker.string_collation_key_cache_fallback); +} +CATCH + TEST_F(AggExecutorTestRunner, FineGrainedShuffle) try { diff --git a/dbms/src/Interpreters/Aggregator.cpp b/dbms/src/Interpreters/Aggregator.cpp index 17436b927c9..e34e93b2e06 100644 --- a/dbms/src/Interpreters/Aggregator.cpp +++ b/dbms/src/Interpreters/Aggregator.cpp @@ -641,6 +641,8 @@ void Aggregator::executeImplInner( { auto * aggregates_pool = result.aggregates_pool; typename Method::State state(agg_process_info.key_columns, key_sizes, collators); + if (!result.string_collation_key_cache_fallback) + state.initCollationSortKeyCache(agg_process_info.end_row - agg_process_info.start_row); // For key_serialized, memory allocation and key serialization will be batch-wise. // For key_string, collation decode will be batch-wise. @@ -673,6 +675,8 @@ void Aggregator::executeImplInner( aggregates_pool, agg_process_info); } + + result.string_collation_key_cache_fallback |= state.hasCollationSortKeyCacheFallback(); } template diff --git a/dbms/src/Interpreters/Aggregator.h b/dbms/src/Interpreters/Aggregator.h index 516190cccc4..415eef0caec 100644 --- a/dbms/src/Interpreters/Aggregator.h +++ b/dbms/src/Interpreters/Aggregator.h @@ -478,6 +478,9 @@ struct AggregatedDataVariants : private boost::noncopyable // This is done both for better performance and because currently, batch and non-batch methods are not compatible. bool batch_get_key_holder = false; + // Once a block exceeds the sort-key cache distinct limit, this worker skips cache attempts for later blocks. + bool string_collation_key_cache_fallback = false; + using AggregationMethod_key8 = AggregationMethodOneNumber; using AggregationMethod_key16 = AggregationMethodOneNumber; using AggregationMethod_key32 = AggregationMethodOneNumber; From 403bfa274a4e23a34321bf91c1cfe977c4db432d Mon Sep 17 00:00:00 2001 From: ChangRui-Ryan Date: Wed, 15 Jul 2026 11:23:07 +0800 Subject: [PATCH 2/3] executor: Fix cached string collation sort key size --- dbms/src/Common/ColumnsHashing.h | 20 +++----- .../tests/gtest_aggregation_executor.cpp | 49 +++++++++++++++++++ 2 files changed, 56 insertions(+), 13 deletions(-) diff --git a/dbms/src/Common/ColumnsHashing.h b/dbms/src/Common/ColumnsHashing.h index d74fb3dc72f..3d804dc4da3 100644 --- a/dbms/src/Common/ColumnsHashing.h +++ b/dbms/src/Common/ColumnsHashing.h @@ -128,8 +128,10 @@ class KeyStringBatchHandlerBase } const auto sort_key_index = cached_sort_keys.size(); - cached_sort_keys.emplace_back(); - auto sort_key = collator.sortKey(raw_key.data, raw_key.size, cached_sort_keys.back()); + auto & cached_sort_key = cached_sort_keys.emplace_back(); + const auto sort_key_size = collator.sortKey(raw_key.data, raw_key.size, cached_sort_key).size; + // Collators may resize the container to their worst-case output size. + cached_sort_key.resize(sort_key_size); typename decltype(raw_to_sort_key_index)::LookupResult it; bool inserted = false; @@ -137,7 +139,7 @@ class KeyStringBatchHandlerBase RUNTIME_CHECK(inserted); it->getMapped() = sort_key_index; - return sort_key; + return {cached_sort_key.data(), cached_sort_key.size()}; } template @@ -177,17 +179,9 @@ class KeyStringBatchHandlerBase { const auto & derived_collator = *static_cast(collator); if (collation_sort_key_cache_active) - prepareNextBatchWithCollator( - chars, - offsets, - cur_batch_size, - derived_collator); + prepareNextBatchWithCollator(chars, offsets, cur_batch_size, derived_collator); else - prepareNextBatchWithCollator( - chars, - offsets, - cur_batch_size, - derived_collator); + prepareNextBatchWithCollator(chars, offsets, cur_batch_size, derived_collator); } else { diff --git a/dbms/src/Flash/tests/gtest_aggregation_executor.cpp b/dbms/src/Flash/tests/gtest_aggregation_executor.cpp index 88cfe9dcc6b..338da574b08 100644 --- a/dbms/src/Flash/tests/gtest_aggregation_executor.cpp +++ b/dbms/src/Flash/tests/gtest_aggregation_executor.cpp @@ -21,6 +21,8 @@ #include #include +#include + namespace DB { namespace FailPoints @@ -1299,6 +1301,53 @@ try } CATCH +TEST_F(AggExecutorTestRunner, StringCollationKeyCacheEffectiveSortKeySize) +try +{ + const String db_name = "test_db"; + constexpr size_t rows = 1024; + + auto execute_case = [&](Int64 collator_id, + const String & collation_name, + const String & tbl_name, + const std::array & raw_keys) { + DB::MockColumnInfoVec table_column_infos{ + {"key", TiDB::TP::TypeString, false, Poco::Dynamic::Var(collation_name)}}; + std::vector keys(rows); + for (size_t i = 0; i < rows; ++i) + keys[i] = raw_keys[i % raw_keys.size()]; + + context.setCollation(collator_id); + context.context->setSetting("group_by_collation_sensitive", Field(static_cast(1))); + context.context->setSetting("max_block_size", Field(static_cast(rows))); + context.addMockTable({db_name, tbl_name}, table_column_infos, {toVec("key", keys)}); + + auto request = buildDAGRequest( + {db_name, tbl_name}, + {Count(lit(Field(static_cast(1))))}, + {col("key")}, + {"count(1)"}); + executeAndAssertColumnsEqual(request, {toVec("count(1)", ColumnWithUInt64{rows / 2, rows / 2})}); + }; + + execute_case( + TiDB::ITiDBCollator::UTF8_UNICODE_CI, + "utf8_unicode_ci", + "string_collation_key_cache_unicode_ci_effective_size", + {"Alpha", "Beta"}); + execute_case( + TiDB::ITiDBCollator::UTF8MB4_GENERAL_CI, + "utf8mb4_general_ci", + "string_collation_key_cache_general_ci_effective_size", + {"Alpha ", "Beta "}); + execute_case( + TiDB::ITiDBCollator::UTF8MB4_0900_AI_CI, + "utf8mb4_0900_ai_ci", + "string_collation_key_cache_0900_ai_ci_effective_size", + {"Alpha", "Alpha "}); +} +CATCH + TEST_F(AggExecutorTestRunner, StringCollationKeyCacheGeneralCILowAndHighNDV) try { From 7e5436739144ec9bbb09837a434d4e22b9ac7fed Mon Sep 17 00:00:00 2001 From: ChangRui-Ryan Date: Thu, 16 Jul 2026 22:53:37 +0800 Subject: [PATCH 3/3] executor: Reserve string collation key cache index --- dbms/src/Common/ColumnsHashing.h | 1 + 1 file changed, 1 insertion(+) diff --git a/dbms/src/Common/ColumnsHashing.h b/dbms/src/Common/ColumnsHashing.h index 3d804dc4da3..637032dfec9 100644 --- a/dbms/src/Common/ColumnsHashing.h +++ b/dbms/src/Common/ColumnsHashing.h @@ -207,6 +207,7 @@ class KeyStringBatchHandlerBase collation_sort_key_cache_active = true; collation_sort_key_cache_distinct_limit = rows / max_collation_sort_key_cache_distinct_ratio; + raw_to_sort_key_index.reserve(collation_sort_key_cache_distinct_limit + 1); cached_sort_keys.reserve(collation_sort_key_cache_distinct_limit + 1); }