feat: Add listRowsFast in RowContainer (facebookincubator#15698)

Summary:
Adds an optional `useListRowIndex` flag and listRowsFast method in `RowContainer`.
When enabled, RowContainer maintains a row-pointer index (`rowPointers_`) and
uses `listRowsFast()` to iterate rows without scanning allocation ranges or checking
free/probe flags. It is intended to be used in SortBuffer and SortInputSpiller to improve performance.

Pull Request resolved: facebookincubator#15698

Reviewed By: Yuhta

Differential Revision: D88497347

Pulled By: zacw7

fbshipit-source-id: 4cec8eeb6b49fd7a80e97933cb0dd2da038a8a70

Author: duanmeng

velox/exec/GroupingSet.cpp

@@ -1102,6 +1102,7 @@ bool GroupingSet::getOutputWithSpill(
           false,
           false,
           false,
+          false,
           pool_);
 
       initializeAggregates(aggregates_, *mergeRows_, false);
@@ -1437,6 +1438,7 @@ void GroupingSet::abandonPartialAggregation() {
       false,
       false,
       false,
+      false,
       pool_);
   initializeAggregates(aggregates_, *intermediateRows_, true);
   table_.reset();

velox/exec/HashTable.cpp

@@ -75,6 +75,7 @@ HashTable<ignoreNullKeys>::HashTable(
       isJoinBuild,
       hasProbedFlag,
       hashMode_ != HashMode::kHash,
+      /*useListRowIndex=*/false,
       pool);
   nextOffset_ = rows_->nextOffset();
 }

velox/exec/RowContainer.cpp

@@ -136,14 +136,17 @@ RowContainer::RowContainer(
     bool isJoinBuild,
     bool hasProbedFlag,
     bool hasNormalizedKeys,
+    bool useListRowIndex,
     memory::MemoryPool* pool)
     : keyTypes_(keyTypes),
       nullableKeys_(nullableKeys),
       isJoinBuild_(isJoinBuild),
       hasNormalizedKeys_(hasNormalizedKeys),
+      useListRowIndex_(useListRowIndex),
       stringAllocator_(std::make_unique<HashStringAllocator>(pool)),
       accumulators_(accumulators),
