Support non-recursive CTE in TiFlash (#10086)

ref #10085

Author: xzhangxian1008

dbms/src/Common/TiFlashMetrics.h

@@ -95,6 +95,8 @@ static_assert(RAFT_REGION_BIG_WRITE_THRES * 4 < RAFT_REGION_BIG_WRITE_MAX, "Inva
       F(type_exchange_receiver, {"type", "exchange_receiver"}),                                                                     \
       F(type_projection, {"type", "projection"}),                                                                                   \
       F(type_partition_ts, {"type", "partition_table_scan"}),                                                                       \
+      F(type_cte_sink, {"type", "cte_sink"}),                                                                                       \
+      F(type_cte_source, {"type", "cte_source"}),                                                                                   \
       F(type_window, {"type", "window"}),                                                                                           \
       F(type_window_sort, {"type", "window_sort"}),                                                                                 \
       F(type_expand, {"type", "expand"}))                                                                                           \
@@ -775,7 +777,8 @@ static_assert(RAFT_REGION_BIG_WRITE_THRES * 4 < RAFT_REGION_BIG_WRITE_MAX, "Inva
       F(type_wait_on_tunnel_sender_write, {"type", "wait_on_tunnel_sender_write"}),                                                 \
       F(type_wait_on_join_build, {"type", "wait_on_join_build"}),                                                                   \
       F(type_wait_on_join_probe, {"type", "wait_on_join_probe"}),                                                                   \
-      F(type_wait_on_result_queue_write, {"type", "wait_on_result_queue_write"}))                                                   \
+      F(type_wait_on_result_queue_write, {"type", "wait_on_result_queue_write"}),                                                   \
+      F(type_type_wait_on_cte_read, {"type", "type_wait_on_cte_read"}))                                                             \
     M(tiflash_pipeline_task_duration_seconds,                                                                                       \
       "Bucketed histogram of pipeline task duration in seconds",                                                                    \
       Histogram, /* these command usually cost several hundred milliseconds to several seconds, increase the start bucket to 5ms */ \

dbms/src/Common/WeakHash.h

@@ -28,7 +28,7 @@ class WeakHash32
 public:
     using Container = PaddedPODArray<UInt32>;
 
-    static constexpr UInt32 initial_hash = ~UInt32(0);
+    static constexpr UInt32 initial_hash = ~static_cast<UInt32>(0);
 
     explicit WeakHash32(size_t size)
         : data(size, initial_hash)

dbms/src/Flash/Coprocessor/DAGContext.h

@@ -14,6 +14,9 @@
 
 #pragma once
 
+#include <cstddef>
+#include <mutex>
+#include <unordered_map>
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wunused-parameter"
 #ifdef __clang__
@@ -23,6 +26,7 @@
 #pragma GCC diagnostic pop
 
 #include <Common/ConcurrentBoundedQueue.h>
+#include <Common/Exception.h>
 #include <Common/Logger.h>
 #include <Core/QueryOperatorSpillContexts.h>
 #include <Core/TaskOperatorSpillContexts.h>
@@ -55,6 +59,7 @@ class CoprocessorReader;
 using CoprocessorReaderPtr = std::shared_ptr<CoprocessorReader>;
 
 class AutoSpillTrigger;
+class CTE;
 
 struct JoinProfileInfo
 {
@@ -359,6 +364,70 @@ class DAGContext
 
     MPPReceiverSetPtr getMPPReceiverSet() const { return mpp_receiver_set; }
 
+    String getQueryIDAndCTEIDForSink()
+    {
+        std::lock_guard<std::mutex> lock(this->cte_mu);
+        return this->query_id_and_cte_id_for_sink;
+    }
+
+    String getQueryIDAndCTEIDForSource(size_t cte_id)
+    {
+        std::lock_guard<std::mutex> lock(this->cte_mu);
+        return this->query_id_and_cte_id_for_sources[cte_id];
+    }
+
+    void setQueryIDAndCTEIDForSink(const String & query_id_and_cte_id)
+    {
+        std::lock_guard<std::mutex> lock(this->cte_mu);
+
+        // MPP Task has only one CTESink, it's impossible to set query_id_and_cte_id_for_sink twice
+        RUNTIME_CHECK(this->query_id_and_cte_id_for_sink.empty(), this->query_id_and_cte_id_for_sink);
+        this->query_id_and_cte_id_for_sink = query_id_and_cte_id;
+    }
+
+    void addQueryIDAndCTEIDForSource(size_t cte_id, const String & query_id_and_cte_id)
+    {
+        std::lock_guard<std::mutex> lock(this->cte_mu);
+        auto iter = this->query_id_and_cte_id_for_sources.find(cte_id);
+        if (iter != this->query_id_and_cte_id_for_sources.end())
+        {
+            RUNTIME_CHECK_MSG(iter->second == query_id_and_cte_id, "{} vs {}", iter->second, query_id_and_cte_id);
+            return;
+        }
+        this->query_id_and_cte_id_for_sources.insert(std::make_pair(cte_id, query_id_and_cte_id));
+    }
+
+    std::shared_ptr<CTE> getCTESink()
+    {
+        std::lock_guard<std::mutex> lock(this->cte_mu);
+        return this->sink_cte;
+    }
+
+    std::unordered_map<size_t, std::shared_ptr<CTE>> getCTESource()
+    {
+        std::lock_guard<std::mutex> lock(this->cte_mu);
+        return this->source_ctes;
+    }
+
+    void setCTESink(std::shared_ptr<CTE> & cte)
+    {
+        std::lock_guard<std::mutex> lock(this->cte_mu);
+        RUNTIME_CHECK(!this->sink_cte);
+        this->sink_cte = cte;
+    }
+
+    void addCTESource(size_t cte_id, std::shared_ptr<CTE> & cte)
+    {
+        std::lock_guard<std::mutex> lock(this->cte_mu);
+        auto iter = this->source_ctes.find(cte_id);
+        if (iter != this->source_ctes.end())
+        {
+            RUNTIME_CHECK(iter->second.get() == cte.get());
+            return;
+        }
+        this->source_ctes.insert(std::make_pair(cte_id, cte));
+    }
+
 public:
     DAGRequest dag_request;
     /// Some existing code inherited from Clickhouse assume that each query must have a valid query string and query ast,
@@ -475,6 +544,13 @@ class DAGContext
     UInt64 connection_id;
     // It's the session alias between mysql client and tidb
     String connection_alias;
+
+    String query_id_and_cte_id_for_sink;
+    std::unordered_map<size_t, String> query_id_and_cte_id_for_sources;
+
+    std::mutex cte_mu;
+    std::shared_ptr<CTE> sink_cte;
+    std::unordered_map<size_t, std::shared_ptr<CTE>> source_ctes;
 };
 
 } // namespace DB

