move relation execution into directive

Author: saimukkamala

wrangler-api/src/main/java/io/cdap/wrangler/api/Directive.java

@@ -16,6 +16,9 @@
 
 package io.cdap.wrangler.api;
 
+import io.cdap.cdap.etl.api.relational.LinearRelationalTransform;
+import io.cdap.cdap.etl.api.relational.Relation;
+import io.cdap.cdap.etl.api.relational.RelationalTranformContext;
 import io.cdap.wrangler.api.parser.UsageDefinition;
 
 import java.util.List;
@@ -51,7 +54,8 @@
  *   }
  * </code>
  */
-public interface Directive extends Executor<List<Row>, List<Row>>, EntityMetrics {
+public interface Directive extends Executor<List<Row>, List<Row>>, EntityMetrics,
+    LinearRelationalTransform {
   /**
    * This defines a interface variable that is static and final for specify
    * the {@code type} of the plugin this interface would provide.
@@ -126,4 +130,11 @@ default List<EntityCountMetric> getCountMetrics() {
     // no op
     return null;
   }
+
+  @Override
+  default Relation transform(RelationalTranformContext relationalTranformContext,
+      Relation relation) {
+    // no-op
+    return relation;
+  }
 }

wrangler-api/src/main/java/io/cdap/wrangler/api/RelationalDirective.java

@@ -16,15 +16,19 @@
 
 package io.cdap.wrangler.api;
 
+import io.cdap.cdap.etl.api.relational.LinearRelationalTransform;
+import io.cdap.cdap.etl.api.relational.Relation;
+import io.cdap.cdap.etl.api.relational.RelationalTranformContext;
+
 /**
  * Directive interface which supports Relational transformations
  */
-public interface RelationalDirective extends Directive {
-
-  /**
-   * returns sql expression
-   * @return expression
-   */
-  String getSQL();
+public interface RelationalDirective extends Directive, LinearRelationalTransform {
 
+  @Override
+  default Relation transform(RelationalTranformContext relationalTranformContext,
+      Relation relation) {
+    // no-op
+    return relation;
+  }
 }

wrangler-core/src/main/java/io/cdap/directives/column/Drop.java

@@ -19,6 +19,8 @@
 import io.cdap.cdap.api.annotation.Description;
 import io.cdap.cdap.api.annotation.Name;
 import io.cdap.cdap.api.annotation.Plugin;
+import io.cdap.cdap.etl.api.relational.Relation;
+import io.cdap.cdap.etl.api.relational.RelationalTranformContext;
 import io.cdap.wrangler.api.Arguments;
 import io.cdap.wrangler.api.Directive;
 import io.cdap.wrangler.api.DirectiveExecutionException;
@@ -44,7 +46,7 @@
 @Name(Drop.NAME)
 @Categories(categories = { "column"})
 @Description("Drop one or more columns.")
-public class Drop implements RelationalDirective, Lineage {
+public class Drop implements Directive, Lineage {
   public static final String NAME = "drop";
 
   // Columns to be dropped.
@@ -91,12 +93,11 @@ public Mutation lineage() {
   }
 
   @Override
-  public String getSQL() {
-    String sql = "DROP COLUMN ";
-    for (String col : columns) {
-      sql += col + ",";
+  public Relation transform(RelationalTranformContext relationalTranformContext,
+      Relation relation) {
+    for (String col: columns) {
+      relation = relation.dropColumn(col);
     }
-    sql = sql.substring(0, sql.length() - 1);
-    return sql;
+    return relation;
   }
 }

wrangler-transform/src/main/java/io/cdap/wrangler/Wrangler.java

@@ -589,8 +589,6 @@ && checkPreconditionNotEmpty(true))) {
     Expression filterExpression = expressionFactory.get().compile(config.getPreconditionSQL());
     Relation filteredRelation = relation.filter(filterExpression);
 
-    ExpressionFactory<String> expFactory = expressionFactory.get();
-
     String recipe = config.getDirectives();
 
     registry = SystemDirectiveRegistry.INSTANCE;
@@ -611,17 +609,8 @@ && checkPreconditionNotEmpty(true))) {
     }
 
     for (Directive directive : directives) {
-      // Expression exp = expFactory.compile(sql);
-      if (!(directive instanceof RelationalDirective)) {
-        throw new RuntimeException("Directive is not relational Directive");
-      }
-      // currently supporting only drop column
-      // SQL will be returned as "DROP COLUMN col1, col2"
-      String sql = ((RelationalDirective) directive).getSQL();
-      List<String> cols = getColumnsOfDropSQL(sql);
-      for (String col : cols) {
-        filteredRelation = filteredRelation.dropColumn(col);
-      }
+      filteredRelation = directive
+          .transform(relationalTranformContext, filteredRelation);
     }
     return filteredRelation;
   }