PECOBLR-1121 Arrow patch to circumvent Arrow issues with JDk 16+. #1156

tejassp-db · 2025-12-22T10:50:38Z

Description

Databricks server shares query results in Arrow format for easy cross language functionality. The JDBC driver experiences compatibility issues with JDK 16 and later versions when processing Arrow results.

This problem arises from stricter encapsulation of internal APIs in newer Java versions, which affects the driver's use of the Apache Arrow result format consumption with the Apache Arrow library. The JDBC driver is used in partner solutions, where they do not have control of the runtime environment, and the workaround of setting JVM arguments is not feasible.

Testing

Tests are added in other stacked PRs.

Additional Notes to the Reviewer

Its a stacked PR.

Patch Arrow to create a Databricks ArrowBuf which allocates memory on the heap and provides access to it through Java methods. This removes the need to specify "--add-opens=java.base/java.nio=ALL-UNNAMED" as JVM args for JDK 16+.

Use native Arrow if available. Otherwise fallback to the patch version.

Remove irrelevant reference counting in patch code. Patch code uses heap memory for arrow operations and reference counting is not required.

Remove redundant todos for accounting.

Patch DecimalUtility to not use unsafe methods to set decimal values on DatabricksArrowBuf.

Add notice to all patched Arrow Java code. In NOTICE file mention Arrow has been patched by Databricks.

On static init failure of MemoryUtil class, it prints a stack trace to stderr. Remove this print, since now we fallback to DatabricksBufferAllocator when this happens. And the error is logged as well.

vikrantpuppala · 2026-01-02T10:15:53Z

src/main/java/org/apache/arrow/memory/ArrowBuf.java

+
+  // ---- Databricks patch start ----
+  private final HistoricalLog historicalLog =
+      DEBUG ? new HistoricalLog(DEBUG_LOG_LENGTH, "ArrowBuf[%d]", id) : null;


do we even need to worry about the historical log? can we not just set it to null? seems like its usage is null checked everywhere anyway?

+1, add a comment on why this is needed

Idea is to have minimal patches w.r.t the Arrow code.

vikrantpuppala · 2026-01-02T10:29:48Z

src/main/java/org/apache/arrow/memory/DatabricksAllocationReservation.java

+    long currentReservation = reservedSize.get();
+    long newReservation = currentReservation + nBytes;
+    if (newReservation > allocator.getHeadroom() + currentReservation) {
+      return false;
+    }
+    reservedSize.addAndGet(nBytes);


this is not thread safe, should we use compareAndSet?

The interface that this class implements AllocationReservation, is explicitly marked as not thread-safe. So we are following contract. Also this code is generally called from a single thread.

vikrantpuppala · 2026-01-02T10:35:48Z

src/main/java/org/apache/arrow/memory/DatabricksArrowBuf.java

+  @Override
+  public ReferenceManager getReferenceManager() {
+    return referenceManager;
+  }
+
+  @Override
+  public long capacity() {
+    return capacity;
+  }


i think a bunch of these methods do not have a changed override behaviour, can we remove them so that it is easy to review and maintain?

Idea is to have all behaviour captured within a single class and make all things explicit. Also it is not dependent on any changes in the base class ArrowBuf.

vikrantpuppala · 2026-01-02T10:36:55Z

src/main/java/org/apache/arrow/memory/DatabricksArrowBuf.java

+    if (capacity > Integer.MAX_VALUE) {
+      throw new IllegalArgumentException(
+          "DatabricksArrowBuf does not support capacity > Integer.MAX_VALUE");
+    }


why this limit?

this is missing from the other constructor

can we reuse constructors

ByteBuffer.allocate constructor takes an integer argument. It is an inherent limitation of Java nio ByteBuffer. In JDBC case the maximum allocation will be a chunk size, which is 20MiB as of today.

Reuse constructor - I will give it more thought and get back.

In one case the constructor is allocating the ByteBuffer, the check has to happen before that. In the other case the allocator is being sliced. Delegating to one single constructor is not clean in this case.

gopalldb · 2026-01-05T05:18:59Z

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowBufferAllocator.java

