Add configurable Arrow memory limit - withArrowMemoryLimitMB (#32)

Author: kczimm

.github/workflows/build.yaml

@@ -41,20 +41,14 @@ jobs:
       if: matrix.os == 'windows-latest'
       run: mvn --% install -DskipTests=true -Dgpg.skip -B -V # tell powershell to stop parsing with --% so it doesn't error with "Unknown lifecycle phase .skip"
 
-    - name: Install Spice (https://install.spiceai.org) (Linux)
-      if: matrix.os == 'ubuntu-latest'
+    - name: Install Spice (https://install.spiceai.org) (Unix)
+      if: matrix.os == 'ubuntu-latest' || matrix.os == 'macos-latest'
       env:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       run: |
         curl https://install.spiceai.org | /bin/bash
         echo "$HOME/.spice/bin" >> $GITHUB_PATH
         $HOME/.spice/bin/spice install
-    
-    - name: Install Spice (https://install.spiceai.org) (MacOS)
-      if: matrix.os == 'macos-latest'
-      run: |
-        brew install spiceai/spiceai/spice
-        brew install spiceai/spiceai/spiced
 
     - name: install Spice (Windows)
       if: matrix.os == 'windows-latest'

README.md

@@ -122,6 +122,16 @@ 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 uses all available memory. You can configure the memory limit using megabytes:
+
+```java
+SpiceClient client = SpiceClient.builder()
+    .withArrowMemoryLimitMB(1024) // 1GB limit
+    .build();
+```
+
 ### Spice.ai Runtime commands
 
 #### Accelerated dataset refresh

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

@@ -127,4 +127,4 @@ public static String getUserAgent() {
                 + System.getProperty("os.version") + " "
                 + osArch + ")";
     }
-}
+}

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

@@ -61,6 +61,8 @@ of this software and associated documentation files (the "Software"), to deal
  */
 public class SpiceClient implements AutoCloseable {
 
+    private static final long BYTES_PER_MB = 1024L * 1024L;
+
     private String appId;
     private String apiKey;
     private URI flightAddress;
@@ -86,13 +88,19 @@ public static SpiceClientBuilder builder() throws URISyntaxException {
      *                      application
      * @param apiKey        the API key used for authentication with Spice.ai
      *                      services
-     * @param flightAddress the URI of the flight address for connecting to Spice.ai
+     * @param flightAddress the URI of the flight address for connecting to
+     *                      Spice.ai
      *                      services
      * @param httpAddress   the URI of the Spice.ai runtime HTTP address
      * 
-     * @param maxRetries    the maximum number of connection retries for the client
+     * @param maxRetries    the maximum number of connection retries for the
+     *                      client
+     * @param userAgent     the user agent string
+     * @param memoryLimitMB the memory limit in megabytes 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 memoryLimitMB) {
         this.appId = appId;
         this.apiKey = apiKey;
         this.maxRetries = maxRetries;
@@ -108,7 +116,12 @@ 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));
+        // Convert megabytes to bytes for RootAllocator:
+        // https://arrow.apache.org/java/main/reference/org.apache.arrow.memory.core/org/apache/arrow/memory/RootAllocator.html
+        long memoryLimitBytes = (memoryLimitMB > Long.MAX_VALUE / BYTES_PER_MB)
+                ? Long.MAX_VALUE
+                : memoryLimitMB * BYTES_PER_MB;
+        Builder builder = FlightClient.builder(new RootAllocator(memoryLimitBytes), 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 memoryLimitMB = Long.MAX_VALUE; // Default is all available memory.
 
     /**
      * Constructs a new SpiceClientBuilder instance
@@ -139,12 +140,39 @@ public SpiceClientBuilder withMaxRetries(int maxRetries) {
         return this;
     }
 
+    /**
+     * Sets the memory limit for Apache Arrow allocator in megabytes.
+     * This controls the maximum amount of off-heap memory that can be allocated
+     * for Arrow Flight operations. If not set, the allocator will use all available
+     * memory.
+     *
+     * @param memoryLimitMB Maximum memory limit in megabytes. Default is all
+     *                      available memory.
+     *                      Must be positive.
+     * @return The current instance of SpiceClientBuilder for method chaining.
+     * @throws IllegalArgumentException if memoryLimitMB is not positive
+     *
+     * @see org.apache.arrow.memory.RootAllocator
+     */
+    public SpiceClientBuilder withArrowMemoryLimitMB(long memoryLimitMB) {
+        if (memoryLimitMB <= 0) {
+            throw new IllegalArgumentException("Memory limit must be positive, got: " + memoryLimitMB + " MB");
+        }
+        if (memoryLimitMB > Long.MAX_VALUE / 1024L / 1024L) {
+            throw new IllegalArgumentException(
+                    "Memory limit is too large: " + memoryLimitMB + " MB");
+        }
+
+        this.memoryLimitMB = memoryLimitMB;
+        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, memoryLimitMB);
     }
-}
+}

