diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java
index d9e392d9..28202284 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentMethodArgumentResolver.java
@@ -96,6 +96,9 @@ final class EntityArgumentMethodArgumentResolver extends ArgumentMethodArgumentR
 	}
 
 
+	/**
+	 * Wrap the DataFetchingEnvironment to also make the representation map available.
+	 */
 	static class EntityDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
 
 		private final Map<String, Object> representation;
@@ -111,6 +114,10 @@ final class EntityArgumentMethodArgumentResolver extends ArgumentMethodArgumentR
 	}
 
 
+	/**
+	 * Wrap the DataFetchingEnvironment to also make representation maps available
+	 * for batched invocations.
+	 */
 	static class EntityBatchDataFetchingEnvironment extends DelegatingDataFetchingEnvironment {
 
 		private final List<Map<String, Object>> representations;
diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentsMethodArgumentResolver.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentsMethodArgumentResolver.java
index 898c1070..675f58c6 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentsMethodArgumentResolver.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityArgumentsMethodArgumentResolver.java
@@ -30,8 +30,8 @@ import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
 import org.springframework.util.Assert;
 
 /**
- * Resolver for the representation map of an entity, or for all representations
- * for the target schema type (batched handler methods).
+ * Resolver that exposes the raw representation input {@link Map},
+ * or {@link List} of maps for batched invocations.
  *
  * @author Rossen Stoyanchev
  */