[ML] Add support for per-partition categorization jobs

x-pack/plugins/ml/common/types/anomalies.ts

@@ -57,3 +57,20 @@ export interface AnomaliesTableRecord {
 }
 
 export type PartitionFieldsType = typeof PARTITION_FIELDS[number];
+
+export interface AnomalyCategorizerStatsDoc {
+  [key: string]: any;
+  job_id: string;
+  result_type: 'categorizer_stats';
+  partition_field_name?: string;
+  partition_field_value?: string;
+  categorized_doc_count: number;
+  total_category_count: number;
+  frequent_category_count: number;
+  rare_category_count: number;
+  dead_category_count: number;
+  failed_category_count: number;
+  categorization_status: 'ok' | 'warn';
+  log_time: number;
+  timestamp: number;
+}

x-pack/plugins/ml/common/types/anomaly_detection_jobs/job.ts

@@ -93,6 +93,6 @@ export interface CustomRule {
 }
 
 export interface PerPartitionCategorization {
-  enabled: boolean;
+  enabled?: boolean;
   stop_on_warn?: boolean;
 }

x-pack/plugins/ml/public/application/explorer/explorer.js

@@ -30,6 +30,7 @@ import {
   EuiPanel,
   EuiAccordion,
   EuiBadge,
+  EuiCallOut,
 } from '@elastic/eui';
 
 import { AnnotationFlyout } from '../components/annotations/annotation_flyout';
