add maxMemory for RootAllocator

Author: kczimm

README.md

@@ -122,6 +122,18 @@ SpiceClient client = SpiceClient.builder()
 Retries are performed for connection and system internal errors. It is the SDK user's responsibility to properly
 handle other errors, for example RESOURCE_EXHAUSTED (HTTP 429).
 
+### Memory Configuration
+
+The `SpiceClient` uses an Arrow `RootAllocator` for managing off-heap memory. By default, it allows unlimited memory allocation. To prevent OOM issues in constrained environments, you can limit the maximum memory:
+
+```java
+SpiceClient client = SpiceClient.builder()
+    .withMaxMemory(1024 * 1024 * 1024L) // 1GB limit
+    .build();
+```
+
+You can also set the limit via the `SPICE_MAX_MEMORY` environment variable (in bytes). If both are provided, the builder method takes precedence.
+
 ### Spice.ai Runtime commands
 
 #### Accelerated dataset refresh

src/main/java/ai/spice/Config.java

@@ -93,6 +93,27 @@ public static URI getCloudHttpAddressUri() throws URISyntaxException {
         return new URI(CLOUD_HTTP_ADDRESS);
     }
 
+    /**
+     * Returns the maximum memory allocation for the Arrow RootAllocator from the
+     * SPICE_MAX_MEMORY environment variable.
+     *
+     * @return the maximum memory in bytes, or Long.MAX_VALUE if not set or invalid.
+     */
+    public static long getMaxMemory() {
+        String maxMemoryStr = System.getenv("SPICE_MAX_MEMORY");
+        if (maxMemoryStr != null && !maxMemoryStr.isEmpty()) {
+            try {
+                long maxMemory = Long.parseLong(maxMemoryStr);
+                if (maxMemory > 0) {
+                    return maxMemory;
+                }
+            } catch (NumberFormatException e) {
+                // Invalid value, fall back to default
+            }
+        }
+        return Long.MAX_VALUE;
+    }
+
     /**
      * Returns the Spice SDK user agent for this system, including the package
      * version, system OS, version and arch.
@@ -127,4 +148,4 @@ public static String getUserAgent() {
                 + System.getProperty("os.version") + " "
                 + osArch + ")";
     }
-}
+}

src/main/java/ai/spice/SpiceClient.java

@@ -91,8 +91,10 @@ public static SpiceClientBuilder builder() throws URISyntaxException {
      * @param httpAddress   the URI of the Spice.ai runtime HTTP address
      * 
      * @param maxRetries    the maximum number of connection retries for the client
+     * @param userAgent     the user agent string
+     * @param maxMemory     the maximum memory allocation for the Arrow RootAllocator
      */
-    public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries, String userAgent) {
+    public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddress, int maxRetries, String userAgent, long maxMemory) {
         this.appId = appId;
         this.apiKey = apiKey;
         this.maxRetries = maxRetries;
@@ -108,7 +110,7 @@ public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddre
             this.flightAddress = flightAddress;
         }
 
-        Builder builder = FlightClient.builder(new RootAllocator(Long.MAX_VALUE), new Location(this.flightAddress));
+        Builder builder = FlightClient.builder(new RootAllocator(maxMemory), new Location(this.flightAddress));
 
         if (Strings.isNullOrEmpty(apiKey)) {
             this.flightClient = new FlightSqlClient(builder.build());

src/main/java/ai/spice/SpiceClientBuilder.java

@@ -38,6 +38,7 @@ public class SpiceClientBuilder {
     private URI flightAddress;
     private URI httpAddress;
     private int maxRetries = 3;
+    private long maxMemory = Long.MAX_VALUE;
 
     /**
      * Constructs a new SpiceClientBuilder instance
@@ -47,6 +48,7 @@ public class SpiceClientBuilder {
     SpiceClientBuilder() throws URISyntaxException {
         this.flightAddress = Config.getLocalFlightAddressUri();
         this.httpAddress = Config.getLocalHttpAddressUri();
+        this.maxMemory = Config.getMaxMemory();
     }
 
     /**
@@ -139,12 +141,26 @@ public SpiceClientBuilder withMaxRetries(int maxRetries) {
         return this;
     }
 
+    /**
+     * Sets the maximum memory allocation for the Arrow RootAllocator.
+     *
+     * @param maxMemory The maximum memory in bytes for the RootAllocator (must be > 0)
+     * @return The current instance of SpiceClientBuilder for method chaining.
+     */
+    public SpiceClientBuilder withMaxMemory(long maxMemory) {
+        if (maxMemory <= 0) {
+            throw new IllegalArgumentException("maxMemory must be greater than 0");
+        }
+        this.maxMemory = maxMemory;
+        return this;
+    }
+
     /**
      * Creates SpiceClient with provided parameters.
      *
      * @return The SpiceClient instance
      */
     public SpiceClient build() {
-        return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent);
+        return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, maxMemory);
     }
-}
+}