-      rows_(pool) {
+      rows_(pool),
+      rowPointers_(StlAllocator<char*>(stringAllocator_.get())) {
   // Compute the layout of the payload row.  The row has keys, null flags,
   // accumulators, dependent fields. All fields are fixed width. If variable
   // width data is referenced, this is done with StringView(for VARCHAR) and
@@ -279,6 +282,10 @@ char* RowContainer::newRow() {
     if (normalizedKeySize_) {
       ++numRowsWithNormalizedKey_;
     }
+
+    if (useListRowIndex_) {
+      rowPointers_.push_back(row);
+    }
   }
   return initializeRow(row, false /* reuse */);
 }
@@ -964,6 +971,8 @@ void RowContainer::clear() {
   hasDuplicateRows_ = false;
 
   rows_.clear();
+  rowPointers_.clear();
+  rowPointers_.shrink_to_fit();
   stringAllocator_->clear();
   numRows_ = 0;
   numRowsWithNormalizedKey_ = 0;
@@ -1025,7 +1034,8 @@ std::optional<int64_t> RowContainer::estimateRowSize() const {
   }
   int64_t freeBytes = rows_.freeBytes() + fixedRowSize_ * numFreeRows_;
   int64_t usedSize = rows_.allocatedBytes() - freeBytes +
-      stringAllocator_->retainedSize() - stringAllocator_->freeSpace();
+      stringAllocator_->retainedSize() - stringAllocator_->freeSpace() -
+      rowPointers_.capacity() * sizeof(char*);
   int64_t rowSize = usedSize / numRows_;
   VELOX_CHECK_GT(
       rowSize, 0, "Estimated row size of the RowContainer must be positive.");

velox/exec/RowContainer.h

@@ -83,6 +83,8 @@ struct RowContainerIterator {
   char* rowBegin{nullptr};
   /// First byte after the end of the range containing 'currentRow'.
   char* endOfRun{nullptr};
+  /// Cursor of the list row operation.
+  int32_t listRowCursor{0};
 
   /// Returns the current row, skipping a possible normalized key below the
   /// first byte of row.
@@ -273,6 +275,21 @@ class RowContainer {
       const std::vector<TypePtr>& keyTypes,
       const std::vector<TypePtr>& dependentTypes,
       memory::MemoryPool* pool)
+      : RowContainer(
+            keyTypes,
+            dependentTypes,
+            /*useListRowIndex=*/false,
+            pool) {}
+
+  /// If 'useListRowIndex' is true, the container maintains an internal array of
+  /// row pointers so that listRowsFast() can return rows without scanning
+  /// underlying allocations or checking free/probe flags. It is intended to be
+  /// used in SortBuffer and SortInputSpiller to improve performance.
+  RowContainer(
+      const std::vector<TypePtr>& keyTypes,
+      const std::vector<TypePtr>& dependentTypes,
+      bool useListRowIndex,
+      memory::MemoryPool* pool)
       : RowContainer(
             keyTypes,
             true, // nullableKeys
@@ -282,6 +299,7 @@ class RowContainer {
             false, // isJoinBuild
             false, // hasProbedFlag
             false, // hasNormalizedKey
+            useListRowIndex,
             pool) {}
 
   ~RowContainer();
@@ -313,6 +331,7 @@ class RowContainer {
       bool isJoinBuild,
       bool hasProbedFlag,
       bool hasNormalizedKey,
+      bool useListRowIndex,
       memory::MemoryPool* pool);
 
   /// Allocates a new row and initializes possible aggregates to null.
@@ -637,6 +656,20 @@ class RowContainer {
     return count;
   }
 
+  /// Fast path for `listRows` that returns `rowPointers_` directly. Used by
+  /// `SortBuffer` and `SortInputSpiller`, so it skips checking the free and
+  /// probe flags.
+  int32_t listRowsFast(RowContainerIterator* iter, int32_t maxRows, char** rows)
+      const {
+    int32_t count = 0;
+    while (count < maxRows && iter->listRowCursor < rowPointers_.size()) {
+      char* row = rowPointers_[iter->listRowCursor];
+      rows[count++] = row;
+      ++iter->listRowCursor;
+    }
+    return count;
+  }
+
   /// Extracts up to 'maxRows' rows starting at the position of 'iter'. A
   /// default constructed or reset iter starts at the beginning. Returns the
   /// number of rows written to 'rows'. Returns 0 when at end. Stops after the
@@ -651,6 +684,9 @@ class RowContainer {
 
   int32_t listRows(RowContainerIterator* iter, int32_t maxRows, char** rows)
       const {
+    if (useListRowIndex_) {
+      return listRowsFast(iter, maxRows, rows);
+    }
     return listRows<ProbeType::kAll>(iter, maxRows, kUnlimited, rows);
   }
 
@@ -791,6 +827,10 @@ class RowContainer {
     return 0;
   }
 
+  const std::vector<char*, StlAllocator<char*>>& testingRowPointers() const {
+    return rowPointers_;
+  }
+
   memory::MemoryPool* pool() const {
     return stringAllocator_->pool();
   }
@@ -1471,7 +1511,8 @@ class RowContainer {
   const bool isJoinBuild_;
   // True if normalized keys are enabled in initial state.
   const bool hasNormalizedKeys_;
-
+  // True if use 'listRowsFast'.
+  const bool useListRowIndex_;
   const std::unique_ptr<HashStringAllocator> stringAllocator_;
 
   // Indicates if we can add new row to this row container. It is set to false
@@ -1527,6 +1568,7 @@ class RowContainer {
   uint64_t numFreeRows_ = 0;
 
   memory::AllocationPool rows_;
+  std::vector<char*, StlAllocator<char*>> rowPointers_;
 
   int alignment_ = 1;
 

velox/exec/SortBuffer.cpp

@@ -74,7 +74,7 @@ SortBuffer::SortBuffer(
   }
 
   data_ = std::make_unique<RowContainer>(
-      sortedColumnTypes, nonSortedColumnTypes, pool_);
+      sortedColumnTypes, nonSortedColumnTypes, /*useListRowIndex=*/true, pool_);
   spillerStoreType_ =
       ROW(std::move(sortedSpillColumnNames), std::move(sortedSpillColumnTypes));
 }
@@ -128,6 +128,7 @@ void SortBuffer::noMoreInput() {
     updateEstimatedOutputRowSize();
     // Sort the pointers to the rows in RowContainer (data_) instead of sorting
     // the rows.
+    // TODO: Reuse 'RowContainer::rowPointers_'.
     sortedRows_.resize(numInputRows_);
     RowContainerIterator iter;
     data_->listRows(&iter, numInputRows_, sortedRows_.data());

velox/exec/Spiller.cpp

@@ -101,8 +101,9 @@ bool SpillerBase::fillSpillRuns(RowContainerIterator* iterator) {
 
     uint64_t totalRows{0};
     for (;;) {
-      const auto numRows = container_->listRows(
-          iterator, rows.size(), RowContainer::kUnlimited, rows.data());
+      // TODO: Reuse 'RowContainer::rowPointers_'.
+      const auto numRows =
+          container_->listRows(iterator, rows.size(), rows.data());
       if (numRows == 0) {
         lastRun = true;
         break;

velox/exec/StreamingAggregation.cpp

@@ -425,6 +425,7 @@ std::unique_ptr<RowContainer> StreamingAggregation::makeRowContainer(
       false,
       false,
       false,
+      false,
       pool());
 }
 

velox/exec/tests/AggregateSpillBenchmarkBase.cpp

@@ -36,10 +36,11 @@ std::unique_ptr<RowContainer> makeRowContainer(
       true, // nullableKeys
       std::vector<Accumulator>{},
       dependentTypes,
-      false, // hasNext
-      false, // isJoinBuild
-      false, // hasProbedFlag
-      false, // hasNormalizedKey
+      /*hasNext=*/false,
+      /*isJoinBuild=*/false,
+      /*hasProbedFlag=*/false,
+      /*hasNormalizedKey=*/false,
+      /*useListRowIndex=*/false,
       pool.get());
 }
 

velox/exec/tests/AggregationTest.cpp

@@ -384,6 +384,7 @@ class AggregationTest : public OperatorTestBase {
         false,
         true,
         true,
+        false,
         pool_.get());
   }
 
@@ -4028,6 +4029,7 @@ TEST_F(AggregationTest, destroyAfterPartialInitialization) {
       false, // isJoinBuild
       false, // hasProbedFlag
       false, // hasNormalizedKeys
+      false, // useListRowIndex
       pool());
   const auto rowColumn = rows.columnAt(0);
   agg.setOffsets(