feat(explain-plan-helper): add support for indexes in stages COMPASS-5878

packages/explain-plan-helper/src/get-execution-stats.ts

@@ -4,7 +4,7 @@ import {
   getStageCursorKey,
 } from 'mongodb-explain-compat';
 
-import type { Stage } from './index';
+import type { Stage, IndexInformation } from './index';
 
 export type ExecutionStats = {
   executionSuccess: boolean;
@@ -14,6 +14,7 @@ export type ExecutionStats = {
   totalDocsExamined: number;
   executionStages: Stage;
   allPlansExecution: unknown[];
+  stageIndexes?: IndexInformation[];
 };
 
 export const getExecutionStats = (explain: Stage): ExecutionStats => {
@@ -38,11 +39,13 @@ const _getUnshardedAggregationStats = (explain: Stage): ExecutionStats => {
 
   const stats = _getFindStats(firstStage[cursorKey]);
   stats.nReturned = lastStage.nReturned;
+  stats.stageIndexes = getIndexesFromStages(explain.stages);
   stats.executionTimeMillis = getExecutionTime(stats, explain.stages);
   return stats;
 };
 const _getShardedAggregationStats = (explain: Stage): ExecutionStats => {
   const shardStats = [];
+  let stageIndexes: IndexInformation[] = [];
   for (const shardName in explain.shards) {
     const stats = explain.shards[shardName].stages
       ? _getUnshardedAggregationStats(explain.shards[shardName])
@@ -52,6 +55,9 @@ const _getShardedAggregationStats = (explain: Stage): ExecutionStats => {
       shardName,
       ...stats,
     });
+    stageIndexes = stageIndexes.concat(
+      getIndexesFromStages(explain.shards[shardName].stages ?? [], shardName)
+    );
   }
 
   const nReturned = sumArrayProp(shardStats, 'nReturned');
@@ -65,6 +71,7 @@ const _getShardedAggregationStats = (explain: Stage): ExecutionStats => {
     totalDocsExamined,
     allPlansExecution: [],
     executionSuccess: true,
+    stageIndexes,
     executionStages: {
       stage: shardStats.length === 1 ? 'SINGLE_SHARD' : 'SHARD_MERGE',
       nReturned,
@@ -85,6 +92,21 @@ function sumArrayProp<T>(arr: T[], prop: keyof T): number {
   return arr.reduce((acc, x) => acc + Number(x[prop] ?? 0), 0);
 }
 
+/**
+ *
+ * @param stages List of all stages in the explain plan
+ * @param shard Shard name
+ * @returns Indexes used in the stages
+ */
+function getIndexesFromStages(
+  stages: Stage[],
+  shard?: string
+): IndexInformation[] {
+  return stages
+    .reduce((acc: string[], x: Stage) => acc.concat(x?.indexesUsed ?? []), [])
+    .map((index: string) => ({ index, shard: shard ?? null }));
+}
+
 function getExecutionTime(stats: ExecutionStats, stages: Stage[]): number {
   // We sum executionTimeMillisEstimate from all the stages, except for first
   // as we use its executationStats.executionTimeMillis

packages/explain-plan-helper/src/index.spec.ts

@@ -288,6 +288,20 @@ describe('explain-plan-plan', function () {
         });
       });
 
+      describe('stages with indexesUsed', function () {
+        beforeEach(async function () {
+          plan = await loadExplainFixture(
+            'aggregate_$lookup_with_indexes.json'
+          );
+        });
+
+        it('should have usedIndexes in executionStats object', function () {
+          expect(plan.executionStats.stageIndexes).to.deep.equal([
+            { index: '_id_', shard: null },
+          ]);
+        });
+      });
+
       describe('execution time', function () {
         beforeEach(async function () {
           plan = await loadExplainFixture('aggregate_stages_with_no_time.json');
@@ -506,6 +520,20 @@ describe('explain-plan-plan', function () {
           expect(stage1).to.equal(stage2);
         });
       });
+      describe('stages with indexesUsed', function () {
+        let plan: ExplainPlan;
+        beforeEach(async function () {
+          plan = await loadExplainFixture(
+            'sharded_aggregate_$lookup_with_indexes.json'
+          );
+        });
+
+        it('should have usedIndexes in executionStats object', function () {
+          expect(plan.executionStats.stageIndexes).to.deep.equal([
+            { index: '_id_', shard: 'shard2' },
+          ]);
+        });
+      });
     });
   });
 });

