merge output blocks if need in hash join

dbms/src/Core/Block.cpp

@@ -514,6 +514,36 @@ static ReturnType checkBlockStructure(const Block & lhs, const Block & rhs, cons
     return ReturnType(true);
 }
 
+Block mergeBlocks(Blocks && blocks)
+{
+    assert(!blocks.empty());
+    auto & sample_block = blocks[0];
+    size_t result_rows = 0;
+    for (const auto & block : blocks)
+    {
+        result_rows += block.rows();
+    }
+
+    MutableColumns dst_columns(sample_block.columns());
+
+    for (size_t i = 0; i < sample_block.columns(); i++)
+    {
+        dst_columns[i] = (*std::move(sample_block.getByPosition(i).column)).mutate();
+        dst_columns[i]->reserve(result_rows);
+    }
+
+    for (size_t i = 1; i < blocks.size(); ++i)
+    {
+        if (blocks[i].rows() > 0)
+        {
+            for (size_t column = 0; column < blocks[i].columns(); column++)
+            {
+                dst_columns[column]->insertRangeFrom(*blocks[i].getByPosition(column).column, 0, blocks[i].rows());
+            }
+        }
+    }
+    return sample_block.cloneWithColumns(std::move(dst_columns));
+}
 
 bool blocksHaveEqualStructure(const Block & lhs, const Block & rhs)
 {

dbms/src/Core/Block.h

@@ -149,6 +149,7 @@ class Block
       */
     void updateHash(SipHash & hash) const;
 
+
 private:
     void eraseImpl(size_t position);
     void initializeIndexByName();
@@ -157,6 +158,7 @@ class Block
 using Blocks = std::vector<Block>;
 using BlocksList = std::list<Block>;
 
+Block mergeBlocks(Blocks && blocks);
 
 /// Compare number of columns, data types, column types, column names, and values of constant columns.
 bool blocksHaveEqualStructure(const Block & lhs, const Block & rhs);

dbms/src/DataStreams/HashJoinProbeBlockInputStream.cpp

@@ -25,6 +25,7 @@ HashJoinProbeBlockInputStream::HashJoinProbeBlockInputStream(
     : log(Logger::get(req_id))
     , join(join_)
     , probe_process_info(max_block_size)
+    , join_finished(false)
 {
     children.push_back(input);
 
@@ -67,17 +68,81 @@ Block HashJoinProbeBlockInputStream::getHeader() const
 
 Block HashJoinProbeBlockInputStream::readImpl()
 {
-    if (probe_process_info.all_rows_joined_finish)
+    // if join finished, return {}
+    if (join_finished)
+    {
+        return Block{};
+    }
+
+    result_blocks.clear();
+    size_t output_rows = 0;
+
+    // if over_limit_block is not null, we need to push it into result_blocks first.
+    if (over_limit_block)
+    {
+        output_rows += over_limit_block.rows();
+        result_blocks.push_back(std::move(over_limit_block));
+        over_limit_block = Block{};
+    }
+
+    while (output_rows <= probe_process_info.max_block_size)
+    {
+        Block result_block = getOutputBlock(probe_process_info);
+
+        if (!result_block)
+        {
+            // if result blocks is not empty, merge and return them, then mark join finished.
+            if (!result_blocks.empty())
+            {
+                join_finished = true;
+                return mergeResultBlocks(std::move(result_blocks));
+            }
+            // if result blocks is empty, return result block directly.
+            return result_block;
+        }
+        size_t current_rows = result_block.rows();
+
+        if (!output_rows || output_rows + current_rows <= probe_process_info.max_block_size)
+        {
+            result_blocks.push_back(result_block);
+        }
+        else
+        {
+            // if output_rows + current_rows > max block size, put the current result block into over_limit_block and handle it in next read.
+            over_limit_block = result_block;
+        }
+        output_rows += current_rows;
+    }
+
+    return mergeResultBlocks(std::move(result_blocks));
+}
+
+Block HashJoinProbeBlockInputStream::getOutputBlock(ProbeProcessInfo & probe_process_info_) const
+{
+    if (probe_process_info_.all_rows_joined_finish)
     {
         Block block = children.back()->read();
         if (!block)
+        {
             return block;
+        }
         join->checkTypes(block);
-        probe_process_info.resetBlock(std::move(block));
+        probe_process_info_.resetBlock(std::move(block));
     }
 
-    return join->joinBlock(probe_process_info);
+    return join->joinBlock(probe_process_info_);
 }
 
+Block HashJoinProbeBlockInputStream::mergeResultBlocks(Blocks && result_blocks)
+{
+    if (result_blocks.size() == 1)
+    {
+        return result_blocks[0];
+    }
+    else
+    {
+        return mergeBlocks(std::move(result_blocks));
+    }
+}
 
 } // namespace DB

dbms/src/DataStreams/HashJoinProbeBlockInputStream.h

@@ -44,6 +44,8 @@ class HashJoinProbeBlockInputStream : public IProfilingBlockInputStream
     String getName() const override { return name; }
     Block getTotals() override;
     Block getHeader() const override;
+    Block getOutputBlock(ProbeProcessInfo & probe_process_info_) const;
+    static Block mergeResultBlocks(Blocks && result_blocks);
 
 protected:
     Block readImpl() override;
@@ -52,6 +54,9 @@ class HashJoinProbeBlockInputStream : public IProfilingBlockInputStream
     const LoggerPtr log;
     JoinPtr join;
     ProbeProcessInfo probe_process_info;
+    Blocks result_blocks;
+    Block over_limit_block;
+    bool join_finished;
 };
 
 } // namespace DB

dbms/src/Flash/tests/gtest_join_executor.cpp

