arangodb
diff --git a/‎arangod/Aql/IResearchViewExecutor.h
Lines changed: 103 additions & 17 deletions b/‎arangod/Aql/IResearchViewExecutor.h
Lines changed: 103 additions & 17 deletions
@@ -30,6 +30,7 @@
 #include "Aql/AqlFunctionsInternalCache.h"
 #include "Aql/RegisterInfos.h"
 #include "Aql/VarInfoMap.h"
+#include "Containers/FlatHashSet.h"
 #include "IResearch/ExpressionFilter.h"
 #include "IResearch/IResearchFilterFactory.h"
 #include "IResearch/IResearchExpressionContext.h"
@@ -45,6 +46,7 @@
 
 #include <formats/formats.hpp>
 #include <index/heap_iterator.hpp>
+#include <utils/empty.hpp>
 
 #include <utility>
 #include <variant>
@@ -88,8 +90,8 @@ class IResearchViewExecutorInfos {
       iresearch::IResearchViewNode::ViewValuesRegisters&&
           outNonMaterializedViewRegs,
       iresearch::CountApproximate, iresearch::FilterOptimization,
-      std::vector<std::pair<size_t, bool>> scorersSort, size_t scorersSortLimit,
-      iresearch::SearchMeta const* meta);
+      std::vector<iresearch::HeapSortElement> const& heapSort,
+      size_t heapSortLimit, iresearch::SearchMeta const* meta);
 
   auto getDocumentRegister() const noexcept -> RegisterId;
 
