Add new generateQueryFragments option to query planner config (#2958)

If enabled, the query planner will extract inline fragments into
fragment definitions before sending queries to subgraphs. This can
significantly reduce the size of the query sent to subgraphs, but may
increase the time it takes to plan the query.

This is a reimplementation of
#2893 that applies
transformations on the selection set instead of AST.

Credit to @samuelAndalon for suggesting the new approach. Resolves
#2892

.changeset/spicy-falcons-learn.md

@@ -0,0 +1,9 @@
+---
+"apollo-federation-integration-testsuite": minor
+"@apollo/query-planner": minor
+"@apollo/federation-internals": minor
+---
+
+Add new `generateQueryFragments` option to query planner config
+
+If enabled, the query planner will extract inline fragments into fragment definitions before sending queries to subgraphs. This can significantly reduce the size of the query sent to subgraphs, but may increase the time it takes to plan the query.

internals-js/src/operations.ts

@@ -973,6 +973,18 @@ export class Operation {
     return this.withUpdatedSelectionSetAndFragments(optimizedSelection, finalFragments ?? undefined);
   }
 
+  generateQueryFragments(): Operation {
+    const [minimizedSelectionSet, fragments] = this.selectionSet.minimizeSelectionSet();
+    return new Operation(
+      this.schema,
+      this.rootKind,
+      minimizedSelectionSet,
+      this.variableDefinitions,
+      fragments,
+      this.name,
+    );
+  }
+
   expandAllFragments(): Operation {
     // We clear up the fragments since we've expanded all.
     // Also note that expanding fragment usually generate unecessary fragments/inefficient selections, so it
@@ -1535,6 +1547,59 @@ export class SelectionSet {
     this._selections = mapValues(keyedSelections);
   }
 
+  /**
+   * Takes a selection set and extracts inline fragments into named fragments,
+   * reusing generated named fragments when possible.
+   */
+  minimizeSelectionSet(
+    namedFragments: NamedFragments = new NamedFragments(),
+    seenSelections: Map<string, [SelectionSet, NamedFragmentDefinition][]> = new Map(),
+  ): [SelectionSet, NamedFragments] {
+    const minimizedSelectionSet = this.lazyMap((selection) => {
+      if (selection.kind === 'FragmentSelection' && selection.element.typeCondition && selection.element.appliedDirectives.length === 0 && selection.selectionSet) {
+        // No proper hash code, so we use a unique enough number that's cheap to
+        // compute and handle collisions as necessary.
+        const mockHashCode = `on${selection.element.typeCondition}` + selection.selectionSet.selections().length;
+        const equivalentSelectionSetCandidates = seenSelections.get(mockHashCode);
+        if (equivalentSelectionSetCandidates) {
+          // See if any candidates have an equivalent selection set, i.e. {x y} and {y x}.
+          const match = equivalentSelectionSetCandidates.find(([candidateSet]) => candidateSet.equals(selection.selectionSet!));
+          if (match) {
+            // If we found a match, we can reuse the fragment (but we still need
+            // to create a new FragmentSpread since parent types may differ).
+            return new FragmentSpreadSelection(this.parentType, namedFragments, match[1], []);
+          }
+        }
+
+        // No match, so we need to create a new fragment. First, we minimize the
+        // selection set before creating the fragment with it.
+        const [minimizedSelectionSet] = selection.selectionSet.minimizeSelectionSet(namedFragments, seenSelections);
+        const fragmentDefinition = new NamedFragmentDefinition(
+          this.parentType.schema(),
+          `_generated_${mockHashCode}_${equivalentSelectionSetCandidates?.length ?? 0}`,
+          selection.element.typeCondition
+        ).setSelectionSet(minimizedSelectionSet);
+
+        // Create a new "hash code" bucket or add to the existing one.
+        if (!equivalentSelectionSetCandidates) {
+          seenSelections.set(mockHashCode, [[selection.selectionSet, fragmentDefinition]]);
+          namedFragments.add(fragmentDefinition);
+        } else {
+          equivalentSelectionSetCandidates.push([selection.selectionSet, fragmentDefinition]);
+        }
+
+        return new FragmentSpreadSelection(this.parentType, namedFragments, fragmentDefinition, []);
+      } else if (selection.kind === 'FieldSelection') {
+        if (selection.selectionSet) {
+          selection = selection.withUpdatedSelectionSet(selection.selectionSet.minimizeSelectionSet(namedFragments, seenSelections)[0]);
+        }
+      }
+      return selection;
+    });
+
+    return [minimizedSelectionSet, namedFragments];
+  }
+
   selectionsInReverseOrder(): readonly Selection[] {
     const length = this._selections.length;
     const reversed = new Array<Selection>(length);
@@ -2413,7 +2478,7 @@ export function selectionOfElement(element: OperationElement, subSelection?: Sel
   return element.kind === 'Field' ? new FieldSelection(element, subSelection) : new InlineFragmentSelection(element, subSelection!);
 }
 
-export type Selection = FieldSelection | FragmentSelection;
+export type Selection = FieldSelection | FragmentSelection | FragmentSpreadSelection;
 abstract class AbstractSelection<TElement extends OperationElement, TIsLeaf extends undefined | never, TOwnType extends AbstractSelection<TElement, TIsLeaf, TOwnType>> {
   constructor(
     readonly element: TElement,

query-planner-js/src/__tests__/buildPlan.test.ts

@@ -5042,6 +5042,190 @@ describe('Named fragments preservation', () => {
   });
 });
 
+describe('Fragment autogeneration', () => {
+  const subgraph = {
+    name: 'Subgraph1',
+    typeDefs: gql`
+      type Query {
+        t: T
+        t2: T
+      }
+
+      union T = A | B
+
+      type A {
+        x: Int
+        y: Int
+        t: T
+      }
+
+      type B {
+        z: Int
+      }
+    `,
+  };
+
+  it('respects generateQueryFragments option', () => {
+    const [api, queryPlanner] = composeAndCreatePlannerWithOptions([subgraph], {
+      generateQueryFragments: true,
+    });
+    const operation = operationFromDocument(
+      api,
+      gql`
+        query {
+          t {
+            ... on A {
+              x
+              y
+            }
+            ... on B {
+              z
+            }
+          }
+        }
+      `,
+    );
+
+    const plan = queryPlanner.buildQueryPlan(operation);
+
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Fetch(service: "Subgraph1") {
+          {
+            t {
+              __typename
+              ..._generated_onA2_0
+              ..._generated_onB1_0
+            }
+          }
+          
+          fragment _generated_onA2_0 on A {
+            x
+            y
+          }
+          
+          fragment _generated_onB1_0 on B {
+            z
+          }
+        },
+      }
+    `);
+  });
+
+  it('handles nested fragment generation', () => {
+    const [api, queryPlanner] = composeAndCreatePlannerWithOptions([subgraph], {
+      generateQueryFragments: true,
+    });
+    const operation = operationFromDocument(
+      api,
+      gql`
+        query {
+          t {
+            ... on A {
+              x
+              y
+              t {
+                ... on A {
+                  x
+                }
+                ... on B {
+                  z
+                }
+              }
+            }
+            ... on B {
+              z
+            }
+          }
+        }
+      `,
+    );
+
+    const plan = queryPlanner.buildQueryPlan(operation);
+
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Fetch(service: "Subgraph1") {
+          {
+            t {
+              __typename
+              ..._generated_onA3_0
+              ..._generated_onB1_0
+            }
+          }
+          
+          fragment _generated_onA1_0 on A {
+            x
+          }
+          
+          fragment _generated_onB1_0 on B {
+            z
+          }
+          
+          fragment _generated_onA3_0 on A {
+            x
+            y
+            t {
+              __typename
+              ..._generated_onA1_0
+              ..._generated_onB1_0
+            }
+          }
+        },
+      }
+    `);
+  });
+
+  it("identifies and reuses equivalent fragments that aren't identical", () => {
+    const [api, queryPlanner] = composeAndCreatePlannerWithOptions([subgraph], {
+      generateQueryFragments: true,
+    });
+    const operation = operationFromDocument(
+      api,
+      gql`
+        query {
+          t {
+            ... on A {
+              x
+              y
+            }
+          }
+          t2 {
+            ... on A {
+              y
+              x
+            }
+          }
+        }
+      `,
+    );
+
+    const plan = queryPlanner.buildQueryPlan(operation);
+
+    expect(plan).toMatchInlineSnapshot(`
+      QueryPlan {
+        Fetch(service: "Subgraph1") {
+          {
+            t {
+              __typename
+              ..._generated_onA2_0
+            }
+            t2 {
+              __typename
+              ..._generated_onA2_0
+            }
+          }
+          
+          fragment _generated_onA2_0 on A {
+            x
+            y
+          }
+        },
+      }
+    `);
+  });
+});
+
 test('works with key chains', () => {
   const subgraph1 = {
     name: 'Subgraph1',