diff --git a/velox/dwio/common/BufferUtil.h b/velox/dwio/common/BufferUtil.h index 738d81efa167..9bedcbc7b1c7 100644 --- a/velox/dwio/common/BufferUtil.h +++ b/velox/dwio/common/BufferUtil.h @@ -23,18 +23,18 @@ namespace facebook::velox::dwio::common { template inline void ensureCapacity( BufferPtr& data, - size_t capacity, + size_t numElements, velox::memory::MemoryPool* pool, bool preserveOldData = false, bool clearBits = false) { size_t oldSize = 0; + size_t newCapacity = BaseVector::byteSize(numElements); if (!data) { - data = AlignedBuffer::allocate(capacity, pool); + data = AlignedBuffer::allocate(numElements, pool); } else { oldSize = data->size(); - if (!data->isMutable() || - data->capacity() < BaseVector::byteSize(capacity)) { - auto newData = AlignedBuffer::allocate(capacity, pool); + if (!data->isMutable() || data->capacity() < newCapacity) { + auto newData = AlignedBuffer::allocate(numElements, pool); if (preserveOldData) { std::memcpy( newData->template asMutable(), @@ -45,9 +45,11 @@ inline void ensureCapacity( } } - if (clearBits) { + if (clearBits && newCapacity > oldSize) { std::memset( - (void*)(data->asMutable() + oldSize), 0L, capacity - oldSize); + (void*)(data->asMutable() + oldSize), + 0L, + newCapacity - oldSize); } }