diff --git a/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java b/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java index 7a2531e4ae3..e041e7e02be 100644 --- a/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java +++ b/java/memory/src/main/java/org/apache/arrow/memory/Accountant.java @@ -255,7 +255,9 @@ public long getHeadroom() { return localHeadroom; } - return Math.min(localHeadroom, parent.getHeadroom()); + // Amount of reserved memory left on top of what parent has + long reservedHeadroom = Math.max(0, reservation - locallyHeldMemory.get()); + return Math.min(localHeadroom, parent.getHeadroom() + reservedHeadroom); } } diff --git a/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java b/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java index a45de487625..db5d858e3a7 100644 --- a/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java +++ b/java/memory/src/test/java/org/apache/arrow/memory/TestAccountant.java @@ -80,7 +80,7 @@ private void ensureAccurateReservations(Accountant outsideParent) { final Accountant child = new Accountant(parent, 2, Long.MAX_VALUE); assertEquals(2, parent.getAllocatedMemory()); - + assertEquals(10, child.getHeadroom()); { AllocationOutcome first = child.allocateBytes(1); assertEquals(AllocationOutcome.SUCCESS, first); @@ -139,7 +139,7 @@ private void ensureAccurateReservations(Accountant outsideParent) { child.releaseBytes(9); assertEquals(1, child.getAllocatedMemory()); - assertEquals(8, child.getHeadroom()); + assertEquals(9, child.getHeadroom()); // back to reservation size assertEquals(2, parent.getAllocatedMemory()); @@ -164,7 +164,7 @@ private void ensureAccurateReservations(Accountant outsideParent) { child.releaseBytes(11); assertEquals(child.getAllocatedMemory(), 0); assertEquals(parent.getAllocatedMemory(), 2); - assertEquals(8, child.getHeadroom()); + assertEquals(10, child.getHeadroom()); assertEquals(8, parent.getHeadroom()); child.close();