Transformations: Account for group by / count when assessing if calculation is needed (#110546)

* Add logic for groupby count and add tests

* remove skip

* Use better and accurate test naming

* deconstruct options, add test for group by / count

* Have different values for test

* Remove only and change variables to allow optional chaining based on field

* add back destructuring

packages/grafana-data/src/transformations/transformers/groupBy.test.ts

@@ -5,7 +5,7 @@ import { mockTransformationsRegistry } from '../../utils/tests/mockTransformatio
 import { ReducerID } from '../fieldReducer';
 import { transformDataFrame } from '../transformDataFrame';
 
-import { GroupByOperationID, groupByTransformer, GroupByTransformerOptions } from './groupBy';
+import { GroupByOperationID, groupByTransformer, GroupByTransformerOptions, shouldCalculateField } from './groupBy';
 import { DataTransformerID } from './ids';
 
 // returns a simple group by / reducer pair
@@ -492,4 +492,71 @@ describe('GroupBy transformer', () => {
       expect(result[0].fields).toEqual(expected);
     });
   });
+
+  it('should calculate count on a grouped field when selected', async () => {
+    const testSeries = toDataFrame({
+      name: 'A',
+      fields: [
+        { name: 'category', type: FieldType.string, values: ['A', 'A', 'B', 'B', 'B'], config: {} },
+        { name: 'values', type: FieldType.number, values: [1, 2, 3, 4, 5], config: {} },
+      ],
+    });
+
+    let cfg = {
+      id: DataTransformerID.groupBy,
+      options: {
+        fields: {
+          category: {
+            operation: GroupByOperationID.groupBy,
+            aggregations: [ReducerID.count],
+          },
+          values: {
+            operation: undefined,
+          },
+        },
+      },
+    };
+
+    await expect(transformDataFrame([cfg], [testSeries])).toEmitValuesWith((received) => {
+      const result = received[0];
+      const expected: Field[] = [
+        {
+          name: 'category',
+          type: FieldType.string,
+          values: ['A', 'B'],
+          config: {},
+        },
+        {
+          name: 'category (count)',
+          type: FieldType.number,
+          values: [2, 3],
+          config: {},
+        },
+      ];
+      expect(result[0].fields).toEqual(expected);
+    });
+  });
+});
+
+describe('shouldCalculateField()', () => {
+  it.each([
+    [GroupByOperationID.aggregate, [], false],
+    [GroupByOperationID.aggregate, [ReducerID.count], true],
+    [GroupByOperationID.aggregate, [ReducerID.sum, ReducerID.count], true],
+    [GroupByOperationID.groupBy, [], false],
+    [GroupByOperationID.groupBy, [ReducerID.count], true],
+    [GroupByOperationID.groupBy, [ReducerID.sum], false],
+    [GroupByOperationID.groupBy, [ReducerID.sum, ReducerID.count], false],
+  ])('when provided operation %s and aggregations %s, should return %s', (operation, aggregations, expected) => {
+    const field: Field = {
+      name: 'testField',
+      type: FieldType.string,
+      config: {},
+      values: [],
+    };
+    const options: GroupByTransformerOptions = {
+      fields: { testField: { aggregations, operation } },
+    };
+    expect(shouldCalculateField(field, options)).toBe(expected);
+  });
 });

packages/grafana-data/src/transformations/transformers/groupBy.ts

@@ -144,13 +144,20 @@ export const groupByTransformer: DataTransformerInfo<GroupByTransformerOptions>
     ),
 };
 
-const shouldCalculateField = (field: Field, options: GroupByTransformerOptions): boolean => {
+// exported for test
+export const shouldCalculateField = (field: Field, options: GroupByTransformerOptions): boolean => {
   const fieldName = getFieldDisplayName(field);
-  return (
-    options?.fields[fieldName]?.operation === GroupByOperationID.aggregate &&
-    Array.isArray(options?.fields[fieldName]?.aggregations) &&
-    options?.fields[fieldName].aggregations.length > 0
-  );
+  const { operation, aggregations = [] } = options.fields[fieldName] ?? {};
+
+  if (!Array.isArray(aggregations)) {
+    return false;
+  } else if (operation === GroupByOperationID.aggregate) {
+    return aggregations.length > 0;
+  } else if (operation === GroupByOperationID.groupBy) {
+    return aggregations.length === 1 && aggregations[0] === ReducerID.count;
+  } else {
+    return false;
+  }
 };
 
 /**