feat(explore): Add support for count_unique in eap

static/app/utils/fields/index.ts

@@ -829,6 +829,7 @@ export const ALLOWED_EXPLORE_VISUALIZE_FIELDS: SpanIndexedField[] = [
 
 export const ALLOWED_EXPLORE_VISUALIZE_AGGREGATES: AggregationKey[] = [
   AggregationKey.COUNT, // DO NOT RE-ORDER: the first element is used as the default
+  AggregationKey.COUNT_UNIQUE,
   AggregationKey.AVG,
   AggregationKey.P50,
   AggregationKey.P75,

static/app/views/alerts/rules/metric/eapField.spec.tsx

@@ -82,18 +82,6 @@ describe('EAPField', () => {
     }
 
     render(<Component />);
-    expect(fieldsMock).toHaveBeenCalledWith(
-      `/organizations/${organization.slug}/spans/fields/`,
-      expect.objectContaining({
-        query: expect.objectContaining({type: 'number'}),
-      })
-    );
-    expect(fieldsMock).toHaveBeenCalledWith(
-      `/organizations/${organization.slug}/spans/fields/`,
-      expect.objectContaining({
-        query: expect.objectContaining({type: 'string'}),
-      })
-    );
 
     // switch from count(spans) -> max(span.duration)
     await userEvent.click(screen.getByText('count'));
@@ -111,4 +99,37 @@ describe('EAPField', () => {
     expect(screen.getByText('count')).toBeInTheDocument();
     expect(screen.getByText('spans')).toBeInTheDocument();
   });
+
+  it('defaults count_unique argument to user', async function () {
+    function Component() {
+      const [aggregate, setAggregate] = useState('count(span.duration)');
+      return (
+        <SpanTagsProvider dataset={DiscoverDatasets.SPANS_EAP} enabled>
+          <EAPField aggregate={aggregate} onChange={setAggregate} />
+        </SpanTagsProvider>
+      );
+    }
+
+    render(<Component />);
+
+    // switch from count(spans) -> count_unique(span.op)
+    await userEvent.click(screen.getByText('count'));
+    await userEvent.click(await screen.findByText('count_unique'));
+    expect(screen.getByText('count_unique')).toBeInTheDocument();
+    expect(screen.getByText('span.op')).toBeInTheDocument();
+
+    // switch from count_unique(span.op) -> avg(span.self_time)
+    await userEvent.click(screen.getByText('count_unique'));
+    await userEvent.click(await screen.findByText('avg'));
+    await userEvent.click(screen.getByText('span.duration'));
+    await userEvent.click(await screen.findByText('span.self_time'));
+    expect(screen.getByText('avg')).toBeInTheDocument();
+    expect(screen.getByText('span.self_time')).toBeInTheDocument();
+
+    // switch from avg(span.self_time) -> count_unique(span.op)
+    await userEvent.click(screen.getByText('avg'));
+    await userEvent.click(await screen.findByText('count_unique'));
+    expect(screen.getByText('count_unique')).toBeInTheDocument();
+    expect(screen.getByText('span.op')).toBeInTheDocument();
+  });
 });

static/app/views/alerts/rules/metric/eapField.tsx

@@ -7,16 +7,17 @@ import {space} from 'sentry/styles/space';
 import type {TagCollection} from 'sentry/types/group';
 import {defined} from 'sentry/utils';
 import {parseFunction, prettifyTagKey} from 'sentry/utils/discover/fields';
-import {ALLOWED_EXPLORE_VISUALIZE_AGGREGATES} from 'sentry/utils/fields';
+import {AggregationKey, ALLOWED_EXPLORE_VISUALIZE_AGGREGATES} from 'sentry/utils/fields';
 import {
   DEFAULT_VISUALIZATION,
-  DEFAULT_VISUALIZATION_AGGREGATE,
   DEFAULT_VISUALIZATION_FIELD,
+  updateVisualizeAggregate,
 } from 'sentry/views/explore/contexts/pageParamsContext/visualizes';
 import {useSpanTags} from 'sentry/views/explore/contexts/spanTagsContext';
 
+export const DEFAULT_EAP_AGGREGATION = 'count';
 export const DEFAULT_EAP_FIELD = 'span.duration';
-export const DEFAULT_EAP_METRICS_ALERT_FIELD = `count(${DEFAULT_EAP_FIELD})`;
+export const DEFAULT_EAP_METRICS_ALERT_FIELD = `${DEFAULT_EAP_AGGREGATION}(${DEFAULT_EAP_FIELD})`;
 
 interface Props {
   aggregate: string;
@@ -48,7 +49,10 @@ function EAPField({aggregate, onChange}: Props) {
   // compatibility, we don't want to lock down all `count` queries immediately.
   const lockOptions = aggregate === DEFAULT_VISUALIZATION;
 
-  const {tags: storedTags} = useSpanTags('number');
+  const {tags: storedStringTags} = useSpanTags('string');
+  const {tags: storedNumberTags} = useSpanTags('number');
+  const storedTags =
+    aggregation === AggregationKey.COUNT_UNIQUE ? storedStringTags : storedNumberTags;
   const numberTags: TagCollection = useMemo(() => {
     const availableTags: TagCollection = storedTags;
     if (field && !defined(storedTags[field])) {
@@ -84,13 +88,14 @@ function EAPField({aggregate, onChange}: Props) {
 
   const handleOperationChange = useCallback(
     (option: any) => {
-      const newAggregate =
-        option.value === DEFAULT_VISUALIZATION_AGGREGATE
-          ? DEFAULT_VISUALIZATION
-          : `${option.value}(${field || DEFAULT_EAP_FIELD})`;
+      const newAggregate = updateVisualizeAggregate({
+        newAggregate: option.value,
+        oldAggregate: aggregation || DEFAULT_EAP_AGGREGATION,
+        oldArgument: field || DEFAULT_EAP_FIELD,
+      });
       onChange(newAggregate, {});
     },
-    [field, onChange]
+    [aggregation, field, onChange]
   );
 
   // As SelectControl does not support an options size limit out of the box

static/app/views/dashboards/datasetConfig/spans.tsx

@@ -21,7 +21,7 @@ import {
   doDiscoverQuery,
 } from 'sentry/utils/discover/genericDiscoverQuery';
 import {DiscoverDatasets} from 'sentry/utils/discover/types';
-import {ALLOWED_EXPLORE_VISUALIZE_AGGREGATES} from 'sentry/utils/fields';
+import {AggregationKey, ALLOWED_EXPLORE_VISUALIZE_AGGREGATES} from 'sentry/utils/fields';
 import type {MEPState} from 'sentry/utils/performance/contexts/metricsEnhancedSetting';
 import type {OnDemandControlContext} from 'sentry/utils/performance/contexts/onDemandControl';
 import {
@@ -177,15 +177,22 @@ function _isNotNumericTag(option: FieldValueOption) {
   return true;
 }
 
-function filterAggregateParams(option: FieldValueOption) {
+function filterAggregateParams(option: FieldValueOption, fieldValue?: QueryFieldValue) {
   // Allow for unknown values to be used for aggregate functions
   // This supports showing the tag value even if it's not in the current
   // set of tags.
   if ('unknown' in option.value.meta && option.value.meta.unknown) {
     return true;
   }
+
+  const expectedDataType =
+    fieldValue?.kind === 'function' &&
+    fieldValue?.function[0] === AggregationKey.COUNT_UNIQUE
+      ? 'string'
+      : 'number';
+
   if ('dataType' in option.value.meta) {
-    return option.value.meta.dataType === 'number';
+    return option.value.meta.dataType === expectedDataType;
   }
   return true;
 }

static/app/views/dashboards/widgetBuilder/components/visualize/index.spec.tsx

@@ -42,8 +42,19 @@ describe('Visualize', () => {
         return {tags, isLoading: false};
       }
 
+      const tags: TagCollection = {
+        'span.op': {
+          key: 'span.op',
+          name: 'span.op',
+        },
+        'span.description': {
+          key: 'span.description',
+          name: 'span.description',
+        },
+      };
+
       return {
-        tags: {},
+        tags,
         isLoading: false,
       };
     });
@@ -1379,4 +1390,70 @@ describe('Visualize', () => {
       'spans'
     );
   });
+
+  it('defaults count_unique argument to user', async function () {
+    render(
+      <WidgetBuilderProvider>
+        <Visualize />
+      </WidgetBuilderProvider>,
+      {
+        organization,
+        router: RouterFixture({
+          location: LocationFixture({
+            query: {
+              dataset: WidgetType.SPANS,
+              displayType: DisplayType.LINE,
+              yAxis: ['count(span.duration)'],
+            },
+          }),
+        }),
+      }
+    );
+
+    expect(
+      await screen.findByRole('button', {name: 'Aggregate Selection'})
+    ).toBeEnabled();
+
+    expect(
+      await screen.findByRole('button', {name: 'Aggregate Selection'})
+    ).toHaveTextContent('count');
+    expect(
+      await screen.findByRole('button', {name: 'Column Selection'})
+    ).toHaveTextContent('spans');
+
+    await userEvent.click(
+      await screen.findByRole('button', {name: 'Aggregate Selection'})
+    );
+    await userEvent.click(await screen.findByRole('option', {name: 'count_unique'}));
+
+    expect(
+      await screen.findByRole('button', {name: 'Aggregate Selection'})
+    ).toHaveTextContent('count_unique');
+    expect(
+      await screen.findByRole('button', {name: 'Column Selection'})
+    ).toHaveTextContent('span.op');
+
+    await userEvent.click(
+      await screen.findByRole('button', {name: 'Aggregate Selection'})
+    );
+    await userEvent.click(await screen.findByRole('option', {name: 'avg'}));
+
+    expect(
+      await screen.findByRole('button', {name: 'Aggregate Selection'})
+    ).toHaveTextContent('avg');
+    expect(
+      await screen.findByRole('button', {name: 'Column Selection'})
+    ).toHaveTextContent('span.duration');
+
+    await userEvent.click(
+      await screen.findByRole('button', {name: 'Aggregate Selection'})
+    );
+    await userEvent.click(await screen.findByRole('option', {name: 'count_unique'}));
+    expect(
+      await screen.findByRole('button', {name: 'Aggregate Selection'})
+    ).toHaveTextContent('count_unique');
+    expect(
+      await screen.findByRole('button', {name: 'Column Selection'})
+    ).toHaveTextContent('span.op');
+  });
 });

static/app/views/dashboards/widgetBuilder/components/visualize/selectRow.tsx

@@ -16,6 +16,7 @@ import {
   parseFunction,
   type QueryFieldValue,
 } from 'sentry/utils/discover/fields';
+import {AggregationKey} from 'sentry/utils/fields';
 import useOrganization from 'sentry/utils/useOrganization';
 import {getDatasetConfig} from 'sentry/views/dashboards/datasetConfig/base';
 import {DisplayType, WidgetType} from 'sentry/views/dashboards/types';
@@ -36,6 +37,7 @@ import {
   DEFAULT_VISUALIZATION_AGGREGATE,
   DEFAULT_VISUALIZATION_FIELD,
 } from 'sentry/views/explore/contexts/pageParamsContext/visualizes';
+import {SpanIndexedField} from 'sentry/views/insights/types';
 
 type AggregateFunction = [
   AggregationKeyWithAlias,
@@ -266,6 +268,7 @@ export function SelectRow({
             openColumnSelect();
           } else {
             if (currentField.kind === FieldValueKind.FUNCTION) {
+              const originalFunction = currentField.function[0];
               // Handle setting an aggregate from an aggregate
               currentField.function[0] = parseAggregateFromValueKey(
                 dropdownSelection.value as string
@@ -279,6 +282,30 @@ export function SelectRow({
               ) {
                 currentField.function[1] = DEFAULT_VISUALIZATION_FIELD;
 
+                // Wipe out the remaining parameters that are unnecessary
+                for (let i = 1; i < MAX_FUNCTION_PARAMETERS; i++) {
+                  currentField.function[i + 1] = undefined;
+                }
+              } else if (
+                // when switching to the count_unique aggregate, we want to reset the
+                // field to the default
+                state.dataset === WidgetType.SPANS &&
+                currentField.function[0] === AggregationKey.COUNT_UNIQUE
+              ) {
+                currentField.function[1] = SpanIndexedField.SPAN_OP;
+
+                // Wipe out the remaining parameters that are unnecessary
+                for (let i = 1; i < MAX_FUNCTION_PARAMETERS; i++) {
+                  currentField.function[i + 1] = undefined;
+                }
+              } else if (
+                // when switching away from the count_unique aggregate, we want to reset the
+                // field to the default
+                state.dataset === WidgetType.SPANS &&
+                originalFunction === AggregationKey.COUNT_UNIQUE
+              ) {
+                currentField.function[1] = DEFAULT_VISUALIZATION_FIELD;
+
                 // Wipe out the remaining parameters that are unnecessary
                 for (let i = 1; i < MAX_FUNCTION_PARAMETERS; i++) {
                   currentField.function[i + 1] = undefined;

static/app/views/dashboards/widgetBuilder/widgetBuilderSortBy.spec.tsx

@@ -969,7 +969,7 @@ describe('WidgetBuilder', function () {
       await selectEvent.openMenu(screen.getAllByText('count(\u2026)')[1]!);
 
       // 12 options in the dropdown
-      expect(screen.queryAllByTestId('menu-list-item-label')).toHaveLength(12);
+      expect(screen.queryAllByTestId('menu-list-item-label')).toHaveLength(13);
 
       // Appears once in the y-axis section, dropdown, and in the sort by field
       expect(await screen.findAllByText('count(\u2026)')).toHaveLength(3);

static/app/views/explore/contexts/pageParamsContext/visualizes.tsx

@@ -5,11 +5,13 @@ import type {Organization} from 'sentry/types/organization';
 import {defined} from 'sentry/utils';
 import {parseFunction} from 'sentry/utils/discover/fields';
 import {
+  AggregationKey,
   ALLOWED_EXPLORE_VISUALIZE_AGGREGATES,
   ALLOWED_EXPLORE_VISUALIZE_FIELDS,
 } from 'sentry/utils/fields';
 import {decodeList} from 'sentry/utils/queryString';
 import {ChartType} from 'sentry/views/insights/common/components/chart';
+import {SpanIndexedField} from 'sentry/views/insights/types';
 
 export const MAX_VISUALIZES = 4;
 
@@ -100,3 +102,34 @@ export function updateLocationWithVisualizes(
     delete location.query.visualize;
   }
 }
+
+export function updateVisualizeAggregate({
+  newAggregate,
+  oldAggregate,
+  oldArgument,
+}: {
+  newAggregate: string;
+  oldAggregate: string;
+  oldArgument: string;
+}): string {
+  // the default aggregate only has 1 allowed field
+  if (newAggregate === DEFAULT_VISUALIZATION_AGGREGATE) {
+    return DEFAULT_VISUALIZATION;
+  }
+
+  // count_unique uses a different set of fields
+  if (newAggregate === AggregationKey.COUNT_UNIQUE) {
+    // The general thing to do here is to valid the the argument type
+    // and carry the argument if it's the same type, reset to a default
+    // if it's not the same type. Just hard coding it for now for simplicity
+    // as `count_unique` is the only aggregate that takes a string.
+    return `${AggregationKey.COUNT_UNIQUE}(${SpanIndexedField.SPAN_OP})`;
+  }
+
+  // switching away from count_unique means we need to reset the field
+  if (oldAggregate === AggregationKey.COUNT_UNIQUE) {
+    return `${newAggregate}(${DEFAULT_VISUALIZATION_FIELD})`;
+  }
+
+  return `${newAggregate}(${oldArgument})`;
+}