test updating idMap based on new agg order

elastic · Jun 14, 2022 · 78bfdce · 78bfdce
1 parent 5d61742
commit 78bfdce
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 12 deletions.
diff --git a/x-pack/plugins/lens/public/indexpattern_datasource/indexpattern.test.ts b/x-pack/plugins/lens/public/indexpattern_datasource/indexpattern.test.ts
@@ -1009,6 +1009,80 @@ describe('IndexPattern Data Source', () => {
       optimizeMock.mockRestore();
     });
 
+    it('should update anticipated esAggs column IDs based on the order of the optimized agg expression builders', () => {
+      const queryBaseState: DataViewBaseState = {
+        currentIndexPatternId: '1',
+        layers: {
+          first: {
+            indexPatternId: '1',
+            columns: {
+              col1: {
+                label: 'timestamp',
+                dataType: 'date',
+                operationType: 'date_histogram',
+                sourceField: 'timestamp',
+                isBucketed: true,
+                scale: 'interval',
+                params: {
+                  interval: 'auto',
+                  includeEmptyRows: true,
+                  dropPartials: false,
+                },
+              } as DateHistogramIndexPatternColumn,
+              col2: {
+                label: '95th percentile of bytes',
+                dataType: 'number',
+                operationType: 'percentile',
+                sourceField: 'bytes',
+                isBucketed: false,
+                scale: 'ratio',
+                params: {
+                  percentile: 95,
+                },
+              } as PercentileIndexPatternColumn,
+              col3: {
+                label: 'Count of records',
+                dataType: 'number',
+                isBucketed: false,
+                sourceField: '___records___',
+                operationType: 'count',
+                timeScale: 'h',
+              },
+              col4: {
+                label: 'Count of records2',
+                dataType: 'number',
+                isBucketed: false,
+                sourceField: '___records___',
+                operationType: 'count',
+                timeScale: 'h',
+              },
+            },
+            columnOrder: ['col1', 'col2', 'col3', 'col4'],
+            incompleteColumns: {},
+          },
+        },
+      };
+
+      const state = enrichBaseState(queryBaseState);
+
+      const optimizeMock = jest
+        .spyOn(operationDefinitionMap.percentile, 'optimizeEsAggs')
+        .mockImplementation((aggs, esAggsIdMap) => {
+          // change the order of the aggregations
+          return { aggs: aggs.reverse(), esAggsIdMap };
+        });
+
+      const ast = indexPatternDatasource.toExpression(state, 'first') as Ast;
+
+      expect(operationDefinitionMap.percentile.optimizeEsAggs).toHaveBeenCalledTimes(1);
+
+      const idMap = JSON.parse(ast.chain[2].arguments.idMap as unknown as string);
+
+      expect(Object.keys(idMap)).toEqual(['col-0-3', 'col-1-2', 'col-2-1', 'col-3-0']);
+
+      optimizeMock.mockRestore();
+    });
+
     describe('references', () => {
       beforeEach(() => {
         // @ts-expect-error we are inserting an invalid type

diff --git a/x-pack/plugins/lens/public/indexpattern_datasource/to_expression.ts b/x-pack/plugins/lens/public/indexpattern_datasource/to_expression.ts
@@ -217,30 +217,29 @@ function getExpressionForLayer(
       col-1-2:    column3
       col-3-3.98: column4 (98th percentile)
     */
-    const newEsAggsIdToOriginalColumn: Record<string, OriginalColumn[]> = {};
+
+    const updatedEsAggsIdMap: Record<string, OriginalColumn[]> = {};
+    const extractEsAggId = (id: string) => id.split('.')[0].split('-')[2];
+    const updatePositionIndex = (currentId: string, newIndex: number) => {
+      const [fullId, percentile] = currentId.split('.');
+      const idParts = fullId.split('-');
+      idParts[1] = String(newIndex);
+      return idParts.join('-') + (percentile ? `.${percentile}` : '');
+    };
     let counter = 0;
     aggs.forEach((builder) => {
-      const extractEsAggId = (id: string) => id.split('.')[0].split('-')[2];
-      const updatePositionIndex = (currentId: string, newIndex: number) => {
-        const [fullId, percentile] = currentId.split('.');
-        const idParts = fullId.split('-');
-        idParts[1] = String(newIndex);
-        return idParts.join('-') + (percentile ? `.${percentile}` : '');
-      };
-
       const esAggId = builder.functions[0].getArgument('id')?.[0];
       const matchingEsAggColumnIds = Object.keys(esAggsIdMap).filter(
         (id) => extractEsAggId(id) === esAggId
       );
 
       matchingEsAggColumnIds.forEach((currentId) => {
         const currentColumn = esAggsIdMap[currentId][0];
-        // TODO what if multiple columns are mapped to, and only one is bucketed?
         const aggIndex = window.ELASTIC_LENS_DELAY_SECONDS
           ? counter + (currentColumn.isBucketed ? 0 : 1)
           : counter;
         const newId = updatePositionIndex(currentId, aggIndex);
-        newEsAggsIdToOriginalColumn[newId] = esAggsIdMap[currentId];
+        updatedEsAggsIdMap[newId] = esAggsIdMap[currentId];
 
         counter++;
       });
@@ -363,7 +362,7 @@ function getExpressionForLayer(
           type: 'function',
           function: 'lens_map_to_original_columns',
           arguments: {
-            idMap: [JSON.stringify(newEsAggsIdToOriginalColumn)],
+            idMap: [JSON.stringify(updatedEsAggsIdMap)],
           },
         },
         ...expressions,