feat: Impl sort key for LocalShuffleWriter (#26547)

Summary:

Adds sorting capability to LocalShuffleWriter to enable sorted shuffle operations in Presto native execution. When sort keys are provided, shuffle data is sorted before writing to disk.

Differential Revision: D86322593

Author: duxiao1212

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

@@ -15,6 +15,8 @@
 #include "presto_cpp/external/json/nlohmann/json.hpp"
 #include "presto_cpp/main/common/Configs.h"
 
+#include <boost/range/algorithm/sort.hpp>
+
 using namespace facebook::velox::exec;
 using namespace facebook::velox;
 
@@ -39,15 +41,32 @@ inline std::string createShuffleFileName(
       id);
 }
 
-// This file is used to indicate that the shuffle system is ready to be used for
-// reading (acts as a sync point between readers if needed). Mostly used for
-// test purposes.
-const static std::string kReadyForReadFilename = "readyForRead";
+inline void appendShuffleRow(
+    char* FOLLY_NONNULL buffer,
+    std::string_view sortKey,
+    std::string_view data) {
+  // Write memory layout: [keySize][dataSize][key][data]
+  *reinterpret_cast<TRowSize*>(buffer) =
+      folly::Endian::big(static_cast<TRowSize>(sortKey.size()));
+  buffer += kUint32Size;
+
+  *reinterpret_cast<TRowSize*>(buffer) =
+      folly::Endian::big(static_cast<TRowSize>(data.size()));
+  buffer += kUint32Size;
+
+  if (!sortKey.empty()) {
+    memcpy(buffer, sortKey.data(), sortKey.size());
+    buffer += sortKey.size();
+  }
+
+  memcpy(buffer, data.data(), data.size());
+}
 } // namespace
 
