feat: impl sort key for LocalShuffleReader (#26620)

Summary:

Implement sorted shuffle k-way merge for LocalShuffleReader, when it's sortedShuffle.

Added k-way merge support using TreeOfLosers to efficiently merge multiple sorted shuffle files. The reader streams data from sorted files and returns merged results in sorted order.

Differential Revision: D86888221

Author: duxiao1212

presto-native-execution/presto_cpp/main/operators/LocalShuffle.cpp

@@ -113,8 +113,9 @@ std::vector<RowMetadata> extractAndSortRowMetadata(
           const char* lhsKey = buffer + lhs.rowStart + (kUint32Size * 2);
           const char* rhsKey = buffer + rhs.rowStart + (kUint32Size * 2);
           return compareKeys(
-              std::string_view(lhsKey, lhs.keySize),
-              std::string_view(rhsKey, rhs.keySize));
+                     std::string_view(lhsKey, lhs.keySize),
+                     std::string_view(rhsKey, rhs.keySize)) ==
+              std::strong_ordering::less;
         });
   }
   return rows;
@@ -277,6 +278,17 @@ void LocalShuffleWriter::collect(
       "key '{}' must be empty for non-sorted shuffle",
       key);
   const auto rowSize = this->rowSize(key.size(), data.size());
+
+  VELOX_CHECK_LE(
+      rowSize,
+      kDefaultStreamReaderBufferSize,
+      "Row size {} bytes (keySize={}, dataSize={}) exceeds maximum reader buffer size {} bytes. "
+      "This row would be unreadable.",
+      rowSize,
+      key.size(),
+      data.size(),
+      kDefaultStreamReaderBufferSize);
+
   auto& buffer = inProgressPartitions_[partition];
   if (buffer == nullptr) {
     buffer = AlignedBuffer::allocate<char>(
@@ -305,6 +317,116 @@ void LocalShuffleWriter::noMoreData(bool success) {
   }
 }
 
+LocalShuffleReader::SortedShuffleFileStreamReader::
+    SortedShuffleFileStreamReader(
+        const std::string& filePath,
+        uint16_t streamIdx,
+        velox::memory::MemoryPool* pool,
+        size_t bufferSize)
+    : pool_(pool),
+      filePath_(filePath),
+      streamIdx_(streamIdx) {
+  auto fs = velox::filesystems::getFileSystem(filePath, nullptr);
+  file_ = fs->openFileForRead(filePath);
+  fileSize_ = file_->size();
+  // Set bufferSize_ to actual allocated size: min of requested buffer size and file size
+  bufferSize_ = std::min(bufferSize, fileSize_);
+  buffer_ = velox::AlignedBuffer::allocate<char>(bufferSize_, pool_);
+  if (!nextRow()) {
+    eof_ = true;
+  }
+}
+
+bool LocalShuffleReader::SortedShuffleFileStreamReader::next() {
+  if (eof_) {
+    return false;
+  }
+  eof_ = !nextRow();
+  return !eof_;
+}
+
+void LocalShuffleReader::SortedShuffleFileStreamReader::fillBuffer() {
+  // Compact buffer: move remaining data to the beginning
+  if (bufferPos_ > 0 && bufferPos_ < bytesInBuffer_) {
+    const size_t remaining = bytesInBuffer_ - bufferPos_;
+    std::memmove(
+        buffer_->asMutable<char>(),
+        buffer_->as<char>() + bufferPos_,
+        remaining);
+    bytesInBuffer_ = remaining;
+  } else if (bufferPos_ >= bytesInBuffer_) {
+    bytesInBuffer_ = 0;
+  }
+
+  bufferPos_ = 0;
+
+  // Read more data from file to fill remaining buffer space
+  if (bytesInBuffer_ < bufferSize_ && filePos_ < fileSize_) {
+    const size_t toRead =
+        std::min(bufferSize_ - bytesInBuffer_, fileSize_ - filePos_);
+
+    if (toRead > 0) {
+      const std::string_view bytesRead = file_->pread(
+          filePos_, toRead, buffer_->asMutable<char>() + bytesInBuffer_);
+      bytesInBuffer_ += bytesRead.size();
+      filePos_ += bytesRead.size();
+    }
+  }
+}
+
+bool LocalShuffleReader::SortedShuffleFileStreamReader::nextRow() {
+  constexpr size_t kHeaderSize = sizeof(TRowSize) * 2;
+  if (bufferPos_ + kHeaderSize > bytesInBuffer_) {
+    fillBuffer();
+    if (bytesInBuffer_ < kHeaderSize) {
+      // If we have some bytes but not enough for a header, the file is
+      // corrupted
+      VELOX_CHECK_EQ(
+          bytesInBuffer_,
+          0,
+          "Corrupted shuffle file {}: incomplete header at file position {}, "
+          "expected {} bytes but only {} bytes available",
+          filePath_,
+          filePos_,
+          kHeaderSize,
+          bytesInBuffer_);
+
+      return false;
+    }
+  }
+
+  const char* pos = buffer_->as<char>() + bufferPos_;
+  const TRowSize keySize =
+      folly::Endian::big(*reinterpret_cast<const TRowSize*>(pos));
+  const TRowSize dataSize = folly::Endian::big(
+      *reinterpret_cast<const TRowSize*>(pos + sizeof(TRowSize)));
+
+  bufferPos_ += kHeaderSize;
+  const size_t rowSize = keySize + dataSize;
+  // Ensure buffer contains complete entry (rowSize must fit in bufferSize)
+  if (bufferPos_ + rowSize > bytesInBuffer_) {
+    fillBuffer();
+    // If still not enough data after fill, the file is corrupted
+    VELOX_CHECK_GE(
+        bytesInBuffer_,
+        rowSize,
+        "Corrupted shuffle file {}: incomplete entry at file position {}, "
+        "expected {} bytes (keySize={}, dataSize={}) but only {} bytes available",
+        filePath_,
+        filePos_,
+        rowSize,
+        keySize,
+        dataSize,
+        bytesInBuffer_);
+  }
+
+  pos = buffer_->as<char>() + bufferPos_;
+  currentKey_.assign(pos, keySize);
+  currentData_.assign(pos + keySize, dataSize);
+  bufferPos_ += rowSize;
+  return true;
+}
+
 LocalShuffleReader::LocalShuffleReader(
     const std::string& rootPath,
     const std::string& queryId,
@@ -319,22 +441,120 @@ LocalShuffleReader::LocalShuffleReader(
   fileSystem_ = velox::filesystems::getFileSystem(rootPath_, nullptr);
 }
 
-folly::SemiFuture<std::vector<std::unique_ptr<ReadBatch>>>
-LocalShuffleReader::next(uint64_t maxBytes) {
-  if (readPartitionFiles_.empty()) {
-    readPartitionFiles_ = getReadPartitionFiles();
+std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextSorted(
+    uint64_t maxBytes) {
+  std::vector<std::unique_ptr<ReadBatch>> batches;
+
+  if (!mergeInitialized_) {
+    if (readPartitionFiles_.empty()) {
+      mergeInitialized_ = true;
+      return batches;
+    }
+
+    std::vector<std::unique_ptr<velox::MergeStream>> streams;
+    streams.reserve(readPartitionFiles_.size());
+    uint16_t streamIndex = 0;
+    for (const auto& filename : readPartitionFiles_) {
+      VELOX_CHECK(
+          !filename.empty(),
+          "Invalid empty shuffle file path for query {}, partitions: [{}]",
+          queryId_,
+          folly::join(", ", partitionIds_));
+      auto reader = std::make_unique<SortedShuffleFileStreamReader>(
+          filename, streamIndex, pool_);
+      if (reader->hasData()) {
+        streams.push_back(std::move(reader));
+        ++streamIndex;
+      }
+    }
+    if (!streams.empty()) {
+      merge_ =
+          std::make_unique<velox::TreeOfLosers<velox::MergeStream, uint16_t>>(
+              std::move(streams));
+    }
+
+    mergeInitialized_ = true;
+  }
+
+  auto batchBuffer = AlignedBuffer::allocate<char>(maxBytes, pool_, 0);
+  std::vector<std::string_view> rows;
+  size_t totalBytes = 0;
+  size_t bufferUsed = 0;
+
+  VELOX_CHECK(
+      merge_ || readPartitionFiles_.empty(),
+      "Failed to initialize merge: no valid shuffle files found");
+
+  if (merge_ == nullptr) {
+    return batches;
   }
 
+  while (auto* stream = merge_->next()) {
+    auto* reader = dynamic_cast<SortedShuffleFileStreamReader*>(stream);
+    const auto data = reader->currentData();
+    const size_t rowSize = kUint32Size + data.size();
+
+    // Ensure single row doesn't exceed maxBytes limit
+    VELOX_CHECK_LE(
+        rowSize,
+        maxBytes,
+        "Single row size {} bytes exceeds maxBytes limit {} bytes. "
+        "This indicates data corruption or misconfiguration.",
+        rowSize,
+        maxBytes);
+
+    // Check if adding this row would exceed maxBytes limit
+    if (totalBytes > 0 && totalBytes + rowSize > maxBytes) {
+      batches.push_back(
+          std::make_unique<ReadBatch>(std::move(rows), std::move(batchBuffer)));
+      return batches;
+    }
+
+    // Write row: [dataSize][data]
+    char* writePos = batchBuffer->asMutable<char>() + bufferUsed;
+    *reinterpret_cast<TRowSize*>(writePos) =
+        folly::Endian::big(static_cast<TRowSize>(data.size()));
+
+    if (!data.empty()) {
+      memcpy(writePos + sizeof(TRowSize), data.data(), data.size());
+    }
+
+    rows.emplace_back(batchBuffer->as<char>() + bufferUsed, rowSize);
+    bufferUsed += rowSize;
+    totalBytes += rowSize;
+
+    reader->next();
+  }
+
+  if (!rows.empty()) {
+    batches.push_back(
+        std::make_unique<ReadBatch>(std::move(rows), std::move(batchBuffer)));
+  }
+
+  return batches;
+}
+
+std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextUnsorted(
+    uint64_t maxBytes) {
   std::vector<std::unique_ptr<ReadBatch>> batches;
   uint64_t totalBytes{0};
-  // Read files until we reach maxBytes limit or run out of files.
+
   while (readPartitionFileIndex_ < readPartitionFiles_.size()) {
     const auto filename = readPartitionFiles_[readPartitionFileIndex_];
     auto file = fileSystem_->openFileForRead(filename);
     const auto fileSize = file->size();
 
-    // Stop if adding this file would exceed maxBytes (unless we haven't read
-    // any files yet)
+    // Ensure single file doesn't exceed maxBytes limit
+    // This guarantees we can always make progress and prevents OOM
+    VELOX_CHECK_LE(
+        fileSize,
+        maxBytes,
+        "Shuffle file {} size {} bytes exceeds maxBytes limit {} bytes. "
+        "This indicates data corruption or misconfiguration.",
+        filename,
+        fileSize,
+        maxBytes);
+
     if (!batches.empty() && totalBytes + fileSize > maxBytes) {
       break;
     }
@@ -343,7 +563,6 @@ LocalShuffleReader::next(uint64_t maxBytes) {
     file->pread(0, fileSize, buffer->asMutable<void>());
     ++readPartitionFileIndex_;
 
-    // Parse the buffer to extract individual rows
     const char* data = buffer->as<char>();
     const auto parsedRows = extractRowMetadata(data, fileSize, sortedShuffle_);
     std::vector<std::string_view> rows;
@@ -357,7 +576,16 @@ LocalShuffleReader::next(uint64_t maxBytes) {
         std::make_unique<ReadBatch>(std::move(rows), std::move(buffer)));
   }
 
-  return folly::makeSemiFuture(std::move(batches));
+  return batches;
+}
+
+folly::SemiFuture<std::vector<std::unique_ptr<ReadBatch>>>
+LocalShuffleReader::next(uint64_t maxBytes) {
+  if (readPartitionFiles_.empty()) {
+    readPartitionFiles_ = getReadPartitionFiles();
+  }
+  return folly::makeSemiFuture(
+      sortedShuffle_ ? nextSorted(maxBytes) : nextUnsorted(maxBytes));
 }
 
 void LocalShuffleReader::noMoreData(bool success) {
@@ -381,7 +609,7 @@ std::vector<std::string> LocalShuffleReader::getReadPartitionFiles() const {
         fmt::format("{}/{}_{}_", trimmedRootPath, queryId_, partitionId);
     auto files = fileSystem_->list(fmt::format("{}/", rootPath_));
     for (const auto& file : files) {
-      if (file.starts_with(prefix)) {
+      if (file.find(prefix) == 0) {
         partitionFiles.push_back(file);
       }
     }