[ML] Transform: Fix percentiles and terms sub-aggs (#138821)

* update percentiles agg

* types

* optional isValid

* terms agg

* update getPivotDropdownOptions tests

* update getAggConfigFromEsAgg percentiles tests

* add terms tests

* rename isPivotAggsConfigWithUiBase type guard

* remove non-null assertion

x-pack/plugins/transform/common/types/pivot_aggs.ts

@@ -5,8 +5,8 @@
  * 2.0.
  */
 
+import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { AggName } from './aggregations';
-import { EsFieldName } from './fields';
 
 export const PIVOT_SUPPORTED_AGGS = {
   AVG: 'avg',
@@ -23,11 +23,7 @@ export const PIVOT_SUPPORTED_AGGS = {
 
 export type PivotSupportedAggs = typeof PIVOT_SUPPORTED_AGGS[keyof typeof PIVOT_SUPPORTED_AGGS];
 
-export type PivotAgg = {
-  [key in PivotSupportedAggs]?: {
-    field: EsFieldName;
-  };
-};
+export type PivotAgg = estypes.AggregationsAggregationContainer;
 
 export type PivotAggDict = {
   [key in AggName]: PivotAgg;

x-pack/plugins/transform/public/app/common/index.ts

@@ -37,7 +37,7 @@ export type {
 } from './pivot_aggs';
 export {
   getEsAggFromAggConfig,
-  isPivotAggsConfigWithUiSupport,
+  isPivotAggsConfigWithUiBase,
   isPivotAggsConfigPercentiles,
   isPivotAggsConfigTerms,
   PERCENTILES_AGG_DEFAULT_PERCENTS,

x-pack/plugins/transform/public/app/common/pivot_aggs.test.ts

@@ -68,7 +68,7 @@ describe('getAggConfigFromEsAgg', () => {
     });
   });
 
-  test('should resolve percentiles agg in sub-aggregations', () => {
+  test('should resolve percentiles and terms agg in sub-aggregations', () => {
     const esConfig = {
       filter: {
         exists: {
@@ -82,6 +82,12 @@ describe('getAggConfigFromEsAgg', () => {
             percents: [1, 5, 25, 50, 75, 95, 99],
           },
         },
+        check_terms: {
+          terms: {
+            field: 'products.base_price',
+            size: 3,
+          },
+        },
       },
     };
 
@@ -93,6 +99,20 @@ describe('getAggConfigFromEsAgg', () => {
       dropDownName: 'products.base_price.percentiles',
       field: 'products.base_price',
       parentAgg: result,
+      aggConfig: {
+        percents: '1,5,25,50,75,95,99',
+      },
+    });
+
+    expect(result.subAggs!.check_terms).toMatchObject({
+      agg: 'terms',
+      aggName: 'check_terms',
+      dropDownName: 'check_terms',
+      field: 'products.base_price',
+      parentAgg: result,
+      aggConfig: {
+        size: 3,
+      },
     });
   });
 

x-pack/plugins/transform/public/app/common/pivot_aggs.ts

@@ -13,8 +13,8 @@ import { isPopulatedObject } from '@kbn/ml-is-populated-object';
 import type { AggName } from '../../../common/types/aggregations';
 import type { Dictionary } from '../../../common/types/common';
 import type { EsFieldName } from '../../../common/types/fields';
-import type { PivotAgg, PivotSupportedAggs } from '../../../common/types/pivot_aggs';
-import { PIVOT_SUPPORTED_AGGS } from '../../../common/types/pivot_aggs';
+import type { PivotSupportedAggs } from '../../../common/types/pivot_aggs';
+import { PIVOT_SUPPORTED_AGGS, PivotAgg } from '../../../common/types/pivot_aggs';
 
 import { getAggFormConfig } from '../sections/create_transform/components/step_define/common/get_agg_form_config';
 import { PivotAggsConfigFilter } from '../sections/create_transform/components/step_define/common/filter_agg/types';
@@ -115,7 +115,7 @@ export const TOP_METRICS_SPECIAL_SORT_FIELDS = {
   _SCORE: '_score',
 } as const;
 
-export const isSpecialSortField = (sortField: unknown) => {
+export const isSpecialSortField = (sortField: unknown): sortField is string => {
   return Object.values(TOP_METRICS_SPECIAL_SORT_FIELDS).some((v) => v === sortField);
 };
 
@@ -205,19 +205,21 @@ export interface PivotAggsConfigWithUiBase extends PivotAggsConfigBase {
   field: EsFieldName | EsFieldName[] | null;
 }
 
-export interface PivotAggsConfigWithExtra<T> extends PivotAggsConfigWithUiBase {
+export interface PivotAggsConfigWithExtra<T, ESConfig extends { [key: string]: any }>
+  extends PivotAggsConfigWithUiBase {
   /** Form component */
   AggFormComponent: FC<{
     aggConfig: Partial<T>;
     onChange: (arg: Partial<T>) => void;
     selectedField: string;
+    isValid?: boolean;
   }>;
   /** Aggregation specific configuration */
   aggConfig: Partial<T>;
   /** Set UI configuration from ES aggregation definition */
-  setUiConfigFromEs: (arg: { [key: string]: any }) => void;
+  setUiConfigFromEs: (arg: ESConfig) => void;
   /** Converts UI agg config form to ES agg request object */
-  getEsAggConfig: () => { [key: string]: any } | null;
+  getEsAggConfig: () => ESConfig | null;
   /** Indicates if the configuration is valid */
   isValid: () => boolean;
   /** Provides aggregation name generated based on the configuration */
@@ -242,7 +244,7 @@ export type PivotAggsConfigWithUiSupport =
   | PivotAggsConfigTerms
   | PivotAggsConfigWithExtendedForm;
 
-export function isPivotAggsConfigWithUiSupport(arg: unknown): arg is PivotAggsConfigWithUiSupport {
+export function isPivotAggsConfigWithUiBase(arg: unknown): arg is PivotAggsConfigWithUiBase {
   return (
     isPopulatedObject(arg, ['agg', 'aggName', 'dropDownName', 'field']) &&
     isPivotSupportedAggs(arg.agg)
@@ -256,8 +258,7 @@ type PivotAggsConfigWithExtendedForm = PivotAggsConfigFilter | PivotAggsConfigTo
 
 export function isPivotAggsWithExtendedForm(arg: unknown): arg is PivotAggsConfigWithExtendedForm {
   return (
-    (isPopulatedObject(arg) && arg.hasOwnProperty('setUiConfigFromEs')) ||
-    isPopulatedObject(arg, ['AggFormComponent'])
+    isPopulatedObject(arg, ['setUiConfigFromEs']) || isPopulatedObject(arg, ['AggFormComponent'])
   );
 }
 
@@ -288,15 +289,15 @@ export type PivotAggsConfigDict = Dictionary<PivotAggsConfig>;
 export function getEsAggFromAggConfig(
   pivotAggsConfig: PivotAggsConfigBase | PivotAggsConfigWithExtendedForm
 ): PivotAgg | null {
-  let esAgg: { [key: string]: any } | null = { ...pivotAggsConfig };
+  let esAgg: { [key: string]: any } = { ...pivotAggsConfig };
 
   delete esAgg.agg;
   delete esAgg.aggName;
   delete esAgg.dropDownName;
   delete esAgg.parentAgg;
 
   if (isPivotAggsWithExtendedForm(pivotAggsConfig)) {
-    esAgg = pivotAggsConfig.getEsAggConfig();
+    esAgg = pivotAggsConfig.getEsAggConfig() as PivotAgg;
 
     if (esAgg === null) {
       return null;
@@ -305,16 +306,16 @@ export function getEsAggFromAggConfig(
 
   const result = {
     [pivotAggsConfig.agg]: esAgg,
-  };
+  } as PivotAgg;
 
   if (
-    isPivotAggsConfigWithUiSupport(pivotAggsConfig) &&
+    isPivotAggsConfigWithUiBase(pivotAggsConfig) &&
     pivotAggsConfig.subAggs !== undefined &&
     Object.keys(pivotAggsConfig.subAggs).length > 0
   ) {
     result.aggs = {};
     for (const subAggConfig of Object.values(pivotAggsConfig.subAggs)) {
-      result.aggs[subAggConfig.aggName] = getEsAggFromAggConfig(subAggConfig);
+      result.aggs[subAggConfig.aggName] = getEsAggFromAggConfig(subAggConfig) as PivotAgg;
     }
   }
 

x-pack/plugins/transform/public/app/sections/create_transform/components/aggregation_list/agg_label_form.tsx

@@ -14,7 +14,7 @@ import { EuiButtonIcon, EuiFlexGroup, EuiFlexItem, EuiPopover, EuiTextColor } fr
 import { AggName } from '../../../../../../common/types/aggregations';
 
 import {
-  isPivotAggsConfigWithUiSupport,
+  isPivotAggsConfigWithUiBase,
   PivotAggsConfig,
   PivotAggsConfigWithUiSupportDict,
 } from '../../../../common';
@@ -50,7 +50,7 @@ export const AggLabelForm: React.FC<Props> = ({
   const helperText = isPivotAggsWithExtendedForm(item) && item.helperText && item.helperText();
 
   const isSubAggSupported =
-    isPivotAggsConfigWithUiSupport(item) &&
+    isPivotAggsConfigWithUiBase(item) &&
     item.isSubAggsSupported &&
     (isPivotAggsWithExtendedForm(item) ? item.isValid() : true);