packages/explain-plan-helper/src/index.ts

@@ -47,7 +47,7 @@ export class ExplainPlan {
     const ixscan = this.findAllStagesByName('IXSCAN');
     // special case for IDHACK stage, using the _id_ index.
     const idhack = this.findStageByName('IDHACK');
-    const ret: IndexInformation[] = [];
+    const ret: IndexInformation[] = this.executionStats.stageIndexes ?? [];
     for (const stage of [...ixscan, idhack]) {
       if (!stage) continue;
       let shard: string | null = null;

packages/explain-plan-helper/test/fixtures/aggregate_$lookup_with_indexes.json

@@ -0,0 +1,102 @@
+{
+  "explainVersion": "1",
+  "stages": [
+    {
+      "$cursor": {
+        "queryPlanner": {
+          "namespace": "kashmir.folks",
+          "indexFilterSet": false,
+          "parsedQuery": {},
+          "queryHash": "8B3D4AB8",
+          "planCacheKey": "D542626C",
+          "maxIndexedOrSolutionsReached": false,
+          "maxIndexedAndSolutionsReached": false,
+          "maxScansToExplodeReached": false,
+          "winningPlan": {
+            "stage": "COLLSCAN",
+            "direction": "forward"
+          },
+          "rejectedPlans": []
+        },
+        "executionStats": {
+          "executionSuccess": true,
+          "nReturned": 6,
+          "executionTimeMillis": 0,
+          "totalKeysExamined": 0,
+          "totalDocsExamined": 6,
+          "executionStages": {
+            "stage": "COLLSCAN",
+            "nReturned": 6,
+            "executionTimeMillisEstimate": 0,
+            "works": 8,
+            "advanced": 6,
+            "needTime": 1,
+            "needYield": 0,
+            "saveState": 1,
+            "restoreState": 1,
+            "isEOF": 1,
+            "direction": "forward",
+            "docsExamined": 6
+          },
+          "allPlansExecution": []
+        }
+      },
+      "nReturned": 6,
+      "executionTimeMillisEstimate": 0
+    },
+    {
+      "$lookup": {
+        "from": "districts",
+        "as": "district",
+        "localField": "district_id",
+        "foreignField": "_id",
+        "unwinding": {
+          "preserveNullAndEmptyArrays": false
+        }
+      },
+      "totalDocsExamined": 5,
+      "totalKeysExamined": 5,
+      "collectionScans": 0,
+      "indexesUsed": ["_id_"],
+      "nReturned": 6,
+      "executionTimeMillisEstimate": 0
+    }
+  ],
+  "serverInfo": {
+    "host": "101740a051f6",
+    "port": 27017,
+    "version": "5.0.8",
+    "gitVersion": "c87e1c23421bf79614baf500fda6622bd90f674e"
+  },
+  "serverParameters": {
+    "internalQueryFacetBufferSizeBytes": 104857600,
+    "internalQueryFacetMaxOutputDocSizeBytes": 104857600,
+    "internalLookupStageIntermediateDocumentMaxSizeBytes": 104857600,
+    "internalDocumentSourceGroupMaxMemoryBytes": 104857600,
+    "internalQueryMaxBlockingSortMemoryUsageBytes": 104857600,
+    "internalQueryProhibitBlockingMergeOnMongoS": 0,
+    "internalQueryMaxAddToSetBytes": 104857600,
+    "internalDocumentSourceSetWindowFieldsMaxMemoryBytes": 104857600
+  },
+  "command": {
+    "aggregate": "folks",
+    "pipeline": [
+      {
+        "$lookup": {
+          "from": "districts",
+          "localField": "district_id",
+          "foreignField": "_id",
+          "as": "district"
+        }
+      },
+      {
+        "$unwind": "$district"
+      }
+    ],
+    "allowDiskUse": true,
+    "cursor": {},
+    "maxTimeMS": 60000,
+    "$db": "kashmir"
+  },
+  "ok": 1
+}