misc: Refactor ShuffleTest (prestodb#26643)

Summary:
Pull Request resolved: prestodb#26643

Completely replace TestShuffleWriter/Reader with local shuffle, and refactor the ShuffleTest class

Differential Revision: D87160791

Author: duxiao1212

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

@@ -353,12 +353,30 @@ void LocalShuffleWriter::collect(
     int32_t partition,
     std::string_view key,
     std::string_view data) {
+  LOG(INFO) << "LocalShuffleWriter::collect START partition=" << partition
+            << " sortedShuffle=" << sortedShuffle_
+            << " keySize=" << key.size()
+            << " dataSize=" << data.size();
   VELOX_CHECK_LT(partition, numPartitions_);
+  // For non-sorted shuffle, key must be empty
+  // For sorted shuffle, key is used for sorting
   VELOX_CHECK(
       sortedShuffle_ || key.empty(),
-      "key '{}' must be empty for non-sorted shuffle",
-      key);
+      "key must be empty for non-sorted shuffle, got key size {}",
+      key.size());
+
+  // Testvalue injection point for exception testing
+  velox::common::testutil::TestValue::adjust(
+      "facebook::presto::operators::LocalShuffleWriter::collect", this);
+
+  // Log the actual data content for debugging
+  if (data.size() >= 16) {
+    LOG(INFO) << "LocalShuffleWriter::collect data first 16 bytes (hex): "
+              << folly::hexlify(std::string_view(data.data(), std::min(data.size(), size_t(16))));
+  }
+
   const auto rowSize = this->rowSize(key.size(), data.size());
+  LOG(INFO) << "LocalShuffleWriter::collect rowSize=" << rowSize << " keySize=" << key.size();
 
   auto& buffer = inProgressPartitions_[partition];
   if (buffer == nullptr) {
@@ -377,15 +395,19 @@ void LocalShuffleWriter::collect(
 }
 
 void LocalShuffleWriter::noMoreData(bool success) {
+  LOG(INFO) << "LocalShuffleWriter::noMoreData START success=" << success
+            << " numPartitions=" << numPartitions_;
   // Delete all shuffle files on failure.
   if (!success) {
     cleanup();
   }
   for (auto i = 0; i < numPartitions_; ++i) {
     if (inProgressSizes_[i] > 0) {
+      LOG(INFO) << "LocalShuffleWriter::noMoreData writing block for partition " << i;
       writeBlock(i);
     }
   }
+  LOG(INFO) << "LocalShuffleWriter::noMoreData COMPLETE";
 }
 
 LocalShuffleReader::LocalShuffleReader(
@@ -403,9 +425,11 @@ LocalShuffleReader::LocalShuffleReader(
 }
 
 void LocalShuffleReader::initialize() {
+  LOG(INFO) << "LocalShuffleReader::initialize START sortedShuffle=" << sortedShuffle_;
   VELOX_CHECK(!initialized_, "LocalShuffleReader already initialized");
 
   readPartitionFiles_ = getReadPartitionFiles();
+  LOG(INFO) << "LocalShuffleReader::initialize found " << readPartitionFiles_.size() << " files";
 
   if (sortedShuffle_ && !readPartitionFiles_.empty()) {
     std::vector<std::unique_ptr<velox::MergeStream>> streams;
@@ -432,13 +456,15 @@ void LocalShuffleReader::initialize() {
   }
 
   initialized_ = true;
+  LOG(INFO) << "LocalShuffleReader::initialize COMPLETE";
 }
 
 std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextSorted(
     uint64_t maxBytes) {
   std::vector<std::unique_ptr<ReadBatch>> batches;
 
   if (merge_ == nullptr) {
+    LOG(INFO) << "LocalShuffleReader::nextSorted merge is null, returning empty";
     return batches;
   }
 
@@ -449,37 +475,42 @@ std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextSorted(
   while (auto* stream = merge_->next()) {
     auto* reader = dynamic_cast<SortedFileInputStream*>(stream);
     const auto data = reader->currentData();
-    const auto rowSize = kUint32Size + data.size();
+    const auto key = reader->currentKey();
+
+    LOG(INFO) << "LocalShuffleReader::nextSorted processing row: keySize=" << key.size()
+              << " dataSize=" << data.size() << " bufferUsed=" << bufferUsed;
 
-    // With the current row the bufferUsed byte will exceed the maxBytes
-    if (bufferUsed + rowSize > maxBytes) {
+    // For sorted shuffle: data is already CompactRow serialized data
+    // Return it as-is, NO size prefix (same as CoscoShuffleReader)
+    if (bufferUsed + data.size() > maxBytes) {
       if (bufferUsed > 0) {
-        // We have some rows already, return them to release the memory
+        LOG(INFO) << "LocalShuffleReader::nextSorted returning batch with " << rows.size() << " rows";
         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);
+      // Single row exceeds buffer - allocate larger buffer
+      batchBuffer = velox::AlignedBuffer::allocate<char>(data.size(), pool_, 0);
       bufferUsed = 0;
     }
 
-    // Write row: [dataSize][data]
+    // Copy data as-is without size prefix
     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());
+      memcpy(writePos, data.data(), data.size());
     }
 
-    rows.emplace_back(batchBuffer->as<char>() + bufferUsed, rowSize);
-    bufferUsed += rowSize;
+    LOG(INFO) << "LocalShuffleReader::nextSorted wrote row at offset " << bufferUsed
+              << " size=" << data.size() << " (no prefix, like CoscoShuffle)";
+
+    rows.emplace_back(batchBuffer->as<char>() + bufferUsed, data.size());
+    bufferUsed += data.size();
 
     reader->next();
   }
 
   if (!rows.empty()) {
+    LOG(INFO) << "LocalShuffleReader::nextSorted final batch with " << rows.size() << " rows";
     batches.push_back(
         std::make_unique<ReadBatch>(std::move(rows), std::move(batchBuffer)));
   }
@@ -528,11 +559,19 @@ std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextUnsorted(
 
 folly::SemiFuture<std::vector<std::unique_ptr<ReadBatch>>>
 LocalShuffleReader::next(uint64_t maxBytes) {
+  LOG(INFO) << "LocalShuffleReader::next START maxBytes=" << maxBytes
+            << " sortedShuffle=" << sortedShuffle_;
   VELOX_CHECK(
       initialized_,
       "LocalShuffleReader::initialize() must be called before next()");
-  return folly::makeSemiFuture(
-      sortedShuffle_ ? nextSorted(maxBytes) : nextUnsorted(maxBytes));
+
+  // Testvalue injection point for exception testing
+  velox::common::testutil::TestValue::adjust(
+      "facebook::presto::operators::LocalShuffleReader::next", this);
+
+  auto result = sortedShuffle_ ? nextSorted(maxBytes) : nextUnsorted(maxBytes);
+  LOG(INFO) << "LocalShuffleReader::next COMPLETE returned " << result.size() << " batches";
+  return folly::makeSemiFuture(std::move(result));
 }
 
 void LocalShuffleReader::noMoreData(bool success) {
@@ -576,42 +615,64 @@ std::shared_ptr<ShuffleReader> LocalPersistentShuffleFactory::createReader(
     const std::string& serializedStr,
     const int32_t /*partition*/,
     velox::memory::MemoryPool* pool) {
+  LOG(INFO) << "LocalPersistentShuffleFactory::createReader START";
   const operators::LocalShuffleReadInfo readInfo =
       operators::LocalShuffleReadInfo::deserialize(serializedStr);
+  // Check if sortedShuffle field is present in the JSON
+  bool sortedShuffle = false;
+  try {
+    const auto jsonReadInfo = json::parse(serializedStr);
+    if (jsonReadInfo.contains("sortedShuffle")) {
+      jsonReadInfo.at("sortedShuffle").get_to(sortedShuffle);
+    }
+  } catch (const std::exception& /*e*/) {
+    // If parsing fails or field doesn't exist, default to false
+    sortedShuffle = false;
+  }
+  LOG(INFO) << "LocalPersistentShuffleFactory::createReader sortedShuffle=" << sortedShuffle
+            << " rootPath=" << readInfo.rootPath;
   auto reader = std::make_shared<LocalShuffleReader>(
       readInfo.rootPath,
       readInfo.queryId,
       readInfo.partitionIds,
-      /*sortShuffle=*/false, // default to false for now
+      sortedShuffle,
       pool);
   reader->initialize();
+  LOG(INFO) << "LocalPersistentShuffleFactory::createReader COMPLETE";
   return reader;
 }
 
 std::shared_ptr<ShuffleWriter> LocalPersistentShuffleFactory::createWriter(
     const std::string& serializedStr,
     velox::memory::MemoryPool* pool) {
+  LOG(INFO) << "LocalPersistentShuffleFactory::createWriter START";
   static const uint64_t maxBytesPerPartition =
       SystemConfig::instance()->localShuffleMaxPartitionBytes();
   const operators::LocalShuffleWriteInfo writeInfo =
       operators::LocalShuffleWriteInfo::deserialize(serializedStr);
-  return std::make_shared<LocalShuffleWriter>(
+  // Check if sortedShuffle field is present in the JSON
+  bool sortedShuffle = false;
+  try {
+    const auto jsonWriteInfo = json::parse(serializedStr);
+    if (jsonWriteInfo.contains("sortedShuffle")) {
+      jsonWriteInfo.at("sortedShuffle").get_to(sortedShuffle);
+    }
+  } catch (const std::exception& /*e*/) {
+    // If parsing fails or field doesn't exist, default to false
+    sortedShuffle = false;
+  }
+  LOG(INFO) << "LocalPersistentShuffleFactory::createWriter sortedShuffle=" << sortedShuffle
+            << " numPartitions=" << writeInfo.numPartitions
+            << " rootPath=" << writeInfo.rootPath;
+  auto writer = std::make_shared<LocalShuffleWriter>(
       writeInfo.rootPath,
       writeInfo.queryId,
       writeInfo.shuffleId,
       writeInfo.numPartitions,
       maxBytesPerPartition,
-      /*sortedShuffle=*/false, // default to false for now
+      sortedShuffle,
       pool);
+  LOG(INFO) << "LocalPersistentShuffleFactory::createWriter COMPLETE";
+  return writer;
 }
-
-// Testing function to expose extractRowMetadata for tests.
-// This will be removed after reader changes.
-std::vector<RowMetadata> testingExtractRowMetadata(
-    const char* buffer,
-    size_t bufferSize,
-    bool sortedShuffle) {
-  return extractRowMetadata(buffer, bufferSize, sortedShuffle);
-}
-
 } // namespace facebook::presto::operators

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

@@ -53,12 +53,6 @@ inline std::strong_ordering compareKeys(
       reinterpret_cast<const uint8_t*>(key2.data() + key2.size()));
 }
 
-// Testing function to expose extractRowMetadata for tests.
-std::vector<RowMetadata> testingExtractRowMetadata(
-    const char* buffer,
-    size_t bufferSize,
-    bool sortedShuffle);
-
 // LocalShuffleWriteInfo is used for containing shuffle write information.
 // This struct is a 1:1 strict API mapping to
 // presto-spark-base/src/main/java/com/facebook/presto/spark/execution/PrestoSparkLocalShuffleWriteInfo.java
@@ -75,7 +69,6 @@ struct LocalShuffleWriteInfo {
   /// Structures are assumed to be encoded in JSON format.
   static LocalShuffleWriteInfo deserialize(const std::string& info);
 };
-
 // LocalShuffleReadInfo is used for containing shuffle read metadata
 // This struct is a 1:1 strict API mapping to
 // presto-spark-base/src/main/java/com/facebook/presto/spark/execution/PrestoSparkLocalShuffleReadInfo.java.
@@ -87,8 +80,6 @@ struct LocalShuffleReadInfo {
   std::string queryId;
   std::vector<std::string> partitionIds;
 
-  /// Deserializes shuffle information that is used by LocalPersistentShuffle.
-  /// Structures are assumed to be encoded in JSON format.
   static LocalShuffleReadInfo deserialize(const std::string& info);
 };
 

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

@@ -86,6 +86,11 @@ RowVectorPtr ShuffleRead::getOutput() {
       auto* batch = checked_pointer_cast<ShuffleRowBatch>(page.get());
       const auto& rows = batch->rows();
       for (const auto& row : rows) {
+        LOG(INFO) << "ShuffleRead::getOutput adding row to rows_: size=" << row.size();
+        if (row.size() >= 16) {
+          LOG(INFO) << "ShuffleRead::getOutput row data (hex): "
+                    << folly::hexlify(std::string_view(row.data(), std::min(row.size(), size_t(32))));
+        }
         rows_.emplace_back(row);
       }
     }