@@ -700,7 +700,7 @@ CATCH
 TEST_F(JoinExecutorTestRunner, SplitJoinResult)
 try
 {
-    context.addMockTable("split_test", "t1", {{"a", TiDB::TP::TypeLong}}, {toVec<Int32>("a", {1, 1, 1, 1, 1, 1, 1, 1, 1, 1})});
+    context.addMockTable("split_test", "t1", {{"a", TiDB::TP::TypeLong}, {"b", TiDB::TP::TypeLong}}, {toVec<Int32>("a", {1, 1, 1, 1, 1, 1, 1, 1, 1, 1}), toVec<Int32>("b", {1, 1, 3, 3, 1, 1, 3, 3, 1, 3})});
     context.addMockTable("split_test", "t2", {{"a", TiDB::TP::TypeLong}}, {toVec<Int32>("a", {1, 1, 1, 1, 1})});
 
     auto request = context
@@ -720,6 +720,25 @@ try
             ASSERT_EQ(expect[i][j], blocks[j].rows());
         }
     }
+
+    // with other condition
+    const auto cond = gt(col("b"), lit(Field(static_cast<Int64>(2))));
+    request = context
+                  .scan("split_test", "t1")
+                  .join(context.scan("split_test", "t2"), tipb::JoinType::TypeInnerJoin, {col("a")}, {}, {}, {cond}, {})
+
+                  .build(context);
+    expect = {{5, 5, 5, 5, 5}, {5, 5, 5, 5, 5}, {5, 5, 5, 5, 5}, {25}, {25}, {25}, {25}, {25}};
+    for (size_t i = 0; i < block_sizes.size(); ++i)
+    {
+        context.context.setSetting("max_block_size", Field(static_cast<UInt64>(block_sizes[i])));
+        auto blocks = getExecuteStreamsReturnBlocks(request);
+        ASSERT_EQ(expect[i].size(), blocks.size());
+        for (size_t j = 0; j < blocks.size(); ++j)
+        {
+            ASSERT_EQ(expect[i][j], blocks[j].rows());
+        }
+    }
 }
 CATCH
 

dbms/src/Storages/DeltaMerge/tests/gtest_segment_test_basic.cpp

@@ -301,7 +301,7 @@ Block SegmentTestBasic::prepareWriteBlock(Int64 start_key, Int64 end_key, bool i
         is_deleted);
 }
 
-Block mergeBlocks(std::vector<Block> && blocks)
+Block sortMergeBlocks(std::vector<Block> && blocks)
 {
     auto accumulated_block = std::move(blocks[0]);
 
@@ -391,7 +391,7 @@ Block SegmentTestBasic::prepareWriteBlockInSegmentRange(PageId segment_id, UInt6
                   remaining_rows);
     }
 
-    return mergeBlocks(std::move(blocks));
+    return sortMergeBlocks(std::move(blocks));
 }
 
 void SegmentTestBasic::writeSegment(PageId segment_id, UInt64 write_rows, std::optional<Int64> start_at)

dbms/src/TestUtils/ExecutorTestUtils.cpp

@@ -166,7 +166,7 @@ void ExecutorTest::executeAndAssertRowsEqual(const std::shared_ptr<tipb::DAGRequ
     });
 }
 
-Block mergeBlocks(Blocks blocks)
+Block mergeBlocksForTest(Blocks blocks)
 {
     if (blocks.empty())
         return {};
@@ -214,7 +214,7 @@ DB::ColumnsWithTypeAndName readBlocks(std::vector<BlockInputStreamPtr> streams)
     Blocks actual_blocks;
     for (const auto & stream : streams)
         readStream(actual_blocks, stream);
-    return mergeBlocks(actual_blocks).getColumnsWithTypeAndName();
+    return mergeBlocksForTest(actual_blocks).getColumnsWithTypeAndName();
 }
 
 void ExecutorTest::enablePlanner(bool is_enable)
@@ -231,7 +231,7 @@ DB::ColumnsWithTypeAndName ExecutorTest::executeStreams(const std::shared_ptr<ti
     // Currently, don't care about regions information in tests.
     Blocks blocks;
     queryExecute(context.context, /*internal=*/true)->execute([&blocks](const Block & block) { blocks.push_back(block); }).verify();
-    return mergeBlocks(blocks).getColumnsWithTypeAndName();
+    return mergeBlocksForTest(blocks).getColumnsWithTypeAndName();
 }
 
 Blocks ExecutorTest::getExecuteStreamsReturnBlocks(const std::shared_ptr<tipb::DAGRequest> & request, size_t concurrency)

dbms/src/TestUtils/ExecutorTestUtils.h

@@ -30,7 +30,7 @@ TiDB::TP dataTypeToTP(const DataTypePtr & type);
 
 ColumnsWithTypeAndName readBlock(BlockInputStreamPtr stream);
 ColumnsWithTypeAndName readBlocks(std::vector<BlockInputStreamPtr> streams);
-Block mergeBlocks(Blocks blocks);
+Block mergeBlocksForTest(Blocks blocks);
 
 
 #define WRAP_FOR_DIS_ENABLE_PLANNER_BEGIN \

dbms/src/TestUtils/MPPTaskTestUtils.cpp

@@ -93,7 +93,7 @@ ColumnsWithTypeAndName extractColumns(Context & context, const std::shared_ptr<t
     auto schema = getSelectSchema(context);
     for (const auto & chunk : dag_response->chunks())
         blocks.emplace_back(codec->decode(chunk.rows_data(), schema));
-    return mergeBlocks(blocks).getColumnsWithTypeAndName();
+    return mergeBlocksForTest(blocks).getColumnsWithTypeAndName();
 }
 
 ColumnsWithTypeAndName MPPTaskTestUtils::executeCoprocessorTask(std::shared_ptr<tipb::DAGRequest> & dag_request)