Skip to content

Commit 84428f9

Browse files
Huameng (Michael) Jiangfacebook-github-bot
Huameng (Michael) Jiang
authored and
facebook-github-bot
committed
Allow copying vector to a new memory pool (facebookincubator#10647)
Summary: Pull Request resolved: facebookincubator#10647 Until we have a proper solution to transfer vector memory ownership from velox/koski to downstream components, we extend the vector api to currently copy the vector with a downstream memory pool. Reviewed By: Yuhta Differential Revision: D60616369 fbshipit-source-id: 13515ae32ef9beadc2500e2c050f40244d46dc8e
1 parent 5cff4e5 commit 84428f9

10 files changed

+78
-42
lines changed

‎velox/vector/BaseVector.h

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -464,9 +464,11 @@ class BaseVector {
464464
const vector_size_t* toSourceRow);
465465

466466
/// Utility for making a deep copy of a whole vector.
467-
static VectorPtr copy(const BaseVector& vector) {
468-
auto result =
469-
BaseVector::create(vector.type(), vector.size(), vector.pool());
467+
static VectorPtr copy(
468+
const BaseVector& vector,
469+
velox::memory::MemoryPool* pool = nullptr) {
470+
auto result = BaseVector::create(
471+
vector.type(), vector.size(), pool ? pool : vector.pool());
470472
result->copy(&vector, 0, 0, vector.size());
471473
return result;
472474
}
@@ -497,10 +499,11 @@ class BaseVector {
497499
}
498500

499501
/// This makes a deep copy of the Vector allocating new child Vectors and
500-
// Buffers recursively. Unlike copy, this preserves encodings recursively.
501-
virtual VectorPtr copyPreserveEncodings() const = 0;
502+
/// Buffers recursively. Unlike copy, this preserves encodings recursively.
503+
virtual VectorPtr copyPreserveEncodings(
504+
velox::memory::MemoryPool* pool = nullptr) const = 0;
502505

503-
// Construct a zero-copy slice of the vector with the indicated offset and
506+
/// Construct a zero-copy slice of the vector with the indicated offset and
504507
/// length.
505508
virtual VectorPtr slice(vector_size_t offset, vector_size_t length) const = 0;
506509

‎velox/vector/BiasVector.h

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -163,13 +163,15 @@ class BiasVector : public SimpleVector<T> {
163163
VELOX_NYI();
164164
}
165165

166-
VectorPtr copyPreserveEncodings() const override {
166+
VectorPtr copyPreserveEncodings(
167+
velox::memory::MemoryPool* pool = nullptr) const override {
168+
auto selfPool = pool ? pool : BaseVector::pool_;
167169
return std::make_shared<BiasVector<T>>(
168-
BaseVector::pool_,
169-
AlignedBuffer::copy(BaseVector::pool_, BaseVector::nulls_),
170+
selfPool,
171+
AlignedBuffer::copy(selfPool, BaseVector::nulls_),
170172
BaseVector::length_,
171173
valueType_,
172-
AlignedBuffer::copy(BaseVector::pool_, values_),
174+
AlignedBuffer::copy(selfPool, values_),
173175
bias_,
174176
SimpleVector<T>::stats_,
175177
BaseVector::distinctValueCount_,

‎velox/vector/ComplexVector.h

Lines changed: 23 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -168,17 +168,19 @@ class RowVector : public BaseVector {
168168
const BaseVector* source,
169169
const folly::Range<const CopyRange*>& ranges) override;
170170

171-
VectorPtr copyPreserveEncodings() const override {
171+
VectorPtr copyPreserveEncodings(
172+
velox::memory::MemoryPool* pool = nullptr) const override {
172173
std::vector<VectorPtr> copiedChildren(children_.size());
173174

174175
for (auto i = 0; i < children_.size(); ++i) {
175-
copiedChildren[i] = children_[i]->copyPreserveEncodings();
176+
copiedChildren[i] = children_[i]->copyPreserveEncodings(pool);
176177
}
177178

179+
auto selfPool = pool ? pool : pool_;
178180
return std::make_shared<RowVector>(
179-
pool_,
181+
selfPool,
180182
type_,
181-
AlignedBuffer::copy(pool_, nulls_),
183+
AlignedBuffer::copy(selfPool, nulls_),
182184
length_,
183185
copiedChildren,
184186
nullCount_);
@@ -466,15 +468,17 @@ class ArrayVector : public ArrayVectorBase {
466468
const BaseVector* source,
467469
const folly::Range<const CopyRange*>& ranges) override;
468470

469-
VectorPtr copyPreserveEncodings() const override {
471+
VectorPtr copyPreserveEncodings(
472+
velox::memory::MemoryPool* pool = nullptr) const override {
473+
auto selfPool = pool ? pool : pool_;
470474
return std::make_shared<ArrayVector>(
471-
pool_,
475+
selfPool,
472476
type_,
473-
AlignedBuffer::copy(pool_, nulls_),
477+
AlignedBuffer::copy(selfPool, nulls_),
474478
length_,
475-
AlignedBuffer::copy(pool_, offsets_),
476-
AlignedBuffer::copy(pool_, sizes_),
477-
elements_->copyPreserveEncodings(),
479+
AlignedBuffer::copy(selfPool, offsets_),
480+
AlignedBuffer::copy(selfPool, sizes_),
481+
elements_->copyPreserveEncodings(pool),
478482
nullCount_);
479483
}
480484

@@ -607,16 +611,18 @@ class MapVector : public ArrayVectorBase {
607611
const BaseVector* source,
608612
const folly::Range<const CopyRange*>& ranges) override;
609613

610-
VectorPtr copyPreserveEncodings() const override {
614+
VectorPtr copyPreserveEncodings(
615+
velox::memory::MemoryPool* pool = nullptr) const override {
616+
auto selfPool = pool ? pool : pool_;
611617
return std::make_shared<MapVector>(
612-
pool_,
618+
selfPool,
613619
type_,
614-
AlignedBuffer::copy(pool_, nulls_),
620+
AlignedBuffer::copy(selfPool, nulls_),
615621
length_,
616-
AlignedBuffer::copy(pool_, offsets_),
617-
AlignedBuffer::copy(pool_, sizes_),
618-
keys_->copyPreserveEncodings(),
619-
values_->copyPreserveEncodings(),
622+
AlignedBuffer::copy(selfPool, offsets_),
623+
AlignedBuffer::copy(selfPool, sizes_),
624+
keys_->copyPreserveEncodings(pool),
625+
values_->copyPreserveEncodings(pool),
620626
nullCount_,
621627
sortedKeys_);
622628
}

‎velox/vector/ConstantVector.h

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -356,18 +356,20 @@ class ConstantVector final : public SimpleVector<T> {
356356
}
357357
}
358358

359-
VectorPtr copyPreserveEncodings() const override {
359+
VectorPtr copyPreserveEncodings(
360+
velox::memory::MemoryPool* pool = nullptr) const override {
361+
auto selfPool = pool ? pool : BaseVector::pool_;
360362
if (valueVector_) {
361363
return std::make_shared<ConstantVector<T>>(
362-
BaseVector::pool_,
364+
selfPool,
363365
BaseVector::length_,
364366
index_,
365-
valueVector_->copyPreserveEncodings(),
367+
valueVector_->copyPreserveEncodings(pool),
366368
SimpleVector<T>::stats_);
367369
}
368370

369371
return std::make_shared<ConstantVector<T>>(
370-
BaseVector::pool_,
372+
selfPool,
371373
BaseVector::length_,
372374
isNull_,
373375
BaseVector::type_,

‎velox/vector/DictionaryVector.h

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -227,13 +227,15 @@ class DictionaryVector : public SimpleVector<T> {
227227

228228
void validate(const VectorValidateOptions& options) const override;
229229

230-
VectorPtr copyPreserveEncodings() const override {
230+
VectorPtr copyPreserveEncodings(
231+
velox::memory::MemoryPool* pool = nullptr) const override {
232+
auto selfPool = pool ? pool : BaseVector::pool_;
231233
return std::make_shared<DictionaryVector<T>>(
232-
BaseVector::pool_,
233-
AlignedBuffer::copy(BaseVector::pool_, BaseVector::nulls_),
234+
selfPool,
235+
AlignedBuffer::copy(selfPool, BaseVector::nulls_),
234236
BaseVector::length_,
235237
dictionaryValues_->copyPreserveEncodings(),
236-
AlignedBuffer::copy(BaseVector::pool_, indices_),
238+
AlignedBuffer::copy(selfPool, indices_),
237239
SimpleVector<T>::stats_,
238240
BaseVector::distinctValueCount_,
239241
BaseVector::nullCount_,

‎velox/vector/FlatVector.h

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -272,9 +272,10 @@ class FlatVector final : public SimpleVector<T> {
272272
const BaseVector* source,
273273
const folly::Range<const BaseVector::CopyRange*>& ranges) override;
274274

275-
VectorPtr copyPreserveEncodings() const override {
275+
VectorPtr copyPreserveEncodings(
276+
velox::memory::MemoryPool* pool = nullptr) const override {
276277
return std::make_shared<FlatVector<T>>(
277-
BaseVector::pool_,
278+
pool ? pool : BaseVector::pool_,
278279
BaseVector::type_,
279280
AlignedBuffer::copy(BaseVector::pool_, BaseVector::nulls_),
280281
BaseVector::length_,

‎velox/vector/FunctionVector.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,8 @@ class FunctionVector : public BaseVector {
196196
VELOX_NYI();
197197
}
198198

199-
VectorPtr copyPreserveEncodings() const override {
199+
VectorPtr copyPreserveEncodings(
200+
velox::memory::MemoryPool* /* pool */ = nullptr) const override {
200201
VELOX_UNSUPPORTED("copyPreserveEncodings not defined for FunctionVector");
201202
}
202203

‎velox/vector/LazyVector.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,8 @@ class LazyVector : public BaseVector {
259259

260260
void validate(const VectorValidateOptions& options) const override;
261261

262-
VectorPtr copyPreserveEncodings() const override {
262+
VectorPtr copyPreserveEncodings(
263+
velox::memory::MemoryPool* /* pool */ = nullptr) const override {
263264
VELOX_UNSUPPORTED("copyPreserveEncodings not defined for LazyVector");
264265
}
265266

‎velox/vector/SequenceVector.h

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -198,12 +198,14 @@ class SequenceVector : public SimpleVector<T> {
198198
return false;
199199
}
200200

201-
VectorPtr copyPreserveEncodings() const override {
201+
VectorPtr copyPreserveEncodings(
202+
velox::memory::MemoryPool* pool = nullptr) const override {
203+
auto selfPool = pool ? pool : BaseVector::pool_;
202204
return std::make_shared<SequenceVector<T>>(
203-
BaseVector::pool_,
205+
selfPool,
204206
BaseVector::length_,
205207
sequenceValues_->copyPreserveEncodings(),
206-
AlignedBuffer::copy(BaseVector::pool_, sequenceLengths_),
208+
AlignedBuffer::copy(selfPool, sequenceLengths_),
207209
SimpleVector<T>::stats_,
208210
BaseVector::distinctValueCount_,
209211
BaseVector::nullCount_,

‎velox/vector/tests/CopyPreserveEncodingsTest.cpp

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -349,4 +349,20 @@ TEST_F(CopyPreserveEncodingsTest, sequenceNoNulls) {
349349
assertEqualVectors(sequenceVector, copy);
350350
validateCopyPreserveEncodings(sequenceVector, copy);
351351
}
352+
353+
TEST_F(CopyPreserveEncodingsTest, newMemoryPool) {
354+
auto dictionaryVector = vectorMaker_.dictionaryVector(generateIntInput());
355+
356+
auto sourcePool = dictionaryVector->pool();
357+
auto targetPool = memory::memoryManager()->addLeafPool();
358+
auto preCopySrcMemory = sourcePool->usedBytes();
359+
ASSERT_EQ(0, targetPool->usedBytes());
360+
361+
auto copy = dictionaryVector->copyPreserveEncodings(targetPool.get());
362+
assertEqualVectors(dictionaryVector, copy);
363+
validateCopyPreserveEncodings(dictionaryVector, copy);
364+
365+
EXPECT_EQ(preCopySrcMemory, sourcePool->usedBytes());
366+
EXPECT_EQ(preCopySrcMemory, targetPool->usedBytes());
367+
}
352368
} // namespace facebook::velox::test

0 commit comments

Comments
 (0)