src/test/java/ai/spice/MemoryConfigurationTest.java

@@ -0,0 +1,80 @@
+package ai.spice;
+
+import junit.framework.TestCase;
+
+/**
+ * Test memory configuration functionality
+ */
+public class MemoryConfigurationTest extends TestCase {
+
+    public void testMemoryLimitMBConfiguration() throws Exception {
+        try {
+            // Test that memory limit in MB is properly configured
+            SpiceClient client = SpiceClient.builder()
+                    .withArrowMemoryLimitMB(128) // 128 MB
+                    .build();
+
+            // If we reach here without exception, the configuration worked
+            assertTrue("Memory configuration should not throw exception", true);
+
+            client.close();
+        } catch (Exception e) {
+            fail("Should not throw exception: " + e.getMessage());
+        }
+    }
+
+    public void testInvalidMemoryLimitMB() throws Exception {
+        try {
+            SpiceClient.builder()
+                    .withArrowMemoryLimitMB(0) // Invalid: must be positive
+                    .build();
+            fail("Should throw IllegalArgumentException for zero memory limit");
+        } catch (IllegalArgumentException e) {
+            // Expected exception
+            assertTrue("Should throw IllegalArgumentException for zero memory limit", true);
+        }
+    }
+
+    public void testNegativeMemoryLimitMB() throws Exception {
+        try {
+            SpiceClient.builder()
+                    .withArrowMemoryLimitMB(-100) // Invalid: negative
+                    .build();
+            fail("Should throw IllegalArgumentException for negative memory limit");
+        } catch (IllegalArgumentException e) {
+            assertTrue("Should throw IllegalArgumentException for negative memory limit", true);
+        }
+    }
+
+    public void testOverflowProtection() throws Exception {
+        // Test overflow protection - values that would cause overflow when converted to
+        // bytes
+        long maxSafeMB = Long.MAX_VALUE / (1024L * 1024L);
+        long overflowValue = maxSafeMB + 1;
+
+        try {
+            SpiceClient.builder()
+                    .withArrowMemoryLimitMB(overflowValue) // Would cause overflow
+                    .build();
+            fail("Should throw IllegalArgumentException for overflow-causing memory limit");
+        } catch (IllegalArgumentException e) {
+            assertTrue("Should throw IllegalArgumentException for overflow protection", true);
+        }
+    }
+
+    public void testMaxSafeMemoryLimit() throws Exception {
+        // Test that the maximum safe value works without throwing
+        long maxSafeMB = Long.MAX_VALUE / (1024L * 1024L);
+
+        try {
+            SpiceClient client = SpiceClient.builder()
+                    .withArrowMemoryLimitMB(maxSafeMB) // Maximum safe value
+                    .build();
+
+            assertTrue("Maximum safe memory limit should work", true);
+            client.close();
+        } catch (Exception e) {
+            fail("Should not throw exception: " + e.getMessage());
+        }
+    }
+}