@@ -132,9 +134,9 @@ class IResearchViewExecutorInfos {
 
   iresearch::IResearchViewStoredValues const& storedValues() const noexcept;
 
-  size_t scoresSortLimit() const noexcept { return _scorersSortLimit; }
+  size_t heapSortLimit() const noexcept { return _heapSortLimit; }
 
-  auto scoresSort() const noexcept { return std::span{_scorersSort}; }
+  auto heapSort() const noexcept { return std::span{_heapSort}; }
 
   size_t scoreRegistersCount() const noexcept { return _scoreRegistersCount; }
 
@@ -160,8 +162,8 @@ class IResearchViewExecutorInfos {
   iresearch::IResearchViewNode::ViewValuesRegisters _outNonMaterializedViewRegs;
   iresearch::CountApproximate _countApproximate;
   iresearch::FilterOptimization _filterOptimization;
-  std::vector<std::pair<size_t, bool>> _scorersSort;
-  size_t _scorersSortLimit;
+  std::vector<iresearch::HeapSortElement> const& _heapSort;
+  size_t _heapSortLimit;
   iresearch::SearchMeta const* _meta;
   int const _depth;
   bool _filterConditionIsEmpty;
@@ -246,6 +248,11 @@ class ScoreIterator {
   size_t _numScores;
 };
 
+union HeapSortValue {
+  irs::score_t score;
+  velocypack::Slice slice;
+};
+
 // Holds and encapsulates the data read from the iresearch index.
 template<typename ValueType, bool copyStored>
 class IndexReadBuffer {
@@ -274,8 +281,35 @@ class IndexReadBuffer {
 
   ScoreIterator getScores(IndexReadBufferEntry bufferEntry) noexcept;
 
-  void setScoresSort(std::span<std::pair<size_t, bool> const> s) noexcept {
-    _scoresSort = s;
+  void setHeapSort(std::span<iresearch::HeapSortElement const> s,
+                   iresearch::IResearchViewNode::ViewValuesRegisters const&
+                       storedValues) noexcept {
+    TRI_ASSERT(_heapSort.empty());
+    TRI_ASSERT(_heapOnlyColumnsCount == 0);
+    auto const storedValuesCount = storedValues.size();
+    size_t j = 0;
+    for (auto const& c : s) {
+      auto& sort = _heapSort.emplace_back(BufferHeapSortElement{c});
+      if (c.isScore()) {
+        continue;
+      }
+      auto it = storedValues.find(c.source);
+      if (it != storedValues.end()) {
+        sort.container = &_storedValuesBuffer;
+        sort.multiplier = storedValuesCount;
+        sort.offset =
+            static_cast<size_t>(std::distance(storedValues.begin(), it));
+      } else {
+        sort.container = &_heapOnlyStoredValuesBuffer;
+        sort.offset = j++;
+        ++_heapOnlyColumnsCount;
+      }
+    }
+    for (auto& c : _heapSort) {
+      if (c.container == &_heapOnlyStoredValuesBuffer) {
+        c.multiplier = _heapOnlyColumnsCount;
+      }
+    }
   }
 
   irs::score_t* pushNoneScores(size_t count);
@@ -288,7 +322,9 @@ class IndexReadBuffer {
     _searchDocs.emplace_back(segment, docId);
   }
 
-  void pushSortedValue(StorageSnapshot const& snapshot, ValueType&& value,
+  template<typename ColumnReaderProvider>
+  void pushSortedValue(ColumnReaderProvider& columnReader,
+                       StorageSnapshot const& snapshot, ValueType&& value,
                        std::span<float_t const> scores, irs::score& score,
                        irs::score_t& threshold);
 
@@ -315,14 +351,22 @@ class IndexReadBuffer {
   // before and after.
   void assertSizeCoherence() const noexcept;
 
-  size_t memoryUsage(size_t maxSize) const noexcept {
+  size_t memoryUsage(size_t maxSize, size_t scores,
+                     size_t stored) const noexcept {
     auto res =
         maxSize * sizeof(typename decltype(_keyBuffer)::value_type) +
-        maxSize * sizeof(typename decltype(_scoreBuffer)::value_type) +
+        scores * maxSize * sizeof(typename decltype(_scoreBuffer)::value_type) +
         maxSize * sizeof(typename decltype(_searchDocs)::value_type) +
-        maxSize * sizeof(typename decltype(_storedValuesBuffer)::value_type);
-    if (!_scoresSort.empty()) {
+        stored * maxSize *
+            sizeof(typename decltype(_storedValuesBuffer)::value_type);
+    if (!_heapSort.empty()) {
       res += maxSize * sizeof(typename decltype(_rows)::value_type);
+      res += _heapOnlyColumnsCount * maxSize *
+             sizeof(typename decltype(_heapOnlyStoredValuesBuffer)::value_type);
+      res += _heapOnlyColumnsCount *
+             sizeof(typename decltype(_currentDocumentBuffer)::value_type);
+      res += (maxSize * _heapSort.size()) *
+             sizeof(typename decltype(_heapSortValues)::value_type);
     }
     return res;
   }
@@ -331,17 +375,27 @@ class IndexReadBuffer {
                                      size_t stored) {
     TRI_ASSERT(_storedValuesBuffer.empty());
     if (_keyBuffer.capacity() < atMost) {
-      auto newMemoryUsage = memoryUsage(atMost);
+      auto newMemoryUsage = memoryUsage(atMost, scores, stored);
       auto tracked = _memoryTracker.tracked();
       if (newMemoryUsage > tracked) {
         _memoryTracker.increase(newMemoryUsage - tracked);
       }
       _keyBuffer.reserve(atMost);
       _searchDocs.reserve(atMost);
       _scoreBuffer.reserve(atMost * scores);
-      _storedValuesBuffer.reserve(atMost * stored);
-      if (!_scoresSort.empty()) {
+      if (!_heapSort.empty()) {
         _rows.reserve(atMost);
+        _currentDocumentBuffer.reserve(_heapSort.size());
+        // resize is important here as we want
+        // indexed access for setting values
+        _storedValuesBuffer.resize(atMost * stored);
+        _heapOnlyStoredValuesBuffer.resize(_heapOnlyColumnsCount * atMost);
+        if (!scores) {
+          // save ourselves 1 "if" during pushing values
+          _scoreBuffer.push_back(irs::score_t{});
+        }
+      } else {
+        _storedValuesBuffer.reserve(atMost * stored);
       }
     }
     _maxSize = atMost;
@@ -370,6 +424,12 @@ class IndexReadBuffer {
 
   StoredValuesContainer const& getStoredValues() const noexcept;
 
+  struct BufferHeapSortElement : iresearch::HeapSortElement {
+    StoredValuesContainer* container{};
+    size_t offset{};
+    size_t multiplier{};
+  };
+
  private:
   // _keyBuffer, _scoreBuffer, _storedValuesBuffer together hold all the
   // information read from the iresearch index.
@@ -390,14 +450,39 @@ class IndexReadBuffer {
     StorageSnapshot const* snapshot;
     ValueType value;
   };
+
+  template<typename ColumnReaderProvider>
+  velocypack::Slice readHeapSortColumn(iresearch::HeapSortElement const& cmp,
+                                       irs::doc_id_t doc,
+                                       ColumnReaderProvider& readerProvider,
+                                       size_t readerSlot);
+  template<bool fullHeap, typename ColumnReaderProvider>
+  void finalizeHeapSortDocument(size_t idx, irs::doc_id_t,
+                                std::span<float_t const> scores,
+                                ColumnReaderProvider& readerProvider);
+
   std::vector<BufferValueType> _keyBuffer;
   // FIXME(gnusi): compile time
   std::vector<iresearch::SearchDoc> _searchDocs;
   std::vector<irs::score_t> _scoreBuffer;
   StoredValuesContainer _storedValuesBuffer;
+  // Heap Sort facilities
+  // atMost * <num heap only values>
+  StoredValuesContainer _heapOnlyStoredValuesBuffer;
+  // <num heap only values>
+  std::vector<ColumnIterator> _heapOnlyStoredValuesReaders;
+  // <num heap only values>
+  StoredValuesContainer _currentDocumentBuffer;
+  IRS_NO_UNIQUE_ADDRESS
+  irs::utils::Need<!copyStored, std::vector<velocypack::Slice>>
+    
A2F4
  _currentDocumentSlices;
+  std::vector<HeapSortValue> _heapSortValues;
+  std::vector<BufferHeapSortElement> _heapSort;
+  size_t _heapOnlyColumnsCount{0};
+  size_t _currentReaderOffset{std::numeric_limits<size_t>::max()};
+
   size_t _numScoreRegisters;
   size_t _keyBaseIdx;
-  std::span<std::pair<size_t, bool> const> _scoresSort;
   std::vector<size_t> _rows;
   size_t _maxSize;
   size_t _heapSizeLeft;
@@ -585,6 +670,7 @@ class IResearchViewExecutorBase {
   irs::Scorers _scorers;
   irs::WandContext _wand;
   std::vector<ColumnIterator> _storedValuesReaders;
+  containers::FlatHashSet<ptrdiff_t> _storedColumnsMask;
   std::array<char, arangodb::iresearch::kSearchDocBufSize> _buf;
   bool _isInitialized;
   // new mangling only: