feat: Impl sort key for LocalShuffleWriter (#26547)

duxiao1212 · facebook-github-bot · commit b75bf65ad5fb · 2025-11-10T23:06:38.000-08:00
Summary:

Impl sort key for LocalShuffleWriter

Differential Revision: D86322593
diff --git a/presto-native-execution/presto_cpp/main/operators/LocalShuffle.cpp b/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,41 @@ 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 size_t rowSize(size_t keySize, size_t dataSize) noexcept {
+  return (kUint32Size * 2) + keySize + dataSize;
+}
+
+inline void appendShuffleRow(
+    char* FOLLY_NONNULL buffer,
+    std::string_view sortKey,
+    std::string_view data) {
+  // Write memory layout: [keySize][dataSize][key][data]
+  *(TRowSize*)(buffer) =
+      folly::Endian::big(static_cast<TRowSize>(sortKey.size()));
+  buffer += kUint32Size;
+
+  *(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());
+}
+
+void writeBufferToFile(velox::WriteFile* file, const char* data, size_t size) {
+  VELOX_CHECK_NOT_NULL(file);
+  file->append(std::string_view(data, size));
+  file->close();
+}
+
 } // 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 +87,53 @@ 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));
+    }
+    file->close();
+  } else {
+    writeBufferToFile(file.get(), buffer->as<char>(), bufferSize);
+  }
+
+  // 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,47 +159,36 @@ 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_);
 
-  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.
-  }
+  const auto size = rowSize(key.size(), data.size());
 
-  // Allocate buffer if needed.
-  if (buffer == nullptr) {
-    buffer = AlignedBuffer::allocate<char>(
-        std::max((uint64_t)size, maxBytesPerPartition_), pool_);
-    inProgressSizes_[partition] = 0;
-    inProgressPartitions_[partition] = buffer;
+  if (inProgressPartitions_[partition] == nullptr) {
+    inProgressPartitions_[partition] =
+        AlignedBuffer::allocate<char>(maxBytesPerPartition_, pool_, 0);
+  }
+  if (inProgressSizes_[partition] + size >= maxBytesPerPartition_) {
+    writeBlock(partition);
   }
 
-  // Copy data.
-  auto offset = inProgressSizes_[partition];
-  auto rawBuffer = buffer->asMutable<char>() + offset;
+  auto* rawBuffer = inProgressPartitions_[partition]->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();
+
+  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] += size;
 }
@@ -144,11 +200,10 @@ 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,
diff --git a/presto-native-execution/presto_cpp/main/operators/LocalShuffle.h b/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,6 +120,29 @@ 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};
+    }
+  } __attribute__((aligned(32)));
+
+  std::vector<std::vector<RowOffset>> inProgressRowOffsets_;
+
+  void sortRows(
+      std::vector<RowOffset>& offsets,
+      const char* FOLLY_NONNULL buffer);
 };
 
 class LocalShuffleReader : public ShuffleReader {
diff --git a/presto-native-execution/presto_cpp/main/operators/LocalShuffleUtils.cpp b/presto-native-execution/presto_cpp/main/operators/LocalShuffleUtils.cpp
@@ -0,0 +1,73 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "presto_cpp/main/operators/LocalShuffleUtils.h"
+#include "velox/common/base/Exceptions.h"
+#include "folly/lang/Bits.h"
+
+namespace facebook::presto::operators::detail {
+
+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()));
+}
+
+ParsedShuffleData parseShuffleRows(const char* buffer, size_t totalSize) {
+  ParsedShuffleData result;
+  size_t offset = 0;
+
+  while (offset + static_cast<size_t>(2 * kUint32Size) <= totalSize) {
+    // Read KeySize (big-endian)
+    const TRowSize keySize =
+        folly::Endian::big(*reinterpret_cast<const TRowSize*>(buffer + offset));
+    offset += kUint32Size;
+
+    // Read DataSize (big-endian)
+    const TRowSize dataSize =
+        folly::Endian::big(*reinterpret_cast<const TRowSize*>(buffer + offset));
+    offset += kUint32Size;
+
+    // Validate keySize
+    if (keySize > 0) {
+      VELOX_CHECK_LE(
+          offset + keySize,
+          totalSize,
+          "Invalid row data, offset + keySize: {} exceeds totalSize: {}",
+          offset + keySize,
+          totalSize);
+    }
+
+    // Validate total size
+    VELOX_CHECK_LE(
+        offset + keySize + dataSize,
+        totalSize,
+        "Invalid row data, offset + keySize + dataSize: {} exceeds totalSize: {}",
+        offset + keySize + dataSize,
+        totalSize);
+
+    // Extract key and data as string_views
+    result.keys.emplace_back(buffer + offset, keySize);
+    offset += keySize;
+
+    result.data.emplace_back(buffer + offset, dataSize);
+    offset += dataSize;
+  }
+
+  return result;
+}
+
+} // namespace facebook::presto::operators::detail
diff --git a/presto-native-execution/presto_cpp/main/operators/LocalShuffleUtils.h b/presto-native-execution/presto_cpp/main/operators/LocalShuffleUtils.h
@@ -0,0 +1,42 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cstdint>
+#include <string_view>
+#include <vector>
+
+namespace facebook::presto::operators {
+
+using TRowSize = uint32_t;
+constexpr uint32_t kUint32Size = static_cast<uint32_t>(sizeof(uint32_t));
+
+namespace detail {
+
+/// Structure to hold parsed shuffle data with separate keys and data vectors
+struct ParsedShuffleData {
+  std::vector<std::string_view> keys;
+  std::vector<std::string_view> data;
+};
+
+/// Compares two keys lexicographically (byte-by-byte comparison)
+/// Returns true if key1 < key2
+bool compareKeys(std::string_view key1, std::string_view key2) noexcept;
+
+/// Parses shuffle data in the format: [keySize][dataSize][key][data]
+/// This format is used when sort keys are present in the shuffle data
+ParsedShuffleData parseShuffleRows(const char* buffer, size_t totalSize);
+
+} // namespace detail
+} // namespace facebook::presto::operators
diff --git a/presto-native-execution/presto_cpp/main/operators/tests/ShuffleTest.cpp b/presto-native-execution/presto_cpp/main/operators/tests/ShuffleTest.cpp
