support fine grained shuffle for window function

dbms/src/DataStreams/TiRemoteBlockInputStream.h

@@ -58,6 +58,11 @@ class TiRemoteBlockInputStream : public IProfilingBlockInputStream
 
     uint64_t total_rows;
 
+    // For fine grained shuffle, sender will partition data into muiltiple streams by hashing.
+    // ExchangeReceiverBlockInputStream only need to read its own stream, i.e., streams[stream_id].
+    // CoprocessorBlockInputStream doesn't take care of this.
+    UInt32 stream_id;
+
     void initRemoteExecutionSummaries(tipb::SelectResponse & resp, size_t index)
     {
         for (const auto & execution_summary : resp.execution_summaries())
@@ -120,7 +125,7 @@ class TiRemoteBlockInputStream : public IProfilingBlockInputStream
 
     bool fetchRemoteResult()
     {
-        auto result = remote_reader->nextResult(block_queue, sample_block);
+        auto result = remote_reader->nextResult(block_queue, sample_block, stream_id);
         if (result.meet_error)
         {
             LOG_FMT_WARNING(log, "remote reader meets error: {}", result.error_msg);
@@ -168,13 +173,14 @@ class TiRemoteBlockInputStream : public IProfilingBlockInputStream
     }
 
 public:
-    TiRemoteBlockInputStream(std::shared_ptr<RemoteReader> remote_reader_, const String & req_id, const String & executor_id)
+    TiRemoteBlockInputStream(std::shared_ptr<RemoteReader> remote_reader_, const String & req_id, const String & executor_id, UInt32 stream_id_)
         : remote_reader(remote_reader_)
         , source_num(remote_reader->getSourceNum())
         , name(fmt::format("TiRemoteBlockInputStream({})", RemoteReader::name))
         , execution_summaries_inited(source_num)
         , log(Logger::get(name, req_id, executor_id))
         , total_rows(0)
+        , stream_id(stream_id_)
     {
         // generate sample block
         ColumnsWithTypeAndName columns;

dbms/src/Debug/dbgFuncCoprocessor.cpp

@@ -290,8 +290,9 @@ BlockInputStreamPtr executeQuery(Context & context, RegionID region_id, const DA
                 tipb_exchange_receiver.encoded_task_meta_size(),
                 10,
                 /*req_id=*/"",
-                /*executor_id=*/"");
-        BlockInputStreamPtr ret = std::make_shared<ExchangeReceiverInputStream>(exchange_receiver, /*req_id=*/"", /*executor_id=*/"");
+                /*executor_id=*/"",
+                /*fine_grained_shuffle_stream_count=*/0);
+        BlockInputStreamPtr ret = std::make_shared<ExchangeReceiverInputStream>(exchange_receiver, /*req_id=*/"", /*executor_id=*/"", /*stream_id*/0);
         return ret;
     }
     else

dbms/src/Flash/Coprocessor/CoprocessorReader.h

@@ -139,7 +139,7 @@ class CoprocessorReader
         return detail;
     }
 
-    CoprocessorReaderResult nextResult(std::queue<Block> & block_queue, const Block & header)
+    CoprocessorReaderResult nextResult(std::queue<Block> & block_queue, const Block & header, UInt32 /* stream_id */)
     {
         auto && [result, has_next] = resp_iter.next();
         if (!result.error.empty())

dbms/src/Flash/Coprocessor/DAGContext.cpp

@@ -228,7 +228,8 @@ void DAGContext::initExchangeReceiverIfMPP(Context & context, size_t max_streams
                     executor.exchange_receiver().encoded_task_meta_size(),
                     max_streams,
                     log->identifier(),
-                    executor_id);
+                    executor_id,
+                    fine_grained_shuffle_stream_count);
                 mpp_exchange_receiver_map[executor_id] = exchange_receiver;
                 new_thread_count_of_exchange_receiver += exchange_receiver->computeNewThreadCount();
             }

dbms/src/Flash/Coprocessor/DAGContext.h

@@ -109,6 +109,8 @@ constexpr UInt64 NO_ENGINE_SUBSTITUTION = 1ul << 30ul;
 constexpr UInt64 ALLOW_INVALID_DATES = 1ul << 32ul;
 } // namespace TiDBSQLMode
 
