Merge de229f8 into 264201f

Author: ivanmorozov333

ydb/core/tx/columnshard/engines/reader/plain_reader/fetched_data.h

@@ -109,4 +109,16 @@ class TFetchedData {
 
 };
 
+class TFetchedResult {
+private:
+    YDB_READONLY_DEF(std::shared_ptr<arrow::RecordBatch>, Batch);
+    YDB_READONLY_DEF(std::shared_ptr<NArrow::TColumnFilter>, NotAppliedFilter);
+public:
+    TFetchedResult(std::unique_ptr<TFetchedData>&& data)
+        : Batch(data->GetBatch())
+        , NotAppliedFilter(data->GetNotAppliedFilter()) {
+
+    }
+};
+
 }

ydb/core/tx/columnshard/engines/reader/plain_reader/fetching.cpp

@@ -16,18 +16,21 @@ bool TStepAction::DoApply(IDataReader& /*owner*/) const {
 bool TStepAction::DoExecute() {
     while (Step) {
         if (Source->IsEmptyData()) {
+            Source->Finalize();
             FinishedFlag = true;
             return true;
         }
         if (!Step->ExecuteInplace(Source, Step)) {
             return true;
         }
         if (Source->IsEmptyData()) {
+            Source->Finalize();
             FinishedFlag = true;
             return true;
         }
         Step = Step->GetNextStep();
     }
+    Source->Finalize();
     FinishedFlag = true;
     return true;
 }

ydb/core/tx/columnshard/engines/reader/plain_reader/interval.cpp

@@ -60,7 +60,7 @@ class TMergeTask: public NColumnShard::IDataTasksProcessor::ITask {
     }
     virtual bool DoExecute() override {
         if (MergingContext->IsExclusiveInterval()) {
-            ResultBatch = Sources.begin()->second->GetStageData().GetBatch();
+            ResultBatch = Sources.begin()->second->GetStageResult().GetBatch();
             if (ResultBatch && ResultBatch->num_rows()) {
                 LastPK = Sources.begin()->second->GetLastPK();
                 ResultBatch = NArrow::ExtractColumnsValidate(ResultBatch, Context->GetProgramInputColumns()->GetColumnNamesVector());
@@ -81,8 +81,8 @@ class TMergeTask: public NColumnShard::IDataTasksProcessor::ITask {
         }
         std::shared_ptr<NIndexedReader::TMergePartialStream> merger = Context->BuildMerger();
         for (auto&& [_, i] : Sources) {
-            if (auto rb = i->GetStageData().GetBatch()) {
-                merger->AddSource(rb, i->GetStageData().GetNotAppliedFilter());
+            if (auto rb = i->GetStageResult().GetBatch()) {
+                merger->AddSource(rb, i->GetStageResult().GetNotAppliedFilter());
             }
         }
         AFL_VERIFY(merger->GetSourcesCount() <= Sources.size());

ydb/core/tx/columnshard/engines/reader/plain_reader/source.cpp

@@ -15,7 +15,7 @@ namespace NKikimr::NOlap::NPlainReader {
 void IDataSource::InitFetchingPlan(const std::shared_ptr<IFetchingStep>& fetchingFirstStep, const std::shared_ptr<IDataSource>& sourcePtr, const bool isExclusive) {
     AFL_VERIFY(fetchingFirstStep);
     if (AtomicCas(&FilterStageFlag, 1, 0)) {
-        StageData = std::make_shared<TFetchedData>(isExclusive);
+        StageData = std::make_unique<TFetchedData>(isExclusive);
         AFL_DEBUG(NKikimrServices::TX_COLUMNSHARD_SCAN)("InitFetchingPlan", fetchingFirstStep->DebugString())("source_idx", SourceIdx);
         NActors::TLogContextGuard logGuard(NActors::TLogContextBuilder::Build()("source", SourceIdx)("method", "InitFetchingPlan"));
         if (IsAborted()) {

ydb/core/tx/columnshard/engines/reader/plain_reader/source.h

@@ -36,7 +36,8 @@ class IDataSource {
 protected:
     THashMap<ui32, TFetchingInterval*> Intervals;
 
-    std::shared_ptr<TFetchedData> StageData;
+    std::unique_ptr<TFetchedData> StageData;
+    std::unique_ptr<TFetchedResult> StageResult;
 
     TAtomic FilterStageFlag = 0;
     bool IsReadyFlag = false;
@@ -51,8 +52,17 @@ class IDataSource {
     virtual void DoAbort() = 0;
     virtual void DoApplyIndex(const NIndexes::TIndexCheckerContainer& indexMeta) = 0;
 public:
+    const TFetchedResult& GetStageResult() const {
+        AFL_VERIFY(!!StageResult);
+        return *StageResult;
+    }
+
     void SetIsReady();
 
+    void Finalize() {
+        StageResult = std::make_unique<TFetchedResult>(std::move(StageData));
+    }
+
     bool IsEmptyData() const {
         return GetStageData().IsEmpty();
     }