From 39f2159d744ef05ed63f8d05624769235b4a9d44 Mon Sep 17 00:00:00 2001
From: Rong Ma <rong.ma@intel.com>
Date: Fri, 6 Sep 2024 09:12:27 +0000
Subject: [PATCH] fix up

---
 velox/serializers/CompactRowSerializer.cpp          | 10 +++++++---
 .../serializers/tests/CompactRowSerializerTest.cpp  | 13 +++++++++++--
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/velox/serializers/CompactRowSerializer.cpp b/velox/serializers/CompactRowSerializer.cpp
index a2579bde9efd..e63814f4f3dc 100644
--- a/velox/serializers/CompactRowSerializer.cpp
+++ b/velox/serializers/CompactRowSerializer.cpp
@@ -75,9 +75,13 @@ class CompactRowVectorSerializer : public IterativeVectorSerializer {
     vector_size_t index = 0;
     for (const auto& range : ranges) {
       if (range.size == 1) {
-        *(TRowSize*)(rawBuffer) = folly::Endian::big(rowSize[index++]);
-        static const auto offset = sizeof(TRowSize);
-        row.serialize(range.begin, range.size, rawBuffer, &offset);
+        // Fast path for single-row serialization.
+        *(TRowSize*)(rawBuffer + offset) = folly::Endian::big(rowSize[index]);
+        static const auto rowSizeOffset = sizeof(TRowSize);
+        row.serialize(
+            range.begin, range.size, rawBuffer + offset, &rowSizeOffset);
+        offset += rowSize[index] + sizeof(TRowSize);
+        index++;
       } else {
         raw_vector<size_t> offsets(range.size);
         for (auto i = 0; i < range.size; ++i) {
diff --git a/velox/serializers/tests/CompactRowSerializerTest.cpp b/velox/serializers/tests/CompactRowSerializerTest.cpp
index 834619a3b3dc..c0a091b5e0b4 100644
--- a/velox/serializers/tests/CompactRowSerializerTest.cpp
+++ b/velox/serializers/tests/CompactRowSerializerTest.cpp
@@ -36,8 +36,16 @@ class CompactRowSerializerTest : public ::testing::Test,
   void serialize(RowVectorPtr rowVector, std::ostream* output) {
     auto numRows = rowVector->size();
 
+    // Serialize with different range size.
     std::vector<IndexRange> rows;
-    rows.push_back(IndexRange{0, numRows});
+    vector_size_t offset = 0;
+    vector_size_t rangeSize = 1;
+    while (offset < numRows) {
+      auto size = std::min<vector_size_t>(rangeSize, numRows - offset);
+      rows.push_back(IndexRange{offset, size});
+      offset += size;
+      rangeSize = checkedMultiply<vector_size_t>(rangeSize, 2);
+    }
 
     auto arena = std::make_unique<StreamArena>(pool_.get());
     auto rowType = asRowType(rowVector->type());
@@ -45,7 +53,8 @@ class CompactRowSerializerTest : public ::testing::Test,
         serde_->createIterativeSerializer(rowType, numRows, arena.get());
 
     Scratch scratch;
-    serializer->append(rowVector, folly::Range(rows.data(), 1), scratch);
+    serializer->append(
+        rowVector, folly::Range(rows.data(), rows.size()), scratch);
     auto size = serializer->maxSerializedSize();
     OStreamOutputStream out(output);
     serializer->flush(&out);