mergeResponses: use map to find frames to combine (#112855)

* mergeResponses: use map to find frames to combine

* Remove console

public/app/plugins/datasource/loki/mergeResponses.ts

@@ -13,18 +13,52 @@ import {
 
 import { LOADING_FRAME_NAME } from './querySplitting';
 
+function getFrameKey(frame: DataFrame): string | undefined {
+  // Metric range query data
+  if (frame.meta?.type === DataFrameType.TimeSeriesMulti) {
+    const field = frame.fields.find((f) => f.type === FieldType.number);
+    if (!field) {
+      throw new Error(`Unable to find number field on sharded dataframe!`);
+    }
+    let key = '';
+    if (frame.refId) {
+      key += frame.refId;
+    }
+    if (frame.name) {
+      key += frame.name;
+    }
+    if (field.labels) {
+      key += JSON.stringify(field.labels);
+    }
+    return key !== '' ? key : undefined;
+  }
+  return frame.refId ?? frame.name;
+}
+
 export function combineResponses(currentResponse: DataQueryResponse | null, newResponse: DataQueryResponse) {
   if (!currentResponse) {
     return cloneQueryResponse(newResponse);
   }
 
-  newResponse.data.forEach((newFrame) => {
+  const currentResponseLabelsMap = new Map<string, DataFrame>();
-    const currentFrame = currentResponse.data.find((frame) => shouldCombine(frame, newFrame));
+  currentResponse.data.forEach((frame: DataFrame) => {
-    if (!currentFrame) {
+    const key = getFrameKey(frame);
+    // It is expected that all frames contain a refId or a name, but since the type allows for it
+    // we need to account for possibly undefined cases.
+    if (key) {
+      currentResponseLabelsMap.set(key, frame);
+    }
+  });
+
+  newResponse.data.forEach((newFrame: DataFrame) => {
+    let currentFrame: DataFrame | undefined = undefined;
+    const key = getFrameKey(newFrame);
+    if (key !== undefined && currentResponseLabelsMap.has(key)) {
+      currentFrame = currentResponseLabelsMap.get(key);
+      mergeFrames(currentFrame!, newFrame);
+    } else {
       currentResponse.data.push(cloneDataFrame(newFrame));
-      return;
     }
-    mergeFrames(currentFrame, newFrame);
   });
 
   const mergedErrors = [...(currentResponse.errors ?? []), ...(newResponse.errors ?? [])];