+inline bool enableFineGrainedShuffle(uint32_t stream_count) { return stream_count > 0; }
+
 /// A context used to track the information that needs to be passed around during DAG planning.
 class DAGContext
 {
@@ -125,6 +127,8 @@ class DAGContext
         , max_recorded_error_count(getMaxErrorCount(*dag_request))
         , warnings(max_recorded_error_count)
         , warning_count(0)
+        , fine_grained_shuffle_stream_count(0)
+        , fine_grained_shuffle_batch_size(0)
     {
         assert(dag_request->has_root_executor() || dag_request->executors_size() > 0);
         return_executor_id = dag_request->root_executor().has_executor_id() || dag_request->executors(0).has_executor_id();
@@ -147,6 +151,8 @@ class DAGContext
         , max_recorded_error_count(getMaxErrorCount(*dag_request))
         , warnings(max_recorded_error_count)
         , warning_count(0)
+        , fine_grained_shuffle_stream_count(0)
+        , fine_grained_shuffle_batch_size(0)
     {
         assert(dag_request->has_root_executor() && dag_request->root_executor().has_executor_id());
 
@@ -167,6 +173,8 @@ class DAGContext
         , max_recorded_error_count(max_error_count_)
         , warnings(max_recorded_error_count)
         , warning_count(0)
+        , fine_grained_shuffle_stream_count(0)
+        , fine_grained_shuffle_batch_size(0)
         , is_test(true)
     {}
 
@@ -341,6 +349,12 @@ class DAGContext
     std::vector<tipb::FieldType> output_field_types;
     std::vector<Int32> output_offsets;
 
+    bool enableFineGrainedShuffle() const { return ::DB::enableFineGrainedShuffle(fine_grained_shuffle_stream_count); }
+    uint32_t getFineGrainedShuffleStreamCount() const { return fine_grained_shuffle_stream_count; }
+    Int64 getFineGrainedShuffleBatchSize() const { return fine_grained_shuffle_batch_size; }
+    void setFineGrainedShuffleStreamCount(uint32_t count) { fine_grained_shuffle_stream_count = count; }
+    void setFineGrainedShuffleBatchSize(Int64 size) { fine_grained_shuffle_batch_size = size; }
+
 private:
     void initExecutorIdToJoinIdMap();
     void initOutputInfo();
@@ -375,6 +389,9 @@ class DAGContext
     /// The order of the vector is also the order of the subquery.
     std::vector<SubqueriesForSets> subqueries;
 
+    uint32_t fine_grained_shuffle_stream_count;
+    Int64 fine_grained_shuffle_batch_size;
+
     bool is_test = false; /// switch for test, do not use it in production.
     std::unordered_map<String, ColumnsWithTypeAndName> columns_for_test_map; /// <exector_id, columns>, for multiple sources
 };

dbms/src/Flash/Coprocessor/DAGDriver.cpp

@@ -137,7 +137,9 @@ try
             context.getSettingsRef().dag_records_per_chunk,
             context.getSettingsRef().batch_send_min_limit,
             true,
-            dag_context);
+            dag_context,
+            /*fine_grained_shuffle_stream_count=*/0,
+            /*fine_grained_shuffle_batch_size=*/0);
         dag_output_stream = std::make_shared<DAGBlockOutputStream>(streams.in->getHeader(), std::move(response_writer));
         copyData(*streams.in, *dag_output_stream);
     }

dbms/src/Flash/Coprocessor/DAGQueryBlockInterpreter.cpp

