non sorted collection usage (#16804)

Author: ivanmorozov333

ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/collections.cpp

@@ -92,7 +92,7 @@ ISourcesCollection::ISourcesCollection(const std::shared_ptr<TSpecialReadContext
     }
 }
 
-std::shared_ptr<NKikimr::NOlap::IScanCursor> TNotSortedFullScanCollection::DoBuildCursor(
+std::shared_ptr<NKikimr::NOlap::IScanCursor> TNotSortedCollection::DoBuildCursor(
     const std::shared_ptr<IDataSource>& source, const ui32 readyRecords) const {
     return std::make_shared<TNotSortedSimpleScanCursor>(source->GetSourceId(), readyRecords);
 }

ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/collections.h

@@ -46,6 +46,7 @@ class ISourcesCollection {
     }
 
     void OnSourceFinished(const std::shared_ptr<IDataSource>& source) {
+        AFL_VERIFY(source);
         SourcesInFlightCount.Dec();
         DoOnSourceFinished(source);
     }
@@ -61,11 +62,14 @@ class ISourcesCollection {
     ISourcesCollection(const std::shared_ptr<TSpecialReadContext>& context);
 };
 
-class TNotSortedFullScanCollection: public ISourcesCollection {
+class TNotSortedCollection: public ISourcesCollection {
 private:
     using TBase = ISourcesCollection;
+    std::optional<ui32> Limit;
+    ui32 InFlightLimit = 1;
     std::deque<TSourceConstructor> Sources;
     TPositiveControlInteger InFlightCount;
+    ui32 FetchedCount = 0;
     virtual void DoClear() override {
         Sources.clear();
     }
@@ -81,16 +85,31 @@ class TNotSortedFullScanCollection: public ISourcesCollection {
         return result;
     }
     virtual bool DoCheckInFlightLimits() const override {
-        return InFlightCount < GetMaxInFlight();
+        return InFlightCount < InFlightLimit;
     }
-    virtual void DoOnSourceFinished(const std::shared_ptr<IDataSource>& /*source*/) override {
+    virtual void DoOnSourceFinished(const std::shared_ptr<IDataSource>& source) override {
+        if (!source->GetResultRecordsCount() && InFlightLimit * 2 < GetMaxInFlight()) {
+            InFlightLimit *= 2;
+        }
+        FetchedCount += source->GetResultRecordsCount();
+        if (Limit && *Limit <= FetchedCount && Sources.size()) {
+            AFL_NOTICE(NKikimrServices::TX_COLUMNSHARD)("event", "limit_exhausted")("limit", Limit)("fetched", FetchedCount);
+            Sources.clear();
+        }
         InFlightCount.Dec();
     }
 
 public:
-    TNotSortedFullScanCollection(const std::shared_ptr<TSpecialReadContext>& context, std::deque<TSourceConstructor>&& sources,
-        const std::shared_ptr<IScanCursor>& cursor)
-        : TBase(context) {
+    TNotSortedCollection(const std::shared_ptr<TSpecialReadContext>& context, std::deque<TSourceConstructor>&& sources,
+        const std::shared_ptr<IScanCursor>& cursor, const std::optional<ui32> limit)
+        : TBase(context)
+        , Limit(limit)
+    {
+        if (Limit) {
+            InFlightLimit = 1;
+        } else {
+            InFlightLimit = GetMaxInFlight();
+        }
         if (cursor && cursor->IsInitialized()) {
             while (sources.size()) {
                 bool usage = false;

ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/plain_read_data.cpp

@@ -21,7 +21,6 @@ TPlainReadData::TPlainReadData(const std::shared_ptr<TReadContext>& context)
 
         sources.emplace_back(TSourceConstructor(sourceIdx++, i, context));
     }
-    std::make_heap(sources.begin(), sources.end());
     Scanner = std::make_shared<TScanHead>(std::move(sources), SpecialReadContext);
 
     auto& stats = GetReadMetadata()->ReadStats;

ydb/core/tx/columnshard/engines/reader/simple_reader/iterator/scanner.cpp

@@ -66,15 +66,17 @@ TConclusionStatus TScanHead::Start() {
 
 TScanHead::TScanHead(std::deque<TSourceConstructor>&& sources, const std::shared_ptr<TSpecialReadContext>& context)
     : Context(context) {
-    if (Context->GetReadMetadata()->HasLimit()) {
-        SourcesCollection =
-            std::make_unique<TScanWithLimitCollection>(Context, std::move(sources), context->GetCommonContext()->GetScanCursor());
-    } else if (Context->GetReadMetadata()->IsSorted()) {
-        SourcesCollection =
-            std::make_unique<TSortedFullScanCollection>(Context, std::move(sources), context->GetCommonContext()->GetScanCursor());
+    if (Context->GetReadMetadata()->IsSorted()) {
+        if (Context->GetReadMetadata()->HasLimit()) {
+            SourcesCollection =
+                std::make_unique<TScanWithLimitCollection>(Context, std::move(sources), context->GetCommonContext()->GetScanCursor());
+        } else {
+            SourcesCollection =
+                std::make_unique<TSortedFullScanCollection>(Context, std::move(sources), context->GetCommonContext()->GetScanCursor());
+        }
     } else {
-        SourcesCollection =
-            std::make_unique<TNotSortedFullScanCollection>(Context, std::move(sources), context->GetCommonContext()->GetScanCursor());
+        SourcesCollection = std::make_unique<TNotSortedCollection>(
+            Context, std::move(sources), context->GetCommonContext()->GetScanCursor(), Context->GetReadMetadata()->GetLimitRobustOptional());
     }
 }