refactor: Change native pos API to return BaseSerializedPage (#26692)

Summary:

Use BaseSerializedPage directly from shuffle. This allows seamless handle of shuffle data all the way to ShuffleRead(subclass of Exchange operator)

Reviewed By: xiaoxmeng

Differential Revision: D87852058

Author: tanjialiang

presto-native-execution/presto_cpp/main/PrestoExchangeSource.cpp

@@ -323,7 +323,7 @@ void PrestoExchangeSource::processDataResponse(
         contentLength, 0, "next token is not set in non-empty data response");
   }
 
-  std::unique_ptr<exec::SerializedPage> page;
+  std::unique_ptr<exec::SerializedPageBase> page;
   const bool empty = response->empty();
   if (!empty) {
     std::vector<std::unique_ptr<folly::IOBuf>> iobufs;
@@ -351,7 +351,7 @@ void PrestoExchangeSource::processDataResponse(
     }
 
     if (enableBufferCopy_) {
-      page = std::make_unique<exec::SerializedPage>(
+      page = std::make_unique<exec::PrestoSerializedPage>(
           std::move(singleChain), [pool = pool_](folly::IOBuf& iobuf) {
             int64_t freedBytes{0};
             // Free the backed memory from MemoryAllocator on page dtor
@@ -365,7 +365,7 @@ void PrestoExchangeSource::processDataResponse(
             PrestoExchangeSource::updateMemoryUsage(-freedBytes);
           });
     } else {
-      page = std::make_unique<exec::SerializedPage>(
+      page = std::make_unique<exec::PrestoSerializedPage>(
           std::move(singleChain), [totalBytes](folly::IOBuf& iobuf) {
             PrestoExchangeSource::updateMemoryUsage(-totalBytes);
           });

presto-native-execution/presto_cpp/main/PrestoServer.cpp

@@ -41,6 +41,7 @@
 #include "presto_cpp/main/operators/PartitionAndSerialize.h"
 #include "presto_cpp/main/operators/ShuffleExchangeSource.h"
 #include "presto_cpp/main/operators/ShuffleRead.h"
+#include "presto_cpp/main/operators/ShuffleWrite.h"
 #include "presto_cpp/main/types/PrestoToVeloxQueryPlan.h"
 #include "presto_cpp/main/types/VeloxPlanConversion.h"
 #include "velox/common/base/Counters.h"

presto-native-execution/presto_cpp/main/SessionProperties.cpp

@@ -345,7 +345,7 @@ SessionProperties::SessionProperties() {
       "creating tiny SerializedPages. For "
       "PartitionedOutputNode::Kind::kPartitioned, PartitionedOutput operator"
       "would buffer up to that number of bytes / number of destinations for "
-      "each destination before producing a SerializedPage.",
+      "each destination before producing a SerializedPageBase.",
       BIGINT(),
       false,
       QueryConfig::kMaxPartitionedOutputBufferSize,

presto-native-execution/presto_cpp/main/SessionProperties.h

@@ -268,7 +268,7 @@ class SessionProperties {
   /// creating tiny SerializedPages. For
   /// PartitionedOutputNode::Kind::kPartitioned, PartitionedOutput operator
   /// would buffer up to that number of bytes / number of destinations for each
-  /// destination before producing a SerializedPage.
+  /// destination before producing a SerializedPageBase.
   static constexpr const char* kMaxPartitionedOutputBufferSize =
       "native_max_page_partitioning_buffer_size";
 

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

@@ -42,14 +42,14 @@ BroadcastExchangeSource::request(
 
   return folly::makeTryWith([&]() -> Response {
     int64_t totalBytes = 0;
-    std::vector<std::unique_ptr<velox::exec::SerializedPage>> pages;
+    std::vector<std::unique_ptr<velox::exec::SerializedPageBase>> pages;
 
     while (totalBytes < maxBytes && reader_->hasNext()) {
       auto buffer = reader_->next();
       VELOX_CHECK_NOT_NULL(buffer);
 
       auto ioBuf = folly::IOBuf::wrapBuffer(buffer->as<char>(), buffer->size());
-      auto page = std::make_unique<velox::exec::SerializedPage>(
+      auto page = std::make_unique<velox::exec::PrestoSerializedPage>(
           std::move(ioBuf), [buffer](auto& /*unused*/) {});
       pages.push_back(std::move(page));
 

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

@@ -105,6 +105,30 @@ class SortedFileInputStream final : public velox::common::FileInputStream,
   std::string currentValue_;
 };
 
+class LocalShuffleSerializedPage : public ShuffleSerializedPage {
+ public:
+  LocalShuffleSerializedPage(
+      const std::vector<std::string_view>& rows,
+      velox::BufferPtr buffer)
+      : rows_{std::move(rows)}, buffer_{std::move(buffer)} {}
+
+  const std::vector<std::string_view>& rows() override {
+    return rows_;
+  }
+
+  uint64_t size() const override {
+    return buffer_->size();
+  }
+
+  std::optional<int64_t> numRows() const override {
+    return rows_.size();
+  }
+
+ private:
+  const std::vector<std::string_view> rows_;
+  const velox::BufferPtr buffer_;
+};
+
 std::vector<RowMetadata>
 extractRowMetadata(const char* buffer, size_t bufferSize, bool sortedShuffle) {
   std::vector<RowMetadata> rows;
@@ -450,9 +474,9 @@ void LocalShuffleReader::initSortedShuffleRead() {
   }
 }
 
-std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextSorted(
-    uint64_t maxBytes) {
-  std::vector<std::unique_ptr<ReadBatch>> batches;
+std::vector<std::unique_ptr<ShuffleSerializedPage>>
+LocalShuffleReader::nextSorted(uint64_t maxBytes) {
+  std::vector<std::unique_ptr<ShuffleSerializedPage>> batches;
 
   if (merge_ == nullptr) {
     return batches;
@@ -469,7 +493,7 @@ std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextSorted(
     if (bufferUsed + data.size() > maxBytes) {
       if (bufferUsed > 0) {
         batches.push_back(
-            std::make_unique<ReadBatch>(
+            std::make_unique<LocalShuffleSerializedPage>(
                 std::move(rows), std::move(batchBuffer)));
         return batches;
       }
@@ -489,15 +513,16 @@ std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextSorted(
 
   if (!rows.empty()) {
     batches.push_back(
-        std::make_unique<ReadBatch>(std::move(rows), std::move(batchBuffer)));
+        std::make_unique<LocalShuffleSerializedPage>(
+            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;
+std::vector<std::unique_ptr<ShuffleSerializedPage>>
+LocalShuffleReader::nextUnsorted(uint64_t maxBytes) {
+  std::vector<std::unique_ptr<ShuffleSerializedPage>> batches;
   uint64_t totalBytes{0};
 
   while (readPartitionFileIndex_ < readPartitionFiles_.size()) {
@@ -527,13 +552,14 @@ std::vector<std::unique_ptr<ReadBatch>> LocalShuffleReader::nextUnsorted(
 
     totalBytes += fileSize;
     batches.push_back(
-        std::make_unique<ReadBatch>(std::move(rows), std::move(buffer)));
+        std::make_unique<LocalShuffleSerializedPage>(
+            std::move(rows), std::move(buffer)));
   }
 
   return batches;
 }
 
-folly::SemiFuture<std::vector<std::unique_ptr<ReadBatch>>>
+folly::SemiFuture<std::vector<std::unique_ptr<ShuffleSerializedPage>>>
 LocalShuffleReader::next(uint64_t maxBytes) {
   VELOX_CHECK(
       initialized_,

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

@@ -173,7 +173,7 @@ class LocalShuffleReader : public ShuffleReader {
   /// For sorted shuffle, this opens all shuffle files and prepares k-way merge.
   void initialize();
 
-  folly::SemiFuture<std::vector<std::unique_ptr<ReadBatch>>> next(
+  folly::SemiFuture<std::vector<std::unique_ptr<ShuffleSerializedPage>>> next(
       uint64_t maxBytes) override;
 
   void noMoreData(bool success) override;
@@ -192,10 +192,12 @@ class LocalShuffleReader : public ShuffleReader {
   void initSortedShuffleRead();
 
   // Reads sorted shuffle data using k-way merge with TreeOfLosers.
-  std::vector<std::unique_ptr<ReadBatch>> nextSorted(uint64_t maxBytes);
+  std::vector<std::unique_ptr<ShuffleSerializedPage>> nextSorted(
+      uint64_t maxBytes);
 
   // Reads unsorted shuffle data in batch-based file reading.
-  std::vector<std::unique_ptr<ReadBatch>> nextUnsorted(uint64_t maxBytes);
+  std::vector<std::unique_ptr<ShuffleSerializedPage>> nextUnsorted(
+      uint64_t maxBytes);
 
   const std::string rootPath_;
   const std::string queryId_;

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

@@ -35,31 +35,29 @@ ShuffleExchangeSource::request(
     std::chrono::microseconds /*maxWait*/) {
   auto nextBatch = [this, maxBytes]() {
     return std::move(shuffleReader_->next(maxBytes))
-        .deferValue([this](std::vector<std::unique_ptr<ReadBatch>>&& batches) {
-          std::vector<velox::ContinuePromise> promises;
-          int64_t totalBytes{0};
-          {
-            std::lock_guard<std::mutex> l(queue_->mutex());
-            if (batches.empty()) {
-              atEnd_ = true;
-              queue_->enqueueLocked(nullptr, promises);
-            } else {
-              for (auto& batch : batches) {
-                totalBytes = batch->data->size();
-                VELOX_CHECK_LE(totalBytes, std::numeric_limits<int32_t>::max());
-                ++numBatches_;
-                queue_->enqueueLocked(
-                    std::make_unique<ShuffleRowBatch>(std::move(batch)),
-                    promises);
+        .deferValue(
+            [this](
+                std::vector<std::unique_ptr<ShuffleSerializedPage>>&& batches) {
+              std::vector<velox::ContinuePromise> promises;
+              int64_t totalBytes{0};
+              {
+                std::lock_guard<std::mutex> l(queue_->mutex());
+                if (batches.empty()) {
+                  atEnd_ = true;
+                  queue_->enqueueLocked(nullptr, promises);
+                } else {
+                  for (auto& batch : batches) {
+                    ++numBatches_;
+                    queue_->enqueueLocked(std::move(batch), promises);
+                  }
+                }
               }
-            }
-          }
 
-          for (auto& promise : promises) {
-            promise.setValue();
-          }
-          return folly::makeFuture(Response{totalBytes, atEnd_});
-        })
+              for (auto& promise : promises) {
+                promise.setValue();
+              }
+              return folly::makeFuture(Response{totalBytes, atEnd_});
+            })
         .deferError(
             [](folly::exception_wrapper e) mutable
                 -> ShuffleExchangeSource::Response {

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

@@ -14,32 +14,11 @@
 #pragma once
 
 #include "presto_cpp/main/operators/ShuffleInterface.h"
-#include "presto_cpp/main/operators/ShuffleWrite.h"
 #include "velox/core/PlanNode.h"
-#include "velox/exec/Exchange.h"
 #include "velox/exec/Operator.h"
 
 namespace facebook::presto::operators {
 
-class ShuffleRowBatch : public velox::exec::SerializedPage {
- public:
-  explicit ShuffleRowBatch(
-      std::unique_ptr<ReadBatch> rowBatch)
-      : velox::exec::
-            SerializedPage{folly::IOBuf::wrapBuffer(
-                  rowBatch->data->as<char>(), rowBatch->data->size()), nullptr, rowBatch->rows.size()},
-        rowBatch_{std::move(rowBatch)} {}
-
-  ~ShuffleRowBatch() override {}
-
-  const std::vector<std::string_view>& rows() const {
-    return rowBatch_->rows;
-  }
-
- private:
-  const std::unique_ptr<ReadBatch> rowBatch_;
-};
-
 class ShuffleExchangeSource : public velox::exec::ExchangeSource {
  public:
   ShuffleExchangeSource(

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

@@ -14,8 +14,13 @@
 #pragma once
 
 #include <fmt/format.h>
+#include "velox/exec/Exchange.h"
 #include "velox/exec/Operator.h"
 
+namespace facebook::velox {
+class ByteInputStream;
+}
+
 namespace facebook::presto::operators {
 
 class ShuffleWriter {
@@ -44,12 +49,21 @@ class ShuffleWriter {
   }
 };
 
-struct ReadBatch {
-  std::vector<std::string_view> rows;
-  velox::BufferPtr data;
+class ShuffleSerializedPage : public velox::exec::SerializedPageBase {
+ public:
+  ShuffleSerializedPage() = default;
+  ~ShuffleSerializedPage() override = default;
+
+  std::unique_ptr<velox::ByteInputStream> prepareStreamForDeserialize()
+      override {
+    VELOX_UNSUPPORTED();
+  }
+
+  std::unique_ptr<folly::IOBuf> getIOBuf() const override {
+    VELOX_UNSUPPORTED();
+  }
 
-  ReadBatch(std::vector<std::string_view>&& _rows, velox::BufferPtr&& _data)
-      : rows{std::move(_rows)}, data{std::move(_data)} {}
+  virtual const std::vector<std::string_view>& rows() = 0;
 };
 
 class ShuffleReader {
@@ -58,10 +72,11 @@ class ShuffleReader {
 
   /// Fetch the next batch of rows from the shuffle reader.
   /// @param bytes Maximum number of bytes to read in this batch.
-  /// @return A semi-future resolving to a vector of ReadBatch pointers, where
-  /// each ReadBatch contains rows and associated data buffers.
-  virtual folly::SemiFuture<std::vector<std::unique_ptr<ReadBatch>>> next(
-      uint64_t maxBytes) = 0;
+  /// @return A semi-future resolving to a vector of ShuffleSerializedPage
+  /// pointers, where each ShuffleSerializedPage contains rows and associated
+  /// data buffers.
+  virtual folly::SemiFuture<std::vector<std::unique_ptr<ShuffleSerializedPage>>>
+  next(uint64_t maxBytes) = 0;
 
   /// Tell the shuffle system the reader is done. May be called with 'success'
   /// true before reading all the data. This happens when a query has a LIMIT or