Process streams of partition tables one by one in MultiplexInputStream

dbms/src/DataStreams/MultiplexInputStream.h

@@ -38,38 +38,25 @@ class MultiPartitionStreamPool
         if (cur_streams.empty())
             return;
         std::unique_lock lk(mu);
-        streams_queue_by_partition.push_back(std::make_shared<std::queue<std::shared_ptr<IBlockInputStream>>>());
-        for (const auto & stream : cur_streams)
-            streams_queue_by_partition.back()->push(stream);
         added_streams.insert(added_streams.end(), cur_streams.begin(), cur_streams.end());
     }
 
-    std::shared_ptr<IBlockInputStream> pickOne()
+    BlockInputStreamPtr pickOne()
     {
         std::unique_lock lk(mu);
-        if (streams_queue_by_partition.empty())
+        if (added_streams.empty())
             return nullptr;
-        if (streams_queue_id >= static_cast<int>(streams_queue_by_partition.size()))
-            streams_queue_id = 0;
-
-        auto & q = *streams_queue_by_partition[streams_queue_id];
-        std::shared_ptr<IBlockInputStream> ret = nullptr;
-        assert(!q.empty());
-        ret = q.front();
-        q.pop();
-        if (q.empty())
-            streams_queue_id = removeQueue(streams_queue_id);
-        else
-            streams_queue_id = nextQueueId(streams_queue_id);
+
+        auto ret = added_streams.front();
+        added_streams.pop_front();
         return ret;
     }
 
-    int exportAddedStreams(BlockInputStreams & ret_streams)
+    void exportAddedStreams(BlockInputStreams & ret_streams)
     {
         std::unique_lock lk(mu);
         for (auto & stream : added_streams)
             ret_streams.push_back(stream);
-        return added_streams.size();
     }
 
     int addedStreamsCnt()
@@ -79,40 +66,7 @@ class MultiPartitionStreamPool
     }
 
 private:
-    int removeQueue(int queue_id)
-    {
-        streams_queue_by_partition[queue_id] = nullptr;
-        if (queue_id != static_cast<int>(streams_queue_by_partition.size()) - 1)
-        {
-            swap(streams_queue_by_partition[queue_id], streams_queue_by_partition.back());
-            streams_queue_by_partition.pop_back();
-            return queue_id;
-        }
-        else
-        {
-            streams_queue_by_partition.pop_back();
-            return 0;
-        }
-    }
-
-    int nextQueueId(int queue_id) const
-    {
-        if (queue_id + 1 < static_cast<int>(streams_queue_by_partition.size()))
-            return queue_id + 1;
-        else
-            return 0;
-    }
-
-    static void swap(
-        std::shared_ptr<std::queue<std::shared_ptr<IBlockInputStream>>> & a,
-        std::shared_ptr<std::queue<std::shared_ptr<IBlockInputStream>>> & b)
-    {
-        a.swap(b);
-    }
-
-    std::vector<std::shared_ptr<std::queue<std::shared_ptr<IBlockInputStream>>>> streams_queue_by_partition;
-    std::vector<std::shared_ptr<IBlockInputStream>> added_streams;
-    int streams_queue_id = 0;
+    std::deque<BlockInputStreamPtr> added_streams;
     std::mutex mu;
 };