fix up

marin-ma · Oct 31, 2024 · fa507d1 · fa507d1
1 parent e192d02
commit fa507d1
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git 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
@@ -36,16 +36,25 @@ 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());
     auto serializer =
         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);