Skip to content

Commit 7524325

Browse files
siddharthteotiasiddharthteotia
committed
added tests
1 parent 66620f1 commit 7524325

6 files changed

Lines changed: 410 additions & 7 deletions

File tree

‎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: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,7 @@ public void reAlloc() {
214214
baseSize = (long)currentBufferCapacity;
215215
}
216216
long newAllocationSize = baseSize * 2L;
217-
newAllocationSize = Long.highestOneBit(newAllocationSize - 1) << 1;
217+
newAllocationSize = BaseAllocator.nextPowerOfTwo(newAllocationSize);
218218

219219
if (newAllocationSize > MAX_ALLOCATION_SIZE) {
220220
throw new OversizedAllocationException("Unable to expand the buffer. Max allowed buffer size is reached.");

‎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: 2 additions & 1 deletion
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;
@@ -214,7 +215,7 @@ public void reAlloc() {
214215
baseSize = (long)currentBufferCapacity;
215216
}
216217
long newAllocationSize = baseSize * 2L;
217-
newAllocationSize = Long.highestOneBit(newAllocationSize - 1) << 1;
218+
newAllocationSize = BaseAllocator.nextPowerOfTwo(newAllocationSize);
218219

219220
if (newAllocationSize > MAX_ALLOCATION_SIZE) {
220221
throw new OversizedAllocationException("Requested amount of memory is more than max allowed allocation size");

‎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)