fix(ingestion/spark): Platform instance and column level lineage fix (#…

…10843)

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

metadata-ingestion/docs/sources/databricks/README.md

@@ -11,7 +11,7 @@ The alternative way to integrate is via the Hive connector. The [Hive starter re
 
 ## Databricks Spark
 
-To complete the picture, we recommend adding push-based ingestion from your Spark jobs to see real-time activity and lineage between your Databricks tables and your Spark jobs. Use the Spark agent to push metadata to DataHub using the instructions [here](../../../../metadata-integration/java/spark-lineage/README.md#configuration-instructions-databricks).
+To complete the picture, we recommend adding push-based ingestion from your Spark jobs to see real-time activity and lineage between your Databricks tables and your Spark jobs. Use the Spark agent to push metadata to DataHub using the instructions [here](../../../../metadata-integration/java/spark-lineage-beta/README.md#configuration-instructions-databricks).
 
 ## Watch the DataHub Talk at the Data and AI Summit 2022
 

metadata-integration/java/openlineage-converter/src/main/java/io/datahubproject/openlineage/config/DatahubOpenlineageConfig.java

@@ -33,6 +33,7 @@ public class DatahubOpenlineageConfig {
   @Builder.Default private String hivePlatformAlias = "hive";
   @Builder.Default private Map<String, String> urnAliases = new HashMap<>();
   @Builder.Default private final boolean disableSymlinkResolution = false;
+  @Builder.Default private final boolean lowerCaseDatasetUrns = false;
 
   public List<PathSpec> getPathSpecsForPlatform(String platform) {
     if ((pathSpecs == null) || (pathSpecs.isEmpty())) {

metadata-integration/java/openlineage-converter/src/main/java/io/datahubproject/openlineage/converter/OpenLineageToDataHub.java

@@ -4,6 +4,7 @@
 import com.linkedin.common.DataPlatformInstance;
 import com.linkedin.common.Edge;
 import com.linkedin.common.EdgeArray;
+import com.linkedin.common.FabricType;
 import com.linkedin.common.GlobalTags;
 import com.linkedin.common.Owner;
 import com.linkedin.common.OwnerArray;
@@ -57,6 +58,8 @@
 import io.datahubproject.openlineage.dataset.DatahubJob;
 import io.datahubproject.openlineage.dataset.HdfsPathDataset;
 import io.datahubproject.openlineage.dataset.HdfsPlatform;
+import io.datahubproject.openlineage.dataset.PathSpec;
+import io.datahubproject.openlineage.utils.DatahubUtils;
 import io.openlineage.client.OpenLineage;
 import io.openlineage.client.OpenLineageClientUtils;
 import java.io.IOException;
@@ -151,6 +154,11 @@ public static Optional<DatasetUrn> convertOpenlineageDatasetToDatasetUrn(
   private static Optional<DatasetUrn> getDatasetUrnFromOlDataset(
       String namespace, String datasetName, DatahubOpenlineageConfig mappingConfig) {
     String platform;
+    if (mappingConfig.isLowerCaseDatasetUrns()) {
+      namespace = namespace.toLowerCase();
+      datasetName = datasetName.toLowerCase();
+    }
+
     if (namespace.contains(SCHEME_SEPARATOR)) {
       try {
         URI datasetUri;
@@ -183,12 +191,45 @@ private static Optional<DatasetUrn> getDatasetUrnFromOlDataset(
       platform = namespace;
     }
 
-    if (mappingConfig.getCommonDatasetPlatformInstance() != null) {
-      datasetName = mappingConfig.getCommonDatasetPlatformInstance() + "." + datasetName;
+    String platformInstance = getPlatformInstance(mappingConfig, platform);
+    FabricType env = getEnv(mappingConfig, platform);
+    return Optional.of(DatahubUtils.createDatasetUrn(platform, platformInstance, datasetName, env));
+  }
+
+  private static FabricType getEnv(DatahubOpenlineageConfig mappingConfig, String platform) {
+    FabricType fabricType = mappingConfig.getFabricType();
+    if (mappingConfig.getPathSpecs() != null
+        && mappingConfig.getPathSpecs().containsKey(platform)) {
+      List<PathSpec> path_specs = mappingConfig.getPathSpecs().get(platform);
+      for (PathSpec pathSpec : path_specs) {
+        if (pathSpec.getEnv().isPresent()) {
+          try {
+            fabricType = FabricType.valueOf(pathSpec.getEnv().get());
+            return fabricType;
+          } catch (IllegalArgumentException e) {
+            log.warn("Invalid environment value: {}", pathSpec.getEnv());
+          }
+        }
+      }
     }
+    return fabricType;
+  }
 
-    return Optional.of(
-        new DatasetUrn(new DataPlatformUrn(platform), datasetName, mappingConfig.getFabricType()));
+  private static String getPlatformInstance(
+      DatahubOpenlineageConfig mappingConfig, String platform) {
+    // Use the platform instance from the path spec if it is present otherwise use the one from the
+    // commonDatasetPlatformInstance
+    String platformInstance = mappingConfig.getCommonDatasetPlatformInstance();
+    if (mappingConfig.getPathSpecs() != null
+        && mappingConfig.getPathSpecs().containsKey(platform)) {
+      List<PathSpec> path_specs = mappingConfig.getPathSpecs().get(platform);
+      for (PathSpec pathSpec : path_specs) {
+        if (pathSpec.getPlatformInstance().isPresent()) {
+          return pathSpec.getPlatformInstance().get();
+        }
+      }
+    }
+    return platformInstance;
   }
 
   public static GlobalTags generateTags(List<String> tags) {

metadata-integration/java/openlineage-converter/src/main/java/io/datahubproject/openlineage/dataset/DatahubJob.java

@@ -280,11 +280,11 @@ private Pair<UrnArray, EdgeArray> processDownstreams(
                     for (Urn downstream :
                         Objects.requireNonNull(fineGrainedLineage.getDownstreams())) {
                       upstreamLineagePatchBuilder.addFineGrainedUpstreamField(
-                          downstream,
+                          upstream,
                           fineGrainedLineage.getConfidenceScore(),
                           StringUtils.defaultIfEmpty(
                               fineGrainedLineage.getTransformOperation(), "TRANSFORM"),
-                          upstream,
+                          downstream,
                           null);
                     }
                   }

metadata-integration/java/openlineage-converter/src/main/java/io/datahubproject/openlineage/dataset/PathSpec.java

@@ -14,7 +14,7 @@
 public class PathSpec {
   final String alias;
   final String platform;
-  @Builder.Default final String env = "PROD";
+  @Builder.Default final Optional<String> env = Optional.empty();
   final List<String> pathSpecList;
   @Builder.Default final Optional<String> platformInstance = Optional.empty();
 }

metadata-integration/java/spark-lineage-beta/README.md

@@ -24,15 +24,15 @@ When running jobs using spark-submit, the agent needs to be configured in the co
 
 ```text
 #Configuring DataHub spark agent jar
-spark.jars.packages                          io.acryl:acryl-spark-lineage:0.2.11
+spark.jars.packages                          io.acryl:acryl-spark-lineage:0.2.13
 spark.extraListeners                         datahub.spark.DatahubSparkListener
 spark.datahub.rest.server                    http://localhost:8080
 ```
 
 ## spark-submit command line
 
 ```sh
-spark-submit --packages io.acryl:acryl-spark-lineage:0.2.11 --conf "spark.extraListeners=datahub.spark.DatahubSparkListener" my_spark_job_to_run.py
+spark-submit --packages io.acryl:acryl-spark-lineage:0.2.13 --conf "spark.extraListeners=datahub.spark.DatahubSparkListener" my_spark_job_to_run.py
 ```
 
 ### Configuration Instructions:  Amazon EMR
@@ -41,7 +41,7 @@ Set the following spark-defaults configuration properties as it
 stated [here](https://docs.aws.amazon.com/emr/latest/ReleaseGuide/emr-spark-configure.html)
 
 ```text
-spark.jars.packages                          io.acryl:acryl-spark-lineage:0.2.11
+spark.jars.packages                          io.acryl:acryl-spark-lineage:0.2.13
 spark.extraListeners                         datahub.spark.DatahubSparkListener
 spark.datahub.rest.server                    https://your_datahub_host/gms
 #If you have authentication set up then you also need to specify the Datahub access token
@@ -56,7 +56,7 @@ When running interactive jobs from a notebook, the listener can be configured wh
 spark = SparkSession.builder
 .master("spark://spark-master:7077")
 .appName("test-application")
-.config("spark.jars.packages", "io.acryl:acryl-spark-lineage:0.2.11")
+.config("spark.jars.packages", "io.acryl:acryl-spark-lineage:0.2.13")
 .config("spark.extraListeners", "datahub.spark.DatahubSparkListener")
 .config("spark.datahub.rest.server", "http://localhost:8080")
 .enableHiveSupport()
@@ -79,7 +79,7 @@ appName("test-application")
 config("spark.master","spark://spark-master:7077")
         .
 
-config("spark.jars.packages","io.acryl:acryl-spark-lineage:0.2.11")
+config("spark.jars.packages","io.acryl:acryl-spark-lineage:0.2.13")
         .
 
 config("spark.extraListeners","datahub.spark.DatahubSparkListener")
@@ -159,7 +159,7 @@ information like tokens.
 
 | Field                                                               | Required | Default | Description                                                                                                                                                                               |
 |---------------------------------------------------------------------|----------|---------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| spark.jars.packages                                                 | ✅        |         | Set with latest/required version  io.acryl:datahub-spark-lineage:0.8.23                                                                                                                   |
+| spark.jars.packages                                                 | ✅        |         | Set with latest/required version  io.acryl:acryl-spark-lineage:0.2.13                                                                                                                     |
 | spark.extraListeners                                                | ✅        |         | datahub.spark.DatahubSparkListener                                                                                                                                                        |
 | spark.datahub.rest.server                                           | ✅        |         | Datahub server url  eg:<http://localhost:8080>                                                                                                                                            |
 | spark.datahub.rest.token                                            |          |         | Authentication token.                                                                                                                                                                     |
@@ -181,9 +181,10 @@ information like tokens.
 | spark.datahub.partition_regexp_pattern                              |          |         | Strip partition part from the path if path end matches with the specified regexp. Example `year=.*/month=.*/day=.*`                                                                       |
 | spark.datahub.tags                                                  |          |         | Comma separated list of tags to attach to the DataFlow                                                                                                                                    |
 | spark.datahub.domains                                               |          |         | Comma separated list of domain urns to attach to the DataFlow                                                                                                                             |
-| spark.datahub.stage_metadata_coalescing                             |          |         | Normally it coalesce and send metadata at the onApplicationEnd event which is never called on Databricks or on Glue. You should enable this on Databricks if you want coalesced run .     |
-| spark.datahub.patch.enabled                                         |          | false   | Set this to true to send lineage as a patch, which appends rather than overwrites existing Dataset lineage edges. By default it is enabled.                                               |
-|
+| spark.datahub.stage_metadata_coalescing                             |          |         | Normally it coalesces and sends metadata at the onApplicationEnd event which is never called on Databricks or on Glue. You should enable this on Databricks if you want coalesced run.     |
+| spark.datahub.patch.enabled                                         |          | false   | Set this to true to send lineage as a patch, which appends rather than overwrites existing Dataset lineage edges. By default, it is disabled.                                              |
+| spark.datahub.metadata.dataset.lowerCaseUrns                        |          | false   | Set this to true to lowercase dataset urns. By default, it is disabled.                                                                                                                    |
+| spark.datahub.disableSymlinkResolution                              |          | false   | Set this to true if you prefer using the s3 location instead of the Hive table. By default, it is disabled.                          |
 
 ## What to Expect: The Metadata Model
 
@@ -343,3 +344,15 @@ Use Java 8 to build the project. The project uses Gradle as the build tool. To b
 ```
 ## Known limitations
 
++ 
+## Changelog
+
+### Version 0.2.12
+- Silencing some chatty warnings in RddPathUtils
+
+### Version 0.2.12
+
+- Add option to lowercase dataset URNs
+- Add option to set platform instance and/or env per platform with `spark.datahub.platform.<platform_name>.env` and `spark.datahub.platform.<platform_name>.platform_instance` config parameter
+- Fixing platform instance setting for datasets when `spark.datahub.metadata.dataset.platformInstance` is set
+- Fixing column level lineage support when patch is enabled

metadata-integration/java/spark-lineage-beta/src/main/java/datahub/spark/conf/SparkConfigParser.java

@@ -42,6 +42,8 @@ public class SparkConfigParser {
   public static final String DATAHUB_FLOW_NAME = "flow_name";
   public static final String DATASET_ENV_KEY = "metadata.dataset.env";
   public static final String DATASET_HIVE_PLATFORM_ALIAS = "metadata.dataset.hivePlatformAlias";
+  public static final String DATASET_LOWERCASE_URNS = "metadata.dataset.lowerCaseUrns";
+
   public static final String DATASET_MATERIALIZE_KEY = "metadata.dataset.materialize";
   public static final String DATASET_PLATFORM_INSTANCE_KEY = "metadata.dataset.platformInstance";
   public static final String DATASET_INCLUDE_SCHEMA_METADATA =
@@ -152,6 +154,7 @@ public static DatahubOpenlineageConfig sparkConfigToDatahubOpenlineageConf(
     builder.hivePlatformAlias(SparkConfigParser.getHivePlatformAlias(sparkConfig));
     builder.usePatch(SparkConfigParser.isPatchEnabled(sparkConfig));
     builder.disableSymlinkResolution(SparkConfigParser.isDisableSymlinkResolution(sparkConfig));
+    builder.lowerCaseDatasetUrns(SparkConfigParser.isLowerCaseDatasetUrns(sparkConfig));
     try {
       String parentJob = SparkConfigParser.getParentJobKey(sparkConfig);
       if (parentJob != null) {
@@ -246,15 +249,18 @@ public static Map<String, List<PathSpec>> getPathSpecListMap(Config datahubConfi
           pathSpecBuilder.alias(pathSpecKey);
           pathSpecBuilder.platform(key);
           if (datahubConfig.hasPath(aliasKey + ".env")) {
-            pathSpecBuilder.env(datahubConfig.getString(aliasKey + ".env"));
+            pathSpecBuilder.env(Optional.ofNullable(datahubConfig.getString(aliasKey + ".env")));
           }
-          if (datahubConfig.hasPath(aliasKey + ".platformInstance")) {
+          if (datahubConfig.hasPath(aliasKey + "." + PLATFORM_INSTANCE_KEY)) {
             pathSpecBuilder.platformInstance(
-                Optional.ofNullable(datahubConfig.getString(aliasKey + ".platformInstance")));
+                Optional.ofNullable(
+                    datahubConfig.getString(aliasKey + "." + PLATFORM_INSTANCE_KEY)));
+          }
+          if (datahubConfig.hasPath(aliasKey + "." + PATH_SPEC_LIST_KEY)) {
+            pathSpecBuilder.pathSpecList(
+                Arrays.asList(
+                    datahubConfig.getString(aliasKey + "." + PATH_SPEC_LIST_KEY).split(",")));
           }
-          pathSpecBuilder.pathSpecList(
-              Arrays.asList(datahubConfig.getString(aliasKey + "." + pathSpecKey).split(",")));
-
           platformSpecs.add(pathSpecBuilder.build());
         }
         pathSpecMap.put(key, platformSpecs);
@@ -264,8 +270,8 @@ public static Map<String, List<PathSpec>> getPathSpecListMap(Config datahubConfi
   }
 
   public static String getPlatformInstance(Config pathSpecConfig) {
-    return pathSpecConfig.hasPath(PLATFORM_INSTANCE_KEY)
-        ? pathSpecConfig.getString(PLATFORM_INSTANCE_KEY)
+    return pathSpecConfig.hasPath(PIPELINE_PLATFORM_INSTANCE_KEY)
+        ? pathSpecConfig.getString(PIPELINE_PLATFORM_INSTANCE_KEY)
         : null;
   }
 
@@ -341,4 +347,9 @@ public static boolean isEmitCoalescePeriodically(Config datahubConfig) {
     return datahubConfig.hasPath(STAGE_METADATA_COALESCING)
         && datahubConfig.getBoolean(STAGE_METADATA_COALESCING);
   }
+
+  public static boolean isLowerCaseDatasetUrns(Config datahubConfig) {
+    return datahubConfig.hasPath(DATASET_LOWERCASE_URNS)
+        && datahubConfig.getBoolean(DATASET_LOWERCASE_URNS);
+  }
 }