dbms/src/Flash/Coprocessor/GenSchemaAndColumn.cpp

@@ -61,6 +61,11 @@ String genNameForExchangeReceiver(Int32 col_index)
     return fmt::format("exchange_receiver_{}", col_index);
 }
 
+String genNameForCTESource(Int32 cte_id, Int32 col_index)
+{
+    return fmt::format("cte_source_{}_{}", cte_id, col_index);
+}
+
 NamesAndTypes genNamesAndTypes(const TiDB::ColumnInfos & column_infos, const StringRef & column_prefix)
 {
     NamesAndTypes names_and_types;

dbms/src/Flash/Coprocessor/GenSchemaAndColumn.h

@@ -27,6 +27,7 @@ namespace DB
 NamesAndTypes genNamesAndTypesForExchangeReceiver(const TiDBTableScan & table_scan);
 NamesAndTypes genNamesAndTypesForTableScan(const TiDBTableScan & table_scan);
 String genNameForExchangeReceiver(Int32 col_index);
+String genNameForCTESource(Int32 cte_id, Int32 col_index);
 
 NamesAndTypes genNamesAndTypes(const TiDBTableScan & table_scan, const StringRef & column_prefix);
 NamesAndTypes genNamesAndTypes(const TiDB::ColumnInfos & column_infos, const StringRef & column_prefix);

dbms/src/Flash/Executor/PipelineExecutor.cpp

@@ -14,6 +14,7 @@
 
 #include <Flash/Coprocessor/DAGContext.h>
 #include <Flash/Executor/PipelineExecutor.h>
+#include <Flash/Executor/PipelineExecutorContext.h>
 #include <Flash/Pipeline/Pipeline.h>
 #include <Flash/Pipeline/Schedule/Events/Event.h>
 #include <Flash/Planner/PhysicalPlan.h>

dbms/src/Flash/Executor/PipelineExecutorContext.cpp

@@ -20,6 +20,7 @@
 #include <Flash/Mpp/Utils.h>
 #include <Flash/Pipeline/Schedule/TaskScheduler.h>
 #include <Flash/Pipeline/Schedule/Tasks/OneTimeNotifyFuture.h>
+#include <Operators/CTE.h>
 #include <Operators/SharedQueue.h>
 
 #include <exception>
@@ -185,6 +186,14 @@ void PipelineExecutorContext::cancel()
             // pipeline tasks waiting in the WAIT_FOR_NOTIFY state from never being notified.
             if (dag_context->tunnel_set)
                 dag_context->tunnel_set->close(getTrimmedErrMsg(), false);
+
+            auto cte = dag_context->getCTESink();
+            if (cte)
+                cte->notifyCancel<false>(getTrimmedErrMsg());
+
+            for (auto & p : dag_context->getCTESource())
+                p.second->notifyCancel<false>(getTrimmedErrMsg());
+
             if (auto mpp_receiver_set = dag_context->getMPPReceiverSet(); mpp_receiver_set)
                 mpp_receiver_set->cancel();
         }

