fix lazy vector doesn't get flattened

marin-ma · Mar 28, 2024 · c2d28e4 · c2d28e4
1 parent 3aa020d
commit c2d28e4
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 23 deletions.
diff --git a/velox/vector/BaseVector.cpp b/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;

diff --git a/velox/vector/LazyVector.cpp b/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
diff --git a/velox/vector/LazyVector.h b/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.

diff --git a/velox/vector/tests/VectorTest.cpp b/velox/vector/tests/VectorTest.cpp
@@ -2754,6 +2754,11 @@ 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);