[8.x] [ML] Transforms: Improve data grid memoization. (#195394) (#195975)

# Backport This will backport the following commits from `main` to `8.x`: - [[ML] Transforms: Improve data grid memoization. (#195394)](https://github.com/elastic/kibana/pull/195394)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Walter Rafelsberger <walter.rafelsberger@elastic.co>
2025-04-24 01:38:56 -04:00 · 2024-10-12 06:59:01 +11:00 · 2024-10-12 06:59:01 +11:00 · c9b0d86381
commit c9b0d86381
parent bf56b4ea3a
6 changed files with 163 additions and 175 deletions
--- a/x-pack/packages/ml/data_grid/components/data_grid.tsx
+++ b/x-pack/packages/ml/data_grid/components/data_grid.tsx
@ -154,7 +154,9 @@ export const DataGrid: FC<Props> = memo(

    const wrapperEl = useRef<HTMLDivElement>(null);

-    if (status === INDEX_STATUS.LOADED && data.length === 0) {
+    // `UNUSED` occurs if we were not able to identify populated fields, because
+    // then the query to fetch actual docs would not have triggered yet.
+    if ((status === INDEX_STATUS.UNUSED || status === INDEX_STATUS.LOADED) && data.length === 0) {
      return (
        <div data-test-subj={`${dataTestSubj} empty`}>
          {isWithHeader(props) && <DataGridTitle title={props.title} />}
--- a/x-pack/packages/ml/data_grid/hooks/use_data_grid.tsx
+++ b/x-pack/packages/ml/data_grid/hooks/use_data_grid.tsx
@ -168,35 +168,54 @@ export const useDataGrid = (
    }
  }, [chartsVisible, rowCount, rowCountRelation]);

-  return {
-    ccsWarning,
-    chartsVisible,
-    chartsButtonVisible: true,
-    columnsWithCharts,
-    errorMessage,
-    invalidSortingColumnns,
-    noDataMessage,
-    onChangeItemsPerPage,
-    onChangePage,
-    onSort,
-    pagination,
-    resetPagination,
-    rowCount,
-    rowCountRelation,
-    setColumnCharts,
-    setCcsWarning,
-    setErrorMessage,
-    setNoDataMessage,
-    setPagination,
-    setRowCountInfo,
-    setSortingColumns,
-    setStatus,
-    setTableItems,
-    setVisibleColumns,
-    sortingColumns,
-    status,
-    tableItems,
-    toggleChartVisibility,
-    visibleColumns,
-  };
+  return useMemo(
+    () => ({
+      ccsWarning,
+      chartsVisible,
+      chartsButtonVisible: true,
+      columnsWithCharts,
+      errorMessage,
+      invalidSortingColumnns,
+      noDataMessage,
+      onChangeItemsPerPage,
+      onChangePage,
+      onSort,
+      pagination,
+      resetPagination,
+      rowCount,
+      rowCountRelation,
+      setColumnCharts,
+      setCcsWarning,
+      setErrorMessage,
+      setNoDataMessage,
+      setPagination,
+      setRowCountInfo,
+      setSortingColumns,
+      setStatus,
+      setTableItems,
+      setVisibleColumns,
+      sortingColumns,
+      status,
+      tableItems,
+      toggleChartVisibility,
+      visibleColumns,
+    }),
+    // custom comparison
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+    [
+      ccsWarning,
+      chartsVisible,
+      columnsWithCharts,
+      errorMessage,
+      invalidSortingColumnns,
+      noDataMessage,
+      pagination,
+      rowCount,
+      rowCountRelation,
+      sortingColumns,
+      status,
+      tableItems,
+      visibleColumns,
+    ]
+  );
 };
--- a/x-pack/plugins/transform/public/app/mocks/app_dependencies.tsx
+++ b/x-pack/plugins/transform/public/app/mocks/app_dependencies.tsx
@ -67,6 +67,9 @@ dataStart.search.search = jest.fn(({ params }: IKibanaSearchRequest) => {
  });
 }) as ISearchGeneric;

