merge join result block if needed

dbms/src/Flash/tests/gtest_join_executor.cpp

@@ -800,6 +800,65 @@ try
 }
 CATCH
 
+TEST_F(JoinExecutorTestRunner, MergeAfterSplit)
+try
+{
+    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", {2, 2, 2, 2, 2, 2, 2, 2, 2, 2})});
+    context.addMockTable("split_test", "t2", {{"a", TiDB::TP::TypeLong}, {"c", TiDB::TP::TypeLong}}, {toVec<Int32>("a", {1, 1, 1, 1, 1}), toVec<Int32>("c", {1, 2, 3, 4, 5})});
+
+    std::vector<size_t> block_sizes{
+        1,
+        2,
+        7,
+        25,
+        49,
+        50,
+        51,
+        DEFAULT_BLOCK_SIZE};
+    auto join_types = {tipb::JoinType::TypeInnerJoin, tipb::JoinType::TypeSemiJoin};
+    std::vector<std::vector<std::vector<size_t>>> expects{
+        {
+            {1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+            {1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+            {4, 3, 2, 1},
+            {5, 5},
+            {9, 1},
+            {10},
+            {10},
+            {10},
+        },
+        {
+            {1, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+            {2, 2, 2, 2, 2},
+            {7, 3},
+            {10},
+            {10},
+            {10},
+            {10},
+            {10},
+        },
+    };
+    for (size_t index = 0; index < join_types.size(); index++)
+    {
+        auto request = context
+                           .scan("split_test", "t1")
+                           .join(context.scan("split_test", "t2"), *(join_types.begin() + index), {col("a")}, {}, {}, {gt(col("b"), col("c"))}, {})
+                           .build(context);
+        auto & expect = expects[index];
+
+        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
 
 TEST_F(JoinExecutorTestRunner, SpillToDisk)
 try

dbms/src/Interpreters/Join.cpp

@@ -130,6 +130,8 @@ Join::Join(
     , match_helper_name(match_helper_name)
     , kind(kind_)
     , strictness(strictness_)
+    , original_strictness(strictness)
+    , may_probe_side_expanded_after_join(mayProbeSideExpandedAfterJoin(kind, strictness))
     , key_names_left(key_names_left_)
     , key_names_right(key_names_right_)
     , build_concurrency(0)
@@ -138,7 +140,6 @@ Join::Join(
     , active_probe_threads(0)
     , collators(collators_)
     , non_equal_conditions(non_equal_conditions_)
-    , original_strictness(strictness)
     , max_block_size(max_block_size_)
     , max_bytes_before_external_join(max_bytes_before_external_join_)
     , build_spill_config(build_spill_config_)
@@ -821,8 +822,10 @@ void Join::handleOtherConditions(Block & block, std::unique_ptr<IColumn::Filter>
     throw Exception("Logical error: unknown combination of JOIN", ErrorCodes::LOGICAL_ERROR);
 }
 
-Block Join::joinBlockHash(ProbeProcessInfo & probe_process_info) const
+Block Join::doJoinBlockHash(ProbeProcessInfo & probe_process_info) const
 {
+    probe_process_info.updateStartRow();
+    /// this makes a copy of `probe_process_info.block`
     Block block = probe_process_info.block;
     size_t keys_size = key_names_left.size();
 
@@ -966,6 +969,26 @@ Block Join::joinBlockHash(ProbeProcessInfo & probe_process_info) const
     return block;
 }
 
+Block Join::joinBlockHash(ProbeProcessInfo & probe_process_info) const
+{
+    std::vector<Block> result_blocks;
+    size_t result_rows = 0;
+    while (true)
+    {
+        auto block = doJoinBlockHash(probe_process_info);
+        assert(block);
+        result_rows += block.rows();
+        result_blocks.push_back(std::move(block));
+        /// exit the while loop if
+        /// 1. probe_process_info.all_rows_joined_finish is true, which means all the rows in current block is processed
+        /// 2. the block may be expanded after join and result_rows exceeds the min_result_block_size
+        if (probe_process_info.all_rows_joined_finish || (may_probe_side_expanded_after_join && result_rows >= probe_process_info.min_result_block_size))
+            break;
+    }
+    assert(!result_blocks.empty());
+    return vstackBlocks(std::move(result_blocks));
+}
+
 namespace
 {
 template <ASTTableJoin::Kind KIND, ASTTableJoin::Strictness STRICTNESS>
@@ -1245,6 +1268,8 @@ Block Join::joinBlockCross(ProbeProcessInfo & probe_process_info) const
         DISPATCH(false)
     }
 #undef DISPATCH
+    /// todo control the returned block size for cross join
+    probe_process_info.all_rows_joined_finish = true;
     return block;
 }
 
@@ -1307,6 +1332,9 @@ Block Join::joinBlockNullAware(ProbeProcessInfo & probe_process_info) const
 
     FAIL_POINT_TRIGGER_EXCEPTION(FailPoints::random_join_prob_failpoint);
 
+    /// Null aware join never expand the left block, just handle the whole block at one time is enough
+    probe_process_info.all_rows_joined_finish = true;
+
     return block;
 }
 
@@ -1585,6 +1613,7 @@ void Join::finishOneNonJoin(size_t partition_index)
 
 Block Join::joinBlock(ProbeProcessInfo & probe_process_info, bool dry_run) const
 {
+    assert(!probe_process_info.all_rows_joined_finish);
     if unlikely (dry_run)
     {
         assert(probe_process_info.block.rows() == 0);
@@ -1600,8 +1629,6 @@ Block Join::joinBlock(ProbeProcessInfo & probe_process_info, bool dry_run) const
     }
     std::shared_lock lock(rwlock);
 
-    probe_process_info.updateStartRow();
-
     Block block{};
 
     using enum ASTTableJoin::Strictness;
@@ -1629,11 +1656,6 @@ Block Join::joinBlock(ProbeProcessInfo & probe_process_info, bool dry_run) const
         block.getByName(match_helper_name).column = ColumnNullable::create(std::move(col_non_matched), std::move(nullable_column->getNullMapColumnPtr()));
     }
 
-    if (isCrossJoin(kind) || isNullAwareSemiFamily(kind))
-    {
-        probe_process_info.all_rows_joined_finish = true;
-    }
-
     return block;
 }