diff --git a/java/memory/src/main/java/org/apache/arrow/memory/AllocationManager.java b/java/memory/src/main/java/org/apache/arrow/memory/AllocationManager.java index 43ee9c108d9..491196c426d 100644 --- a/java/memory/src/main/java/org/apache/arrow/memory/AllocationManager.java +++ b/java/memory/src/main/java/org/apache/arrow/memory/AllocationManager.java @@ -332,7 +332,7 @@ public int decrement(int decrement) { * @return The ledger associated with the BufferAllocator. */ public BufferLedger getLedgerForAllocator(BufferAllocator allocator) { - return associate((BaseAllocator) allocator); + return associate(allocator.unwrap(BaseAllocator.class)); } /** 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 f1503c902d0..3c149ce73f7 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 @@ -783,6 +783,16 @@ public static enum Verbosity { } } + @SuppressWarnings("unchecked") + @Override + public T unwrap(Class c) { + if (BaseAllocator.class.isAssignableFrom(c)) { + return (T) this; + } + + throw new UnsupportedOperationException("Unable to unwrap type to class: " + c.getName()); + } + public static boolean isDebug() { return DEBUG; } diff --git a/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java b/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java index 356a3416cbf..d46d680e83c 100644 --- a/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java +++ b/java/memory/src/main/java/org/apache/arrow/memory/BufferAllocator.java @@ -79,6 +79,14 @@ public interface BufferAllocator extends AutoCloseable { @Override public void close(); + /** + * Unwrap the class so that exposes the provided interface, if possible. Otherwise, throw Exception. + * @param c + * The class or interface that you want this class to implement/extend. + * @return The instance of that class related to 'this' + */ + public T unwrap(Class c); + /** * Returns the amount of memory currently allocated from this allocator. *