[SPARK-5932] Updated memoryStringToMb to use new interfaces

Author: Ilya Ganelin

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

@@ -1080,18 +1080,9 @@ private[spark] object Utils extends Logging {
    * Convert a Java memory parameter passed to -Xmx (such as 300m or 1g) to a number of megabytes.
    */
   def memoryStringToMb(str: String): Int = {
-    val lower = str.toLowerCase
-    if (lower.endsWith("k")) {
-      (lower.substring(0, lower.length-1).toLong / 1024).toInt
-    } else if (lower.endsWith("m")) {
-      lower.substring(0, lower.length-1).toInt
-    } else if (lower.endsWith("g")) {
-      lower.substring(0, lower.length-1).toInt * 1024
-    } else if (lower.endsWith("t")) {
-      lower.substring(0, lower.length-1).toInt * 1024 * 1024
-    } else {// no suffix, so it's just a number in bytes
-      (lower.toLong / 1024 / 1024).toInt
-    }
+    // Convert to bytes, rather than directly to MB, because when no units are specified the unit
+    // is assumed to be bytes
+    (JavaUtils.byteStringAsBytes(str) / 1048576.0).toInt
   }
 
   /**

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

@@ -91,16 +91,20 @@ class UtilsSuite extends FunSuite with ResetSystemProperties {
 
     assert(Utils.byteStringAsGB("1") === 1)
     assert(Utils.byteStringAsGB("1g") === 1)
-    assert(Utils.byteStringAsGB("1000m") === 1)
-    assert(Utils.byteStringAsGB("1000000k") === 1)
+    assert(Utils.byteStringAsGB("1023m") === 0)
+    assert(Utils.byteStringAsGB("1024m") === 1)
+    assert(Utils.byteStringAsGB("1048575k") === 0)
+    assert(Utils.byteStringAsGB("1048576k") === 1)
     assert(Utils.byteStringAsGB("1k") === 0)
-    assert(Utils.byteStringAsGB("1023m") === 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("1048575b") === 0)
+    assert(Utils.byteStringAsMB("1048576b") === 1)
+    assert(Utils.byteStringAsMB("1023k") === 0)
+    assert(Utils.byteStringAsMB("1024k") === 1)
     assert(Utils.byteStringAsMB("3645k") === 3)
     assert(Utils.byteStringAsMB("1g") === ByteUnit.GB.toMB(1))
     assert(Utils.byteStringAsMB("1t") === ByteUnit.TB.toMB(1))

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

@@ -42,8 +42,7 @@ public double convert(double d, ByteUnit u) {
 
   /**
    * <pre>
-   * Kilobyte (kB)
-   * 10^3 Byte = 1.000 Byte
+   * Kilobyte (kB) = 1024 Byte
    */
   KB {
     @Override
@@ -59,8 +58,7 @@ public double convert(double d, ByteUnit u) {
 
   /**
    * <pre>
-   * Megabyte (MB)
-   * 10^6 Byte = 1.000.000 Byte
+   * Megabyte (MB) = 1024 * 1024 Byte
    */
   MB {
     @Override
@@ -76,8 +74,7 @@ public double convert(double d, ByteUnit u) {
 
   /**
    * <pre>
-   * Gigabyte (GB)
-   * 10^9 Byte = 1.000.000.000 Byte
+   * Gigabyte (GB) = 1024 * 1024 * 1024 Byte
    */
   GB {
     @Override
@@ -91,8 +88,7 @@ public double toBytes(double d) {
 
   /**
    * <pre>
-   * Terabyte (TB)
-   * 10^12 Byte = 1.000.000.000.000 Byte
+   * Terabyte (TB) = 1024 * 1024 * 1024 * 1024 Byte
    */
   TB {
     @Override
@@ -108,8 +104,7 @@ public double convert(double d, ByteUnit u) {
 
   /**
    * <pre>
-   * Petabyte (PB)
-   * 10^15 Byte = 1.000.000.000.000.000 Byte
+   * Petabyte (PB) = 1024 * 1024 * 1024 * 1024 * 1024 Byte
    */
   PB {
     @Override
@@ -123,11 +118,11 @@ public double convert(double d, ByteUnit u) {
     }
   };
 
-  static final double C_KB = Math.pow(10d, 3d);
-  static final double C_MB = Math.pow(10d, 6d);
-  static final double C_GB = Math.pow(10d, 9d);
-  static final double C_TB = Math.pow(10d, 12d);
-  static final double C_PB = Math.pow(10d, 15d);
+  static final double C_KB = 1024d;
+  static final double C_MB = Math.pow(1024d, 2d);
+  static final double C_GB = Math.pow(1024d, 3d);
+  static final double C_TB = Math.pow(1024d, 4d);
+  static final double C_PB = Math.pow(1024d, 5d);
 
   private static final double MAX = Double.MAX_VALUE;