Fix BaseVector::flattenVector for Lazy input (facebookincubator#9282)

Summary:
The current implementation in `flattenVector` doesn't modify the loaded vector in `LazyVector`. Because the `loadedVector` is not a reference, the flattening doesn't take effect.

https://github.com/facebookincubator/velox/blob/0c86a0aedb90bfa966087def7551a40fc3571a01/velox/vector/BaseVector.cpp#L866-L871

Fixes facebookincubator#8697

Pull Request resolved: facebookincubator#9282

Reviewed By: kgpai

Differential Revision: D55597691

Pulled By: bikramSingh91

fbshipit-source-id: 2d1f5f25bc986479f58c7a1b5555ea57d6b0be3b

Author: marin-ma

velox/vector/BaseVector.cpp

@@ -864,7 +864,7 @@ void BaseVector::flattenVector(VectorPtr& vector) {
       return;
     }
     case VectorEncoding::Simple::LAZY: {
-      auto loadedVector =
+      auto& loadedVector =
           vector->asUnchecked<LazyVector>()->loadedVectorShared();
       BaseVector::flattenVector(loadedVector);
       return;

velox/vector/LazyVector.cpp

@@ -271,4 +271,28 @@ void LazyVector::load(RowSet rows, ValueHook* hook) const {
   }
 }
 
+void LazyVector::loadVectorInternal() const {
+  if (!allLoaded_) {
+    if (!vector_) {
+      vector_ = BaseVector::create(type_, 0, pool_);
+    }
+    SelectivityVector allRows(BaseVector::length_);
+    loader_->load(allRows, nullptr, size(), &vector_);
+    VELOX_CHECK(vector_);
+    if (vector_->encoding() == VectorEncoding::Simple::LAZY) {
+      vector_ = vector_->asUnchecked<LazyVector>()->loadedVectorShared();
+    } else {
+      // If the load produced a wrapper, load the wrapped vector.
+      vector_->loadedVector();
+    }
+    allLoaded_ = true;
+    const_cast<LazyVector*>(this)->BaseVector::nulls_ = vector_->nulls_;
+    if (BaseVector::nulls_) {
+      const_cast<LazyVector*>(this)->BaseVector::rawNulls_ =
+          BaseVector::nulls_->as<uint64_t>();
+    }
+  } else {
+    VELOX_CHECK(vector_);
+  }
+}
 } // namespace facebook::velox

velox/vector/LazyVector.h

@@ -188,29 +188,13 @@ class LazyVector : public BaseVector {
   // Returns a shared_ptr to the vector holding the values. If vector is not
   // loaded, loads all the rows, otherwise returns the loaded vector which can
   // have partially loaded rows.
+  VectorPtr& loadedVectorShared() {
+    loadVectorInternal();
+    return vector_;
+  }
+
   const VectorPtr& loadedVectorShared() const {
-    if (!allLoaded_) {
-      if (!vector_) {
-        vector_ = BaseVector::create(type_, 0, pool_);
-      }
-      SelectivityVector allRows(BaseVector::length_);
-      loader_->load(allRows, nullptr, size(), &vector_);
-      VELOX_CHECK(vector_);
-      if (vector_->encoding() == VectorEncoding::Simple::LAZY) {
-        vector_ = vector_->asUnchecked<LazyVector>()->loadedVectorShared();
-      } else {
-        // If the load produced a wrapper, load the wrapped vector.
-        vector_->loadedVector();
-      }
-      allLoaded_ = true;
-      const_cast<LazyVector*>(this)->BaseVector::nulls_ = vector_->nulls_;
-      if (BaseVector::nulls_) {
-        const_cast<LazyVector*>(this)->BaseVector::rawNulls_ =
-            BaseVector::nulls_->as<uint64_t>();
-      }
-    } else {
-      VELOX_CHECK(vector_);
-    }
+    loadVectorInternal();
     return vector_;
   }
 
@@ -287,6 +271,8 @@ class LazyVector : public BaseVector {
       const SelectivityVector& rows,
       SelectivityVector& baseRows);
 
+  void loadVectorInternal() const;
+
   std::unique_ptr<VectorLoader> loader_;
 
   // True if all values are loaded.

velox/vector/tests/VectorTest.cpp

@@ -2754,6 +2754,12 @@ TEST_F(VectorTest, flattenVector) {
   test(dictionary, false);
   EXPECT_TRUE(dictionary->isFlatEncoding());
 
+  VectorPtr lazyDictionary =
+      wrapInLazyDictionary(makeFlatVector<int32_t>({1, 2, 3}));
+  test(lazyDictionary, true);
+  EXPECT_TRUE(lazyDictionary->isLazy());
+  EXPECT_TRUE(lazyDictionary->loadedVector()->isFlatEncoding());
+
   // Array with constant elements.
   auto* arrayVector = array->as<ArrayVector>();
   arrayVector->elements() = BaseVector::wrapInConstant(100, 1, flat);