arangodb · Dronplane · Aug 14, 2023 · Apr 25, 2023 · Apr 26, 2023 · Jul 18, 2023
diff --git a/arangod/Aql/IResearchViewExecutor.h b/arangod/Aql/IResearchViewExecutor.h
@@ -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"
@@ -88,8 +89,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 +133,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 +161,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 +247,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 +280,25 @@ 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 {
+    _heapSort = s;
+    TRI_ASSERT(_heapUsedStoredColumns.empty());
+    TRI_ASSERT(_heapOnlyColumnsCount == 0);
+    for (auto const& c : _heapSort) {
+      if (c.isScore()) {
+        continue;
+      }
+      auto it = storedValues.find(c.source);
+      if (it != storedValues.end()) {
+        _heapUsedStoredColumns.emplace(
+            c.source,
+            static_cast<size_t>(std::distance(storedValues.begin(), it)));
+      } else {
+        ++_heapOnlyColumnsCount;
+      }
+    }
   }
 
   irs::score_t* pushNoneScores(size_t count);
@@ -288,7 +311,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 +340,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) +
-        maxSize * sizeof(typename decltype(_searchDocs)::value_type) +
-        maxSize * sizeof(typename decltype(_storedValuesBuffer)::value_type);
-    if (!_scoresSort.empty()) {
+        scores * maxSize * sizeof(typename decltype(_scoreBuffer)::value_type) +
+        maxSize * sizeof(typename decltype(_searchDocs)::value_type);
+    if (!_heapSort.empty()) {
+      res += stored * maxSize *
+             sizeof(typename decltype(_storedValuesBuffer)::value_type);
       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 +364,21 @@ 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);
       }
     }
     _maxSize = atMost;
@@ -390,14 +427,44 @@ 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;
+  StoredValuesContainer _storedValuesBuffer;  // FIXME: ensure atMost rows are
+                                              // allocated from the beginning!
+
+  // Heap Sort facilities
+  // maps stored column to index in _storedValuesBuffer
+  containers::FlatHashMap<ptrdiff_t, size_t> _heapUsedStoredColumns;
+  // atMost * <num heap only values>
+  StoredValuesContainer _heapOnlyStoredValuesBuffer;
+  // <num heap only values>
+  std::vector<ColumnIterator> _heapOnlyStoredValuesReaders;
+  // <num heap only values>
+  StoredValuesContainer _currentDocumentBuffer;
+  using DocumentSlicesContainer =
+      typename std::conditional<copyStored, bool,
+                                std::vector<velocypack::Slice>>::type;
+  DocumentSlicesContainer _currentDocumentSlices;
+  std::vector<HeapSortValue> _heapSortValues;
+  std::span<iresearch::HeapSortElement const> _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 +652,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: