Optimize data size of Broadcast / Passthrough exchange operator

dbms/CMakeLists.txt

@@ -115,6 +115,7 @@ check_then_add_sources_compile_flag (
     src/Columns/ColumnsCommon.cpp
     src/Columns/ColumnVector.cpp
     src/DataTypes/DataTypeString.cpp
+    src/Interpreters/Join.cpp
 )
 
 list (APPEND clickhouse_common_io_sources ${CONFIG_BUILD})

dbms/src/Common/TiFlashMetrics.h

@@ -92,9 +92,16 @@ namespace DB
         F(type_hash_none_compression_local, {"type", "hash_none_compression_local"}), /*the local exchange data size by hash partition with no compression*/ \
         F(type_hash_lz4_compression, {"type", "hash_lz4_compression"}), /*the exchange data size by hash partition with lz4 compression*/ \
         F(type_hash_zstd_compression, {"type", "hash_zstd_compression"}), /*the exchange data size by hash partition with zstd compression*/ \
-        F(type_broadcast_passthrough_original, {"type", "broadcast_passthrough_original"}), /*the original exchange data size by broadcast/passthough*/ \
-        F(type_broadcast_passthrough_none_compression_local, {"type", "broadcast_passthrough_none_compression_local"}), /*the local exchange data size by broadcast/passthough with no compression*/ \
-        F(type_broadcast_passthrough_none_compression_remote, {"type", "broadcast_passthrough_none_compression_remote"}), /*the remote exchange data size by broadcast/passthough with no compression*/ \
+        F(type_broadcast_original, {"type", "broadcast_original"}), \
+        F(type_broadcast_none_compression_local, {"type", "broadcast_none_compression_local"}), \
+        F(type_broadcast_none_compression_remote, {"type", "broadcast_none_compression_remote"}), \
+        F(type_broadcast_lz4_compression, {"type", "broadcast_lz4_compression"}), \
+        F(type_broadcast_zstd_compression, {"type", "broadcast_zstd_compression"}), \
+        F(type_passthrough_original, {"type", "passthrough_original"}), \
+        F(type_passthrough_none_compression_local, {"type", "passthrough_none_compression_local"}), \
+        F(type_passthrough_none_compression_remote, {"type", "passthrough_none_compression_remote"}), \
+        F(type_passthrough_lz4_compression, {"type", "passthrough_lz4_compression"}), \
+        F(type_passthrough_zstd_compression, {"type", "passthrough_zstd_compression"}), \
     ) \
     M(tiflash_schema_version, "Current version of tiflash cached schema", Gauge)                                                          \
     M(tiflash_schema_applying, "Whether the schema is applying or not (holding lock)", Gauge)                                             \

dbms/src/Flash/Coprocessor/CHBlockChunkCodecV1.cpp

@@ -248,6 +248,10 @@ struct CHBlockChunkCodecV1Impl
     {
         return encodeImpl(blocks, compression_method);
     }
+    CHBlockChunkCodecV1::EncodeRes encode(std::vector<Block> && blocks, CompressionMethod compression_method)
+    {
+        return encodeImpl(std::move(blocks), compression_method);
+    }
 
     static const ColumnPtr & toColumnPtr(const Columns & c, size_t index)
     {
@@ -269,6 +273,10 @@ struct CHBlockChunkCodecV1Impl
     {
         return block.getByPosition(index).column;
     }
+    static ColumnPtr toColumnPtr(Block && block, size_t index)
+    {
+        return std::move(block.getByPosition(index).column);
+    }
 
     template <typename ColumnsHolder>
     static size_t getRows(ColumnsHolder && columns_holder)
@@ -349,6 +357,13 @@ struct CHBlockChunkCodecV1Impl
         return encodeColumnImpl(block, ostr_ptr);
     }
     void encodeColumn(const std::vector<Block> & blocks, WriteBuffer * ostr_ptr)
+    {
+        for (auto && block : blocks)
+        {
+            encodeColumnImpl(block, ostr_ptr);
+        }
+    }
+    void encodeColumn(std::vector<Block> && blocks, WriteBuffer * ostr_ptr)
     {
         for (auto && block : blocks)
         {
@@ -495,6 +510,19 @@ CHBlockChunkCodecV1::EncodeRes CHBlockChunkCodecV1::encode(const std::vector<Blo
     return CHBlockChunkCodecV1Impl{*this}.encode(blocks, compression_method);
 }
 
+CHBlockChunkCodecV1::EncodeRes CHBlockChunkCodecV1::encode(std::vector<Block> && blocks, CompressionMethod compression_method, bool check_schema)
+{
+    if (check_schema)
+    {
+        for (auto && block : blocks)
+        {
+            checkSchema(header, block);
+        }
+    }
+
+    return CHBlockChunkCodecV1Impl{*this}.encode(std::move(blocks), compression_method);
+}
+
 static Block decodeCompression(const Block & header, ReadBuffer & istr)
 {
     size_t decoded_rows{};
@@ -504,6 +532,22 @@ static Block decodeCompression(const Block & header, ReadBuffer & istr)
     return decoded_block;
 }
 
+template <typename Buffer>
+extern size_t CompressionEncode(
+    std::string_view source,
+    const CompressionSettings & compression_settings,
+    Buffer & compressed_buffer);
+
+CHBlockChunkCodecV1::EncodeRes CHBlockChunkCodecV1::encode(std::string_view str, CompressionMethod compression_method)
+{
+    assert(compression_method != CompressionMethod::NONE);
+
+    String compressed_buffer;
+    size_t compressed_size = CompressionEncode(str, CompressionSettings(compression_method), compressed_buffer);
+    compressed_buffer.resize(compressed_size);
+    return compressed_buffer;
+}
+
 Block CHBlockChunkCodecV1::decode(const Block & header, std::string_view str)
 {
     assert(!str.empty());

dbms/src/Flash/Coprocessor/CHBlockChunkCodecV1.h

@@ -53,7 +53,9 @@ struct CHBlockChunkCodecV1 : boost::noncopyable
     EncodeRes encode(std::vector<Columns> && columns, CompressionMethod compression_method);
     EncodeRes encode(const Block & block, CompressionMethod compression_method, bool check_schema = true);
     EncodeRes encode(const std::vector<Block> & blocks, CompressionMethod compression_method, bool check_schema = true);
+    EncodeRes encode(std::vector<Block> && blocks, CompressionMethod compression_method, bool check_schema = true);
     //
+    static EncodeRes encode(std::string_view str, CompressionMethod compression_method);
     static Block decode(const Block & header, std::string_view str);
 };
 

dbms/src/Flash/Coprocessor/tests/gtest_block_chunk_codec.cpp

@@ -126,6 +126,41 @@ TEST(CHBlockChunkCodec, ChunkCodecV1)
             auto decoded_block = CHBlockChunkCodecV1::decode(header, str);
             ASSERT_EQ(total_rows, decoded_block.rows());
         }
+        {
+            std::vector<Block> blocks_to_move;
+            blocks_to_move.reserve(blocks.size());
+            for (auto && block : blocks)
+            {
+                blocks_to_move.emplace_back(block);
+            }
+            for (auto && block : blocks_to_move)
+            {
+                for (auto && col : block)
+                {
+                    ASSERT_TRUE(col.column);
+                }
+            }
+            auto codec = CHBlockChunkCodecV1{
+                header,
+            };
+            auto str = codec.encode(std::move(blocks_to_move), mode);
+            for (auto && block : blocks_to_move)
+            {
+                ASSERT_EQ(block.rows(), 0);
+            }
+            ASSERT_FALSE(str.empty());
+            ASSERT_EQ(codec.encoded_rows, total_rows);
+
+            if (mode == CompressionMethod::NONE)
+                ASSERT_EQ(codec.compressed_size, 0);
+            else
+                ASSERT_NE(codec.compressed_size, 0);
+
+            ASSERT_NE(codec.original_size, 0);
+
+            auto decoded_block = CHBlockChunkCodecV1::decode(header, str);
+            ASSERT_EQ(total_rows, decoded_block.rows());
+        }
         {
             auto columns = prepareBlock(rows).getColumns();
             auto codec = CHBlockChunkCodecV1{
@@ -179,5 +214,19 @@ TEST(CHBlockChunkCodec, ChunkCodecV1)
         }
         test_enocde_release_data(std::move(batch_columns), header, total_rows);
     }
+    {
+        auto source_str = CHBlockChunkCodecV1{header}.encode(blocks.front(), CompressionMethod::NONE);
+        ASSERT_FALSE(source_str.empty());
+        ASSERT_EQ(static_cast<CompressionMethodByte>(source_str[0]), CompressionMethodByte::NONE);
+
+        for (auto mode : {CompressionMethod::LZ4, CompressionMethod::ZSTD})
+        {
+            auto compressed_str_a = CHBlockChunkCodecV1::encode({&source_str[1], source_str.size() - 1}, mode);
+            auto compressed_str_b = CHBlockChunkCodecV1{header}.encode(blocks.front(), mode);
+
+            ASSERT_EQ(compressed_str_a, compressed_str_b);
+        }
+    }
 }
