memory control correction

ydb-platform · Nov 27, 2024 · b8e882a · b8e882a
1 parent 8793419
commit b8e882a
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 28 deletions.
diff --git a/ydb/core/tx/columnshard/engines/reader/common/result.h b/ydb/core/tx/columnshard/engines/reader/common/result.h
@@ -12,7 +12,7 @@ namespace NKikimr::NOlap::NReader {
 // Represents a batch of rows produced by ASC or DESC scan with applied filters and partial aggregation
 class TPartialReadResult: public TNonCopyable {
 private:
-    YDB_READONLY_DEF(std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>, ResourcesGuard);
+    YDB_READONLY_DEF(std::vector<std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>>, ResourceGuards);
     YDB_READONLY_DEF(std::shared_ptr<NGroupedMemoryManager::TGroupGuard>, GroupGuard);
     NArrow::TShardedRecordBatch ResultBatch;
 
@@ -54,11 +54,11 @@ class TPartialReadResult: public TNonCopyable {
         return ScanCursor;
     }
 
-    explicit TPartialReadResult(std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>&& resourcesGuard,
-        std::shared_ptr<NGroupedMemoryManager::TGroupGuard>&& gGuard, const NArrow::TShardedRecordBatch& batch,
+    explicit TPartialReadResult(const std::vector<std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>>& resourceGuards,
+        const std::shared_ptr<NGroupedMemoryManager::TGroupGuard>& gGuard, const NArrow::TShardedRecordBatch& batch,
         const std::shared_ptr<IScanCursor>& scanCursor, const std::optional<ui32> notFinishedIntervalIdx)
-        : ResourcesGuard(std::move(resourcesGuard))
-        , GroupGuard(std::move(gGuard))
+        : ResourceGuards(resourceGuards)
+        , GroupGuard(gGuard)
         , ResultBatch(batch)
         , ScanCursor(scanCursor)
         , NotFinishedIntervalIdx(notFinishedIntervalIdx) {
@@ -68,7 +68,7 @@ class TPartialReadResult: public TNonCopyable {
 
     explicit TPartialReadResult(const NArrow::TShardedRecordBatch& batch, const std::shared_ptr<IScanCursor>& scanCursor,
         const std::optional<ui32> notFinishedIntervalIdx)
-        : TPartialReadResult(nullptr, nullptr, batch, scanCursor, notFinishedIntervalIdx) {
+        : TPartialReadResult({}, nullptr, batch, scanCursor, notFinishedIntervalIdx) {
     }
 };
 

diff --git a/ydb/core/tx/columnshard/engines/reader/plain_reader/iterator/scanner.cpp b/ydb/core/tx/columnshard/engines/reader/plain_reader/iterator/scanner.cpp
@@ -27,7 +27,8 @@ void TScanHead::OnIntervalResult(std::shared_ptr<NGroupedMemoryManager::TAllocat
         } else {
             gGuard = itInterval->second->GetGroupGuard();
         }
-        AFL_VERIFY(ReadyIntervals.emplace(intervalIdx, std::make_shared<TPartialReadResult>(std::move(allocationGuard), std::move(gGuard), *newBatch, 
+        std::vector<std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>> guards = { std::move(allocationGuard) };
+        AFL_VERIFY(ReadyIntervals.emplace(intervalIdx, std::make_shared<TPartialReadResult>(guards, std::move(gGuard), *newBatch,
             std::make_shared<TPlainScanCursor>(lastPK), callbackIdxSubscriver)).second);
     } else {
         AFL_VERIFY(ReadyIntervals.emplace(intervalIdx, nullptr).second);

diff --git a/ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/context.cpp b/ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/context.cpp
@@ -96,27 +96,27 @@ class TColumnsAccumulator {
         const bool sequential) {
         auto actualColumns = columns - AssemblerReadyColumns;
         AssemblerReadyColumns = AssemblerReadyColumns + columns;
-        if (!actualColumns.IsEmpty()) {
-            auto actualSet = std::make_shared<TColumnsSet>(actualColumns.GetColumnIds(), FullSchema);
-            if (sequential) {
-                const auto notSequentialColumnIds = GuaranteeNotOptional->Intersect(*actualSet);
-                if (notSequentialColumnIds.size()) {
-                    script.Allocation(notSequentialColumnIds, stage, EMemType::Raw);
-                    std::shared_ptr<TColumnsSet> cross = actualSet->BuildSamePtr(notSequentialColumnIds);
-                    script.AddStep<TAssemblerStep>(cross, purposeId);
-                    *actualSet = *actualSet - *cross;
-                }
-                if (!actualSet->IsEmpty()) {
-                    script.Allocation(notSequentialColumnIds, stage, EMemType::RawSequential);
-                    script.AddStep<TOptionalAssemblerStep>(actualSet, purposeId);
-                }
-            } else {
-                script.Allocation(actualColumns.GetColumnIds(), stage, EMemType::Raw);
-                script.AddStep<TAssemblerStep>(actualSet, purposeId);
+        if (actualColumns.IsEmpty()) {
+            return false;
+        }
+        auto actualSet = std::make_shared<TColumnsSet>(actualColumns.GetColumnIds(), FullSchema);
+        if (sequential) {
+            const auto notSequentialColumnIds = GuaranteeNotOptional->Intersect(*actualSet);
+            if (notSequentialColumnIds.size()) {
+                script.Allocation(notSequentialColumnIds, stage, EMemType::Raw);
+                std::shared_ptr<TColumnsSet> cross = actualSet->BuildSamePtr(notSequentialColumnIds);
+                script.AddStep<TAssemblerStep>(cross, purposeId);
+                *actualSet = *actualSet - *cross;
             }
-            return true;
+            if (!actualSet->IsEmpty()) {
+                script.Allocation(notSequentialColumnIds, stage, EMemType::RawSequential);
+                script.AddStep<TOptionalAssemblerStep>(actualSet, purposeId);
+            }
+        } else {
+            script.Allocation(actualColumns.GetColumnIds(), stage, EMemType::Raw);
+            script.AddStep<TAssemblerStep>(actualSet, purposeId);
         }
-        return false;
+        return true;
     }
 };
 

diff --git a/ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/scanner.cpp b/ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/scanner.cpp
@@ -36,7 +36,8 @@ void TScanHead::OnSourceReady(const std::shared_ptr<IDataSource>& source, std::s
                 "source_idx", frontSource->GetSourceIdx())("table", table->num_rows());
             auto cursor =
                 std::make_shared<TSimpleScanCursor>(frontSource->GetStartPKRecordBatch(), frontSource->GetSourceId(), startIndex + recordsCount);
-            reader.OnIntervalResult(std::make_shared<TPartialReadResult>(nullptr, nullptr, table, cursor, sourceIdxToContinue));
+            reader.OnIntervalResult(
+                std::make_shared<TPartialReadResult>(frontSource->GetResourceGuards(), frontSource->GetGroupGuard(), table, cursor, sourceIdxToContinue));
         } else if (sourceIdxToContinue) {
             AFL_DEBUG(NKikimrServices::TX_COLUMNSHARD)("event", "continue_source")("source_id", frontSource->GetSourceId())(
                 "source_idx", frontSource->GetSourceIdx());

diff --git a/ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/source.h b/ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/source.h
@@ -60,7 +60,7 @@ class IDataSource: public ICursorEntity {
     YDB_READONLY(bool, HasDeletions, false);
     virtual NJson::TJsonValue DoDebugJson() const = 0;
     std::shared_ptr<TFetchingScript> FetchingPlan;
-    std::vector<std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>> ResourceGuards;
+    YDB_READONLY_DEF(std::vector<std::shared_ptr<NGroupedMemoryManager::TAllocationGuard>>, ResourceGuards);
     YDB_READONLY(TPKRangeFilter::EUsageClass, UsageClass, TPKRangeFilter::EUsageClass::PartialUsage);
     bool ProcessingStarted = false;
     bool IsStartedByCursor = false;
@@ -99,6 +99,11 @@ class IDataSource: public ICursorEntity {
         return IsStartedByCursor;
     }
 
+    const std::shared_ptr<NGroupedMemoryManager::TGroupGuard>& GetGroupGuard() const {
+        AFL_VERIFY(SourceGroupGuard);
+        return SourceGroupGuard;
+    }
+
     ui64 GetMemoryGroupId() const {
         AFL_VERIFY(SourceGroupGuard);
         return SourceGroupGuard->GetGroupId();