+
 LocalShuffleWriter::LocalShuffleWriter(
-    const std::string& rootPath,
-    const std::string& queryId,
+    std::string rootPath,
+    std::string queryId,
     uint32_t shuffleId,
     uint32_t numPartitions,
     uint64_t maxBytesPerPartition,
@@ -59,14 +78,54 @@ LocalShuffleWriter::LocalShuffleWriter(
       rootPath_(std::move(rootPath)),
       queryId_(std::move(queryId)),
       shuffleId_(shuffleId) {
-  // Use resize/assign instead of resize(size, val).
-  inProgressPartitions_.resize(numPartitions_);
   inProgressPartitions_.assign(numPartitions_, nullptr);
-  inProgressSizes_.resize(numPartitions_);
   inProgressSizes_.assign(numPartitions_, 0);
+  inProgressRowOffsets_.assign(numPartitions_, {});
   fileSystem_ = velox::filesystems::getFileSystem(rootPath_, nullptr);
 }
 
+void LocalShuffleWriter::writeBlock(int32_t partition) {
+  auto& buffer = inProgressPartitions_[partition];
+  const auto bufferSize = inProgressSizes_[partition];
+  if (!buffer || bufferSize == 0) {
+    return;
+  }
+
+  auto file = getNextOutputFile(partition);
+  auto& offsets = inProgressRowOffsets_[partition];
+
+  const bool needsSorting = !offsets.empty() && offsets[0].keySize > 0;
+
+  if (needsSorting) {
+    const char* sourceData = buffer->as<char>();
+    sortRows(offsets, sourceData);
+
+    for (const auto& offset : offsets) {
+      const char* rowStart = sourceData + offset.keyOffset - (kUint32Size * 2);
+      const size_t rowLen =
+          (kUint32Size * 2) + offset.keySize + offset.dataSize;
+      file->append(std::string_view(rowStart, rowLen));
+    }
+  } else {
+    file->append(std::string_view(buffer->as<char>(), bufferSize));
+  }
+  file->close();
+
+  // Reset for buffer reuse
+  inProgressSizes_[partition] = 0;
+  inProgressRowOffsets_[partition].clear();
+}
+
+void LocalShuffleWriter::sortRows(
+    std::vector<RowOffset>& offsets,
+    const char* bufferData) {
+  boost::range::sort(
+      offsets, [bufferData](const RowOffset& lhs, const RowOffset& rhs) {
+        return detail::compareKeys(
+            lhs.getKey(bufferData), rhs.getKey(bufferData));
+      });
+}
+
 std::unique_ptr<velox::WriteFile> LocalShuffleWriter::getNextOutputFile(
     int32_t partition) {
   auto filename = nextAvailablePartitionFileName(rootPath_, partition);
@@ -92,49 +151,42 @@ std::string LocalShuffleWriter::nextAvailablePartitionFileName(
   return filename;
 }
 
-void LocalShuffleWriter::storePartitionBlock(int32_t partition) {
-  auto& buffer = inProgressPartitions_[partition];
-  auto file = getNextOutputFile(partition);
-  file->append(
-      std::string_view(buffer->as<char>(), inProgressSizes_[partition]));
-  file->close();
-  inProgressPartitions_[partition].reset();
-  inProgressSizes_[partition] = 0;
-}
-
 void LocalShuffleWriter::collect(
     int32_t partition,
-    std::string_view /* key */,
+    std::string_view key,
     std::string_view data) {
-  using TRowSize = uint32_t;
+  VELOX_CHECK_LT(partition, numPartitions_);
+
+  const auto rowSize = (kUint32Size * 2) + key.size() + data.size();
 
   auto& buffer = inProgressPartitions_[partition];
-  const TRowSize rowSize = data.size();
-  const auto size = sizeof(TRowSize) + rowSize;
-
-  // Check if there is enough space in the buffer.
-  if ((buffer != nullptr) &&
-      (inProgressSizes_[partition] + size >= buffer->capacity())) {
-    storePartitionBlock(partition);
-    // NOTE: the referenced 'buffer' will be reset in storePartitionBlock.
-  }
 
-  // Allocate buffer if needed.
   if (buffer == nullptr) {
     buffer = AlignedBuffer::allocate<char>(
-        std::max((uint64_t)size, maxBytesPerPartition_), pool_);
+        std::max(static_cast<uint64_t>(rowSize), maxBytesPerPartition_),
+        pool_,
+        0);
     inProgressSizes_[partition] = 0;
-    inProgressPartitions_[partition] = buffer;
+  } else if (inProgressSizes_[partition] + rowSize >= buffer->capacity()) {
+    writeBlock(partition);
   }
 
-  // Copy data.
-  auto offset = inProgressSizes_[partition];
-  auto rawBuffer = buffer->asMutable<char>() + offset;
+  auto* rawBuffer = buffer->asMutable<char>();
+  auto* writePos = rawBuffer + inProgressSizes_[partition];
 
-  *(TRowSize*)(rawBuffer) = folly::Endian::big(rowSize);
-  ::memcpy(rawBuffer + sizeof(TRowSize), data.data(), rowSize);
+  appendShuffleRow(writePos, key, data);
+  const auto currentOffset = inProgressSizes_[partition];
+  const auto keyOffset = currentOffset + (kUint32Size * 2);
+  const auto dataOffset = keyOffset + key.size();
 
-  inProgressSizes_[partition] += size;
+  inProgressRowOffsets_[partition].emplace_back(
+      RowOffset{
+          .keyOffset = keyOffset,
+          .dataOffset = dataOffset,
+          .keySize = static_cast<uint32_t>(key.size()),
+          .dataSize = static_cast<uint32_t>(data.size())});
+
+  inProgressSizes_[partition] += rowSize;
 }
 
 void LocalShuffleWriter::noMoreData(bool success) {
@@ -144,27 +196,24 @@ void LocalShuffleWriter::noMoreData(bool success) {
   }
   for (auto i = 0; i < numPartitions_; ++i) {
     if (inProgressSizes_[i] > 0) {
-      storePartitionBlock(i);
+      writeBlock(i);
     }
   }
 }
-
 LocalShuffleReader::LocalShuffleReader(
-    const std::string& rootPath,
-    const std::string& queryId,
+    std::string rootPath,
+    std::string queryId,
     std::vector<std::string> partitionIds,
     velox::memory::MemoryPool* pool)
-    : rootPath_(rootPath),
-      queryId_(queryId),
+    : rootPath_(std::move(rootPath)),
+      queryId_(std::move(queryId)),
       partitionIds_(std::move(partitionIds)),
       pool_(pool) {
   fileSystem_ = velox::filesystems::getFileSystem(rootPath_, nullptr);
 }
 
 folly::SemiFuture<std::vector<std::unique_ptr<ReadBatch>>>
 LocalShuffleReader::next(uint64_t maxBytes) {
-  using TRowSize = uint32_t;
-
   if (readPartitionFiles_.empty()) {
     readPartitionFiles_ = getReadPartitionFiles();
   }
@@ -188,21 +237,34 @@ LocalShuffleReader::next(uint64_t maxBytes) {
     ++readPartitionFileIndex_;
 
     // Parse the buffer to extract individual rows.
-    // Each row is stored as: | row-size (4 bytes) | row-data (row-size bytes) |
+    // Each row is stored as: [keySize][dataSize][key][data]
     std::vector<std::string_view> rows;
     const char* data = buffer->as<char>();
     size_t offset = 0;
     const size_t totalSize = buffer->size();
 
-    while (offset + sizeof(TRowSize) <= totalSize) {
+    while (offset + (kUint32Size * 2) <= totalSize) {
       // Read row size (stored in big endian).
-      const TRowSize rowSize = folly::Endian::big(*(TRowSize*)(data + offset));
-      offset += sizeof(TRowSize);
+      const TRowSize keySize = folly::Endian::big(*reinterpret_cast<const TRowSize*>(data + offset));
+
+      offset += kUint32Size;
+      const TRowSize dataSize = folly::Endian::big(*reinterpret_cast<const TRowSize*>(data + offset));
+      offset += kUint32Size;
+
+      VELOX_CHECK_LE(
+          offset + keySize + dataSize,
+          totalSize,
+          "Corrupted shuffle data: expected {} bytes for row (offset={}, keySize={}, dataSize={}) but only {} bytes available in buffer",
+          offset + keySize + dataSize,
+          offset,
+          keySize,
+          dataSize,
+          totalSize);
 
-      VELOX_CHECK_LE(offset + rowSize, totalSize, "Invalid row data: row size");
       // Create a Row with empty key and the row data as value.
-      rows.emplace_back(std::string_view{data + offset, rowSize});
-      offset += rowSize;
+      offset += keySize; // Skip  sort key
+      rows.emplace_back(data + offset, dataSize);
+      offset += dataSize;
     }
 
     totalBytes += fileSize;
@@ -234,7 +296,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.find(prefix) == 0) {
+      if (file.starts_with(prefix)) {
         partitionFiles.push_back(file);
       }
     }

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

@@ -13,6 +13,7 @@
  */
 #pragma once
 
+#include "presto_cpp/main/operators/LocalShuffleUtils.h"
 #include "presto_cpp/main/operators/ShuffleInterface.h"
 #include "velox/buffer/Buffer.h"
 #include "velox/common/file/File.h"
@@ -72,17 +73,15 @@ struct LocalShuffleReadInfo {
 class LocalShuffleWriter : public ShuffleWriter {
  public:
   LocalShuffleWriter(
-      const std::string& rootPath,
-      const std::string& queryId,
+      std::string rootPath,
+      std::string queryId,
       uint32_t shuffleId,
       uint32_t numPartitions,
       uint64_t maxBytesPerPartition,
       velox::memory::MemoryPool* pool);
 
-  void collect(
-      int32_t partition,
-      std::string_view /* key */,
-      std::string_view data) override;
+  void collect(int32_t partition, std::string_view key, std::string_view data)
+      override;
 
   void noMoreData(bool success) override;
 
@@ -97,7 +96,7 @@ class LocalShuffleWriter : public ShuffleWriter {
   std::unique_ptr<velox::WriteFile> getNextOutputFile(int32_t partition);
 
   // Writes the in-progress block to the given partition.
-  void storePartitionBlock(int32_t partition);
+  void writeBlock(int32_t partition);
 
   // Deletes all the files in the root directory.
   void cleanup();
@@ -121,13 +120,36 @@ class LocalShuffleWriter : public ShuffleWriter {
   std::vector<velox::BufferPtr> inProgressPartitions_;
   std::vector<size_t> inProgressSizes_;
   std::shared_ptr<velox::filesystems::FileSystem> fileSystem_;
+
+  // For sorted shuffle, track the key and data offsets for each row for
+  // efficient sorting.
+  struct RowOffset {
+    size_t keyOffset;
+    size_t dataOffset;
+    uint32_t keySize;
+    uint32_t dataSize;
+
+    inline std::string_view getKey(const char* FOLLY_NONNULL buffer) const {
+      return {buffer + keyOffset, keySize};
+    }
+
+    inline std::string_view getData(const char* FOLLY_NONNULL buffer) const {
+      return {buffer + dataOffset, dataSize};
+    }
+  };
+
+  std::vector<std::vector<RowOffset>> inProgressRowOffsets_;
+
+  void sortRows(
+      std::vector<RowOffset>& offsets,
+      const char* FOLLY_NONNULL buffer);
 };
 
 class LocalShuffleReader : public ShuffleReader {
  public:
   LocalShuffleReader(
-      const std::string& rootPath,
-      const std::string& queryId,
+      std::string rootPath,
+      std::string queryId,
       std::vector<std::string> partitionIds,
       velox::memory::MemoryPool* pool);