Relational directives

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,
+    DirectiveRelationalTransform {
   /**
    * 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,5 @@ default List<EntityCountMetric> getCountMetrics() {
     // no op
     return null;
   }
+
 }

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright © 2023 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;
+
+import io.cdap.cdap.etl.api.relational.InvalidRelation;
+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.cdap.etl.api.relational.RelationalTransform;
+
+/**
+ * {@link DirectiveRelationalTransform} provides relational transform support for
+ * wrangler directives.
+ */
+public interface DirectiveRelationalTransform extends LinearRelationalTransform {
+
+    /**
+     * Implementation of linear relational transform for each supported directive.
+     *
+     * @param relationalTranformContext transformation context with engine, input and output parameters
+     * @param relation input relation upon which the transformation is applied.
+     * @return transformed relation as the output relation. By default, returns an Invalid relation
+     * for unsupported directives.
+     */
+    default Relation transform(RelationalTranformContext relationalTranformContext,
+                               Relation relation) {
+    return new InvalidRelation("SQL execution for the directive is currently not supported.");
+    }
+
+    /**
+     * Indicates whether the directive is supported by relational transformation or not.
+     *
+     * @return boolean value for the directive SQL support.
+     * By default, returns false, indicating that the directive is currently not supported.
+     */
+    default boolean isSQLSupported() {
+        return false;
+    }
+
+}

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

@@ -0,0 +1,34 @@
+/*
+ *  Copyright © 2017-2019 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;
+
+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, LinearRelationalTransform {
+
+  @Override
+  default Relation transform(RelationalTranformContext relationalTranformContext,
+      Relation relation) {
+    // no-op
+    return relation;
+  }
+}

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

@@ -19,6 +19,10 @@
 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.ExpressionFactory;
+import io.cdap.cdap.etl.api.relational.InvalidRelation;
+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;
@@ -28,11 +32,11 @@
 import io.cdap.wrangler.api.Row;
 import io.cdap.wrangler.api.annotations.Categories;
 import io.cdap.wrangler.api.lineage.Lineage;
-import io.cdap.wrangler.api.lineage.Many;
 import io.cdap.wrangler.api.lineage.Mutation;
 import io.cdap.wrangler.api.parser.ColumnName;
 import io.cdap.wrangler.api.parser.TokenType;
 import io.cdap.wrangler.api.parser.UsageDefinition;
+import io.cdap.wrangler.utils.SqlExpressionGenerator;
 
 import java.util.List;
 
@@ -110,4 +114,20 @@ public Mutation lineage() {
       .conditional(source.value(), destination.value())
       .build();
   }
+  @Override
+  public Relation transform(RelationalTranformContext relationalTranformContext,
+                            Relation relation) {
+    java.util.Optional<ExpressionFactory<String>> expressionFactory = SqlExpressionGenerator
+            .getExpressionFactory(relationalTranformContext);
+    if (!expressionFactory.isPresent()) {
+      return new InvalidRelation("Cannot find an Expression Factory");
+    }
+    return relation.setColumn(destination.value(), expressionFactory.get().compile(source.value()));
+  }
+
+  @Override
+  public boolean isSQLSupported() {
+    return true;
+  }
+
 }

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

@@ -19,6 +19,10 @@
 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.ExpressionFactory;
+import io.cdap.cdap.etl.api.relational.InvalidRelation;
+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;
@@ -32,6 +36,7 @@
 import io.cdap.wrangler.api.parser.ColumnNameList;
 import io.cdap.wrangler.api.parser.TokenType;
 import io.cdap.wrangler.api.parser.UsageDefinition;
+import io.cdap.wrangler.utils.SqlExpressionGenerator;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -101,4 +106,18 @@ public Mutation lineage() {
       .relation(Many.columns(columns), targetColumn)
       .build();
   }
+
+  @Override
+  public Relation transform(RelationalTranformContext relationalTranformContext,
+                            Relation relation) {
+    java.util.Optional<ExpressionFactory<String>> expressionFactory = SqlExpressionGenerator
+            .getExpressionFactory(relationalTranformContext);
+    if (!expressionFactory.isPresent()) {
+      return new InvalidRelation("Cannot find an Expression Factory");
+    }
+    String getColumnString = String.join(",", columns);
+    return relation.setColumn(targetColumn, expressionFactory.get().compile(String
+            .format("struct(%s)", getColumnString)));
+  }
+
 }

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

@@ -19,11 +19,14 @@
 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;
 import io.cdap.wrangler.api.DirectiveParseException;
 import io.cdap.wrangler.api.ExecutorContext;
+import io.cdap.wrangler.api.RelationalDirective;
 import io.cdap.wrangler.api.Row;
 import io.cdap.wrangler.api.annotations.Categories;
 import io.cdap.wrangler.api.lineage.Lineage;
@@ -88,4 +91,18 @@ public Mutation lineage() {
       .drop(Many.of(columns))
       .build();
   }
+
+  @Override
+  public Relation transform(RelationalTranformContext relationalTranformContext,
+      Relation relation) {
+    for (String col: columns) {
+      relation = relation.dropColumn(col);
+    }
+    return relation;
+  }
+
+  @Override
+  public boolean isSQLSupported() {
+    return true;
+  }
 }

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

@@ -19,6 +19,11 @@
 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.Expression;
+import io.cdap.cdap.etl.api.relational.ExpressionFactory;
+import io.cdap.cdap.etl.api.relational.InvalidRelation;
+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;
@@ -32,9 +37,12 @@
 import io.cdap.wrangler.api.parser.ColumnNameList;
 import io.cdap.wrangler.api.parser.TokenType;
 import io.cdap.wrangler.api.parser.UsageDefinition;
+import io.cdap.wrangler.utils.SqlExpressionGenerator;
 
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 
 /**
@@ -93,4 +101,22 @@ public Mutation lineage() {
     keep.forEach(column -> builder.relation(column, column));
     return builder.build();
   }
+  @Override
+  public Relation transform(RelationalTranformContext relationalTranformContext,
+                            Relation relation) {
+    Optional<ExpressionFactory<String>> expressionFactory = SqlExpressionGenerator
+            .getExpressionFactory(relationalTranformContext);
+    if (!expressionFactory.isPresent()) {
+      return new InvalidRelation("Cannot find an Expression Factory");
+    }
+    Map<String, Expression> keepCol = SqlExpressionGenerator
+            .generateColumnExpMap(keep, expressionFactory.get());
+    return relation.select(keepCol);
+  }
+
+  @Override
+  public boolean isSQLSupported() {
+    return true;
+  }
+
 }

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

@@ -19,6 +19,10 @@
 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.ExpressionFactory;
+import io.cdap.cdap.etl.api.relational.InvalidRelation;
+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;
@@ -33,10 +37,12 @@
 import io.cdap.wrangler.api.parser.Text;
 import io.cdap.wrangler.api.parser.TokenType;
 import io.cdap.wrangler.api.parser.UsageDefinition;
+import io.cdap.wrangler.utils.SqlExpressionGenerator;
 import org.apache.commons.lang3.StringEscapeUtils;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 
 /**
  * A directive for merging two columns and creates a third column.
@@ -108,4 +114,20 @@ public Mutation lineage() {
       .relation(Many.columns(col1, col2), Many.of(col1, col2, dest))
       .build();
   }
+  public Relation transform(RelationalTranformContext relationalTranformContext,
+                            Relation relation) {
+    Optional<ExpressionFactory<String>> expressionFactory = SqlExpressionGenerator
+            .getExpressionFactory(relationalTranformContext);
+    if (!expressionFactory.isPresent()) {
+      return new InvalidRelation("Cannot find an Expression Factory");
+    }
+    return relation.setColumn(dest, expressionFactory.get()
+            .compile(String.format("CONCAT(%s,'%s',%s)", col1, delimiter, col2)));
+  }
+
+  @Override
+  public boolean isSQLSupported() {
+    return true;
+  }
+
 }

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

@@ -19,6 +19,10 @@
 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.ExpressionFactory;
+import io.cdap.cdap.etl.api.relational.InvalidRelation;
+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;
@@ -31,8 +35,10 @@
 import io.cdap.wrangler.api.parser.TokenType;
 import io.cdap.wrangler.api.parser.UsageDefinition;
 import io.cdap.wrangler.utils.ColumnConverter;
+import io.cdap.wrangler.utils.SqlExpressionGenerator;
 
 import java.util.List;
+import java.util.Optional;
 
 /**
  * A directive for renaming columns.
@@ -82,4 +88,21 @@ public Mutation lineage() {
       .relation(source, target)
       .build();
   }
+  @Override
+  public Relation transform(RelationalTranformContext relationalTranformContext,
+                            Relation relation) {
+    Optional<ExpressionFactory<String>> expressionFactory = SqlExpressionGenerator
+            .getExpressionFactory(relationalTranformContext);
+    if (!expressionFactory.isPresent()) {
+      return new InvalidRelation("Cannot find an Expression Factory");
+    }
+    relation = relation.setColumn(target.value(), expressionFactory.get().compile(source.value()));
+    return relation.dropColumn(source.value());
+  }
+
+  @Override
+  public boolean isSQLSupported() {
+    return true;
+  }
+
 }