@@ -351,10 +351,18 @@ void DAGQueryBlockInterpreter::executeWindow(
 {
     executeExpression(pipeline, window_description.before_window, "before window");
 
-    /// If there are several streams, we merge them into one
-    executeUnion(pipeline, max_streams, log, false, "merge into one for window input");
-    assert(pipeline.streams.size() == 1);
-    pipeline.firstStream() = std::make_shared<WindowBlockInputStream>(pipeline.firstStream(), window_description, log->identifier());
+    if (dagContext().enableFineGrainedShuffle())
+    {
+        // fine_grained_shuffle_stream_count is logical, no need to be equal to real stream count.
+        pipeline.transform([&](auto & stream) { stream = std::make_shared<WindowBlockInputStream>(stream, window_description, log->identifier()); });
+    }
+    else
+    {
+        /// If there are several streams, we merge them into one
+        executeUnion(pipeline, max_streams, log, false, "merge into one for window input");
+        assert(pipeline.streams.size() == 1);
+        pipeline.firstStream() = std::make_shared<WindowBlockInputStream>(pipeline.firstStream(), window_description, log->identifier());
+    }
 }
 
 void DAGQueryBlockInterpreter::executeAggregation(
@@ -434,16 +442,16 @@ void DAGQueryBlockInterpreter::executeExpression(DAGPipeline & pipeline, const E
 
 void DAGQueryBlockInterpreter::executeWindowOrder(DAGPipeline & pipeline, SortDescription sort_desc)
 {
-    orderStreams(pipeline, sort_desc, 0);
+    orderStreams(pipeline, sort_desc, 0, dagContext().enableFineGrainedShuffle());
 }
 
 void DAGQueryBlockInterpreter::executeOrder(DAGPipeline & pipeline, const NamesAndTypes & order_columns)
 {
     Int64 limit = query_block.limit_or_topn->topn().limit();
-    orderStreams(pipeline, getSortDescription(order_columns, query_block.limit_or_topn->topn().order_by()), limit);
+    orderStreams(pipeline, getSortDescription(order_columns, query_block.limit_or_topn->topn().order_by()), limit, false);
 }
 
-void DAGQueryBlockInterpreter::orderStreams(DAGPipeline & pipeline, SortDescription order_descr, Int64 limit)
+void DAGQueryBlockInterpreter::orderStreams(DAGPipeline & pipeline, SortDescription order_descr, Int64 limit, bool enable_fine_grained_shuffle)
 {
     const Settings & settings = context.getSettingsRef();
 
@@ -459,18 +467,33 @@ void DAGQueryBlockInterpreter::orderStreams(DAGPipeline & pipeline, SortDescript
         stream = sorting_stream;
     });
 
-    /// If there are several streams, we merge them into one
-    executeUnion(pipeline, max_streams, log, false, "for partial order");
-
-    /// Merge the sorted blocks.
-    pipeline.firstStream() = std::make_shared<MergeSortingBlockInputStream>(
-        pipeline.firstStream(),
-        order_descr,
-        settings.max_block_size,
-        limit,
-        settings.max_bytes_before_external_sort,
-        context.getTemporaryPath(),
-        log->identifier());
+    if (enable_fine_grained_shuffle)
+    {
+        pipeline.transform([&](auto & stream) { stream = std::make_shared<MergeSortingBlockInputStream>(
+                    stream,
+                    order_descr,
+                    settings.max_block_size,
+                    limit,
+                    settings.max_bytes_before_external_sort,
+                    context.getTemporaryPath(),
+                    log->identifier());
+                });
+    }
+    else
+    {
+        /// If there are several streams, we merge them into one
+        executeUnion(pipeline, max_streams, log, false, "for partial order");
+
+        /// Merge the sorted blocks.
+        pipeline.firstStream() = std::make_shared<MergeSortingBlockInputStream>(
+            pipeline.firstStream(),
+            order_descr,
+            settings.max_block_size,
+            limit,
+            settings.max_bytes_before_external_sort,
+            context.getTemporaryPath(),
+            log->identifier());
+    }
 }
 
 void DAGQueryBlockInterpreter::recordProfileStreams(DAGPipeline & pipeline, const String & key)
@@ -486,9 +509,12 @@ void DAGQueryBlockInterpreter::handleExchangeReceiver(DAGPipeline & pipeline)
         throw Exception("Can not find exchange receiver for " + query_block.source_name, ErrorCodes::LOGICAL_ERROR);
     // todo choose a more reasonable stream number
     auto & exchange_receiver_io_input_streams = dagContext().getInBoundIOInputStreamsMap()[query_block.source_name];
+    size_t stream_id = 0;
     for (size_t i = 0; i < max_streams; ++i)
     {
-        BlockInputStreamPtr stream = std::make_shared<ExchangeReceiverInputStream>(it->second, log->identifier(), query_block.source_name);
+        if (dagContext().enableFineGrainedShuffle())
+            stream_id = i;
+        BlockInputStreamPtr stream = std::make_shared<ExchangeReceiverInputStream>(it->second, log->identifier(), query_block.source_name, stream_id);
         exchange_receiver_io_input_streams.push_back(stream);
         stream = std::make_shared<SquashingBlockInputStream>(stream, 8192, 0, log->identifier());
         stream->setExtraInfo("squashing after exchange receiver");
@@ -743,6 +769,7 @@ void DAGQueryBlockInterpreter::handleExchangeSender(DAGPipeline & pipeline)
     std::vector<Int64> partition_col_ids = ExchangeSenderInterpreterHelper::genPartitionColIds(exchange_sender);
     TiDB::TiDBCollators partition_col_collators = ExchangeSenderInterpreterHelper::genPartitionColCollators(exchange_sender);
     int stream_id = 0;
+
     pipeline.transform([&](auto & stream) {
         // construct writer
         std::unique_ptr<DAGResponseWriter> response_writer = std::make_unique<StreamingDAGResponseWriter<MPPTunnelSetPtr>>(
@@ -753,7 +780,9 @@ void DAGQueryBlockInterpreter::handleExchangeSender(DAGPipeline & pipeline)
             context.getSettingsRef().dag_records_per_chunk,
             context.getSettingsRef().batch_send_min_limit,
             stream_id++ == 0, /// only one stream needs to sending execution summaries for the last response
-            dagContext());
+            dagContext(),
+            dagContext().getFineGrainedShuffleStreamCount(),
+            dagContext().getFineGrainedShuffleBatchSize());
         stream = std::make_shared<ExchangeSenderBlockInputStream>(stream, std::move(response_writer), log->identifier());
     });
 }
@@ -783,4 +812,4 @@ BlockInputStreams DAGQueryBlockInterpreter::execute()
 
     return pipeline.streams;
 }