@@ -204,7 +205,7 @@ export class Explorer extends React.Component {
   updateLanguage = (language) => this.setState({ language });
 
   render() {
-    const { showCharts, severity } = this.props;
+    const { showCharts, severity, stoppedPartitions } = this.props;
 
     const {
       annotations,
@@ -297,6 +298,21 @@ export class Explorer extends React.Component {
 
           <div className={mainColumnClasses}>
             <EuiSpacer size="m" />
+
+            {stoppedPartitions && (
+              <EuiCallOut
+                size={'s'}
+                title={i18n.translate('xpack.ml.explorer.stoppedPartitionsExistCallout', {
+                  defaultMessage:
+                    'There may be fewer results than there could have been because stop_on_warn is turned on. Both categorization and subsequent anomaly detection have stopped for some partitions in {jobsWithStoppedPartitions, plural, one {job} other {jobs}} [{stoppedPartitions}] where the categorization status has changed to warn.',
+                  values: {
+                    jobsWithStoppedPartitions: stoppedPartitions.length,
+                    stoppedPartitions: stoppedPartitions.join(', '),
+                  },
+                })}
+              />
+            )}
+
             <AnomalyTimeline
               explorerState={this.props.explorerState}
               setSelectedCells={this.props.setSelectedCells}

x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/categorization_job_creator.ts

@@ -75,6 +75,11 @@ export class CategorizationJobCreator extends JobCreator {
   private _createDetector(agg: Aggregation, field: Field) {
     const dtr: Detector = createBasicDetector(agg, field);
     dtr.by_field_name = mlCategory.id;
+
+    // API requires if per_partition_categorization is enabled, add partition field to the detector
+    if (this.perPartitionCategorization && this.categorizationPerPartitionField !== null) {
+      dtr.partition_field_name = this.categorizationPerPartitionField;
+    }
     this._addDetector(dtr, agg, mlCategory);
   }
 

x-pack/plugins/ml/public/application/jobs/new_job/common/job_creator/job_creator.ts

@@ -57,6 +57,7 @@ export class JobCreator {
   private _stopAllRefreshPolls: {
     stop: boolean;
   } = { stop: false };
+  private _partitionField: string | null;
 
   protected _wizardInitialized$ = new BehaviorSubject<boolean>(false);
   public wizardInitialized$ = this._wizardInitialized$.asObservable();
@@ -81,6 +82,7 @@ export class JobCreator {
     }
 
     this._datafeed_config.query = query;
+    this._partitionField = null;
   }
 
   public get type(): JOB_TYPE {
@@ -622,6 +624,56 @@ export class JobCreator {
     return JSON.stringify(this._datafeed_config, null, 2);
   }
 
+  private _initPerPartitionCategorization() {
+    if (this._job_config.analysis_config.per_partition_categorization === undefined) {
+      this._job_config.analysis_config.per_partition_categorization = {};
+    }
+    if (this._job_config.analysis_config.per_partition_categorization?.enabled === undefined) {
+      this._job_config.analysis_config.per_partition_categorization!.enabled = false;
+    }
+    if (this._job_config.analysis_config.per_partition_categorization?.stop_on_warn === undefined) {
+      this._job_config.analysis_config.per_partition_categorization!.stop_on_warn = false;
+    }
+  }
+
+  public get perPartitionCategorization() {
+    return this._job_config.analysis_config.per_partition_categorization?.enabled === true;
+  }
+
+  public set perPartitionCategorization(enabled: boolean) {
+    this._initPerPartitionCategorization();
+    this._job_config.analysis_config.per_partition_categorization!.enabled = enabled;
+  }
+
+  public get partitionStopOnWarn() {
+    return this._job_config.analysis_config.per_partition_categorization?.stop_on_warn === true;
+  }
+
+  public set partitionStopOnWarn(enabled: boolean) {
+    this._initPerPartitionCategorization();
+    this._job_config.analysis_config.per_partition_categorization!.stop_on_warn = enabled;
+  }
+
+  public get categorizationPerPartitionField() {
+    return this._partitionField;
+  }
+
+  public set categorizationPerPartitionField(fieldName: string | null) {
+    if (fieldName === null) {
+      this._detectors.forEach((detector) => {
+        delete detector.partition_field_name;
+      });
+      this._partitionField = null;
+    } else {
+      if (this._partitionField !== fieldName) {
+        this._partitionField = fieldName;
+        this._detectors.forEach((detector) => {
+          detector.partition_field_name = fieldName;
+        });
+      }
+    }
+  }
+
   protected _overrideConfigs(job: Job, datafeed: Datafeed) {
     this._job_config = job;
     this._datafeed_config = datafeed;

x-pack/plugins/ml/public/application/jobs/new_job/pages/components/pick_fields_step/components/advanced_view/extra.tsx

@@ -4,13 +4,19 @@
  * you may not use this file except in compliance with the Elastic License.
  */
 
-import React, { Fragment, FC } from 'react';
+import React, { Fragment, FC, useContext } from 'react';
 import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
 
 import { SummaryCountField } from '../summary_count_field';
 import { CategorizationField } from '../categorization_field';
+import { CategorizationPerPartitionField } from '../categorization_partition_field';
+import { JobCreatorContext } from '../../../job_creator_context';
+import { isAdvancedJobCreator } from '../../../../../common/job_creator';
 
 export const ExtraSettings: FC = () => {
+  const { jobCreator } = useContext(JobCreatorContext);
+  const showCategorizationPerPartitionField =
+    isAdvancedJobCreator(jobCreator) && jobCreator.categorizationFieldName !== null;
   return (
     <Fragment>
       <EuiFlexGroup gutterSize="xl">
@@ -21,6 +27,13 @@ export const ExtraSettings: FC = () => {
           <SummaryCountField />
         </EuiFlexItem>
       </EuiFlexGroup>
+      {showCategorizationPerPartitionField && (
+        <EuiFlexGroup gutterSize="xl">
+          <EuiFlexItem>
+            <CategorizationPerPartitionField />
+          </EuiFlexItem>
+        </EuiFlexGroup>
+      )}
     </Fragment>
   );
 };

x-pack/plugins/ml/public/application/jobs/new_job/pages/components/pick_fields_step/components/categorization_partition_field/categorization_per_partition.tsx

@@ -0,0 +1,97 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React, { FC, useContext, useEffect, useState } from 'react';
+import { EuiFormRow } from '@elastic/eui';
+import { i18n } from '@kbn/i18n';
+import { JobCreatorContext } from '../../../job_creator_context';
+import {
+  AdvancedJobCreator,
+  CategorizationJobCreator,
+  isCategorizationJobCreator,
+} from '../../../../../common/job_creator';
+import { newJobCapsService } from '../../../../../../../services/new_job_capabilities_service';
+import { Description } from './description';
+import { CategorizationPerPartitionSwitch } from './categorization_per_partition_switch';
+import { CategorizationPerPartitionFieldSelect } from './categorization_per_partition_input';
+import { CategorizationPerPartitionStopOnWarnSwitch } from './categorization_stop_on_warn_switch';
+
+export const CategorizationPerPartitionField: FC = () => {
+  const { jobCreator: jc, jobCreatorUpdate, jobCreatorUpdated } = useContext(JobCreatorContext);
+  const jobCreator = jc as AdvancedJobCreator | CategorizationJobCreator;
+  const [enablePerPartitionCategorization, setEnablePerPartitionCategorization] = useState(false);
+  const [categorizationPartitionFieldName, setCategorizationPartitionFieldName] = useState<
+    string | null
+  >(jobCreator.categorizationPerPartitionField);
+
+  const { catFields } = newJobCapsService;
+
+  const filteredCategories = catFields.filter((c) => c.id !== jobCreator.categorizationFieldName);
+
+  useEffect(() => {
+    jobCreator.categorizationPerPartitionField = categorizationPartitionFieldName;
+    jobCreatorUpdate();
+  }, [categorizationPartitionFieldName]);
+
+  useEffect(() => {
+    // set the first item in category as partition field by default
+    // because API requires partition_field to be defined in each detector with mlcategory
+    // if per-partition categorization is enabled
+    if (
+      jobCreator.perPartitionCategorization &&
+      jobCreator.categorizationPerPartitionField === null
+    ) {
+      jobCreator.categorizationPerPartitionField = filteredCategories[0].id;
+    }
+    setCategorizationPartitionFieldName(jobCreator.categorizationPerPartitionField);
+    setEnablePerPartitionCategorization(jobCreator.perPartitionCategorization);
+  }, [jobCreatorUpdated]);
+
+  const isCategorizationJob = isCategorizationJobCreator(jobCreator);
+  return (
+    <Description isOptional={isCategorizationJob === false}>
+      <EuiFormRow
+        label={i18n.translate(
+          'xpack.ml.newJob.wizard.extraStep.categorizationJob.perPartitionCategorizationLabel',
+          {
+            defaultMessage: 'Enable per-partition categorization',
+          }
+        )}
+      >
+        <CategorizationPerPartitionSwitch />
+      </EuiFormRow>
+
+      {enablePerPartitionCategorization && (
+        <>
+          <EuiFormRow
+            label={i18n.translate(
+              'xpack.ml.newJob.wizard.extraStep.categorizationJob.stopOnWarnLabel',
+              {
+                defaultMessage: 'Stop on warn',
+              }
+            )}
+          >
+            <CategorizationPerPartitionStopOnWarnSwitch />
+          </EuiFormRow>
+          <EuiFormRow
+            label={i18n.translate(
+              'xpack.ml.newJob.wizard.extraStep.categorizationJob.partitionFieldLabel',
+              {
+                defaultMessage: 'Partition field',
+              }
+            )}
+          >
+            <CategorizationPerPartitionFieldSelect
+              fields={filteredCategories}
+              changeHandler={setCategorizationPartitionFieldName}
+              selectedField={categorizationPartitionFieldName || ''}
+            />
+          </EuiFormRow>
+        </>
+      )}
+    </Description>
+  );
+};

x-pack/plugins/ml/public/application/jobs/new_job/pages/components/pick_fields_step/components/categorization_partition_field/categorization_per_partition_input.tsx

@@ -0,0 +1,54 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React, { FC, useContext } from 'react';
+import { EuiComboBox, EuiComboBoxOptionOption } from '@elastic/eui';
+
+import { JobCreatorContext } from '../../../job_creator_context';
+import { Field } from '../../../../../../../../../common/types/fields';
+import { createFieldOptions } from '../../../../../common/job_creator/util/general';
+
+interface Props {
+  fields: Field[];
+  changeHandler(i: string | null): void;
+  selectedField: string | null;
+}
+
+export const CategorizationPerPartitionFieldSelect: FC<Props> = ({
+  fields,
+  changeHandler,
+  selectedField,
+}) => {
+  const { jobCreator } = useContext(JobCreatorContext);
+  const options: EuiComboBoxOptionOption[] = [
+    ...createFieldOptions(fields, jobCreator.additionalFields),
+  ];
+
+  const selection: EuiComboBoxOptionOption[] = [];
+  if (selectedField !== null) {
+    selection.push({ label: selectedField });
+  }
+
+  function onChange(selectedOptions: EuiComboBoxOptionOption[]) {
+    const option = selectedOptions[0];
+    if (typeof option !== 'undefined') {
+      changeHandler(option.label);
+    } else {
+      changeHandler(null);
+    }
+  }
+
+  return (
+    <EuiComboBox
+      singleSelection={{ asPlainText: true }}
+      options={options}
+      selectedOptions={selection}
+      onChange={onChange}
+      isClearable={true}
+      data-test-subj="mlCategorizationPerPartitionFieldNameSelect"
+    />
+  );
+};

x-pack/plugins/ml/public/application/jobs/new_job/pages/components/pick_fields_step/components/categorization_partition_field/categorization_per_partition_switch.tsx

@@ -0,0 +1,45 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import React, { FC, useContext, useEffect, useState } from 'react';
+import { i18n } from '@kbn/i18n';
+import { EuiSwitch } from '@elastic/eui';
+import { JobCreatorContext } from '../../../job_creator_context';
+import { AdvancedJobCreator, CategorizationJobCreator } from '../../../../../common/job_creator';
+
+export const CategorizationPerPartitionSwitch: FC = () => {
+  const { jobCreator: jc, jobCreatorUpdate } = useContext(JobCreatorContext);
+  const jobCreator = jc as AdvancedJobCreator | CategorizationJobCreator;
+  const [enablePerPartitionCategorization, setEnablePerPartitionCategorization] = useState(
+    jobCreator.perPartitionCategorization
+  );
+
+  const toggleEnablePerPartitionCategorization = () =>
+    setEnablePerPartitionCategorization(!enablePerPartitionCategorization);
+
+  useEffect(() => {
+    // also turn off stop on warn if per_partition_categorization is turned off
+    if (enablePerPartitionCategorization === false) {
+      jobCreator.partitionStopOnWarn = false;
+    }
+
+    jobCreator.perPartitionCategorization = enablePerPartitionCategorization;
+    jobCreatorUpdate();
+  }, [enablePerPartitionCategorization]);
+
+  return (
+    <EuiSwitch
+      name="switch"
+      disabled={false}
+      checked={enablePerPartitionCategorization}
+      onChange={toggleEnablePerPartitionCategorization}
+      data-test-subj="mlJobWizardSwitchCategorizationPerPartitionField"
+      label={i18n.translate('xpack.ml.newJob.wizard.perPartitionCategorizationSwitchLabel', {
+        defaultMessage: 'Enable per-partition categorization',
+      })}
+    />
+  );
+};