Merging parallel format parsing

dbms/programs/client/Client.cpp

@@ -1112,7 +1112,14 @@ class Client : public Poco::Util::Application
             /// Check if server send Exception packet
             auto packet_type = connection->checkPacket();
             if (packet_type && *packet_type == Protocol::Server::Exception)
+            {
+                /*
+                 * We're exiting with error, so it makes sense to kill the
+                 * input stream without waiting for it to complete.
+                 */
+                async_block_input->cancel(true);
                 return;
+            }
 
             connection->sendData(block);
             processed_rows += block.rows();

dbms/src/Core/Settings.h

@@ -111,6 +111,9 @@ struct Settings : public SettingsCollection<Settings>
     M(SettingBool, distributed_group_by_no_merge, false, "Do not merge aggregation states from different servers for distributed query processing - in case it is for certain that there are different keys on different shards.", 0) \
     M(SettingBool, optimize_skip_unused_shards, false, "Assumes that data is distributed by sharding_key. Optimization to skip unused shards if SELECT query filters by sharding_key.", 0) \
     \
+    M(SettingBool, input_format_parallel_parsing, true, "Enable parallel parsing for some data formats.", 0) \
+    M(SettingUInt64, min_chunk_bytes_for_parallel_parsing, (1024 * 1024), "The minimum chunk size in bytes, which each thread will parse in parallel.", 0) \
+    \
     M(SettingUInt64, merge_tree_min_rows_for_concurrent_read, (20 * 8192), "If at least as many lines are read from one file, the reading can be parallelized.", 0) \
     M(SettingUInt64, merge_tree_min_bytes_for_concurrent_read, (24 * 10 * 1024 * 1024), "If at least as many bytes are read from one file, the reading can be parallelized.", 0) \
     M(SettingUInt64, merge_tree_min_rows_for_seek, 0, "You can skip reading more than that number of rows at the price of one seek per file.", 0) \

dbms/src/DataStreams/ParallelParsingBlockInputStream.cpp

