Merge pull request #18802 from ClickHouse/backport/20.12/18776

Backport #18776 to 20.12: Fix the issue with async Distributed INSERTs and network_compression_method
ClickHouse · Jan 7, 2021 · d6bd1bd · d6bd1bd
2 parents 461fb29 + a55dd0c
commit d6bd1bd
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 5 deletions.
diff --git a/src/Compression/CompressedReadBuffer.h b/src/Compression/CompressedReadBuffer.h
@@ -16,8 +16,8 @@ class CompressedReadBuffer : public CompressedReadBufferBase, public BufferWithO
     bool nextImpl() override;
 
 public:
-    CompressedReadBuffer(ReadBuffer & in_)
-        : CompressedReadBufferBase(&in_), BufferWithOwnMemory<ReadBuffer>(0)
+    CompressedReadBuffer(ReadBuffer & in_, bool allow_different_codecs_ = false)
+        : CompressedReadBufferBase(&in_, allow_different_codecs_), BufferWithOwnMemory<ReadBuffer>(0)
     {
     }
 

diff --git a/src/Compression/CompressionFactory.cpp b/src/Compression/CompressionFactory.cpp
@@ -46,6 +46,9 @@ CompressionCodecPtr CompressionCodecFactory::get(const String & family_name, std
 
 void CompressionCodecFactory::validateCodec(const String & family_name, std::optional<int> level, bool sanity_check) const
 {
+    if (family_name.empty())
+        throw Exception("Compression codec name cannot be empty", ErrorCodes::BAD_ARGUMENTS);
+
     if (level)
     {
         auto literal = std::make_shared<ASTLiteral>(static_cast<UInt64>(*level));

diff --git a/src/Server/TCPHandler.cpp b/src/Server/TCPHandler.cpp
@@ -1101,7 +1101,7 @@ void TCPHandler::receiveUnexpectedData()
     std::shared_ptr<ReadBuffer> maybe_compressed_in;
 
     if (last_block_in.compression == Protocol::Compression::Enable)
-        maybe_compressed_in = std::make_shared<CompressedReadBuffer>(*in);
+        maybe_compressed_in = std::make_shared<CompressedReadBuffer>(*in, /* allow_different_codecs */ true);
     else
         maybe_compressed_in = in;
 
@@ -1118,8 +1118,11 @@ void TCPHandler::initBlockInput()
 {
     if (!state.block_in)
     {
+        /// 'allow_different_codecs' is set to true, because some parts of compressed data can be precompressed in advance
+        /// with another codec that the rest of the data. Example: data sent by Distributed tables.
+
         if (state.compression == Protocol::Compression::Enable)
-            state.maybe_compressed_in = std::make_shared<CompressedReadBuffer>(*in);
+            state.maybe_compressed_in = std::make_shared<CompressedReadBuffer>(*in, /* allow_different_codecs */ true);
         else
             state.maybe_compressed_in = in;
 

diff --git a/src/Storages/Distributed/DistributedBlockOutputStream.cpp b/src/Storages/Distributed/DistributedBlockOutputStream.cpp
@@ -563,6 +563,17 @@ void DistributedBlockOutputStream::writeToLocal(const Block & block, const size_
 
 void DistributedBlockOutputStream::writeToShard(const Block & block, const std::vector<std::string> & dir_names)
 {
+    const auto & settings = context.getSettingsRef();
+
+    std::string compression_method = Poco::toUpper(settings.network_compression_method.toString());
+    std::optional<int> compression_level;
+
+    if (compression_method == "ZSTD")
+        compression_level = settings.network_zstd_compression_level;
+
+    CompressionCodecFactory::instance().validateCodec(compression_method, compression_level, !settings.allow_suspicious_codecs);
+    CompressionCodecPtr compression_codec = CompressionCodecFactory::instance().get(compression_method, compression_level);
+
     /// tmp directory is used to ensure atomicity of transactions
     /// and keep monitor thread out from reading incomplete data
     std::string first_file_tmp_path{};
@@ -588,7 +599,7 @@ void DistributedBlockOutputStream::writeToShard(const Block & block, const std::
         /// Write batch to temporary location
         {
             WriteBufferFromFile out{first_file_tmp_path};
-            CompressedWriteBuffer compress{out};
+            CompressedWriteBuffer compress{out, compression_codec};
             NativeBlockOutputStream stream{compress, DBMS_TCP_PROTOCOL_VERSION, block.cloneEmpty()};
 
             /// Prepare the header.

diff --git a/tests/queries/0_stateless/01640_distributed_async_insert_compression.reference b/tests/queries/0_stateless/01640_distributed_async_insert_compression.reference
@@ -0,0 +1,2 @@
+256
+512
diff --git a/tests/queries/0_stateless/01640_distributed_async_insert_compression.sql b/tests/queries/0_stateless/01640_distributed_async_insert_compression.sql
@@ -0,0 +1,16 @@
+DROP TABLE IF EXISTS local;
+DROP TABLE IF EXISTS distributed;
+
+CREATE TABLE local (x UInt8) ENGINE = Memory;
+CREATE TABLE distributed AS local ENGINE = Distributed(test_cluster_two_shards, currentDatabase(), local, x);
+
+SET insert_distributed_sync = 0, network_compression_method = 'zstd';
+
+INSERT INTO distributed SELECT number FROM numbers(256);
+SYSTEM FLUSH DISTRIBUTED distributed;
+
+SELECT count() FROM local;
+SELECT count() FROM distributed;
+
+DROP TABLE local;
+DROP TABLE distributed;