You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by ks...@apache.org on 2020/10/19 18:01:09 UTC

[arrow] 14/30: ARROW-9475: [Java] Clean up usages of BaseAllocator, use BufferAllocator in…

This is an automated email from the ASF dual-hosted git repository.

kszucs pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git

commit cb5814605d170c0c5ccc812216b0149041cf7015
Author: Hongze Zhang <ho...@intel.com>
AuthorDate: Thu Oct 15 22:21:23 2020 -0700

    ARROW-9475: [Java] Clean up usages of BaseAllocator, use BufferAllocator in…
    
    …stead
    
    Issue link: https://issues.apache.org/jira/browse/ARROW-9475.
    
    Closes #7768 from zhztheplayer/ARROW-9475
    
    Authored-by: Hongze Zhang <ho...@intel.com>
    Signed-off-by: Micah Kornfield <em...@gmail.com>
---
 .../java/org/apache/arrow/memory/Accountant.java   |  3 +-
 .../org/apache/arrow/memory/AllocationManager.java | 39 ++++++++++++----------
 .../org/apache/arrow/memory/BaseAllocator.java     | 16 ++++++---
 .../org/apache/arrow/memory/BufferAllocator.java   | 32 ++++++++++++++++++
 .../java/org/apache/arrow/memory/BufferLedger.java | 22 ++++++------
 .../memory/DefaultAllocationManagerFactory.java    |  2 +-
 .../memory/DefaultAllocationManagerFactory.java    |  2 +-
 .../arrow/memory/NettyAllocationManager.java       |  6 ++--
 .../org/apache/arrow/memory/TestBaseAllocator.java |  2 +-
 .../arrow/memory/TestNettyAllocationManager.java   |  2 +-
 .../memory/DefaultAllocationManagerFactory.java    |  2 +-
 .../arrow/memory/UnsafeAllocationManager.java      |  4 +--
 12 files changed, 87 insertions(+), 45 deletions(-)

diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
index da93511..42dac7b 100644
--- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
+++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java
@@ -140,7 +140,7 @@ class Accountant implements AutoCloseable {
    * @param size to increase
    * @return Whether the allocation fit within limits.
    */
-  boolean forceAllocate(long size) {
+  public boolean forceAllocate(long size) {
     final AllocationOutcome.Status outcome = allocate(size, true, true, null);
     return outcome.isOk();
   }
@@ -220,7 +220,6 @@ class Accountant implements AutoCloseable {
       final long actualToReleaseToParent = Math.min(size, possibleAmountToReleaseToParent);
       parent.releaseBytes(actualToReleaseToParent);
     }
-
   }
 
   public boolean isOverLimit() {
diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java
index c61d041..9c7cfa9 100644
--- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java
+++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java
@@ -47,11 +47,11 @@ public abstract class AllocationManager {
 
   private static final AtomicLong MANAGER_ID_GENERATOR = new AtomicLong(0);
 
-  private final RootAllocator root;
+  private final BufferAllocator root;
   private final long allocatorManagerId = MANAGER_ID_GENERATOR.incrementAndGet();
   // ARROW-1627 Trying to minimize memory overhead caused by previously used IdentityHashMap
   // see JIRA for details
-  private final LowCostIdentityHashMap<BaseAllocator, BufferLedger> map = new LowCostIdentityHashMap<>();
+  private final LowCostIdentityHashMap<BufferAllocator, BufferLedger> map = new LowCostIdentityHashMap<>();
   private final long amCreationTime = System.nanoTime();
 
   // The ReferenceManager created at the time of creation of this AllocationManager
@@ -60,11 +60,11 @@ public abstract class AllocationManager {
   private volatile BufferLedger owningLedger;
   private volatile long amDestructionTime = 0;
 
-  protected AllocationManager(BaseAllocator accountingAllocator) {
+  protected AllocationManager(BufferAllocator accountingAllocator) {
     Preconditions.checkNotNull(accountingAllocator);
     accountingAllocator.assertOpen();
 
-    this.root = accountingAllocator.root;
+    this.root = accountingAllocator.getRoot();
 
     // we do a no retain association since our creator will want to retrieve the newly created
     // ledger and will create a reference count at that point
@@ -87,13 +87,13 @@ public abstract class AllocationManager {
    * @return The reference manager (new or existing) that associates the underlying
    *         buffer to this new ledger.
    */
-  BufferLedger associate(final BaseAllocator allocator) {
+  BufferLedger associate(final BufferAllocator allocator) {
     return associate(allocator, true);
   }
 
-  private BufferLedger associate(final BaseAllocator allocator, final boolean retain) {
+  private BufferLedger associate(final BufferAllocator allocator, final boolean retain) {
     allocator.assertOpen();
-    Preconditions.checkState(root == allocator.root,
+    Preconditions.checkState(root == allocator.getRoot(),
           "A buffer can only be associated between two allocators that share the same root");
 
     synchronized (this) {
@@ -118,9 +118,11 @@ public abstract class AllocationManager {
       Preconditions.checkState(oldLedger == null,
           "Detected inconsistent state: A reference manager already exists for this allocator");
 
-      // needed for debugging only: keep a pointer to reference manager inside allocator
-      // to dump state, verify allocator state etc
-      allocator.associateLedger(ledger);
+      if (allocator instanceof BaseAllocator) {
+        // needed for debugging only: keep a pointer to reference manager inside allocator
+        // to dump state, verify allocator state etc
+        ((BaseAllocator) allocator).associateLedger(ledger);
+      }
       return ledger;
     }
   }
@@ -133,7 +135,7 @@ public abstract class AllocationManager {
    * calling ReferenceManager drops to 0.
    */
   void release(final BufferLedger ledger) {
-    final BaseAllocator allocator = (BaseAllocator) ledger.getAllocator();
+    final BufferAllocator allocator = ledger.getAllocator();
     allocator.assertOpen();
 
     // remove the <BaseAllocator, BufferLedger> mapping for the allocator
@@ -142,9 +144,12 @@ public abstract class AllocationManager {
         "Expecting a mapping for allocator and reference manager");
     final BufferLedger oldLedger = map.remove(allocator);
 
-    // needed for debug only: tell the allocator that AllocationManager is removing a
-    // reference manager associated with this particular allocator
-    ((BaseAllocator) oldLedger.getAllocator()).dissociateLedger(oldLedger);
+    BufferAllocator oldAllocator = oldLedger.getAllocator();
+    if (oldAllocator instanceof BaseAllocator) {
+      // needed for debug only: tell the allocator that AllocationManager is removing a
+      // reference manager associated with this particular allocator
+      ((BaseAllocator) oldAllocator).dissociateLedger(oldLedger);
+    }
 
     if (oldLedger == owningLedger) {
       // the release call was made by the owning reference manager
@@ -152,10 +157,10 @@ public abstract class AllocationManager {
         // the only <allocator, reference manager> mapping was for the owner
         // which now has been removed, it implies we can safely destroy the
         // underlying memory chunk as it is no longer being referenced
-        ((BaseAllocator) oldLedger.getAllocator()).releaseBytes(getSize());
+        oldAllocator.releaseBytes(getSize());
         // free the memory chunk associated with the allocation manager
         release0();
-        ((BaseAllocator) oldLedger.getAllocator()).getListener().onRelease(getSize());
+        oldAllocator.getListener().onRelease(getSize());
         amDestructionTime = System.nanoTime();
         owningLedger = null;
       } else {
@@ -209,7 +214,7 @@ public abstract class AllocationManager {
      * @param size Size (in bytes) of memory managed by the AllocationManager
      * @return The created AllocationManager used by this allocator
      */
-    AllocationManager create(BaseAllocator accountingAllocator, long size);
+    AllocationManager create(BufferAllocator accountingAllocator, long size);
 
     ArrowBuf empty();
   }
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 81f6649..246b221 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
@@ -61,10 +61,10 @@ abstract class BaseAllocator extends Accountant implements BufferAllocator {
   public static final Config DEFAULT_CONFIG = ImmutableConfig.builder().build();
 
   // Package exposed for sharing between AllocatorManger and BaseAllocator objects
-  final String name;
-  final RootAllocator root;
+  private final String name;
+  private final RootAllocator root;
   private final Object DEBUG_LOCK = DEBUG ? new Object() : null;
-  final AllocationListener listener;
+  private final AllocationListener listener;
   private final BaseAllocator parentAllocator;
   private final Map<BaseAllocator, Object> childAllocators;
   private final ArrowBuf empty;
@@ -124,7 +124,8 @@ abstract class BaseAllocator extends Accountant implements BufferAllocator {
     this.roundingPolicy = config.getRoundingPolicy();
   }
 
-  AllocationListener getListener() {
+  @Override
+  public AllocationListener getListener() {
     return listener;
   }
 
@@ -315,6 +316,11 @@ abstract class BaseAllocator extends Accountant implements BufferAllocator {
   }
 
   @Override
+  public BufferAllocator getRoot() {
+    return root;
+  }
+
+  @Override
   public BufferAllocator newChildAllocator(
       final String name,
       final long initReservation,
@@ -343,7 +349,7 @@ abstract class BaseAllocator extends Accountant implements BufferAllocator {
       synchronized (DEBUG_LOCK) {
         childAllocators.put(childAllocator, childAllocator);
         historicalLog.recordEvent("allocator[%s] created new child allocator[%s]", name,
-            childAllocator.name);
+            childAllocator.getName());
       }
     } else {
       childAllocators.put(childAllocator, childAllocator);
diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
index aa1f856..8fbf6f7 100644
--- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
+++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java
@@ -50,6 +50,14 @@ public interface BufferAllocator extends AutoCloseable {
   ArrowBuf buffer(long size, BufferManager manager);
 
   /**
+   * Get the root allocator of this allocator. If this allocator is already a root, return
+   * this directly.
+   *
+   * @return The root allocator
+   */
+  BufferAllocator getRoot();
+
+  /**
    * Create a new child allocator.
    *
    * @param name            the name of the allocator.
@@ -127,6 +135,30 @@ public interface BufferAllocator extends AutoCloseable {
   long getHeadroom();
 
   /**
+   * Forcibly allocate bytes. Returns whether the allocation fit within limits.
+   *
+   * @param size to increase
+   * @return Whether the allocation fit within limits.
+   */
+  boolean forceAllocate(long size);
+
+
+  /**
+   * Release bytes from this allocator.
+   *
+   * @param size to release
+   */
+  void releaseBytes(long size);
+
+  /**
+   * Returns the allocation listener used by this allocator.
+   *
+   * @return the {@link AllocationListener} instance. Or {@link AllocationListener#NOOP} by default if no listener
+   *         is configured when this allocator was created.
+   */
+  AllocationListener getListener();
+
+  /**
    * Returns the parent allocator.
    *
    * @return parent allocator
diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
index 9fa4de7..48b3e18 100644
--- a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
+++ b/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java
@@ -31,7 +31,7 @@ import org.apache.arrow.util.Preconditions;
  * ArrowBufs managed by this reference manager share a common
  * fate (same reference count).
  */
-public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, ReferenceManager {
+public class BufferLedger implements ValueWithKeyIncluded<BufferAllocator>, ReferenceManager {
   private final IdentityHashMap<ArrowBuf, Object> buffers =
           BaseAllocator.DEBUG ? new IdentityHashMap<>() : null;
   private static final AtomicLong LEDGER_ID_GENERATOR = new AtomicLong(0);
@@ -41,14 +41,14 @@ public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, Refere
   // manage request for retain
   // correctly
   private final long lCreationTime = System.nanoTime();
-  private final BaseAllocator allocator;
+  private final BufferAllocator allocator;
   private final AllocationManager allocationManager;
   private final HistoricalLog historicalLog =
       BaseAllocator.DEBUG ? new HistoricalLog(BaseAllocator.DEBUG_LOG_LENGTH,
         "BufferLedger[%d]", 1) : null;
   private volatile long lDestructionTime = 0;
 
-  BufferLedger(final BaseAllocator allocator, final AllocationManager allocationManager) {
+  BufferLedger(final BufferAllocator allocator, final AllocationManager allocationManager) {
     this.allocator = allocator;
     this.allocationManager = allocationManager;
   }
@@ -57,7 +57,7 @@ public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, Refere
     return this == allocationManager.getOwningLedger();
   }
 
-  public BaseAllocator getKey() {
+  public BufferAllocator getKey() {
     return allocator;
   }
 
@@ -238,7 +238,7 @@ public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, Refere
               "ArrowBuf(BufferLedger, BufferAllocator[%s], " +
                       "UnsafeDirectLittleEndian[identityHashCode == " +
                       "%d](%s)) => ledger hc == %d",
-              allocator.name, System.identityHashCode(derivedBuf), derivedBuf.toString(),
+              allocator.getName(), System.identityHashCode(derivedBuf), derivedBuf.toString(),
               System.identityHashCode(this));
 
       synchronized (buffers) {
@@ -275,7 +275,7 @@ public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, Refere
       historicalLog.recordEvent(
           "ArrowBuf(BufferLedger, BufferAllocator[%s], " +
           "UnsafeDirectLittleEndian[identityHashCode == " + "%d](%s)) => ledger hc == %d",
-          allocator.name, System.identityHashCode(buf), buf.toString(),
+          allocator.getName(), System.identityHashCode(buf), buf.toString(),
           System.identityHashCode(this));
 
       synchronized (buffers) {
@@ -317,7 +317,7 @@ public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, Refere
     // alternatively, if there was already a mapping for <buffer allocator, ref manager> in
     // allocation manager, the ref count of the new buffer will be targetrefmanager.refcount() + 1
     // and this will be true for all the existing buffers currently managed by targetrefmanager
-    final BufferLedger targetRefManager = allocationManager.associate((BaseAllocator) target);
+    final BufferLedger targetRefManager = allocationManager.associate(target);
     // create a new ArrowBuf to associate with new allocator and target ref manager
     final long targetBufLength = srcBuffer.capacity();
     ArrowBuf targetArrowBuf = targetRefManager.deriveBuffer(srcBuffer, 0, targetBufLength);
@@ -336,8 +336,8 @@ public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, Refere
   boolean transferBalance(final ReferenceManager targetReferenceManager) {
     Preconditions.checkArgument(targetReferenceManager != null,
         "Expecting valid target reference manager");
-    final BaseAllocator targetAllocator = (BaseAllocator) targetReferenceManager.getAllocator();
-    Preconditions.checkArgument(allocator.root == targetAllocator.root,
+    final BufferAllocator targetAllocator = targetReferenceManager.getAllocator();
+    Preconditions.checkArgument(allocator.getRoot() == targetAllocator.getRoot(),
         "You can only transfer between two allocators that share the same root.");
 
     allocator.assertOpen();
@@ -411,7 +411,7 @@ public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, Refere
     // alternatively, if there was already a mapping for <buffer allocator, ref manager> in
     // allocation manager, the ref count of the new buffer will be targetrefmanager.refcount() + 1
     // and this will be true for all the existing buffers currently managed by targetrefmanager
-    final BufferLedger targetRefManager = allocationManager.associate((BaseAllocator) target);
+    final BufferLedger targetRefManager = allocationManager.associate(target);
     // create a new ArrowBuf to associate with new allocator and target ref manager
     final long targetBufLength = srcBuffer.capacity();
     final ArrowBuf targetArrowBuf = targetRefManager.deriveBuffer(srcBuffer, 0, targetBufLength);
@@ -486,7 +486,7 @@ public class BufferLedger implements ValueWithKeyIncluded<BaseAllocator>, Refere
       .append("ledger[")
       .append(ledgerId)
       .append("] allocator: ")
-      .append(allocator.name)
+      .append(allocator.getName())
       .append("), isOwning: ")
       .append(", size: ")
       .append(", references: ")
diff --git a/java/memory/memory-core/src/test/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java b/java/memory/memory-core/src/test/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
index e455310..bfe4965 100644
--- a/java/memory/memory-core/src/test/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
+++ b/java/memory/memory-core/src/test/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
@@ -34,7 +34,7 @@ public class DefaultAllocationManagerFactory implements AllocationManager.Factor
       MemoryUtil.UNSAFE.allocateMemory(0));
 
   @Override
-  public AllocationManager create(BaseAllocator accountingAllocator, long size) {
+  public AllocationManager create(BufferAllocator accountingAllocator, long size) {
     return new AllocationManager(accountingAllocator) {
       private final long allocatedSize = size;
       private final long address = MemoryUtil.UNSAFE.allocateMemory(size);
diff --git a/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java b/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
index 15651a3..10cfb5c 100644
--- a/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
+++ b/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
@@ -26,7 +26,7 @@ public class DefaultAllocationManagerFactory implements AllocationManager.Factor
   public static final AllocationManager.Factory FACTORY = NettyAllocationManager.FACTORY;
 
   @Override
-  public AllocationManager create(BaseAllocator accountingAllocator, long size) {
+  public AllocationManager create(BufferAllocator accountingAllocator, long size) {
     return FACTORY.create(accountingAllocator, size);
   }
 
diff --git a/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/NettyAllocationManager.java b/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/NettyAllocationManager.java
index 45bd5d9..2000477 100644
--- a/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/NettyAllocationManager.java
+++ b/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/NettyAllocationManager.java
@@ -30,7 +30,7 @@ public class NettyAllocationManager extends AllocationManager {
   public static final AllocationManager.Factory FACTORY = new AllocationManager.Factory() {
 
     @Override
-    public AllocationManager create(BaseAllocator accountingAllocator, long size) {
+    public AllocationManager create(BufferAllocator accountingAllocator, long size) {
       return new NettyAllocationManager(accountingAllocator, size);
     }
 
@@ -65,7 +65,7 @@ public class NettyAllocationManager extends AllocationManager {
    */
   private final int allocationCutOffValue;
 
-  NettyAllocationManager(BaseAllocator accountingAllocator, long requestedSize, int allocationCutOffValue) {
+  NettyAllocationManager(BufferAllocator accountingAllocator, long requestedSize, int allocationCutOffValue) {
     super(accountingAllocator);
     this.allocationCutOffValue = allocationCutOffValue;
 
@@ -80,7 +80,7 @@ public class NettyAllocationManager extends AllocationManager {
     }
   }
 
-  NettyAllocationManager(BaseAllocator accountingAllocator, long requestedSize) {
+  NettyAllocationManager(BufferAllocator accountingAllocator, long requestedSize) {
     this(accountingAllocator, requestedSize, DEFAULT_ALLOCATION_CUTOFF_VALUE);
   }
 
diff --git a/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java b/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java
index a42e272..ef49e41 100644
--- a/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java
+++ b/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java
@@ -393,7 +393,7 @@ public class TestBaseAllocator {
         .maxAllocation(MAX_ALLOCATION)
         .allocationManagerFactory(new AllocationManager.Factory() {
           @Override
-          public AllocationManager create(BaseAllocator accountingAllocator, long requestedSize) {
+          public AllocationManager create(BufferAllocator accountingAllocator, long requestedSize) {
             return new AllocationManager(accountingAllocator) {
               private final Unsafe unsafe = getUnsafe();
               private final long address = unsafe.allocateMemory(requestedSize);
diff --git a/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestNettyAllocationManager.java b/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestNettyAllocationManager.java
index f386ea6..1b64cd7 100644
--- a/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestNettyAllocationManager.java
+++ b/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestNettyAllocationManager.java
@@ -35,7 +35,7 @@ public class TestNettyAllocationManager {
     return new RootAllocator(BaseAllocator.configBuilder()
         .allocationManagerFactory(new AllocationManager.Factory() {
           @Override
-          public AllocationManager create(BaseAllocator accountingAllocator, long size) {
+          public AllocationManager create(BufferAllocator accountingAllocator, long size) {
             return new NettyAllocationManager(accountingAllocator, size, CUSTOMIZED_ALLOCATION_CUTOFF_VALUE);
           }
 
diff --git a/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java b/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
index 3963c18..720c3d0 100644
--- a/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
+++ b/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java
@@ -26,7 +26,7 @@ public class DefaultAllocationManagerFactory implements AllocationManager.Factor
   public static final AllocationManager.Factory FACTORY = UnsafeAllocationManager.FACTORY;
 
   @Override
-  public AllocationManager create(BaseAllocator accountingAllocator, long size) {
+  public AllocationManager create(BufferAllocator accountingAllocator, long size) {
     return FACTORY.create(accountingAllocator, size);
   }
 
diff --git a/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java b/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java
index f975653..b10aba3 100644
--- a/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java
+++ b/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java
@@ -32,7 +32,7 @@ public final class UnsafeAllocationManager extends AllocationManager {
 
   public static final AllocationManager.Factory FACTORY = new Factory() {
     @Override
-    public AllocationManager create(BaseAllocator accountingAllocator, long size) {
+    public AllocationManager create(BufferAllocator accountingAllocator, long size) {
       return new UnsafeAllocationManager(accountingAllocator, size);
     }
 
@@ -46,7 +46,7 @@ public final class UnsafeAllocationManager extends AllocationManager {
 
   private final long allocatedAddress;
 
-  UnsafeAllocationManager(BaseAllocator accountingAllocator, long requestedSize) {
+  UnsafeAllocationManager(BufferAllocator accountingAllocator, long requestedSize) {
     super(accountingAllocator);
     allocatedAddress = MemoryUtil.UNSAFE.allocateMemory(requestedSize);
     allocatedSize = requestedSize;