[Lens] [ES|QL] Add helper to map the value to the corresponding column

x-pack/platform/plugins/shared/lens/public/app_plugin/shared/edit_on_the_fly/helpers.test.ts

@@ -17,7 +17,7 @@ import {
   mockAllSuggestions,
 } from '../../../mocks';
 import { suggestionsApi } from '../../../lens_suggestions_api';
-import { getSuggestions, getGridAttrs } from './helpers';
+import { buildDisplayRowsFromEsqlValues, getGridAttrs, getSuggestions } from './helpers';
 
 const mockSuggestionApi = suggestionsApi as jest.Mock;
 const mockFetchData = getESQLResults as jest.Mock;
@@ -67,6 +67,44 @@ jest.mock('@kbn/esql-utils', () => {
 });
 
 describe('Lens inline editing helpers', () => {
+  describe('buildDisplayRowsFromEsqlValues', () => {
+    it('returns values unchanged when value and display columns match in order', () => {
+      const valueColumns = [
+        { name: 'a', type: 'double' },
+        { name: 'b', type: 'integer' },
+      ];
+      const values = [
+        [1, 2],
+        [3, 4],
+      ];
+
+      expect(
+        buildDisplayRowsFromEsqlValues({
+          displayColumns: valueColumns,
+          valueColumns,
+          values,
+        })
+      ).toEqual(values);
+    });
+
+    it('maps row cells by column name when all_columns is a superset of columns', () => {
+      const displayColumns = [
+        { name: 'count', type: 'double' },
+        { name: 'max_value', type: 'integer' },
+      ];
+      const valueColumns = [{ name: 'max_value', type: 'integer' }];
+      const values = [[500]];
+
+      expect(
+        buildDisplayRowsFromEsqlValues({
+          displayColumns,
+          valueColumns,
+          values,
+        })
+      ).toEqual([[null, 500]]);
+    });
+  });
+
   describe('getSuggestions', () => {
     const query = {
       esql: 'from index1 | limit 10 | stats average = avg(bytes)',
@@ -247,6 +285,39 @@ describe('Lens inline editing helpers', () => {
       expect(gridAttributes.columns).toStrictEqual(emptyColumns);
     });
 
+    it('passes all_columns to formatESQLColumns and expands values by name when columns is a subset', async () => {
+      const allColumnsRaw = [
+        { name: 'count', type: 'double' },
+        { name: 'max_value', type: 'integer' },
+      ];
+      const subsetColumns = [{ name: 'max_value', type: 'integer' }];
+      const formattedColumns = [
+        { name: 'count', id: 'count', meta: { type: 'number', esType: 'double' } },
+        { name: 'max_value', id: 'max_value', meta: { type: 'number', esType: 'integer' } },
+      ];
+
+      mockFetchData.mockImplementation(() => ({
+        response: {
+          all_columns: allColumnsRaw,
+          columns: subsetColumns,
+          values: [[500]],
+        },
+      }));
+      mockformatESQLColumns.mockReturnValueOnce(formattedColumns);
+
+      const gridAttributes = await getGridAttrs(
+        query,
+        dataviewSpecArr,
+        startDependencies.data,
+        startDependencies.http,
+        uiSettingsMock
+      );
+
+      expect(mockformatESQLColumns).toHaveBeenCalledWith(allColumnsRaw);
+      expect(gridAttributes.rows).toEqual([[null, 500]]);
+      expect(gridAttributes.columns).toStrictEqual(formattedColumns);
+    });
+
     it('falls back to getESQLAdHocDataview when spec has no timeFieldName', async () => {
       dataViews.create.mockClear();
       mockFetchData.mockImplementation(() => ({

x-pack/platform/plugins/shared/lens/public/app_plugin/shared/edit_on_the_fly/helpers.ts

@@ -15,7 +15,7 @@ import { type AggregateQuery, buildEsQuery } from '@kbn/es-query';
 import type { CoreStart, IUiSettingsClient } from '@kbn/core/public';
 import { getEsQueryConfig, UI_SETTINGS } from '@kbn/data-plugin/public';
 import type { ESQLControlVariable } from '@kbn/esql-types';
-import type { ESQLRow } from '@kbn/es-types';
+import type { ESQLColumn, ESQLRow } from '@kbn/es-types';
 import { getLensAttributesFromSuggestion, mapVisToChartType } from '@kbn/visualization-utils';
 import type { DataViewSpec } from '@kbn/data-views-plugin/public';
 import type { DataView } from '@kbn/data-views-plugin/common';
@@ -33,6 +33,31 @@ export interface ESQLDataGridAttrs {
   columns: DatatableColumn[];
 }
 
+const columnsMatchInOrder = (a: ESQLColumn[], b: ESQLColumn[]) => {
+  return a.length === b.length && a.every((col, i) => col.name === b[i]?.name);
+};
+
+export const buildDisplayRowsFromEsqlValues = ({
+  displayColumns,
+  valueColumns,
+  values,
+}: {
+  displayColumns: ESQLColumn[];
+  valueColumns: ESQLColumn[];
+  values: ESQLRow[];
+}): ESQLRow[] => {
+  if (columnsMatchInOrder(valueColumns, displayColumns)) {
+    return values;
+  }
+
+  // Pre-compute which value column index each display column maps to (-1 if missing)
+  const valueIndexPerGridColumn = displayColumns.map((col) =>
+    valueColumns.findIndex((v) => v.name === col.name)
+  );
+  // For each row, pick values by index; fill null for columns with no data
+  return values.map((row) => valueIndexPerGridColumn.map((i) => (i >= 0 ? row[i] : null)));
+};
+
 const getDSLFilter = (
   queryService: DataPublicPluginStart['query'],
   uiSettings: IUiSettingsClient,
@@ -94,17 +119,16 @@ export const getGridAttrs = async (
     timezone,
   });
 
-  let queryColumns = results.response.columns;
-  // Use all_columns property if it exists in the payload
+  const { all_columns: allColumns = [], columns: valueColumns = [], values } = results.response;
+  // Use `all_columns` property if it exists in the payload,
   // which has all columns regardless if they have data or not
-  if (results.response.all_columns) {
-    queryColumns = results.response.all_columns;
-  }
+  const displayColumns = allColumns.length > 0 ? allColumns : valueColumns;
 
-  const columns = formatESQLColumns(queryColumns);
+  const rows = buildDisplayRowsFromEsqlValues({ displayColumns, valueColumns, values });
+  const columns = formatESQLColumns(displayColumns);
 
   return {
-    rows: results.response.values,
+    rows,
     dataView,
     columns,
   };