Prometheus: properly de-accumulate multi-heatmap responses (#53688)

* Implement workaround to #3373, grouping heatmaps by query before sorting prevents bug when calculating a heatmap with multiple values for the same label on the x-axis

* Group heatmap frames prior to sort by the concatenation of their values, excluding quantile (le)

* add unit tests for multiple query and multi-dimensional query

* replace final le value with more apt quartile label

* unify quantile labels

Co-authored-by: Leon Sorokin <leeoniya@gmail.com>

public/app/plugins/datasource/prometheus/result_transformer.test.ts

@@ -291,6 +291,245 @@ describe('Prometheus Result Transformer', () => {
       expect(series.data[0].fields[3].values.toArray()).toEqual([10, 0, 10]);
     });
 
+    it('results with heatmap format from multiple queries should be correctly transformed', () => {
+      const options = {
+        targets: [
+          {
+            format: 'heatmap',
+            refId: 'A',
+          },
+          {
+            format: 'heatmap',
+            refId: 'B',
+          },
+        ],
+      } as unknown as DataQueryRequest<PromQuery>;
+      const response = {
+        state: 'Done',
+        data: [
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [10, 10, 0],
+                labels: { le: '1' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [20, 10, 30],
+                labels: { le: '2' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [30, 10, 40],
+                labels: { le: '+Inf' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'B',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [10, 10, 0],
+                labels: { le: '1' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'B',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [20, 10, 30],
+                labels: { le: '2' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'B',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [30, 10, 40],
+                labels: { le: '+Inf' },
+              },
+            ],
+          }),
+        ],
+      } as unknown as DataQueryResponse;
+
+      const series = transformV2(response, options, {});
+      expect(series.data[0].fields.length).toEqual(4);
+      expect(series.data[0].fields[1].values.toArray()).toEqual([10, 10, 0]);
+      expect(series.data[0].fields[2].values.toArray()).toEqual([10, 0, 30]);
+      expect(series.data[0].fields[3].values.toArray()).toEqual([10, 0, 10]);
+    });
+
+    it('results with heatmap format and multiple histograms should be grouped and de-accumulated by non-le labels', () => {
+      const options = {
+        targets: [
+          {
+            format: 'heatmap',
+            refId: 'A',
+          },
+        ],
+      } as unknown as DataQueryRequest<PromQuery>;
+      const response = {
+        state: 'Done',
+        data: [
+          // 10
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [10, 10, 0],
+                labels: { le: '1', additionalProperty: '10' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [20, 10, 30],
+                labels: { le: '2', additionalProperty: '10' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [30, 10, 40],
+                labels: { le: '+Inf', additionalProperty: '10' },
+              },
+            ],
+          }),
+          // 20
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [0, 10, 10],
+                labels: { le: '1', additionalProperty: '20' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [20, 10, 40],
+                labels: { le: '2', additionalProperty: '20' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [30, 10, 60],
+                labels: { le: '+Inf', additionalProperty: '20' },
+              },
+            ],
+          }),
+          // 30
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [30, 30, 60],
+                labels: { le: '1', additionalProperty: '30' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [30, 40, 60],
+                labels: { le: '2', additionalProperty: '30' },
+              },
+            ],
+          }),
+          new MutableDataFrame({
+            refId: 'A',
+            fields: [
+              { name: 'Time', type: FieldType.time, values: [6, 5, 4] },
+              {
+                name: 'Value',
+                type: FieldType.number,
+                values: [40, 40, 60],
+                labels: { le: '+Inf', additionalProperty: '30' },
+              },
+            ],
+          }),
+        ],
+      } as unknown as DataQueryResponse;
+
+      const series = transformV2(response, options, {});
+      expect(series.data[0].fields.length).toEqual(4);
+      expect(series.data[0].fields[1].values.toArray()).toEqual([10, 10, 0]);
+      expect(series.data[0].fields[2].values.toArray()).toEqual([10, 0, 30]);
+      expect(series.data[0].fields[3].values.toArray()).toEqual([10, 0, 10]);
+
+      expect(series.data[1].fields[1].values.toArray()).toEqual([0, 10, 10]);
+      expect(series.data[1].fields[2].values.toArray()).toEqual([20, 0, 30]);
+      expect(series.data[1].fields[3].values.toArray()).toEqual([10, 0, 20]);
+
+      expect(series.data[2].fields[1].values.toArray()).toEqual([30, 30, 60]);
+      expect(series.data[2].fields[2].values.toArray()).toEqual([0, 10, 0]);
+      expect(series.data[2].fields[3].values.toArray()).toEqual([10, 0, 0]);
+    });
+
     it('Retains exemplar frames when data returned is a heatmap', () => {
       const options = {
         targets: [

public/app/plugins/datasource/prometheus/result_transformer.ts

@@ -1,5 +1,5 @@
 import { descending, deviation } from 'd3';
-import { groupBy, partition } from 'lodash';
+import { flatten, forOwn, groupBy, partition } from 'lodash';
 
 import {
   ArrayDataFrame,
@@ -102,9 +102,39 @@ export function transformV2(
     (df) => isHeatmapResult(df, request)
   );
 
-  const processedHeatmapFrames = mergeHeatmapFrames(
+  // Group heatmaps by query
-    transformToHistogramOverTime(heatmapResults.sort(sortSeriesByLabel))
+  const heatmapResultsGroupedByQuery = groupBy<DataFrame>(heatmapResults, (h) => h.refId);
-  );
+
+  // Initialize empty array to push grouped histogram frames to
+  let processedHeatmapResultsGroupedByQuery: DataFrame[][] = [];
+
+  // Iterate through every query in this heatmap
+  for (const query in heatmapResultsGroupedByQuery) {
+    // Get reference to dataFrames for heatmap
+    const heatmapResultsGroup = heatmapResultsGroupedByQuery[query];
+
+    // Create a new grouping by iterating through the data frames...
+    const heatmapResultsGroupedByValues = groupBy<DataFrame>(heatmapResultsGroup, (dataFrame) => {
+      // Each data frame has `Time` and `Value` properties, we want to get the values
+      const values = dataFrame.fields.find((field) => field.name === TIME_SERIES_VALUE_FIELD_NAME);
+      // Specific functionality for special "le" quantile heatmap value, we know if this value exists, that we do not want to calculate the heatmap density across data frames from the same quartile
+      if (values?.labels && HISTOGRAM_QUANTILE_LABEL_NAME in values.labels) {
+        const { le, ...notLE } = values?.labels;
+        return Object.values(notLE).join();
+      }
+
+      // Return a string made from the concatenation of this frame's values to represent a grouping in the query
+      return Object.values(values?.labels ?? []).join();
+    });
+
+    // Then iterate through the resultant object
+    forOwn(heatmapResultsGroupedByValues, (dataFrames, key) => {
+      // Sort frames within each grouping
+      const sortedHeatmap = dataFrames.sort(sortSeriesByLabel);
+      // And push the sorted grouping with the rest
+      processedHeatmapResultsGroupedByQuery.push(mergeHeatmapFrames(transformToHistogramOverTime(sortedHeatmap)));
+    });
+  }
 
   // Everything else is processed as time_series result and graph preferredVisualisationType
   const otherFrames = framesWithoutTableHeatmapsAndExemplars.map((dataFrame) => {
@@ -118,12 +148,16 @@ export function transformV2(
     return df;
   });
 
+  const flattenedProcessedHeatmapFrames = flatten(processedHeatmapResultsGroupedByQuery);
+
   return {
     ...response,
-    data: [...otherFrames, ...processedTableFrames, ...processedHeatmapFrames, ...processedExemplarFrames],
+    data: [...otherFrames, ...processedTableFrames, ...flattenedProcessedHeatmapFrames, ...processedExemplarFrames],
   };
 }
 
+const HISTOGRAM_QUANTILE_LABEL_NAME = 'le';
+
 export function transformDFToTable(dfs: DataFrame[]): DataFrame[] {
   // If no dataFrames or if 1 dataFrames with no values, return original dataFrame
   if (dfs.length === 0 || (dfs.length === 1 && dfs[0].length === 0)) {
@@ -151,7 +185,7 @@ export function transformDFToTable(dfs: DataFrame[]): DataFrame[] {
         .forEach((label) => {
           // If we don't have label in labelFields, add it
           if (!labelFields.some((l) => l.name === label)) {
-            const numberField = label === 'le';
+            const numberField = label === HISTOGRAM_QUANTILE_LABEL_NAME;
             labelFields.push({
               name: label,
               config: { filterable: true },
@@ -441,7 +475,7 @@ function transformMetricDataToTable(md: MatrixOrVectorResult[], options: Transfo
     .map((label) => {
       // Labels have string field type, otherwise table tries to figure out the type which can result in unexpected results
       // Only "le" label has a number field type
-      const numberField = label === 'le';
+      const numberField = label === HISTOGRAM_QUANTILE_LABEL_NAME;
       return {
         name: label,
         config: { filterable: true },
@@ -475,7 +509,7 @@ function transformMetricDataToTable(md: MatrixOrVectorResult[], options: Transfo
 
 function getLabelValue(metric: PromMetric, label: string): string | number {
   if (metric.hasOwnProperty(label)) {
-    if (label === 'le') {
+    if (label === HISTOGRAM_QUANTILE_LABEL_NAME) {
       return parseSampleValue(metric[label]);
     }
     return metric[label];