From 8ad89a82d4522ab84d32068ef3d8d9a063165d3b Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Mon, 13 Apr 2020 02:08:32 +0300
Subject: [PATCH] Support Nested types in "generateRandom"

---
 src/Storages/StorageGenerateRandom.cpp        | 28 +++++++++++++++----
 .../01128_generate_random_nested.reference    |  2 ++
 .../01128_generate_random_nested.sql          |  8 ++++++
 3 files changed, 32 insertions(+), 6 deletions(-)
 create mode 100644 tests/queries/0_stateless/01128_generate_random_nested.reference
 create mode 100644 tests/queries/0_stateless/01128_generate_random_nested.sql

diff --git a/src/Storages/StorageGenerateRandom.cpp b/src/Storages/StorageGenerateRandom.cpp
index d0772254045..70b84c076b7 100644
--- a/src/Storages/StorageGenerateRandom.cpp
+++ b/src/Storages/StorageGenerateRandom.cpp
@@ -14,6 +14,7 @@
 #include <DataTypes/DataTypeArray.h>
 #include <DataTypes/DataTypeString.h>
 #include <DataTypes/DataTypeFixedString.h>
+#include <DataTypes/NestedUtils.h>
 #include <Columns/ColumnArray.h>
 #include <Columns/ColumnFixedString.h>
 #include <Columns/ColumnString.h>
@@ -57,7 +58,12 @@ void fillBufferWithRandomData(char * __restrict data, size_t size, pcg64 & rng)
 
 
 ColumnPtr fillColumnWithRandomData(
-    const DataTypePtr type, UInt64 limit, UInt64 max_array_length, UInt64 max_string_length, pcg64 & rng, const Context & context)
+    const DataTypePtr type,
+    UInt64 limit,
+    UInt64 max_array_length,
+    UInt64 max_string_length,
+    pcg64 & rng,
+    const Context & context)
 {
     TypeIndex idx = type->getTypeId();
 
@@ -340,14 +346,24 @@ public:
 protected:
     Chunk generate() override
     {
+        /// To support Nested types, we will collect them to single Array of Tuple.
+        auto names_and_types = Nested::collect(block_header.getNamesAndTypesList());
+
         Columns columns;
-        columns.reserve(block_header.columns());
-        DataTypes types = block_header.getDataTypes();
+        columns.reserve(names_and_types.size());
 
-        for (const auto & type : types)
-            columns.emplace_back(fillColumnWithRandomData(type, block_size, max_array_length, max_string_length, rng, context));
+        Block compact_block;
+        for (const auto & elem : names_and_types)
+        {
+            compact_block.insert(
+            {
+                fillColumnWithRandomData(elem.type, block_size, max_array_length, max_string_length, rng, context),
+                elem.type,
+                elem.name
+            });
+        }
 
-        return {std::move(columns), block_size};
+        return {Nested::flatten(compact_block).getColumns(), block_size};
     }
 
 private:
diff --git a/tests/queries/0_stateless/01128_generate_random_nested.reference b/tests/queries/0_stateless/01128_generate_random_nested.reference
new file mode 100644
index 00000000000..d9d2b251702
--- /dev/null
+++ b/tests/queries/0_stateless/01128_generate_random_nested.reference
@@ -0,0 +1,2 @@
+100	12366141706519416319
+109	2990700419202507835
diff --git a/tests/queries/0_stateless/01128_generate_random_nested.sql b/tests/queries/0_stateless/01128_generate_random_nested.sql
new file mode 100644
index 00000000000..2af52e69893
--- /dev/null
+++ b/tests/queries/0_stateless/01128_generate_random_nested.sql
@@ -0,0 +1,8 @@
+DROP TABLE IF EXISTS mass_table_312;
+CREATE TABLE mass_table_312 (d Date DEFAULT '2000-01-01', x UInt64, n Nested(a String, b String)) ENGINE = MergeTree(d, x, 1);
+INSERT INTO mass_table_312 SELECT * FROM generateRandom('`d` Date,`x` UInt64,`n.a` Array(String),`n.b` Array(String)', 1, 10, 2) LIMIT 100;
+
+SELECT count(), sum(cityHash64(*)) FROM mass_table_312;
+SELECT count(), sum(cityHash64(*)) FROM mass_table_312 ARRAY JOIN n;
+
+DROP TABLE mass_table_312;