feat(patch): support fine grained lineage patches (datahub-project#9408)

Co-authored-by: Harshal Sheth <hsheth2@gmail.com>

entity-registry/src/main/java/com/linkedin/metadata/models/registry/template/dataset/UpstreamLineageTemplate.java

@@ -1,20 +1,41 @@
 package com.linkedin.metadata.models.registry.template.dataset;
 
+import static com.fasterxml.jackson.databind.node.JsonNodeFactory.*;
+import static com.linkedin.metadata.Constants.*;
+
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.node.ArrayNode;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.google.common.collect.Streams;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.dataset.FineGrainedLineageArray;
 import com.linkedin.dataset.UpstreamArray;
 import com.linkedin.dataset.UpstreamLineage;
-import com.linkedin.metadata.models.registry.template.ArrayMergingTemplate;
+import com.linkedin.metadata.models.registry.template.CompoundKeyTemplate;
 import java.util.Collections;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
-public class UpstreamLineageTemplate implements ArrayMergingTemplate<UpstreamLineage> {
+public class UpstreamLineageTemplate extends CompoundKeyTemplate<UpstreamLineage> {
 
+ // Fields
  private static final String UPSTREAMS_FIELD_NAME = "upstreams";
  private static final String DATASET_FIELD_NAME = "dataset";
+ private static final String FINE_GRAINED_LINEAGES_FIELD_NAME = "fineGrainedLineages";
+ private static final String FINE_GRAINED_UPSTREAM_TYPE = "upstreamType";
+ private static final String FINE_GRAINED_UPSTREAMS = "upstreams";
+ private static final String FINE_GRAINED_DOWNSTREAM_TYPE = "downstreamType";
+ private static final String FINE_GRAINED_DOWNSTREAMS = "downstreams";
+ private static final String FINE_GRAINED_TRANSFORMATION_OPERATION = "transformOperation";
+ private static final String FINE_GRAINED_CONFIDENCE_SCORE = "confidenceScore";
 
- // TODO: Fine Grained Lineages not patchable at this time, they don't have a well established key
+ // Template support
+ private static final String NONE_TRANSFORMATION_TYPE = "NONE";
+ private static final Float DEFAULT_CONFIDENCE_SCORE = 1.0f;
 
  @Override
  public UpstreamLineage getSubtype(RecordTemplate recordTemplate) throws ClassCastException {
@@ -42,14 +63,250 @@ public UpstreamLineage getDefault() {
  @Nonnull
  @Override
  public JsonNode transformFields(JsonNode baseNode) {
- return arrayFieldToMap(
- baseNode, UPSTREAMS_FIELD_NAME, Collections.singletonList(DATASET_FIELD_NAME));
+ JsonNode transformedNode =
+ arrayFieldToMap(
+ baseNode, UPSTREAMS_FIELD_NAME, Collections.singletonList(DATASET_FIELD_NAME));
+ ((ObjectNode) transformedNode)
+ .set(
+ FINE_GRAINED_LINEAGES_FIELD_NAME,
+ combineAndTransformFineGrainedLineages(
+ transformedNode.get(FINE_GRAINED_LINEAGES_FIELD_NAME)));
+
+ return transformedNode;
  }
 
  @Nonnull
  @Override
  public JsonNode rebaseFields(JsonNode patched) {
- return transformedMapToArray(
- patched, UPSTREAMS_FIELD_NAME, Collections.singletonList(DATASET_FIELD_NAME));
+ JsonNode rebasedNode =
+ transformedMapToArray(
+ patched, UPSTREAMS_FIELD_NAME, Collections.singletonList(DATASET_FIELD_NAME));
+ ((ObjectNode) rebasedNode)
+ .set(
+ FINE_GRAINED_LINEAGES_FIELD_NAME,
+ reconstructFineGrainedLineages(rebasedNode.get(FINE_GRAINED_LINEAGES_FIELD_NAME)));
+ return rebasedNode;
+ }
+
+ /**
+ * Combines fine grained lineage array into a map using upstream and downstream types as keys,
+ * defaulting when not present. Due to this construction, patches will look like: path:
+ * /fineGrainedLineages/TRANSFORMATION_OPERATION/(upstreamType || downstreamType)/TYPE/FIELD_URN,
+ * op: ADD/REMOVE, value: float (confidenceScore) Due to the way FineGrainedLineage was designed
+ * it doesn't necessarily have a consistent key we can reference, so this specialized method
+ * mimics the arrayFieldToMap of the super class with the specialization that it does not put the
+ * full value of the aspect at the end of the key, just the particular array. This prevents
+ * unintended overwrites through improper MCP construction that is technically allowed by the
+ * schema when combining under fields that form the natural key.
+ *
+ * @param fineGrainedLineages the fine grained lineage array node
+ * @return the modified {@link JsonNode} with array fields transformed to maps
+ */
+ private JsonNode combineAndTransformFineGrainedLineages(@Nullable JsonNode fineGrainedLineages) {
+ ObjectNode mapNode = instance.objectNode();
+ if (!(fineGrainedLineages instanceof ArrayNode) || fineGrainedLineages.isEmpty()) {
+ return mapNode;
+ }
+ JsonNode lineageCopy = fineGrainedLineages.deepCopy();
+
+ lineageCopy
+ .elements()
+ .forEachRemaining(
+ node -> {
+ JsonNode nodeClone = node.deepCopy();
+ String transformationOperation =
+ nodeClone.has(FINE_GRAINED_TRANSFORMATION_OPERATION)
+ ? nodeClone.get(FINE_GRAINED_TRANSFORMATION_OPERATION).asText()
+ : NONE_TRANSFORMATION_TYPE;
+
+ if (!mapNode.has(transformationOperation)) {
+ mapNode.set(transformationOperation, instance.objectNode());
+ }
+ ObjectNode transformationOperationNode =
+ (ObjectNode) mapNode.get(transformationOperation);
+
+ Float confidenceScore =
+ nodeClone.has(FINE_GRAINED_CONFIDENCE_SCORE)
+ ? nodeClone.get(FINE_GRAINED_CONFIDENCE_SCORE).floatValue()
+ : DEFAULT_CONFIDENCE_SCORE;
+
+ String upstreamType =
+ nodeClone.has(FINE_GRAINED_UPSTREAM_TYPE)
+ ? nodeClone.get(FINE_GRAINED_UPSTREAM_TYPE).asText()
+ : null;
+ String downstreamType =
+ nodeClone.has(FINE_GRAINED_DOWNSTREAM_TYPE)
+ ? nodeClone.get(FINE_GRAINED_DOWNSTREAM_TYPE).asText()
+ : null;
+ ArrayNode upstreams =
+ nodeClone.has(FINE_GRAINED_UPSTREAMS)
+ ? (ArrayNode) nodeClone.get(FINE_GRAINED_UPSTREAMS)
+ : null;
+ ArrayNode downstreams =
+ nodeClone.has(FINE_GRAINED_DOWNSTREAMS)
+ ? (ArrayNode) nodeClone.get(FINE_GRAINED_DOWNSTREAMS)
+ : null;
+
+ // Handle upstreams
+ if (upstreamType == null) {
+ // Determine default type
+ Urn upstreamUrn =
+ upstreams != null ? UrnUtils.getUrn(upstreams.get(0).asText()) : null;
+ if (upstreamUrn != null
+ && SCHEMA_FIELD_ENTITY_NAME.equals(upstreamUrn.getEntityType())) {
+ upstreamType = FINE_GRAINED_LINEAGE_FIELD_SET_TYPE;
+ } else {
+ upstreamType = FINE_GRAINED_LINEAGE_DATASET_TYPE;
+ }
+ }
+ if (!transformationOperationNode.has(FINE_GRAINED_UPSTREAM_TYPE)) {
+ transformationOperationNode.set(FINE_GRAINED_UPSTREAM_TYPE, instance.objectNode());
+ }
+ ObjectNode upstreamTypeNode =
+ (ObjectNode) transformationOperationNode.get(FINE_GRAINED_UPSTREAM_TYPE);
+ if (!upstreamTypeNode.has(upstreamType)) {
+ upstreamTypeNode.set(upstreamType, instance.objectNode());
+ }
+ if (upstreams != null) {
+ addUrnsToSubType(upstreamTypeNode, upstreams, upstreamType, confidenceScore);
+ }
+
+ // Handle downstreams
+ if (downstreamType == null) {
+ // Determine default type
+ if (downstreams != null && downstreams.size() > 1) {
+ downstreamType = FINE_GRAINED_LINEAGE_FIELD_SET_TYPE;
+ } else {
+ downstreamType = FINE_GRAINED_LINEAGE_FIELD_TYPE;
+ }
+ }
+ if (!transformationOperationNode.has(FINE_GRAINED_DOWNSTREAM_TYPE)) {
+ transformationOperationNode.set(
+ FINE_GRAINED_DOWNSTREAM_TYPE, instance.objectNode());
+ }
+ ObjectNode downstreamTypeNode =
+ (ObjectNode) transformationOperationNode.get(FINE_GRAINED_DOWNSTREAM_TYPE);
+ if (!downstreamTypeNode.has(downstreamType)) {
+ downstreamTypeNode.set(downstreamType, instance.objectNode());
+ }
+ if (downstreams != null) {
+ addUrnsToSubType(downstreamTypeNode, downstreams, downstreamType, confidenceScore);
+ }
+ });
+ return mapNode;
+ }
+
+ private void addUrnsToSubType(
+ JsonNode superType, ArrayNode urnsList, String subType, Float confidenceScore) {
+ ObjectNode upstreamSubTypeNode = (ObjectNode) superType.get(subType);
+ // Will overwrite repeat urns with different confidence scores with the most recently seen
+ upstreamSubTypeNode.setAll(
+ Streams.stream(urnsList.elements())
+ .map(JsonNode::asText)
+ .distinct()
+ .collect(Collectors.toMap(urn -> urn, urn -> instance.numberNode(confidenceScore))));
+ }
+
+ /**
+ * Takes the transformed fine grained lineages map from pre-processing and reconstructs an array
+ * of FineGrainedLineages Avoids producing side effects by copying nodes, use resulting node and
+ * not the original
+ *
+ * @param transformedFineGrainedLineages the transformed fine grained lineage map
+ * @return the modified {@link JsonNode} formatted consistent with the original schema
+ */
+ private ArrayNode reconstructFineGrainedLineages(JsonNode transformedFineGrainedLineages) {
+ if (transformedFineGrainedLineages instanceof ArrayNode) {
+ // We already have an ArrayNode, no need to transform. This happens during `replace`
+ // operations
+ return (ArrayNode) transformedFineGrainedLineages;
+ }
+ ObjectNode mapNode = (ObjectNode) transformedFineGrainedLineages;
+ ArrayNode arrayNode = instance.arrayNode();
+
+ mapNode
+ .fieldNames()
+ .forEachRemaining(
+ transformationOperation -> {
+ final ObjectNode transformationOperationNode =
+ (ObjectNode) mapNode.get(transformationOperation);
+ final ObjectNode upstreamType =
+ transformationOperationNode.has(FINE_GRAINED_UPSTREAM_TYPE)
+ ? (ObjectNode) transformationOperationNode.get(FINE_GRAINED_UPSTREAM_TYPE)
+ : instance.objectNode();
+ final ObjectNode downstreamType =
+ transformationOperationNode.has(FINE_GRAINED_DOWNSTREAM_TYPE)
+ ? (ObjectNode) transformationOperationNode.get(FINE_GRAINED_DOWNSTREAM_TYPE)
+ : instance.objectNode();
+
+ // Handle upstreams
+ if (!upstreamType.isEmpty()) {
+ populateTypeNode(
+ upstreamType,
+ transformationOperation,
+ FINE_GRAINED_UPSTREAM_TYPE,
+ FINE_GRAINED_UPSTREAMS,
+ FINE_GRAINED_DOWNSTREAM_TYPE,
+ arrayNode);
+ }
+
+ // Handle downstreams
+ if (!downstreamType.isEmpty()) {
+ populateTypeNode(
+ downstreamType,
+ transformationOperation,
+ FINE_GRAINED_DOWNSTREAM_TYPE,
+ FINE_GRAINED_DOWNSTREAMS,
+ FINE_GRAINED_UPSTREAM_TYPE,
+ arrayNode);
+ }
+ });
+
+ return arrayNode;
+ }
+
+ private void populateTypeNode(
+ JsonNode typeNode,
+ String transformationOperation,
+ String typeName,
+ String arrayTypeName,
+ String defaultTypeName,
+ ArrayNode arrayNode) {
+ typeNode
+ .fieldNames()
+ .forEachRemaining(
+ subTypeName -> {
+ ObjectNode subType = (ObjectNode) typeNode.get(subTypeName);
+ if (!subType.isEmpty()) {
+ ObjectNode fineGrainedLineage = instance.objectNode();
+ AtomicReference<Float> minimumConfidenceScore = new AtomicReference<>(1.0f);
+
+ fineGrainedLineage.put(typeName, subTypeName);
+ fineGrainedLineage.put(
+ FINE_GRAINED_TRANSFORMATION_OPERATION, transformationOperation);
+ // Array to actually be filled out
+ fineGrainedLineage.set(arrayTypeName, instance.arrayNode());
+ // Added to pass model validation, because we have no way of appropriately pairing
+ // upstreams and downstreams
+ // within fine grained lineages consistently due to being able to have multiple
+ // downstream types paired with a single
+ // transform operation, we just set a default type because it's a required property
+ fineGrainedLineage.put(defaultTypeName, FINE_GRAINED_LINEAGE_FIELD_SET_TYPE);
+ subType
+ .fieldNames()
+ .forEachRemaining(
+ subTypeKey -> {
+ ((ArrayNode) fineGrainedLineage.get(arrayTypeName)).add(subTypeKey);
+ Float scoreValue = subType.get(subTypeKey).floatValue();
+ if (scoreValue <= minimumConfidenceScore.get()) {
+ minimumConfidenceScore.set(scoreValue);
+ fineGrainedLineage.set(
+ FINE_GRAINED_CONFIDENCE_SCORE,
+ instance.numberNode(minimumConfidenceScore.get()));
+ }
+ });
+ arrayNode.add(fineGrainedLineage);
+ }
+ });
  }
 }