From f7fcc14f690d5ba4e321d0feba26876655ebe86c Mon Sep 17 00:00:00 2001
From: Leon Sorokin <leeoniya@gmail.com>
Date: Wed, 23 Oct 2024 17:14:06 -0500
Subject: [PATCH] Transformations: Add field from calc - avoid frame joining in
 some cases (#93602)

* Transformations: Add field from calc - avoid frame joining in some cases

* Add test for multiple query, 1 field, replace all

* omit frames that had no binary ops performed on them

---------

Co-authored-by: drew08t <drew08@gmail.com>
---
 .../transformers/calculateField.test.ts       | 138 ++++++++++++++++++
 .../transformers/calculateField.ts            |  33 ++++-
 2 files changed, 164 insertions(+), 7 deletions(-)

diff --git a/packages/grafana-data/src/transformations/transformers/calculateField.test.ts b/packages/grafana-data/src/transformations/transformers/calculateField.test.ts
index 3091be651bd..c5e5911ce88 100644
--- a/packages/grafana-data/src/transformations/transformers/calculateField.test.ts
+++ b/packages/grafana-data/src/transformations/transformers/calculateField.test.ts
@@ -41,6 +41,15 @@ describe('calculateField transformer w/ timeseries', () => {
     mockTransformationsRegistry([calculateFieldTransformer]);
   });
 
+  beforeEach(() => {
+    seriesA.fields.forEach((f) => {
+      delete f.state;
+    });
+    seriesBC.fields.forEach((f) => {
+      delete f.state;
+    });
+  });
+
   it('will filter and alias', async () => {
     const cfg = {
       id: DataTransformerID.calculateField,
@@ -198,6 +207,56 @@ describe('calculateField transformer w/ timeseries', () => {
     });
   });
 
+  it('multiple queries + field + static number', async () => {
+    const cfg = {
+      id: DataTransformerID.calculateField,
+      options: {
+        mode: CalculateFieldMode.BinaryOperation,
+        binary: {
+          left: 'B',
+          operator: BinaryOperationID.Add,
+          right: '2',
+        },
+        replaceFields: true,
+      },
+    };
+
+    await expect(transformDataFrame([cfg], [seriesA, seriesBC])).toEmitValuesWith((received) => {
+      const data = received[0];
+      expect(data).toMatchInlineSnapshot(`
+        [
+          {
+            "fields": [
+              {
+                "config": {},
+                "name": "TheTime",
+                "state": {
+                  "displayName": "TheTime",
+                  "multipleFrames": true,
+                },
+                "type": "time",
+                "values": [
+                  1000,
+                  2000,
+                ],
+              },
+              {
+                "config": {},
+                "name": "B + 2",
+                "type": "number",
+                "values": [
+                  4,
+                  202,
+                ],
+              },
+            ],
+            "length": 2,
+          },
+        ]
+      `);
+    });
+  });
+
   it('all numbers + static number', async () => {
     const cfg = {
       id: DataTransformerID.calculateField,
@@ -231,6 +290,85 @@ describe('calculateField transformer w/ timeseries', () => {
     });
   });
 
+  it('all numbers + static number (multi-frame, avoids join)', async () => {
+    const cfg = {
+      id: DataTransformerID.calculateField,
+      options: {
+        mode: CalculateFieldMode.BinaryOperation,
+        binary: {
+          left: { matcher: { id: FieldMatcherID.byType, options: FieldType.number } },
+          operator: BinaryOperationID.Add,
+          right: '2',
+        },
+        replaceFields: true,
+      },
+    };
+
+    await expect(transformDataFrame([cfg], [seriesA, seriesBC])).toEmitValuesWith((received) => {
+      const data = received[0];
+
+      expect(data).toMatchInlineSnapshot(`
+        [
+          {
+            "fields": [
+              {
+                "config": {},
+                "name": "TheTime",
+                "type": "time",
+                "values": [
+                  1000,
+                  2000,
+                ],
+              },
+              {
+                "config": {},
+                "name": "A + 2",
+                "type": "number",
+                "values": [
+                  3,
+                  102,
+                ],
+              },
+            ],
+            "length": 2,
+          },
+          {
+            "fields": [
+              {
+                "config": {},
+                "name": "TheTime",
+                "type": "time",
+                "values": [
+                  1000,
+                  2000,
+                ],
+              },
+              {
+                "config": {},
+                "name": "B + 2",
+                "type": "number",
+                "values": [
+                  4,
+                  202,
+                ],
+              },
+              {
+                "config": {},
+                "name": "C + 2",
+                "type": "number",
+                "values": [
+                  5,
+                  302,
+                ],
+              },
+            ],
+            "length": 2,
+          },
+        ]
+      `);
+    });
+  });
+
   it('all numbers + field number', async () => {
     const cfg = {
       id: DataTransformerID.calculateField,
diff --git a/packages/grafana-data/src/transformations/transformers/calculateField.ts b/packages/grafana-data/src/transformations/transformers/calculateField.ts
index 72a136c3980..1a1ada7a5b8 100644
--- a/packages/grafana-data/src/transformations/transformers/calculateField.ts
+++ b/packages/grafana-data/src/transformations/transformers/calculateField.ts
@@ -128,10 +128,16 @@ export const calculateFieldTransformer: DataTransformerInfo<CalculateFieldTransf
     },
   },
   operator: (options, ctx) => (outerSource) => {
-    const operator =
-      options && options.timeSeries !== false
-        ? ensureColumnsTransformer.operator(null, ctx)
-        : noopTransformer.operator({}, ctx);
+    const mode = options.mode ?? CalculateFieldMode.ReduceRow;
+
+    const asTimeSeries = options.timeSeries !== false;
+    const isBinaryFixed = mode === CalculateFieldMode.BinaryOperation && options.binary?.right.fixed != null;
+
+    const needsSingleFrame = asTimeSeries && !isBinaryFixed;
+
+    const operator = needsSingleFrame
+      ? ensureColumnsTransformer.operator(null, ctx)
+      : noopTransformer.operator({}, ctx);
 
     if (options.alias != null) {
       options.alias = ctx.interpolate(options.alias);
@@ -140,7 +146,6 @@ export const calculateFieldTransformer: DataTransformerInfo<CalculateFieldTransf
     return outerSource.pipe(
       operator,
       map((data) => {
-        const mode = options.mode ?? CalculateFieldMode.ReduceRow;
         let creator: ValuesCreator | undefined = undefined;
 
         switch (mode) {
@@ -172,9 +177,10 @@ export const calculateFieldTransformer: DataTransformerInfo<CalculateFieldTransf
             if (binaryOptions.left?.matcher?.id && binaryOptions.left?.matcher.id === FieldMatcherID.byType) {
               const fieldType = binaryOptions.left.matcher.options;
               const operator = binaryOperators.getIfExists(binaryOptions.operator);
-              return data.map((frame) => {
+              const outFrames = data.map((frame) => {
                 const { timeField } = getTimeField(frame);
                 const newFields: Field[] = [];
+                let didAddNewFields = false;
                 if (timeField && options.timeSeries !== false) {
                   newFields.push(timeField);
                 }
@@ -206,10 +212,18 @@ export const calculateFieldTransformer: DataTransformerInfo<CalculateFieldTransf
                       values: arr,
                     };
                     newFields.push(newField);
+                    didAddNewFields = true;
                   }
                 });
+
+                if (options.replaceFields && !didAddNewFields) {
+                  return undefined;
+                }
+
                 return { ...frame, fields: newFields };
               });
+
+              return outFrames.filter((frame) => frame != null);
             } else {
               creator = getBinaryCreator(defaults(binaryOptions, defaultBinaryOptions), data, ctx);
             }
@@ -242,10 +256,14 @@ export const calculateFieldTransformer: DataTransformerInfo<CalculateFieldTransf
           return data;
         }
 
-        return data.map((frame) => {
+        const outFrames = data.map((frame) => {
           // delegate field creation to the specific function
           const values = creator!(frame);
           if (!values) {
+            // if nothing was done to frame, omit it when replacing fields
+            if (options.replaceFields) {
+              return undefined;
+            }
             return frame;
           }
 
@@ -273,6 +291,7 @@ export const calculateFieldTransformer: DataTransformerInfo<CalculateFieldTransf
             fields,
           };
         });
+        return outFrames.filter((frame) => frame != null);
       })
     );
   },