From e81f573349535f5496093d179ad8eaa12e53908d Mon Sep 17 00:00:00 2001
From: Laurent Goujon <laurent@dremio.com>
Date: Mon, 18 Apr 2016 11:07:22 -0700
Subject: [PATCH] [ARROW-113] Make BaseValueVector#MAX_ALLOCATION_SIZE
 configurable

Some of the tests are based on the assumption that the JVM can allocate at least
2GB of memory, which is not a common occurence (JVM usually defaults at 512MB).
Current Travis CI VM only have 3GB of memory total, which would have make challenging
to run some of the tests on them

Add a system property to change BaseValueVector.MAX_ALLOCATION_SIZE to allow to use
a much smaller value during tests.
---
 .../apache/arrow/vector/BaseValueVector.java  | 14 ++++----
 .../apache/arrow/vector/TestValueVector.java  | 36 +++++++++++++++----
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/java/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java b/java/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java
index 8bca3c00537..932e6f13caf 100644
--- a/java/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java
+++ b/java/vector/src/main/java/org/apache/arrow/vector/BaseValueVector.java
@@ -17,23 +17,24 @@
  */
 package org.apache.arrow.vector;
 
-import io.netty.buffer.ArrowBuf;
-
 import java.util.Iterator;
 
-import com.google.common.base.Preconditions;
-import com.google.common.collect.Iterators;
-
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.types.MaterializedField;
 import org.apache.arrow.vector.util.TransferPair;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Iterators;
+
+import io.netty.buffer.ArrowBuf;
+
 public abstract class BaseValueVector implements ValueVector {
   private static final Logger logger = LoggerFactory.getLogger(BaseValueVector.class);
 
-  public static final int MAX_ALLOCATION_SIZE = Integer.MAX_VALUE;
+  public static final String MAX_ALLOCATION_SIZE_PROPERTY = "arrow.vector.max_allocation_bytes";
+  public static final int MAX_ALLOCATION_SIZE = Integer.getInteger(MAX_ALLOCATION_SIZE_PROPERTY, Integer.MAX_VALUE);
   public static final int INITIAL_VALUE_ALLOCATION = 4096;
 
   protected final BufferAllocator allocator;
@@ -99,6 +100,7 @@ protected BaseMutator() { }
     public void generateTestData(int values) {}
 
     //TODO: consider making mutator stateless(if possible) on another issue.
+    @Override
     public void reset() {}
   }
 
diff --git a/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java b/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java
index ac3eebe98ea..b5c4509c8b5 100644
--- a/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java
+++ b/java/vector/src/test/java/org/apache/arrow/vector/TestValueVector.java
@@ -23,16 +23,12 @@
 
 import java.nio.charset.Charset;
 
+import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.complex.ListVector;
 import org.apache.arrow.vector.complex.MapVector;
 import org.apache.arrow.vector.complex.RepeatedListVector;
 import org.apache.arrow.vector.complex.RepeatedMapVector;
-import org.apache.arrow.vector.types.MaterializedField;
-import org.apache.arrow.vector.types.Types;
-import org.apache.arrow.vector.types.Types.MinorType;
-import org.apache.arrow.vector.util.BasicTypeHelper;
-import org.apache.arrow.vector.util.OversizedAllocationException;
 import org.apache.arrow.vector.holders.BitHolder;
 import org.apache.arrow.vector.holders.IntHolder;
 import org.apache.arrow.vector.holders.NullableFloat4Holder;
@@ -44,10 +40,16 @@
 import org.apache.arrow.vector.holders.RepeatedVarBinaryHolder;
 import org.apache.arrow.vector.holders.UInt4Holder;
 import org.apache.arrow.vector.holders.VarCharHolder;
-import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.types.MaterializedField;
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.Types.MinorType;
+import org.apache.arrow.vector.util.BasicTypeHelper;
+import org.apache.arrow.vector.util.OversizedAllocationException;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExternalResource;
 
 
 public class TestValueVector {
@@ -57,6 +59,28 @@ public class TestValueVector {
 
   private BufferAllocator allocator;
 
+  // Rule to adjust MAX_ALLOCATION_SIZE and restore it back after the tests
+  @Rule
+  public final ExternalResource rule = new ExternalResource() {
+    private final String systemValue = System.getProperty(BaseValueVector.MAX_ALLOCATION_SIZE_PROPERTY);
+    private final String testValue = Long.toString(32*1024*1024);
+
+    @Override
+    protected void before() throws Throwable {
+      System.setProperty(BaseValueVector.MAX_ALLOCATION_SIZE_PROPERTY, testValue);
+    }
+
+    @Override
+    protected void after() {
+      if (systemValue != null) {
+        System.setProperty(BaseValueVector.MAX_ALLOCATION_SIZE_PROPERTY, systemValue);
+      }
+      else {
+        System.clearProperty(BaseValueVector.MAX_ALLOCATION_SIZE_PROPERTY);
+      }
+    }
+  };
+
   @Before
   public void init() {
     allocator = new RootAllocator(Long.MAX_VALUE);