From 07e779858fa6c5e5744d5d5617ff5da9e50eea30 Mon Sep 17 00:00:00 2001
From: Kevin Wilfong <kevinwilfong@meta.com>
Date: Tue, 2 Jul 2024 15:06:06 -0700
Subject: [PATCH] StringWriter needs to implement finalizeNull (#10376)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/10376

When StringWriter commits it calls proxy_.prepareForReuse to update the state of the proxy, either
advancing it if a value is being written or resetting it if a null is being written.

StringWriter needs to implement finalizeNull to do the same in case its parent commits a null.  When
a writer for a complex type commits a null, it invokes finalizeNull on its children to reset their state.
StringWriter does not implement this today which means, e.g. if it is the writer for the elements of
an Array, if the code has written an uncommitted String to the Array when the Array is committed as
null, the next Array that gets written will end up with its StringWriter still holding the state from that
last string in the previous Array, meaning its contents could appear at the beginning of the first string
in that Array.

We were also not calling proxy_.prepareForReuse in commitNull (only commit(false)) which resulted
in a similar issue, this change also fixes that.

This partially addresses the issue identified in https://github.com/facebookincubator/velox/issues/10162

Reviewed By: mbasmanova

Differential Revision: D59291362
---
 velox/expression/ComplexWriterTypes.h       |  2 ++
 velox/expression/VectorWriters.h            |  7 +++++-
 velox/expression/tests/StringWriterTest.cpp | 24 +++++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/velox/expression/ComplexWriterTypes.h b/velox/expression/ComplexWriterTypes.h
index c7cb91c233a..1874d8c014b 100644
--- a/velox/expression/ComplexWriterTypes.h
+++ b/velox/expression/ComplexWriterTypes.h
@@ -50,6 +50,8 @@ class VectorWriterBase {
   virtual void commit(bool isSet) = 0;
   virtual void ensureSize(size_t size) = 0;
   virtual void finish() {}
+  // Implementations that write variable length data or complex types should
+  // override this to reset their state and that of their children.
   virtual void finalizeNull() {}
   virtual ~VectorWriterBase() {}
   vector_size_t offset_ = 0;
diff --git a/velox/expression/VectorWriters.h b/velox/expression/VectorWriters.h
index df89bdf4836..fef8e6977b6 100644
--- a/velox/expression/VectorWriters.h
+++ b/velox/expression/VectorWriters.h
@@ -373,16 +373,21 @@ struct VectorWriter<
 
   void commitNull() {
     proxy_.vector_->setNull(proxy_.offset_, true);
+    finalizeNull();
+  }
+
+  void finalizeNull() override {
+    proxy_.prepareForReuse(false);
   }
 
   void commit(bool isSet) override {
     // this code path is called when the slice is top-level
     if (isSet) {
       proxy_.finalize();
+      proxy_.prepareForReuse(true);
     } else {
       commitNull();
     }
-    proxy_.prepareForReuse(isSet);
   }
 
   void setOffset(vector_size_t offset) override {
diff --git a/velox/expression/tests/StringWriterTest.cpp b/velox/expression/tests/StringWriterTest.cpp
index 9f8543d95fa..aec9d9cacb8 100644
--- a/velox/expression/tests/StringWriterTest.cpp
+++ b/velox/expression/tests/StringWriterTest.cpp
@@ -18,9 +18,11 @@
 #include <glog/logging.h>
 #include "folly/Range.h"
 #include "gtest/gtest.h"
+#include "velox/expression/VectorWriters.h"
 #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
 
 namespace facebook::velox::expressions::test {
+using namespace facebook::velox::test;
 
 class StringWriterTest : public functions::test::FunctionBaseTest {};
 
@@ -103,4 +105,26 @@ TEST_F(StringWriterTest, copyFromCString) {
 
   ASSERT_EQ(vector->valueAt(0), "1 2 3 4 5 "_sv);
 }
+
+TEST_F(StringWriterTest, vectorWriter) {
+  auto vector = makeFlatVector<StringView>(3);
+  exec::VectorWriter<Varchar> writer;
+  writer.init(*vector);
+  writer.setOffset(0);
+  writer.current().copy_from("1 2 3");
+  writer.commitNull();
+
+  writer.setOffset(1);
+  writer.current().copy_from("4 5 6");
+  writer.commit(true);
+
+  writer.setOffset(2);
+  writer.current().copy_from("7 8 9");
+  writer.commit(false);
+  writer.finish();
+
+  auto expected = std::vector<std::optional<std::string>>{
+      std::nullopt, "4 5 6", std::nullopt};
+  assertEqualVectors(vector, makeNullableFlatVector(expected));
+}
 } // namespace facebook::velox::expressions::test