From fc16fb04075add316c07f8c487ff27afc61f7157 Mon Sep 17 00:00:00 2001
From: Matias Chomicki <matyax@gmail.com>
Date: Thu, 30 Mar 2023 10:29:19 +0200
Subject: [PATCH] Loki Query Chunking: Group queries by resolution (#65353)

* Loki chunking: group queries by resolution

* Update unit tests

* Add chunked + grouped test case
---
 .../datasource/loki/queryChunking.test.ts     | 69 +++++++++++++++++++
 .../plugins/datasource/loki/queryChunking.ts  | 46 +++++++------
 2 files changed, 95 insertions(+), 20 deletions(-)

diff --git a/public/app/plugins/datasource/loki/queryChunking.test.ts b/public/app/plugins/datasource/loki/queryChunking.test.ts
index f57307e9cde..126decd19f9 100644
--- a/public/app/plugins/datasource/loki/queryChunking.test.ts
+++ b/public/app/plugins/datasource/loki/queryChunking.test.ts
@@ -260,4 +260,73 @@ describe('runQueryInChunks()', () => {
       });
     });
   });
+
+  describe('Splitting targets based on resolution', () => {
+    const range1d = {
+      from: dateTime('2023-02-08T05:00:00.000Z'),
+      to: dateTime('2023-02-09T05:00:00.000Z'),
+      raw: {
+        from: dateTime('2023-02-08T05:00:00.000Z'),
+        to: dateTime('2023-02-09T05:00:00.000Z'),
+      },
+    };
+    test('Groups logs queries by resolution', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [
+          { expr: '{a="b"}', refId: 'A', resolution: 3 },
+          { expr: '{a="b"}', refId: 'B', resolution: 5 },
+        ],
+        range: range1d,
+      });
+      await expect(runQueryInChunks(datasource, request)).toEmitValuesWith(() => {
+        // A, B
+        expect(datasource.runQuery).toHaveBeenCalledTimes(2);
+      });
+    });
+    test('Groups metric queries by resolution', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [
+          { expr: 'count_over_time({a="b"}[1m])', refId: 'A', resolution: 3 },
+          { expr: 'count_over_time{a="b"}[1m])', refId: 'B', resolution: 5 },
+        ],
+        range: range1d,
+      });
+      await expect(runQueryInChunks(datasource, request)).toEmitValuesWith(() => {
+        // A, B
+        expect(datasource.runQuery).toHaveBeenCalledTimes(2);
+      });
+    });
+    test('Groups mixed queries by resolution', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [
+          { expr: '{a="b"}', refId: 'A', resolution: 3 },
+          { expr: '{a="b"}', refId: 'B', resolution: 5 },
+          { expr: 'count_over_time({a="b"}[1m])', refId: 'C', resolution: 3 },
+          { expr: 'count_over_time{a="b"}[1m])', refId: 'D', resolution: 5 },
+          { expr: '{a="b"}', refId: 'E', resolution: 5, queryType: LokiQueryType.Instant },
+        ],
+        range: range1d,
+      });
+      await expect(runQueryInChunks(datasource, request)).toEmitValuesWith(() => {
+        // A, B, C, D, E
+        expect(datasource.runQuery).toHaveBeenCalledTimes(5);
+      });
+    });
+    test('Chunked groups mixed queries by resolution', async () => {
+      const request = getQueryOptions<LokiQuery>({
+        targets: [
+          { expr: '{a="b"}', refId: 'A', resolution: 3 },
+          { expr: '{a="b"}', refId: 'B', resolution: 5 },
+          { expr: 'count_over_time({a="b"}[1m])', refId: 'C', resolution: 3 },
+          { expr: 'count_over_time{a="b"}[1m])', refId: 'D', resolution: 5 },
+          { expr: '{a="b"}', refId: 'E', resolution: 5, queryType: LokiQueryType.Instant },
+        ],
+        range, // 3 days
+      });
+      await expect(runQueryInChunks(datasource, request)).toEmitValuesWith(() => {
+        // 3 * A, 3 * B, 3 * C, 3 * D, 1 * E
+        expect(datasource.runQuery).toHaveBeenCalledTimes(13);
+      });
+    });
+  });
 });
diff --git a/public/app/plugins/datasource/loki/queryChunking.ts b/public/app/plugins/datasource/loki/queryChunking.ts
index 525387123d1..0489f0af487 100644
--- a/public/app/plugins/datasource/loki/queryChunking.ts
+++ b/public/app/plugins/datasource/loki/queryChunking.ts
@@ -184,29 +184,35 @@ export function runQueryInChunks(datasource: LokiDatasource, request: DataQueryR
 
   const requests: LokiGroupedRequest = [];
   for (const [chunkRangeMs, queries] of Object.entries(rangePartitionedLogQueries)) {
-    requests.push({
-      request: { ...request, targets: queries },
-      partition: partitionTimeRange(
-        true,
-        request.range,
-        request.intervalMs,
-        queries[0].resolution ?? 1,
-        Number(chunkRangeMs)
-      ),
-    });
+    const resolutionPartition = groupBy(queries, (query) => query.resolution || 1);
+    for (const resolution in resolutionPartition) {
+      requests.push({
+        request: { ...request, targets: resolutionPartition[resolution] },
+        partition: partitionTimeRange(
+          true,
+          request.range,
+          request.intervalMs,
+          Number(resolution),
+          Number(chunkRangeMs)
+        ),
+      });
+    }
   }
 
   for (const [chunkRangeMs, queries] of Object.entries(rangePartitionedMetricQueries)) {
-    requests.push({
-      request: { ...request, targets: queries },
-      partition: partitionTimeRange(
-        false,
-        request.range,
-        request.intervalMs,
-        queries[0].resolution ?? 1,
-        Number(chunkRangeMs)
-      ),
-    });
+    const resolutionPartition = groupBy(queries, (query) => query.resolution || 1);
+    for (const resolution in resolutionPartition) {
+      requests.push({
+        request: { ...request, targets: resolutionPartition[resolution] },
+        partition: partitionTimeRange(
+          false,
+          request.range,
+          request.intervalMs,
+          Number(resolution),
+          Number(chunkRangeMs)
+        ),
+      });
+    }
   }
 
   if (instantQueries.length) {