diff --git a/java/memory/src/main/java/io/netty/buffer/ArrowBuf.java b/java/memory/src/main/java/io/netty/buffer/ArrowBuf.java index 36890bbdf540d..f430755dc8ca5 100644 --- a/java/memory/src/main/java/io/netty/buffer/ArrowBuf.java +++ b/java/memory/src/main/java/io/netty/buffer/ArrowBuf.java @@ -29,6 +29,7 @@ import org.apache.arrow.memory.BaseAllocator.Verbosity; import org.apache.arrow.memory.BoundsChecking; import org.apache.arrow.memory.BufferLedger; +import org.apache.arrow.memory.BufferManager; import org.apache.arrow.memory.ReferenceManager; import org.apache.arrow.memory.util.HistoricalLog; import org.apache.arrow.util.Preconditions; @@ -71,6 +72,7 @@ public final class ArrowBuf implements AutoCloseable { private static final int LOG_BYTES_PER_ROW = 10; private final long id = idGenerator.incrementAndGet(); private final ReferenceManager referenceManager; + private final BufferManager bufferManager; private final long addr; private final boolean isEmpty; private int readerIndex; @@ -87,10 +89,12 @@ public final class ArrowBuf implements AutoCloseable { */ public ArrowBuf( final ReferenceManager referenceManager, + final BufferManager bufferManager, final int length, final long memoryAddress, boolean isEmpty) { this.referenceManager = referenceManager; + this.bufferManager = bufferManager; this.isEmpty = isEmpty; this.addr = memoryAddress; this.length = length; @@ -1003,18 +1007,6 @@ public long getId() { return id; } - /** Returns all ledger information with stack traces as a string. */ - public String toVerboseString() { - if (isEmpty) { - return toString(); - } - - StringBuilder sb = new StringBuilder(); - // TODO SIDD - //referenceManager.print(sb, 0, Verbosity.LOG_WITH_STACKTRACE); - return sb.toString(); - } - /** * Prints information of this buffer into sb at the given * indentation and verbosity level. @@ -1083,4 +1075,21 @@ public ArrowBuf setZero(int index, int length) { return this; } } + + /** + * Returns this if size is less then {@link #capacity()}, otherwise + * delegates to {@link BufferManager#replace(ArrowBuf, int)} to get a new buffer. + */ + public ArrowBuf reallocIfNeeded(final int size) { + Preconditions.checkArgument(size >= 0, "reallocation size must be non-negative"); + if (this.capacity() >= size) { + return this; + } + if (bufferManager != null) { + return bufferManager.replace(this, size); + } else { + throw new UnsupportedOperationException( + "Realloc is only available in the context of operator's UDFs"); + } + } } diff --git a/java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java b/java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java index bcb022017715a..9c26a2782503d 100644 --- a/java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java +++ b/java/memory/src/main/java/org/apache/arrow/memory/BaseAllocator.java @@ -261,7 +261,7 @@ public ArrowBuf buffer(final int initialRequestSize) { } private ArrowBuf createEmpty() { - return new ArrowBuf(ReferenceManager.NO_OP, 0, 0, true); + return new ArrowBuf(ReferenceManager.NO_OP, null, 0, 0, true); } @Override diff --git a/java/memory/src/main/java/org/apache/arrow/memory/BufferLedger.java b/java/memory/src/main/java/org/apache/arrow/memory/BufferLedger.java index 4f335a8cd249f..0f0e37b38f9c3 100644 --- a/java/memory/src/main/java/org/apache/arrow/memory/BufferLedger.java +++ b/java/memory/src/main/java/org/apache/arrow/memory/BufferLedger.java @@ -36,8 +36,15 @@ * fate (same reference count). */ public class BufferLedger implements ValueWithKeyIncluded, ReferenceManager { - - private final IdentityHashMap buffers = BaseAllocator.DEBUG ? new IdentityHashMap<>() : null; + // since ArrowBuf interface has been cleaned to just include the length + // of memory chunk it has access to and starting virtual address in the chunk, + // ArrowBuf no longer tracks its offset within the chunk and that info + // is kept by the ReferenceManager/BufferLedger here. thus we need this map + // to track offsets of ArrowBufs managed by this ledger. the downside is + // that there could be potential increase in heap overhead as earlier + // map was created in debug mode only but now it will always be created + // per instance creation of BufferLedger + private final IdentityHashMap buffers = new IdentityHashMap<>(); private static final AtomicLong LEDGER_ID_GENERATOR = new AtomicLong(0); // unique ID assigned to each ledger private final long ledgerId = LEDGER_ID_GENERATOR.incrementAndGet(); @@ -236,6 +243,7 @@ public ArrowBuf deriveBuffer(final ArrowBuf sourceBuffer, int index, int length) // create new ArrowBuf derivedBuf = new ArrowBuf( this, + null, length, // length (in bytes) in the underlying memory chunk for this new ArrowBuf startAddress, // starting byte address in the underlying memory for this new ArrowBuf false); @@ -270,15 +278,13 @@ public ArrowBuf deriveBuffer(final ArrowBuf sourceBuffer, int index, int length) * with this BufferLedger */ ArrowBuf newArrowBuf(final int length, final BufferManager manager) { - // TODO: SIDD I believe we can remove BufferManager. I have removed it from ArrowBuf. I don't think its needed - // need to evaluate the use/significance of buffer manager allocator.assertOpen(); // the start virtual address of the ArrowBuf will be same as address of memory chunk final long startAddress = allocationManager.getMemoryChunk().memoryAddress(); // create ArrowBuf - final ArrowBuf buf = new ArrowBuf(this, length, startAddress, false); + final ArrowBuf buf = new ArrowBuf(this, manager, length, startAddress, false); // store the offset (within the underlying memory chunk) of new buffer synchronized (buffers) { @@ -502,16 +508,13 @@ public int getAccountedSize() { * @param verbosity The level of verbosity to print. */ void print(StringBuilder sb, int indent, BaseAllocator.Verbosity verbosity) { - // TODO SIDD: figure out how to make the debug logging work with new changes indent(sb, indent) .append("ledger[") .append(ledgerId) .append("] allocator: ") .append(allocator.name) .append("), isOwning: ") - //.append(owningLedger == this) .append(", size: ") - //.append(size) .append(", references: ") .append(bufRefCnt.get()) .append(", life: ") @@ -519,11 +522,7 @@ void print(StringBuilder sb, int indent, BaseAllocator.Verbosity verbosity) { .append("..") .append(lDestructionTime) .append(", allocatorManager: [") - //.append(AllocationManager.this.allocatorManagerId) .append(", life: "); - //.append(amCreationTime) - //.append("..") - //.append(amDestructionTime); if (!BaseAllocator.DEBUG) { sb.append("]\n"); diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java b/java/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java index d9c8578aaca44..f62371d7525fd 100644 --- a/java/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java +++ b/java/vector/src/test/java/org/apache/arrow/vector/TestBitVectorHelper.java @@ -30,7 +30,7 @@ public class TestBitVectorHelper { public void testGetNullCount() throws Exception { // test case 1, 1 null value for 0b110 ArrowBuf validityBuffer = new ArrowBuf( - ReferenceManager.NO_OP,3, new PooledByteBufAllocatorL().empty.memoryAddress(), true); + ReferenceManager.NO_OP, null,3, new PooledByteBufAllocatorL().empty.memoryAddress(), true); // we set validity buffer to be 0b10110, but only have 3 items with 1st item is null validityBuffer.setByte(0, 0b10110); @@ -40,7 +40,7 @@ public void testGetNullCount() throws Exception { // test case 2, no null value for 0xFF validityBuffer = new ArrowBuf( - ReferenceManager.NO_OP, 8, new PooledByteBufAllocatorL().empty.memoryAddress(), true); + ReferenceManager.NO_OP, null,8, new PooledByteBufAllocatorL().empty.memoryAddress(), true); validityBuffer.setByte(0, 0xFF); count = BitVectorHelper.getNullCount(validityBuffer, 8); @@ -48,7 +48,7 @@ public void testGetNullCount() throws Exception { // test case 3, 1 null value for 0x7F validityBuffer = new ArrowBuf( - ReferenceManager.NO_OP, 8, new PooledByteBufAllocatorL().empty.memoryAddress(), true); + ReferenceManager.NO_OP, null, 8, new PooledByteBufAllocatorL().empty.memoryAddress(), true); validityBuffer.setByte(0, 0x7F); count = BitVectorHelper.getNullCount(validityBuffer, 8); @@ -56,7 +56,7 @@ public void testGetNullCount() throws Exception { // test case 4, validity buffer has multiple bytes, 11 items validityBuffer = new ArrowBuf( - ReferenceManager.NO_OP, 11, new PooledByteBufAllocatorL().empty.memoryAddress(), true); + ReferenceManager.NO_OP, null,11, new PooledByteBufAllocatorL().empty.memoryAddress(), true); validityBuffer.setByte(0, 0b10101010); validityBuffer.setByte(1, 0b01010101);