+// Replace mock to support tests for `use_index_data`.
+coreSetup.http.post = jest.fn().mockResolvedValue([]);
+
 const appDependencies: AppDependencies = {
  analytics: coreStart.analytics,
  application: coreStart.application,
--- a/x-pack/plugins/transform/public/app/hooks/use_index_data.test.tsx
+++ b/x-pack/plugins/transform/public/app/hooks/use_index_data.test.tsx
@ -12,7 +12,7 @@ import { renderHook } from '@testing-library/react-hooks';

 import { __IntlProvider as IntlProvider } from '@kbn/i18n-react';
 import type { CoreSetup } from '@kbn/core/public';
-import { DataGrid, type UseIndexDataReturnType } from '@kbn/ml-data-grid';
+import { DataGrid, type UseIndexDataReturnType, INDEX_STATUS } from '@kbn/ml-data-grid';
 import type { RuntimeMappings } from '@kbn/ml-runtime-field-utils';
 import type { SimpleQuery } from '@kbn/ml-query-utils';

@ -40,6 +40,37 @@ const runtimeMappings: RuntimeMappings = {
 const queryClient = new QueryClient();

 describe('Transform: useIndexData()', () => {
+  test('empty populatedFields does not trigger loading', async () => {
+    const wrapper: FC<PropsWithChildren<unknown>> = ({ children }) => (
+      <QueryClientProvider client={queryClient}>
+        <IntlProvider locale="en">{children}</IntlProvider>
+      </QueryClientProvider>
+    );
+
+    const { result, waitForNextUpdate } = renderHook(
+      () =>
+        useIndexData({
+          dataView: {
+            id: 'the-id',
+            getIndexPattern: () => 'the-index-pattern',
+            fields: [],
+          } as unknown as SearchItems['dataView'],
+          query,
+          combinedRuntimeMappings: runtimeMappings,
+          populatedFields: [],
+        }),
+      { wrapper }
+    );
+
+    const IndexObj: UseIndexDataReturnType = result.current;
+
+    await waitForNextUpdate();
+
+    expect(IndexObj.errorMessage).toBe('');
+    expect(IndexObj.status).toBe(INDEX_STATUS.UNUSED);
+    expect(IndexObj.tableItems).toEqual([]);
+  });
+
  test('dataView set triggers loading', async () => {
    const wrapper: FC<PropsWithChildren<unknown>> = ({ children }) => (
      <QueryClientProvider client={queryClient}>
@ -49,15 +80,29 @@ describe('Transform: useIndexData()', () => {

    const { result, waitForNextUpdate } = renderHook(
      () =>
-        useIndexData(
-          {
+        useIndexData({
+          dataView: {
            id: 'the-id',
            getIndexPattern: () => 'the-index-pattern',
-            fields: [],
+            metaFields: [],
+            // minimal mock of DataView fields (array with getByName method)
+            fields: new (class DataViewFields extends Array<{ name: string }> {
+              getByName(id: string) {
+                return this.find((d) => d.name === id);
+              }
+            })(
+              {
+                name: 'the-populated-field',
+              },
+              {
+                name: 'the-unpopulated-field',
+              }
+            ),
          } as unknown as SearchItems['dataView'],
          query,
-          runtimeMappings
-        ),
+          combinedRuntimeMappings: runtimeMappings,
+          populatedFields: ['the-populated-field'],
+        }),
      { wrapper }
    );

@ -66,7 +111,7 @@ describe('Transform: useIndexData()', () => {
    await waitForNextUpdate();

    expect(IndexObj.errorMessage).toBe('');
-    expect(IndexObj.status).toBe(1);
+    expect(IndexObj.status).toBe(INDEX_STATUS.LOADING);
    expect(IndexObj.tableItems).toEqual([]);
  });
 });
@ -81,7 +126,12 @@ describe('Transform: <DataGrid /> with useIndexData()', () => {

    const Wrapper = () => {
      const props = {
-        ...useIndexData(dataView, { match_all: {} }, runtimeMappings),
+        ...useIndexData({
+          dataView,
+          query: { match_all: {} },
+          combinedRuntimeMappings: runtimeMappings,
+          populatedFields: ['the-populated-field'],
+        }),
        copyToClipboard: 'the-copy-to-clipboard-code',
        copyToClipboardDescription: 'the-copy-to-clipboard-description',
        dataTestSubj: 'the-data-test-subj',
@ -109,42 +159,4 @@ describe('Transform: <DataGrid /> with useIndexData()', () => {
      ).not.toBeInTheDocument();
    });
  });
-
-  test('Cross-cluster search warning', async () => {
-    // Arrange
-    const dataView = {
-      getIndexPattern: () => 'remote:the-index-pattern-title',
-      fields: [] as any[],
-    } as SearchItems['dataView'];
-
-    const Wrapper = () => {
-      const props = {
-        ...useIndexData(dataView, { match_all: {} }, runtimeMappings),
-        copyToClipboard: 'the-copy-to-clipboard-code',
-        copyToClipboardDescription: 'the-copy-to-clipboard-description',
-        dataTestSubj: 'the-data-test-subj',
-        title: 'the-index-preview-title',
-        toastNotifications: {} as CoreSetup['notifications']['toasts'],
-      };
-
-      return <DataGrid {...props} />;
-    };
-
-    render(
-      <QueryClientProvider client={queryClient}>
-        <IntlProvider locale="en">
-          <Wrapper />
-        </IntlProvider>
-      </QueryClientProvider>
-    );
-
-    // Act
-    // Assert
-    await waitFor(() => {
-      expect(screen.queryByText('the-index-preview-title')).toBeInTheDocument();
-      expect(
-        screen.queryByText('Cross-cluster search returned no fields data.')
-      ).toBeInTheDocument();
-    });
-  });
 });
--- a/x-pack/plugins/transform/public/app/hooks/use_index_data.ts
+++ b/x-pack/plugins/transform/public/app/hooks/use_index_data.ts
@ -43,13 +43,16 @@ import type { SearchItems } from './use_search_items';
 import { useGetHistogramsForFields } from './use_get_histograms_for_fields';
 import { useDataSearch } from './use_data_search';

-export const useIndexData = (
-  dataView: SearchItems['dataView'],
-  query: TransformConfigQuery,
-  combinedRuntimeMappings?: StepDefineExposedState['runtimeMappings'],
-  timeRangeMs?: TimeRangeMs,
-  populatedFields?: string[]
-): UseIndexDataReturnType => {
+interface UseIndexDataOptions {
+  dataView: SearchItems['dataView'];
+  query: TransformConfigQuery;
+  populatedFields: string[];
+  combinedRuntimeMappings?: StepDefineExposedState['runtimeMappings'];
+  timeRangeMs?: TimeRangeMs;
+}
+
+export const useIndexData = (options: UseIndexDataOptions): UseIndexDataReturnType => {
+  const { dataView, query, populatedFields, combinedRuntimeMappings, timeRangeMs } = options;
  const { analytics } = useAppDependencies();

  // Store the performance metric's start time using a ref
@ -59,11 +62,10 @@ export const useIndexData = (
  const indexPattern = useMemo(() => dataView.getIndexPattern(), [dataView]);
  const toastNotifications = useToastNotifications();

-  const baseFilterCriteria = buildBaseFilterCriteria(
-    dataView.timeFieldName,
-    timeRangeMs?.from,
-    timeRangeMs?.to,
-    query
+  const baseFilterCriteria = useMemo(
+    () =>
+      buildBaseFilterCriteria(dataView.timeFieldName, timeRangeMs?.from, timeRangeMs?.to, query),
+    [dataView.timeFieldName, query, timeRangeMs]
  );

  const defaultQuery = useMemo(
@ -71,85 +73,23 @@ export const useIndexData = (
    [baseFilterCriteria, dataView, timeRangeMs]
  );

-  const queryWithBaseFilterCriteria = {
-    bool: {
-      filter: baseFilterCriteria,
-    },
-  };
-
-  // 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 {
-    error: dataViewFieldsError,
-    data: dataViewFieldsData,
-    isError: dataViewFieldsIsError,
-    isLoading: dataViewFieldsIsLoading,
-  } = useDataSearch(
-    {
-      index: indexPattern,
-      body: {
-        fields: ['*'],
-        _source: false,
-        query: {
-          function_score: {
-            query: defaultQuery,
-            random_score: {},
-          },
-        },
-        size: 500,
+  const queryWithBaseFilterCriteria = useMemo(
+    () => ({
+      bool: {
+        filter: baseFilterCriteria,
      },
-    },
-    // Check whether fetching should be enabled
-    // If populatedFields are not provided, make own request to calculate
-    !Array.isArray(populatedFields) &&
-      !(dataView.timeFieldName !== undefined && timeRangeMs === undefined)
+    }),
+    [baseFilterCriteria]
  );

-  useEffect(() => {
-    if (dataViewFieldsIsLoading && !dataViewFieldsIsError) {
-      setErrorMessage('');
-      setStatus(INDEX_STATUS.LOADING);
-    } else if (dataViewFieldsError !== null) {
-      setErrorMessage(getErrorMessage(dataViewFieldsError));
-      setStatus(INDEX_STATUS.ERROR);
-    } else if (
-      !dataViewFieldsIsLoading &&
-      !dataViewFieldsIsError &&
-      dataViewFieldsData !== undefined
-    ) {
-      const isCrossClusterSearch = indexPattern.includes(':');
-      const isMissingFields = dataViewFieldsData.hits.hits.every(
-        (d) => typeof d.fields === 'undefined'
-      );
-
-      setCcsWarning(isCrossClusterSearch && isMissingFields);
-      setStatus(INDEX_STATUS.LOADED);
-    }
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [dataViewFieldsData, dataViewFieldsError, dataViewFieldsIsError, dataViewFieldsIsLoading]);
-
  const dataViewFields = useMemo(() => {
-    let allPopulatedFields = Array.isArray(populatedFields) ? populatedFields : [];
-
-    if (populatedFields === undefined && dataViewFieldsData) {
-      // Get all field names for each returned doc and flatten it
-      // to a list of unique field names used across all docs.
-      const docs = dataViewFieldsData.hits.hits.map((d) => getProcessedFields(d.fields ?? {}));
-      allPopulatedFields = [...new Set(docs.map(Object.keys).flat(1))];
-    }
-
+    const allPopulatedFields = Array.isArray(populatedFields) ? populatedFields : [];
    const allDataViewFields = getFieldsFromKibanaDataView(dataView);
    return allPopulatedFields.filter((d) => allDataViewFields.includes(d)).sort();
    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [dataViewFieldsData, populatedFields]);
+  }, [populatedFields]);

  const columns: EuiDataGridColumn[] = useMemo(() => {
-    if (typeof dataViewFields === 'undefined') {
-      return [];
-    }
-
    let result: Array<{ id: string; schema: string | undefined }> = [];

    // Get the the runtime fields that are defined from API field and index patterns
@ -206,7 +146,9 @@ export const useIndexData = (
    error: dataGridDataError,
    data: dataGridData,
    isError: dataGridDataIsError,
-    isLoading: dataGridDataIsLoading,
+    // React Query v4 has the weird behavior that `isLoading` is `true` on load
+    // when a query is disabled, that's why we need to use `isFetching` here.
+    isFetching: dataGridDataIsLoading,
  } = useDataSearch(
    {
      index: indexPattern,
@ -223,7 +165,7 @@ export const useIndexData = (
      },
    },
    // Check whether fetching should be enabled
-    dataViewFields !== undefined
+    dataViewFields.length > 0
  );

  useEffect(() => {
@ -309,7 +251,7 @@ export const useIndexData = (

  if (
    dataGrid.status === INDEX_STATUS.LOADED &&
-    dataViewFields !== undefined &&
+    dataViewFields.length > 0 &&
    Array.isArray(histogramsForFieldsData) &&
    histogramsForFieldsData.length > 0 &&
    loadIndexDataStartTime.current !== undefined
@ -325,8 +267,11 @@ export const useIndexData = (
    });
  }

-  return {
-    ...dataGrid,
-    renderCellValue,
-  };
+  return useMemo(
+    () => ({
+      ...dataGrid,
+      renderCellValue,
+    }),
+    [dataGrid, renderCellValue]
+  );
 };
--- a/x-pack/plugins/transform/public/app/sections/create_transform/components/step_define/step_define_form.tsx
+++ b/x-pack/plugins/transform/public/app/sections/create_transform/components/step_define/step_define_form.tsx
@ -126,16 +126,23 @@ export const StepDefineForm: FC<StepDefineFormProps> = React.memo((props) => {
  const { runtimeMappings } = stepDefineForm.runtimeMappingsEditor.state;

  const fieldStatsContext = useFieldStatsFlyoutContext();
-  const indexPreviewProps = {
-    ...useIndexData(
-      dataView,
-      transformConfigQuery,
-      runtimeMappings,
-      timeRangeMs,
+
+  const populatedFields = useMemo(
+    () =>
      isPopulatedFields(fieldStatsContext?.populatedFields)
        ? [...fieldStatsContext.populatedFields]
-        : []
-    ),
+        : [],
+    [fieldStatsContext?.populatedFields]
+  );
+
+  const indexPreviewProps = {
+    ...useIndexData({
+      dataView,
+      query: transformConfigQuery,
+      populatedFields,
+      combinedRuntimeMappings: runtimeMappings,
+      timeRangeMs,
+    }),
    dataTestSubj: 'transformIndexPreview',
    toastNotifications,
  };