feat: Add PartitioningSerializer for Presto serialization

This commit introduces PartitioningSerializer, which buffers RowVectors across
multiple append() calls, partitions rows in-place using PartitionedVector,
and on flush() serializes each non-empty partition into a Presto
wire-format IOBuf. The serializer has no dependency on velox_exec —
it returns raw folly::IOBuf objects, leaving SerializedPage creation
to the caller.

Author: yingsu00

velox/serializers/CMakeLists.txt

@@ -25,6 +25,7 @@ velox_link_libraries(
 velox_add_library(
   velox_presto_serializer
   CompactRowSerializer.cpp
+  PartitioningSerializer.cpp
   PrestoSerializer.cpp
   UnsafeRowSerializer.cpp
   PrestoBatchVectorSerializer.cpp

velox/serializers/PartitioningSerializer.cpp

@@ -0,0 +1,109 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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 "velox/serializers/PartitioningSerializer.h"
+
+#include "velox/common/memory/ByteStream.h"
+#include "velox/vector/VectorStream.h"
+
+namespace facebook::velox::serializer::presto {
+
+PartitioningSerializer::PartitioningSerializer(
+    RowTypePtr inputType,
+    uint32_t numPartitions,
+    const SerdeOpts& opts,
+    memory::MemoryPool* pool)
+    : type_(std::move(inputType)),
+      numPartitions_(numPartitions),
+      opts_(opts),
+      pool_(pool),
+      rowsPerPartition_(numPartitions, 0) {
+  VELOX_CHECK_GT(numPartitions_, 0);
+  VELOX_CHECK_NOT_NULL(pool_);
+}
+
+void PartitioningSerializer::append(
+    const RowVectorPtr& input,
+    const std::vector<uint32_t>& partitions) {
+  VELOX_CHECK_NOT_NULL(input);
+  VELOX_CHECK_EQ(
+      input->size(),
+      partitions.size(),
+      "partitions.size() must equal input->size()");
+
+  if (input->size() == 0) {
+    return;
+  }
+
+  PartitionBuildContext ctx;
+  auto pv = PartitionedVector::create(
+      std::static_pointer_cast<BaseVector>(input),
+      partitions,
+      numPartitions_,
+      ctx,
+      pool_);
+
+  const auto* offsets = pv->rawPartitionOffsets();
+  vector_size_t prevOffset = 0;
+  for (uint32_t p = 0; p < numPartitions_; ++p) {
+    rowsPerPartition_[p] += offsets[p] - prevOffset;
+    prevOffset = offsets[p];
+  }
+
+  partitionedPages_.push_back(std::move(pv));
+  bytesBuffered_ += input->retainedSize();
+  rowsBuffered_ += static_cast<int64_t>(input->size());
+}
+
+std::map<uint32_t, std::unique_ptr<folly::IOBuf>>
+PartitioningSerializer::flush() {
+  if (partitionedPages_.empty()) {
+    return {};
+  }
+
+  std::map<uint32_t, std::unique_ptr<folly::IOBuf>> result;
+
+  for (uint32_t p = 0; p < numPartitions_; ++p) {
+    const auto numRows = rowsPerPartition_[p];
+    if (numRows == 0) {
+      continue;
+    }
+
+    auto streamGroup = std::make_unique<VectorStreamGroup>(pool_, &serde_);
+    streamGroup->createStreamTree(type_, numRows, &opts_);
+
+    for (auto& pv : partitionedPages_) {
+      auto slice = std::dynamic_pointer_cast<RowVector>(pv->partitionAt(p));
+      if (slice && slice->size() > 0) {
+        streamGroup->append(slice);
+      }
+    }
+
+    const int32_t estimatedSize =
+        std::max(static_cast<int32_t>(bytesBuffered_ / numPartitions_), 1024);
+    IOBufOutputStream out(*pool_, nullptr, estimatedSize);
+    streamGroup->flush(&out);
+    result[p] = out.getIOBuf();
+  }
+
+  partitionedPages_.clear();
+  std::fill(rowsPerPartition_.begin(), rowsPerPartition_.end(), 0);
+  bytesBuffered_ = 0;
+  rowsBuffered_ = 0;
+
+  return result;
+}
+
+} // namespace facebook::velox::serializer::presto

velox/serializers/PartitioningSerializer.h

@@ -0,0 +1,84 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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 <map>
+
+#include <folly/io/IOBuf.h>
+
+#include "velox/serializers/PrestoSerializer.h"
+#include "velox/vector/PartitionedVector.h"
+
+namespace facebook::velox::serializer::presto {
+
+using SerdeOpts = PrestoVectorSerde::PrestoOptions;
+
+/// Buffers one or more RowVectors and serializes them into per-partition
+/// Presto-format pages. Rows are assigned to partitions via the partitions
+/// array passed to append(). The serializer accumulates partitioned data
+/// across multiple append() calls and emits all buffered data on flush().
+///
+/// Thread-safety: not thread-safe.
+class PartitioningSerializer {
+ public:
+  /// Constructs a serializer for the given input type, number of partitions,
+  /// and serde options. All serialized output uses the provided memory pool.
+  PartitioningSerializer(
+      RowTypePtr inputType,
+      uint32_t numPartitions,
+      const SerdeOpts& opts,
+      memory::MemoryPool* pool);
+
+  /// Appends 'input' to the buffer. partitions[i] must be in [0,
+  /// numPartitions) and gives the destination partition for row i. The size
+  /// of 'partitions' must equal input->size().
+  void append(
+      const RowVectorPtr& input,
+      const std::vector<uint32_t>& partitions);
+
+  /// Serializes all buffered data into Presto wire format. Returns one IOBuf
+  /// per non-empty partition, keyed by partition index. The caller is
+  /// responsible for wrapping the IOBuf in an appropriate page type (e.g.
+  /// PrestoSerializedPage). Clears the internal buffer after flushing.
+  std::map<uint32_t, std::unique_ptr<folly::IOBuf>> flush();
+
+  /// Returns the total in-memory bytes of all buffered RowVectors.
+  int64_t bytesBuffered() const {
+    return bytesBuffered_;
+  }
+
+  /// Returns the total number of rows in all buffered RowVectors.
+  int64_t rowsBuffered() const {
+    return rowsBuffered_;
+  }
+
+ private:
+  const RowTypePtr type_;
+  const uint32_t numPartitions_;
+  const SerdeOpts opts_;
+  PrestoVectorSerde serde_;
+  memory::MemoryPool* const pool_;
+
+  std::vector<PartitionedVectorPtr> partitionedPages_;
+
+  /// Cumulative number of rows assigned to each partition across all appends.
+  std::vector<vector_size_t> rowsPerPartition_;
+
+  int64_t bytesBuffered_{0};
+  int64_t rowsBuffered_{0};
+};
+
+} // namespace facebook::velox::serializer::presto

velox/serializers/tests/CMakeLists.txt

@@ -35,6 +35,7 @@ target_link_libraries(
 add_executable(
   velox_serializer_test
   CompactRowSerializerTest.cpp
+  PartitioningSerializerTest.cpp
   PrestoOutputStreamListenerTest.cpp
   PrestoSerializerTest.cpp
   SerializedPageFileTest.cpp