From a9ee446de478d161cb103d677cd1685de4585f52 Mon Sep 17 00:00:00 2001
From: Leon Sorokin <leeoniya@gmail.com>
Date: Mon, 7 Mar 2022 20:51:11 -0600
Subject: [PATCH] HeatmapNG: fix Prometheus frame-per-bucket heatmap response
 rendering (#46309)

---
 .../transformers/calculateHeatmap/heatmap.ts  | 92 +++++++++++--------
 .../app/plugins/panel/heatmap-new/fields.ts   | 11 +--
 2 files changed, 60 insertions(+), 43 deletions(-)

diff --git a/public/app/features/transformers/calculateHeatmap/heatmap.ts b/public/app/features/transformers/calculateHeatmap/heatmap.ts
index 1ca19f436ec..080f289aa6f 100644
--- a/public/app/features/transformers/calculateHeatmap/heatmap.ts
+++ b/public/app/features/transformers/calculateHeatmap/heatmap.ts
@@ -46,12 +46,28 @@ export function sortAscStrInf(aName?: string | null, bName?: string | null) {
 }
 
 /** Given existing buckets, create a values style frame */
+// Assumes frames have already been sorted ASC and de-accumulated.
 export function createHeatmapFromBuckets(frames: DataFrame[]): DataFrame {
-  frames = frames.slice();
+  // assumes all Time fields are identical
+  // TODO: handle null-filling w/ fields[0].config.interval?
+  const xField = frames[0].fields[0];
+  const xValues = xField.values.toArray();
+  const yField = frames[0].fields[1];
 
-  // sort ASC by frame.name (Prometheus bucket bound)
-  // or use frame.fields[1].config.displayNameFromDS ?
-  frames.sort((a, b) => sortAscStrInf(a.name, b.name));
+  // similar to initBins() below
+  const len = xValues.length * frames.length;
+  const xs = new Array(len);
+  const ys = new Array(len);
+  const counts2 = new Array(len);
+
+  const counts = frames.map((frame) => frame.fields[1].values.toArray().slice());
+
+  // transpose
+  counts.forEach((bucketCounts, bi) => {
+    for (let i = 0; i < bucketCounts.length; i++) {
+      counts2[counts.length * i + bi] = bucketCounts[i];
+    }
+  });
 
   const bucketBounds = frames.map((frame, i) => {
     return i; // until we have y ordinal scales working for facets/scatter
@@ -70,39 +86,6 @@ export function createHeatmapFromBuckets(frames: DataFrame[]): DataFrame {
     */
   });
 
-  // assumes all Time fields are identical
-  // TODO: handle null-filling w/ fields[0].config.interval?
-  const xField = frames[0].fields[0];
-  const xValues = xField.values.toArray();
-  const yField = frames[0].fields[1];
-
-  // similar to initBins() below
-  const len = xValues.length * bucketBounds.length;
-  const xs = new Array(len);
-  const ys = new Array(len);
-  const counts2 = new Array(len);
-
-  // cumulative counts
-  const counts = frames.map((frame) => frame.fields[1].values.toArray().slice());
-
-  // de-accumulate
-  counts.reverse();
-  counts.forEach((bucketCounts, bi) => {
-    if (bi < counts.length - 1) {
-      for (let i = 0; i < bucketCounts.length; i++) {
-        bucketCounts[i] -= counts[bi + 1][i];
-      }
-    }
-  });
-  counts.reverse();
-
-  // transpose
-  counts.forEach((bucketCounts, bi) => {
-    for (let i = 0; i < bucketCounts.length; i++) {
-      counts2[counts.length * i + bi] = bucketCounts[i];
-    }
-  });
-
   // fill flat/repeating array
   for (let i = 0, yi = 0, xi = 0; i < len; yi = ++i % bucketBounds.length) {
     ys[i] = bucketBounds[yi];
@@ -144,6 +127,41 @@ export function createHeatmapFromBuckets(frames: DataFrame[]): DataFrame {
   };
 }
 
+// Sorts frames ASC by numeric bucket name and de-accumulates values in each frame's Value field [1]
+// similar to Prometheus result_transformer.ts -> transformToHistogramOverTime()
+export function prepBucketFrames(frames: DataFrame[]): DataFrame[] {
+  frames = frames.slice();
+
+  // sort ASC by frame.name (Prometheus bucket bound)
+  // or use frame.fields[1].config.displayNameFromDS ?
+  frames.sort((a, b) => sortAscStrInf(a.name, b.name));
+
+  // cumulative counts
+  const counts = frames.map((frame) => frame.fields[1].values.toArray().slice());
+
+  // de-accumulate
+  counts.reverse();
+  counts.forEach((bucketCounts, bi) => {
+    if (bi < counts.length - 1) {
+      for (let i = 0; i < bucketCounts.length; i++) {
+        bucketCounts[i] -= counts[bi + 1][i];
+      }
+    }
+  });
+  counts.reverse();
+
+  return frames.map((frame, i) => ({
+    ...frame,
+    fields: [
+      frame.fields[0],
+      {
+        ...frame.fields[1],
+        values: new ArrayVector(counts[i]),
+      },
+    ],
+  }));
+}
+
 export function calculateHeatmapFromData(frames: DataFrame[], options: HeatmapCalculationOptions): DataFrame {
   //console.time('calculateHeatmapFromData');
 
diff --git a/public/app/plugins/panel/heatmap-new/fields.ts b/public/app/plugins/panel/heatmap-new/fields.ts
index 7b7f02b8af4..e1c0cf6f197 100644
--- a/public/app/plugins/panel/heatmap-new/fields.ts
+++ b/public/app/plugins/panel/heatmap-new/fields.ts
@@ -1,9 +1,5 @@
 import { DataFrame, DataFrameType, getDisplayProcessor, GrafanaTheme2 } from '@grafana/data';
-import {
-  calculateHeatmapFromData,
-  createHeatmapFromBuckets,
-  sortAscStrInf,
-} from 'app/features/transformers/calculateHeatmap/heatmap';
+import { calculateHeatmapFromData, createHeatmapFromBuckets } from 'app/features/transformers/calculateHeatmap/heatmap';
 import { HeatmapSourceMode, PanelOptions } from './models.gen';
 
 export const enum BucketLayout {
@@ -59,8 +55,11 @@ export function prepareHeatmapData(
   // detect a frame-per-bucket heatmap frame
   // TODO: improve heuristic? infer from fields[1].labels.le === '+Inf' ?
   if (frames[0].meta?.custom?.resultType === 'matrix' && frames.some((f) => f.name?.startsWith('+Inf'))) {
+    // already done by the Prometheus datasource frontend
+    //frames = prepBucketFrames(frames);
+
     return {
-      yAxisValues: frames.map((f) => f.name ?? null).sort(sortAscStrInf),
+      yAxisValues: frames.map((f) => f.name ?? null),
       ...getHeatmapData(createHeatmapFromBuckets(frames), theme),
     };
   }