Assignment: Enhance Wrangler with Byte Size and Time Duration Units

wrangler-api/src/main/java/io/cdap/wrangler/api/parser/ByteSize.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright © 2024 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.wrangler.api.parser;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import io.cdap.wrangler.api.annotations.PublicEvolving;
+
+/**
+ * Represents a byte size value with units (KB, MB, GB, TB).
+ */
+@PublicEvolving
+public class ByteSize implements Token {
+  private final String value;
+  private final long bytes;
+
+  public ByteSize(String value) {
+    this.value = value;
+    this.bytes = parseBytes(value);
+  }
+
+  private static long parseBytes(String value) {
+    String number = value.replaceAll("[^0-9.]", "");
+    String unit = value.replaceAll("[0-9.]", "").toLowerCase();
+    double size = Double.parseDouble(number);
+
+    switch (unit) {
+      case "kb":
+        return (long) (size * 1024);
+      case "mb":
+        return (long) (size * 1024 * 1024); 
+      case "gb":
+        return (long) (size * 1024 * 1024 * 1024);
+      case "tb":
+        return (long) (size * 1024L * 1024L * 1024L * 1024L);
+      default:
+        return (long) size; // Base unit bytes
+    }
+  }
+
+  @Override
+  public String value() {
+    return value;
+  }
+
+  public long getBytes() {
+    return bytes;
+  }
+
+  @Override
+  public TokenType type() {
+    return TokenType.BYTE_SIZE;
+  }
+
+  @Override
+  public JsonElement toJson() {
+    JsonObject object = new JsonObject();
+    object.addProperty("type", TokenType.BYTE_SIZE.name());
+    object.addProperty("value", value);
+    object.addProperty("bytes", bytes);
+    return object;
+  }
+}

wrangler-api/src/main/java/io/cdap/wrangler/api/parser/TimeDuration.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright © 2024 Cask Data, Inc.
+ *
+ * Licensed 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 io.cdap.wrangler.api.parser;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonObject;
+import io.cdap.wrangler.api.annotations.PublicEvolving;
+
+/**
+ * Represents a time duration value with units (ms, s, m, h, d).
+ */
+@PublicEvolving
+public class TimeDuration implements Token {
+  private final String value;
+  private final long milliseconds;
+
+  public TimeDuration(String value) {
+    this.value = value;
+    this.milliseconds = parseMilliseconds(value);
+  }
+
+  private static long parseMilliseconds(String value) {
+    String number = value.replaceAll("[^0-9.]", "");
+    String unit = value.replaceAll("[0-9.]", "").toLowerCase();
+    double duration = Double.parseDouble(number);
+
+    switch (unit) {
+      case "ms":
+        return (long) duration;
+      case "s":
+        return (long) (duration * 1000);
+      case "m":
+        return (long) (duration * 60 * 1000);
+      case "h":
+        return (long) (duration * 60 * 60 * 1000);
+      case "d":
+        return (long) (duration * 24 * 60 * 60 * 1000);
+      default:
+        return (long) duration; // Default to milliseconds
+    }
+  }
+
+  @Override
+  public String value() {
+    return value;
+  }
+
+  public long getMilliseconds() {
+    return milliseconds;
+  }
+
+  @Override
+  public TokenType type() {
+    return TokenType.TIME_DURATION;
+  }
+
+  @Override
+  public JsonElement toJson() {
+    JsonObject object = new JsonObject();
+    object.addProperty("type", TokenType.TIME_DURATION.name());
+    object.addProperty("value", value);
+    object.addProperty("milliseconds", milliseconds);
+    return object;
+  }
+}

wrangler-api/src/main/java/io/cdap/wrangler/api/parser/TokenType.java

@@ -148,6 +148,18 @@ public enum TokenType implements Serializable {
    */
   RANGES,
 
+  /**
+   * Represents the enumerated type for byte size values with units (KB, MB, GB etc).
+   * E.g. "10KB", "1.5MB", "2GB"
+   */
+  BYTE_SIZE,
+
+  /**
+   * Represents the enumerated type for time duration values with units (ms, s, m, h, d).
+   * E.g. "100ms", "5s", "2h"
+   */
+  TIME_DURATION,
+
   /**
    * Represents the enumerated type for the object of type {@code String} with restrictions
    * on characters that can be present in a string.

wrangler-core/src/main/antlr4/io/cdap/wrangler/parser/Directives.g4

@@ -64,6 +64,8 @@ directive
     | stringList
     | numberRanges
     | properties
+    | byteSize     // New rule for byte size
+    | timeDuration // New rule for time duration
   )*?
   ;
 
@@ -195,6 +197,13 @@ identifierList
  : Identifier (',' Identifier)*
  ;
 
+byteSize
+ : ByteSize
+ ;
+
+timeDuration
+ : TimeDuration
+ ;
 
 /*
  * Following are the Lexer Rules used for tokenizing the recipe.
@@ -247,7 +256,6 @@ BackSlash: '\\';
 Dollar   : '$';
 Tilde    : '~';
 
-
 Bool
  : 'true'
  | 'false'
@@ -257,6 +265,29 @@ Number
  : Int ('.' Digit*)?
  ;
 
+ByteSize
+ : Int ('.' Digit*)? ByteUnit
+ ;
+
+TimeDuration
+ : Int ('.' Digit*)? TimeUnit
+ ;
+
+fragment ByteUnit
+ : [kK][bB]    // KB
+ | [mM][bB]    // MB
+ | [gG][bB]    // GB
+ | [tT][bB]    // TB
+ ;
+
+fragment TimeUnit
+ : 'ms'        // milliseconds
+ | 's'         // seconds 
+ | 'm'         // minutes
+ | 'h'         // hours
+ | 'd'         // days
+ ;
+
 Identifier
  : [a-zA-Z_\-] [a-zA-Z_0-9\-]*
  ;