fix: respect disabled field data budget

Signed-off-by: Shawn Wang <shawn.wang@zilliz.com>

Author: sparknack

internal/core/src/segcore/MemoryPlannerTest.cpp

@@ -20,10 +20,12 @@
 #include <cstdint>
 #include <future>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
 #include "arrow/api.h"
+#include "cachinglayer/LoadingOverheadTracker.h"
 #include <folly/CancellationToken.h>
 #include "folly/ScopeGuard.h"
 #include "common/EasyAssert.h"
@@ -209,6 +211,52 @@ TEST(FieldDataLoadBatchSplitTargetBytes, CapsTargetByConfiguredBudget) {
     EXPECT_EQ(FieldDataLoadBatchSplitTargetBytes(), 8 * MB);
 }
 
+TEST(FieldDataLoadingOverheadUpperBound, UsesUnlimitedWhenBudgetDisabled) {
+    auto& budget =
+        milvus::storage::TransientMemoryBudget::GetLoadTransientBudget();
+    auto old_capacity = budget.CapacityBytes();
+    auto cleanup = folly::makeGuard(
+        [&budget, old_capacity]() { budget.SetCapacityBytes(old_capacity); });
+
+    budget.SetCapacityBytes(0);
+
+    auto upper_bound =
+        FieldDataLoadingOverheadUpperBound(/*max_memory_overhead=*/128);
+
+    EXPECT_EQ(
+        upper_bound.memory_bytes,
+        milvus::cachinglayer::LoadingOverheadTracker::kUnlimited.memory_bytes);
+    EXPECT_EQ(
+        upper_bound.file_bytes,
+        milvus::cachinglayer::LoadingOverheadTracker::kUnlimited.file_bytes);
+}
+
+TEST(FieldDataLoadingOverheadUpperBound, UsesBudgetWithMaxOverheadFloor) {
+    constexpr int64_t MB = 1 << 20;
+    auto& budget =
+        milvus::storage::TransientMemoryBudget::GetLoadTransientBudget();
+    auto old_capacity = budget.CapacityBytes();
+    auto cleanup = folly::makeGuard(
+        [&budget, old_capacity]() { budget.SetCapacityBytes(old_capacity); });
+
+    budget.SetCapacityBytes(8 * MB);
+
+    auto memory_only =
+        FieldDataLoadingOverheadUpperBound(/*max_memory_overhead=*/16 * MB);
+    EXPECT_EQ(memory_only.memory_bytes, 16 * MB);
+    EXPECT_EQ(memory_only.file_bytes, 0);
+
+    auto mmap_smaller_overhead = FieldDataLoadingOverheadUpperBound(
+        /*max_memory_overhead=*/4 * MB, std::optional<int64_t>{2 * MB});
+    EXPECT_EQ(mmap_smaller_overhead.memory_bytes, 8 * MB);
+    EXPECT_EQ(mmap_smaller_overhead.file_bytes, 8 * MB);
+
+    auto mmap_larger_file_overhead = FieldDataLoadingOverheadUpperBound(
+        /*max_memory_overhead=*/4 * MB, std::optional<int64_t>{32 * MB});
+    EXPECT_EQ(mmap_larger_file_overhead.memory_bytes, 8 * MB);
+    EXPECT_EQ(mmap_larger_file_overhead.file_bytes, 32 * MB);
+}
+
 // ---- LoadCellBatchAsync tests ----
 
 namespace {

internal/core/src/segcore/memory_planner.cpp

@@ -27,6 +27,7 @@
 #include <vector>
 
 #include "arrow/api.h"
+#include "cachinglayer/LoadingOverheadTracker.h"
 #include "common/Channel.h"
 #include "common/Common.h"
 #include "common/EasyAssert.h"
@@ -115,6 +116,24 @@ SetFieldDataMaxReadParallelism(int64_t parallelism) {
              FIELD_DATA_MAX_READ_PARALLELISM.load());
 }
 
