diff --git a/velox/vector/BaseVector.cpp b/velox/vector/BaseVector.cpp index 5d382aac6d8a6..143728278abdc 100644 --- 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()->loadedVectorShared(); BaseVector::flattenVector(loadedVector); return; diff --git a/velox/vector/LazyVector.cpp b/velox/vector/LazyVector.cpp index c0b12a15e3956..61663d5998856 100644 --- 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()->loadedVectorShared(); + } else { + // If the load produced a wrapper, load the wrapped vector. + vector_->loadedVector(); + } + allLoaded_ = true; + const_cast(this)->BaseVector::nulls_ = vector_->nulls_; + if (BaseVector::nulls_) { + const_cast(this)->BaseVector::rawNulls_ = + BaseVector::nulls_->as(); + } + } else { + VELOX_CHECK(vector_); + } +} } // namespace facebook::velox diff --git a/velox/vector/LazyVector.h b/velox/vector/LazyVector.h index c448f8c833090..198747cc9e03f 100644 --- 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()->loadedVectorShared(); - } else { - // If the load produced a wrapper, load the wrapped vector. - vector_->loadedVector(); - } - allLoaded_ = true; - const_cast(this)->BaseVector::nulls_ = vector_->nulls_; - if (BaseVector::nulls_) { - const_cast(this)->BaseVector::rawNulls_ = - BaseVector::nulls_->as(); - } - } 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 loader_; // True if all values are loaded. diff --git a/velox/vector/tests/VectorTest.cpp b/velox/vector/tests/VectorTest.cpp index 01e2d4a58623d..42930f48fedfc 100644 --- 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({1, 2, 3})); + test(lazyDictionary, true); + EXPECT_TRUE(lazyDictionary->isLazy()); + EXPECT_TRUE(lazyDictionary->loadedVector()->isFlatEncoding()); + // Array with constant elements. auto* arrayVector = array->as(); arrayVector->elements() = BaseVector::wrapInConstant(100, 1, flat);