Fix calling PrestoSerializer::flush after PrestoSerializer::maxSerial…

…izedSize has wrong result (6973)
oap-project · Oct 10, 2023 · 8324109 · 8324109
1 parent 07afb52
commit 8324109
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/velox/serializers/PrestoSerializer.cpp b/velox/serializers/PrestoSerializer.cpp
@@ -919,10 +919,11 @@ class VectorStream {
           child->flush(out);
         }
         writeInt32(out, nullCount_ + nonNullCount_);
-        if (nullCount_ + nonNullCount_ == 0) {
+        if (!counted_ && nullCount_ + nonNullCount_ == 0) {
           // If nothing was added, there is still one offset in the wire
           // format.
           lengths_.appendOne<int32_t>(0);
+          counted_ = true;
         }
         lengths_.flush(out);
         flushNulls(out);
@@ -931,10 +932,11 @@ class VectorStream {
       case TypeKind::ARRAY:
         children_[0]->flush(out);
         writeInt32(out, nullCount_ + nonNullCount_);
-        if (nullCount_ + nonNullCount_ == 0) {
+        if (!counted_ && nullCount_ + nonNullCount_ == 0) {
           // If nothing was added, there is still one offset in the wire
           // format.
           lengths_.appendOne<int32_t>(0);
+          counted_ = true;
         }
         lengths_.flush(out);
         flushNulls(out);
@@ -946,10 +948,11 @@ class VectorStream {
         // hash table size. -1 means not included in serialization.
         writeInt32(out, -1);
         writeInt32(out, nullCount_ + nonNullCount_);
-        if (nullCount_ + nonNullCount_ == 0) {
+        if (!counted_ && nullCount_ + nonNullCount_ == 0) {
           // If nothing was added, there is still one offset in the wire
           // format.
           lengths_.appendOne<int32_t>(0);
+          counted_ = true;
         }
 
         lengths_.flush(out);
@@ -995,6 +998,7 @@ class VectorStream {
   int32_t nullCount_{0};
   int32_t totalLength_{0};
   bool hasLengths_{false};
+  bool counted_{false};
   ByteRange header_;
   ByteStream nulls_;
   ByteStream lengths_;

diff --git a/velox/serializers/tests/PrestoSerializerTest.cpp b/velox/serializers/tests/PrestoSerializerTest.cpp
@@ -434,6 +434,22 @@ TEST_P(PrestoSerializerTest, ioBufRoundTrip) {
   }
 }
 
+TEST_P(PrestoSerializerTest, roundTrip) {
+  VectorFuzzer::Options opts;
+  opts.timestampPrecision =
+      VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
+  opts.nullRatio = 0.1;
+  VectorFuzzer fuzzer(opts, pool_.get());
+
+  const size_t numRounds = 20;
+
+  for (size_t i = 0; i < numRounds; ++i) {
+    auto rowType = fuzzer.randRowType();
+    auto inputRowVector = fuzzer.fuzzInputRow(rowType);
+    testRoundTrip(inputRowVector);
+  }
+}
+
 INSTANTIATE_TEST_SUITE_P(
     PrestoSerializerTest,
     PrestoSerializerTest,