Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions NOTICE
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ Notice - https://github.com/databricks/databricks-sdk-java/blob/main/NOTICE

apache/arrow - https://github.com/apache/arrow/tree/main
Copyright 2016-2025 The Apache Software Foundation
*This software contains code modified by Databricks, Inc.*
Notice - https://github.com/apache/arrow/blob/main/NOTICE.txt

diffplug/spotless - https://github.com/diffplug/spotless/tree/main
Expand Down
2 changes: 1 addition & 1 deletion pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@
<properties>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
<arrow.version>17.0.0</arrow.version>
<arrow.version>18.3.0</arrow.version>
<commons-lang3.version>3.18.0</commons-lang3.version>
<maven.compiler.source>11</maven.compiler.source>
<maven.compiler.target>11</maven.compiler.target>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.ValueVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamReader;
Expand Down Expand Up @@ -99,7 +98,7 @@ protected AbstractArrowResultChunk(
this.rowOffset = rowOffset;
this.chunkIndex = chunkIndex;
this.statementId = statementId;
this.rootAllocator = new RootAllocator(Integer.MAX_VALUE);
this.rootAllocator = ArrowBufferAllocator.getBufferAllocator();
this.chunkReadyFuture = new CompletableFuture<>();
this.chunkLink = chunkLink;
this.expiryTime = expiryTime;
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
package com.databricks.jdbc.api.impl.arrow;

import com.databricks.jdbc.log.JdbcLogger;
import com.databricks.jdbc.log.JdbcLoggerFactory;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.DatabricksBufferAllocator;
import org.apache.arrow.memory.RootAllocator;

/**
* Creates {@link BufferAllocator} instances.
*
* <p>First tries to create a {@link RootAllocator} which uses off-heap memory and is faster. If
* that fails (usually due to JVM reflection restrictions), falls back to {@link
* DatabricksBufferAllocator} which uses heap memory.
*/
public class ArrowBufferAllocator {
/** Can a {@code RootAllocator} be created in this JVM instance? */
private static final boolean canCreateRootAllocator;

/** Logger instance. */
private static final JdbcLogger LOGGER = JdbcLoggerFactory.getLogger(ArrowBufferAllocator.class);

/* Check if the RootAllocator can be instantiated. */
static {
RootAllocator rootAllocator = null;
try {
rootAllocator = new RootAllocator();
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we were using Integer.MAX_VALUE, not needed now?

} catch (Throwable t) {
LOGGER.info(
"Failed to create RootAllocator, will use DatabricksBufferAllocator as fallback: "
+ t.getMessage());
}

canCreateRootAllocator = rootAllocator != null;
if (rootAllocator != null) {
try {
rootAllocator.close();
} catch (Throwable t) {
LOGGER.warn("RootAllocator could not be closed: " + t.getMessage());
}
}
}

/**
* @return an instance of the {@code BufferAllocator}.
*/
public static BufferAllocator getBufferAllocator() {
// TODO reuse RootAllocators?
// TODO should this method be non-static.
if (canCreateRootAllocator) {
return new RootAllocator();
} else {
return new DatabricksBufferAllocator();
}
}
}
Loading
Loading