Avoid sparsely-populated arrays in query graph cache (#3066)

For condition resolution, we use a cache that supports caching data per query graph node or edge.

This cache's code was pre-allocating arrays that would almost always be sparsely populated, and for large schemas this overhead is significant. This PR changes these arrays into maps.

.changeset/friendly-tables-mix.md

@@ -0,0 +1,7 @@
+---
+"@apollo/composition": patch
+"@apollo/query-planner": patch
+"@apollo/query-graphs": patch
+---
+
+Query graph caches now use maps instead of sparsely-populated arrays for per-subgraph data.

composition-js/src/validate.ts

@@ -705,7 +705,7 @@ class ValidationTraversal {
       conditionResolver: this.conditionResolver,
       overrideConditions: new Map(),
     })));
-    this.previousVisits = new QueryGraphState(supergraphAPI);
+    this.previousVisits = new QueryGraphState();
     this.context = new ValidationContext(
       supergraphSchema,
       subgraphNameToGraphEnumValue,

query-graphs-js/src/conditionsCaching.ts

@@ -1,9 +1,9 @@
 import { SelectionSet, assert } from "@apollo/federation-internals";
 import { ConditionResolution, ConditionResolver, ExcludedConditions, ExcludedDestinations, sameExcludedDestinations } from "./graphPath";
 import { PathContext } from "./pathContext";
-import { Edge, QueryGraph, QueryGraphState } from "./querygraph";
+import { Edge, QueryGraphState } from "./querygraph";
 
-export function cachingConditionResolver(graph: QueryGraph, resolver: ConditionResolver): ConditionResolver {
+export function cachingConditionResolver(resolver: ConditionResolver): ConditionResolver {
   // For every edge having a condition, we cache the resolution its conditions when possible.
   // We save resolution with the set of excluded edges that were used to compute it: the reason we do this is
   // that excluded edges impact the resolution, so we should only used a cached value if we know the excluded
@@ -13,7 +13,7 @@ export function cachingConditionResolver(graph: QueryGraph, resolver: ConditionR
   // when we have no excluded edges, we'd only ever use the cache for the first key of every type. However,
   // as the algorithm always try keys in the same order (the order of the edges in the query graph), including
   // the excluded edges we see on the first ever call is actually the proper thing to do.
-  const cache = new QueryGraphState<undefined, [ConditionResolution, ExcludedDestinations]>(graph);
+  const cache = new QueryGraphState<undefined, [ConditionResolution, ExcludedDestinations]>();
   return (edge: Edge, context: PathContext, excludedDestinations: ExcludedDestinations, excludedConditions: ExcludedConditions, extraConditions?: SelectionSet) => {
     assert(edge.conditions || extraConditions, 'Should not have been called for edge without conditions');
 

query-graphs-js/src/conditionsValidation.ts

@@ -121,5 +121,5 @@ export function simpleValidationConditionResolver({
     // condition is validated. Note that we use a cost of 1 for all conditions as we don't care about efficiency.
     return { satisfied: true, cost: 1 };
   };
-  return withCaching ? cachingConditionResolver(queryGraph, resolver) : resolver;
+  return withCaching ? cachingConditionResolver(resolver) : resolver;
 }

query-graphs-js/src/graphviz.ts

@@ -100,7 +100,7 @@ function addToVizGraph(graph: QueryGraph, vizGraph: GraphBaseModel, noTerminal:
     }
     return vizGraph;
   }
-  const state = new QueryGraphState<NodeModel, EdgeModel>(graph);
+  const state = new QueryGraphState<NodeModel, EdgeModel>();
   const onEdge = function (edge: Edge): boolean {
     const head = edge.head;
     const tail = edge.tail;

query-graphs-js/src/nonTrivialEdgePrecomputing.ts

@@ -2,7 +2,7 @@ import { assert } from "@apollo/federation-internals";
 import { Edge, QueryGraph, QueryGraphState, simpleTraversal } from "./querygraph";
 
 export function preComputeNonTrivialFollowupEdges(graph: QueryGraph): (previousEdge: Edge) => readonly Edge[] {
-  const state = new QueryGraphState<undefined, readonly Edge[]>(graph);
+  const state = new QueryGraphState<undefined, readonly Edge[]>();
   simpleTraversal(graph, () => {}, (edge) => {
     const followupEdges = graph.outEdges(edge.tail);
     state.setEdgeState(edge, computeNonTrivialFollowups(edge, followupEdges));

query-graphs-js/src/querygraph.ts

@@ -476,16 +476,8 @@ export class QueryGraph {
  */
 export class QueryGraphState<VertexState, EdgeState = undefined> {
   // Store some "user" state for each vertex (accessed by index)
-  private readonly verticesStates: (VertexState | undefined)[];
-  private readonly adjacenciesStates: (EdgeState | undefined)[][];
-
-  /**
-   * Creates a new `QueryGraphState` to associate state to the vertices and/or edges of `graph`.
-   */
-  constructor(readonly graph: QueryGraph) {
-    this.verticesStates = new Array(graph.verticesCount());
-    this.adjacenciesStates = new Array(graph.verticesCount());
-  }
+  private readonly verticesStates: Map<number, VertexState> = new Map();
+  private readonly adjacenciesStates: Map<number, Map<number, EdgeState>> = new Map();
 
   /**
    * Associates the provided state to the provided vertex.
@@ -496,7 +488,7 @@ export class QueryGraphState<VertexState, EdgeState = undefined> {
    * @param state - the state/value to associate to `vertex`.
    */
   setVertexState(vertex: Vertex, state: VertexState) {
-    this.verticesStates[vertex.index] = state;
+    this.verticesStates.set(vertex.index, state);
   }
 
   /**
@@ -507,7 +499,7 @@ export class QueryGraphState<VertexState, EdgeState = undefined> {
    *    `QueryGraphState` was created (and its behavior is undefined if it isn't).
    */
   removeVertexState(vertex: Vertex) {
-    this.verticesStates[vertex.index] = undefined;
+    this.verticesStates.delete(vertex.index);
   }
 
   /**
@@ -519,7 +511,7 @@ export class QueryGraphState<VertexState, EdgeState = undefined> {
    * @return the state associated to `vertex`, if any.
    */
   getVertexState(vertex: Vertex): VertexState | undefined {
-    return this.verticesStates[vertex.index];
+    return this.verticesStates.get(vertex.index);
   }
 
   /**
@@ -531,10 +523,12 @@ export class QueryGraphState<VertexState, EdgeState = undefined> {
    * @param state - the state/value to associate to `edge`.
    */
   setEdgeState(edge: Edge, state: EdgeState) {
-    if (!this.adjacenciesStates[edge.head.index]) {
-      this.adjacenciesStates[edge.head.index] = new Array(this.graph.outEdgesCount(edge.head));
+    let edgeMap = this.adjacenciesStates.get(edge.head.index)
+    if (!edgeMap) {
+      edgeMap = new Map();
+      this.adjacenciesStates.set(edge.head.index, edgeMap);
     }
-    this.adjacenciesStates[edge.head.index][edge.index] = state;
+    edgeMap.set(edge.index, state);
   }
 
   /**
@@ -545,7 +539,13 @@ export class QueryGraphState<VertexState, EdgeState = undefined> {
    *    `QueryGraphState` was created (and its behavior is undefined if it isn't).
    */
   removeEdgeState(edge: Edge) {
-    this.adjacenciesStates[edge.head.index][edge.index] = undefined;
+    const edgeMap = this.adjacenciesStates.get(edge.head.index);
+    if (edgeMap) {
+      edgeMap.delete(edge.index);
+      if (edgeMap.size === 0) {
+        this.adjacenciesStates.delete(edge.head.index);
+      }
+    }
   }
 
   /**
@@ -557,16 +557,21 @@ export class QueryGraphState<VertexState, EdgeState = undefined> {
    * @return the state associated to `edge`, if any.
    */
   getEdgeState(edge: Edge): EdgeState | undefined {
-    const forEdge = this.adjacenciesStates[edge.head.index];
-    return forEdge ? forEdge[edge.index] : undefined;
+    return this.adjacenciesStates.get(edge.head.index)?.get(edge.index);
   }
 
   toDebugString(
     vertexMapper: (s: VertexState) => string,
     edgeMapper: (e: EdgeState) => string
   ): string {
-    const vs = this.verticesStates.map((state, idx) => ` ${idx}: ${!state ? "<null>" : vertexMapper(state)}`).join("\n");
-    const es = this.adjacenciesStates.map((adj, vIdx) => adj.map((state, eIdx) => ` ${vIdx}[${eIdx}]: ${!state ? "<null>" : edgeMapper(state)}`).join("\n")).join("\n");
+    const vs = Array.from(this.verticesStates.entries()).sort(([a], [b]) => a - b).map(([idx, state]) =>
+      ` ${idx}: ${!state ? "<null>" : vertexMapper(state)}`
+    ).join("\n");
+    const es = Array.from(this.adjacenciesStates.entries()).sort(([a], [b]) => a - b).map(([vIdx, adj]) =>
+      Array.from(adj.entries()).sort(([a], [b]) => a - b).map(([eIdx, state]) =>
+        ` ${vIdx}[${eIdx}]: ${!state ? "<null>" : edgeMapper(state)}`
+      ).join("\n")
+    ).join("\n");
     return `vertices = {${vs}\n}, edges = {${es}\n}`;
   }
 }

query-planner-js/src/buildPlan.ts

@@ -400,7 +400,6 @@ class QueryPlanningTraversal<RV extends Vertex> {
     this.isTopLevel = isRootVertex(root);
     this.optionsLimit = parameters.config.debug?.pathsLimit;
     this.conditionResolver = cachingConditionResolver(
-      federatedQueryGraph,
       (edge, context, excludedEdges, excludedConditions, extras) => this.resolveConditionPlan(edge, context, excludedEdges, excludedConditions, extras),
     );