From 2c3abcaad12175b2545990e2f37515ba4c270523 Mon Sep 17 00:00:00 2001
From: Maksim Kita <maksim-kita@yandex-team.ru>
Date: Wed, 14 Apr 2021 16:49:38 +0300
Subject: [PATCH] Updated test

---
 src/Dictionaries/DirectDictionary.cpp   | 8 ++++++++
 tests/performance/flat_dictionary.xml   | 6 ++++--
 tests/performance/hashed_dictionary.xml | 3 ++-
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/Dictionaries/DirectDictionary.cpp b/src/Dictionaries/DirectDictionary.cpp
index bacb1a87dc9..ed5da3eead0 100644
--- a/src/Dictionaries/DirectDictionary.cpp
+++ b/src/Dictionaries/DirectDictionary.cpp
@@ -51,6 +51,14 @@ Columns DirectDictionary<dictionary_key_type>::getColumns(
     key_to_fetched_index.reserve(requested_keys.size());
 
     auto fetched_columns_from_storage = request.makeAttributesResultColumns();
+    for (size_t attribute_index = 0; attribute_index < request.attributesSize(); ++attribute_index)
+    {
+        if (!request.shouldFillResultColumnWithIndex(attribute_index))
+            continue;
+
+        auto & fetched_column_from_storage = fetched_columns_from_storage[attribute_index];
+        fetched_column_from_storage->reserve(requested_keys.size());
+    }
 
     size_t fetched_key_index = 0;
 
diff --git a/tests/performance/flat_dictionary.xml b/tests/performance/flat_dictionary.xml
index 56a94358eb9..a80631db541 100644
--- a/tests/performance/flat_dictionary.xml
+++ b/tests/performance/flat_dictionary.xml
@@ -53,7 +53,8 @@
     </substitutions>
 
     <query>
-        SELECT dictGet('default.simple_key_flat_dictionary', {column_name}, rand64() % toUInt64({elements_count}))
+        WITH rand64() % toUInt64({elements_count}) as key
+        SELECT dictGet('default.simple_key_flat_dictionary', {column_name}, key)
         FROM system.numbers
         LIMIT {elements_count}
         FORMAT Null;
@@ -65,7 +66,8 @@
     </query>
 
     <query>
-        SELECT dictHas('default.simple_key_flat_dictionary', rand64() % toUInt64(75000000))
+        WITH rand64() % toUInt64(75000000) as key
+        SELECT dictHas('default.simple_key_flat_dictionary', key)
         FROM system.numbers
         LIMIT 75000000
         FORMAT Null;
diff --git a/tests/performance/hashed_dictionary.xml b/tests/performance/hashed_dictionary.xml
index cd19ba035e5..5cbe1caeb23 100644
--- a/tests/performance/hashed_dictionary.xml
+++ b/tests/performance/hashed_dictionary.xml
@@ -95,7 +95,8 @@
         FORMAT Null;
     </query>
     <query>
-        SELECT * FROM default.simple_key_hashed_dictionary;
+        SELECT * FROM default.simple_key_hashed_dictionary
+        FORMAT Null;
     </query>
     <query>
         WITH rand64() % toUInt64({elements_count}) as key