[SPARK-5932] Added unit tests for ByteString conversion

Author: Ilya Ganelin

core/src/main/scala/org/apache/spark/util/Utils.scala

@@ -1036,6 +1036,46 @@ private[spark] object Utils extends Logging {
     JavaUtils.timeStringAsSec(str)
   }
 
+  /**
+   * Convert a passed byte string (e.g. 50b, 100kb, or 250mb) to bytes for
+   * internal use.
+   *
+   * If no suffix is provided, the passed number is assumed to be in bytes.
+   */
+  def byteStringAsBytes(str: String): Long = {
+    JavaUtils.byteStringAsBytes(str)
+  }
+
+  /**
+   * Convert a passed byte string (e.g. 50b, 100kb, or 250mb) to kilobytes for
+   * internal use.
+   *
+   * If no suffix is provided, the passed number is assumed to be in kilobytes.
+   */
+  def byteStringAsKB(str: String): Long = {
+    JavaUtils.byteStringAsKB(str)
+  }
+
+  /**
+   * Convert a passed byte string (e.g. 50b, 100kb, or 250mb) to megabytes for
+   * internal use.
+   *
+   * If no suffix is provided, the passed number is assumed to be in megabytes.
+   */
+  def byteStringAsMB(str: String): Long = {
+    JavaUtils.byteStringAsMB(str)
+  }
+
+  /**
+   * Convert a passed byte string (e.g. 50b, 100kb, or 250mb, 500gb) to gigabytes for
+   * internal use.
+   *
+   * If no suffix is provided, the passed number is assumed to be in gigabytes.
+   */
+  def byteStringAsGB(str: String): Long = {
+    JavaUtils.byteStringAsGB(str)
+  }
+
   /**
    * Convert a Java memory parameter passed to -Xmx (such as 300m or 1g) to a number of megabytes.
    */

core/src/test/scala/org/apache/spark/util/UtilsSuite.scala

@@ -33,6 +33,7 @@ import org.scalatest.FunSuite
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 
+import org.apache.spark.network.util.ByteUnit
 import org.apache.spark.SparkConf
 
 class UtilsSuite extends FunSuite with ResetSystemProperties {
@@ -63,6 +64,10 @@ class UtilsSuite extends FunSuite with ResetSystemProperties {
     assert(Utils.timeStringAsMs("1d") === TimeUnit.DAYS.toMillis(1))
 
     // Test invalid strings
+    intercept[NumberFormatException] {
+      Utils.timeStringAsMs("600l")
+    }
+    
     intercept[NumberFormatException] {
       Utils.timeStringAsMs("This breaks 600s")
     }
@@ -79,6 +84,64 @@ class UtilsSuite extends FunSuite with ResetSystemProperties {
       Utils.timeStringAsMs("This 123s breaks")
     }
   }
+
+  test("Test byteString conversion") {
+    // Test zero
+    assert(Utils.byteStringAsBytes("0") === 0)
+
+    assert(Utils.byteStringAsGB("1") === 1)
+    assert(Utils.byteStringAsGB("1g") === 1)
+    assert(Utils.byteStringAsGB("1000m") === 1)
+    assert(Utils.byteStringAsGB("1000000k") === 1)
+    assert(Utils.byteStringAsGB("1k") === ByteUnit.KB.toGB(1))
+    assert(Utils.byteStringAsGB("1m") === ByteUnit.MB.toGB(1))
+    assert(Utils.byteStringAsGB("1t") === ByteUnit.TB.toGB(1))
+    assert(Utils.byteStringAsGB("1p") === ByteUnit.PB.toGB(1))
+    
+    assert(Utils.byteStringAsMB("1") === 1)
+    assert(Utils.byteStringAsMB("1m") === 1)
+    assert(Utils.byteStringAsMB("1000k") === 1)
+    assert(Utils.byteStringAsMB("1k") === ByteUnit.KB.toMB(1))
+    assert(Utils.byteStringAsMB("1g") === ByteUnit.GB.toMB(1))
+    assert(Utils.byteStringAsMB("1t") === ByteUnit.TB.toMB(1))
+    assert(Utils.byteStringAsMB("1p") === ByteUnit.PB.toMB(1))
+
+    assert(Utils.byteStringAsKB("1") === 1)
+    assert(Utils.byteStringAsKB("1k") === 1)
+    assert(Utils.byteStringAsKB("1m") === ByteUnit.MB.toKB(1))
+    assert(Utils.byteStringAsKB("1g") === ByteUnit.GB.toKB(1))
+    assert(Utils.byteStringAsKB("1t") === ByteUnit.TB.toKB(1))
+    assert(Utils.byteStringAsKB("1p") === ByteUnit.PB.toKB(1))
+    
+    assert(Utils.byteStringAsBytes("1") === 1)
+    assert(Utils.byteStringAsBytes("1k") === ByteUnit.KB.toBytes(1))
+    assert(Utils.byteStringAsBytes("1m") === ByteUnit.MB.toBytes(1))
+    assert(Utils.byteStringAsBytes("1g") === ByteUnit.GB.toBytes(1))
+    assert(Utils.byteStringAsBytes("1t") === ByteUnit.TB.toBytes(1))
+    assert(Utils.byteStringAsBytes("1p") === ByteUnit.PB.toBytes(1))
+
+    // Test invalid strings
+    intercept[NumberFormatException] {
+      Utils.byteStringAsBytes("500ub")
+    }
+    
+    // Test invalid strings
+    intercept[NumberFormatException] {
+      Utils.byteStringAsBytes("This breaks 600b")
+    }
+
+    intercept[NumberFormatException] {
+      Utils.byteStringAsBytes("This breaks 600")
+    }
+
+    intercept[NumberFormatException] {
+      Utils.byteStringAsBytes("600gb This breaks")
+    }
+
+    intercept[NumberFormatException] {
+      Utils.byteStringAsBytes("This 123mb breaks")
+    }
+  }
 
   test("bytesToString") {
     assert(Utils.bytesToString(10) === "10.0 B")

network/common/src/main/java/org/apache/spark/network/util/JavaUtils.java

@@ -22,14 +22,12 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.concurrent.TimeUnit;
-
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 import com.google.common.base.Charsets;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
-import com.sun.javafx.css.SizeUnits;
 import io.netty.buffer.Unpooled;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -142,11 +140,11 @@ private static boolean isSymlink(File file) throws IOException {
   private static ImmutableMap<String, ByteUnit> byteSuffixes =
     ImmutableMap.<String, ByteUnit>builder()
       .put("b", ByteUnit.BYTE)
-      .put("kb", ByteUnit.KB)
-      .put("mb", ByteUnit.MB)
-      .put("gb", ByteUnit.GB)
-      .put("tb", ByteUnit.TB)
-      .put("pb", ByteUnit.PB)
+      .put("k", ByteUnit.KB)
+      .put("m", ByteUnit.MB)
+      .put("g", ByteUnit.GB)
+      .put("t", ByteUnit.TB)
+      .put("p", ByteUnit.PB)
       .build();
 
   /**
@@ -273,6 +271,4 @@ public static long byteStringAsMB(String str) {
   public static long byteStringAsGB(String str) {
     return parseByteString(str, ByteUnit.GB);
   }
-
-  
 }