dbms/src/Flash/Executor/PipelineExecutorContext.h

@@ -17,6 +17,7 @@
 #include <Common/Logger.h>
 #include <Common/MemoryTracker.h>
 #include <Core/AutoSpillTrigger.h>
+#include <Flash/Coprocessor/DAGContext.h>
 #include <Flash/Executor/ExecutionResult.h>
 #include <Flash/Executor/ResultHandler.h>
 #include <Flash/Executor/ResultQueue_fwd.h>

dbms/src/Flash/Mpp/CTEManager.cpp

@@ -0,0 +1,85 @@
+// Copyright 2025 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <Flash/Mpp/CTEManager.h>
+#include <fmt/core.h>
+#include <tipb/select.pb.h>
+
+#include <mutex>
+#include <utility>
+
+namespace DB
+{
+void CTEManager::releaseCTEBySource(const String & query_id_and_cte_id)
+{
+    std::lock_guard<std::mutex> lock(this->mu);
+    auto iter = this->ctes.find(query_id_and_cte_id);
+    if unlikely (iter == this->ctes.end())
+        // Maybe the task is cancelled and the cte has been released
+        return;
+
+    iter->second->sourceExit();
+    if (iter->second->allExit())
+        this->ctes.erase(iter);
+}
+
+void CTEManager::releaseCTEBySink(const tipb::SelectResponse & resp, const String & query_id_and_cte_id)
+{
+    std::lock_guard<std::mutex> lock(this->mu);
+    auto iter = this->ctes.find(query_id_and_cte_id);
+    if unlikely (iter == this->ctes.end())
+        // Maybe the task is cancelled and the cte has been released
+        return;
+
+    std::shared_ptr<CTE> cte = iter->second;
+    cte->addResp(resp);
+    cte->sinkExit<false>();
+    if (cte->allExit())
+        this->ctes.erase(iter);
+}
+
+void CTEManager::releaseCTE(const String & query_id_and_cte_id)
+{
+    std::lock_guard<std::mutex> lock(this->mu);
+    this->ctes.erase(query_id_and_cte_id);
+}
+
+std::shared_ptr<CTE> CTEManager::getOrCreateCTE(
+    const String & query_id_and_cte_id,
+    Int32 concurrency,
+    Int32 expected_sink_num,
+    Int32 expected_source_num)
+{
+    std::lock_guard<std::mutex> lock(this->mu);
+    auto iter = this->ctes.find(query_id_and_cte_id);
+    std::shared_ptr<CTE> cte;
+    if (iter == this->ctes.end())
+    {
+        cte = std::make_shared<CTE>(concurrency, expected_sink_num, expected_source_num);
+        this->ctes.insert(std::make_pair(query_id_and_cte_id, cte));
+    }
+    else
+    {
+        cte = iter->second;
+    }
+
+    return cte;
+}
+
+bool CTEManager::hasCTEForTest(const String & query_id_and_cte_id)
+{
+    std::lock_guard<std::mutex> lock(this->mu);
+    return !(this->ctes.find(query_id_and_cte_id) == this->ctes.end());
+}
+} // namespace DB

dbms/src/Flash/Mpp/CTEManager.h

@@ -0,0 +1,42 @@
+// Copyright 2025 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <Operators/CTE.h>
+
+#include <memory>
+#include <mutex>
+
+namespace DB
+{
+class CTEManager
+{
+public:
+    std::shared_ptr<CTE> getOrCreateCTE(
+        const String & query_id_and_cte_id,
+        Int32 concurrency,
+        Int32 expected_sink_num,
+        Int32 expected_source_num);
+    void releaseCTEBySource(const String & query_id_and_cte_id);
+    void releaseCTEBySink(const tipb::SelectResponse & resp, const String & query_id_and_cte_id);
+    void releaseCTE(const String & query_id_and_cte_id);
+
+    bool hasCTEForTest(const String & query_id_and_cte_id);
+
+private:
+    std::mutex mu;
+    std::unordered_map<String, std::shared_ptr<CTE>> ctes;
+};
+} // namespace DB