Model and store column lineage in Marquez DB

CHANGELOG.md

@@ -1,9 +1,13 @@
 # Changelog
 
 ## [Unreleased](https://github.com/MarquezProject/marquez/compare/0.26.0...HEAD)
+
+### Added
+* Implemented dataset symlink feature which allows providing multiple names for a dataset and adds edges to lineage graph based on symlinks [`#2066`](https://github.com/MarquezProject/marquez/pull/2066) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
+* Store column lineage facets in separate table [`#2096`](https://github.com/MarquezProject/marquez/pull/2096) [@mzareba382](https://github.com/mzareba382) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
+
 ### Fixed
 * Add support for `parentRun` facet as reported by older Airflow OpenLineage versions [@collado-mike](https://github.com/collado-mike)
-* Implemented dataset symlink feature which allows providing multiple names for a dataset and adds edges to lineage graph based on symlinks [`#2066`](https://github.com/MarquezProject/marquez/pull/2066) [@pawel-big-lebowski](https://github.com/pawel-big-lebowski)
 
 ## [0.26.0](https://github.com/MarquezProject/marquez/compare/0.25.0...0.26.0) - 2022-09-15
 

api/src/main/java/marquez/db/BaseDao.java

@@ -53,4 +53,7 @@ public interface BaseDao extends SqlObject {
 
   @CreateSqlObject
   OpenLineageDao createOpenLineageDao();
+
+  @CreateSqlObject
+  ColumnLineageDao createColumnLineageDao();
 }

api/src/main/java/marquez/db/ColumnLineageDao.java

@@ -0,0 +1,91 @@
+/*
+ * Copyright 2018-2022 contributors to the Marquez project
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package marquez.db;
+
+import java.time.Instant;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import marquez.db.mappers.ColumnLineageRowMapper;
+import marquez.db.models.ColumnLineageRow;
+import org.apache.commons.lang3.tuple.Pair;
+import org.jdbi.v3.sqlobject.config.RegisterRowMapper;
+import org.jdbi.v3.sqlobject.customizer.BindBeanList;
+import org.jdbi.v3.sqlobject.statement.SqlQuery;
+import org.jdbi.v3.sqlobject.statement.SqlUpdate;
+
+@RegisterRowMapper(ColumnLineageRowMapper.class)
+public interface ColumnLineageDao extends BaseDao {
+
+  default List<ColumnLineageRow> upsertColumnLineageRow(
+      UUID outputDatasetVersionUuid,
+      UUID outputDatasetFieldUuid,
+      List<Pair<UUID, UUID>> inputs,
+      String transformationDescription,
+      String transformationType,
+      Instant now) {
+
+    if (inputs.isEmpty()) {
+      return Collections.emptyList();
+    }
+
+    doUpsertColumnLineageRow(
+        inputs.stream()
+            .map(
+                input ->
+                    new ColumnLineageRow(
+                        outputDatasetVersionUuid,
+                        outputDatasetFieldUuid,
+                        input.getLeft(), // input_dataset_version_uuid
+                        input.getRight(), // input_dataset_field_uuid
+                        transformationDescription,
+                        transformationType,
+                        now,
+                        now))
+            .collect(Collectors.toList()));
+    return findColumnLineageByDatasetVersionColumnAndOutputDatasetField(
+        outputDatasetVersionUuid, outputDatasetFieldUuid);
+  }
+
+  @SqlQuery(
+      "SELECT * FROM column_lineage WHERE output_dataset_version_uuid = :datasetVersionUuid AND output_dataset_field_uuid = :outputDatasetFieldUuid")
+  List<ColumnLineageRow> findColumnLineageByDatasetVersionColumnAndOutputDatasetField(
+      UUID datasetVersionUuid, UUID outputDatasetFieldUuid);
+
+  @SqlUpdate(
+      """
+          INSERT INTO column_lineage (
+          output_dataset_version_uuid,
+          output_dataset_field_uuid,
+          input_dataset_version_uuid,
+          input_dataset_field_uuid,
+          transformation_description,
+          transformation_type,
+          created_at,
+          updated_at
+          ) VALUES <values>
+          ON CONFLICT (output_dataset_version_uuid, output_dataset_field_uuid, input_dataset_version_uuid, input_dataset_field_uuid)
+          DO UPDATE SET
+          transformation_description = EXCLUDED.transformation_description,
+          transformation_type = EXCLUDED.transformation_type,
+          updated_at = EXCLUDED.updated_at
+          """)
+  void doUpsertColumnLineageRow(
+      @BindBeanList(
+              propertyNames = {
+                "outputDatasetVersionUuid",
+                "outputDatasetFieldUuid",
+                "inputDatasetVersionUuid",
+                "inputDatasetFieldUuid",
+                "transformationDescription",
+                "transformationType",
+                "createdAt",
+                "updatedAt"
+              },
+              value = "values")
+          List<ColumnLineageRow> rows);
+}

api/src/main/java/marquez/db/Columns.java

@@ -129,6 +129,15 @@ private Columns() {}
   public static final String RUN_UUID = "run_uuid";
   public static final String STATE = "state";
 
+  /* COLUMN LEVEL LINEAGE ROW COLUMNS */
+  public static final String FIELD_NAME = "field_name";
+  public static final String OUTPUT_DATASET_VERSION_UUID = "output_dataset_version_uuid";
+  public static final String OUTPUT_DATASET_FIELD_UUID = "output_dataset_field_uuid";
+  public static final String INPUT_DATASET_FIELD_UUID = "input_dataset_field_uuid";
+  public static final String INPUT_DATASET_VERSION_UUID = "input_dataset_version_uuid";
+  public static final String TRANSFORMATION_DESCRIPTION = "transformation_description";
+  public static final String TRANSFORMATION_TYPE = "transformation_type";
+
   /* LINEAGE EVENT ROW COLUMNS */
   public static final String EVENT = "event";
 

api/src/main/java/marquez/db/DatasetFieldDao.java

@@ -16,8 +16,10 @@
 import marquez.common.models.TagName;
 import marquez.db.mappers.DatasetFieldMapper;
 import marquez.db.mappers.DatasetFieldRowMapper;
+import marquez.db.mappers.FieldDataMapper;
 import marquez.db.models.DatasetFieldRow;
 import marquez.db.models.DatasetRow;
+import marquez.db.models.InputFieldData;
 import marquez.db.models.TagRow;
 import marquez.service.models.Dataset;
 import marquez.service.models.DatasetVersion;
@@ -29,6 +31,7 @@
 
 @RegisterRowMapper(DatasetFieldRowMapper.class)
 @RegisterRowMapper(DatasetFieldMapper.class)
+@RegisterRowMapper(FieldDataMapper.class)
 public interface DatasetFieldDao extends BaseDao {
   @SqlQuery(
       "SELECT EXISTS ("
@@ -101,6 +104,24 @@ default Dataset updateTags(
           + "WHERE fm.dataset_version_uuid = :datasetVersionUuid")
   List<Field> find(UUID datasetVersionUuid);
 
+  @SqlQuery(
+      """
+          SELECT
+            datasets_view.namespace_name as namespace_name,
+            datasets_view.name as dataset_name,
+            dataset_fields.name as field_name,
+            datasets_view.uuid as dataset_uuid,
+            dataset_versions.uuid as dataset_version_uuid,
+            dataset_fields.uuid as dataset_field_uuid
+          FROM dataset_fields
+          JOIN dataset_versions_field_mapping fm ON fm.dataset_field_uuid = dataset_fields.uuid
+          JOIN dataset_versions ON dataset_versions.uuid = fm.dataset_version_uuid
+          JOIN datasets_view ON datasets_view.uuid = dataset_versions.dataset_uuid
+          JOIN runs_input_mapping ON runs_input_mapping.dataset_version_uuid =  dataset_versions.uuid
+          WHERE runs_input_mapping.run_uuid = :runUuid
+          """)
+  List<InputFieldData> findInputFieldsDataAssociatedWithRun(UUID runUuid);
+
   @SqlQuery(
       "INSERT INTO dataset_fields ("
           + "uuid, "

api/src/main/java/marquez/db/OpenLineageDao.java

@@ -13,13 +13,16 @@
 import java.time.ZoneId;
 import java.time.ZonedDateTime;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import marquez.common.Utils;
 import marquez.common.models.DatasetId;
 import marquez.common.models.DatasetName;
@@ -31,10 +34,12 @@
 import marquez.db.DatasetFieldDao.DatasetFieldMapping;
 import marquez.db.JobVersionDao.BagOfJobVersionInfo;
 import marquez.db.mappers.LineageEventMapper;
+import marquez.db.models.ColumnLineageRow;
 import marquez.db.models.DatasetFieldRow;
 import marquez.db.models.DatasetRow;
 import marquez.db.models.DatasetSymlinkRow;
 import marquez.db.models.DatasetVersionRow;
+import marquez.db.models.InputFieldData;
 import marquez.db.models.JobContextRow;
 import marquez.db.models.JobRow;
 import marquez.db.models.NamespaceRow;
@@ -56,6 +61,7 @@
 import marquez.service.models.LineageEvent.RunFacet;
 import marquez.service.models.LineageEvent.SchemaDatasetFacet;
 import marquez.service.models.LineageEvent.SchemaField;
+import org.apache.commons.lang3.tuple.Pair;
 import org.jdbi.v3.sqlobject.config.RegisterRowMapper;
 import org.jdbi.v3.sqlobject.statement.SqlQuery;
 import org.jdbi.v3.sqlobject.statement.SqlUpdate;
@@ -131,6 +137,7 @@ default UpdateLineageRow updateBaseMarquezModel(LineageEvent event, ObjectMapper
     RunDao runDao = createRunDao();
     RunArgsDao runArgsDao = createRunArgsDao();
     RunStateDao runStateDao = createRunStateDao();
+    ColumnLineageDao columnLineageDao = createColumnLineageDao();
 
     Instant now = event.getEventTime().withZoneSameInstant(ZoneId.of("UTC")).toInstant();
 
@@ -323,7 +330,8 @@ default UpdateLineageRow updateBaseMarquezModel(LineageEvent event, ObjectMapper
                 datasetDao,
                 datasetVersionDao,
                 datasetFieldDao,
-                runDao);
+                runDao,
+                columnLineageDao);
         datasetInputs.add(record);
       }
     }
@@ -345,7 +353,8 @@ default UpdateLineageRow updateBaseMarquezModel(LineageEvent event, ObjectMapper
                 datasetDao,
                 datasetVersionDao,
                 datasetFieldDao,
-                runDao);
+                runDao,
+                columnLineageDao);
         datasetOutputs.add(record);
       }
     }
@@ -541,7 +550,8 @@ default DatasetRecord upsertLineageDataset(
       DatasetDao datasetDao,
       DatasetVersionDao datasetVersionDao,
       DatasetFieldDao datasetFieldDao,
-      RunDao runDao) {
+      RunDao runDao,
+      ColumnLineageDao columnLineageDao) {
     NamespaceRow dsNamespace =
         namespaceDao.upsertNamespaceRow(
             UUID.randomUUID(), now, ds.getNamespace(), DEFAULT_NAMESPACE_OWNER);
@@ -662,6 +672,7 @@ default DatasetRecord upsertLineageDataset(
                   return row;
                 });
     List<DatasetFieldMapping> datasetFieldMappings = new ArrayList<>();
+    List<DatasetFieldRow> datasetFields = new ArrayList<>();
     if (fields != null) {
       for (SchemaField field : fields) {
         DatasetFieldRow datasetFieldRow =
@@ -672,6 +683,7 @@ default DatasetRecord upsertLineageDataset(
                 field.getType(),
                 field.getDescription(),
                 datasetRow.getUuid());
+        datasetFields.add(datasetFieldRow);
         datasetFieldMappings.add(
             new DatasetFieldMapping(datasetVersionRow.getUuid(), datasetFieldRow.getUuid()));
       }
@@ -690,7 +702,85 @@ default DatasetRecord upsertLineageDataset(
       }
     }
 
-    return new DatasetRecord(datasetRow, datasetVersionRow, datasetNamespace);
+    List<ColumnLineageRow> columnLineageRows = Collections.emptyList();
+    if (!isInput) {
+      columnLineageRows =
+          upsertColumnLineage(
+              runUuid,
+              ds,
+              now,
+              datasetFields,
+              columnLineageDao,
+              datasetFieldDao,
+              datasetVersionRow);
+    }
+
+    return new DatasetRecord(datasetRow, datasetVersionRow, datasetNamespace, columnLineageRows);
+  }
+
+  private List<ColumnLineageRow> upsertColumnLineage(
+      UUID runUuid,
+      Dataset ds,
+      Instant now,
+      List<DatasetFieldRow> datasetFields,
+      ColumnLineageDao columnLineageDao,
+      DatasetFieldDao datasetFieldDao,
+      DatasetVersionRow datasetVersionRow) {
+    // get all the fields related to this particular run
+    List<InputFieldData> runFields = datasetFieldDao.findInputFieldsDataAssociatedWithRun(runUuid);
+
+    return Optional.ofNullable(ds.getFacets())
+        .map(DatasetFacets::getColumnLineage)
+        .map(LineageEvent.ColumnLineageFacet::getOutputColumnsList)
+        .stream()
+        .flatMap(list -> list.stream())
+        .flatMap(
+            outputColumn -> {
+              Optional<DatasetFieldRow> outputField =
+                  datasetFields.stream()
+                      .filter(dfr -> dfr.getName().equals(outputColumn.getName()))
+                      .findAny();
+
+              if (outputField.isEmpty()) {
+                Logger log = LoggerFactory.getLogger(OpenLineageDao.class);
+                log.error(
+                    "Cannot produce column lineage for missing output field in output dataset: {}",
+                    outputColumn.getName());
+                return Stream.empty();
+              }
+
+              // get field uuids of input columns related to this run
+              List<Pair<UUID, UUID>> inputFields =
+                  runFields.stream()
+                      .filter(
+                          fieldData ->
+                              outputColumn.getInputFields().stream()
+                                  .filter(
+                                      of ->
+                                          of.getDatasetNamespace().equals(fieldData.getNamespace())
+                                              && of.getDatasetName()
+                                                  .equals(fieldData.getDatasetName())
+                                              && of.getFieldName().equals(fieldData.getField()))
+                                  .findAny()
+                                  .isPresent())
+                      .map(
+                          fieldData ->
+                              Pair.of(
+                                  fieldData.getDatasetVersionUuid(),
+                                  fieldData.getDatasetFieldUuid()))
+                      .collect(Collectors.toList());
+
+              return columnLineageDao
+                  .upsertColumnLineageRow(
+                      datasetVersionRow.getUuid(),
+                      outputField.get().getUuid(),
+                      inputFields,
+                      outputColumn.getTransformationDescription(),
+                      outputColumn.getTransformationType(),
+                      now)
+                  .stream();
+            })
+        .collect(Collectors.toList());
   }
 
   default String formatDatasetName(String name) {

api/src/main/java/marquez/db/mappers/ColumnLineageRowMapper.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2018-2022 contributors to the Marquez project
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package marquez.db.mappers;
+
+import static marquez.db.Columns.TRANSFORMATION_DESCRIPTION;
+import static marquez.db.Columns.TRANSFORMATION_TYPE;
+import static marquez.db.Columns.stringOrThrow;
+import static marquez.db.Columns.timestampOrThrow;
+import static marquez.db.Columns.uuidOrThrow;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import lombok.NonNull;
+import marquez.db.Columns;
+import marquez.db.models.ColumnLineageRow;
+import org.jdbi.v3.core.mapper.RowMapper;
+import org.jdbi.v3.core.statement.StatementContext;
+
+public class ColumnLineageRowMapper implements RowMapper<ColumnLineageRow> {
+
+  @Override
+  public ColumnLineageRow map(@NonNull ResultSet results, @NonNull StatementContext context)
+      throws SQLException {
+    return new ColumnLineageRow(
+        uuidOrThrow(results, Columns.OUTPUT_DATASET_VERSION_UUID),
+        uuidOrThrow(results, Columns.OUTPUT_DATASET_FIELD_UUID),
+        uuidOrThrow(results, Columns.INPUT_DATASET_VERSION_UUID),
+        uuidOrThrow(results, Columns.INPUT_DATASET_FIELD_UUID),
+        stringOrThrow(results, TRANSFORMATION_DESCRIPTION),
+        stringOrThrow(results, TRANSFORMATION_TYPE),
+        timestampOrThrow(results, Columns.CREATED_AT),
+        timestampOrThrow(results, Columns.UPDATED_AT));
+  }
+}