Skip to content

Commit 2706b7f

Browse files
siddharthteotiawesm
siddharthteotia
authored and
wesm
committed
ARROW-1533: [JAVA] realloc should consider the existing buffer capacity for computing target memory requirement
cc @jacques-n , This is same as #1097 The latter one was closed as I had to rename the branch correctly and use the correct JIRA number. Author: siddharth <siddharth@dremio.com> Closes #1112 from siddharthteotia/ARROW-1533 and squashes the following commits: 4c97be4 [siddharth] ARROW-1533: realloc should consider the existing buffer capacity for computing target memory requirement
1 parent aebc412 commit 2706b7f

File tree

6 files changed

+511
-9
lines changed

6 files changed

+511
-9
lines changed

‎java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ private static String createErrorMsg(final BufferAllocator allocator, final int
133133
* @param val An integer value.
134134
* @return The closest power of two of that value.
135135
*/
136-
static int nextPowerOfTwo(int val) {
136+
public static int nextPowerOfTwo(int val) {
137137
int highestBit = Integer.highestOneBit(val);
138138
if (highestBit == val) {
139139
return val;
@@ -142,6 +142,21 @@ static int nextPowerOfTwo(int val) {
142142
}
143143
}
144144

145+
/**
146+
* Rounds up the provided value to the nearest power of two.
147+
*
148+
* @param val A long value.
149+
* @return The closest power of two of that value.
150+
*/
151+
public static long nextPowerOfTwo(long val) {
152+
long highestBit = Long.highestOneBit(val);
153+
if (highestBit == val) {
154+
return val;
155+
} else {
156+
return highestBit << 1;
157+
}
158+
}
159+
145160
public static StringBuilder indent(StringBuilder sb, int indent) {
146161
final char[] indentation = new char[indent * 2];
147162
Arrays.fill(indentation, ' ');

‎java/vector/src/main/codegen/templates/FixedValueVectors.java‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -208,14 +208,21 @@ private void allocateBytes(final long size) {
208208
* @throws org.apache.arrow.memory.OutOfMemoryException if it can't allocate the new buffer
209209
*/
210210
public void reAlloc() {
211-
final long newAllocationSize = allocationSizeInBytes * 2L;
212-
if (newAllocationSize > MAX_ALLOCATION_SIZE) {
211+
long baseSize = allocationSizeInBytes;
212+
final int currentBufferCapacity = data.capacity();
213+
if (baseSize < (long)currentBufferCapacity) {
214+
baseSize = (long)currentBufferCapacity;
215+
}
216+
long newAllocationSize = baseSize * 2L;
217+
newAllocationSize = BaseAllocator.nextPowerOfTwo(newAllocationSize);
218+
219+
if (newAllocationSize > MAX_ALLOCATION_SIZE) {
213220
throw new OversizedAllocationException("Unable to expand the buffer. Max allowed buffer size is reached.");
214221
}
215222

216223
logger.debug("Reallocating vector [{}]. # of bytes: [{}] -> [{}]", name, allocationSizeInBytes, newAllocationSize);
217224
final ArrowBuf newBuf = allocator.buffer((int)newAllocationSize);
218-
newBuf.setBytes(0, data, 0, data.capacity());
225+
newBuf.setBytes(0, data, 0, currentBufferCapacity);
219226
final int halfNewCapacity = newBuf.capacity() / 2;
220227
newBuf.setZero(halfNewCapacity, halfNewCapacity);
221228
newBuf.writerIndex(data.writerIndex());

‎java/vector/src/main/codegen/templates/VariableLengthVectors.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -370,13 +370,20 @@ public void reset() {
370370
}
371371

372372
public void reAlloc() {
373-
final long newAllocationSize = allocationSizeInBytes*2L;
373+
long baseSize = allocationSizeInBytes;
374+
final int currentBufferCapacity = data.capacity();
375+
if (baseSize < (long)currentBufferCapacity) {
376+
baseSize = (long)currentBufferCapacity;
377+
}
378+
long newAllocationSize = baseSize * 2L;
379+
newAllocationSize = BaseAllocator.nextPowerOfTwo(newAllocationSize);
380+
374381
if (newAllocationSize > MAX_ALLOCATION_SIZE) {
375382
throw new OversizedAllocationException("Unable to expand the buffer. Max allowed buffer size is reached.");
376383
}
377384

378385
final ArrowBuf newBuf = allocator.buffer((int)newAllocationSize);
379-
newBuf.setBytes(0, data, 0, data.capacity());
386+
newBuf.setBytes(0, data, 0, currentBufferCapacity);
380387
data.release();
381388
data = newBuf;
382389
allocationSizeInBytes = (int)newAllocationSize;

‎java/vector/src/main/java/org/apache/arrow/vector/BitVector.java‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
package org.apache.arrow.vector;
2020

2121
import org.apache.arrow.memory.BufferAllocator;
22+
import org.apache.arrow.memory.BaseAllocator;
2223
import org.apache.arrow.memory.OutOfMemoryException;
2324
import org.apache.arrow.vector.complex.reader.FieldReader;
2425
import org.apache.arrow.vector.holders.BitHolder;
@@ -208,15 +209,22 @@ private void allocateBytes(final long size) {
208209
* Allocate new buffer with double capacity, and copy data into the new buffer. Replace vector's buffer with new buffer, and release old one
209210
*/
210211
public void reAlloc() {
211-
final long newAllocationSize = allocationSizeInBytes * 2L;
212+
long baseSize = allocationSizeInBytes;
213+
final int currentBufferCapacity = data.capacity();
214+
if (baseSize < (long)currentBufferCapacity) {
215+
baseSize = (long)currentBufferCapacity;
216+
}
217+
long newAllocationSize = baseSize * 2L;
218+
newAllocationSize = BaseAllocator.nextPowerOfTwo(newAllocationSize);
219+
212220
if (newAllocationSize > MAX_ALLOCATION_SIZE) {
213221
throw new OversizedAllocationException("Requested amount of memory is more than max allowed allocation size");
214222
}
215223

216224
final int curSize = (int) newAllocationSize;
217225
final ArrowBuf newBuf = allocator.buffer(curSize);
218226
newBuf.setZero(0, newBuf.capacity());
219-
newBuf.setBytes(0, data, 0, data.capacity());
227+
newBuf.setBytes(0, data, 0, currentBufferCapacity);
220228
data.release();
221229
data = newBuf;
222230
allocationSizeInBytes = curSize;

‎java/vector/src/test/java/org/apache/arrow/vector/TestBitVector.java‎

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static org.junit.Assert.assertEquals;
2222
import static org.junit.Assert.assertNull;
2323
import static org.junit.Assert.assertTrue;
24+
import static org.junit.Assert.assertFalse;
2425

2526
import org.apache.arrow.memory.BufferAllocator;
2627
import org.apache.arrow.memory.RootAllocator;
@@ -234,6 +235,198 @@ public void testSplitAndTransfer2() throws Exception {
234235
}
235236
}
236237

238+
@Test
239+
public void testReallocAfterVectorTransfer1() {
240+
try (final BitVector vector = new BitVector(EMPTY_SCHEMA_PATH, allocator)) {
241+
vector.allocateNew(4096);
242+
int valueCapacity = vector.getValueCapacity();
243+
assertEquals(4096, valueCapacity);
244+
245+
final BitVector.Mutator mutator = vector.getMutator();
246+
final BitVector.Accessor accessor = vector.getAccessor();
247+
248+
for (int i = 0; i < valueCapacity; i++) {
249+
if ((i & 1) == 1) {
250+
mutator.setToOne(i);
251+
}
252+
}
253+
254+
for (int i = 0; i < valueCapacity; i++) {
255+
int val = accessor.get(i);
256+
if ((i & 1) == 1) {
257+
assertEquals("unexpected cleared bit at index: " + i, 1, val);
258+
}
259+
else {
260+
assertEquals("unexpected set bit at index: " + i, 0, val);
261+
}
262+
}
263+
264+
/* trigger first realloc */
265+
mutator.setSafeToOne(valueCapacity);
266+
assertEquals(valueCapacity * 2, vector.getValueCapacity());
267+
268+
for (int i = valueCapacity; i < valueCapacity*2; i++) {
269+
if ((i & 1) == 1) {
270+
mutator.setToOne(i);
271+
}
272+
}
273+
274+
for (int i = 0; i < valueCapacity*2; i++) {
275+
int val = accessor.get(i);
276+
if (((i & 1) == 1) || (i == valueCapacity)) {
277+
assertEquals("unexpected cleared bit at index: " + i, 1, val);
278+
}
279+
else {
280+
assertEquals("unexpected set bit at index: " + i, 0, val);
281+
}
282+
}
283+
284+
/* trigger second realloc */
285+
mutator.setSafeToOne(valueCapacity*2);
286+
assertEquals(valueCapacity * 4, vector.getValueCapacity());
287+
288+
for (int i = valueCapacity*2; i < valueCapacity*4; i++) {
289+
if ((i & 1) == 1) {
290+
mutator.setToOne(i);
291+
}
292+
}
293+
294+
for (int i = 0; i < valueCapacity*4; i++) {
295+
int val = accessor.get(i);
296+
if (((i & 1) == 1) || (i == valueCapacity) || (i == valueCapacity*2)) {
297+
assertEquals("unexpected cleared bit at index: " + i, 1, val);
298+
}
299+
else {
300+
assertEquals("unexpected set bit at index: " + i, 0, val);
301+
}
302+
}
303+
304+
/* now transfer the vector */
305+
TransferPair transferPair = vector.getTransferPair(allocator);
306+
transferPair.transfer();
307+
final BitVector toVector = (BitVector)transferPair.getTo();
308+
final BitVector.Accessor toAccessor = toVector.getAccessor();
309+
final BitVector.Mutator toMutator = toVector.getMutator();
310+
311+
assertEquals(valueCapacity * 4, toVector.getValueCapacity());
312+
313+
/* realloc the toVector */
314+
toMutator.setSafeToOne(valueCapacity * 4);
315+
316+
for (int i = 0; i < toVector.getValueCapacity(); i++) {
317+
int val = toAccessor.get(i);
318+
if (i <= valueCapacity * 4) {
319+
if (((i & 1) == 1) || (i == valueCapacity) ||
320+
(i == valueCapacity*2) || (i == valueCapacity*4)) {
321+
assertEquals("unexpected cleared bit at index: " + i, 1, val);
322+
}
323+
else {
324+
assertEquals("unexpected set bit at index: " + i, 0, val);
325+
}
326+
}
327+
else {
328+
assertEquals("unexpected set bit at index: " + i, 0, val);
329+
}
330+
}
331+
332+
toVector.close();
333+
}
334+
}
335+
336+
@Test
337+
public void testReallocAfterVectorTransfer2() {
338+
try (final NullableBitVector vector = new NullableBitVector(EMPTY_SCHEMA_PATH, allocator)) {
339+
vector.allocateNew(4096);
340+
int valueCapacity = vector.getValueCapacity();
341+
assertEquals(4096, valueCapacity);
342+
343+
final NullableBitVector.Mutator mutator = vector.getMutator();
344+
final NullableBitVector.Accessor accessor = vector.getAccessor();
345+
346+
for (int i = 0; i < valueCapacity; i++) {
347+
if ((i & 1) == 1) {
348+
mutator.set(i, 1);
349+
}
350+
}
351+
352+
for (int i = 0; i < valueCapacity; i++) {
353+
if ((i & 1) == 1) {
354+
assertFalse("unexpected cleared bit at index: " + i, accessor.isNull(i));
355+
}
356+
else {
357+
assertTrue("unexpected set bit at index: " + i, accessor.isNull(i));
358+
}
359+
}
360+
361+
/* trigger first realloc */
362+
mutator.setSafe(valueCapacity, 1, 1);
363+
assertEquals(valueCapacity * 2, vector.getValueCapacity());
364+
365+
for (int i = valueCapacity; i < valueCapacity*2; i++) {
366+
if ((i & 1) == 1) {
367+
mutator.set(i, 1);
368+
}
369+
}
370+
371+
for (int i = 0; i < valueCapacity*2; i++) {
372+
if (((i & 1) == 1) || (i == valueCapacity)) {
373+
assertFalse("unexpected cleared bit at index: " + i, accessor.isNull(i));
374+
}
375+
else {
376+
assertTrue("unexpected set bit at index: " + i, accessor.isNull(i));
377+
}
378+
}
379+
380+
/* trigger second realloc */
381+
mutator.setSafe(valueCapacity*2, 1, 1);
382+
assertEquals(valueCapacity * 4, vector.getValueCapacity());
383+
384+
for (int i = valueCapacity*2; i < valueCapacity*4; i++) {
385+
if ((i & 1) == 1) {
386+
mutator.set(i, 1);
387+
}
388+
}
389+
390+
for (int i = 0; i < valueCapacity*4; i++) {
391+
if (((i & 1) == 1) || (i == valueCapacity) || (i == valueCapacity*2)) {
392+
assertFalse("unexpected cleared bit at index: " + i, accessor.isNull(i));
393+
}
394+
else {
395+
assertTrue("unexpected set bit at index: " + i, accessor.isNull(i));
396+
}
397+
}
398+
399+
/* now transfer the vector */
400+
TransferPair transferPair = vector.getTransferPair(allocator);
401+
transferPair.transfer();
402+
final NullableBitVector toVector = (NullableBitVector)transferPair.getTo();
403+
final NullableBitVector.Accessor toAccessor = toVector.getAccessor();
404+
final NullableBitVector.Mutator toMutator = toVector.getMutator();
405+
406+
assertEquals(valueCapacity * 4, toVector.getValueCapacity());
407+
408+
/* realloc the toVector */
409+
toMutator.setSafe(valueCapacity * 4, 1, 1);
410+
411+
for (int i = 0; i < toVector.getValueCapacity(); i++) {
412+
if (i <= valueCapacity * 4) {
413+
if (((i & 1) == 1) || (i == valueCapacity) ||
414+
(i == valueCapacity*2) || (i == valueCapacity*4)) {
415+
assertFalse("unexpected cleared bit at index: " + i, toAccessor.isNull(i));
416+
}
417+
else {
418+
assertTrue("unexpected set bit at index: " + i, toAccessor.isNull(i));
419+
}
420+
}
421+
else {
422+
assertTrue("unexpected set bit at index: " + i, toAccessor.isNull(i));
423+
}
424+
}
425+
426+
toVector.close();
427+
}
428+
}
429+
237430
@Test
238431
public void testBitVector() {
239432
// Create a new value vector for 1024 integers

0 commit comments

Comments
 (0)