+
 } // namespace DB::tests

dbms/src/Flash/Coprocessor/tests/gtest_ti_remote_block_inputstream.cpp

@@ -86,17 +86,48 @@ struct MockWriter
         return summary;
     }
 
-    void broadcastOrPassThroughWrite(Blocks & blocks)
+    void broadcastOrPassThroughWriteV0(Blocks & blocks)
     {
         auto && packet = MPPTunnelSetHelper::ToPacketV0(blocks, result_field_types);
         ++total_packets;
         if (!packet)
             return;
 
-        if (!packet->packet.chunks().empty())
-            total_bytes += packet->packet.ByteSizeLong();
+        total_bytes += packet->packet.ByteSizeLong();
         queue->push(std::move(packet));
     }
+
+    void broadcastWrite(Blocks & blocks)
+    {
+        return broadcastOrPassThroughWriteV0(blocks);
+    }
+    void passThroughWrite(Blocks & blocks)
+    {
+        return broadcastOrPassThroughWriteV0(blocks);
+    }
+    void broadcastOrPassThroughWrite(Blocks & blocks, MPPDataPacketVersion version, CompressionMethod compression_method)
+    {
+        if (version == MPPDataPacketV0)
+            return broadcastOrPassThroughWriteV0(blocks);
+
+        size_t original_size{};
+        auto && packet = MPPTunnelSetHelper::ToPacket(std::move(blocks), version, compression_method, original_size);
+        ++total_packets;
+        if (!packet)
+            return;
+
+        total_bytes += packet->packet.ByteSizeLong();
+        queue->push(std::move(packet));
+    }
+    void broadcastWrite(Blocks & blocks, MPPDataPacketVersion version, CompressionMethod compression_method)
+    {
+        return broadcastOrPassThroughWrite(blocks, version, compression_method);
+    }
+    void passThroughWrite(Blocks & blocks, MPPDataPacketVersion version, CompressionMethod compression_method)
+    {
+        return broadcastOrPassThroughWrite(blocks, version, compression_method);
+    }
+
     void write(tipb::SelectResponse & response)
     {
         if (add_summary)
@@ -124,10 +155,6 @@ struct MockWriter
         write(tmp);
     }
     uint16_t getPartitionNum() const { return 1; }
-    bool isLocal(size_t index) const
-    {
-        return index == 0;
-    }
 
     std::vector<tipb::FieldType> result_field_types;
 
@@ -361,7 +388,10 @@ class TestTiRemoteBlockInputStream : public testing::Test
         auto dag_writer = std::make_shared<BroadcastOrPassThroughWriter<MockWriterPtr>>(
             writer,
             batch_send_min_limit,
-            *dag_context_ptr);
+            *dag_context_ptr,
+            MPPDataPacketVersion::MPPDataPacketV1,
+            tipb::CompressionMode::FAST,
+            tipb::ExchangeType::Broadcast);
 
         // 2. encode all blocks
         for (const auto & block : source_blocks)

dbms/src/Flash/Mpp/BroadcastOrPassThroughWriter.cpp

@@ -14,6 +14,7 @@
 
 #include <Common/TiFlashException.h>
 #include <Flash/Coprocessor/CHBlockChunkCodec.h>
+#include <Flash/Coprocessor/CHBlockChunkCodecV1.h>
 #include <Flash/Coprocessor/DAGContext.h>
 #include <Flash/Mpp/BroadcastOrPassThroughWriter.h>
 #include <Flash/Mpp/MPPTunnelSet.h>
