Adds count field to field picker (#2231) (#2237)

Signed-off-by: Ashwin Pc <ashwinpc@amazon.com>

Signed-off-by: Ashwin Pc <ashwinpc@amazon.com>
(cherry picked from commit aaa35c3)

Co-authored-by: Ashwin P Chandran <ashwinpc@amazon.com>

src/plugins/wizard/public/application/components/data_tab/field_selector.tsx

@@ -73,6 +73,14 @@ export const FieldSelector = () => {
         </form>
       </div>
       <div className="wizFieldSelector__fieldGroups">
+        {/* Count Field */}
+        <FieldSelectorField
+          field={{
+            displayName: 'Count',
+            scripted: false,
+            type: 'number',
+          }}
+        />
         <FieldGroup
           id="categoricalFields"
           header="Categorical Fields"

src/plugins/wizard/public/application/components/data_tab/field_selector_field.tsx

@@ -32,28 +32,22 @@ import React, { useState } from 'react';
 import { IndexPatternField } from '../../../../../data/public';
 import { FieldButton, FieldIcon } from '../../../../../opensearch_dashboards_react/public';
 import { useDrag } from '../../utils/drag_drop/drag_drop_context';
+import { COUNT_FIELD } from '../../utils/drag_drop/types';
 
 import './field_selector_field.scss';
 
 export interface FieldSelectorFieldProps {
-  field: IndexPatternField;
+  field: Partial<IndexPatternField> & Pick<IndexPatternField, 'displayName' | 'type' | 'scripted'>;
 }
 
-export type FieldDragData = Pick<IndexPatternField, 'name' | 'displayName' | 'type'>;
-
 // TODO:
 // 1. Add field sections (Available fields, popular fields from src/plugins/discover/public/application/components/sidebar/discover_sidebar.tsx)
 // 2. Add popover for fields stats from discover as well
 export const FieldSelectorField = ({ field }: FieldSelectorFieldProps) => {
-  const { displayName, type, name } = field;
   const [infoIsOpen, setOpen] = useState(false);
   const [dragProps] = useDrag({
     namespace: 'field-data',
-    value: {
-      displayName,
-      name,
-      type,
-    },
+    value: field.name || COUNT_FIELD,
   });
 
   function togglePopover() {

src/plugins/wizard/public/application/components/data_tab/use/use_dropbox.tsx

@@ -7,7 +7,7 @@ import { useState, useEffect, useCallback, useMemo } from 'react';
 import { cloneDeep } from 'lodash';
 import { BucketAggType, IndexPatternField, propFilter } from '../../../../../../data/common';
 import { Schema } from '../../../../../../vis_default_editor/public';
-import { FieldDragDataType } from '../../../utils/drag_drop/types';
+import { COUNT_FIELD, FieldDragDataType } from '../../../utils/drag_drop/types';
 import { useTypedDispatch, useTypedSelector } from '../../../utils/state_management';
 import { DropboxDisplay, DropboxProps } from '../dropbox';
 import { useDrop } from '../../../utils/drag_drop';
@@ -119,10 +119,10 @@ export const useDropbox = (props: UseDropboxProps): DropboxProps => {
     (data: FieldDragDataType['value']) => {
       if (!data || !validAggTypes.length) return;
 
-      const { name: fieldName } = data;
+      const fieldName = data;
       const schemaAggTypes = (schema.defaults as any).aggTypes;
       const allowedAggTypes = schemaAggTypes
-        ? schemaAggTypes.filter((type) => validAggTypes.includes(type))
+        ? schemaAggTypes.filter((type: string) => validAggTypes.includes(type))
         : [];
 
       aggConfigs?.createAggConfig({
@@ -160,23 +160,31 @@ export const useDropbox = (props: UseDropboxProps): DropboxProps => {
   useEffect(() => {
     const getValidAggTypes = () => {
       if (!dragData || schema.group === 'none') return [];
+      const isCountField = dragData === COUNT_FIELD;
 
-      const indexField = getIndexPatternField(dragData.name, indexPattern?.fields ?? []);
+      const indexField = isCountField
+        ? { type: 'count' }
+        : getIndexPatternField(dragData, indexPattern?.fields ?? []);
 
       if (!indexField) return [];
 
       // Get all aggTypes allowed by the schema and get a list of all the aggTypes that the dragged index field can use
       const aggTypes = aggService.types.getAll();
-      const allowedAggTypes = filterByName(aggTypes[schema.group], schema.aggFilter);
+      // `types` can be either a Bucket or Metric aggType, but both types have the name property.
+      const allowedAggTypes = filterByName(
+        aggTypes[schema.group] as BucketAggType[],
+        schema.aggFilter
+      );
 
       return (
         allowedAggTypes
           .filter((aggType) => {
             const allowedFieldTypes = aggType.paramByName('field')?.filterFieldTypes;
             return filterByType([indexField], allowedFieldTypes).length !== 0;
           })
+          .filter((aggType) => (isCountField ? true : aggType.name !== 'count'))
           // `types` can be either a Bucket or Metric aggType, but both types have the name property.
-          .map((agg) => (agg as BucketAggType).name)
+          .map((aggType) => (aggType as BucketAggType).name)
       );
     };
 

src/plugins/wizard/public/application/utils/drag_drop/types.ts

@@ -3,15 +3,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { IndexPatternField } from '../../../../../data/common';
+import { IndexPatternField, METRIC_TYPES } from '../../../../../data/common';
+
+export const COUNT_FIELD = Symbol.for(METRIC_TYPES.COUNT);
 
 export interface EmptyDragDataType {
   namespace: null;
   value: null;
 }
 export interface FieldDragDataType {
   namespace: 'field-data';
-  value: Pick<IndexPatternField, 'name' | 'displayName' | 'type'> | null;
+  value: IndexPatternField['name'] | null | typeof COUNT_FIELD;
 }
 
 export type DragDataType = EmptyDragDataType | FieldDragDataType;