Improves the performance of the table ES|QL visualization (#187142)

## Summary Closes #187089 Fixes the performance of the ES|QL charts: - the panel - the table visualization The ES|QL charts become imperformant for a big amount of fields. I am changing the implementation a bit in order to render (and hold to cache) only the fields that take part in the visualization and not all colums ### Checklist - [ ] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios
elastic · Jul 2, 2024 · e73eb1d · e73eb1d
1 parent d183092
commit e73eb1d
Show file tree

Hide file tree

Showing 12 changed files with 174 additions and 65 deletions.
diff --git a/x-pack/plugins/lens/common/expressions/map_to_columns/map_to_columns.test.ts b/x-pack/plugins/lens/common/expressions/map_to_columns/map_to_columns.test.ts
@@ -279,4 +279,56 @@ describe('map_to_columns', () => {
       }
     `);
   });
+
+  describe('map_to_columns_text_based', () => {
+    it('should keep columns that exist in idMap only', async () => {
+      const input: Datatable = {
+        type: 'datatable',
+        columns: [
+          { id: 'a', name: 'A', meta: { type: 'number' } },
+          { id: 'b', name: 'B', meta: { type: 'number' } },
+          { id: 'c', name: 'C', meta: { type: 'string' } },
+        ],
+        rows: [
+          { a: 1, b: 2, c: '3' },
+          { a: 3, b: 4, c: '5' },
+          { a: 5, b: 6, c: '7' },
+          { a: 7, b: 8, c: '9' },
+        ],
+      };
+
+      const idMap = {
+        a: [
+          {
+            id: 'a',
+            label: 'A',
+          },
+        ],
+        b: [
+          {
+            id: 'b',
+            label: 'B',
+          },
+        ],
+      };
+
+      const result = await mapToColumns.fn(
+        input,
+        { idMap: JSON.stringify(idMap), isTextBased: true },
+        createMockExecutionContext()
+      );
+
+      expect(result.columns).toStrictEqual([
+        { id: 'a', name: 'A', meta: { type: 'number' } },
+        { id: 'b', name: 'B', meta: { type: 'number' } },
+      ]);
+
+      expect(result.rows).toStrictEqual([
+        { a: 1, b: 2 },
+        { a: 3, b: 4 },
+        { a: 5, b: 6 },
+        { a: 7, b: 8 },
+      ]);
+    });
+  });
 });
diff --git a/x-pack/plugins/lens/common/expressions/map_to_columns/map_to_columns.ts b/x-pack/plugins/lens/common/expressions/map_to_columns/map_to_columns.ts
@@ -22,11 +22,21 @@ export const mapToColumns: MapToColumnsExpressionFunction = {
           'A JSON encoded object in which keys are the datatable column ids and values are the Lens column definitions. Any datatable columns not mentioned within the ID map will be kept unmapped.',
       }),
     },
+    isTextBased: {
+      types: ['boolean'],
+      help: i18n.translate('xpack.lens.functions.mapToColumns.isESQL.help', {
+        defaultMessage: 'An optional flag to indicate if this is about the text based datasource.',
+      }),
+    },
   },
   inputTypes: ['datatable'],
   async fn(...args) {
     /** Build optimization: prevent adding extra code into initial bundle **/
     const { mapToOriginalColumns } = await import('./map_to_columns_fn');
-    return mapToOriginalColumns(...args);
+    const { mapToOriginalColumnsTextBased } = await import('./map_to_columns_fn_textbased');
+
+    return args?.[1]?.isTextBased
+      ? mapToOriginalColumnsTextBased(...args)
+      : mapToOriginalColumns(...args);
   },
 };
diff --git a/x-pack/plugins/lens/common/expressions/map_to_columns/map_to_columns_fn_textbased.ts b/x-pack/plugins/lens/common/expressions/map_to_columns/map_to_columns_fn_textbased.ts
@@ -0,0 +1,38 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { OriginalColumn, MapToColumnsExpressionFunction } from './types';
+
+export const mapToOriginalColumnsTextBased: MapToColumnsExpressionFunction['fn'] = (
+  data,
+  { idMap: encodedIdMap }
+) => {
+  const idMap = JSON.parse(encodedIdMap) as Record<string, OriginalColumn[]>;
+
+  return {
+    ...data,
+    rows: data.rows.map((row) => {
+      const mappedRow: Record<string, unknown> = {};
+
+      for (const id in row) {
+        if (id in idMap) {
+          for (const cachedEntry of idMap[id]) {
+            mappedRow[cachedEntry.id] = row[id]; // <= I wrote idMap rather than mappedRow
+          }
+        }
+      }
+
+      return mappedRow;
+    }),
+    columns: data.columns.flatMap((column) => {
+      if (!(column.id in idMap)) {
+        return [];
+      }
+      return idMap[column.id].map((originalColumn) => ({ ...column, id: originalColumn.id }));
+    }),
+  };
+};
diff --git a/x-pack/plugins/lens/common/expressions/map_to_columns/types.ts b/x-pack/plugins/lens/common/expressions/map_to_columns/types.ts
@@ -17,6 +17,7 @@ export type MapToColumnsExpressionFunction = ExpressionFunctionDefinition<
   Datatable,
   {
     idMap: string;
+    isTextBased?: boolean;
   },
   Datatable | Promise<Datatable>
 >;
diff --git a/x-pack/plugins/lens/public/app_plugin/shared/edit_on_the_fly/helpers.ts b/x-pack/plugins/lens/public/app_plugin/shared/edit_on_the_fly/helpers.ts
@@ -42,7 +42,7 @@ export const getSuggestions = async (
       signal: abortController?.signal,
     });
     const context = {
-      dataViewSpec: dataView?.toSpec(),
+      dataViewSpec: dataView?.toSpec(false),
       fieldName: '',
       textBasedColumns: columns,
       query,

diff --git a/x-pack/plugins/lens/public/app_plugin/shared/edit_on_the_fly/lens_configuration_flyout.tsx b/x-pack/plugins/lens/public/app_plugin/shared/edit_on_the_fly/lens_configuration_flyout.tsx
@@ -117,8 +117,7 @@ export function LensEditConfigurationFlyout({
         // there are cases where a query can return a big amount of columns
         // at this case we don't suggest all columns in a table but the first
         // MAX_NUM_OF_COLUMNS
-        const columns = Object.keys(table.rows?.[0]) ?? [];
-        setSuggestsLimitedColumns(columns.length >= MAX_NUM_OF_COLUMNS);
+        setSuggestsLimitedColumns(table.columns.length >= MAX_NUM_OF_COLUMNS);
         layers.forEach((layer) => {
           activeData[layer] = table;
         });

diff --git a/x-pack/plugins/lens/public/datasources/text_based/components/dimension_editor.tsx b/x-pack/plugins/lens/public/datasources/text_based/components/dimension_editor.tsx
@@ -5,15 +5,15 @@
  * 2.0.
  */
 
-import React from 'react';
+import React, { useEffect, useMemo, useState } from 'react';
 import { i18n } from '@kbn/i18n';
 import { EuiFormRow } from '@elastic/eui';
 import { euiThemeVars } from '@kbn/ui-theme';
-import type { ExpressionsStart } from '@kbn/expressions-plugin/public';
+import type { ExpressionsStart, DatatableColumn } from '@kbn/expressions-plugin/public';
+import { fetchFieldsFromESQL } from '@kbn/text-based-editor';
 import type { DatasourceDimensionEditorProps, DataType } from '../../../types';
 import { FieldSelect } from './field_select';
 import type { TextBasedPrivateState } from '../types';
-import { retrieveLayerColumnsFromCache, getColumnsFromCache } from '../fieldlist_cache';
 import { isNotNumeric, isNumeric } from '../utils';
 
 export type TextBasedDimensionEditorProps =
@@ -22,30 +22,55 @@ export type TextBasedDimensionEditorProps =
   };
 
 export function TextBasedDimensionEditor(props: TextBasedDimensionEditorProps) {
+  const [allColumns, setAllColumns] = useState<DatatableColumn[]>([]);
   const query = props.state.layers[props.layerId]?.query;
 
-  const allColumns = retrieveLayerColumnsFromCache(
-    props.state.layers[props.layerId]?.columns ?? [],
-    query
-  );
-  const allFields = query ? getColumnsFromCache(query) : [];
+  useEffect(() => {
+    // in case the columns are not in the cache, I refetch them
+    async function fetchColumns() {
+      if (query) {
+        const table = await fetchFieldsFromESQL(
+          { esql: `${query.esql} | limit 0` },
+          props.expressions
+        );
+        if (table) {
+          setAllColumns(table.columns);
+        }
+      }
+    }
+    fetchColumns();
+  }, [props.expressions, query]);
+
   const hasNumberTypeColumns = allColumns?.some(isNumeric);
-  const fields = allFields.map((col) => {
-    return {
-      id: col.id,
-      name: col.name,
-      meta: col?.meta ?? { type: 'number' },
-      compatible:
-        props.isMetricDimension && hasNumberTypeColumns
-          ? props.filterOperations({
-              dataType: col?.meta?.type as DataType,
-              isBucketed: Boolean(isNotNumeric(col)),
-              scale: 'ordinal',
-            })
-          : true,
-    };
-  });
-  const selectedField = allColumns?.find((column) => column.columnId === props.columnId);
+  const fields = useMemo(() => {
+    return allColumns.map((col) => {
+      return {
+        id: col.id,
+        name: col.name,
+        meta: col?.meta ?? { type: 'number' },
+        compatible:
+          props.isMetricDimension && hasNumberTypeColumns
+            ? props.filterOperations({
+                dataType: col?.meta?.type as DataType,
+                isBucketed: Boolean(isNotNumeric(col)),
+                scale: 'ordinal',
+              })
+            : true,
+      };
+    });
+  }, [allColumns, hasNumberTypeColumns, props]);
+
+  const selectedField = useMemo(() => {
+    const field = fields?.find((column) => column.id === props.columnId);
+    if (field) {
+      return {
+        fieldName: field.name,
+        meta: field.meta,
+        columnId: field.id,
+      };
+    }
+    return undefined;
+  }, [fields, props.columnId]);
 
   return (
     <>

diff --git a/x-pack/plugins/lens/public/datasources/text_based/components/dimension_trigger.tsx b/x-pack/plugins/lens/public/datasources/text_based/components/dimension_trigger.tsx
@@ -5,54 +5,25 @@
  * 2.0.
  */
 
-import React, { useEffect, useState } from 'react';
+import React from 'react';
 import { i18n } from '@kbn/i18n';
-import { fetchFieldsFromESQL } from '@kbn/text-based-editor';
 import { DimensionTrigger } from '@kbn/visualization-ui-components';
 import type { ExpressionsStart } from '@kbn/expressions-plugin/public';
 import type { DatasourceDimensionTriggerProps } from '../../../types';
 import type { TextBasedPrivateState } from '../types';
-import {
-  getColumnsFromCache,
-  addColumnsToCache,
-  retrieveLayerColumnsFromCache,
-} from '../fieldlist_cache';
 
 export type TextBasedDimensionTrigger = DatasourceDimensionTriggerProps<TextBasedPrivateState> & {
   columnLabelMap: Record<string, string>;
   expressions: ExpressionsStart;
 };
 
 export function TextBasedDimensionTrigger(props: TextBasedDimensionTrigger) {
-  const [dataHasLoaded, setDataHasLoaded] = useState(false);
-  const query = props.state.layers[props.layerId]?.query;
-  useEffect(() => {
-    // in case the columns are not in the cache, I refetch them
-    async function fetchColumns() {
-      const fieldList = query ? getColumnsFromCache(query) : [];
+  const customLabel: string | undefined = props.columnLabelMap[props.columnId];
 
-      if (fieldList.length === 0 && query) {
-        const table = await fetchFieldsFromESQL(query, props.expressions);
-        if (table) {
-          addColumnsToCache(query, table.columns);
-        }
-      }
-      setDataHasLoaded(true);
-    }
-    fetchColumns();
-  }, [props.expressions, query]);
-  const allColumns = dataHasLoaded
-    ? retrieveLayerColumnsFromCache(props.state.layers[props.layerId]?.columns ?? [], query)
-    : [];
-  const selectedField = allColumns?.find((column) => column.columnId === props.columnId);
-  let customLabel: string | undefined = props.columnLabelMap[props.columnId];
-  if (!customLabel) {
-    customLabel = selectedField?.fieldName;
-  }
   return (
     <DimensionTrigger
       id={props.columnId}
-      color={customLabel && selectedField ? 'primary' : 'danger'}
+      color={customLabel ? 'primary' : 'danger'}
       dataTestSubj="lns-dimensionTrigger-textBased"
       label={
         customLabel ??

diff --git a/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts b/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts
@@ -747,6 +747,9 @@ describe('Textbased Data Source', () => {
                 "idMap": Array [
                   "{\\"Test 1\\":[{\\"id\\":\\"a\\",\\"label\\":\\"Test 1\\"}],\\"Test 2\\":[{\\"id\\":\\"b\\",\\"label\\":\\"Test 2\\"}]}",
                 ],
+                "isTextBased": Array [
+                  true,
+                ],
               },
               "function": "lens_map_to_columns",
               "type": "function",

diff --git a/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.tsx b/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.tsx
@@ -11,7 +11,7 @@ import { CoreStart } from '@kbn/core/public';
 import { IStorageWrapper } from '@kbn/kibana-utils-plugin/public';
 import { AggregateQuery, isOfAggregateQueryType, getAggregateQueryMode } from '@kbn/es-query';
 import type { SavedObjectReference } from '@kbn/core/public';
-import type { ExpressionsStart } from '@kbn/expressions-plugin/public';
+import type { ExpressionsStart, DatatableColumn } from '@kbn/expressions-plugin/public';
 import type { DataViewsPublicPluginStart, DataView } from '@kbn/data-views-plugin/public';
 import type { DataPublicPluginStart } from '@kbn/data-plugin/public';
 import memoizeOne from 'memoize-one';
@@ -180,6 +180,15 @@ export function getTextBasedDatasource({
     return Object.entries(state.layers)?.flatMap(([id, layer]) => {
       const allColumns = retrieveLayerColumnsFromCache(layer.columns, layer.query);
 
+      if (!allColumns.length && layer.query) {
+        const layerColumns = layer.columns.map((c) => ({
+          id: c.columnId,
+          name: c.fieldName,
+          meta: c.meta,
+        })) as DatatableColumn[];
+        addColumnsToCache(layer.query, layerColumns);
+      }
+
       const unchangedSuggestionTable = getUnchangedSuggestionTable(state, allColumns, id);
 
       // we are trying here to cover the most common cases for the charts we offer
@@ -214,7 +223,7 @@ export function getTextBasedDatasource({
     if (fieldName) return [];
     if (context && 'dataViewSpec' in context && context.dataViewSpec.title && context.query) {
       const newLayerId = generateId();
-      const textBasedQueryColumns = context.textBasedColumns ?? [];
+      const textBasedQueryColumns = context.textBasedColumns?.slice(0, MAX_NUM_OF_COLUMNS) ?? [];
       // Number fields are assigned automatically as metrics (!isBucketed). There are cases where the query
       // will not return number fields. In these cases we want to suggest a datatable
       // Datatable works differently in this case. On the metrics dimension can be all type of fields
@@ -258,7 +267,7 @@ export function getTextBasedDatasource({
           [newLayerId]: {
             index,
             query,
-            columns: newColumns.slice(0, MAX_NUM_OF_COLUMNS) ?? [],
+            columns: newColumns ?? [],
             timeField: context.dataViewSpec.timeFieldName,
           },
         },
@@ -275,7 +284,7 @@ export function getTextBasedDatasource({
             notAssignedMetrics: !hasNumberTypeColumns,
             layerId: newLayerId,
             columns:
-              newColumns?.slice(0, MAX_NUM_OF_COLUMNS)?.map((f) => {
+              newColumns?.map((f) => {
                 return {
                   columnId: f.columnId,
                   operation: {

diff --git a/x-pack/plugins/lens/public/datasources/text_based/to_expression.ts b/x-pack/plugins/lens/public/datasources/text_based/to_expression.ts
@@ -59,6 +59,7 @@ function getExpressionForLayer(
       function: 'lens_map_to_columns',
       arguments: {
         idMap: [JSON.stringify(idMapper)],
+        isTextBased: [true],
       },
     });
     return textBasedQueryToAst;

diff --git a/x-pack/plugins/lens/public/trigger_actions/open_lens_config/create_action_helpers.ts b/x-pack/plugins/lens/public/trigger_actions/open_lens_config/create_action_helpers.ts
@@ -78,7 +78,7 @@ export async function executeCreateAction({
   });
 
   const context = {
-    dataViewSpec: dataView.toSpec(),
+    dataViewSpec: dataView.toSpec(false),
     fieldName: '',
     textBasedColumns: columns,
     query: defaultEsqlQuery,