-} // namespace DB
+} // namespace DB

dbms/src/Flash/Coprocessor/DAGQueryBlockInterpreter.h

@@ -69,7 +69,7 @@ class DAGQueryBlockInterpreter
     void executeWhere(DAGPipeline & pipeline, const ExpressionActionsPtr & expressionActionsPtr, String & filter_column, const String & extra_info = "");
     void executeExpression(DAGPipeline & pipeline, const ExpressionActionsPtr & expressionActionsPtr, const String & extra_info = "");
     void executeWindowOrder(DAGPipeline & pipeline, SortDescription sort_desc);
-    void orderStreams(DAGPipeline & pipeline, SortDescription order_descr, Int64 limit);
+    void orderStreams(DAGPipeline & pipeline, SortDescription order_descr, Int64 limit, bool enable_fine_grained_shuffle);
     void executeOrder(DAGPipeline & pipeline, const NamesAndTypes & order_columns);
     void executeLimit(DAGPipeline & pipeline);
     void executeWindow(

dbms/src/Flash/Coprocessor/DAGStorageInterpreter.cpp

@@ -486,7 +486,7 @@ void DAGStorageInterpreter::buildRemoteStreams(std::vector<RemoteRequest> && rem
         std::vector<pingcap::coprocessor::copTask> tasks(all_tasks.begin() + task_start, all_tasks.begin() + task_end);
 
         auto coprocessor_reader = std::make_shared<CoprocessorReader>(schema, cluster, tasks, has_enforce_encode_type, 1);
-        BlockInputStreamPtr input = std::make_shared<CoprocessorBlockInputStream>(coprocessor_reader, log->identifier(), table_scan.getTableScanExecutorID());
+        BlockInputStreamPtr input = std::make_shared<CoprocessorBlockInputStream>(coprocessor_reader, log->identifier(), table_scan.getTableScanExecutorID(), /*stream_id=*/0);
         pipeline.streams.push_back(input);
         task_start = task_end;
     }

dbms/src/Flash/Coprocessor/DecodeDetail.h

@@ -21,8 +21,12 @@ namespace DB
 /// Detail of the packet that decoding in TiRemoteInputStream.RemoteReader.decodeChunks()
 struct DecodeDetail
 {
+    // For fine grained shuffle, each ExchangeReceiver/thread will decode its own blocks.
+    // So this is the row number of partial blocks of the original packet.
+    // This will be the row number of all blocks of the original packet if it's not fine grained shuffle.
     Int64 rows = 0;
-    // byte size of origin packet.
+
+    // Total byte size of the origin packet, even for fine grained shuffle.
     Int64 packet_bytes = 0;
 };
-} // namespace DB
+} // namespace DB