support broadcast join

dbms/src/DataStreams/CoprocessorBlockInputStream.h

@@ -29,10 +29,9 @@ class CoprocessorBlockInputStream : public IProfilingBlockInputStream
     }
 
 public:
-    CoprocessorBlockInputStream(pingcap::kv::Cluster * cluster_, const pingcap::coprocessor::Request & req_, const DAGSchema & schema_,
-        int concurrency, pingcap::kv::StoreType store_type)
-        : req(req_),
-          resp_iter(pingcap::coprocessor::Client::send(cluster_, &req, concurrency, store_type)),
+    CoprocessorBlockInputStream(
+        pingcap::kv::Cluster * cluster_, std::vector<pingcap::coprocessor::copTask> tasks, const DAGSchema & schema_, int concurrency)
+        : resp_iter(std::move(tasks), cluster_, concurrency, &Logger::get("pingcap/coprocessor")),
           schema(schema_),
           sample_block(getSampleBlock()),
           log(&Logger::get("pingcap/coprocessor"))
@@ -52,8 +51,7 @@ class CoprocessorBlockInputStream : public IProfilingBlockInputStream
             if (!has_next)
                 return {};
         }
-
-        auto chunk = std::move(chunk_queue.front());
+        auto chunk = chunk_queue.front();
         chunk_queue.pop();
         switch (resp->encode_type())
         {
@@ -101,7 +99,6 @@ class CoprocessorBlockInputStream : public IProfilingBlockInputStream
         return true;
     }
 
-    pingcap::coprocessor::Request req;
     pingcap::coprocessor::ResponseIter resp_iter;
     DAGSchema schema;
 

dbms/src/DataStreams/CreatingSetsBlockInputStream.cpp

@@ -1,9 +1,10 @@
-#include <Interpreters/Set.h>
-#include <Interpreters/Join.h>
-#include <DataStreams/materializeBlock.h>
-#include <DataStreams/IBlockOutputStream.h>
 #include <DataStreams/CreatingSetsBlockInputStream.h>
+#include <DataStreams/IBlockOutputStream.h>
+#include <DataStreams/materializeBlock.h>
+#include <Interpreters/Join.h>
+#include <Interpreters/Set.h>
 #include <Storages/IStorage.h>
+
 #include <iomanip>
 
 
@@ -12,25 +13,38 @@ namespace DB
 
 namespace ErrorCodes
 {
-    extern const int SET_SIZE_LIMIT_EXCEEDED;
+extern const int SET_SIZE_LIMIT_EXCEEDED;
 }
 
+CreatingSetsBlockInputStream::CreatingSetsBlockInputStream(const BlockInputStreamPtr & input,
+    std::vector<SubqueriesForSets> && subqueries_for_sets_list_,
+    const SizeLimits & network_transfer_limits)
+    : subqueries_for_sets_list(std::move(subqueries_for_sets_list_)), network_transfer_limits(network_transfer_limits)
+{
+    init(input);
+}
 
 CreatingSetsBlockInputStream::CreatingSetsBlockInputStream(
-    const BlockInputStreamPtr & input,
-    const SubqueriesForSets & subqueries_for_sets_,
-    const SizeLimits & network_transfer_limits)
-    : subqueries_for_sets(subqueries_for_sets_),
-    network_transfer_limits(network_transfer_limits)
+    const BlockInputStreamPtr & input, const SubqueriesForSets & subqueries_for_sets, const SizeLimits & network_transfer_limits)
+    : network_transfer_limits(network_transfer_limits)
 {
-    for (auto & elem : subqueries_for_sets)
+    subqueries_for_sets_list.push_back(subqueries_for_sets);
+    init(input);
+}
+
+void CreatingSetsBlockInputStream::init(const BlockInputStreamPtr & input)
+{
+    for (auto & subqueries_for_sets : subqueries_for_sets_list)
     {
-        if (elem.second.source)
+        for (auto & elem : subqueries_for_sets)
         {
-            children.push_back(elem.second.source);
+            if (elem.second.source)
+            {
+                children.push_back(elem.second.source);
 
-            if (elem.second.set)
-                elem.second.set->setHeader(elem.second.source->getHeader());
+                if (elem.second.set)
+                    elem.second.set->setHeader(elem.second.source->getHeader());
+            }
         }
     }
 
@@ -51,10 +65,7 @@ Block CreatingSetsBlockInputStream::readImpl()
 }
 
 