@@ -24,13 +25,41 @@ template <class ExchangeWriterPtr>
 BroadcastOrPassThroughWriter<ExchangeWriterPtr>::BroadcastOrPassThroughWriter(
     ExchangeWriterPtr writer_,
     Int64 batch_send_min_limit_,
-    DAGContext & dag_context_)
+    DAGContext & dag_context_,
+    MPPDataPacketVersion data_codec_version_,
+    tipb::CompressionMode compression_mode_,
+    tipb::ExchangeType exchange_type_)
     : DAGResponseWriter(/*records_per_chunk=*/-1, dag_context_)
     , batch_send_min_limit(batch_send_min_limit_)
     , writer(writer_)
+    , exchange_type(exchange_type_)
+    , data_codec_version(data_codec_version_)
+    , compression_method(ToInternalCompressionMethod(compression_mode_))
 {
     rows_in_blocks = 0;
     RUNTIME_CHECK(dag_context.encode_type == tipb::EncodeType::TypeCHBlock);
+    RUNTIME_CHECK(exchange_type == tipb::ExchangeType::Broadcast || exchange_type == tipb::ExchangeType::PassThrough);
+
+    switch (data_codec_version)
+    {
+    case MPPDataPacketV0:
+        break;
+    case MPPDataPacketV1:
+    default:
+    {
+        // make `batch_send_min_limit` always GT 0
+        if (batch_send_min_limit <= 0)
+        {
+            // set upper limit if not specified
+            batch_send_min_limit = 8 * 1024 /* 8K */;
+        }
+        for (const auto & field_type : dag_context.result_field_types)
+        {
+            expected_types.emplace_back(getDataTypeByFieldTypeForComputingLayer(field_type));
+        }
+        break;
+    }
+    }
 }
 
 template <class ExchangeWriterPtr>
@@ -60,10 +89,20 @@ void BroadcastOrPassThroughWriter<ExchangeWriterPtr>::write(const Block & block)
 template <class ExchangeWriterPtr>
 void BroadcastOrPassThroughWriter<ExchangeWriterPtr>::writeBlocks()
 {
-    if (unlikely(blocks.empty()))
+    if unlikely (blocks.empty())
         return;
 
-    writer->broadcastOrPassThroughWrite(blocks);
+    // check schema
+    if (!expected_types.empty())
+    {
+        for (auto && block : blocks)
+            assertBlockSchema(expected_types, block, "BroadcastOrPassThroughWriter");
+    }
+
+    if (exchange_type == tipb::ExchangeType::Broadcast)
+        writer->broadcastWrite(blocks, data_codec_version, compression_method);
+    else
+        writer->passThroughWrite(blocks, data_codec_version, compression_method);
     blocks.clear();
     rows_in_blocks = 0;
 }

dbms/src/Flash/Mpp/BroadcastOrPassThroughWriter.h

@@ -22,6 +22,8 @@
 namespace DB
 {
 class DAGContext;
+enum class CompressionMethod;
+enum MPPDataPacketVersion : int64_t;
 
 template <class ExchangeWriterPtr>
 class BroadcastOrPassThroughWriter : public DAGResponseWriter
@@ -30,7 +32,10 @@ class BroadcastOrPassThroughWriter : public DAGResponseWriter
     BroadcastOrPassThroughWriter(
         ExchangeWriterPtr writer_,
         Int64 batch_send_min_limit_,
-        DAGContext & dag_context_);
+        DAGContext & dag_context_,
+        MPPDataPacketVersion data_codec_version_,
+        tipb::CompressionMode compression_mode_,
+        tipb::ExchangeType exchange_type_);
     void write(const Block & block) override;
     void flush() override;
 
@@ -42,6 +47,12 @@ class BroadcastOrPassThroughWriter : public DAGResponseWriter
     ExchangeWriterPtr writer;
     std::vector<Block> blocks;
     size_t rows_in_blocks;
+    const tipb::ExchangeType exchange_type;
+
+    // support data compression
+    DataTypes expected_types;
+    MPPDataPacketVersion data_codec_version;
+    CompressionMethod compression_method{};
 };
 
 } // namespace DB