[SPARK-5932] Added byte string conversion to Jav utils

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

@@ -0,0 +1,172 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.network.util;
+
+/**
+ * Code based on https://github.com/fabian-barney/Utils
+ * Copyright 2011 Fabian Barney
+ *
+ * @author Fabian Barney
+ */
+public enum ByteUnit {
+  /**
+   * <pre>
+   * Byte (B)
+   * 1 Byte
+   */
+  BYTE {
+    @Override
+    public double toBytes(double d) {
+      return d;
+    }
+
+    @Override
+    public double convert(double d, ByteUnit u) {
+      return u.toBytes(d);
+    }
+  },
+
+  /**
+   * <pre>
+   * Kilobyte (kB)
+   * 10^3 Byte = 1.000 Byte
+   */
+  KB {
+    @Override
+    public double toBytes(double d) {
+      return safeMulti(d, C_KB);
+    }
+
+    @Override
+    public double convert(double d, ByteUnit u) {
+      return u.toKB(d);
+    }
+  },
+
+  /**
+   * <pre>
+   * Megabyte (MB)
+   * 10^6 Byte = 1.000.000 Byte
+   */
+  MB {
+    @Override
+    public double toBytes(double d) {
+      return safeMulti(d, C_MB);
+    }
+
+    @Override
+    public double convert(double d, ByteUnit u) {
+      return u.toMB(d);
+    }
+  },
+
+  /**
+   * <pre>
+   * Gigabyte (GB)
+   * 10^9 Byte = 1.000.000.000 Byte
+   */
+  GB {
+    @Override
+    public double toBytes(double d) {
+      return safeMulti(d, C_GB);
+    }
+
+    @Override
+    public double convert(double d, ByteUnit u) {
+      return u.toGB(d);
+    }
+  },
+
+  /**
+   * <pre>
+   * Terabyte (TB)
+   * 10^12 Byte = 1.000.000.000.000 Byte
+   */
+  TB {
+    @Override
+    public double toBytes(double d) {
+      return safeMulti(d, C_TB);
+    }
+
+    @Override
+    public double convert(double d, ByteUnit u) {
+      return u.toTB(d);
+    }
+  },
+
+  /**
+   * <pre>
+   * Petabyte (PB)
+   * 10^15 Byte = 1.000.000.000.000.000 Byte
+   */
+  PB {
+    @Override
+    public double toBytes(double d) {
+      return safeMulti(d, C_PB);
+    }
+
+    @Override
+    public double convert(double d, ByteUnit u) {
+      return u.toPB(d);
+    }
+  };
+
+  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);
+
+  private static final double MAX = Double.MAX_VALUE;
+
+  static double safeMulti(double d, double multi) {
+    double limit = MAX / multi;
+
+    if (d > limit) {
+      return Double.MAX_VALUE;
+    }
+    if (d < -limit) {
+      return Double.MIN_VALUE;
+    }
+
+    return d * multi;
+  }
+
+  public abstract double toBytes(double d);
+
+  public final double toKB(double d) {
+    return toBytes(d) / C_KB;
+  }
+
+  public final double toMB(double d) {
+    return toBytes(d) / C_MB;
+  }
+
+  public final double toGB(double d) {
+    return toBytes(d) / C_GB;
+  }
+
+  public final double toTB(double d) {
+    return toBytes(d) / C_TB;
+  }
+
+  public final double toPB(double d) {
+    return toBytes(d) / C_PB;
+  }
+
+  public abstract double convert(double d, ByteUnit u);
+}

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

@@ -22,12 +22,14 @@
 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;
@@ -137,6 +139,16 @@ private static boolean isSymlink(File file) throws IOException {
       .put("d", TimeUnit.DAYS)
       .build();
 
+  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)
+      .build();
+
   /**
    * Convert a passed time string (e.g. 50s, 100ms, or 250us) to a time count for
    * internal use. If no suffix is provided a direct conversion is attempted.
@@ -164,7 +176,7 @@ private static long parseTimeString(String str, TimeUnit unit) {
       return unit.convert(val, suffix != null ? timeSuffixes.get(suffix) : unit);
     } catch (NumberFormatException e) {
       String timeError = "Time must be specified as seconds (s), " +
-              "milliseconds (ms), microseconds (us), minutes (m or min) hour (h), or day (d). " +
+              "milliseconds (ms), microseconds (us), minutes (m or min), hour (h), or day (d). " +
               "E.g. 50s, 100ms, or 250us.";
 
       throw new NumberFormatException(timeError + "\n" + e.getMessage());
@@ -186,5 +198,81 @@ public static long timeStringAsMs(String str) {
   public static long timeStringAsSec(String str) {
     return parseTimeString(str, TimeUnit.SECONDS);
   }
+
+  /**
+   * Convert a passed byte string (e.g. 50b, 100kb, or 250mb) to a ByteUnit for
+   * internal use. If no suffix is provided a direct conversion of the provided default is 
+   * attempted.
+   */
+  private static long parseByteString(String str, ByteUnit unit) {
+    String lower = str.toLowerCase().trim();
+
+    try {
+      String suffix;
+      long val;
+      Matcher m = Pattern.compile("([0-9]+)([a-z]+)?").matcher(lower);
+      if (m.matches()) {
+        val = Long.parseLong(m.group(1));
+        suffix = m.group(2);
+      } else {
+        throw new NumberFormatException("Failed to parse byte string: " + str);
+      }
+
+      // Check for invalid suffixes
+      if (suffix != null && !byteSuffixes.containsKey(suffix)) {
+        throw new NumberFormatException("Invalid suffix: \"" + suffix + "\"");
+      }
+
+      // If suffix is valid use that, otherwise none was provided and use the default passed
+      return (long) unit.convert(val, suffix != null ? byteSuffixes.get(suffix) : unit);
+    } catch (NumberFormatException e) {
+      String timeError = "Size must be specified as bytes (b), " +
+        "kilobytes (kb), megabytes (mb), gigabytes (gb), terabytes (tb), or petabytes(pb). " +
+        "E.g. 50b, 100kb, or 250mb.";
+
+      throw new NumberFormatException(timeError + "\n" + e.getMessage());
+    }
+  }
+
+  /**
+   * 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.
+   */
+  public static long byteStringAsBytes(String str) {
+    return parseByteString(str, ByteUnit.BYTE);
+  }
+
+  /**
+   * 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.
+   */
+  public static long byteStringAsKB(String str) {
+    return parseByteString(str, ByteUnit.KB);
+  }
+
+  /**
+   * 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.
+   */
+  public static long byteStringAsMB(String str) {
+    return parseByteString(str, ByteUnit.MB);
+  }
 
+  /**
+   * Convert a passed byte string (e.g. 50b, 100kb, or 250mb) to gigabytes for
+   * internal use.
+   *
+   * If no suffix is provided, the passed number is assumed to be in gigabytes.
+   */
+  public static long byteStringAsGB(String str) {
+    return parseByteString(str, ByteUnit.GB);
+  }
+
+
 }