Support Non-recursive CTE spill in TiFlash

dbms/src/Common/TiFlashMetrics.h

@@ -778,7 +778,8 @@ static_assert(RAFT_REGION_BIG_WRITE_THRES * 4 < RAFT_REGION_BIG_WRITE_MAX, "Inva
       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_type_wait_on_cte_read, {"type", "type_wait_on_cte_read"}))                                                             \
+      F(type_wait_on_cte_read, {"type", "type_wait_on_cte_read"}),                                                                  \
+      F(type_wait_on_cte_io, {"type", "type_wait_on_cte_io"}))                                                                      \
     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/Core/Spiller.h

@@ -161,7 +161,7 @@ class Spiller
     /// file is read, otherwise, the spilled file will be released when destruct the spiller. Currently, all the spilled
     /// file can be released on restore since it is only read once, but in the future if SharedScan(shared cte) need spill,
     /// the data may be restored multiple times and release_spilled_file_on_restore need to be set to false.
-    const bool release_spilled_file_on_restore;
+    bool release_spilled_file_on_restore;
     bool enable_append_write = false;
 };
 

dbms/src/DataStreams/NativeBlockOutputStream.cpp

@@ -140,5 +140,4 @@ void NativeBlockOutputStream::write(const Block & block)
         }
     }
 }
-
 } // namespace DB

dbms/src/Flash/FlashService.cpp

@@ -151,6 +151,7 @@ void updateSettingsFromTiDB(const grpc::ServerContext * grpc_context, ContextPtr
         std::make_pair("tidb_max_bytes_before_tiflash_external_join", "max_bytes_before_external_join"),
         std::make_pair("tidb_max_bytes_before_tiflash_external_group_by", "max_bytes_before_external_group_by"),
         std::make_pair("tidb_max_bytes_before_tiflash_external_sort", "max_bytes_before_external_sort"),
+        std::make_pair("tidb_max_bytes_before_tiflash_cte_spill", "max_bytes_before_cte_spill"),
         std::make_pair("tiflash_mem_quota_query_per_node", "max_memory_usage"),
         std::make_pair("tiflash_query_spill_ratio", "auto_memory_revoke_trigger_threshold"),
         std::make_pair("tiflash_use_hash_join_v2", "enable_hash_join_v2"),
@@ -188,6 +189,11 @@ void updateSettingsForAutoSpill(ContextPtr & context, const LoggerPtr & log)
             need_log_warning = true;
             context->setSetting("max_bytes_before_external_join", "0");
         }
+        if (context->getSettingsRef().max_bytes_before_cte_spill > 0)
+        {
+            need_log_warning = true;
+            context->setSetting("max_bytes_before_cte_spill", "0");
+        }
         if (need_log_warning)
             LOG_WARNING(log, "auto spill is enabled, so per operator's memory threshold is disabled");
     }

dbms/src/Flash/Mpp/MPPTask.cpp

@@ -43,7 +43,6 @@
 #include <ext/scope_guard.h>
 #include <magic_enum.hpp>
 
-
 namespace DB
 {
 namespace ErrorCodes

dbms/src/Flash/Pipeline/Schedule/Reactor/WaitReactor.cpp

@@ -40,7 +40,8 @@ WaitReactor::WaitReactor(TaskScheduler & scheduler_)
     GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_wait_on_join_probe).Set(0);
     GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_wait_on_result_queue_write).Set(0);
     GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_wait_on_shared_queue_write).Set(0);
-    GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_type_wait_on_cte_read).Set(0);
+    GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_wait_on_cte_read).Set(0);
+    GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_wait_on_cte_io).Set(0);
     thread = std::thread(&WaitReactor::loop, this);
 }
 

dbms/src/Flash/Pipeline/Schedule/Tasks/PipeConditionVariable.h

@@ -58,7 +58,10 @@ class PipeConditionVariable
             GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_wait_on_shared_queue_write).Increment(change);
             break;
         case NotifyType::WAIT_ON_CTE_READ:
-            GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_type_wait_on_cte_read).Increment(change);
+            GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_wait_on_cte_read).Increment(change);
+            break;
+        case NotifyType::WAIT_ON_CTE_IO:
+            GET_METRIC(tiflash_pipeline_wait_on_notify_tasks, type_wait_on_cte_io).Increment(change);
             break;
         case NotifyType::WAIT_ON_NOTHING:
             throw Exception("task notify type should be set before register or notify");

dbms/src/Flash/Pipeline/Schedule/Tasks/Task.h

@@ -59,6 +59,7 @@ enum class NotifyType
     WAIT_ON_JOIN_PROBE_FINISH,
     WAIT_ON_RESULT_QUEUE_WRITE,
     WAIT_ON_CTE_READ,
+    WAIT_ON_CTE_IO,
     WAIT_ON_NOTHING,
 };
 

dbms/src/Flash/Planner/Plans/PhysicalCTESink.cpp

@@ -52,6 +52,23 @@ void PhysicalCTESink::buildPipelineExecGroupImpl(
         builder.setSinkOp(std::make_unique<CTESinkOp>(exec_context, log->identifier(), cte, id));
         id++;
     });
