[ML] Fix application of runtime fields to index search. (#98718)

- A previous PR introduced a regression where only runtime fields would show up in the analytics wizard's source index preview. The code for transforms and analytics is a bit different so this regression didn't occur in transforms.
- This PR fixes the problem and cleans up use_index_data.ts for the analytics wizard to remove some duplicate code to determine runtime field mappings.
- Async fetch functions have been refactored to named function expressions and moved inside their corresponding useEffect calls (this change caused most of the diff).
- combinedRuntimeMappings has been moved to an outer useMemo so it doesn't have to be generated in multiple places.
- getIndexData has been renamed to fetchIndexData to indicate it's an async call getting remote data and to be in line with the other function names.

x-pack/plugins/ml/public/application/data_frame_analytics/pages/analytics_creation/hooks/use_index_data.ts

@@ -50,9 +50,10 @@ function getRuntimeFieldColumns(runtimeMappings: RuntimeMappings) {
   });
 }
 
-function getInitialColumns(indexPattern: IndexPattern, fieldsFilter: string[]) {
+function getIndexPatternColumns(indexPattern: IndexPattern, fieldsFilter: string[]) {
   const { fields } = newJobCapsServiceAnalytics;
-  const columns = fields
+
+  return fields
     .filter((field) => fieldsFilter.includes(field.name))
     .map((field) => {
       const schema =
@@ -65,26 +66,6 @@ function getInitialColumns(indexPattern: IndexPattern, fieldsFilter: string[]) {
         isRuntimeFieldColumn: false,
       };
     });
-
-  // Add runtime fields defined in index pattern to columns
-  if (indexPattern) {
-    const computedFields = indexPattern?.getComputedFields();
-
-    if (isRuntimeMappings(computedFields.runtimeFields)) {
-      Object.keys(computedFields.runtimeFields).forEach((runtimeField) => {
-        const schema = getDataGridSchemaFromESFieldType(
-          computedFields.runtimeFields[runtimeField].type
-        );
-        columns.push({
-          id: runtimeField,
-          schema,
-          isExpandable: schema !== 'boolean',
-          isRuntimeFieldColumn: true,
-        });
-      });
-    }
-  }
-  return columns;
 }
 
 export const useIndexData = (
@@ -93,57 +74,71 @@ export const useIndexData = (
   toastNotifications: CoreSetup['notifications']['toasts'],
   runtimeMappings?: RuntimeMappings
 ): UseIndexDataReturnType => {
-  const [indexPatternFields, setIndexPatternFields] = useState<string[]>();
-
   // Fetch 500 random documents to determine populated fields.
   // This is a workaround to avoid passing potentially thousands of unpopulated fields
   // (for example, as part of filebeat/metricbeat/ECS based indices)
   // to the data grid component which would significantly slow down the page.
-  const fetchDataGridSampleDocuments = async function () {
-    setErrorMessage('');
-    setStatus(INDEX_STATUS.LOADING);
-
-    const esSearchRequest = {
-      index: indexPattern.title,
-      body: {
-        fields: ['*'],
-        _source: false,
-        query: {
-          function_score: {
-            query: { match_all: {} },
-            random_score: {},
+  const [indexPatternFields, setIndexPatternFields] = useState<string[]>();
+  useEffect(() => {
+    async function fetchDataGridSampleDocuments() {
+      setErrorMessage('');
+      setStatus(INDEX_STATUS.LOADING);
+
+      const esSearchRequest = {
+        index: indexPattern.title,
+        body: {
+          fields: ['*'],
+          _source: false,
+          query: {
+            function_score: {
+              query: { match_all: {} },
+              random_score: {},
+            },
           },
+          size: 500,
         },
-        size: 500,
-      },
-    };
-
-    try {
-      const resp: IndexSearchResponse = await ml.esSearch(esSearchRequest);
-      const docs = resp.hits.hits.map((d) => getProcessedFields(d.fields ?? {}));
-
-      // Get all field names for each returned doc and flatten it
-      // to a list of unique field names used across all docs.
-      const allKibanaIndexPatternFields = getFieldsFromKibanaIndexPattern(indexPattern);
-      const populatedFields = [...new Set(docs.map(Object.keys).flat(1))]
-        .filter((d) => allKibanaIndexPatternFields.includes(d))
-        .sort();
-
-      setStatus(INDEX_STATUS.LOADED);
-      setIndexPatternFields(populatedFields);
-    } catch (e) {
-      setErrorMessage(extractErrorMessage(e));
-      setStatus(INDEX_STATUS.ERROR);
+      };
+
+      try {
+        const resp: IndexSearchResponse = await ml.esSearch(esSearchRequest);
+        const docs = resp.hits.hits.map((d) => getProcessedFields(d.fields ?? {}));
+
+        // Get all field names for each returned doc and flatten it
+        // to a list of unique field names used across all docs.
+        const allKibanaIndexPatternFields = getFieldsFromKibanaIndexPattern(indexPattern);
+        const populatedFields = [...new Set(docs.map(Object.keys).flat(1))]
+          .filter((d) => allKibanaIndexPatternFields.includes(d))
+          .sort();
+
+        setStatus(INDEX_STATUS.LOADED);
+        setIndexPatternFields(populatedFields);
+      } catch (e) {
+        setErrorMessage(extractErrorMessage(e));
+        setStatus(INDEX_STATUS.ERROR);
+      }
     }
-  };
 
-  useEffect(() => {
     fetchDataGridSampleDocuments();
   }, []);
 
-  const [columns, setColumns] = useState<MLEuiDataGridColumn[]>(
-    getInitialColumns(indexPattern, indexPatternFields ?? [])
+  // To be used for data grid column selection
+  // and will be applied to doc and chart queries.
+  const combinedRuntimeMappings = useMemo(
+    () => getCombinedRuntimeMappings(indexPattern, runtimeMappings),
+    [indexPattern, runtimeMappings]
   );
+
+  // Available data grid columns, will be a combination of index pattern and runtime fields.
+  const [columns, setColumns] = useState<MLEuiDataGridColumn[]>([]);
+  useEffect(() => {
+    if (Array.isArray(indexPatternFields)) {
+      setColumns([
+        ...getIndexPatternColumns(indexPattern, indexPatternFields),
+        ...(combinedRuntimeMappings ? getRuntimeFieldColumns(combinedRuntimeMappings) : []),
+      ]);
+    }
+  }, [indexPattern, indexPatternFields, combinedRuntimeMappings]);
+
   const dataGrid = useDataGrid(columns);
 
   const {
@@ -163,95 +158,87 @@ export const useIndexData = (
     // custom comparison
   }, [JSON.stringify(query)]);
 
-  const getIndexData = async function () {
-    setErrorMessage('');
-    setStatus(INDEX_STATUS.LOADING);
-
-    const combinedRuntimeMappings = getCombinedRuntimeMappings(indexPattern, runtimeMappings);
-
-    const sort: EsSorting = sortingColumns.reduce((s, column) => {
-      s[column.id] = { order: column.direction };
-      return s;
-    }, {} as EsSorting);
-    const esSearchRequest = {
-      index: indexPattern.title,
-      body: {
-        query,
-        from: pagination.pageIndex * pagination.pageSize,
-        size: pagination.pageSize,
-        fields: ['*'],
-        _source: false,
-        ...(Object.keys(sort).length > 0 ? { sort } : {}),
-        ...(isRuntimeMappings(combinedRuntimeMappings)
-          ? { runtime_mappings: combinedRuntimeMappings }
-          : {}),
-      },
-    };
-
-    try {
-      const resp: IndexSearchResponse = await ml.esSearch(esSearchRequest);
-      const docs = resp.hits.hits.map((d) => getProcessedFields(d.fields ?? {}));
-
-      if (isRuntimeMappings(runtimeMappings)) {
-        // remove old runtime field from columns
-        const updatedColumns = columns.filter((col) => col.isRuntimeFieldColumn === false);
-        setColumns([
-          ...updatedColumns,
-          ...(combinedRuntimeMappings ? getRuntimeFieldColumns(combinedRuntimeMappings) : []),
-        ]);
-      } else {
-        setColumns(getInitialColumns(indexPattern, indexPatternFields ?? []));
+  useEffect(() => {
+    async function fetchIndexData() {
+      setErrorMessage('');
+      setStatus(INDEX_STATUS.LOADING);
+
+      const sort: EsSorting = sortingColumns.reduce((s, column) => {
+        s[column.id] = { order: column.direction };
+        return s;
+      }, {} as EsSorting);
+      const esSearchRequest = {
+        index: indexPattern.title,
+        body: {
+          query,
+          from: pagination.pageIndex * pagination.pageSize,
+          size: pagination.pageSize,
+          fields: [
+            ...(indexPatternFields ?? []),
+            ...(isRuntimeMappings(combinedRuntimeMappings)
+              ? Object.keys(combinedRuntimeMappings)
+              : []),
+          ],
+          _source: false,
+          ...(Object.keys(sort).length > 0 ? { sort } : {}),
+          ...(isRuntimeMappings(combinedRuntimeMappings)
+            ? { runtime_mappings: combinedRuntimeMappings }
+            : {}),
+        },
+      };
+
+      try {
+        const resp: IndexSearchResponse = await ml.esSearch(esSearchRequest);
+        const docs = resp.hits.hits.map((d) => getProcessedFields(d.fields ?? {}));
+
+        setRowCount(typeof resp.hits.total === 'number' ? resp.hits.total : resp.hits.total.value);
+        setRowCountRelation(
+          typeof resp.hits.total === 'number'
+            ? ('eq' as estypes.TotalHitsRelation)
+            : resp.hits.total.relation
+        );
+        setTableItems(docs);
+        setStatus(INDEX_STATUS.LOADED);
+      } catch (e) {
+        setErrorMessage(extractErrorMessage(e));
+        setStatus(INDEX_STATUS.ERROR);
       }
-      setRowCount(typeof resp.hits.total === 'number' ? resp.hits.total : resp.hits.total.value);
-      setRowCountRelation(
-        typeof resp.hits.total === 'number'
-          ? ('eq' as estypes.TotalHitsRelation)
-          : resp.hits.total.relation
-      );
-      setTableItems(docs);
-      setStatus(INDEX_STATUS.LOADED);
-    } catch (e) {
-      setErrorMessage(extractErrorMessage(e));
-      setStatus(INDEX_STATUS.ERROR);
     }
-  };
 
-  useEffect(() => {
-    if (query !== undefined) {
-      getIndexData();
+    if (indexPatternFields !== undefined && query !== undefined) {
+      fetchIndexData();
     }
     // custom comparison
   }, [
     indexPattern.title,
     indexPatternFields,
-    JSON.stringify([query, pagination, sortingColumns, runtimeMappings]),
+    JSON.stringify([query, pagination, sortingColumns, combinedRuntimeMappings]),
   ]);
 
   const dataLoader = useMemo(() => new DataLoader(indexPattern, toastNotifications), [
     indexPattern,
   ]);
 
-  const fetchColumnChartsData = async function (fieldHistogramsQuery: Record<string, any>) {
-    const combinedRuntimeMappings = getCombinedRuntimeMappings(indexPattern, runtimeMappings);
-    try {
-      const columnChartsData = await dataLoader.loadFieldHistograms(
-        columns
-          .filter((cT) => dataGrid.visibleColumns.includes(cT.id))
-          .map((cT) => ({
-            fieldName: cT.id,
-            type: getFieldType(cT.schema),
-          })),
-        fieldHistogramsQuery,
-        DEFAULT_SAMPLER_SHARD_SIZE,
-        combinedRuntimeMappings
-      );
-      dataGrid.setColumnCharts(columnChartsData);
-    } catch (e) {
-      showDataGridColumnChartErrorMessageToast(e, toastNotifications);
+  useEffect(() => {
+    async function fetchColumnChartsData(fieldHistogramsQuery: Record<string, any>) {
+      try {
+        const columnChartsData = await dataLoader.loadFieldHistograms(
+          columns
+            .filter((cT) => dataGrid.visibleColumns.includes(cT.id))
+            .map((cT) => ({
+              fieldName: cT.id,
+              type: getFieldType(cT.schema),
+            })),
+          fieldHistogramsQuery,
+          DEFAULT_SAMPLER_SHARD_SIZE,
+          combinedRuntimeMappings
+        );
+        dataGrid.setColumnCharts(columnChartsData);
+      } catch (e) {
+        showDataGridColumnChartErrorMessageToast(e, toastNotifications);
+      }
     }
-  };
 
-  useEffect(() => {
     if (dataGrid.chartsVisible && query !== undefined) {
       fetchColumnChartsData(query);
     }