feat(gms) Add aspect mappings to graphql fields to improve request efficiency

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/AspectMappingRegistry.java

@@ -0,0 +1,130 @@
+package com.linkedin.datahub.graphql;
+
+import graphql.language.ArrayValue;
+import graphql.language.StringValue;
+import graphql.schema.GraphQLArgument;
+import graphql.schema.GraphQLDirective;
+import graphql.schema.GraphQLObjectType;
+import graphql.schema.GraphQLSchema;
+import java.util.*;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Maps GraphQL field selections to the minimum set of aspects needed to resolve them, enabling
+ * performance optimization by fetching only required aspects instead of all aspects.
+ *
+ * <p>This class scans the GraphQL schema for two directives: - @aspectMapping(aspects:
+ * ["aspectName"]) - declares which aspects a field needs - @noAspects - indicates a field needs no
+ * aspects (computed fields, custom resolvers)
+ *
+ * <p>To use in entity types, add one line to batchLoad: Set<String> aspects =
+ * AspectUtils.getOptimizedAspects(context, "Dataset", ALL_ASPECTS, "datasetKey");
+ *
+ * <p>If any field lacks a mapping directive, getRequiredAspects returns null and the entity type
+ * falls back to fetching all aspects for safety.
+ */
+@Slf4j
+public class AspectMappingRegistry {
+  private final Map<String, Set<String>> fieldToAspects = new HashMap<>();
+
+  public AspectMappingRegistry(GraphQLSchema schema) {
+    buildMappingFromSchema(schema);
+  }
+
+  private void buildMappingFromSchema(GraphQLSchema schema) {
+    schema
+        .getTypeMap()
+        .values()
+        .forEach(
+            type -> {
+              if (type instanceof GraphQLObjectType) {
+                GraphQLObjectType objectType = (GraphQLObjectType) type;
+                String typeName = objectType.getName();
+
+                objectType
+                    .getFieldDefinitions()
+                    .forEach(
+                        field -> {
+                          String fieldName = field.getName();
+                          GraphQLDirective aspectsDirective = field.getDirective("aspectMapping");
+                          GraphQLDirective noAspectsDirective = field.getDirective("noAspects");
+
+                          if (aspectsDirective != null) {
+                            GraphQLArgument aspectsArg = aspectsDirective.getArgument("aspects");
+                            if (aspectsArg != null
+                                && aspectsArg.getArgumentValue().getValue() instanceof ArrayValue) {
+                              ArrayValue aspectsArray =
+                                  (ArrayValue) aspectsArg.getArgumentValue().getValue();
+                              Set<String> aspects =
+                                  aspectsArray.getValues().stream()
+                                      .map(value -> ((StringValue) value).getValue())
+                                      .collect(Collectors.toSet());
+
+                              String key = typeName + "." + fieldName;
+                              fieldToAspects.put(key, aspects);
+                              log.debug(
+                                  "Mapped {}.{} to aspects: {}", typeName, fieldName, aspects);
+                            }
+                          } else if (noAspectsDirective != null) {
+                            String key = typeName + "." + fieldName;
+                            fieldToAspects.put(key, new HashSet<>());
+                            log.debug(
+                                "Mapped {}.{} to to request no specific aspects.",
+                                typeName,
+                                fieldName);
+                          }
+                        });
+              }
+            });
+
+    log.info("Built aspect mapping registry with {} field mappings", fieldToAspects.size());
+  }
+
+  /**
+   * Get required aspects for the given fields on a type. Returns null if any field is unmapped
+   * (fallback to all aspects).
+   *
+   * <p>This method filters the selection set to only include fields that directly belong to the
+   * specified type, regardless of where that type appears in the query tree. This allows it to work
+   * correctly for both top-level queries and nested entities (e.g., Dataset inside SearchResult).
+   */
+  @Nullable
+  public Set<String> getRequiredAspects(
+      String typeName, List<graphql.schema.SelectedField> requestedFields) {
+    Set<String> aspects = new HashSet<>();
+
+    for (graphql.schema.SelectedField field : requestedFields) {
+      String fieldName = field.getName();
+
+      // Skip introspection fields
+      if (fieldName.startsWith("__")) {
+        continue;
+      }
+
+      // Only process fields that belong to the target type
+      // getObjectTypeNames() returns the set of types this field belongs to (accounting for
+      // interfaces/unions)
+      if (!field.getObjectTypeNames().contains(typeName)) {
+        continue;
+      }
+
+      String key = typeName + "." + fieldName;
+      Set<String> fieldAspects = fieldToAspects.get(key);
+
+      if (fieldAspects != null) {
+        aspects.addAll(fieldAspects);
+        log.debug("Field {} mapped to aspects: {}", key, fieldAspects);
+      } else {
+        // Unmapped field - fallback to all aspects to be conservative
+        log.debug(
+            "Field {} has no @aspectMapping or @noAspects directives, will fetch all aspects", key);
+        return null;
+      }
+    }
+
+    log.debug("Computed required aspects for {}: {}", typeName, aspects);
+    return aspects.isEmpty() ? Collections.emptySet() : aspects;
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/QueryContext.java

@@ -4,7 +4,9 @@
 import com.datahub.authentication.Authentication;
 import com.datahub.plugins.auth.authorization.Authorizer;
 import com.linkedin.metadata.config.DataHubAppConfiguration;
+import graphql.schema.DataFetchingEnvironment;
 import io.datahubproject.metadata.context.OperationContext;
+import javax.annotation.Nullable;
 
 /** Provided as input to GraphQL resolvers; used to carry information about GQL request context. */
 public interface QueryContext {
@@ -34,4 +36,38 @@ default String getActorUrn() {
   OperationContext getOperationContext();
 
   DataHubAppConfiguration getDataHubAppConfig();
+
+  /**
+   * Returns the {@link DataFetchingEnvironment} associated with the current GraphQL request. This
+   * provides access to the GraphQL query structure, requested fields, and other execution context.
+   *
+   * @return the DataFetchingEnvironment, or null if not available
+   */
+  @Nullable
+  DataFetchingEnvironment getDataFetchingEnvironment();
+
+  /**
+   * Sets the {@link DataFetchingEnvironment} for the current GraphQL request. This is typically
+   * called by GraphQL resolvers to provide access to the execution context.
+   *
+   * @param environment the DataFetchingEnvironment to associate with this context
+   */
+  void setDataFetchingEnvironment(@Nullable DataFetchingEnvironment environment);
+
+  /**
+   * Returns the {@link AspectMappingRegistry} for optimizing aspect fetching based on GraphQL field
+   * selections.
+   *
+   * @return the AspectMappingRegistry, or null if not available
+   */
+  @Nullable
+  AspectMappingRegistry getAspectMappingRegistry();
+
+  /**
+   * Sets the {@link AspectMappingRegistry} for this context. This is typically called during
+   * context initialization.
+   *
+   * @param aspectMappingRegistry the AspectMappingRegistry to use
+   */
+  void setAspectMappingRegistry(@Nullable AspectMappingRegistry aspectMappingRegistry);
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/EntityTypeResolver.java

@@ -2,6 +2,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
+import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.Entity;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
@@ -44,6 +45,12 @@ private boolean isOnlySelectingIdentityFields(DataFetchingEnvironment environmen
 
   @Override
   public CompletableFuture get(DataFetchingEnvironment environment) {
+    // Set the DataFetchingEnvironment in the QueryContext for access in batchLoad methods
+    QueryContext context = environment.getContext();
+    if (context != null) {
+      context.setDataFetchingEnvironment(environment);
+    }
+
     final Entity resolvedEntity = _entityProvider.apply(environment);
     if (resolvedEntity == null) {
       return CompletableFuture.completedFuture(null);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/load/LoadableTypeResolver.java

@@ -1,5 +1,6 @@
 package com.linkedin.datahub.graphql.resolvers.load;
 
+import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.types.LoadableType;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
@@ -32,6 +33,12 @@ public LoadableTypeResolver(
 
   @Override
   public CompletableFuture<T> get(DataFetchingEnvironment environment) {
+    // Set the DataFetchingEnvironment in the QueryContext for access in batchLoad methods
+    QueryContext context = environment.getContext();
+    if (context != null) {
+      context.setDataFetchingEnvironment(environment);
+    }
+
     final K key = _keyProvider.apply(environment);
     if (key == null) {
       return null;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataset/DatasetType.java

@@ -34,6 +34,7 @@
 import com.linkedin.datahub.graphql.types.mappers.BrowsePathsMapper;
 import com.linkedin.datahub.graphql.types.mappers.BrowseResultMapper;
 import com.linkedin.datahub.graphql.types.mappers.UrnSearchResultsMapper;
+import com.linkedin.datahub.graphql.util.AspectUtils;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
@@ -56,7 +57,9 @@
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import lombok.extern.slf4j.Slf4j;
 
+@Slf4j
 public class DatasetType
     implements SearchableEntityType<Dataset, String>,
         BrowsableEntityType<Dataset, String>,
@@ -97,6 +100,7 @@ public class DatasetType
 
   private static final Set<String> FACET_FIELDS = ImmutableSet.of("origin", "platform");
   private static final String ENTITY_NAME = "dataset";
+  private static final String KEY_ASPECT = "datasetKey";
 
   private final EntityClient entityClient;
 
@@ -135,12 +139,16 @@ public List<DataFetcherResult<Dataset>> batchLoad(
     try {
       final List<Urn> urns = urnStrs.stream().map(UrnUtils::getUrn).collect(Collectors.toList());
 
+      // Determine optimal aspects to fetch based on GraphQL field selections
+      Set<String> aspectsToResolve =
+          AspectUtils.getOptimizedAspects(context, "Dataset", ASPECTS_TO_RESOLVE, KEY_ASPECT);
+
       final Map<Urn, EntityResponse> datasetMap =
           entityClient.batchGetV2(
               context.getOperationContext(),
               Constants.DATASET_ENTITY_NAME,
               new HashSet<>(urns),
-              ASPECTS_TO_RESOLVE);
+              aspectsToResolve);
 
       final List<EntityResponse> gmsResults = new ArrayList<>(urnStrs.size());
       for (Urn urn : urns) {

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/util/AspectUtils.java

@@ -0,0 +1,74 @@
+package com.linkedin.datahub.graphql.util;
+
+import com.linkedin.datahub.graphql.QueryContext;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+import javax.annotation.Nonnull;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Utility methods for optimizing aspect fetching in GraphQL entity types by determining which
+ * aspects need to be fetched based on requested fields.
+ */
+@Slf4j
+public class AspectUtils {
+
+  private AspectUtils() {}
+
+  /**
+   * Determines optimal aspects to fetch based on GraphQL field selections. Falls back to
+   * defaultAspects if optimization isn't possible (missing registry, unmapped fields, etc).
+   *
+   * <p>Usage in entity type batchLoad: Set<String> aspects =
+   * AspectUtils.getOptimizedAspects(context, "Dataset", ALL_ASPECTS, "datasetKey");
+   *
+   * @param context the QueryContext containing AspectMappingRegistry and DataFetchingEnvironment
+   * @param entityTypeName the GraphQL type name (e.g., "Dataset", "CorpUser")
+   * @param defaultAspects the full set of aspects to use as fallback
+   * @param alwaysIncludeAspects aspects to always include (e.g., key aspects)
+   * @return optimized aspect set, or defaultAspects if optimization isn't possible
+   */
+  @Nonnull
+  public static Set<String> getOptimizedAspects(
+      @Nonnull final QueryContext context,
+      @Nonnull final String entityTypeName,
+      @Nonnull final Set<String> defaultAspects,
+      @Nonnull final String... alwaysIncludeAspects) {
+
+    // Check if we have the necessary context for optimization
+    if (context.getDataFetchingEnvironment() == null
+        || context.getAspectMappingRegistry() == null) {
+      log.debug(
+          "DataFetchingEnvironment or AspectMappingRegistry not available for {}, fetching all aspects",
+          entityTypeName);
+      return defaultAspects;
+    }
+
+    // Attempt to determine required aspects from GraphQL field selections
+    Set<String> requiredAspects =
+        context
+            .getAspectMappingRegistry()
+            .getRequiredAspects(
+                entityTypeName, context.getDataFetchingEnvironment().getSelectionSet().getFields());
+
+    // If we couldn't determine required aspects (e.g., unmapped field), fall back to all aspects
+    if (requiredAspects == null) {
+      log.debug(
+          "Could not determine required aspects for {}, falling back to fetching all aspects",
+          entityTypeName);
+      return defaultAspects;
+    }
+
+    // Successfully optimized - build the minimal aspect set
+    Set<String> optimizedAspects = new HashSet<>(requiredAspects);
+
+    // Add any aspects that should always be included
+    if (alwaysIncludeAspects != null && alwaysIncludeAspects.length > 0) {
+      Collections.addAll(optimizedAspects, alwaysIncludeAspects);
+    }
+
+    log.info("Fetching optimized aspect set for {}: {}", entityTypeName, optimizedAspects);
+    return optimizedAspects;
+  }
+}