Update

Author: sgrebnov

README.md

@@ -124,11 +124,11 @@ 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 limits memory to 50% of the JVM's maximum heap size or 1GB, whichever is lower. To prevent OOM issues in constrained environments, you can limit the maximum memory:
+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()
-    .withMaxMemory(1024 * 1024 * 1024L) // 1GB limit
+    .withMemoryLimitMB(1024) // 1GB limit
     .build();
 ```
 

pom.xml

@@ -53,11 +53,6 @@
             <artifactId>slf4j-simple</artifactId>
             <version>2.0.16</version>
         </dependency>
-        <dependency>
-            <groupId>com.github.oshi</groupId>
-            <artifactId>oshi-core</artifactId>
-            <version>6.9.1</version>
-        </dependency>
         <dependency>
             <groupId>junit</groupId>
             <artifactId>junit</artifactId>

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,15 +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 maxMemory     the maximum memory allocation for the Arrow RootAllocator
+     * @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, long maxMemory) {
+    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;
@@ -110,7 +116,12 @@ public SpiceClient(String appId, String apiKey, URI flightAddress, URI httpAddre
             this.flightAddress = flightAddress;
         }
 
-        Builder builder = FlightClient.builder(new RootAllocator(maxMemory), 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

@@ -26,8 +26,6 @@ of this software and associated documentation files (the "Software"), to deal
 import java.net.URISyntaxException;
 
 import com.google.common.base.Strings;
-import oshi.SystemInfo;
-import oshi.hardware.GlobalMemory;
 
 /**
  * Builder class for creating instances of SpiceClient.
@@ -40,33 +38,7 @@ public class SpiceClientBuilder {
     private URI flightAddress;
     private URI httpAddress;
     private int maxRetries = 3;
-    private long maxMemory = Long.MAX_VALUE;
-
-    /**
-     * Calculates the default maximum memory for the Arrow RootAllocator.
-     * Returns 50% of the system's total physical memory or 1GB, whichever is lower.
-     * Falls back to JVM heap size if system memory cannot be determined.
-     *
-     * @return the default maximum memory in bytes.
-     */
-    private static long calculateDefaultMaxMemory() {
-        long oneGB = 1024L * 1024 * 1024;
-        try {
-            SystemInfo systemInfo = new SystemInfo();
-            GlobalMemory memory = systemInfo.getHardware().getMemory();
-            long totalMemory = memory.getTotal();
-            if (totalMemory > 0) {
-                long halfMemory = totalMemory / 2;
-                return Math.min(halfMemory, oneGB);
-            }
-        } catch (Exception e) {
-            // Fallback to JVM heap if OSHI fails
-        }
-
-        long jvmMaxMemory = Runtime.getRuntime().maxMemory();
-        long halfMemory = jvmMaxMemory / 2;
-        return Math.min(halfMemory, oneGB);
-    }
+    private long memoryLimitMB = Long.MAX_VALUE; // Default is all available memory.
 
     /**
      * Constructs a new SpiceClientBuilder instance
@@ -76,7 +48,6 @@ private static long calculateDefaultMaxMemory() {
     SpiceClientBuilder() throws URISyntaxException {
         this.flightAddress = Config.getLocalFlightAddressUri();
         this.httpAddress = Config.getLocalHttpAddressUri();
-        this.maxMemory = calculateDefaultMaxMemory();
     }
 
     /**
@@ -170,16 +141,22 @@ public SpiceClientBuilder withMaxRetries(int maxRetries) {
     }
 
     /**
-     * Sets the maximum memory allocation for the Arrow RootAllocator.
+     * 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 maxMemory The maximum memory in bytes for the RootAllocator (must be > 0)
+     * @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 withMaxMemory(long maxMemory) {
-        if (maxMemory <= 0) {
-            throw new IllegalArgumentException("maxMemory must be greater than 0");
+    public SpiceClientBuilder withMemoryLimitMB(long memoryLimitMB) {
+        if (memoryLimitMB <= 0) {
+            throw new IllegalArgumentException("Memory limit must be positive, got: " + memoryLimitMB + " MB");
         }
-        this.maxMemory = maxMemory;
+        this.memoryLimitMB = memoryLimitMB;
         return this;
     }
 
@@ -189,6 +166,6 @@ public SpiceClientBuilder withMaxMemory(long maxMemory) {
      * @return The SpiceClient instance
      */
     public SpiceClient build() {
-        return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, maxMemory);
+        return new SpiceClient(appId, apiKey, flightAddress, httpAddress, maxRetries, userAgent, memoryLimitMB);
     }
 }

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

@@ -0,0 +1,51 @@
+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()
+                .withMemoryLimitMB(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) {
+            // We expect some exceptions due to no local Spice instance,
+            // but not IllegalArgumentException from memory configuration
+            assertFalse("Should not throw IllegalArgumentException for valid memory config", 
+                       e instanceof IllegalArgumentException);
+        }
+    }
+
+    public void testInvalidMemoryLimitMB() throws Exception {
+        try {
+            SpiceClient.builder()
+                .withMemoryLimitMB(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()
+                .withMemoryLimitMB(-100) // Invalid: negative
+                .build();
+            fail("Should throw IllegalArgumentException for negative memory limit");
+        } catch (IllegalArgumentException e) {
+            assertTrue("Should throw IllegalArgumentException for negative memory limit", true);
+        }
+    }
+}