Merge pull request #252 from datastax/issue/242-partitionfile-outofmemoryerror

Fix OOM issue # 242

Author: pravinbhat

src/main/java/com/datastax/cdm/job/SplitPartitions.java

@@ -24,16 +24,17 @@
 import java.io.*;
 import java.math.BigInteger;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
 public class SplitPartitions {
 
     public static Logger logger = LoggerFactory.getLogger(SplitPartitions.class.getName());
+    private static final int MAX_NUM_PARTS_FOR_PARTITION_FILE = 10;
 
     public static List<Partition> getRandomSubPartitions(int numSplits, BigInteger min, BigInteger max, int coveragePercent) {
         logger.info("ThreadID: {} Splitting min: {} max: {}", Thread.currentThread().getId(), min, max);
@@ -48,16 +49,23 @@ public static List<Partition> getRandomSubPartitions(int numSplits, BigInteger m
     public static List<Partition> getSubPartitionsFromFile(int numSplits, String inputFilename) throws IOException {
         logger.info("ThreadID: {} Splitting partitions in file: {} using a split-size of {}"
                 , Thread.currentThread().getId(), inputFilename, numSplits);
+        if (numSplits > 10) {
+            logger.warn("Resetting spark.cdm.perfops.numParts value of {} to max allowed value of {} when using a partition file: {}",
+                    numSplits, MAX_NUM_PARTS_FOR_PARTITION_FILE, inputFilename);
+            numSplits = MAX_NUM_PARTS_FOR_PARTITION_FILE;
+        }
         List<Partition> partitions = new ArrayList<Partition>();
         BufferedReader reader = getfileReader(inputFilename);
         String line = null;
+        PartitionMinMax pMinMax;
         while ((line = reader.readLine()) != null) {
-            if (line.startsWith("#")) {
-                continue;
-            }
-            String[] minMax = line.split(",");
             try {
-                partitions.addAll(getSubPartitions(numSplits, new BigInteger(minMax[0]), new BigInteger(minMax[1]), 100));
+                pMinMax = new PartitionMinMax(line);
+                if (pMinMax.hasError) {
+                    logger.error("Skipping " + pMinMax.error);
+                    continue;
+                }
+                partitions.addAll(getSubPartitions(numSplits, pMinMax.min, pMinMax.max, 100));
             } catch (Exception e) {
                 logger.error("Skipping partition: {}", line, e);
             }
@@ -66,6 +74,26 @@ public static List<Partition> getSubPartitionsFromFile(int numSplits, String inp
         return partitions;
     }
 
+    static class PartitionMinMax {
+        static final Pattern pat = Pattern.compile("^-?[0-9]*,-?[0-9]*");
+        public BigInteger min;
+        public BigInteger max;
+        public boolean hasError = false;
+        public String error;
+
+        public PartitionMinMax(String line) {
+            line = line.replaceAll(" ", "");
+            if (!pat.matcher(line).matches()) {
+                error = "Invaliding partition line: " + line;
+                hasError = true;
+                return;
+            }
+            String[] minMax = line.split(",");
+            min = new BigInteger(minMax[0]);
+            max = new BigInteger(minMax[1]);
+        }
+    }
+
     public static List<PKRows> getRowPartsFromFile(int numSplits, String inputFilename) throws IOException {
         logger.info("ThreadID: {} Splitting rows in file: {} using a split-size of {}"
                 , Thread.currentThread().getId(), inputFilename, numSplits);

src/test/java/com/datastax/cdm/job/SplitPartitionsTest.java

@@ -14,17 +14,15 @@
 
 package com.datastax.cdm.job;
 
-import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.io.IOException;
 import java.math.BigInteger;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.List;
 import java.util.stream.Stream;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.*;
 
 public class SplitPartitionsTest {
 
@@ -44,11 +42,52 @@ void getRandomSubPartitionsTestOver100() {
                 BigInteger.valueOf(44), 200);
         assertEquals(8, partitions.size());
     }
+
+    @Test
+    void getSubPartitionsFromFileNoFileTest() throws IOException {
+        Exception excp = assertThrows(RuntimeException.class,
+                () -> SplitPartitions.getSubPartitionsFromFile(8, "partfile.csv"));
+        assertTrue(excp.getMessage().contains("No 'partfile.csv' file found!! Add this file in the current folder & rerun!"));
+    }
+
+    @Test
+    void getSubPartitionsFromFileTest() throws IOException {
+        List<SplitPartitions.Partition> partitions = SplitPartitions.getSubPartitionsFromFile(5, "./src/resources/partitions.csv");
+        assertEquals(25, partitions.size());
+    }
+
+    @Test
+    void getSubPartitionsFromHighNumPartTest() throws IOException {
+        List<SplitPartitions.Partition> partitions = SplitPartitions.getSubPartitionsFromFile(1000, "./src/resources/partitions.csv");
+        assertEquals(50, partitions.size());
+    }
+
+
     @Test
     void batchesTest() {
         List<String> mutable_list = Arrays.asList("e1", "e2", "e3", "e4", "e5", "e6");
         Stream<List<String>> out = SplitPartitions.batches(mutable_list, 2);
         assertEquals(3, out.count());
     }
 
+    @Test
+    void PartitionMinMaxBlankTest() {
+        assertTrue((new SplitPartitions.PartitionMinMax("")).hasError);
+    }
+
+    @Test
+    void PartitionMinMaxInvalidTest() {
+        assertTrue((new SplitPartitions.PartitionMinMax(" # min, max")).hasError);
+    }
+
+    @Test
+    void PartitionMinMaxValidTest() {
+        assertTrue(!(new SplitPartitions.PartitionMinMax(" -123, 456")).hasError);
+    }
+
+    @Test
+    void PartitionMinMaxValidMinMaxTest() {
+        assertEquals(BigInteger.valueOf(-507900353496146534l), (new SplitPartitions.PartitionMinMax(" -507900353496146534, 456")).min);
+        assertEquals(BigInteger.valueOf(9101008634499147643l), (new SplitPartitions.PartitionMinMax(" -507900353496146534,9101008634499147643")).max);
+    }
 }