misc: Refactor ShuffleTest (prestodb#26643)

Summary:

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

@@ -218,6 +218,16 @@ inline std::string createShuffleFileName(
 }
 } // namespace
 
+std::string LocalShuffleWriteInfo::serialize() const {
+  json obj;
+  obj["rootPath"] = rootPath;
+  obj["queryId"] = queryId;
+  obj["shuffleId"] = shuffleId;
+  obj["numPartitions"] = numPartitions;
+  obj["sortedShuffle"] = sortedShuffle;
+  return obj.dump();
+}
+
 LocalShuffleWriteInfo LocalShuffleWriteInfo::deserialize(
     const std::string& info) {
   const auto jsonReadInfo = json::parse(info);
@@ -230,6 +240,15 @@ LocalShuffleWriteInfo LocalShuffleWriteInfo::deserialize(
   return shuffleInfo;
 }
 
+std::string LocalShuffleReadInfo::serialize() const {
+  json obj;
+  obj["rootPath"] = rootPath;
+  obj["queryId"] = queryId;
+  obj["partitionIds"] = partitionIds;
+  obj["sortedShuffle"] = sortedShuffle;
+  return obj.dump();
+}
+
 LocalShuffleReadInfo LocalShuffleReadInfo::deserialize(
     const std::string& info) {
   const auto jsonReadInfo = json::parse(info);
@@ -357,6 +376,8 @@ void LocalShuffleWriter::collect(
       sortedShuffle_ || key.empty(),
       "key '{}' must be empty for non-sorted shuffle",
       key);
+  velox::common::testutil::TestValue::adjust(
+      "facebook::presto::operators::LocalShuffleWriter::collect", this);
 
   const auto rowSize = this->rowSize(key.size(), data.size());
   auto& buffer = inProgressPartitions_[partition];
@@ -521,6 +542,8 @@ LocalShuffleReader::next(uint64_t maxBytes) {
   VELOX_CHECK(
       initialized_,
       "LocalShuffleReader::initialize() must be called before next()");
+  velox::common::testutil::TestValue::adjust(
+      "facebook::presto::operators::LocalShuffleReader::next", this);
 
   return folly::makeSemiFuture(
       sortedShuffle_ ? nextSorted(maxBytes) : nextUnsorted(maxBytes));
@@ -597,13 +620,4 @@ std::shared_ptr<ShuffleWriter> LocalPersistentShuffleFactory::createWriter(
       writeInfo.sortedShuffle,
       pool);
 }
-
-// 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 bool compareKeys(std::string_view key1, std::string_view key2) noexcept {
   return key1.size() < 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
@@ -72,11 +66,13 @@ struct LocalShuffleWriteInfo {
   uint32_t shuffleId;
   bool sortedShuffle;
 
+  /// Serializes shuffle information to JSON format.
+  std::string serialize() const;
+
   /// Deserializes shuffle information that is used by LocalPersistentShuffle.
   /// 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.
@@ -89,8 +85,9 @@ struct LocalShuffleReadInfo {
   std::vector<std::string> partitionIds;
   bool sortedShuffle;
 
-  /// Deserializes shuffle information that is used by LocalPersistentShuffle.
-  /// Structures are assumed to be encoded in JSON format.
+  /// Serializes shuffle information to JSON format.
+  std::string serialize() const;
+
   static LocalShuffleReadInfo deserialize(const std::string& info);
 };