Adding may fail if the allocator is not allowed to consume any more space. + * + * @param nBytes the number of bytes to add + * @return true if the addition is possible, false otherwise + * @throws IllegalStateException if called after buffer() is used to allocate the reservation + */ + boolean add(long nBytes); + /** * Requests a reservation of additional space. * @@ -44,9 +57,21 @@ public interface AllocationReservation extends AutoCloseable { * * @param nBytes the amount to reserve * @return true if the reservation can be satisfied, false otherwise + * @deprecated use {@link #reserve(long)} instead */ + @Deprecated(forRemoval = true) boolean reserve(int nBytes); + /** + * Requests a reservation of additional space. + * + *
The implementation of the allocator's inner class provides this. + * + * @param nBytes the amount to reserve + * @return true if the reservation can be satisfied, false otherwise + */ + boolean reserve(long nBytes); + /** * Allocate a buffer whose size is the total of all the add()s made. * @@ -65,6 +90,13 @@ public interface AllocationReservation extends AutoCloseable { */ int getSize(); + /** + * Get the current size of the reservation (the sum of all the add()s) as a long value. + * + * @return size of the current reservation + */ + long getSizeLong(); + /** * Return whether or not the reservation has been used. * diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java index dd6375e910b92..20a89d0b7bf18 100644 --- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java +++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java @@ -29,6 +29,7 @@ import org.apache.arrow.memory.util.AssertionUtil; import org.apache.arrow.memory.util.CommonUtil; import org.apache.arrow.memory.util.HistoricalLog; +import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.util.Preconditions; import org.checkerframework.checker.initialization.qual.Initialized; import org.checkerframework.checker.nullness.qual.KeyFor; @@ -860,7 +861,7 @@ RoundingPolicy getRoundingPolicy() { public class Reservation implements AllocationReservation { private final @Nullable HistoricalLog historicalLog; - private int nBytes = 0; + private long nBytes = 0; private boolean used = false; private boolean closed = false; @@ -888,8 +889,15 @@ public Reservation() { } } + @SuppressWarnings({"removal", "InlineMeSuggester"}) + @Deprecated(forRemoval = true) @Override public boolean add(final int nBytes) { + return add((long) nBytes); + } + + @Override + public boolean add(final long nBytes) { assertOpen(); Preconditions.checkArgument(nBytes >= 0, "nBytes(%d) < 0", nBytes); @@ -906,7 +914,7 @@ public boolean add(final int nBytes) { // modifying this behavior so that we maintain what we reserve and what the user asked for // and make sure to only // round to power of two as necessary. - final int nBytesTwo = CommonUtil.nextPowerOfTwo(nBytes); + final long nBytesTwo = CommonUtil.nextPowerOfTwo(nBytes); if (!reserve(nBytesTwo)) { return false; } @@ -929,6 +937,11 @@ public ArrowBuf allocateBuffer() { @Override public int getSize() { + return LargeMemoryUtil.checkedCastToInt(nBytes); + } + + @Override + public long getSizeLong() { return nBytes; } @@ -978,8 +991,15 @@ public void close() { closed = true; } + @SuppressWarnings({"removal", "InlineMeSuggester"}) + @Deprecated(forRemoval = true) @Override public boolean reserve(int nBytes) { + return reserve((long) nBytes); + } + + @Override + public boolean reserve(long nBytes) { assertOpen(); final AllocationOutcome outcome = BaseAllocator.this.allocateBytes(nBytes); @@ -999,7 +1019,7 @@ public boolean reserve(int nBytes) { * @param nBytes the size of the buffer requested * @return the buffer, or null, if the request cannot be satisfied */ - private ArrowBuf allocate(int nBytes) { + private ArrowBuf allocate(long nBytes) { assertOpen(); boolean success = false; @@ -1033,7 +1053,7 @@ private ArrowBuf allocate(int nBytes) { * * @param nBytes the size of the reservation */ - private void releaseReservation(int nBytes) { + private void releaseReservation(long nBytes) { assertOpen(); releaseBytes(nBytes); diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ChildAllocator.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ChildAllocator.java index f8dd7e1d1cb38..50f33d3f021c7 100644 --- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ChildAllocator.java +++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ChildAllocator.java @@ -17,9 +17,9 @@ package org.apache.arrow.memory; /** - * Child allocator class. Only slightly different from the {@see RootAllocator}, in that these can't - * be created directly, but must be obtained from {@link BufferAllocator#newChildAllocator(String, - * AllocationListener, long, long)}. + * Child allocator class. Only slightly different from the {@link RootAllocator}, in that these + * can't be created directly, but must be obtained from {@link + * BufferAllocator#newChildAllocator(String, AllocationListener, long, long)}. * *
Child allocators can only be created by the root, or other children, so this class is package * private. diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/DefaultRoundingPolicy.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/DefaultRoundingPolicy.java index 289b10634d84e..90e8a1d5eca77 100644 --- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/DefaultRoundingPolicy.java +++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/DefaultRoundingPolicy.java @@ -34,13 +34,13 @@ public class DefaultRoundingPolicy implements RoundingPolicy { * *
It was copied from {@link io.netty.buffer.PooledByteBufAllocator}. */ - private static final int MIN_PAGE_SIZE = 4096; + private static final long MIN_PAGE_SIZE = 4096; - private static final int MAX_CHUNK_SIZE = (int) (((long) Integer.MAX_VALUE + 1) / 2); + private static final long MAX_CHUNK_SIZE = ((long) Integer.MAX_VALUE + 1) / 2; private static final long DEFAULT_CHUNK_SIZE; static { - int defaultPageSize = Integer.getInteger("org.apache.memory.allocator.pageSize", 8192); + long defaultPageSize = Long.getLong("org.apache.memory.allocator.pageSize", 8192); try { validateAndCalculatePageShifts(defaultPageSize); } catch (Throwable t) { @@ -60,7 +60,7 @@ public class DefaultRoundingPolicy implements RoundingPolicy { } } - private static int validateAndCalculatePageShifts(int pageSize) { + private static long validateAndCalculatePageShifts(long pageSize) { if (pageSize < MIN_PAGE_SIZE) { throw new IllegalArgumentException( "pageSize: " + pageSize + " (expected: " + MIN_PAGE_SIZE + ")"); @@ -71,17 +71,17 @@ private static int validateAndCalculatePageShifts(int pageSize) { } // Logarithm base 2. At this point we know that pageSize is a power of two. - return Integer.SIZE - 1 - Integer.numberOfLeadingZeros(pageSize); + return Long.SIZE - 1L - Long.numberOfLeadingZeros(pageSize); } - private static int validateAndCalculateChunkSize(int pageSize, int maxOrder) { + private static long validateAndCalculateChunkSize(long pageSize, int maxOrder) { if (maxOrder > 14) { throw new IllegalArgumentException("maxOrder: " + maxOrder + " (expected: 0-14)"); } // Ensure the resulting chunkSize does not overflow. - int chunkSize = pageSize; - for (int i = maxOrder; i > 0; i--) { + long chunkSize = pageSize; + for (long i = maxOrder; i > 0; i--) { if (chunkSize > MAX_CHUNK_SIZE / 2) { throw new IllegalArgumentException( String.format( diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/SegmentRoundingPolicy.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/SegmentRoundingPolicy.java index f501cfedd168d..89db736e6a0f9 100644 --- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/SegmentRoundingPolicy.java +++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/SegmentRoundingPolicy.java @@ -16,6 +16,8 @@ */ package org.apache.arrow.memory.rounding; +import com.google.errorprone.annotations.InlineMe; +import org.apache.arrow.memory.util.LargeMemoryUtil; import org.apache.arrow.util.Preconditions; /** The rounding policy that each buffer size must a multiple of the segment size. */ @@ -28,7 +30,7 @@ public class SegmentRoundingPolicy implements RoundingPolicy { * The segment size. It must be at least {@link SegmentRoundingPolicy#MIN_SEGMENT_SIZE}, and be a * power of 2. */ - private int segmentSize; + private long segmentSize; /** * Constructor for the segment rounding policy. @@ -36,8 +38,22 @@ public class SegmentRoundingPolicy implements RoundingPolicy { * @param segmentSize the segment size. * @throws IllegalArgumentException if the segment size is smaller than {@link * SegmentRoundingPolicy#MIN_SEGMENT_SIZE}, or is not a power of 2. + * @deprecated use {@link SegmentRoundingPolicy#SegmentRoundingPolicy(long)} instead. */ + @Deprecated(forRemoval = true) + @InlineMe(replacement = "this((long) segmentSize)") public SegmentRoundingPolicy(int segmentSize) { + this((long) segmentSize); + } + + /** + * Constructor for the segment rounding policy. + * + * @param segmentSize the segment size. + * @throws IllegalArgumentException if the segment size is smaller than {@link + * SegmentRoundingPolicy#MIN_SEGMENT_SIZE}, or is not a power of 2. + */ + public SegmentRoundingPolicy(long segmentSize) { Preconditions.checkArgument( segmentSize >= MIN_SEGMENT_SIZE, "The segment size cannot be smaller than %s", @@ -52,7 +68,12 @@ public long getRoundedSize(long requestSize) { return (requestSize + (segmentSize - 1)) / segmentSize * segmentSize; } + @Deprecated(forRemoval = true) public int getSegmentSize() { + return LargeMemoryUtil.checkedCastToInt(segmentSize); + } + + public long getSegmentSizeAsLong() { return segmentSize; } } diff --git a/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java b/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java index a5fbc67c48f5c..87e9316964dfc 100644 --- a/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java +++ b/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java @@ -315,7 +315,7 @@ public void testRootAllocator_createChildDontClose() throws Exception { @Test public void testSegmentAllocator() { - RoundingPolicy policy = new SegmentRoundingPolicy(1024); + RoundingPolicy policy = new SegmentRoundingPolicy(1024L); try (RootAllocator allocator = new RootAllocator(AllocationListener.NOOP, 1024 * 1024, policy)) { ArrowBuf buf = allocator.buffer(798); @@ -334,7 +334,7 @@ public void testSegmentAllocator() { @Test public void testSegmentAllocator_childAllocator() { - RoundingPolicy policy = new SegmentRoundingPolicy(1024); + RoundingPolicy policy = new SegmentRoundingPolicy(1024L); try (RootAllocator allocator = new RootAllocator(AllocationListener.NOOP, 1024 * 1024, policy); BufferAllocator childAllocator = allocator.newChildAllocator("child", 0, 512 * 1024)) { @@ -357,14 +357,14 @@ public void testSegmentAllocator_childAllocator() { @Test public void testSegmentAllocator_smallSegment() { IllegalArgumentException e = - assertThrows(IllegalArgumentException.class, () -> new SegmentRoundingPolicy(128)); + assertThrows(IllegalArgumentException.class, () -> new SegmentRoundingPolicy(128L)); assertEquals("The segment size cannot be smaller than 1024", e.getMessage()); } @Test public void testSegmentAllocator_segmentSizeNotPowerOf2() { IllegalArgumentException e = - assertThrows(IllegalArgumentException.class, () -> new SegmentRoundingPolicy(4097)); + assertThrows(IllegalArgumentException.class, () -> new SegmentRoundingPolicy(4097L)); assertEquals("The segment size must be a power of 2", e.getMessage()); } @@ -957,7 +957,7 @@ public void testAllocator_unclaimedReservation() throws Exception { try (final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("unclaimedReservation", 0, MAX_ALLOCATION)) { try (final AllocationReservation reservation = childAllocator1.newReservation()) { - assertTrue(reservation.add(64)); + assertTrue(reservation.add(64L)); } rootAllocator.verify(); } @@ -972,8 +972,8 @@ public void testAllocator_claimedReservation() throws Exception { rootAllocator.newChildAllocator("claimedReservation", 0, MAX_ALLOCATION)) { try (final AllocationReservation reservation = childAllocator1.newReservation()) { - assertTrue(reservation.add(32)); - assertTrue(reservation.add(32)); + assertTrue(reservation.add(32L)); + assertTrue(reservation.add(32L)); final ArrowBuf arrowBuf = reservation.allocateBuffer(); assertEquals(64, arrowBuf.capacity()); diff --git a/java/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/NettyArrowBuf.java b/java/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/NettyArrowBuf.java index bdad3700cb311..9319d15aaa9a9 100644 --- a/java/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/NettyArrowBuf.java +++ b/java/memory/memory-netty-buffer-patch/src/main/java/io/netty/buffer/NettyArrowBuf.java @@ -38,7 +38,7 @@ public class NettyArrowBuf extends AbstractByteBuf implements AutoCloseable { private final ArrowBuf arrowBuf; private final ArrowByteBufAllocator arrowByteBufAllocator; - private int length; + private long length; private final long address; /** @@ -47,10 +47,24 @@ public class NettyArrowBuf extends AbstractByteBuf implements AutoCloseable { * @param arrowBuf The buffer to wrap. * @param bufferAllocator The allocator for the buffer. * @param length The length of this buffer. + * @deprecated Use {@link #NettyArrowBuf(ArrowBuf, BufferAllocator, long)} instead. */ + @Deprecated(forRemoval = true) public NettyArrowBuf( final ArrowBuf arrowBuf, final BufferAllocator bufferAllocator, final int length) { - super(length); + this(arrowBuf, bufferAllocator, (long) length); + } + + /** + * Constructs a new instance. + * + * @param arrowBuf The buffer to wrap. + * @param bufferAllocator The allocator for the buffer. + * @param length The length of this buffer. + */ + public NettyArrowBuf( + final ArrowBuf arrowBuf, final BufferAllocator bufferAllocator, final long length) { + super((int) length); this.arrowBuf = arrowBuf; this.arrowByteBufAllocator = new ArrowByteBufAllocator(bufferAllocator); this.length = length; diff --git a/java/performance/src/main/java/org/apache/arrow/memory/AllocatorBenchmarks.java b/java/performance/src/main/java/org/apache/arrow/memory/AllocatorBenchmarks.java index f275090aae6bf..1154809cae753 100644 --- a/java/performance/src/main/java/org/apache/arrow/memory/AllocatorBenchmarks.java +++ b/java/performance/src/main/java/org/apache/arrow/memory/AllocatorBenchmarks.java @@ -57,9 +57,9 @@ public void defaultAllocatorBenchmark() { @BenchmarkMode(Mode.AverageTime) @OutputTimeUnit(TimeUnit.MICROSECONDS) public void segmentRoundingPolicyBenchmark() { - final int bufferSize = 1024; + final long bufferSize = 1024L; final int numBuffers = 1024; - final int segmentSize = 1024; + final long segmentSize = 1024L; RoundingPolicy policy = new SegmentRoundingPolicy(segmentSize); try (RootAllocator allocator =