Fix prepareForReuse for string vectors inside array/map/struct (#7465)

Summary:
BaseVector::prepareForReuse used to clear all string buffers for string vectors
inside a complex types vector, effectively prevented reuse of the string vectors.

This happened because BaseVector::prepareForReuse for complex vectors calls
BaseVector::resize(0) for child vectors and FlatVector<StringView>::resize(0) used
to clear all string buffers.

A fix is to modify FlatVector<StringView>::resize(0) to keep at most one string
buffer, just like FlatVector<StringView>::prepareForReuse() does.

Pull Request resolved: #7465

Reviewed By: spershin

Differential Revision: D51097758

Pulled By: mbasmanova

fbshipit-source-id: 7ab1e44c1bc5d4ada391668b4b30a52b845ab6be

velox/vector/FlatVector-inl.h

@@ -437,7 +437,7 @@ void FlatVector<T>::resize(vector_size_t newSize, bool setNotNull) {
       SimpleVector<StringView>::resizeIsAsciiIfNotEmpty(newSize, false);
     }
     if (newSize == 0) {
-      clearStringBuffers();
+      keepAtMostOneStringBuffer();
     }
   } else {
     resizeValues(newSize, std::nullopt);

velox/vector/FlatVector.cpp

@@ -89,18 +89,7 @@ void FlatVector<StringView>::prepareForReuse() {
     rawValues_ = nullptr;
   }
 
-  // Check string buffers. Keep at most one singly-referenced buffer if it is
-  // not too large.
-  if (!stringBuffers_.empty()) {
-    auto& firstBuffer = stringBuffers_.front();
-    if (firstBuffer->isMutable() &&
-        firstBuffer->capacity() <= kMaxStringSizeForReuse) {
-      firstBuffer->setSize(0);
-      setStringBuffers({firstBuffer});
-    } else {
-      clearStringBuffers();
-    }
-  }
+  keepAtMostOneStringBuffer();
 
   // Clear the StringViews to avoid referencing freed memory.
   if (rawValues_) {

velox/vector/FlatVector.h

@@ -500,6 +500,23 @@ class FlatVector final : public SimpleVector<T> {
       vector_size_t newSize,
       const std::optional<T>& initialValue);
 
+  // Check string buffers. Keep at most one singly-referenced buffer if it is
+  // not too large.
+  void keepAtMostOneStringBuffer() {
+    if (stringBuffers_.empty()) {
+      return;
+    }
+
+    auto& firstBuffer = stringBuffers_.front();
+    if (firstBuffer->isMutable() &&
+        firstBuffer->capacity() <= kMaxStringSizeForReuse) {
+      firstBuffer->setSize(0);
+      setStringBuffers({firstBuffer});
+    } else {
+      clearStringBuffers();
+    }
+  }
+
   // Contiguous values.
   // If strings, these are velox::StringViews into memory held by
   // 'stringBuffers_'

velox/vector/tests/VectorPrepareForReuseTest.cpp

@@ -224,6 +224,40 @@ TEST_F(VectorPrepareForReuseTest, arrays) {
   ASSERT_EQ(originalSize, vector->retainedSize());
 }
 
+TEST_F(VectorPrepareForReuseTest, arrayOfStrings) {
+  VectorPtr vector = makeArrayVector<std::string>(
+      1'000,
+      [](auto /*row*/) { return 1; },
+      [](auto row, auto index) {
+        return std::string(20 + index, 'a' + row % 5);
+      });
+  auto originalSize = vector->retainedSize();
+  BaseVector* originalVector = vector.get();
+
+  MemoryAllocationChecker allocationChecker(pool());
+  BaseVector::prepareForReuse(vector, vector->size());
+  ASSERT_EQ(originalVector, vector.get());
+  ASSERT_EQ(originalSize, vector->retainedSize());
+
+  auto* arrayVector = vector->as<ArrayVector>();
+  for (auto i = 0; i < 1'000; i++) {
+    ASSERT_EQ(0, arrayVector->sizeAt(i));
+    ASSERT_EQ(0, arrayVector->offsetAt(i));
+  }
+
+  // Cannot use BaseVector::copy because it is too smart and acquired string
+  // buffers instead of copying the strings.
+  auto* elementsVector = arrayVector->elements()->as<FlatVector<StringView>>();
+  elementsVector->resize(1'000);
+  for (auto i = 0; i < 1'000; i++) {
+    arrayVector->setOffsetAndSize(i, i, 1);
+    std::string newValue(21, 'b' + i % 7);
+    elementsVector->set(i, StringView(newValue));
+  }
+
+  ASSERT_EQ(originalSize, vector->retainedSize());
+}
+
 TEST_F(VectorPrepareForReuseTest, dataDependentFlags) {
   auto size = 10;
 

velox/vector/tests/VectorTest.cpp

@@ -1649,6 +1649,23 @@ TEST_F(VectorTest, resizeStringAsciiness) {
   ASSERT_FALSE(stringVector->isAscii(rows));
 }
 
+TEST_F(VectorTest, resizeZeroString) {
+  auto vector = makeFlatVector<std::string>(
+      {"This is a string",
+       "This is another string",
+       "This is the third string"});
+  ASSERT_EQ(1, vector->stringBuffers().size());
+  ASSERT_LT(0, vector->stringBuffers()[0]->size());
+
+  const auto capacity = vector->stringBuffers()[0]->capacity();
+  ASSERT_GT(capacity, 0);
+
+  vector->resize(0);
+  ASSERT_EQ(1, vector->stringBuffers().size());
+  ASSERT_EQ(0, vector->stringBuffers()[0]->size());
+  ASSERT_EQ(capacity, vector->stringBuffers()[0]->capacity());
+}
+
 TEST_F(VectorTest, copyNoRows) {
   {
     auto source = makeFlatVector<int32_t>({1, 2, 3});

velox/vector/tests/utils/VectorMaker.h

@@ -367,16 +367,16 @@ class VectorMaker {
     auto numElements =
         createOffsetsAndSizes(size, sizeAt, isNullAt, &nulls, &offsets, &sizes);
 
-    auto flatVector =
-        BaseVector::create<FlatVector<T>>(type->childAt(0), numElements, pool_);
+    auto flatVector = BaseVector::create<FlatVector<EvalType<T>>>(
+        type->childAt(0), numElements, pool_);
     vector_size_t currentIndex = 0;
     for (vector_size_t i = 0; i < size; ++i) {
       if (isNullAt && isNullAt(i)) {
         continue;
       }
       for (vector_size_t j = 0; j < sizeAt(i); ++j) {
-        auto ret = valueAt(i, j);
-        flatVector->set(currentIndex, valueAt(i, j));
+        auto value = valueAt(i, j);
+        flatVector->set(currentIndex, EvalType<T>(value));
         currentIndex++;
       }
     }