+
+    const Settings & settings = context.getSettingsRef();
+    SpillConfig spill_config(
+        context.getTemporaryPath(),
+        " ",
+        settings.max_cached_data_bytes_in_spiller,
+        settings.max_spilled_rows_per_file,
+        settings.max_spilled_bytes_per_file,
+        context.getFileProvider(),
+        settings.max_threads,
+        settings.max_block_size);
+
+    cte->initCTESpillContextAndPartitionConfig(
+        spill_config,
+        group_builder.getCurrentHeader(),
+        settings.max_bytes_before_cte_spill,
+        context);
 }
 
 void PhysicalCTESink::finalizeImpl(const Names & parent_require)

dbms/src/Interpreters/CTESpillContext.cpp

@@ -0,0 +1,56 @@
+// 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 <Common/Exception.h>
+#include <Interpreters/CTESpillContext.h>
+
+#include <algorithm>
+#include <mutex>
+#include <utility>
+
+namespace DB
+{
+Int64 CTESpillContext::triggerSpillImpl(Int64 expected_released_memories)
+{
+    size_t partition_num = this->partitions.size();
+    std::vector<std::pair<size_t, size_t>> part_idx_with_mem_usages;
+    part_idx_with_mem_usages.reserve(partition_num);
+    for (size_t i = 0; i < partition_num; i++)
+        part_idx_with_mem_usages.push_back(std::make_pair(this->partitions[i]->memory_usage.load(), i));
+
+    std::sort(
+        part_idx_with_mem_usages.begin(),
+        part_idx_with_mem_usages.end(),
+        [](const std::pair<size_t, size_t> & l, const std::pair<size_t, size_t> & r) { return l.first > r.first; });
+
+    for (const auto & item : part_idx_with_mem_usages)
+    {
+        auto memory_usage = item.first;
+        const std::shared_ptr<CTEPartition> & partition = partitions[item.second];
+        std::lock_guard<std::mutex> aux_lock(*(partition->aux_lock));
+        if (partition->status != CTEPartitionStatus::NORMAL)
+            continue;
+
+        if (memory_usage == 0)
+            continue;
+
+        partition->status = CTEPartitionStatus::NEED_SPILL;
+
+        expected_released_memories = std::max(expected_released_memories - memory_usage, 0);
+        if (expected_released_memories <= 0)
+            return expected_released_memories;
+    }
+    return expected_released_memories;
+}
+} // namespace DB

dbms/src/Interpreters/CTESpillContext.h

@@ -0,0 +1,55 @@
+// 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 <Common/Exception.h>
+#include <Common/Logger.h>
+#include <Core/Block.h>
+#include <Core/OperatorSpillContext.h>
+#include <Operators/CTEPartition.h>
+
+#include <memory>
+#include <mutex>
+
+namespace DB
+{
+class CTESpillContext final : public OperatorSpillContext
+{
+public:
+    CTESpillContext(
+        UInt64 operator_spill_threshold_,
+        const String & query_id_and_cte_id_,
+        std::vector<std::shared_ptr<CTEPartition>> partitions_)
+        : OperatorSpillContext(operator_spill_threshold_, "cte", Logger::get(query_id_and_cte_id_))
+        , partitions(partitions_)
+    {}
+
+    bool supportAutoTriggerSpill() const override { return true; }
+    Int64 triggerSpillImpl(Int64 expected_released_memories) override;
+    LoggerPtr getLog() const { return this->log; }
+
+protected:
+    Int64 getTotalRevocableMemoryImpl() override
+    {
+        Int64 total_memory = 0;
+        for (auto & partition : this->partitions)
+            total_memory += partition->memory_usage.load();
+        return total_memory;
+    }
+
+private:
+    std::vector<std::shared_ptr<CTEPartition>> partitions;
+};
+} // namespace DB

dbms/src/Interpreters/Settings.h

@@ -313,6 +313,7 @@ struct Settings
     M(SettingUInt64, async_cqs, 1, "grpc async cqs")                                                                                                                                                                                    \
     M(SettingUInt64, preallocated_request_count_per_poller, 20, "grpc preallocated_request_count_per_poller")                                                                                                                           \
     M(SettingUInt64, max_bytes_before_external_join, 0, "max bytes used by join before spill, 0 as the default value, 0 means no limit")                                                                                                \
+    M(SettingUInt64, max_bytes_before_cte_spill, 0, "max bytes used by cte before spill, 0 as the default value, 0 means no limit")                                                                                                \
     M(SettingInt64, join_restore_concurrency, 0, "join restore concurrency, negative value means restore join serially, 0 means TiFlash choose restore concurrency automatically, 0 as the default value")                              \
     M(SettingUInt64, max_cached_data_bytes_in_spiller, 1024ULL * 1024 * 20, "Max cached data bytes in spiller before spilling, 20 MB as the default value, 0 means no limit")                                                           \
     M(SettingUInt64, max_spilled_rows_per_file, 200000, "Max spilled data rows per spill file, 200000 as the default value, 0 means no limit.")                                                                                         \