feat: impl sort key for LocalShuffleReader (prestodb#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

@@ -15,17 +15,71 @@
 #include "presto_cpp/external/json/nlohmann/json.hpp"
 #include "presto_cpp/main/common/Configs.h"
 
+#include "velox/common/file/FileInputStream.h"
+
 #include <folly/lang/Bits.h>
+#include <boost/range/algorithm/sort.hpp>
 
-using namespace facebook::velox::exec;
-using namespace facebook::velox;
 
 namespace facebook::presto::operators {
 
 using json = nlohmann::json;
 
 namespace {
 
+/// SortedFileInputStream reads sorted (key, data) pairs from a single
+/// shuffle file with buffered I/O. It extends FileInputStream for efficient
+/// buffered I/O and implements MergeStream interface for k-way merge.
+class SortedFileInputStream final : public velox::common::FileInputStream,
+                                    public velox::MergeStream {
+ public:
+  SortedFileInputStream(
+      const std::string& filePath,
+      uint16_t streamIdx,
+      velox::memory::MemoryPool* pool,
+      size_t bufferSize = kDefaultInputStreamBufferSize);
+
+  ~SortedFileInputStream() override = default;
+
+  /// Advances to next entry. Returns false if EOF reached.
+  bool next();
+
+  std::string_view currentKey() const {
+    return currentKey_;
+  }
+
+  std::string_view currentData() const {
+    return currentData_;
+  }
+
+  bool hasData() const override {
+    return hasCurrentRow_;
+  }
+
+  bool operator<(const velox::MergeStream& other) const override {
+    const auto* otherReader = static_cast<const SortedFileInputStream*>(&other);
+    const auto cmp = compareKeys(currentKey_, otherReader->currentKey_);
+    if (cmp != std::strong_ordering::equal) {
+      return cmp == std::strong_ordering::less;
+    }
+    // Tie-break using stream index for deterministic ordering
+    return streamIdx_ < otherReader->streamIdx_;
+  }
+
+  const std::string& filePath() const {
+    return filePath_;
+  }
+
+ private:
+  bool nextRow();
+
+  const std::string filePath_;
+  const uint16_t streamIdx_;
+  std::string currentKey_;
+  std::string currentData_;
+  bool hasCurrentRow_{false};
+};
+
 std::vector<RowMetadata>
 extractRowMetadata(const char* buffer, size_t bufferSize, bool sortedShuffle) {
   std::vector<RowMetadata> rows;
@@ -91,13 +145,8 @@ extractRowMetadata(const char* buffer, size_t bufferSize, bool sortedShuffle) {
 
 inline std::string_view
 extractRowData(const RowMetadata& row, const char* buffer, bool sortedShuffle) {
-  if (sortedShuffle) {
-    const size_t dataOffset = row.rowStart + (kUint32Size * 2) + row.keySize;
-    return {buffer + dataOffset, row.dataSize};
-  } else {
-    const size_t dataOffset = row.rowStart + kUint32Size;
-    return {buffer + dataOffset, row.dataSize};
-  }
+  const auto dataOffset = row.rowStart + (sortedShuffle ? (kUint32Size * 2) + row.keySize : kUint32Size);
+  return {buffer + dataOffset, row.dataSize};
 }
 
 std::vector<RowMetadata> extractAndSortRowMetadata(
@@ -106,15 +155,15 @@ std::vector<RowMetadata> extractAndSortRowMetadata(
     bool sortedShuffle) {
   auto rows = extractRowMetadata(buffer, bufferSize, sortedShuffle);
   if (!rows.empty() && sortedShuffle) {
-    std::sort(
-        rows.begin(),
-        rows.end(),
+    boost::range::sort(
+        rows,
         [buffer](const RowMetadata& lhs, const RowMetadata& rhs) {
           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,9 +326,10 @@ void LocalShuffleWriter::collect(
       "key '{}' must be empty for non-sorted shuffle",
       key);
   const auto rowSize = this->rowSize(key.size(), data.size());
+
   auto& buffer = inProgressPartitions_[partition];
   if (buffer == nullptr) {
-    buffer = AlignedBuffer::allocate<char>(
+    buffer = velox::AlignedBuffer::allocate<char>(
         std::max(static_cast<uint64_t>(rowSize), maxBytesPerPartition_),
         pool_,
         0);
@@ -305,6 +355,49 @@ void LocalShuffleWriter::noMoreData(bool success) {
   }
 }
 
+SortedFileInputStream::SortedFileInputStream(
+    const std::string& filePath,
+    uint16_t streamIdx,
+    velox::memory::MemoryPool* pool,
+    size_t bufferSize)
+    : velox::common::FileInputStream(
+          velox::filesystems::getFileSystem(filePath, nullptr)
+              ->openFileForRead(filePath),
+          bufferSize,
+          pool),
+      filePath_(filePath),
+      streamIdx_(streamIdx) {
+  // Initialize by reading first row.
+  nextRow();
+}
+
+bool SortedFileInputStream::next() {
+  return nextRow();
+}
+
+bool SortedFileInputStream::nextRow() {
+  if (atEnd()) {
+    currentKey_.clear();
+    currentData_.clear();
+    hasCurrentRow_ = false;
+    return false;
+  }
+  const auto keySize = folly::Endian::big(read<TRowSize>());
+  const auto dataSize = folly::Endian::big(read<TRowSize>());
+
+  currentKey_.resize(keySize);
+  if (keySize > 0) {
+    readBytes(reinterpret_cast<uint8_t*>(currentKey_.data()), keySize);
+  }
+  currentData_.resize(dataSize);
+  if (dataSize > 0) {
+    readBytes(reinterpret_cast<uint8_t*>(currentData_.data()), dataSize);
+  }
+
+  hasCurrentRow_ = true;
+  return true;
+}
+
 LocalShuffleReader::LocalShuffleReader(
     const std::string& rootPath,
     const std::string& queryId,
@@ -319,31 +412,113 @@ 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 streamIdx = 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<SortedFileInputStream>(
+          filename, streamIdx, pool_);
+      if (reader->hasData()) {
+        streams.push_back(std::move(reader));
+        ++streamIdx;
+      }
+    }
+    if (!streams.empty()) {
+      merge_ =
+          std::make_unique<velox::TreeOfLosers<velox::MergeStream, uint16_t>>(
+              std::move(streams));
+    }
+
+    mergeInitialized_ = true;
   }
 
+  if (merge_ == nullptr) {
+    return batches;
+  }
+
+  auto batchBuffer = velox::AlignedBuffer::allocate<char>(maxBytes, pool_, 0);
+  std::vector<std::string_view> rows;
+  uint64_t bufferUsed = 0;
+
+  while (auto* stream = merge_->next()) {
+    auto* reader = dynamic_cast<SortedFileInputStream*>(stream);
+    const auto data = reader->currentData();
+    const auto rowSize = kUint32Size + data.size();
+
+    // With the current row the bufferUsed byte will exceed the maxBytes
+    if (bufferUsed + rowSize > maxBytes) {
+      if (bufferUsed > 0) {
+        // We have some rows already, return them to release the memory
+        batches.push_back(
+            std::make_unique<ReadBatch>(std::move(rows), std::move(batchBuffer)));
+        return batches;
+      }
+      // Single row exceeds buffer - allocate larger buffer for this row
+      batchBuffer = velox::AlignedBuffer::allocate<char>(rowSize, pool_, 0);
+      bufferUsed = 0;
+    }
+
+    // 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;
+
+    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)
+
+    // TODO: Refactor to use streaming I/O with bounded buffer size instead of
+    // loading entire files into memory at once. A streaming approach would
+    // reduce peak memory consumption and enable processing arbitrarily large
+    // shuffle files while maintaining constant memory usage.
     if (!batches.empty() && totalBytes + fileSize > maxBytes) {
       break;
     }
 
-    auto buffer = AlignedBuffer::allocate<char>(fileSize, pool_, 0);
+    auto buffer = velox::AlignedBuffer::allocate<char>(fileSize, pool_, 0);
     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 +532,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) {

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

@@ -18,6 +18,7 @@
 #include <string_view>
 #include <vector>
 
+#include "velox/common/base/TreeOfLosers.h"
 #include "velox/common/file/FileSystems.h"
 #include "velox/exec/Operator.h"
 
@@ -28,24 +29,30 @@ namespace facebook::presto::operators {
 using TRowSize = uint32_t;
 constexpr size_t kUint32Size = sizeof(TRowSize);
 
+// Default buffer size for SortedFileInputStream
+// This buffer is used for streaming reads from shuffle files during k-way
+// merge.
+constexpr uint64_t kDefaultInputStreamBufferSize = 8 * 1024 * 1024; // 8MB
+
 // Metadata describing a serialized row's location and sizes in a buffer
 struct RowMetadata {
   size_t rowStart; // Offset to the start of this row
   uint32_t keySize; // Size of key (0 for non-sorted)
   uint32_t dataSize; // Size of data
 };
 
-// Compare two sort keys lexicographically
-inline bool compareKeys(std::string_view key1, std::string_view key2) noexcept {
-  return std::lexicographical_compare(
-      reinterpret_cast<const unsigned char*>(key1.data()),
-      reinterpret_cast<const unsigned char*>(key1.data() + key1.size()),
-      reinterpret_cast<const unsigned char*>(key2.data()),
-      reinterpret_cast<const unsigned char*>(key2.data() + key2.size()));
+// Compares sort keys lexicographically using three-way comparison.
+inline std::strong_ordering compareKeys(
+    std::string_view key1,
+    std::string_view key2) noexcept {
+  return std::lexicographical_compare_three_way(
+      reinterpret_cast<const uint8_t*>(key1.data()),
+      reinterpret_cast<const uint8_t*>(key1.data() + key1.size()),
+      reinterpret_cast<const uint8_t*>(key2.data()),
+      reinterpret_cast<const uint8_t*>(key2.data() + key2.size()));
 }
 
-// TODO: Testing function to expose extractRowMetadata for tests.
-// This will be removed after reader changes.
+// Testing function to expose extractRowMetadata for tests.
 std::vector<RowMetadata> testingExtractRowMetadata(
     const char* buffer,
     size_t bufferSize,
@@ -180,6 +187,12 @@ class LocalShuffleReader : public ShuffleReader {
   // Returns all created shuffle files for 'partition_'.
   std::vector<std::string> getReadPartitionFiles() const;
 
+  // Reads sorted shuffle data using k-way merge with TreeOfLosers.
+  std::vector<std::unique_ptr<ReadBatch>> nextSorted(uint64_t maxBytes);
+
+  // Reads unsorted shuffle data in batch-based file reading.
+  std::vector<std::unique_ptr<ReadBatch>> nextUnsorted(uint64_t maxBytes);
+
   const std::string rootPath_;
   const std::string queryId_;
   const std::vector<std::string> partitionIds_;
@@ -194,6 +207,10 @@ class LocalShuffleReader : public ShuffleReader {
 
   // The top directory of the shuffle files and its file system.
   std::shared_ptr<velox::filesystems::FileSystem> fileSystem_;
+
+  // Used to merge sorted streams from multiple shuffle files for k-way merge.
+  std::unique_ptr<velox::TreeOfLosers<velox::MergeStream, uint16_t>> merge_;
+  bool mergeInitialized_{false};
 };
 
 class LocalPersistentShuffleFactory : public ShuffleInterfaceFactory {