From bf3ac0ff189b06ca75c5e6b241c9f70f79630f63 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=A1bor=20Farkas?= <gabor.farkas@gmail.com>
Date: Tue, 18 Apr 2023 12:04:51 +0200
Subject: [PATCH] loki: query splitting: labels-based frame matching (#66677)

---
 public/app/plugins/datasource/loki/mocks.ts   | 19 ++++++++++
 .../datasource/loki/responseUtils.test.ts     | 11 ++++++
 .../plugins/datasource/loki/responseUtils.ts  | 35 ++++++++++++++++++-
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/public/app/plugins/datasource/loki/mocks.ts b/public/app/plugins/datasource/loki/mocks.ts
index 584b99a5106..f0b3d288ce0 100644
--- a/public/app/plugins/datasource/loki/mocks.ts
+++ b/public/app/plugins/datasource/loki/mocks.ts
@@ -1,6 +1,7 @@
 import {
   ArrayVector,
   DataFrame,
+  DataFrameType,
   DataSourceInstanceSettings,
   DataSourceSettings,
   FieldType,
@@ -151,6 +152,9 @@ export function getMockFrames() {
       },
     ],
     meta: {
+      custom: {
+        frameType: 'LabeledTimeValues',
+      },
       stats: [
         { displayName: 'Summary: total bytes processed', unit: 'decbytes', value: 11 },
         { displayName: 'Ingester: total reached', value: 1 },
@@ -198,6 +202,9 @@ export function getMockFrames() {
       },
     ],
     meta: {
+      custom: {
+        frameType: 'LabeledTimeValues',
+      },
       stats: [
         { displayName: 'Summary: total bytes processed', unit: 'decbytes', value: 22 },
         { displayName: 'Ingester: total reached', value: 2 },
@@ -220,9 +227,13 @@ export function getMockFrames() {
         type: FieldType.number,
         config: {},
         values: new ArrayVector([5, 4]),
+        labels: {
+          level: 'debug',
+        },
       },
     ],
     meta: {
+      type: DataFrameType.TimeSeriesMulti,
       stats: [
         { displayName: 'Ingester: total reached', value: 1 },
         { displayName: 'Summary: total bytes processed', unit: 'decbytes', value: 11 },
@@ -245,9 +256,13 @@ export function getMockFrames() {
         type: FieldType.number,
         config: {},
         values: new ArrayVector([6, 7]),
+        labels: {
+          level: 'debug',
+        },
       },
     ],
     meta: {
+      type: DataFrameType.TimeSeriesMulti,
       stats: [
         { displayName: 'Ingester: total reached', value: 2 },
         { displayName: 'Summary: total bytes processed', unit: 'decbytes', value: 22 },
@@ -271,9 +286,13 @@ export function getMockFrames() {
         type: FieldType.number,
         config: {},
         values: new ArrayVector([6, 7]),
+        labels: {
+          level: 'error',
+        },
       },
     ],
     meta: {
+      type: DataFrameType.TimeSeriesMulti,
       stats: [
         { displayName: 'Ingester: total reached', value: 2 },
         { displayName: 'Summary: total bytes processed', unit: 'decbytes', value: 33 },
diff --git a/public/app/plugins/datasource/loki/responseUtils.test.ts b/public/app/plugins/datasource/loki/responseUtils.test.ts
index 6fa52a91c3b..81286ccb372 100644
--- a/public/app/plugins/datasource/loki/responseUtils.test.ts
+++ b/public/app/plugins/datasource/loki/responseUtils.test.ts
@@ -191,6 +191,9 @@ describe('combineResponses', () => {
           ],
           length: 4,
           meta: {
+            custom: {
+              frameType: 'LabeledTimeValues',
+            },
             stats: [
               {
                 displayName: 'Summary: total bytes processed',
@@ -228,10 +231,14 @@ describe('combineResponses', () => {
               name: 'Value',
               type: 'number',
               values: new ArrayVector([6, 7, 5, 4]),
+              labels: {
+                level: 'debug',
+              },
             },
           ],
           length: 4,
           meta: {
+            type: 'timeseries-multi',
             stats: [
               {
                 displayName: 'Summary: total bytes processed',
@@ -269,10 +276,14 @@ describe('combineResponses', () => {
               name: 'Value',
               type: 'number',
               values: new ArrayVector([6, 7, 5, 4]),
+              labels: {
+                level: 'debug',
+              },
             },
           ],
           length: 4,
           meta: {
+            type: 'timeseries-multi',
             stats: [
               {
                 displayName: 'Summary: total bytes processed',
diff --git a/public/app/plugins/datasource/loki/responseUtils.ts b/public/app/plugins/datasource/loki/responseUtils.ts
index 279133c9450..bebdd5dea11 100644
--- a/public/app/plugins/datasource/loki/responseUtils.ts
+++ b/public/app/plugins/datasource/loki/responseUtils.ts
@@ -1,6 +1,7 @@
 import {
   ArrayVector,
   DataFrame,
+  DataFrameType,
   DataQueryResponse,
   DataQueryResponseData,
   Field,
@@ -8,6 +9,7 @@ import {
   isValidGoDuration,
   Labels,
   QueryResultMetaStat,
+  shallowCompare,
 } from '@grafana/data';
 
 import { isBytesString } from './languageUtils';
@@ -123,7 +125,38 @@ function shouldCombine(frame1: DataFrame, frame2: DataFrame): boolean {
     return false;
   }
 
-  return frame1.name === frame2.name;
+  const frameType1 = frame1.meta?.type;
+  const frameType2 = frame2.meta?.type;
+
+  if (frameType1 !== frameType2) {
+    // we do not join things that have a different type
+    return false;
+  }
+
+  // metric range query data
+  if (frameType1 === DataFrameType.TimeSeriesMulti) {
+    const field1 = frame1.fields.find((f) => f.type === FieldType.number);
+    const field2 = frame2.fields.find((f) => f.type === FieldType.number);
+    if (field1 === undefined || field2 === undefined) {
+      // should never happen
+      return false;
+    }
+
+    return shallowCompare(field1.labels ?? {}, field2.labels ?? {});
+  }
+
+  // logs query data
+  // logs use a special attribute in the dataframe's "custom" section
+  // because we do not have a good "frametype" value for them yet.
+  const customType1 = frame1.meta?.custom?.frameType;
+  const customType2 = frame2.meta?.custom?.frameType;
+
+  if (customType1 === 'LabeledTimeValues' && customType2 === 'LabeledTimeValues') {
+    return true;
+  }
+
+  // should never reach here
+  return false;
 }
 
 export function combineResponses(currentResult: DataQueryResponse | null, newResult: DataQueryResponse) {