[CDAP-16290] Added FLL for ScalaSparkCompute and Sink plugins

Author: vravish

pom.xml

@@ -29,7 +29,7 @@
     <data.pipeline.parent>system:cdap-data-pipeline[6.0.0-SNAPSHOT,7.0.0-SNAPSHOT)</data.pipeline.parent>
     <data.stream.parent>system:cdap-data-streams[6.0.0-SNAPSHOT,7.0.0-SNAPSHOT)</data.stream.parent>
 
-    <cdap.version>6.0.0-SNAPSHOT</cdap.version>
+    <cdap.version>6.2.0-SNAPSHOT</cdap.version>
     <spark.version>2.1.3</spark.version>
     <logback.version>1.0.9</logback.version>
 
@@ -100,6 +100,11 @@
       <version>4.11</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>io.cdap.plugin</groupId>
+      <artifactId>hydrator-common</artifactId>
+      <version>2.4.0-SNAPSHOT</version>
+    </dependency>
     <dependency>
       <groupId>io.cdap.cdap</groupId>
       <artifactId>hydrator-test</artifactId>

src/main/java/io/cdap/plugin/spark/dynamic/ScalaSparkCompute.java

@@ -28,6 +28,8 @@
 import io.cdap.cdap.etl.api.StageConfigurer;
 import io.cdap.cdap.etl.api.batch.SparkCompute;
 import io.cdap.cdap.etl.api.batch.SparkExecutionPluginContext;
+import io.cdap.cdap.etl.api.batch.SparkPluginContext;
+import io.cdap.plugin.common.TransformLineageRecorderUtils;
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.api.java.function.Function;
 import org.apache.spark.rdd.RDD;
@@ -82,6 +84,16 @@ public void initialize(SparkExecutionPluginContext context) throws Exception {
     isRDD = !codeExecutor.isDataFrame();
   }
 
+  @Override
+  public void prepareRun(SparkPluginContext context) throws Exception {
+    super.prepareRun(context);
+    Schema outSchema = config.schema == null ? context.getInputSchema() : Schema.parseJson(config.schema);
+    context.record(TransformLineageRecorderUtils.generateManyToMany(
+      TransformLineageRecorderUtils.getFields(context.getInputSchema()),
+      TransformLineageRecorderUtils.getFields(outSchema), "sparkCompute",
+      "Transformed fields according to spark computation"));
+  }
+
   @Override
   public JavaRDD<StructuredRecord> transform(SparkExecutionPluginContext context,
                                              JavaRDD<StructuredRecord> javaRDD) throws Exception {

src/main/java/io/cdap/plugin/spark/dynamic/ScalaSparkSink.java

@@ -21,16 +21,20 @@
 import io.cdap.cdap.api.annotation.Name;
 import io.cdap.cdap.api.annotation.Plugin;
 import io.cdap.cdap.api.data.format.StructuredRecord;
+import io.cdap.cdap.api.data.schema.Schema;
 import io.cdap.cdap.api.plugin.PluginConfig;
 import io.cdap.cdap.etl.api.PipelineConfigurer;
 import io.cdap.cdap.etl.api.StageConfigurer;
 import io.cdap.cdap.etl.api.batch.SparkExecutionPluginContext;
 import io.cdap.cdap.etl.api.batch.SparkPluginContext;
 import io.cdap.cdap.etl.api.batch.SparkSink;
+import io.cdap.plugin.common.LineageRecorder;
 import org.apache.spark.api.java.JavaRDD;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.List;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
 /**
@@ -64,7 +68,10 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) throws Ille
 
   @Override
   public void prepareRun(SparkPluginContext sparkPluginContext) throws Exception {
-    // no-op
+    Schema schema = sparkPluginContext.getInputSchema();
+    if (schema != null && schema.getFields() != null) {
+      recordLineage(sparkPluginContext, config.referenceName, schema, "Write", "Wrote to Scala Spark Sink.");
+    }
   }
 
   @Override
@@ -109,11 +116,16 @@ public static final class Config extends PluginConfig {
     @Nullable
     private final Boolean deployCompile;
 
-    public Config(String scalaCode, @Nullable String dependencies,
-                  @Nullable Boolean deployCompile) {
+    @Name("referenceName")
+    @Description("This will be used to uniquely identify this source/sink for lineage, annotating metadata, etc.")
+    public String referenceName;
+
+    public Config(String scalaCode, @Nullable String dependencies, @Nullable Boolean deployCompile,
+                  String referenceName) {
       this.scalaCode = scalaCode;
       this.dependencies = dependencies;
       this.deployCompile = deployCompile;
+      this.referenceName = referenceName;
     }
 
     public String getScalaCode() {
@@ -129,5 +141,19 @@ public String getDependencies() {
     public Boolean getDeployCompile() {
       return deployCompile;
     }
+
+    public String getReferenceName() {
+      return referenceName;
+    }
+  }
+
+  private void recordLineage(SparkPluginContext context, String outputName, Schema tableSchema, String operationName,
+                             String description) {
+    LineageRecorder lineageRecorder = new LineageRecorder(context, outputName);
+    lineageRecorder.createExternalDataset(tableSchema);
+    List<String> fieldNames = tableSchema.getFields().stream().map(Schema.Field::getName).collect(Collectors.toList());
+    if (!fieldNames.isEmpty()) {
+      lineageRecorder.recordWrite(operationName, description, fieldNames);
+    }
   }
 }

src/test/java/io/cdap/plugin/spark/dynamic/ScalaSparkTest.java

@@ -418,7 +418,9 @@ private void testWordCountSink(String code, File outputFolder) throws Exception
     ETLBatchConfig etlConfig = ETLBatchConfig.builder("* * * * *")
       .addStage(new ETLStage("source", MockSource.getPlugin(inputTable, inputSchema)))
       .addStage(new ETLStage("sink", new ETLPlugin("ScalaSparkSink", SparkSink.PLUGIN_TYPE,
-                                                   ImmutableMap.of("scalaCode", code))))
+                                                   ImmutableMap.of(
+                                                     "scalaCode", code,
+                                                     "referenceName", "sink"))))
       .addConnection("source", "sink")
       .build();
 

widgets/ScalaSparkSink-sparksink.json

@@ -7,6 +7,14 @@
     {
       "label": "Spark Scala",
       "properties": [
+        {
+          "widget-type": "textbox",
+          "label": "Reference Name",
+          "name": "referenceName",
+          "widget-attributes": {
+            "placeholder": "Name used to identify this source for lineage"
+          }
+        },
         {
           "widget-type": "scala-editor",
           "label": "Scala",
@@ -31,5 +39,12 @@
         }
       ]
     }
-  ]
+  ],
+  "jump-config": {
+    "datasets": [
+      {
+        "ref-property-name": "referenceName"
+      }
+    ]
+  }
 }