+milvus::cachinglayer::ResourceUsage
+FieldDataLoadingOverheadUpperBound(int64_t max_memory_overhead,
+                                   std::optional<int64_t> max_file_overhead) {
+    auto budget_capacity = static_cast<int64_t>(
+        milvus::storage::TransientMemoryBudget::GetLoadTransientBudget()
+            .CapacityBytes());
+    if (budget_capacity == 0) {
+        return milvus::cachinglayer::LoadingOverheadTracker::kUnlimited;
+    }
+
+    auto memory_ub = std::max<int64_t>(budget_capacity, max_memory_overhead);
+    auto file_ub =
+        max_file_overhead.has_value()
+            ? std::max<int64_t>(budget_capacity, max_file_overhead.value())
+            : int64_t{0};
+    return {memory_ub, file_ub};
+}
+
 namespace {
 
 size_t

internal/core/src/segcore/memory_planner.h

@@ -22,9 +22,11 @@
 #include <functional>
 #include <future>
 #include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
+#include "cachinglayer/Utils.h"
 #include "common/Channel.h"
 #include "common/FieldData.h"
 #include "common/GroupChunk.h"
@@ -140,6 +142,11 @@ SetFieldDataReadWindowBytes(int64_t bytes);
 void
 SetFieldDataMaxReadParallelism(int64_t parallelism);
 
+milvus::cachinglayer::ResourceUsage
+FieldDataLoadingOverheadUpperBound(
+    int64_t max_memory_overhead,
+    std::optional<int64_t> max_file_overhead = std::nullopt);
+
 // A cell specification: identifies a cell's location within a specific file.
 struct CellSpec {
     int64_t cid;              // cell id

internal/core/src/segcore/storagev2translator/GroupChunkTranslator.cpp

@@ -23,6 +23,7 @@
 #include <filesystem>
 #include <limits>
 #include <memory>
+#include <optional>
 #include <string>
 #include <unordered_map>
 #include <vector>
@@ -218,16 +219,10 @@ GroupChunkTranslator::GroupChunkTranslator(
     if (!meta_.chunk_memory_size_.empty()) {
         int64_t max_cell_sz = *std::max_element(
             meta_.chunk_memory_size_.begin(), meta_.chunk_memory_size_.end());
-        auto budget_capacity = static_cast<int64_t>(
-            milvus::storage::TransientMemoryBudget::GetLoadTransientBudget()
-                .CapacityBytes());
         auto max_overhead_size = loading_overhead_bytes(max_cell_sz);
-        auto memory_ub = std::max<int64_t>(budget_capacity, max_overhead_size);
-        auto file_ub = use_mmap_
-                           ? std::max<int64_t>(budget_capacity, max_cell_sz)
-                           : int64_t{0};
-        auto upper_bound =
-            milvus::cachinglayer::ResourceUsage{memory_ub, file_ub};
+        auto upper_bound = milvus::segcore::FieldDataLoadingOverheadUpperBound(
+            max_overhead_size,
+            use_mmap_ ? std::optional<int64_t>{max_cell_sz} : std::nullopt);
         // Keep MCL reservation aligned with the process-wide transient load
         // budget rather than multiplying it by translator type.
         auto group = milvus::segcore::kLoadTransientOverheadGroup;

internal/core/src/segcore/storagev2translator/ManifestGroupTranslator.cpp

@@ -217,16 +217,10 @@ ManifestGroupTranslator::ManifestGroupTranslator(
     if (!meta_.chunk_memory_size_.empty()) {
         int64_t max_cell_sz = *std::max_element(
             meta_.chunk_memory_size_.begin(), meta_.chunk_memory_size_.end());
-        auto budget_capacity = static_cast<int64_t>(
-            milvus::storage::TransientMemoryBudget::GetLoadTransientBudget()
-                .CapacityBytes());
         auto max_overhead_size = loading_overhead_bytes(max_cell_sz);
-        auto memory_ub = std::max<int64_t>(budget_capacity, max_overhead_size);
-        auto file_ub = use_mmap_
-                           ? std::max<int64_t>(budget_capacity, max_cell_sz)
-                           : int64_t{0};
-        auto upper_bound =
-            milvus::cachinglayer::ResourceUsage{memory_ub, file_ub};
+        auto upper_bound = milvus::segcore::FieldDataLoadingOverheadUpperBound(
+            max_overhead_size,
+            use_mmap_ ? std::optional<int64_t>{max_cell_sz} : std::nullopt);
         // Keep MCL reservation aligned with the process-wide transient load
         // budget rather than multiplying it by translator type.
         auto group = milvus::segcore::kLoadTransientOverheadGroup;