-void CreatingSetsBlockInputStream::readPrefixImpl()
-{
-    createAll();
-}
+void CreatingSetsBlockInputStream::readPrefixImpl() { createAll(); }
 
 
 Block CreatingSetsBlockInputStream::getTotals()
@@ -72,27 +83,29 @@ void CreatingSetsBlockInputStream::createAll()
 {
     if (!created)
     {
-        for (auto & elem : subqueries_for_sets)
+        for (auto & subqueries_for_sets : subqueries_for_sets_list)
         {
-            if (elem.second.source) /// There could be prepared in advance Set/Join - no source is specified for them.
+            for (auto & elem : subqueries_for_sets)
             {
-                if (isCancelledOrThrowIfKilled())
-                    return;
+                if (elem.second.source) /// There could be prepared in advance Set/Join - no source is specified for them.
+                {
+                    if (isCancelledOrThrowIfKilled())
+                        return;
 
-                createOne(elem.second);
+                    createOne(elem.second);
+                }
             }
         }
 
         created = true;
     }
 }
 
-
 void CreatingSetsBlockInputStream::createOne(SubqueryForSet & subquery)
 {
-    LOG_TRACE(log, (subquery.set ? "Creating set. " : "")
-        << (subquery.join ? "Creating join. " : "")
-        << (subquery.table ? "Filling temporary table. " : ""));
+    LOG_TRACE(log,
+        (subquery.set ? "Creating set. " : "") << (subquery.join ? "Creating join. " : "")
+                                               << (subquery.table ? "Filling temporary table. " : ""));
     Stopwatch watch;
 
     BlockOutputStreamPtr table_out;
@@ -137,7 +150,8 @@ void CreatingSetsBlockInputStream::createOne(SubqueryForSet & subquery)
             rows_to_transfer += block.rows();
             bytes_to_transfer += block.bytes();
 
-            if (!network_transfer_limits.check(rows_to_transfer, bytes_to_transfer, "IN/JOIN external table", ErrorCodes::SET_SIZE_LIMIT_EXCEEDED))
+            if (!network_transfer_limits.check(
+                    rows_to_transfer, bytes_to_transfer, "IN/JOIN external table", ErrorCodes::SET_SIZE_LIMIT_EXCEEDED))
                 done_with_table = true;
         }
 
@@ -188,4 +202,4 @@ void CreatingSetsBlockInputStream::createOne(SubqueryForSet & subquery)
     }
 }
 
-}
+} // namespace DB

dbms/src/DataStreams/CreatingSetsBlockInputStream.h

@@ -1,11 +1,14 @@
 #pragma once
 
-#include <Poco/Logger.h>
 #include <DataStreams/IProfilingBlockInputStream.h>
-#include <Interpreters/ExpressionAnalyzer.h>    /// SubqueriesForSets
+#include <Interpreters/ExpressionAnalyzer.h> /// SubqueriesForSets
+#include <Poco/Logger.h>
 
 
-namespace Poco { class Logger; }
+namespace Poco
+{
+class Logger;
+}
 
 namespace DB
 {
@@ -18,8 +21,10 @@ class CreatingSetsBlockInputStream : public IProfilingBlockInputStream
 {
 public:
     CreatingSetsBlockInputStream(
-        const BlockInputStreamPtr & input,
-        const SubqueriesForSets & subqueries_for_sets_,
+        const BlockInputStreamPtr & input, const SubqueriesForSets & subqueries_for_sets_, const SizeLimits & network_transfer_limits);
+
+    CreatingSetsBlockInputStream(const BlockInputStreamPtr & input,
+        std::vector<SubqueriesForSets> && subqueries_for_sets_list_,
         const SizeLimits & network_transfer_limits);
 
     String getName() const override { return "CreatingSets"; }
@@ -34,7 +39,9 @@ class CreatingSetsBlockInputStream : public IProfilingBlockInputStream
     void readPrefixImpl() override;
 
 private:
-    SubqueriesForSets subqueries_for_sets;
+    void init(const BlockInputStreamPtr & input);
+
+    std::vector<SubqueriesForSets> subqueries_for_sets_list;
     bool created = false;
 
     SizeLimits network_transfer_limits;
@@ -49,4 +56,4 @@ class CreatingSetsBlockInputStream : public IProfilingBlockInputStream
     void createOne(SubqueryForSet & subquery);
 };
 
-}
+} // namespace DB

