fix(lineagecounts) Include entities that are filtered out due to sibl…

…ing logic in the filtered count of lineage counts (datahub-project#8152)

Co-authored-by: Indy Prentice <[email protected]>

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/EntityCountsResolver.java

@@ -38,7 +38,7 @@ public CompletableFuture<EntityCountResults> get(final DataFetchingEnvironment e
  try {
  // First, get all counts
  Map<String, Long> gmsResult = _entityClient.batchGetTotalEntityCount(
- input.getTypes().stream().map(type -> EntityTypeMapper.getName(type)).collect(Collectors.toList()), context.getAuthentication());
+ input.getTypes().stream().map(EntityTypeMapper::getName).collect(Collectors.toList()), context.getAuthentication());
 
  // bind to a result.
  List<EntityCountResult> resultList = gmsResult.entrySet().stream().map(entry -> {

metadata-io/src/main/java/com/linkedin/metadata/graph/SiblingGraphService.java

@@ -122,77 +122,88 @@ public EntityLineageResult getLineage(@Nonnull Urn entityUrn, @Nonnull LineageDi
  offset = Math.max(0, offset - nextEntityLineage.getTotal());
  count = Math.max(0, count - nextEntityLineage.getCount() - entityLineage.getCount());
 
+ entityLineage.setFiltered(getFiltered(entityLineage) + getFiltered(nextEntityLineage));
  entityLineage = nextEntityLineage;
  };
  }
 
  return ValidationUtils.validateEntityLineageResult(entityLineage, _entityService);
  }
 
+ private int getFiltered(@Nullable EntityLineageResult entityLineageResult) {
+ return (entityLineageResult != null && entityLineageResult.getFiltered() != null ? entityLineageResult.getFiltered() : 0);
+ }
+
  // takes a lineage result and removes any nodes that are siblings of some other node already in the result
  private EntityLineageResult filterLineageResultFromSiblings(
  @Nonnull final Urn urn,
  @Nonnull final Set<Urn> allSiblingsInGroup,
  @Nonnull final EntityLineageResult entityLineageResult,
  @Nullable final EntityLineageResult existingResult
  ) {
+ int numFiltered = 0;
+
  // 1) remove the source entities siblings from this entity's downstreams
- final List<LineageRelationship> filteredRelationships = entityLineageResult.getRelationships()
- .stream()
- .filter(lineageRelationship -> !allSiblingsInGroup.contains(lineageRelationship.getEntity())
- || lineageRelationship.getEntity().equals(urn))
- .collect(Collectors.toList());
+ final Map<Boolean, List<LineageRelationship>> partitionedFilteredRelationships = entityLineageResult.getRelationships()
+ .stream().collect(Collectors.partitioningBy(
+ lineageRelationship -> !allSiblingsInGroup.contains(lineageRelationship.getEntity())
+ || lineageRelationship.getEntity().equals(urn)));
+ numFiltered += partitionedFilteredRelationships.get(Boolean.FALSE).size();
+
+ final List<LineageRelationship> filteredRelationships = partitionedFilteredRelationships.get(Boolean.TRUE);
 
  // 2) filter out existing lineage to avoid duplicates in our combined result
  final Set<Urn> existingUrns = existingResult != null
  ? existingResult.getRelationships().stream().map(LineageRelationship::getEntity).collect(Collectors.toSet())
  : new HashSet<>();
- List<LineageRelationship> uniqueFilteredRelationships = filteredRelationships.stream().filter(
- lineageRelationship -> !existingUrns.contains(lineageRelationship.getEntity())).collect(Collectors.toList());
 
- // 3) combine this entity's lineage with the lineage we've already seen and remove duplicates
+ Map<Boolean, List<LineageRelationship>> partitionedUniqueFilteredRelationships = filteredRelationships.stream().collect(
+ Collectors.partitioningBy(lineageRelationship -> !existingUrns.contains(lineageRelationship.getEntity())));
+ numFiltered += partitionedUniqueFilteredRelationships.get(Boolean.FALSE).size();
+
+ List<LineageRelationship> uniqueFilteredRelationships = partitionedUniqueFilteredRelationships.get(Boolean.TRUE);
+
+ // 3) combine this entity's lineage with the lineage we've already seen
  final List<LineageRelationship> combinedResults = Stream.concat(
  uniqueFilteredRelationships.stream(),
  existingResult != null ? existingResult.getRelationships().stream() : ImmutableList.<LineageRelationship>of().stream())
  .collect(Collectors.toList());
 
  // 4) fetch the siblings of each lineage result
- final Set<Urn> combinedResultUrns = combinedResults.stream().map(result -> result.getEntity()).collect(Collectors.toSet());
+ final Set<Urn> combinedResultUrns = combinedResults.stream().map(LineageRelationship::getEntity).collect(Collectors.toSet());
 
  final Map<Urn, List<RecordTemplate>> siblingAspects =
  _entityService.getLatestAspects(combinedResultUrns, ImmutableSet.of(SIBLINGS_ASPECT_NAME));
 
  // 5) if you are not primary & your sibling is in the results, filter yourself out of the return set
- uniqueFilteredRelationships = combinedResults.stream().filter(result -> {
+ Map<Boolean, List<LineageRelationship>> partitionedFilteredSiblings = combinedResults.stream().collect(Collectors.partitioningBy(result -> {
  Optional<RecordTemplate> optionalSiblingsAspect = siblingAspects.get(result.getEntity()).stream().filter(
  aspect -> aspect instanceof Siblings
  ).findAny();
 
- if (!optionalSiblingsAspect.isPresent()) {
+ if (optionalSiblingsAspect.isEmpty()) {
  return true;
  }
 
-
  final Siblings siblingsAspect = (Siblings) optionalSiblingsAspect.get();
 
  if (siblingsAspect.isPrimary()) {
  return true;
  }
 
  // if you are not primary and your sibling exists in the result set, filter yourself out
- if (siblingsAspect.getSiblings().stream().anyMatch(
-  sibling -> combinedResultUrns.contains(sibling)
- )) {
-  return false;
-  }
-
-  return true;
- }).collect(Collectors.toList());
-
- entityLineageResult.setRelationships(new LineageRelationshipArray(uniqueFilteredRelationships));
- entityLineageResult.setTotal(entityLineageResult.getTotal() + (existingResult != null ? existingResult.getTotal() : 0));
- entityLineageResult.setCount(uniqueFilteredRelationships.size());
- return ValidationUtils.validateEntityLineageResult(entityLineageResult, _entityService);
+ return siblingsAspect.getSiblings().stream().noneMatch(combinedResultUrns::contains);
+ }));
+
+ numFiltered += partitionedFilteredSiblings.get(Boolean.FALSE).size();
+ uniqueFilteredRelationships = partitionedFilteredSiblings.get(Boolean.TRUE);
+
+ EntityLineageResult combinedLineageResult = new EntityLineageResult();
+ combinedLineageResult.setStart(entityLineageResult.getStart());
+ combinedLineageResult.setRelationships(new LineageRelationshipArray(uniqueFilteredRelationships));
+ combinedLineageResult.setTotal(entityLineageResult.getTotal() + (existingResult != null ? existingResult.getTotal() : 0));
+ combinedLineageResult.setCount(uniqueFilteredRelationships.size());
+ combinedLineageResult.setFiltered(numFiltered + getFiltered(existingResult) + getFiltered(entityLineageResult));
+ return ValidationUtils.validateEntityLineageResult(combinedLineageResult, _entityService);
  }
-
 }

metadata-io/src/main/java/com/linkedin/metadata/shared/ValidationUtils.java

@@ -150,7 +150,8 @@ public static EntityLineageResult validateEntityLineageResult(@Nullable final En
  .collect(Collectors.toCollection(LineageRelationshipArray::new));
 
  validatedEntityLineageResult.setFiltered(
- entityLineageResult.getRelationships().size() - validatedRelationships.size());
+ (entityLineageResult.hasFiltered() && entityLineageResult.getFiltered() != null ? entityLineageResult.getFiltered() : 0)
+ + entityLineageResult.getRelationships().size() - validatedRelationships.size());
  validatedEntityLineageResult.setRelationships(validatedRelationships);
 
  return validatedEntityLineageResult;