+  static {
+    RootAllocator rootAllocator = null;
+    try {
+      rootAllocator = new RootAllocator();


we were using Integer.MAX_VALUE, not needed now?

Internally RootAllocator delegates with Integer.MAX_VALUE.

gopalldb · 2026-01-05T06:07:13Z

src/main/java/org/apache/arrow/memory/ArrowBuf.java

+  // ---- to avoid unsafe allocation initialization errors.
+  public static final String DEBUG_ALLOCATOR = "arrow.memory.debug.allocator";
+  public static final int DEBUG_LOG_LENGTH = 6;
+  public static final boolean DEBUG;


nit: rename to better name than debug?

This code is copied verbatim from another Arrow class BaseAllocator.

gopalldb · 2026-01-06T03:07:17Z

src/main/java/org/apache/arrow/memory/util/MemoryUtil.java

+                  + "(See https://arrow.apache.org/docs/java/install.html)",
+              e);
+      // ---- Databricks patch start ----
+      // ---- Remove 'failure.printStackTrace();'


can we log the stack trace?

This code explicitly removes the print of stack trace to stderr to prevent customers from thinking something is broken, since our fix works with the native Arrow being absent as well. The exception is caught and error message is logged in static initiliazer of ArrowBufferAllocator.

gopalldb · 2026-01-06T03:10:34Z

src/main/java/org/apache/arrow/memory/DatabricksAllocationReservation.java

+  @Override
+  public ArrowBuf allocateBuffer() {
+    assertNotUsed();
+    if (!used.compareAndSet(false, true)) {


can we add logging?

Sure. I will add it to more classes as well.

jayantsing-db

In progress.

jayantsing-db · 2026-01-18T17:28:15Z

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowBufferAllocator.java

+    RootAllocator rootAllocator = null;
+    try {
+      rootAllocator = new RootAllocator();
+    } catch (Throwable t) {


Can this be more specific like java.lang.reflect.InaccessibleObjectException? This is so that we know when we want to fallback to custom allocator. Otherwise, we could be falling back for unknown reason.

Different JVM versions throw different exceptions, so code becomes brittle and error prone. Also there is nothing inside RootAllocator that should throw an exception other than Class initialisation failure. This is the catch-all safe option.

jayantsing-db · 2026-01-18T17:45:07Z

src/main/java/com/databricks/jdbc/api/impl/arrow/ArrowBufferAllocator.java

+  static {
+    RootAllocator rootAllocator = null;
+    try {
+      rootAllocator = new RootAllocator();


Is instantiating a root allocator object sufficient/robust signal to fallback to custom allocator? For example, for resource/IO wrapper objects like root allocator, Arrow may choose to just create a lightweight object/handle and the off-heap memory is not unsafe-accessed until a memory is explicitly buffered using that root allocator object. So, just creating a root allocator object may succeed but as the arrow reader proceeds to buffer off-heap memory during runtime, then it fails.

Should we do something like this to be deliberate about the fallback logic?

Class<?> unsafeClass = Class.forName.... Field f = .... f.setAccessible(true);

Awesome catch! Changed the logic to write and check.

Added tests for these as well. See stack-1.

jayantsing-db · 2026-01-18T19:13:54Z

src/main/java/org/apache/arrow/memory/ArrowBuf.java

+  // ---- Databricks patch start ----
+  // ---- Copied verbatim from BaseAllocator. We avoid initializing static fields of BaseAllocator
+  // ---- to avoid unsafe allocation initialization errors.


How do we ensure that base allocator class is not loaded (and hence the breaking static code inside base allocator) elsewhere during runtime in JDK 16+?

Tests should catch this. Please see DatabricksArrowPatchReaderWriterTest

jayantsing-db · 2026-01-18T19:16:04Z

src/main/java/org/apache/arrow/memory/ArrowBuf.java

+  // Initialize this before DEFAULT_CONFIG as DEFAULT_CONFIG will eventually initialize the
+  // allocation manager,
+  // which in turn allocates an ArrowBuf, which requires DEBUG to have been properly initialized


Since this is already within the patch block, we can remove this potentially confusing comment from the BaseAllocator class.

Copied code verbatim. Easier to diff.

jayantsing-db · 2026-01-18T19:32:03Z

src/main/java/org/apache/arrow/memory/ArrowBuf.java

+  private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ArrowBuf.class);
+
+  // ---- Databricks patch start ----
+  // ---- Copied verbatim from BaseAllocator. We avoid initializing static fields of BaseAllocator


qq: What was the procedure to determine that "we need to break the static dependency chain to BaseAllocator" because it ends up using unsafe? Is it based on empirical analysis or some deterministic code scan?

The reason I am asking is that there could be a whole bunch of classes loading in this arrow parsing path which might end up using unsafe and ultimately break in exciting new ways.

You are right. Only way to validate is through tests. Please see PR for stack-1 and stack-3.

jayantsing-db · 2026-01-18T19:47:50Z

src/main/java/org/apache/arrow/memory/util/MemoryUtil.java

+      // This exception will get swallowed, but it's necessary for the static analysis that ensures
+      // the static fields above get initialized


nit: Do we need to ensure that this exception is indeed correctly swallowed? If unsure, should we just stop throwing exception in the patch (although i agree with the above philosophy of avoiding as much patch lines as possible)

This is Arrow code. No changes from our end.

jayantsing-db · 2026-01-18T20:05:53Z

src/main/java/org/apache/arrow/memory/DatabricksReferenceManager.java

+  @Override
+  public ArrowBuf deriveBuffer(ArrowBuf sourceBuffer, long index, long length) {
+    Preconditions.checkArgument(
+        length <= Integer.MAX_VALUE,


Should this be index + length?

index + length <= sourceBuffer.capacity() check is present after this line.

jayantsing-db · 2026-01-18T20:10:21Z

src/main/java/org/apache/arrow/memory/DatabricksReferenceManager.java

+    // Create a new DatabricksArrowBuf sharing the same byte buffer.
+    DatabricksArrowBuf buf = checkBufferType(sourceBuffer);
+    return new DatabricksArrowBuf(
+        this, null, buf.getByteBuffer(), buf.getOffset() + (int) index, length);


This always sets buffer manager as null. The reallocIfNeeded will always throw error on slices.

@Override public ArrowBuf reallocIfNeeded(final long 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"); } }

In principle, this may be OKAY because JDBC is read-only from native bytes and reallocIfNeeded is never called?

The code is exactly as it is in ArrowBuf. Their code path also has cases where bufferManager is null. It is an artifact of how the code is called.

jayantsing-db · 2026-01-18T20:52:40Z

src/main/java/org/apache/arrow/memory/DatabricksReferenceManager.java

+
+  @Override
+  public boolean release(int decrement) {
+    return getRefCount() == 0;


This will always return false. Could this cause any unknown issues? For example, in the chain of calls, when an arrow reader is closed -> vectors are closed -> buffer is released; could this method (always returning false), cause any unintended leaks by short-circuiting any close chain? I understand that this doesn't matter for on-heap arrow bufs.

I have read through the code, it should not.

jayantsing-db · 2026-01-18T20:57:52Z

src/main/java/org/apache/arrow/memory/DatabricksReferenceManager.java

+
+  @Override
+  public ArrowBuf retain(ArrowBuf srcBuffer, BufferAllocator targetAllocator) {
+    DatabricksArrowBuf buf = checkBufferType(srcBuffer);


Should we assert target allocator type too?

targetAllocator is unused in our case.

jayantsing-db · 2026-01-18T21:15:59Z

Quick question (while I am excited to see this go to production) how do we want to handle fragility and maintainability here? My preference would be to put this patch in the arrow-java repo under the Databricks module so we minimize the chances of things breaking or drifting over time.

jayantsing-db · 2026-01-18T21:42:20Z

Quick question (while I am excited to see this go to production) how do we want to handle fragility and maintainability here? My preference would be to put this patch in the arrow-java repo under the Databricks module so we minimize the chances of things breaking or drifting over time.

Adding to the above. We may have already looked at this, but I am calling it out for reference and to make sure I understanding correctly. Did we consider using Netty’s allocator distribution rather than memory-unsafe? From what I can see, all of the unsafe allocators stem from memory-unsafe.

If memory-netty works, we can avoid the additional fragility inherent in the current implementation.

Instantiating RootAllocator is insufficient to check that Unsafe memory operations are permitted in Arrow. Writing to an allocated object to validate that it works.

tejassp-db added 4 commits December 16, 2025 15:58

PECOBLR-1121 Patch Arrow to circumvent JVM args issue.

37d7d15

Patch Arrow to create a Databricks ArrowBuf which allocates memory on the heap and provides access to it through Java methods. This removes the need to specify "--add-opens=java.base/java.nio=ALL-UNNAMED" as JVM args for JDK 16+.

PECOBLR-1121 Use Arrow patch as fallback.

ffd6c1c

Use native Arrow if available. Otherwise fallback to the patch version.

PECOBLR-1121 Simplify patch code.

a78a597

Remove irrelevant reference counting in patch code. Patch code uses heap memory for arrow operations and reference counting is not required.

PECOBLR-1121 Minor refactor.

1654f74

tejassp-db self-assigned this Dec 22, 2025

tejassp-db added 5 commits December 23, 2025 16:11

PECOBLR-1121 Fix todos and fixmes.

42422f1

Remove redundant todos for accounting.

PECOBLR-1121 Fix derive buffer

36c2d3d

PECOBLR-1121 Patch DecimalUtility.

dcdc49a

Patch DecimalUtility to not use unsafe methods to set decimal values on DatabricksArrowBuf.

PECOBLR-1121 Add Apache 2 compliant changes.

7c44728

Add notice to all patched Arrow Java code. In NOTICE file mention Arrow has been patched by Databricks.

PECOBLR-1121 Suppress stack trace print on Arrow class init failure.

44271f9

On static init failure of MemoryUtil class, it prints a stack trace to stderr. Remove this print, since now we fallback to DatabricksBufferAllocator when this happens. And the error is logged as well.

tejassp-db requested review from gopalldb, jayantsing-db, jprakash-db, madhav-db, msrathore-db, samikshya-db and vikrantpuppala January 2, 2026 09:47

tejassp-db marked this pull request as ready for review January 2, 2026 09:48

tejassp-db requested a review from sreekanth-db January 2, 2026 09:53

vikrantpuppala reviewed Jan 2, 2026

View reviewed changes

tejassp-db added 3 commits January 5, 2026 10:27

Merge branch 'main' into PECOBLR-1121/arrow-patch/stack-0

1fe15fb

PECOBLR-1121 Fix TODOs.

1cddb50

PECOBLR-1121 Show cause of RootAllocator load failure.

06762bb

gopalldb reviewed Jan 5, 2026

View reviewed changes

gopalldb reviewed Jan 6, 2026

View reviewed changes

jayantsing-db reviewed Jan 18, 2026

View reviewed changes

tejassp-db added 4 commits January 23, 2026 09:00

PECOBLR-1121 Add log lines to patched Arrow classes.

8ee6fbe

PECOBLR-1121 Formatting.

4b34de5

PECOBLR-1121 Fix incorrect check of RootAllocator.

402d062

Instantiating RootAllocator is insufficient to check that Unsafe memory operations are permitted in Arrow. Writing to an allocated object to validate that it works.

PECOBLR-1121 Pull static initializer into separate function.

215156a

		// This exception will get swallowed, but it's necessary for the static analysis that ensures
		// the static fields above get initialized

PECOBLR-1121 Arrow patch to circumvent Arrow issues with JDk 16+. #1156

Are you sure you want to change the base?

PECOBLR-1121 Arrow patch to circumvent Arrow issues with JDk 16+. #1156

Conversation

tejassp-db commented Dec 22, 2025

Description

Testing

Additional Notes to the Reviewer

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayantsing-db left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tejassp-db Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tejassp-db Jan 23, 2026 •

edited

Loading

jayantsing-db commented Jan 18, 2026 •

edited

Loading