@@ -0,0 +1,203 @@
+#include <DataStreams/ParallelParsingBlockInputStream.h>
+#include "ParallelParsingBlockInputStream.h"
+
+namespace DB
+{
+
+void ParallelParsingBlockInputStream::segmentatorThreadFunction()
+{
+    setThreadName("Segmentator");
+    try
+    {
+        while (!finished)
+        {
+            const auto current_unit_number = segmentator_ticket_number % processing_units.size();
+            auto & unit = processing_units[current_unit_number];
+
+            {
+                std::unique_lock lock(mutex);
+                segmentator_condvar.wait(lock,
+                    [&]{ return unit.status == READY_TO_INSERT || finished; });
+            }
+
+            if (finished)
+            {
+                break;
+            }
+
+            assert(unit.status == READY_TO_INSERT);
+
+            // Segmentating the original input.
+            unit.segment.resize(0);
+
+            const bool have_more_data = file_segmentation_engine(original_buffer,
+                unit.segment, min_chunk_bytes);
+
+            unit.is_last = !have_more_data;
+            unit.status = READY_TO_PARSE;
+            scheduleParserThreadForUnitWithNumber(current_unit_number);
+            ++segmentator_ticket_number;
+
+            if (!have_more_data)
+            {
+                break;
+            }
+        }
+    }
+    catch (...)
+    {
+        onBackgroundException();
+    }
+}
+
+void ParallelParsingBlockInputStream::parserThreadFunction(size_t current_unit_number)
+{
+    try
+    {
+        setThreadName("ChunkParser");
+
+        auto & unit = processing_units[current_unit_number];
+
+        /*
+         * This is kind of suspicious -- the input_process_creator contract with
+         * respect to multithreaded use is not clear, but we hope that it is
+         * just a 'normal' factory class that doesn't have any state, and so we
+         * can use it from multiple threads simultaneously.
+         */
+        ReadBuffer read_buffer(unit.segment.data(), unit.segment.size(), 0);
+        auto parser = std::make_unique<InputStreamFromInputFormat>(
+            input_processor_creator(read_buffer, header, context,
+                row_input_format_params, format_settings));
+
+        unit.block_ext.block.clear();
+        unit.block_ext.block_missing_values.clear();
+
+        // We don't know how many blocks will be. So we have to read them all
+        // until an empty block occured.
+        Block block;
+        while (!finished && (block = parser->read()) != Block())
+        {
+            unit.block_ext.block.emplace_back(block);
+            unit.block_ext.block_missing_values.emplace_back(parser->getMissingValues());
+        }
+
+        // We suppose we will get at least some blocks for a non-empty buffer,
+        // except at the end of file. Also see a matching assert in readImpl().
+        assert(unit.is_last || unit.block_ext.block.size() > 0);
+
+        std::unique_lock lock(mutex);
+        unit.status = READY_TO_READ;
+        reader_condvar.notify_all();
+    }
+    catch (...)
+    {
+        onBackgroundException();
+    }
+}
+
+void ParallelParsingBlockInputStream::onBackgroundException()
+{
+    tryLogCurrentException(__PRETTY_FUNCTION__);
+
+    std::unique_lock lock(mutex);
+    if (!background_exception)
+    {
+        background_exception = std::current_exception();
+    }
+    finished = true;
+    reader_condvar.notify_all();
+    segmentator_condvar.notify_all();
+}
+
+Block ParallelParsingBlockInputStream::readImpl()
+{
+    if (isCancelledOrThrowIfKilled() || finished)
+    {
+        /**
+          * Check for background exception and rethrow it before we return.
+          */
+        std::unique_lock lock(mutex);
+        if (background_exception)
+        {
+            lock.unlock();
+            cancel(false);
+            std::rethrow_exception(background_exception);
+        }
+
+        return Block{};
+    }
+
+    const auto current_unit_number = reader_ticket_number % processing_units.size();
+    auto & unit = processing_units[current_unit_number];
+
+    if (!next_block_in_current_unit.has_value())
+    {
+        // We have read out all the Blocks from the previous Processing Unit,
+        // wait for the current one to become ready.
+        std::unique_lock lock(mutex);
+        reader_condvar.wait(lock, [&](){ return unit.status == READY_TO_READ || finished; });
+
+        if (finished)
+        {
+            /**
+              * Check for background exception and rethrow it before we return.
+              */
+            if (background_exception)
+            {
+                lock.unlock();
+                cancel(false);
+                std::rethrow_exception(background_exception);
+            }
+
+            return Block{};
+        }
+
+        assert(unit.status == READY_TO_READ);
+        next_block_in_current_unit = 0;
+    }
+
+    if (unit.block_ext.block.size() == 0)
+    {
+        /*
+         * Can we get zero blocks for an entire segment, when the format parser
+         * skips it entire content and does not create any blocks? Probably not,
+         * but if we ever do, we should add a loop around the above if, to skip
+         * these. Also see a matching assert in the parser thread.
+         */
+        assert(unit.is_last);
+        finished = true;
+        return Block{};
+    }
+
+    assert(next_block_in_current_unit.value() < unit.block_ext.block.size());
+
+    Block res = std::move(unit.block_ext.block.at(*next_block_in_current_unit));
+    last_block_missing_values = std::move(unit.block_ext.block_missing_values[*next_block_in_current_unit]);
+
+    next_block_in_current_unit.value() += 1;
+
+    if (*next_block_in_current_unit == unit.block_ext.block.size())
+    {
+        // Finished reading this Processing Unit, move to the next one.
+        next_block_in_current_unit.reset();
+        ++reader_ticket_number;
+
+        if (unit.is_last)
+        {
+            // It it was the last unit, we're finished.
+            finished = true;
+        }
+        else
+        {
+            // Pass the unit back to the segmentator.
+            std::unique_lock lock(mutex);
+            unit.status = READY_TO_INSERT;
+            segmentator_condvar.notify_all();
+        }
+    }
+
+    return res;
+}
+
+
+}