dbms/src/DataStreams/IProfilingBlockInputStream.cpp

@@ -103,25 +103,29 @@ Block IProfilingBlockInputStream::read(FilterPtr & res_filter, bool return_filte
 
 void IProfilingBlockInputStream::readPrefix()
 {
+    auto start_time = info.total_stopwatch.elapsed();
     readPrefixImpl();
 
     forEachChild([&] (IBlockInputStream & child)
     {
         child.readPrefix();
         return false;
     });
+    info.updateExecutionTime(info.total_stopwatch.elapsed() - start_time);
 }
 
 
 void IProfilingBlockInputStream::readSuffix()
 {
+    auto start_time = info.total_stopwatch.elapsed();
     forEachChild([&] (IBlockInputStream & child)
     {
         child.readSuffix();
         return false;
     });
 
     readSuffixImpl();
+    info.updateExecutionTime(info.total_stopwatch.elapsed() - start_time);
 }
 
 

dbms/src/DataStreams/NativeBlockInputStream.cpp

@@ -22,6 +22,10 @@ namespace ErrorCodes
     extern const int NOT_IMPLEMENTED;
 }
 
+NativeBlockInputStream::NativeBlockInputStream(ReadBuffer & istr_, UInt64 server_revision_, std::vector<String> && output_names_)
+    : istr(istr_), server_revision(server_revision_), output_names(std::move(output_names_))
+{
+}
 
 NativeBlockInputStream::NativeBlockInputStream(ReadBuffer & istr_, UInt64 server_revision_)
     : istr(istr_), server_revision(server_revision_)
@@ -109,6 +113,11 @@ Block NativeBlockInputStream::readImpl()
         rows = index_block_it->num_rows;
     }
 
+    if (output_names.size() > 0 && output_names.size() != columns)
+        throw Exception("NativeBlockInputStream with explicity output name, but the block column size "
+                        "is not equal to the size of output names", ErrorCodes::LOGICAL_ERROR);
+    bool explicit_output_name = output_names.size() > 0;
+
     for (size_t i = 0; i < columns; ++i)
     {
         if (use_index)
@@ -121,6 +130,8 @@ Block NativeBlockInputStream::readImpl()
 
         /// Name
         readBinary(column.name, istr);
+        if (explicit_output_name)
+            column.name = output_names[i];
 
         /// Type
         String type_name;

dbms/src/DataStreams/NativeBlockInputStream.h

@@ -60,6 +60,8 @@ struct IndexForNativeFormat
 class NativeBlockInputStream : public IProfilingBlockInputStream
 {
 public:
+    /// provide output column names explicitly
+    NativeBlockInputStream(ReadBuffer & istr_, UInt64 server_revision_, std::vector<String> && output_names_);
     /// If a non-zero server_revision is specified, additional block information may be expected and read.
     NativeBlockInputStream(ReadBuffer & istr_, UInt64 server_revision_);
 
@@ -96,6 +98,8 @@ class NativeBlockInputStream : public IProfilingBlockInputStream
 
     PODArray<double> avg_value_size_hints;
 
+    std::vector<String> output_names;
+
     void updateAvgValueSizeHints(const Block & block);
 };
 

dbms/src/Flash/Coprocessor/CHBlockChunkCodec.cpp

@@ -69,11 +69,14 @@ std::unique_ptr<ChunkCodecStream> CHBlockChunkCodec::newCodecStream(const std::v
     return std::make_unique<CHBlockChunkCodecStream>(field_types);
 }
 
-Block CHBlockChunkCodec::decode(const tipb::Chunk & chunk, const DAGSchema &)
+Block CHBlockChunkCodec::decode(const tipb::Chunk & chunk, const DAGSchema & schema)
 {
     const String & row_data = chunk.rows_data();
     ReadBufferFromString read_buffer(row_data);
-    NativeBlockInputStream block_in(read_buffer, 0);
+    std::vector<String> output_names;
+    for (const auto & c : schema)
+        output_names.push_back(c.first);
+    NativeBlockInputStream block_in(read_buffer, 0, std::move(output_names));
     return block_in.read();
 }