From ac2fd0633afc3ef8538cc969acad5f11ccc33347 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 27 Aug 2025 10:00:11 -0600
Subject: [PATCH 01/54] suggest user-defined columns in multi-expression EVAL
 statements

---
 .../src/autocomplete/helper.test.ts           | 30 +++++++++++++++++--
 .../src/autocomplete/helper.ts                | 30 +++++++++++++++++--
 2 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts
index 465f14816d4d9..bdaa985966c16 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts
@@ -7,13 +7,23 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { parse } from '@kbn/esql-ast';
+import { EDITOR_MARKER, parse } from '@kbn/esql-ast';
 import { getQueryForFields } from './helper';
 
 describe('getQueryForFields', () => {
   const assert = (query: string, expected: string) => {
     const { root } = parse(query);
 
+    // Simulate removing the editor marker from the AST
+    // As things are set up today, the marker will be there in the query
+    // string but will be removed from the AST...
+    const lastCommand = root.commands[root.commands.length - 1];
+    if (lastCommand.name === 'eval') {
+      lastCommand.args = lastCommand.args.filter(
+        (arg) => Array.isArray(arg) || arg.name !== EDITOR_MARKER
+      );
+    }
+
     const result = getQueryForFields(query, root);
 
     expect(result).toEqual(expected);
@@ -27,15 +37,29 @@ describe('getQueryForFields', () => {
   it('should convert FORK branches into vanilla queries', () => {
     const query = `FROM index
     | EVAL foo = 1
-    | FORK (STATS field1 | EVAL esql_editor_marker)`;
+    | FORK (STATS field1 | EVAL ${EDITOR_MARKER})`;
     assert(query, 'FROM index | EVAL foo = 1 | STATS field1');
 
     const query2 = `FROM index 
     | EVAL foo = 1
-    | FORK (STATS field1) (LIMIT 10) (WHERE field1 == 1 | EVAL esql_editor_marker)`;
+    | FORK (STATS field1) (LIMIT 10) (WHERE field1 == 1 | EVAL ${EDITOR_MARKER})`;
     assert(query2, 'FROM index | EVAL foo = 1 | WHERE field1 == 1');
   });
 
+  it('should convert multiple EVAL expressions into separate EVAL commands', () => {
+    const query = `FROM index
+    | EVAL foo = 1, bar = foo + 1, baz = ${EDITOR_MARKER}`;
+    assert(query, 'FROM index | EVAL foo = 1 | EVAL bar = foo + 1');
+
+    const query2 = `FROM index
+    | EVAL foo = 1, bar = foo + 1, baz = bar + 1, ${EDITOR_MARKER}`;
+    assert(query2, 'FROM index | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1');
+
+    const query3 = `FROM index
+    | EVAL foo = 1, ${EDITOR_MARKER}`;
+    assert(query3, 'FROM index | EVAL foo = 1');
+  });
+
   it('should return empty string if non-FROM source command', () => {
     assert('ROW field1 = 1', '');
     assert('SHOW INFO', '');
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
index bf3acc1212ca9..b4ce1f6fb1abd 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
@@ -8,7 +8,7 @@
  */
 
 import type { ESQLAstQueryExpression } from '@kbn/esql-ast';
-import { type ESQLCommand, BasicPrettyPrinter } from '@kbn/esql-ast';
+import { type ESQLCommand, BasicPrettyPrinter, Builder, EDITOR_MARKER } from '@kbn/esql-ast';
 
 /**
  * This function is used to build the query that will be used to compute the
@@ -32,13 +32,39 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
      * previous context is equivalent to a query without the FORK command.:
      *
      * Original query: FROM lolz | EVAL foo = 1 | FORK (EVAL bar = 2) (EVAL baz = 3 | WHERE /)
-     * Simplified: FROM lolz | EVAL foo = 1 | EVAL baz = 3 | WHERE /
+     * Simplified:     FROM lolz | EVAL foo = 1 | EVAL baz = 3
      */
     const currentBranch = lastCommand.args[lastCommand.args.length - 1] as ESQLAstQueryExpression;
     const newCommands = commands.slice(0, -1).concat(currentBranch.commands.slice(0, -1));
     return BasicPrettyPrinter.print({ ...root, commands: newCommands });
   }
 
+  if (lastCommand && lastCommand.name === 'eval') {
+    const endsWithComma = queryString.replace(EDITOR_MARKER, '').trim().endsWith(',');
+    if (lastCommand.args.length > 1 || endsWithComma) {
+      /**
+       * If we get here, we know that we have a multi-expression EVAL statement.
+       *
+       * e.g. EVAL foo = 1, bar = foo + 1, baz = bar + 1
+       *
+       * In order for this to work with the caching system which expects field availability to be
+       * delineated by pipes, we need to split the current EVAL command into an equivalent
+       * set of single-expression EVAL commands.
+       *
+       * Original query: FROM lolz | EVAL foo = 1, bar = foo + 1, baz = bar + 1, /
+       * Simplified:     FROM lolz | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1
+       */
+      const individualEVALCommands: ESQLCommand[] = [];
+      for (const expression of lastCommand.args) {
+        individualEVALCommands.push(Builder.command({ name: 'eval', args: [expression] }));
+      }
+      const newCommands = commands
+        .slice(0, -1)
+        .concat(endsWithComma ? individualEVALCommands : individualEVALCommands.slice(0, -1));
+      return BasicPrettyPrinter.print({ ...root, commands: newCommands });
+    }
+  }
+
   // If there is only one source command and it does not require fields, do not
   // fetch fields, hence return an empty string.
   return commands.length === 1 && ['row', 'show'].includes(commands[0].name)

From e532b90e51ca9049cf85f40bed21da16838cf98b Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 27 Aug 2025 10:02:08 -0600
Subject: [PATCH 02/54] casing

---
 .../src/autocomplete/helper.ts                              | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
index b4ce1f6fb1abd..36954fa514f92 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
@@ -54,13 +54,13 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
        * Original query: FROM lolz | EVAL foo = 1, bar = foo + 1, baz = bar + 1, /
        * Simplified:     FROM lolz | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1
        */
-      const individualEVALCommands: ESQLCommand[] = [];
+      const individualEvalCommands: ESQLCommand[] = [];
       for (const expression of lastCommand.args) {
-        individualEVALCommands.push(Builder.command({ name: 'eval', args: [expression] }));
+        individualEvalCommands.push(Builder.command({ name: 'eval', args: [expression] }));
       }
       const newCommands = commands
         .slice(0, -1)
-        .concat(endsWithComma ? individualEVALCommands : individualEVALCommands.slice(0, -1));
+        .concat(endsWithComma ? individualEvalCommands : individualEvalCommands.slice(0, -1));
       return BasicPrettyPrinter.print({ ...root, commands: newCommands });
     }
   }

From 23493d7780fb160a8f24f187d4e86ec77adee25f Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 27 Aug 2025 16:18:55 -0600
Subject: [PATCH 03/54] compute available columns at the level of each command

---
 .../commands/completion/validate.ts           |  1 +
 .../src/shared/helpers.ts                     |  4 +
 .../src/validation/helpers.ts                 | 55 +-----------
 .../src/validation/resources.ts               | 45 +---------
 .../src/validation/validation.ts              | 89 ++++++++++---------
 5 files changed, 54 insertions(+), 140 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
index 2e607329df95f..6e3c9b6c561b8 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
@@ -57,6 +57,7 @@ export const validate = (
   const targetName = targetField?.name || 'completion';
 
   // Sets the target field so the column is recognized after the command is applied
+  // @TODO can we remove now?
   context?.userDefinedColumns.set(targetName, [
     {
       name: targetName,
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index 1eb2d013ad01f..1c0566e1f8616 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -135,6 +135,10 @@ export async function getCurrentQueryAvailableFields(
   const lastCommand = commands[commands.length - 1];
   const commandDefinition = esqlCommandRegistry.getCommandByName(lastCommand.name);
 
+  // @TODO — all logic in collectUserDefinedColumns
+  // should be will be moved to columnsAfter methods;
+  // though it may still be useful to delineate between
+  // user-defined columns and other fields... need to consider this
   // If the command has a columnsAfter function, use it to get the fields
   if (commandDefinition?.methods.columnsAfter) {
     const userDefinedColumns = collectUserDefinedColumns([lastCommand], cacheCopy, query);
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
index 6e542710fd8fc..693c802f594cf 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
@@ -7,62 +7,9 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import {
-  type ESQLAst,
-  type ESQLAstQueryExpression,
-  type ESQLColumn,
-  type ESQLSource,
-  type ESQLCommand,
-  type FunctionDefinition,
-  Walker,
-} from '@kbn/esql-ast';
-import { mutate, synth } from '@kbn/esql-ast';
+import { type ESQLAst, type ESQLCommand, type FunctionDefinition, Walker } from '@kbn/esql-ast';
 import type { ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types';
 
-export function buildQueryForFieldsFromSource(queryString: string, ast: ESQLAst) {
-  const firstCommand = ast[0];
-  if (!firstCommand) return '';
-
-  const sources: ESQLSource[] = [];
-  const metadataFields: ESQLColumn[] = [];
-
-  if (firstCommand.name === 'ts') {
-    const timeseries = firstCommand as ESQLCommand<'ts'>;
-    const tsSources = timeseries.args as ESQLSource[];
-
-    sources.push(...tsSources);
-  } else if (firstCommand.name === 'from') {
-    const fromSources = mutate.commands.from.sources.list(firstCommand as any);
-    const fromMetadataColumns = [...mutate.commands.from.metadata.list(firstCommand as any)].map(
-      ([column]) => column
-    );
-
-    sources.push(...fromSources);
-    if (fromMetadataColumns.length) metadataFields.push(...fromMetadataColumns);
-  }
-
-  const joinSummary = mutate.commands.join.summarize({
-    type: 'query',
-    commands: ast,
-  } as ESQLAstQueryExpression);
-  const joinIndices = joinSummary.map(({ target: { index } }) => index);
-
-  if (joinIndices.length > 0) {
-    sources.push(...joinIndices);
-  }
-
-  if (sources.length === 0) {
-    return queryString.substring(0, firstCommand.location.max + 1);
-  }
-
-  const from =
-    metadataFields.length > 0
-      ? synth.cmd`FROM ${sources} METADATA ${metadataFields}`
-      : synth.cmd`FROM ${sources}`;
-
-  return from.toString();
-}
-
 export function buildQueryForFieldsInPolicies(policies: ESQLPolicy[]) {
   return `from ${policies
     .flatMap(({ sourceIndices }) => sourceIndices)
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
index eec01712771f8..e663109be1e59 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
@@ -16,38 +16,7 @@ import {
   getSourcesHelper,
 } from '../shared/resources_helpers';
 import type { ESQLCallbacks } from '../shared/types';
-import {
-  buildQueryForFieldsForStringSources,
-  buildQueryForFieldsFromSource,
-  buildQueryForFieldsInPolicies,
-  getEnrichCommands,
-} from './helpers';
-
-export async function retrieveFields(
-  queryString: string,
-  commands: ESQLCommand[],
-  callbacks?: ESQLCallbacks
-): Promise<Map<string, ESQLFieldWithMetadata>> {
-  if (!callbacks || commands.length < 1) {
-    return new Map();
-  }
-  // Do not fetch fields, if query has only one source command and that command
-  // does not require fields.
-  if (commands.length === 1) {
-    switch (commands[0].name) {
-      case 'from':
-      case 'show':
-      case 'row': {
-        return new Map();
-      }
-    }
-  }
-  if (commands[0].name === 'row') {
-    return new Map();
-  }
-  const customQuery = buildQueryForFieldsFromSource(queryString, commands);
-  return await getFieldsByTypeHelper(customQuery, callbacks).getFieldsMap();
-}
+import { buildQueryForFieldsInPolicies, getEnrichCommands } from './helpers';
 
 export async function retrievePolicies(
   commands: ESQLCommand[],
@@ -100,15 +69,3 @@ export async function retrievePoliciesFields(
   );
   return await getFieldsByTypeHelper(customQuery, callbacks).getFieldsMap();
 }
-
-export async function retrieveFieldsFromStringSources(
-  queryString: string,
-  commands: ESQLCommand[],
-  callbacks?: ESQLCallbacks
-): Promise<Map<string, ESQLFieldWithMetadata>> {
-  if (!callbacks) {
-    return new Map();
-  }
-  const customQuery = buildQueryForFieldsForStringSources(queryString, commands);
-  return await getFieldsByTypeHelper(customQuery, callbacks).getFieldsMap();
-}
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index 9208016fa2770..5b4ad5ab43a8d 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -8,7 +8,7 @@
  */
 
 import type { ESQLAst, ESQLCommand, ESQLMessage, ErrorTypes } from '@kbn/esql-ast';
-import { EsqlQuery, walk, esqlCommandRegistry } from '@kbn/esql-ast';
+import { EsqlQuery, walk, esqlCommandRegistry, Builder } from '@kbn/esql-ast';
 import { getMessageFromId } from '@kbn/esql-ast/src/definitions/utils';
 import type {
   ESQLFieldWithMetadata,
@@ -19,13 +19,13 @@ import type { LicenseType } from '@kbn/licensing-types';
 import type { ESQLCallbacks } from '../shared/types';
 import { collectUserDefinedColumns } from '../shared/user_defined_columns';
 import {
-  retrieveFields,
-  retrieveFieldsFromStringSources,
   retrievePolicies,
-  retrievePoliciesFields,
+  // retrievePoliciesFields,
   retrieveSources,
 } from './resources';
 import type { ReferenceMaps, ValidationOptions, ValidationResult } from './types';
+import { getQueryForFields } from '../autocomplete/helper';
+import { getFieldsByTypeHelper } from '../shared/resources_helpers';
 
 /**
  * ES|QL validation public API
@@ -110,57 +110,60 @@ async function validateAst(
   const parsingResult = EsqlQuery.fromSrc(queryString);
   const rootCommands = parsingResult.ast.commands;
 
-  const [sources, availableFields, availablePolicies, joinIndices] = await Promise.all([
+  const [sources, availablePolicies, joinIndices] = await Promise.all([
     // retrieve the list of available sources
     retrieveSources(rootCommands, callbacks),
-    // retrieve available fields (if a source command has been defined)
-    retrieveFields(queryString, rootCommands, callbacks),
     // retrieve available policies (if an enrich command has been defined)
     retrievePolicies(rootCommands, callbacks),
     // retrieve indices for join command
     callbacks?.getJoinIndices?.(),
   ]);
 
-  if (availablePolicies.size) {
-    const fieldsFromPoliciesMap = await retrievePoliciesFields(
-      rootCommands,
-      availablePolicies,
-      callbacks
-    );
-    fieldsFromPoliciesMap.forEach((value, key) => availableFields.set(key, value));
-  }
-
-  if (rootCommands.some(({ name }) => ['grok', 'dissect'].includes(name))) {
-    const fieldsFromGrokOrDissect = await retrieveFieldsFromStringSources(
-      queryString,
-      rootCommands,
-      callbacks
-    );
-    fieldsFromGrokOrDissect.forEach((value, key) => {
-      // if the field is already present, do not overwrite it
-      // Note: this can also overlap with some userDefinedColumns
-      if (!availableFields.has(key)) {
-        availableFields.set(key, value);
-      }
-    });
-  }
+  // if (availablePolicies.size) {
+  //   const fieldsFromPoliciesMap = await retrievePoliciesFields(
+  //     rootCommands,
+  //     availablePolicies,
+  //     callbacks
+  //   );
+  //   fieldsFromPoliciesMap.forEach((value, key) => availableFields.set(key, value));
+  // }
 
-  const userDefinedColumns = collectUserDefinedColumns(rootCommands, availableFields, queryString);
-  messages.push(...validateUnsupportedTypeFields(availableFields, rootCommands));
+  const sourceFields = await getFieldsByTypeHelper(
+    queryString.split('|')[0],
+    callbacks
+  ).getFieldsMap();
 
-  const references: ReferenceMaps = {
-    sources,
-    fields: availableFields,
-    policies: availablePolicies,
-    userDefinedColumns,
-    query: queryString,
-    joinIndices: joinIndices?.indices || [],
-  };
+  messages.push(...validateUnsupportedTypeFields(sourceFields, rootCommands));
 
   const license = await callbacks?.getLicense?.();
   const hasMinimumLicenseRequired = license?.hasAtLeast;
-  for (const [_, command] of rootCommands.entries()) {
-    const commandMessages = validateCommand(command, references, rootCommands, {
+  for (let i = 0; i < rootCommands.length; i++) {
+    const partialQuery = queryString.slice(0, rootCommands[i].location.max + 1);
+
+    const previousCommands = rootCommands.slice(0, i + 1);
+    const queryForFields = getQueryForFields(
+      partialQuery,
+      Builder.expression.query(previousCommands)
+    );
+
+    const { getFieldsMap } = getFieldsByTypeHelper(queryForFields, callbacks);
+    const availableFields = await getFieldsMap();
+    const userDefinedColumns = collectUserDefinedColumns(
+      previousCommands,
+      availableFields,
+      queryString
+    );
+
+    const references: ReferenceMaps = {
+      sources,
+      fields: availableFields,
+      policies: availablePolicies,
+      userDefinedColumns,
+      query: queryString,
+      joinIndices: joinIndices?.indices || [],
+    };
+
+    const commandMessages = validateCommand(rootCommands[i], references, rootCommands, {
       ...callbacks,
       hasMinimumLicenseRequired,
     });
@@ -170,6 +173,8 @@ async function validateAst(
   const parserErrors = parsingResult.errors;
 
   /**
+   * @TODO - move deeper
+   *
    * Some changes to the grammar deleted the literal names for some tokens.
    * This is a workaround to restore the literals that were lost.
    *

From 4359764c82cb061b3949715c2ea66ca0242884ef Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 28 Aug 2025 15:21:10 -0600
Subject: [PATCH 04/54] start to unify column types

---
 .../commands/change_point/columns_after.ts    | 55 +++++++++++++------
 .../commands/completion/columns_after.ts      | 28 +++++++---
 .../commands/completion/validate.ts           |  1 +
 .../commands/dissect/columns_after.ts         | 13 ++++-
 .../commands/drop/columns_after.ts            |  4 +-
 .../commands/fork/columns_after.ts            |  3 +-
 .../commands/grok/columns_after.ts            |  6 +-
 .../commands/inlinestats/columns_after.ts     |  4 +-
 .../commands/keep/columns_after.test.ts       |  2 +-
 .../commands/keep/columns_after.ts            |  4 +-
 .../commands/rename/columns_after.ts          |  9 ++-
 .../commands/stats/columns_after.ts           |  4 +-
 .../src/commands_registry/registry.ts         |  6 +-
 .../src/commands_registry/types.ts            |  9 ++-
 .../src/shared/helpers.ts                     |  5 +-
 .../src/shared/resources_helpers.ts           | 10 ++--
 16 files changed, 105 insertions(+), 58 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
index 00680ddee9de7..7b938b60ebf87 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
@@ -7,30 +7,49 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import uniqBy from 'lodash/uniqBy';
-import { type ESQLCommand, type ESQLAstChangePointCommand } from '../../../types';
 import { LeafPrinter } from '../../../pretty_print/leaf_printer';
-import type { ESQLFieldWithMetadata } from '../../types';
-import type { ICommandContext } from '../../types';
+import { type ESQLAstChangePointCommand, type ESQLCommand } from '../../../types';
+import type { ESQLColumnData, ICommandContext } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const { target } = command as ESQLAstChangePointCommand;
 
-  return uniqBy(
-    [
-      ...previousColumns,
-      {
-        name: target ? LeafPrinter.column(target.type) : 'type',
-        type: 'keyword' as const,
-      },
-      {
-        name: target ? LeafPrinter.column(target.pvalue) : 'pvalue',
-        type: 'double' as const,
-      },
-    ],
-    'name'
-  );
+  let typeField: ESQLColumnData;
+  let pvalueField: ESQLColumnData;
+
+  if (target?.type) {
+    typeField = {
+      name: LeafPrinter.column(target.type),
+      type: 'keyword' as const,
+      userDefined: true,
+      location: target.type.location,
+    };
+  } else {
+    typeField = {
+      name: 'type',
+      type: 'keyword' as const,
+      userDefined: false,
+    };
+  }
+
+  if (target?.pvalue) {
+    pvalueField = {
+      name: LeafPrinter.column(target.pvalue),
+      type: 'double' as const,
+      userDefined: true,
+      location: target.pvalue.location,
+    };
+  } else {
+    pvalueField = {
+      name: 'pvalue',
+      type: 'double' as const,
+      userDefined: false,
+    };
+  }
+
+  return uniqBy([...previousColumns, typeField, pvalueField], 'name');
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
index 82e3c43d951be..96c98baa09b0a 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
@@ -7,14 +7,18 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import uniqBy from 'lodash/uniqBy';
-import { type ESQLCommand, type ESQLAstCompletionCommand } from '../../../types';
 import { LeafPrinter } from '../../../pretty_print/leaf_printer';
-import type { ESQLFieldWithMetadata } from '../../types';
-import type { ICommandContext } from '../../types';
+import { type ESQLAstCompletionCommand, type ESQLCommand } from '../../../types';
+import type {
+  ESQLColumnData,
+  ESQLFieldWithMetadata,
+  ESQLUserDefinedColumn,
+  ICommandContext,
+} from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const { targetField } = command as ESQLAstCompletionCommand;
@@ -22,10 +26,18 @@ export const columnsAfter = (
   return uniqBy(
     [
       ...previousColumns,
-      {
-        name: targetField ? LeafPrinter.column(targetField) : 'completion',
-        type: 'keyword' as const,
-      },
+      targetField
+        ? ({
+            name: LeafPrinter.column(targetField),
+            type: 'keyword' as const,
+            userDefined: true,
+            location: targetField.location,
+          } as ESQLUserDefinedColumn)
+        : ({
+            name: 'completion',
+            type: 'keyword' as const,
+            userDefined: false,
+          } as ESQLFieldWithMetadata),
     ],
     'name'
   );
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
index 6e3c9b6c561b8..8be981525df68 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
@@ -63,6 +63,7 @@ export const validate = (
       name: targetName,
       location: targetField?.location || command.location,
       type: 'keyword',
+      userDefined: true,
     },
   ]);
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
index c8a929322fcf3..e7e4fd180ca3a 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
@@ -8,7 +8,7 @@
  */
 import { walk } from '../../../walker';
 import { type ESQLCommand } from '../../../types';
-import type { ESQLFieldWithMetadata } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import type { ICommandContext } from '../../types';
 
 function unquoteTemplate(inputString: string): string {
@@ -35,7 +35,7 @@ export function extractDissectColumnNames(pattern: string): string[] {
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const columns: string[] = [];
@@ -49,6 +49,13 @@ export const columnsAfter = (
 
   return [
     ...previousColumns,
-    ...columns.map((column) => ({ name: column, type: 'keyword' as const })),
+    ...columns.map((column) => {
+      const newColumn: ESQLColumnData = {
+        name: column,
+        type: 'keyword' as const,
+        userDefined: false,
+      };
+      return newColumn;
+    }),
   ];
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
index 8033d2e754b97..11ff6dee5db15 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
@@ -8,12 +8,12 @@
  */
 import { walk } from '../../../walker';
 import { type ESQLCommand } from '../../../types';
-import type { ESQLFieldWithMetadata } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import type { ICommandContext } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const columnsToDrop: string[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index a44beb912b818..ced310189f4e5 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -22,7 +22,8 @@ export const columnsAfter = (
       {
         name: '_fork',
         type: 'keyword' as const,
-      },
+        userDefined: false,
+      } as ESQLFieldWithMetadata,
     ],
     'name'
   );
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
index a2803afe17154..336cc08865754 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
@@ -8,7 +8,7 @@
  */
 import { walk } from '../../../walker';
 import { type ESQLCommand } from '../../../types';
-import type { ESQLFieldWithMetadata } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import type { ICommandContext } from '../../types';
 
 function unquoteTemplate(inputString: string): string {
@@ -50,7 +50,7 @@ export function extractSemanticsFromGrok(pattern: string): string[] {
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const columns: string[] = [];
@@ -64,6 +64,6 @@ export const columnsAfter = (
 
   return [
     ...previousColumns,
-    ...columns.map((column) => ({ name: column, type: 'keyword' as const })),
+    ...columns.map((column) => ({ name: column, type: 'keyword' as const, userDefined: false })),
   ];
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
index 0d3fd9d787926..bab2ca61cdc44 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
@@ -8,7 +8,7 @@
  */
 import uniqBy from 'lodash/uniqBy';
 import type { ESQLCommand } from '../../../types';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData, ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
 import type { ICommandContext } from '../../types';
 import type { FieldType } from '../../../definitions/types';
 
@@ -33,7 +33,7 @@ function transformMapToESQLFields(
 
 export const columnsAfter = (
   _command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const userDefinedColumns =
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
index 9c602ba20e3b2..b68f480056a0d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
@@ -10,7 +10,7 @@ import { synth } from '../../../..';
 import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
 import { columnsAfter } from './columns_after';
 
-describe('FORK', () => {
+describe('KEEP', () => {
   const context = {
     userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
     fields: new Map<string, ESQLFieldWithMetadata>([
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
index 6d0ab16aba080..6e6abc12c046a 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
@@ -8,12 +8,12 @@
  */
 import { walk } from '../../../walker';
 import { type ESQLCommand } from '../../../types';
-import type { ESQLFieldWithMetadata } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import type { ICommandContext } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const columnsToKeep: string[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
index 5c15019ff13df..07a4545835fa5 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
@@ -7,14 +7,13 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import uniqBy from 'lodash/uniqBy';
-import type { ESQLCommand, ESQLFunction, ESQLAstBaseItem } from '../../../types';
 import { isFunctionExpression } from '../../../ast/is';
-import type { ESQLFieldWithMetadata } from '../../types';
-import type { ICommandContext } from '../../types';
+import type { ESQLAstBaseItem, ESQLCommand, ESQLFunction } from '../../../types';
+import type { ESQLColumnData, ICommandContext } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const asRenamePairs: ESQLFunction[] = [];
@@ -51,5 +50,5 @@ export const columnsAfter = (
     return oldColumn; // No rename found, keep the old name
   });
 
-  return uniqBy(newFields, 'name');
+  return uniqBy(newFields, 'name') as ESQLColumnData[];
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
index 1855367c5ec69..a50b337cc45d2 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
@@ -9,7 +9,7 @@
 import uniqBy from 'lodash/uniqBy';
 import type { ESQLCommand } from '../../../types';
 import { walk } from '../../../walker';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData, ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
 import type { ICommandContext } from '../../types';
 import type { FieldType } from '../../../definitions/types';
 import { isColumn } from '../../../ast/is';
@@ -36,7 +36,7 @@ function transformMapToESQLFields(
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   const columns: string[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
index 6bb5037901abd..9fbf19f0c604c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
@@ -8,7 +8,7 @@
  */
 import type { LicenseType } from '@kbn/licensing-types';
 import type { ESQLMessage, ESQLCommand, ESQLAst } from '../types';
-import type { ISuggestionItem, ESQLFieldWithMetadata, ICommandCallbacks } from './types';
+import type { ISuggestionItem, ICommandCallbacks, ESQLColumnData } from './types';
 
 /**
  * Interface defining the methods that each ES|QL command should register.
@@ -59,9 +59,9 @@ export interface ICommandMethods<TContext = any> {
    */
   columnsAfter?: (
     command: ESQLCommand,
-    previousColumns: ESQLFieldWithMetadata[],
+    previousColumns: ESQLColumnData[],
     context?: TContext
-  ) => ESQLFieldWithMetadata[];
+  ) => ESQLColumnData[];
 }
 
 export interface ICommandMetadata {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
index bb7ecf56fefe3..e392e5d2d037d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
@@ -97,9 +97,11 @@ export type GetColumnsByTypeFn = (
   }
 ) => Promise<ISuggestionItem[]>;
 
+// TODO consider not exporting this
 export interface ESQLFieldWithMetadata {
   name: string;
   type: FieldType;
+  userDefined?: false; // TODO consider making required property
   isEcs?: boolean;
   hasConflict?: boolean;
   metadata?: {
@@ -107,15 +109,19 @@ export interface ESQLFieldWithMetadata {
   };
 }
 
+// TODO consider not exporting this
 export interface ESQLUserDefinedColumn {
   name: string;
   // invalid expressions produce columns of type "unknown"
   // also, there are some cases where we can't yet infer the type of
   // a valid expression as with `CASE` which can return union types
   type: SupportedDataType | 'unknown';
-  location: ESQLLocation;
+  userDefined: true;
+  location: ESQLLocation; // TODO should this be optional?
 }
 
+export type ESQLColumnData = ESQLUserDefinedColumn | ESQLFieldWithMetadata;
+
 export interface ESQLPolicy {
   name: string;
   sourceIndices: string[];
@@ -131,6 +137,7 @@ export interface ICommandCallbacks {
 }
 
 export interface ICommandContext {
+  // TODO collapse userDefinedColumns and fields into one
   userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>;
   fields: Map<string, ESQLFieldWithMetadata>;
   sources?: ESQLSourceResult[];
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index 1c0566e1f8616..2b6ed2d0a36bf 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -13,6 +13,7 @@ import {
   type FunctionDefinition,
 } from '@kbn/esql-ast';
 import type {
+  ESQLColumnData,
   ESQLFieldWithMetadata,
   ESQLUserDefinedColumn,
 } from '@kbn/esql-ast/src/commands_registry/types';
@@ -128,9 +129,9 @@ export async function getFieldsFromES(query: string, resourceRetriever?: ESQLCal
 export async function getCurrentQueryAvailableFields(
   query: string,
   commands: ESQLAstCommand[],
-  previousPipeFields: ESQLFieldWithMetadata[]
+  previousPipeFields: ESQLColumnData[]
 ) {
-  const cacheCopy = new Map<string, ESQLFieldWithMetadata>();
+  const cacheCopy = new Map<string, ESQLColumnData>();
   previousPipeFields.forEach((field) => cacheCopy.set(field.name, field));
   const lastCommand = commands[commands.length - 1];
   const commandDefinition = esqlCommandRegistry.getCommandByName(lastCommand.name);
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index b4e15d99eaf1c..285269ec233fb 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -8,14 +8,14 @@
  */
 
 import { parse } from '@kbn/esql-ast';
-import type { ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types';
+import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types';
 import type { ESQLCallbacks } from './types';
 import { getFieldsFromES, getCurrentQueryAvailableFields } from './helpers';
 import { removeLastPipe, processPipes, toSingleLine } from './query_string_utils';
 
 export const NOT_SUGGESTED_TYPES = ['unsupported'];
 
-const cache = new Map<string, ESQLFieldWithMetadata[]>();
+const cache = new Map<string, ESQLColumnData[]>();
 
 // Function to check if a key exists in the cache, ignoring case
 function checkCacheInsensitive(keyToCheck: string) {
@@ -88,7 +88,7 @@ export function getFieldsByTypeHelper(queryText: string, resourceRetriever?: ESQ
     getFieldsByType: async (
       expectedType: Readonly<string> | Readonly<string[]> = 'any',
       ignored: string[] = []
-    ): Promise<ESQLFieldWithMetadata[]> => {
+    ): Promise<ESQLColumnData[]> => {
       const types = Array.isArray(expectedType) ? expectedType : [expectedType];
       await getFields();
       const queryTextForCacheSearch = toSingleLine(queryText);
@@ -105,11 +105,11 @@ export function getFieldsByTypeHelper(queryText: string, resourceRetriever?: ESQ
         }) || []
       );
     },
-    getFieldsMap: async () => {
+    getFieldsMap: async (): Promise<Map<string, ESQLColumnData>> => {
       await getFields();
       const queryTextForCacheSearch = toSingleLine(queryText);
       const cachedFields = getValueInsensitive(queryTextForCacheSearch);
-      const cacheCopy = new Map<string, ESQLFieldWithMetadata>();
+      const cacheCopy = new Map<string, ESQLColumnData>();
       cachedFields?.forEach((field) => cacheCopy.set(field.name, field));
       return cacheCopy;
     },

From 27712b1e15ac333bc23a171adce1e77d1e4e65f0 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 28 Aug 2025 15:59:50 -0600
Subject: [PATCH 05/54] columns-after for EVAL

---
 .../commands/eval/columns_after.ts            | 50 +++++++++++++++++++
 .../commands_registry/commands/eval/index.ts  |  2 +
 .../src/shared/helpers.ts                     |  6 +--
 .../src/shared/resources_helpers.ts           |  8 +--
 .../src/shared/user_defined_columns.ts        |  2 +-
 5 files changed, 59 insertions(+), 9 deletions(-)
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
new file mode 100644
index 0000000000000..c7346f6dc69d4
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
@@ -0,0 +1,50 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { getExpressionType } from '../../../definitions/utils';
+import { isAssignment, isColumn } from '../../../ast/is';
+import type { ESQLAstItem, ESQLCommand } from '../../../types';
+import type { ESQLColumnData, ESQLUserDefinedColumn, ICommandContext } from '../../types';
+
+export const columnsAfter = (
+  command: ESQLCommand,
+  previousColumns: ESQLColumnData[],
+  context?: ICommandContext
+) => {
+  const columnMap = new Map<string, ESQLColumnData>();
+  previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient
+
+  const typeOf = (thing: ESQLAstItem) =>
+    getExpressionType(thing, columnMap, context?.userDefinedColumns);
+
+  const newColumns = [];
+
+  for (const expression of command.args) {
+    if (isAssignment(expression) && isColumn(expression.args[0])) {
+      const name = expression.args[0].parts.join('.');
+      const newColumn: ESQLUserDefinedColumn = {
+        name,
+        type: typeOf(expression.args[1]),
+        location: expression.args[0].location,
+        userDefined: true,
+      };
+      newColumns.push(newColumn);
+    } else if (!Array.isArray(expression)) {
+      const newColumn: ESQLUserDefinedColumn = {
+        name: expression.text,
+        type: typeOf(expression),
+        location: expression.location,
+        userDefined: true,
+      };
+      newColumns.push(newColumn);
+    }
+  }
+
+  return [...previousColumns, ...newColumns];
+};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/index.ts
index 7804a7069183a..18c649ab6e348 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/index.ts
@@ -10,11 +10,13 @@ import { i18n } from '@kbn/i18n';
 import type { ICommandMethods } from '../../registry';
 import { autocomplete } from './autocomplete';
 import { validate } from './validate';
+import { columnsAfter } from './columns_after';
 import type { ICommandContext } from '../../types';
 
 const evalCommandMethods: ICommandMethods<ICommandContext> = {
   autocomplete,
   validate,
+  columnsAfter,
 };
 
 export const evalCommand = {
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index 2b6ed2d0a36bf..72dddc1abfdca 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -126,7 +126,7 @@ export async function getFieldsFromES(query: string, resourceRetriever?: ESQLCal
  * @param previousPipeFields, the fields from the previous pipe
  * @returns a list of fields that are available for the current pipe
  */
-export async function getCurrentQueryAvailableFields(
+export async function getCurrentQueryAvailableColumns(
   query: string,
   commands: ESQLAstCommand[],
   previousPipeFields: ESQLColumnData[]
@@ -142,10 +142,8 @@ export async function getCurrentQueryAvailableFields(
   // user-defined columns and other fields... need to consider this
   // If the command has a columnsAfter function, use it to get the fields
   if (commandDefinition?.methods.columnsAfter) {
-    const userDefinedColumns = collectUserDefinedColumns([lastCommand], cacheCopy, query);
-
     return commandDefinition.methods.columnsAfter(lastCommand, previousPipeFields, {
-      userDefinedColumns,
+      userDefinedColumns: new Map(),
     });
   } else {
     // If the command doesn't have a columnsAfter function, use the default behavior
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index 285269ec233fb..3506a3157646f 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -10,7 +10,7 @@
 import { parse } from '@kbn/esql-ast';
 import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types';
 import type { ESQLCallbacks } from './types';
-import { getFieldsFromES, getCurrentQueryAvailableFields } from './helpers';
+import { getFieldsFromES, getCurrentQueryAvailableColumns } from './helpers';
 import { removeLastPipe, processPipes, toSingleLine } from './query_string_utils';
 
 export const NOT_SUGGESTED_TYPES = ['unsupported'];
@@ -42,7 +42,7 @@ function getValueInsensitive(keyToCheck: string) {
  * for the next time the same query is used.
  * @param queryText
  */
-async function cacheFieldsForQuery(queryText: string) {
+async function cacheColumnsForQuery(queryText: string) {
   const existsInCache = checkCacheInsensitive(queryText);
   if (existsInCache) {
     // this is already in the cache
@@ -53,7 +53,7 @@ async function cacheFieldsForQuery(queryText: string) {
   const fieldsAvailableAfterPreviousCommand = getValueInsensitive(queryTextWithoutLastPipe);
   if (fieldsAvailableAfterPreviousCommand && fieldsAvailableAfterPreviousCommand?.length) {
     const { root } = parse(queryText);
-    const availableFields = await getCurrentQueryAvailableFields(
+    const availableFields = await getCurrentQueryAvailableColumns(
       queryText,
       root.commands,
       fieldsAvailableAfterPreviousCommand
@@ -80,7 +80,7 @@ export function getFieldsByTypeHelper(queryText: string, resourceRetriever?: ESQ
 
     // build fields cache for every partial query
     for (const query of partialQueries) {
-      await cacheFieldsForQuery(query);
+      await cacheColumnsForQuery(query);
     }
   };
 
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
index 0ae6ba0f2a47e..e54fd3b68d464 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
@@ -157,7 +157,7 @@ export function collectUserDefinedColumns(
     })
     .on('visitCommand', (ctx) => {
       const ret = [];
-      if (['row', 'eval', 'stats', 'inlinestats', 'ts', 'rename'].includes(ctx.node.name)) {
+      if (['row', 'stats', 'inlinestats', 'ts', 'rename'].includes(ctx.node.name)) {
         ret.push(...ctx.visitArgs());
       }
       if (['stats', 'inlinestats', 'enrich'].includes(ctx.node.name)) {

From e10bac41fb4f5d45a733ff88c3f9e8d3fb335e27 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 28 Aug 2025 17:03:31 -0600
Subject: [PATCH 06/54] STATS and INLINESTATS columns-after

---
 .../commands/inlinestats/columns_after.ts     | 72 +++++++++-----
 .../commands/stats/columns_after.ts           | 94 ++++++++++---------
 .../src/shared/user_defined_columns.ts        |  6 +-
 3 files changed, 99 insertions(+), 73 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
index bab2ca61cdc44..c5060676a4b08 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
@@ -7,40 +7,62 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import uniqBy from 'lodash/uniqBy';
-import type { ESQLCommand } from '../../../types';
-import type { ESQLColumnData, ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
-import type { ICommandContext } from '../../types';
-import type { FieldType } from '../../../definitions/types';
-
-function transformMapToESQLFields(
-  inputMap: Map<string, ESQLUserDefinedColumn[]>
-): ESQLFieldWithMetadata[] {
-  const esqlFields: ESQLFieldWithMetadata[] = [];
-
-  for (const [, userDefinedColumns] of inputMap) {
-    for (const userDefinedColumn of userDefinedColumns) {
-      if (userDefinedColumn.type) {
-        esqlFields.push({
-          name: userDefinedColumn.name,
-          type: userDefinedColumn.type as FieldType,
-        });
-      }
+import type { SupportedDataType } from '../../../definitions/types';
+import { getExpressionType } from '../../../definitions/utils';
+import type { ESQLAstItem, ESQLCommand, ESQLCommandOption } from '../../../types';
+import type { ESQLColumnData, ESQLUserDefinedColumn, ICommandContext } from '../../types';
+import { isAssignment, isColumn, isOptionNode } from '../../../..';
+
+const getUserDefinedColumns = (
+  command: ESQLCommand | ESQLCommandOption,
+  typeOf: (thing: ESQLAstItem) => SupportedDataType | 'unknown'
+): ESQLUserDefinedColumn[] => {
+  const columns = [];
+
+  for (const expression of command.args) {
+    if (isAssignment(expression) && isColumn(expression.args[0])) {
+      const name = expression.args[0].parts.join('.');
+      const newColumn: ESQLUserDefinedColumn = {
+        name,
+        type: typeOf(expression.args[1]),
+        location: expression.args[0].location,
+        userDefined: true,
+      };
+      columns.push(newColumn);
+      continue;
+    }
+
+    if (isOptionNode(expression) && expression.name === 'by') {
+      columns.push(...getUserDefinedColumns(expression, typeOf));
+      continue;
+    }
+
+    if (!isOptionNode(expression) && !Array.isArray(expression)) {
+      const newColumn: ESQLUserDefinedColumn = {
+        name: expression.text,
+        type: typeOf(expression),
+        location: expression.location,
+        userDefined: true,
+      };
+      columns.push(newColumn);
+      continue;
     }
   }
 
-  return esqlFields;
-}
+  return columns;
+};
 
 export const columnsAfter = (
   _command: ESQLCommand,
   previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
-  const userDefinedColumns =
-    context?.userDefinedColumns ?? new Map<string, ESQLUserDefinedColumn[]>();
+  const columnMap = new Map<string, ESQLColumnData>();
+  previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient
 
-  const arrayOfUserDefinedColumns: ESQLFieldWithMetadata[] =
-    transformMapToESQLFields(userDefinedColumns);
+  const typeOf = (thing: ESQLAstItem) =>
+    getExpressionType(thing, columnMap, context?.userDefinedColumns);
 
-  return uniqBy([...previousColumns, ...arrayOfUserDefinedColumns], 'name');
+  // TODO - is this uniqby helpful? Does it do what we expect?
+  return uniqBy([...previousColumns, ...getUserDefinedColumns(_command, typeOf)], 'name');
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
index a50b337cc45d2..469cc4166478d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
@@ -7,56 +7,62 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import uniqBy from 'lodash/uniqBy';
-import type { ESQLCommand } from '../../../types';
-import { walk } from '../../../walker';
-import type { ESQLColumnData, ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
-import type { ICommandContext } from '../../types';
-import type { FieldType } from '../../../definitions/types';
-import { isColumn } from '../../../ast/is';
-
-function transformMapToESQLFields(
-  inputMap: Map<string, ESQLUserDefinedColumn[]>
-): ESQLFieldWithMetadata[] {
-  const esqlFields: ESQLFieldWithMetadata[] = [];
-
-  for (const [, userDefinedColumns] of inputMap) {
-    for (const userDefinedColumn of userDefinedColumns) {
-      // Only include userDefinedColumns that have a known type
-      if (userDefinedColumn.type) {
-        esqlFields.push({
-          name: userDefinedColumn.name,
-          type: userDefinedColumn.type as FieldType,
-        });
-      }
+import { isAssignment, isColumn, isOptionNode } from '../../../ast/is';
+import type { SupportedDataType } from '../../../definitions/types';
+import { getExpressionType } from '../../../definitions/utils';
+import type { ESQLAstItem, ESQLCommand, ESQLCommandOption } from '../../../types';
+import type { ESQLColumnData, ESQLUserDefinedColumn, ICommandContext } from '../../types';
+
+const getUserDefinedColumns = (
+  command: ESQLCommand | ESQLCommandOption,
+  typeOf: (thing: ESQLAstItem) => SupportedDataType | 'unknown'
+): ESQLUserDefinedColumn[] => {
+  const columns = [];
+
+  for (const expression of command.args) {
+    if (isAssignment(expression) && isColumn(expression.args[0])) {
+      const name = expression.args[0].parts.join('.');
+      const newColumn: ESQLUserDefinedColumn = {
+        name,
+        type: typeOf(expression.args[1]),
+        location: expression.args[0].location,
+        userDefined: true,
+      };
+      columns.push(newColumn);
+      continue;
+    }
+
+    if (isOptionNode(expression) && expression.name === 'by') {
+      columns.push(...getUserDefinedColumns(expression, typeOf));
+      continue;
+    }
+
+    if (!isOptionNode(expression) && !Array.isArray(expression)) {
+      const newColumn: ESQLUserDefinedColumn = {
+        name: expression.text,
+        type: typeOf(expression),
+        location: expression.location,
+        userDefined: true,
+      };
+      columns.push(newColumn);
+      continue;
     }
   }
 
-  return esqlFields;
-}
+  return columns;
+};
 
 export const columnsAfter = (
-  command: ESQLCommand,
+  _command: ESQLCommand,
   previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
-  const columns: string[] = [];
-  const userDefinedColumns =
-    context?.userDefinedColumns ?? new Map<string, ESQLUserDefinedColumn[]>();
-
-  walk(command, {
-    visitCommandOption: (node) => {
-      const args = node.args.filter(isColumn);
-      const breakdownColumns = args.map((arg) => arg.name);
-      columns.push(...breakdownColumns);
-    },
-  });
-
-  const columnsToKeep = previousColumns.filter((field) => {
-    return columns.some((column) => column === field.name);
-  });
-
-  const arrayOfUserDefinedColumns: ESQLFieldWithMetadata[] =
-    transformMapToESQLFields(userDefinedColumns);
-
-  return uniqBy([...columnsToKeep, ...arrayOfUserDefinedColumns], 'name');
+  const columnMap = new Map<string, ESQLColumnData>();
+  previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient
+
+  const typeOf = (thing: ESQLAstItem) =>
+    getExpressionType(thing, columnMap, context?.userDefinedColumns);
+
+  // TODO - is this uniqby helpful? Does it do what we expect?
+  return uniqBy([...previousColumns, ...getUserDefinedColumns(_command, typeOf)], 'name');
 };
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
index e54fd3b68d464..00317fb24f949 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
@@ -141,9 +141,7 @@ export function collectUserDefinedColumns(
       }
     })
     .on('visitCommandOption', (ctx) => {
-      if (ctx.node.name === 'by') {
-        return [...ctx.visitArguments()];
-      } else if (ctx.node.name === 'with') {
+      if (ctx.node.name === 'with') {
         for (const assignFn of ctx.node.args) {
           if (isFunctionExpression(assignFn)) {
             const [newArg, oldArg] = assignFn?.args || [];
@@ -157,7 +155,7 @@ export function collectUserDefinedColumns(
     })
     .on('visitCommand', (ctx) => {
       const ret = [];
-      if (['row', 'stats', 'inlinestats', 'ts', 'rename'].includes(ctx.node.name)) {
+      if (['row', 'ts', 'rename'].includes(ctx.node.name)) {
         ret.push(...ctx.visitArgs());
       }
       if (['stats', 'inlinestats', 'enrich'].includes(ctx.node.name)) {

From b9ddc0ac76b3b5c0b73c8f3e383fdd4b80af2f0f Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 10:00:11 -0600
Subject: [PATCH 07/54] reuse stats columns after in inlinestats

---
 .../inlinestats/columns_after.test.ts         | 142 ------------------
 .../commands/inlinestats/columns_after.ts     |  68 ---------
 .../commands/inlinestats/index.ts             |   2 +-
 3 files changed, 1 insertion(+), 211 deletions(-)
 delete mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
 delete mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
deleted file mode 100644
index cb15a668d4a37..0000000000000
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
+++ /dev/null
@@ -1,142 +0,0 @@
-/*
- * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
-import { columnsAfter } from './columns_after';
-
-// Test data fixtures
-const createPreviousFields = (fields: Array<[string, string]>): ESQLFieldWithMetadata[] =>
-  fields.map(([name, type]) => ({ name, type } as ESQLFieldWithMetadata));
-
-describe('INLINESTATS', () => {
-  const createUserDefinedColumn = (
-    name: string,
-    type: ESQLFieldWithMetadata['type'],
-    location = { min: 0, max: 10 }
-  ): ESQLUserDefinedColumn => ({ name, type, location });
-
-  const createContext = (userDefinedColumns: Array<[string, ESQLUserDefinedColumn[]]> = []) => ({
-    userDefinedColumns: new Map(userDefinedColumns),
-    fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword' }],
-      ['count', { name: 'count', type: 'double' }],
-    ]),
-  });
-
-  const expectColumnsAfter = (
-    command: string,
-    previousFields: ESQLFieldWithMetadata[],
-    userColumns: Array<[string, string]>,
-    expectedResult: Array<[string, string]>
-  ) => {
-    const context = createContext(
-      userColumns.map(([name, type]) => [name, [createUserDefinedColumn(name, type as 'keyword')]])
-    );
-    const result = columnsAfter(synth.cmd(command), previousFields, context);
-    const expected = createPreviousFields(expectedResult);
-    expect(result).toEqual(expected);
-  };
-  it('preserves all previous columns and adds the user defined column, when no grouping is given', () => {
-    const previousFields = createPreviousFields([
-      ['field1', 'keyword'],
-      ['field2', 'double'],
-    ]);
-
-    expectColumnsAfter(
-      'INLINESTATS var0=AVG(field2)',
-      previousFields,
-      [['var0', 'double']],
-      [
-        ['field1', 'keyword'],
-        ['field2', 'double'],
-        ['var0', 'double'],
-      ]
-    );
-  });
-
-  it('preserves all previous columns and adds the escaped column, when no grouping is given', () => {
-    const previousFields = createPreviousFields([
-      ['field1', 'keyword'],
-      ['field2', 'double'],
-    ]);
-
-    expectColumnsAfter(
-      'INLINESTATS AVG(field2)',
-      previousFields,
-      [['AVG(field2)', 'double']],
-      [
-        ['field1', 'keyword'],
-        ['field2', 'double'],
-        ['AVG(field2)', 'double'],
-      ]
-    );
-  });
-
-  it('preserves all previous columns and adds the escaped column, with grouping', () => {
-    const previousFields = createPreviousFields([
-      ['field1', 'keyword'],
-      ['field2', 'double'],
-    ]);
-
-    // Note: Unlike STATS, INLINESTATS doesn't care about BY clause for column preservation
-    expectColumnsAfter(
-      'INLINESTATS AVG(field2) BY field1',
-      previousFields,
-      [['AVG(field2)', 'double']],
-      [
-        ['field1', 'keyword'],
-        ['field2', 'double'],
-        ['AVG(field2)', 'double'],
-      ]
-    );
-  });
-
-  it('preserves all previous columns and adds user defined and grouping columns', () => {
-    const previousFields = createPreviousFields([
-      ['field1', 'keyword'],
-      ['field2', 'double'],
-      ['@timestamp', 'date'],
-    ]);
-
-    expectColumnsAfter(
-      'INLINESTATS AVG(field2) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)',
-      previousFields,
-      [
-        ['AVG(field2)', 'double'],
-        ['buckets', 'unknown'],
-      ],
-      [
-        ['field1', 'keyword'],
-        ['field2', 'double'],
-        ['@timestamp', 'date'],
-        ['AVG(field2)', 'double'],
-        ['buckets', 'unknown'],
-      ]
-    );
-  });
-
-  it('handles duplicate column names by keeping the original column type', () => {
-    const previousFields = createPreviousFields([
-      ['field1', 'keyword'],
-      ['field2', 'double'],
-      ['avg_field', 'integer'], // This will be preserved since it comes first
-    ]);
-
-    expectColumnsAfter(
-      'INLINESTATS avg_field=AVG(field2)',
-      previousFields,
-      [['avg_field', 'double']], // This will be ignored due to duplicate name
-      [
-        ['field1', 'keyword'],
-        ['field2', 'double'],
-        ['avg_field', 'integer'],
-      ]
-    );
-  });
-});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
deleted file mode 100644
index c5060676a4b08..0000000000000
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-import uniqBy from 'lodash/uniqBy';
-import type { SupportedDataType } from '../../../definitions/types';
-import { getExpressionType } from '../../../definitions/utils';
-import type { ESQLAstItem, ESQLCommand, ESQLCommandOption } from '../../../types';
-import type { ESQLColumnData, ESQLUserDefinedColumn, ICommandContext } from '../../types';
-import { isAssignment, isColumn, isOptionNode } from '../../../..';
-
-const getUserDefinedColumns = (
-  command: ESQLCommand | ESQLCommandOption,
-  typeOf: (thing: ESQLAstItem) => SupportedDataType | 'unknown'
-): ESQLUserDefinedColumn[] => {
-  const columns = [];
-
-  for (const expression of command.args) {
-    if (isAssignment(expression) && isColumn(expression.args[0])) {
-      const name = expression.args[0].parts.join('.');
-      const newColumn: ESQLUserDefinedColumn = {
-        name,
-        type: typeOf(expression.args[1]),
-        location: expression.args[0].location,
-        userDefined: true,
-      };
-      columns.push(newColumn);
-      continue;
-    }
-
-    if (isOptionNode(expression) && expression.name === 'by') {
-      columns.push(...getUserDefinedColumns(expression, typeOf));
-      continue;
-    }
-
-    if (!isOptionNode(expression) && !Array.isArray(expression)) {
-      const newColumn: ESQLUserDefinedColumn = {
-        name: expression.text,
-        type: typeOf(expression),
-        location: expression.location,
-        userDefined: true,
-      };
-      columns.push(newColumn);
-      continue;
-    }
-  }
-
-  return columns;
-};
-
-export const columnsAfter = (
-  _command: ESQLCommand,
-  previousColumns: ESQLColumnData[],
-  context?: ICommandContext
-) => {
-  const columnMap = new Map<string, ESQLColumnData>();
-  previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient
-
-  const typeOf = (thing: ESQLAstItem) =>
-    getExpressionType(thing, columnMap, context?.userDefinedColumns);
-
-  // TODO - is this uniqby helpful? Does it do what we expect?
-  return uniqBy([...previousColumns, ...getUserDefinedColumns(_command, typeOf)], 'name');
-};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/index.ts
index 473779526c9aa..23ad88f522354 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/index.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { i18n } from '@kbn/i18n';
-import { columnsAfter } from './columns_after';
+import { columnsAfter } from '../stats/columns_after';
 import { autocomplete } from '../stats/autocomplete';
 import { validate } from '../stats/validate';
 import type { ICommandContext } from '../../types';

From 48657cb466dc13dc7e1a5efa0b99b42b853eee92 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 10:00:27 -0600
Subject: [PATCH 08/54] Clean up

---
 .../commands/grok/columns_after.ts            |  4 +-
 .../src/shared/user_defined_columns.ts        | 65 +------------------
 2 files changed, 5 insertions(+), 64 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
index 336cc08865754..d58af7985f99d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
@@ -64,6 +64,8 @@ export const columnsAfter = (
 
   return [
     ...previousColumns,
-    ...columns.map((column) => ({ name: column, type: 'keyword' as const, userDefined: false })),
+    ...columns.map(
+      (column) => ({ name: column, type: 'keyword' as const, userDefined: false } as ESQLColumnData)
+    ),
   ];
 };
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
index 00317fb24f949..c6414c158df32 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
@@ -7,10 +7,8 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import type { ESQLAst, ESQLAstItem, ESQLCommand, ESQLFunction } from '@kbn/esql-ast';
+import type { ESQLAst, ESQLAstItem, ESQLCommand } from '@kbn/esql-ast';
 import { isColumn, isFunctionExpression } from '@kbn/esql-ast';
-import { getExpressionType } from '@kbn/esql-ast/src/definitions/utils';
-import { EDITOR_MARKER } from '@kbn/esql-ast/src/definitions/constants';
 import { Visitor } from '@kbn/esql-ast/src/visitor';
 import type {
   ESQLUserDefinedColumn,
@@ -72,45 +70,6 @@ export function excludeUserDefinedColumnsFromCurrentCommand(
   return resultUserDefinedColumns;
 }
 
-function addUserDefinedColumnFromAssignment(
-  assignOperation: ESQLFunction,
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>,
-  fields: Map<string, ESQLFieldWithMetadata>
-) {
-  if (isColumn(assignOperation.args[0])) {
-    const rightHandSideArgType = getExpressionType(
-      assignOperation.args[1],
-      fields,
-      userDefinedColumns
-    );
-    addToUserDefinedColumnOccurrences(userDefinedColumns, {
-      name: assignOperation.args[0].parts.join('.'),
-      type: rightHandSideArgType /* fallback to number */,
-      location: assignOperation.args[0].location,
-    });
-  }
-}
-
-function addUserDefinedColumnFromExpression(
-  expressionOperation: ESQLFunction,
-  queryString: string,
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>,
-  fields: Map<string, ESQLFieldWithMetadata>
-) {
-  if (!expressionOperation.text.includes(EDITOR_MARKER)) {
-    const expressionText = queryString.substring(
-      expressionOperation.location.min,
-      expressionOperation.location.max + 1
-    );
-    const expressionType = getExpressionType(expressionOperation, fields, userDefinedColumns);
-    addToUserDefinedColumnOccurrences(userDefinedColumns, {
-      name: expressionText,
-      type: expressionType,
-      location: expressionOperation.location,
-    });
-  }
-}
-
 export function collectUserDefinedColumns(
   ast: ESQLAst,
   fields: Map<string, ESQLFieldWithMetadata>,
@@ -123,23 +82,6 @@ export function collectUserDefinedColumns(
       // TODO - add these as userDefinedColumns
     })
     .on('visitExpression', (_ctx) => {}) // required for the types :shrug:
-    .on('visitFunctionCallExpression', (ctx) => {
-      const node = ctx.node;
-
-      if (node.subtype === 'binary-expression' && node.name === 'where') {
-        ctx.visitArgument(0, undefined);
-        return;
-      }
-
-      if (node.name === 'as') {
-        const [oldArg, newArg] = ctx.node.args;
-        addToUserDefinedColumns(oldArg, newArg, fields, userDefinedColumns);
-      } else if (node.name === '=') {
-        addUserDefinedColumnFromAssignment(node, userDefinedColumns, fields);
-      } else {
-        addUserDefinedColumnFromExpression(node, queryString, userDefinedColumns, fields);
-      }
-    })
     .on('visitCommandOption', (ctx) => {
       if (ctx.node.name === 'with') {
         for (const assignFn of ctx.node.args) {
@@ -155,10 +97,7 @@ export function collectUserDefinedColumns(
     })
     .on('visitCommand', (ctx) => {
       const ret = [];
-      if (['row', 'ts', 'rename'].includes(ctx.node.name)) {
-        ret.push(...ctx.visitArgs());
-      }
-      if (['stats', 'inlinestats', 'enrich'].includes(ctx.node.name)) {
+      if (['enrich'].includes(ctx.node.name)) {
         // BY and WITH can contain userDefinedColumns
         ret.push(...ctx.visitOptions());
       }

From e735671e6b11fb135001b81ebab56b67f0c1142c Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 10:41:22 -0600
Subject: [PATCH 09/54] Unify column lists

---
 .../commands/change_point/validate.ts         | 21 ++++---
 .../commands/completion/autocomplete.ts       |  6 +-
 .../commands/completion/validate.ts           | 20 +++----
 .../commands/dissect/validate.ts              |  4 +-
 .../commands/eval/autocomplete.ts             |  5 +-
 .../commands/eval/columns_after.ts            |  3 +-
 .../commands/fork/validate.ts                 |  2 +-
 .../commands/fuse/validate.ts                 |  6 +-
 .../commands/grok/validate.ts                 |  4 +-
 .../commands_registry/commands/join/utils.ts  |  4 +-
 .../commands/sort/autocomplete.ts             |  7 +--
 .../commands/stats/autocomplete.ts            | 13 +----
 .../commands/stats/columns_after.ts           |  3 +-
 .../commands/where/autocomplete.ts            |  6 +-
 .../src/commands_registry/types.ts            |  6 +-
 .../definitions/utils/autocomplete/columns.ts |  8 +--
 .../definitions/utils/autocomplete/helpers.ts |  9 +--
 .../src/definitions/utils/expressions.ts      | 17 +++---
 .../src/definitions/utils/functions.ts        | 25 ++++----
 .../src/definitions/utils/index.ts            |  2 +-
 .../src/definitions/utils/shared.ts           |  8 +--
 .../kbn-esql-validation-autocomplete/index.ts |  2 +-
 .../src/autocomplete/autocomplete.ts          | 58 +++++++++----------
 .../src/shared/resources_helpers.ts           | 12 ++--
 .../src/validation/resources.ts               |  4 +-
 .../src/validation/types.ts                   |  7 +--
 .../src/validation/validation.ts              | 27 ++++-----
 .../src/languages/esql/lib/hover/hover.ts     | 13 ++---
 28 files changed, 124 insertions(+), 178 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
index 41f09df8e9e42..3dee57530d4aa 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
@@ -25,14 +25,10 @@ export const validate = (
   const valueArg = command.args[0];
   if (isColumn(valueArg)) {
     const columnName = valueArg.name;
-    // look up for columns in userDefinedColumns and existing fields
     let valueColumnType: string | undefined;
-    const userDefinedColumnRef = context?.userDefinedColumns.get(columnName);
-    if (userDefinedColumnRef) {
-      valueColumnType = userDefinedColumnRef.find((v) => v.name === columnName)?.type;
-    } else {
-      const fieldRef = context?.fields.get(columnName);
-      valueColumnType = fieldRef?.type;
+
+    if (context?.columns.has(columnName)) {
+      valueColumnType = context?.columns.get(columnName)?.type;
     }
 
     if (valueColumnType && !isNumericType(valueColumnType)) {
@@ -52,7 +48,7 @@ export const validate = (
   // validate ON column
   const defaultOnColumnName = '@timestamp';
   const onColumn = command.args.find((arg) => isOptionNode(arg) && arg.name === 'on');
-  const hasDefaultOnColumn = context?.fields.has(defaultOnColumnName);
+  const hasDefaultOnColumn = context?.columns.has(defaultOnColumnName);
   if (!onColumn && !hasDefaultOnColumn) {
     messages.push({
       location: command.location,
@@ -71,9 +67,12 @@ export const validate = (
     // populate userDefinedColumns references to prevent the common check from failing with unknown column
     asArg.args.forEach((arg, index) => {
       if (isColumn(arg)) {
-        context?.userDefinedColumns.set(arg.name, [
-          { name: arg.name, location: arg.location, type: index === 0 ? 'keyword' : 'long' },
-        ]);
+        // TODO - can we remove this?
+        context?.columns.set(arg.name, {
+          name: arg.name,
+          location: arg.location,
+          type: index === 0 ? 'keyword' : 'long',
+        });
       }
     });
   }
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts
index cf96358dd49c3..5f05b6bd566e6 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts
@@ -62,11 +62,7 @@ function getPosition(
   }
 
   const expressionRoot = prompt?.text !== EDITOR_MARKER ? prompt : undefined;
-  const expressionType = getExpressionType(
-    expressionRoot,
-    context?.fields,
-    context?.userDefinedColumns
-  );
+  const expressionType = getExpressionType(expressionRoot, context?.columns);
 
   if (isExpressionComplete(expressionType, query)) {
     return CompletionPosition.AFTER_PROMPT;
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
index 8be981525df68..0411530cc207c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
@@ -24,11 +24,7 @@ export const validate = (
 
   const { prompt, location, targetField, inferenceId } = command as ESQLAstCompletionCommand;
 
-  const promptExpressionType = getExpressionType(
-    prompt,
-    context?.fields,
-    context?.userDefinedColumns
-  );
+  const promptExpressionType = getExpressionType(prompt, context?.columns);
 
   if (!supportedPromptTypes.includes(promptExpressionType)) {
     messages.push({
@@ -58,14 +54,12 @@ export const validate = (
 
   // Sets the target field so the column is recognized after the command is applied
   // @TODO can we remove now?
-  context?.userDefinedColumns.set(targetName, [
-    {
-      name: targetName,
-      location: targetField?.location || command.location,
-      type: 'keyword',
-      userDefined: true,
-    },
-  ]);
+  context?.columns.set(targetName, {
+    name: targetName,
+    location: targetField?.location || command.location,
+    type: 'keyword',
+    userDefined: true,
+  });
 
   messages.push(...validateCommandArguments(command, ast, context, callbacks));
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/validate.ts
index d38a39beb44cf..ae689cbad4195 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/validate.ts
@@ -22,9 +22,9 @@ import { validateCommandArguments } from '../../../definitions/utils/validation'
 const validateColumnForGrokDissect = (command: ESQLCommand, context?: ICommandContext) => {
   const acceptedColumnTypes: FieldType[] = ['keyword', 'text'];
   const astCol = command.args[0] as ESQLColumn;
-  const columnRef = context?.fields.get(astCol.name);
+  const columnRef = context?.columns.get(astCol.name);
 
-  if (columnRef && !acceptedColumnTypes.includes(columnRef.type)) {
+  if (columnRef && !acceptedColumnTypes.includes(columnRef.type as FieldType)) {
     return [
       getMessageFromId({
         messageId: 'unsupportedColumnTypeForCommand',
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts
index 4bd50f4cac152..ceae006efafa4 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts
@@ -78,10 +78,7 @@ export async function autocomplete(
 
   if (
     // don't suggest finishing characters if incomplete expression
-    isExpressionComplete(
-      getExpressionType(expressionRoot, context?.fields, context?.userDefinedColumns),
-      innerText
-    ) &&
+    isExpressionComplete(getExpressionType(expressionRoot, context?.columns), innerText) &&
     // don't suggest finishing characters if the expression is a column
     // because "EVAL columnName" is a useless expression
     expressionRoot &&
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
index c7346f6dc69d4..1585c075c194c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
@@ -20,8 +20,7 @@ export const columnsAfter = (
   const columnMap = new Map<string, ESQLColumnData>();
   previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient
 
-  const typeOf = (thing: ESQLAstItem) =>
-    getExpressionType(thing, columnMap, context?.userDefinedColumns);
+  const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, columnMap);
 
   const newColumns = [];
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
index f070903d6aa37..36c103749b34e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
@@ -68,7 +68,7 @@ export const validate = (
     messages.push(errors.tooManyForks(forks[1]));
   }
 
-  context?.fields.set('_fork', {
+  context?.columns.set('_fork', {
     name: '_fork',
     type: 'keyword',
   });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.ts
index c580ccb6699d9..7a8a5464e1d85 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.ts
@@ -17,15 +17,15 @@ export const validate = (
 ): ESQLMessage[] => {
   const messages: ESQLMessage[] = [];
 
-  if (!context?.fields.get('_id')) {
+  if (!context?.columns.get('_id')) {
     messages.push(buildMissingMetadataMessage(command, '_id'));
   }
 
-  if (!context?.fields.get('_index')) {
+  if (!context?.columns.get('_index')) {
     messages.push(buildMissingMetadataMessage(command, '_index'));
   }
 
-  if (!context?.fields.get('_score')) {
+  if (!context?.columns.get('_score')) {
     messages.push(buildMissingMetadataMessage(command, '_score'));
   }
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/validate.ts
index 8b8100afd4545..d522b97c5ae11 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/validate.ts
@@ -21,9 +21,9 @@ export const validate = (
   const messages: ESQLMessage[] = [];
   const acceptedColumnTypes: FieldType[] = ['keyword', 'text'];
   const astCol = command.args[0] as ESQLColumn;
-  const columnRef = context?.fields.get(astCol.name);
+  const columnRef = context?.columns.get(astCol.name);
 
-  if (columnRef && !acceptedColumnTypes.includes(columnRef.type)) {
+  if (columnRef && !acceptedColumnTypes.includes(columnRef.type as FieldType)) {
     messages.push(
       getMessageFromId({
         messageId: 'unsupportedColumnTypeForCommand',
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts
index 7d458a9fe49b1..d8a1ebb396307 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts
@@ -12,7 +12,7 @@ import { unescapeColumnName } from '../../../definitions/utils/shared';
 import * as mutate from '../../../mutate';
 import { LeafPrinter } from '../../../pretty_print/leaf_printer';
 import { pipeCompleteItem, commaCompleteItem } from '../../complete_items';
-import { buildFieldsDefinitionsWithMetadata } from '../../../definitions/utils/functions';
+import { buildColumnSuggestions } from '../../../definitions/utils/functions';
 import type { ICommand } from '../../registry';
 import type { ESQLAstJoinCommand, ESQLCommand, ESQLCommandOption } from '../../../types';
 import type {
@@ -111,7 +111,7 @@ export const getFieldSuggestions = async (
   ]);
 
   const supportsControls = context?.supportsControls ?? false;
-  const joinFields = buildFieldsDefinitionsWithMetadata(
+  const joinFields = buildColumnSuggestions(
     lookupIndexFields.filter((f) => !ignoredFields.includes(f.name)),
     [],
     { supportsControls },
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts
index 05a79c9c5fc3e..5d3c6a260f782 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts
@@ -78,12 +78,7 @@ export async function autocomplete(
 
       const columnExists = (name: string) => _columnExists(name, context);
 
-      if (
-        isExpressionComplete(
-          getExpressionType(expressionRoot, context?.fields, context?.userDefinedColumns),
-          innerText
-        )
-      ) {
+      if (isExpressionComplete(getExpressionType(expressionRoot, context?.columns), innerText)) {
         suggestions.push(
           ...getSuggestionsAfterCompleteExpression(innerText, expressionRoot, columnExists)
         );
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts
index 697541fe6ef9b..a83ddcaadfeb2 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts
@@ -176,11 +176,7 @@ export async function autocomplete(
 
       // Is this a complete boolean expression?
       // If so, we can call it done and suggest a pipe
-      const expressionType = getExpressionType(
-        expressionRoot,
-        context?.fields,
-        context?.userDefinedColumns
-      );
+      const expressionType = getExpressionType(expressionRoot, context?.columns);
       if (expressionType === 'boolean' && isExpressionComplete(expressionType, innerText)) {
         suggestions.push(pipeCompleteItem, { ...commaCompleteItem, text: ', ' }, byCompleteItem);
       }
@@ -306,12 +302,7 @@ async function getExpressionSuggestions({
     suggestions.push(...emptySuggestions);
   }
 
-  if (
-    isExpressionComplete(
-      getExpressionType(expressionRoot, context?.fields, context?.userDefinedColumns),
-      innerText
-    )
-  ) {
+  if (isExpressionComplete(getExpressionType(expressionRoot, context?.columns), innerText)) {
     suggestions.push(...afterCompleteSuggestions);
   }
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
index 469cc4166478d..0fc1502f82084 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
@@ -60,8 +60,7 @@ export const columnsAfter = (
   const columnMap = new Map<string, ESQLColumnData>();
   previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient
 
-  const typeOf = (thing: ESQLAstItem) =>
-    getExpressionType(thing, columnMap, context?.userDefinedColumns);
+  const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, columnMap);
 
   // TODO - is this uniqby helpful? Does it do what we expect?
   return uniqBy([...previousColumns, ...getUserDefinedColumns(_command, typeOf)], 'name');
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.ts
index 0773410264679..f737337ea696e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.ts
@@ -49,11 +49,7 @@ export async function autocomplete(
 
   // Is this a complete boolean expression?
   // If so, we can call it done and suggest a pipe
-  const expressionType = getExpressionType(
-    expressionRoot,
-    context?.fields,
-    context?.userDefinedColumns
-  );
+  const expressionType = getExpressionType(expressionRoot, context?.columns);
   if (expressionType === 'boolean' && isExpressionComplete(expressionType, innerText)) {
     suggestions.push(pipeCompleteItem);
   }
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
index e392e5d2d037d..edefc48b218fe 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
@@ -132,14 +132,12 @@ export interface ESQLPolicy {
 export interface ICommandCallbacks {
   getByType?: GetColumnsByTypeFn;
   getSuggestedUserDefinedColumnName?: (extraFieldNames?: string[] | undefined) => string;
-  getColumnsForQuery?: (query: string) => Promise<ESQLFieldWithMetadata[]>;
+  getColumnsForQuery?: (query: string) => Promise<ESQLColumnData[]>;
   hasMinimumLicenseRequired?: (minimumLicenseRequired: LicenseType) => boolean;
 }
 
 export interface ICommandContext {
-  // TODO collapse userDefinedColumns and fields into one
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>;
-  fields: Map<string, ESQLFieldWithMetadata>;
+  columns: Map<string, ESQLColumnData>;
   sources?: ESQLSourceResult[];
   joinSources?: IndexAutocompleteItem[];
   timeSeriesSources?: IndexAutocompleteItem[];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
index 49175c3f8d2e3..0f921740518ae 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
@@ -63,11 +63,7 @@ function addUserDefinedColumnFromAssignment(
   fields: Map<string, ESQLFieldWithMetadata>
 ) {
   if (isColumn(assignOperation.args[0])) {
-    const rightHandSideArgType = getExpressionType(
-      assignOperation.args[1],
-      fields,
-      userDefinedColumns
-    );
+    const rightHandSideArgType = getExpressionType(assignOperation.args[1], fields);
     addToUserDefinedColumnOccurrences(userDefinedColumns, {
       name: assignOperation.args[0].parts.join('.'),
       type: rightHandSideArgType /* fallback to number */,
@@ -87,7 +83,7 @@ function addUserDefinedColumnFromExpression(
       expressionOperation.location.min,
       expressionOperation.location.max + 1
     );
-    const expressionType = getExpressionType(expressionOperation, fields, userDefinedColumns);
+    const expressionType = getExpressionType(expressionOperation, fields);
     addToUserDefinedColumnOccurrences(userDefinedColumns, {
       name: expressionText,
       type: expressionType,
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
index e4e71dc601b66..5537f876bb4bd 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
@@ -380,11 +380,7 @@ export async function suggestForExpression({
     case 'after_literal':
     case 'after_column':
     case 'after_function':
-      const expressionType = getExpressionType(
-        expressionRoot,
-        context?.fields,
-        context?.userDefinedColumns
-      );
+      const expressionType = getExpressionType(expressionRoot, context?.columns);
 
       if (!isParameterType(expressionType)) {
         break;
@@ -472,8 +468,7 @@ export async function suggestForExpression({
           location,
           rootOperator: rightmostOperator,
           preferredExpressionType,
-          getExpressionType: (expression) =>
-            getExpressionType(expression, context?.fields, context?.userDefinedColumns),
+          getExpressionType: (expression) => getExpressionType(expression, context?.columns),
           getColumnsByType,
           hasMinimumLicenseRequired,
           activeProduct,
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.ts
index fd39885e28e3f..aef58a357413e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.ts
@@ -14,7 +14,7 @@ import {
   isFunctionExpression,
 } from '../../ast/is';
 import type { ESQLAstItem, ESQLFunction } from '../../types';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../commands_registry/types';
+import type { ESQLColumnData } from '../../commands_registry/types';
 import type { SupportedDataType, FunctionDefinition } from '../types';
 import { lastItem } from '../../visitor/utils';
 import { getFunctionDefinition } from './functions';
@@ -127,8 +127,7 @@ export function isExpressionComplete(
  */
 export function getExpressionType(
   root: ESQLAstItem | undefined,
-  fields?: Map<string, ESQLFieldWithMetadata>,
-  userDefinedColumns?: Map<string, ESQLUserDefinedColumn[]>
+  columns?: Map<string, ESQLColumnData>
 ): SupportedDataType | 'unknown' {
   if (!root) {
     return 'unknown';
@@ -138,7 +137,7 @@ export function getExpressionType(
     if (root.length === 0) {
       return 'unknown';
     }
-    return getExpressionType(root[0], fields, userDefinedColumns);
+    return getExpressionType(root[0], columns);
   }
 
   if (isLiteral(root)) {
@@ -163,8 +162,8 @@ export function getExpressionType(
     }
   }
 
-  if (isColumn(root) && fields && userDefinedColumns) {
-    const column = getColumnForASTNode(root, { fields, userDefinedColumns });
+  if (isColumn(root) && columns) {
+    const column = getColumnForASTNode(root, { columns });
     const lastArg = lastItem(root.args);
     // If the last argument is a param, we return its type (param literal type)
     // This is useful for cases like `where ??field`
@@ -181,7 +180,7 @@ export function getExpressionType(
   }
 
   if (root.type === 'list') {
-    return getExpressionType(root.values[0], fields, userDefinedColumns);
+    return getExpressionType(root.values[0], columns);
   }
 
   if (isFunctionExpression(root)) {
@@ -215,7 +214,7 @@ export function getExpressionType(
        * will be null, which we aren't detecting. But this is ok because we consider
        * userDefinedColumns and fields to be nullable anyways and account for that during validation.
        */
-      return getExpressionType(root.args[root.args.length - 1], fields, userDefinedColumns);
+      return getExpressionType(root.args[root.args.length - 1], columns);
     }
 
     const signaturesWithCorrectArity = getSignaturesWithMatchingArity(fnDefinition, root);
@@ -223,7 +222,7 @@ export function getExpressionType(
     if (!signaturesWithCorrectArity.length) {
       return 'unknown';
     }
-    const argTypes = root.args.map((arg) => getExpressionType(arg, fields, userDefinedColumns));
+    const argTypes = root.args.map((arg) => getExpressionType(arg, columns));
 
     // When functions are passed null for any argument, they generally return null
     // This is a special case that is not reflected in our function definitions
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/functions.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/functions.ts
index e2054c7ff8351..986edff787959 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/functions.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/functions.ts
@@ -24,7 +24,7 @@ import { aggFunctionDefinitions } from '../generated/aggregation_functions';
 import { timeSeriesAggFunctionDefinitions } from '../generated/time_series_agg_functions';
 import { groupingFunctionDefinitions } from '../generated/grouping_functions';
 import { scalarFunctionDefinitions } from '../generated/scalar_functions';
-import type { ESQLFieldWithMetadata, ISuggestionItem } from '../../commands_registry/types';
+import type { ESQLColumnData, ISuggestionItem } from '../../commands_registry/types';
 import { TRIGGER_SUGGESTION_COMMAND } from '../../commands_registry/constants';
 import { buildFunctionDocumentation } from './documentation';
 import { getSafeInsertText, getControlSuggestion } from './autocomplete/helpers';
@@ -374,8 +374,8 @@ const getVariablePrefix = (variableType: ESQLVariableType) =>
     ? '??'
     : '?';
 
-export const buildFieldsDefinitionsWithMetadata = (
-  fields: ESQLFieldWithMetadata[],
+export const buildColumnSuggestions = (
+  columns: ESQLColumnData[],
   recommendedFieldsFromExtensions: RecommendedField[] = [],
   options?: {
     advanceCursor?: boolean;
@@ -386,20 +386,23 @@ export const buildFieldsDefinitionsWithMetadata = (
   },
   variables?: ESQLControlVariable[]
 ): ISuggestionItem[] => {
-  const fieldsSuggestions = fields.map((field) => {
-    const fieldType = field.type.charAt(0).toUpperCase() + field.type.slice(1);
-    const titleCaseType = `${field.name} (${fieldType})`;
+  const fieldsSuggestions = columns.map((column) => {
+    const fieldType = column.type.charAt(0).toUpperCase() + column.type.slice(1);
+    const titleCaseType = `${column.name} (${fieldType})`;
     // Check if the field is in the recommended fields from extensions list
     // and if so, mark it as recommended. This also ensures that recommended fields
     // that are registered wrongly, won't be shown as suggestions.
     const fieldIsRecommended = recommendedFieldsFromExtensions.some(
-      (recommendedField) => recommendedField.name === field.name
+      (recommendedField) => recommendedField.name === column.name
+    );
+    const sortText = getFieldsSortText(
+      !column.userDefined && Boolean(column.isEcs),
+      Boolean(fieldIsRecommended)
     );
-    const sortText = getFieldsSortText(Boolean(field.isEcs), Boolean(fieldIsRecommended));
     return {
-      label: field.name,
+      label: column.name,
       text:
-        getSafeInsertText(field.name) +
+        getSafeInsertText(column.name) +
         (options?.addComma ? ',' : '') +
         (options?.advanceCursor ? ' ' : ''),
       kind: 'Variable',
@@ -415,7 +418,7 @@ export const buildFieldsDefinitionsWithMetadata = (
     const userDefinedColumns =
       variables?.filter((variable) => variable.type === variableType) ?? [];
 
-    const controlSuggestions = fields.length
+    const controlSuggestions = columns.length
       ? getControlSuggestion(
           variableType,
           userDefinedColumns?.map((v) => `${getVariablePrefix(variableType)}${v.key}`)
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/index.ts
index d0e53570e115b..3eaa965567fcb 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/index.ts
@@ -16,7 +16,7 @@ export {
 } from './autocomplete/helpers';
 export { getSuggestionsToRightOfOperatorExpression } from './operators';
 export {
-  buildFieldsDefinitionsWithMetadata,
+  buildColumnSuggestions as buildFieldsDefinitionsWithMetadata,
   getFunctionSuggestions,
   getFunctionSignatures,
   getFunctionDefinition,
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/shared.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/shared.ts
index fff968191b220..e046d52e7bd5c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/shared.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/shared.ts
@@ -126,10 +126,10 @@ export function unescapeColumnName(columnName: string) {
  */
 export function getColumnByName(
   columnName: string,
-  { fields, userDefinedColumns }: ICommandContext
+  { columns }: ICommandContext
 ): ESQLFieldWithMetadata | ESQLUserDefinedColumn | undefined {
   const unescaped = unescapeColumnName(columnName);
-  return fields.get(unescaped) || userDefinedColumns.get(unescaped)?.[0];
+  return columns.get(unescaped);
 }
 
 /**
@@ -137,10 +137,10 @@ export function getColumnByName(
  */
 export function getColumnForASTNode(
   node: ESQLColumn | ESQLIdentifier,
-  { fields, userDefinedColumns }: ICommandContext
+  { columns }: ICommandContext
 ): ESQLFieldWithMetadata | ESQLUserDefinedColumn | undefined {
   const formatted = node.type === 'identifier' ? node.name : node.parts.join('.');
-  return getColumnByName(formatted, { fields, userDefinedColumns });
+  return getColumnByName(formatted, { columns });
 }
 
 function hasWildcard(name: string) {
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts
index b3507dca1e0ff..9459f53edd5fe 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts
@@ -24,7 +24,7 @@ export { suggest } from './src/autocomplete/autocomplete';
 export { collectUserDefinedColumns } from './src/shared/user_defined_columns';
 
 export {
-  getFieldsByTypeHelper,
+  getColumnsByTypeHelper as getFieldsByTypeHelper,
   getPolicyHelper,
   getSourcesHelper,
 } from './src/shared/resources_helpers';
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
index 4c06fdc67e9d5..101f3ac131f45 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
@@ -25,23 +25,21 @@ import {
 } from '@kbn/esql-ast/src/definitions/utils';
 import { getRecommendedQueriesSuggestionsFromStaticTemplates } from '@kbn/esql-ast/src/commands_registry/options/recommended_queries';
 import type {
-  ESQLUserDefinedColumn,
-  ESQLFieldWithMetadata,
+  ESQLColumnData,
   GetColumnsByTypeFn,
   ISuggestionItem,
 } from '@kbn/esql-ast/src/commands_registry/types';
 import { ESQLVariableType } from '@kbn/esql-types';
 import type { LicenseType } from '@kbn/licensing-types';
 import { isSourceCommand } from '../shared/helpers';
-import { collectUserDefinedColumns } from '../shared/user_defined_columns';
 import { getAstContext } from '../shared/context';
-import { getFieldsByTypeHelper, getSourcesHelper } from '../shared/resources_helpers';
+import { getColumnsByTypeHelper, getSourcesHelper } from '../shared/resources_helpers';
 import type { ESQLCallbacks } from '../shared/types';
 import { getQueryForFields } from './helper';
 import { mapRecommendedQueriesFromExtensions } from './utils/recommended_queries_helpers';
 import { getCommandContext } from './get_command_context';
 
-type GetFieldsMapFn = () => Promise<Map<string, ESQLFieldWithMetadata>>;
+type GetColumnMapFn = () => Promise<Map<string, ESQLColumnData>>;
 
 export async function suggest(
   fullText: string,
@@ -61,7 +59,7 @@ export async function suggest(
   // build the correct query to fetch the list of fields
   const queryForFields = getQueryForFields(correctedQuery, root);
 
-  const { getFieldsByType, getFieldsMap } = getFieldsByTypeRetriever(
+  const { getColumnsByType, getColumnMap } = getColumnsByTypeRetriever(
     queryForFields.replace(EDITOR_MARKER, ''),
     resourceRetriever,
     innerText
@@ -113,7 +111,7 @@ export async function suggest(
           fromCommand = `FROM ${visibleSources[0].name}`;
         }
 
-        const { getFieldsByType: getFieldsByTypeEmptyState } = getFieldsByTypeRetriever(
+        const { getColumnsByType: getColumnsByTypeEmptyState } = getColumnsByTypeRetriever(
           fromCommand,
           resourceRetriever,
           innerText
@@ -127,7 +125,7 @@ export async function suggest(
 
         const recommendedQueriesSuggestionsFromStaticTemplates =
           await getRecommendedQueriesSuggestionsFromStaticTemplates(
-            getFieldsByTypeEmptyState,
+            getColumnsByTypeEmptyState,
             fromCommand
           );
         recommendedQueriesSuggestions.push(
@@ -161,8 +159,8 @@ export async function suggest(
       fullText,
       ast,
       astContext,
-      getFieldsByType,
-      getFieldsMap,
+      getColumnsByType,
+      getColumnMap,
       resourceRetriever,
       offset,
       hasMinimumLicenseRequired
@@ -172,12 +170,12 @@ export async function suggest(
   return [];
 }
 
-export function getFieldsByTypeRetriever(
+export function getColumnsByTypeRetriever(
   queryForFields: string,
   resourceRetriever?: ESQLCallbacks,
   fullQuery?: string
-): { getFieldsByType: GetColumnsByTypeFn; getFieldsMap: GetFieldsMapFn } {
-  const helpers = getFieldsByTypeHelper(queryForFields, resourceRetriever);
+): { getColumnsByType: GetColumnsByTypeFn; getColumnMap: GetColumnMapFn } {
+  const helpers = getColumnsByTypeHelper(queryForFields, resourceRetriever);
   const getVariables = resourceRetriever?.getVariables;
   const canSuggestVariables = resourceRetriever?.canSuggestVariables?.() ?? false;
 
@@ -185,7 +183,7 @@ export function getFieldsByTypeRetriever(
   const lastCharacterTyped = queryString[queryString.length - 1];
   const lastCharIsQuestionMark = lastCharacterTyped === ESQL_VARIABLES_PREFIX;
   return {
-    getFieldsByType: async (
+    getColumnsByType: async (
       expectedType: Readonly<string> | Readonly<string[]> = 'any',
       ignored: string[] = [],
       options
@@ -199,19 +197,19 @@ export function getFieldsByTypeRetriever(
         recommendedFields: [],
       };
       const recommendedFieldsFromExtensions = editorExtensions.recommendedFields;
-      const fields = await helpers.getFieldsByType(expectedType, ignored);
+      const columns = await helpers.getColumnsByType(expectedType, ignored);
       return buildFieldsDefinitionsWithMetadata(
-        fields,
+        columns,
         recommendedFieldsFromExtensions,
         updatedOptions,
         await getVariables?.()
       );
     },
-    getFieldsMap: helpers.getFieldsMap,
+    getColumnMap: helpers.getColumnMap,
   };
 }
 
-function findNewUserDefinedColumn(userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>) {
+function findNewUserDefinedColumn(userDefinedColumns: Map<string, ESQLColumnData>) {
   let autoGeneratedColumnCounter = 0;
   let name = `col${autoGeneratedColumnCounter++}`;
   while (userDefinedColumns.has(name)) {
@@ -230,7 +228,7 @@ async function getSuggestionsWithinCommandExpression(
     containingFunction?: ESQLFunction;
   },
   getColumnsByType: GetColumnsByTypeFn,
-  getFieldsMap: GetFieldsMapFn,
+  getColumnMap: GetColumnMapFn,
   callbacks?: ESQLCallbacks,
   offset?: number,
   hasMinimumLicenseRequired?: (minimumLicenseRequired: LicenseType) => boolean
@@ -243,23 +241,19 @@ async function getSuggestionsWithinCommandExpression(
   }
 
   // collect all fields + userDefinedColumns to suggest
-  const fieldsMap: Map<string, ESQLFieldWithMetadata> = await getFieldsMap();
-  const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, innerText);
+  const columnMap: Map<string, ESQLColumnData> = await getColumnMap();
+  const references = { columns: columnMap };
 
-  const references = { fields: fieldsMap, userDefinedColumns: anyUserDefinedColumns };
-
-  const getSuggestedUserDefinedColumnName = (extraFieldNames?: string[]) => {
-    if (!extraFieldNames?.length) {
-      return findNewUserDefinedColumn(anyUserDefinedColumns);
+  const getSuggestedUserDefinedColumnName = (extraColumnNames?: string[]) => {
+    if (!extraColumnNames?.length) {
+      return findNewUserDefinedColumn(columnMap);
     }
 
-    const augmentedFieldsMap = new Map(fieldsMap);
-    extraFieldNames.forEach((name) => {
-      augmentedFieldsMap.set(name, { name, type: 'double' });
+    const augmentedColumnsMap = new Map(columnMap);
+    extraColumnNames.forEach((name) => {
+      augmentedColumnsMap.set(name, { name, type: 'double' });
     });
-    return findNewUserDefinedColumn(
-      collectUserDefinedColumns(commands, augmentedFieldsMap, innerText)
-    );
+    return findNewUserDefinedColumn(augmentedColumnsMap);
   };
 
   const additionalCommandContext = await getCommandContext(
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index 3506a3157646f..5ad8ea05f4060 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -62,8 +62,8 @@ async function cacheColumnsForQuery(queryText: string) {
   }
 }
 
-export function getFieldsByTypeHelper(queryText: string, resourceRetriever?: ESQLCallbacks) {
-  const getFields = async () => {
+export function getColumnsByTypeHelper(queryText: string, resourceRetriever?: ESQLCallbacks) {
+  const getColumns = async () => {
     // in some cases (as in the case of ROW or SHOW) the query is not set
     if (!queryText) {
       return;
@@ -85,12 +85,12 @@ export function getFieldsByTypeHelper(queryText: string, resourceRetriever?: ESQ
   };
 
   return {
-    getFieldsByType: async (
+    getColumnsByType: async (
       expectedType: Readonly<string> | Readonly<string[]> = 'any',
       ignored: string[] = []
     ): Promise<ESQLColumnData[]> => {
       const types = Array.isArray(expectedType) ? expectedType : [expectedType];
-      await getFields();
+      await getColumns();
       const queryTextForCacheSearch = toSingleLine(queryText);
       const cachedFields = getValueInsensitive(queryTextForCacheSearch);
       return (
@@ -105,8 +105,8 @@ export function getFieldsByTypeHelper(queryText: string, resourceRetriever?: ESQ
         }) || []
       );
     },
-    getFieldsMap: async (): Promise<Map<string, ESQLColumnData>> => {
-      await getFields();
+    getColumnMap: async (): Promise<Map<string, ESQLColumnData>> => {
+      await getColumns();
       const queryTextForCacheSearch = toSingleLine(queryText);
       const cachedFields = getValueInsensitive(queryTextForCacheSearch);
       const cacheCopy = new Map<string, ESQLColumnData>();
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
index e663109be1e59..9a0b08ad354b6 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
@@ -11,7 +11,7 @@ import { isSource, type ESQLCommand } from '@kbn/esql-ast';
 import type { ESQLFieldWithMetadata, ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types';
 import { createMapFromList, nonNullable } from '../shared/helpers';
 import {
-  getFieldsByTypeHelper,
+  getColumnsByTypeHelper,
   getPolicyHelper,
   getSourcesHelper,
 } from '../shared/resources_helpers';
@@ -67,5 +67,5 @@ export async function retrievePoliciesFields(
   const customQuery = buildQueryForFieldsInPolicies(
     policyNames.map((name) => policies.get(name)) as ESQLPolicy[]
   );
-  return await getFieldsByTypeHelper(customQuery, callbacks).getFieldsMap();
+  return await getColumnsByTypeHelper(customQuery, callbacks).getColumnMap();
 }
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/types.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/types.ts
index 5b690fc68135c..cfa50bfddcc1f 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/types.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/types.ts
@@ -8,15 +8,14 @@
  */
 
 import type { ESQLMessage } from '@kbn/esql-ast';
-import type { ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types';
-import type { ESQLPolicy, ESQLUserDefinedColumn } from '@kbn/esql-ast/src/commands_registry/types';
+import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types';
+import type { ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types';
 import type { IndexAutocompleteItem } from '@kbn/esql-types';
 import type { EditorError } from '../types';
 
 export interface ReferenceMaps {
   sources: Set<string>;
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>;
-  fields: Map<string, ESQLFieldWithMetadata>;
+  columns: Map<string, ESQLColumnData>;
   policies: Map<string, ESQLPolicy>;
   query: string;
   joinIndices: IndexAutocompleteItem[];
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index 6ed045a88a21b..63ce09106597e 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -17,7 +17,6 @@ import type {
 import type { LicenseType } from '@kbn/licensing-types';
 
 import type { ESQLCallbacks } from '../shared/types';
-import { collectUserDefinedColumns } from '../shared/user_defined_columns';
 import {
   retrievePolicies,
   // retrievePoliciesFields,
@@ -25,7 +24,7 @@ import {
 } from './resources';
 import type { ReferenceMaps, ValidationOptions, ValidationResult } from './types';
 import { getQueryForFields } from '../autocomplete/helper';
-import { getFieldsByTypeHelper } from '../shared/resources_helpers';
+import { getColumnsByTypeHelper } from '../shared/resources_helpers';
 
 /**
  * ES|QL validation public API
@@ -128,12 +127,17 @@ async function validateAst(
   //   fieldsFromPoliciesMap.forEach((value, key) => availableFields.set(key, value));
   // }
 
-  const sourceFields = await getFieldsByTypeHelper(
+  const sourceFields = await getColumnsByTypeHelper(
     queryString.split('|')[0],
     callbacks
-  ).getFieldsMap();
+  ).getColumnMap();
 
-  messages.push(...validateUnsupportedTypeFields(sourceFields, rootCommands));
+  messages.push(
+    ...validateUnsupportedTypeFields(
+      sourceFields as Map<string, ESQLFieldWithMetadata>,
+      rootCommands
+    )
+  );
 
   const license = await callbacks?.getLicense?.();
   const hasMinimumLicenseRequired = license?.hasAtLeast;
@@ -146,19 +150,13 @@ async function validateAst(
       Builder.expression.query(previousCommands)
     );
 
-    const { getFieldsMap } = getFieldsByTypeHelper(queryForFields, callbacks);
+    const { getColumnMap: getFieldsMap } = getColumnsByTypeHelper(queryForFields, callbacks);
     const availableFields = await getFieldsMap();
-    const userDefinedColumns = collectUserDefinedColumns(
-      previousCommands,
-      availableFields,
-      queryString
-    );
 
     const references: ReferenceMaps = {
       sources,
-      fields: availableFields,
+      columns: availableFields,
       policies: availablePolicies,
-      userDefinedColumns,
       query: queryString,
       joinIndices: joinIndices?.indices || [],
     };
@@ -227,9 +225,8 @@ function validateCommand(
   }
 
   const context = {
-    fields: references.fields,
+    columns: references.columns,
     policies: references.policies,
-    userDefinedColumns: references.userDefinedColumns,
     sources: [...references.sources].map((source) => ({
       name: source,
     })),
diff --git a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
index c849bdae0dfad..4a5803646b828 100644
--- a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
+++ b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
@@ -12,7 +12,6 @@ import {
   ENRICH_MODES,
   modeDescription,
 } from '@kbn/esql-ast/src/commands_registry/commands/enrich/util';
-import type { ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types';
 import {
   getFunctionDefinition,
   getFunctionSignatures,
@@ -26,7 +25,7 @@ import {
   type ESQLSource,
 } from '@kbn/esql-ast/src/types';
 import { collectUserDefinedColumns, type ESQLCallbacks } from '@kbn/esql-validation-autocomplete';
-import { getFieldsByTypeRetriever } from '@kbn/esql-validation-autocomplete/src/autocomplete/autocomplete';
+import { getColumnsByTypeRetriever } from '@kbn/esql-validation-autocomplete/src/autocomplete/autocomplete';
 import { getQueryForFields } from '@kbn/esql-validation-autocomplete/src/autocomplete/helper';
 import { getPolicyHelper } from '@kbn/esql-validation-autocomplete/src/shared/resources_helpers';
 import { i18n } from '@kbn/i18n';
@@ -187,19 +186,19 @@ async function getHintForFunctionArg(
   offset: number,
   resourceRetriever?: ESQLCallbacks
 ) {
-  const queryForFields = getQueryForFields(query, root);
-  const { getFieldsMap } = getFieldsByTypeRetriever(queryForFields, resourceRetriever);
+  const queryForColumns = getQueryForFields(query, root);
+  const { getColumnMap } = getColumnsByTypeRetriever(queryForColumns, resourceRetriever);
 
   const fnDefinition = getFunctionDefinition(fnNode.name);
   // early exit on no hit
   if (!fnDefinition) {
     return [];
   }
-  const fieldsMap: Map<string, ESQLFieldWithMetadata> = await getFieldsMap();
-  const anyUserDefinedColumns = collectUserDefinedColumns(root.commands, fieldsMap, query);
+  const columnsMap = await getColumnMap();
+  const anyUserDefinedColumns = collectUserDefinedColumns(root.commands, columnsMap, query);
 
   const references = {
-    fields: fieldsMap,
+    columns: columnsMap,
     userDefinedColumns: anyUserDefinedColumns,
   };
 

From 78d31d581b9363855e0e04b62bf75cbdb361d5f7 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 12:30:04 -0600
Subject: [PATCH 10/54] fix most non-test type errors

---
 .../src/__tests__/autocomplete.ts             |  17 +-
 .../src/__tests__/context_fixtures.ts         | 118 ++++++-------
 .../kbn-esql-ast/src/__tests__/validation.ts  |   5 +-
 .../commands/change_point/validate.ts         |   1 +
 .../commands/completion/autocomplete.ts       |   3 +-
 .../commands/enrich/autocomplete.ts           |   4 +-
 .../commands/fork/columns_after.ts            |   4 +-
 .../commands/fork/validate.ts                 |   1 +
 .../commands_registry/commands/join/utils.ts  |   7 +-
 .../commands/keep/columns_after.test.ts       |   8 +-
 .../commands/rename/columns_after.test.ts     |  16 +-
 .../commands/stats/columns_after.test.ts      |  33 ++--
 .../src/commands_registry/types.ts            |   2 +-
 .../definitions/utils/autocomplete/columns.ts | 161 +-----------------
 .../utils/autocomplete/functions.ts           |  43 +++--
 .../definitions/utils/autocomplete/helpers.ts |  42 +----
 .../src/definitions/utils/columns.ts          |  10 +-
 .../src/definitions/utils/expressions.test.ts |  45 ++---
 .../src/definitions/utils/operators.ts        |   2 +-
 .../definitions/utils/validation/commands.ts  |   3 +-
 .../definitions/utils/validation/function.ts  |   4 +-
 .../src/shared/user_defined_columns.ts        | 114 -------------
 22 files changed, 173 insertions(+), 470 deletions(-)
 delete mode 100644 src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts
index d588d9d2ab9ba..d6cfa61ac3ae2 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts
@@ -15,11 +15,10 @@
 import { uniq } from 'lodash';
 import type { LicenseType } from '@kbn/licensing-types';
 import type {
-  ESQLUserDefinedColumn,
-  ESQLFieldWithMetadata,
   ICommandCallbacks,
   ISuggestionItem,
   Location,
+  ESQLColumnData,
 } from '../commands_registry/types';
 import { getLocationFromCommandOrOptionName } from '../commands_registry/types';
 import { aggFunctionDefinitions } from '../definitions/generated/aggregation_functions';
@@ -50,8 +49,7 @@ export const suggest = (
     arg1: ESQLCommand,
     arg2: ICommandCallbacks,
     arg3: {
-      userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>;
-      fields: Map<string, ESQLFieldWithMetadata>;
+      columns: Map<string, ESQLColumnData>;
     },
     arg4?: number
   ) => Promise<ISuggestionItem[]>,
@@ -80,8 +78,7 @@ export const expectSuggestions = async (
     arg1: ESQLCommand,
     arg2: ICommandCallbacks,
     arg3: {
-      userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>;
-      fields: Map<string, ESQLFieldWithMetadata>;
+      columns: Map<string, ESQLColumnData>;
     },
     arg4?: number
   ) => Promise<ISuggestionItem[]>,
@@ -100,12 +97,10 @@ export function getFieldNamesByType(
   _requestedType: Readonly<FieldType | 'any' | Array<FieldType | 'any'>>,
   excludeUserDefined: boolean = false
 ) {
-  const fieldsMap = mockContext.fields;
-  const userDefinedColumnsMap = mockContext.userDefinedColumns;
-  const fields = Array.from(fieldsMap.values());
-  const userDefinedColumns = Array.from(userDefinedColumnsMap.values()).flat();
+  const columnMap = mockContext.columns;
+  const columns = Array.from(columnMap.values());
   const requestedType = Array.isArray(_requestedType) ? _requestedType : [_requestedType];
-  const finalArray = excludeUserDefined ? fields : [...fields, ...userDefinedColumns];
+  const finalArray = excludeUserDefined ? columns.filter((col) => !col.userDefined) : columns;
   return finalArray
     .filter(
       ({ type }) =>
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/context_fixtures.ts b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/context_fixtures.ts
index 8522f18e26891..4106ddce9a729 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/context_fixtures.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/context_fixtures.ts
@@ -8,9 +8,8 @@
  */
 import type { IndexAutocompleteItem, InferenceEndpointAutocompleteItem } from '@kbn/esql-types';
 import type {
-  ESQLFieldWithMetadata,
+  ESQLColumnData,
   ESQLPolicy,
-  ESQLUserDefinedColumn,
   ICommandCallbacks,
   ICommandContext,
 } from '../commands_registry/types';
@@ -115,84 +114,79 @@ export const editorExtensions = {
 };
 
 export const mockContext: ICommandContext = {
-  userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([
+  columns: new Map<string, ESQLColumnData>([
     [
       'var0',
-      [
-        {
-          name: 'var0',
-          type: 'double',
-          location: { min: 0, max: 10 },
-        },
-      ],
+      {
+        name: 'var0',
+        type: 'double',
+        location: { min: 0, max: 10 },
+        userDefined: true,
+      },
     ],
     [
       'col0',
-      [
-        {
-          name: 'col0',
-          type: 'double',
-          location: { min: 0, max: 10 },
-        },
-      ],
+      {
+        name: 'col0',
+        type: 'double',
+        location: { min: 0, max: 10 },
+        userDefined: true,
+      },
     ],
     [
       'prompt',
-      [
-        {
-          name: 'prompt',
-          type: 'keyword',
-          location: { min: 0, max: 10 },
-        },
-      ],
+      {
+        name: 'prompt',
+        type: 'keyword',
+        location: { min: 0, max: 10 },
+        userDefined: true,
+      },
     ],
     [
       'integerPrompt',
-      [
-        {
-          name: 'integerPrompt',
-          type: 'integer',
-          location: { min: 0, max: 10 },
-        },
-      ],
+      {
+        name: 'integerPrompt',
+        type: 'integer',
+        location: { min: 0, max: 10 },
+        userDefined: true,
+      },
     ],
     [
       'ipPrompt',
-      [
-        {
-          name: 'ipPrompt',
-          type: 'ip',
-          location: { min: 0, max: 10 },
-        },
-      ],
+      {
+        name: 'ipPrompt',
+        type: 'ip',
+        location: { min: 0, max: 10 },
+        userDefined: true,
+      },
     ],
     [
       'renamedField',
-      [
-        {
-          name: 'renamedField',
-          type: 'keyword',
-          location: { min: 0, max: 10 },
-        },
-      ],
+      {
+        name: 'renamedField',
+        type: 'keyword',
+        location: { min: 0, max: 10 },
+        userDefined: true,
+      },
     ],
-  ]),
-  fields: new Map<string, ESQLFieldWithMetadata>([
-    ['keywordField', { name: 'keywordField', type: 'keyword' }],
-    ['any#Char$Field', { name: 'any#Char$Field', type: 'keyword' }],
-    ['textField', { name: 'textField', type: 'text' }],
-    ['doubleField', { name: 'doubleField', type: 'double' }],
-    ['integerField', { name: 'integerField', type: 'integer' }],
-    ['counterIntegerField', { name: 'counterIntegerField', type: 'counter_integer' }],
-    ['dateField', { name: 'dateField', type: 'date' }],
-    ['dateNanosField', { name: 'dateNanosField', type: 'date_nanos' }],
-    ['@timestamp', { name: '@timestamp', type: 'date' }],
-    ['ipField', { name: 'ipField', type: 'ip' }],
-    ['booleanField', { name: 'booleanField', type: 'boolean' }],
-    ['geoPointField', { name: 'geoPointField', type: 'geo_point' }],
-    ['geoShapeField', { name: 'geoShapeField', type: 'geo_shape' }],
-    ['versionField', { name: 'versionField', type: 'version' }],
-    ['longField', { name: 'longField', type: 'long' }],
+    ['keywordField', { name: 'keywordField', type: 'keyword', userDefined: false }],
+    ['any#Char$Field', { name: 'any#Char$Field', type: 'keyword', userDefined: false }],
+    ['textField', { name: 'textField', type: 'text', userDefined: false }],
+    ['doubleField', { name: 'doubleField', type: 'double', userDefined: false }],
+    ['integerField', { name: 'integerField', type: 'integer', userDefined: false }],
+    [
+      'counterIntegerField',
+      { name: 'counterIntegerField', type: 'counter_integer', userDefined: false },
+    ],
+    ['dateField', { name: 'dateField', type: 'date', userDefined: false }],
+    ['dateNanosField', { name: 'dateNanosField', type: 'date_nanos', userDefined: false }],
+    ['@timestamp', { name: '@timestamp', type: 'date', userDefined: false }],
+    ['ipField', { name: 'ipField', type: 'ip', userDefined: false }],
+    ['booleanField', { name: 'booleanField', type: 'boolean', userDefined: false }],
+    ['geoPointField', { name: 'geoPointField', type: 'geo_point', userDefined: false }],
+    ['geoShapeField', { name: 'geoShapeField', type: 'geo_shape', userDefined: false }],
+    ['versionField', { name: 'versionField', type: 'version', userDefined: false }],
+    ['longField', { name: 'longField', type: 'long', userDefined: false }],
   ]),
   policies: new Map<string, ESQLPolicy>(policies.map((policy) => [policy.name, policy])),
   sources: indexes.map((name) => ({
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/validation.ts b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/validation.ts
index 859db40e90303..3652598d184d0 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/validation.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import type { ESQLUserDefinedColumn, ESQLFieldWithMetadata } from '../commands_registry/types';
+import type { ESQLColumnData } from '../commands_registry/types';
 import { Parser } from '../parser';
 import type { ESQLCommand, ESQLMessage } from '../types';
 import { mockContext } from './context_fixtures';
@@ -28,8 +28,7 @@ export const expectErrors = (
     arg0: ESQLCommand,
     arg1: ESQLCommand[],
     arg2: {
-      userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>;
-      fields: Map<string, ESQLFieldWithMetadata>;
+      columns: Map<string, ESQLColumnData>;
     }
   ) => any
 ) => {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
index 3dee57530d4aa..20c7ae16bb757 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
@@ -72,6 +72,7 @@ export const validate = (
           name: arg.name,
           location: arg.location,
           type: index === 0 ? 'keyword' : 'long',
+          userDefined: true,
         });
       }
     });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts
index 5f05b6bd566e6..a10f5f8100ad3 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts
@@ -157,8 +157,7 @@ export async function autocomplete(
           callbacks?.getByType,
           {
             functions: true,
-            fields: true,
-            userDefinedColumns: context?.userDefinedColumns,
+            columns: true,
           },
           {},
           callbacks?.hasMinimumLicenseRequired,
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts
index c9d0326cee9ff..29ab1d1a69a96 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts
@@ -41,8 +41,8 @@ export async function autocomplete(
   const innerText = query.substring(0, cursorPosition);
   const pos = getPosition(innerText, command);
   const policies = context?.policies ?? new Map<string, ESQLPolicy>();
-  const fieldsMap = context?.fields ?? new Map<string, string>();
-  const allColumnNames = Array.from(fieldsMap.keys());
+  const columnMap = context?.columns ?? new Map<string, string>();
+  const allColumnNames = Array.from(columnMap.keys());
   const policyName = (
     command.args.find((arg) => !Array.isArray(arg) && arg.type === 'source') as
       | ESQLSource
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index ced310189f4e5..7cbc681b9eee0 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -8,12 +8,12 @@
  */
 import { uniqBy } from 'lodash';
 import { type ESQLCommand } from '../../../types';
-import type { ESQLFieldWithMetadata } from '../../types';
+import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
 import type { ICommandContext } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
-  previousColumns: ESQLFieldWithMetadata[],
+  previousColumns: ESQLColumnData[],
   context?: ICommandContext
 ) => {
   return uniqBy(
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
index 36c103749b34e..677eba017ce7c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
@@ -71,6 +71,7 @@ export const validate = (
   context?.columns.set('_fork', {
     name: '_fork',
     type: 'keyword',
+    userDefined: false,
   });
 
   return messages;
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts
index d8a1ebb396307..ce3986a179bf1 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts
@@ -20,6 +20,7 @@ import type {
   GetColumnsByTypeFn,
   ISuggestionItem,
   ICommandContext,
+  ESQLColumnData,
 } from '../../types';
 import type { JoinCommandPosition, JoinPosition, JoinStaticPosition } from './types';
 import { TRIGGER_SUGGESTION_COMMAND } from '../../constants';
@@ -153,7 +154,7 @@ export const suggestFields = async (
   innerText: string,
   command: ESQLCommand,
   getColumnsByType: GetColumnsByTypeFn,
-  getColumnsForQuery: (query: string) => Promise<ESQLFieldWithMetadata[]>,
+  getColumnsForQuery: (query: string) => Promise<ESQLColumnData[]>,
   context?: ICommandContext
 ) => {
   if (!context) {
@@ -163,7 +164,9 @@ export const suggestFields = async (
   const { suggestions: fieldSuggestions, lookupIndexFieldExists } = await getFieldSuggestions(
     command,
     getColumnsByType,
-    getColumnsForQuery,
+    // this type cast is ok because getFieldSuggestions only ever fetches columns
+    // from a bare FROM clause, so they will always be fields, not user-defined columns
+    getColumnsForQuery as (query: string) => Promise<ESQLFieldWithMetadata[]>,
     context
   );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
index b68f480056a0d..b287b9a3c7541 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
@@ -14,14 +14,14 @@ describe('KEEP', () => {
   const context = {
     userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
     fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword' }],
-      ['count', { name: 'count', type: 'double' }],
+      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+      ['count', { name: 'count', type: 'double', userDefined: false }],
     ]),
   };
   it('should return the correct fields after the command', () => {
     const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
     ] as ESQLFieldWithMetadata[];
 
     const result = columnsAfter(synth.cmd`KEEP field1`, previousCommandFields, context);
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
index ae475e70550d2..5a6977b13b595 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
@@ -14,14 +14,14 @@ describe('RENAME', () => {
   const context = {
     userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
     fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword' }],
-      ['count', { name: 'count', type: 'double' }],
+      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+      ['count', { name: 'count', type: 'double', userDefined: false }],
     ]),
   };
   it('renames the given columns with the new names using AS', () => {
     const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
     ] as ESQLFieldWithMetadata[];
 
     const result = columnsAfter(synth.cmd`RENAME field1 as meow`, previousCommandFields, context);
@@ -41,8 +41,8 @@ describe('RENAME', () => {
     const result = columnsAfter(synth.cmd`RENAME meow = field1`, previousCommandFields, context);
 
     expect(result).toEqual([
-      { name: 'meow', type: 'keyword' },
-      { name: 'field2', type: 'double' },
+      { name: 'meow', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
     ]);
   });
 
@@ -59,8 +59,8 @@ describe('RENAME', () => {
     );
 
     expect(result).toEqual([
-      { name: 'meow', type: 'keyword' },
-      { name: 'woof', type: 'double' },
+      { name: 'meow', type: 'keyword', userDefined: false },
+      { name: 'woof', type: 'double', userDefined: false },
     ]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
index f6d74a8e3ee9b..46a26ceb754d9 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
@@ -25,19 +25,20 @@ describe('STATS', () => {
               name: 'var0',
               type: 'double',
               location: { min: 0, max: 10 },
+              userDefined: true,
             },
           ],
         ],
       ]),
       fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword' }],
-        ['count', { name: 'count', type: 'double' }],
+        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+        ['count', { name: 'count', type: 'double', userDefined: false }],
       ]),
     };
 
     const result = columnsAfter(synth.cmd`STATS var0=AVG(field2)`, previousCommandFields, context);
 
-    expect(result).toEqual([{ name: 'var0', type: 'double' }]);
+    expect(result).toEqual([{ name: 'var0', type: 'double', userDefined: false }]);
   });
 
   it('adds the escaped column, when no grouping is given', () => {
@@ -55,19 +56,20 @@ describe('STATS', () => {
               name: 'AVG(field2)',
               type: 'double',
               location: { min: 0, max: 10 },
+              userDefined: true,
             },
           ],
         ],
       ]),
       fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword' }],
-        ['count', { name: 'count', type: 'double' }],
+        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+        ['count', { name: 'count', type: 'double', userDefined: false }],
       ]),
     };
 
     const result = columnsAfter(synth.cmd`STATS AVG(field2)`, previousCommandFields, context);
 
-    expect(result).toEqual([{ name: 'AVG(field2)', type: 'double' }]);
+    expect(result).toEqual([{ name: 'AVG(field2)', type: 'double', userDefined: false }]);
   });
 
   it('adds the escaped and grouping columns', () => {
@@ -85,13 +87,14 @@ describe('STATS', () => {
               name: 'AVG(field2)',
               type: 'double',
               location: { min: 0, max: 10 },
+              userDefined: true,
             },
           ],
         ],
       ]),
       fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword' }],
-        ['count', { name: 'count', type: 'double' }],
+        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+        ['count', { name: 'count', type: 'double', userDefined: false }],
       ]),
     };
 
@@ -102,8 +105,8 @@ describe('STATS', () => {
     );
 
     expect(result).toEqual([
-      { name: 'field1', type: 'keyword' },
-      { name: 'AVG(field2)', type: 'double' },
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'AVG(field2)', type: 'double', userDefined: false },
     ]);
   });
 
@@ -123,6 +126,7 @@ describe('STATS', () => {
               name: 'AVG(field2)',
               type: 'double',
               location: { min: 0, max: 10 },
+              userDefined: true,
             },
           ],
         ],
@@ -133,13 +137,14 @@ describe('STATS', () => {
               name: 'buckets',
               type: 'unknown',
               location: { min: 0, max: 10 },
+              userDefined: true,
             },
           ],
         ],
       ]),
       fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword' }],
-        ['count', { name: 'count', type: 'double' }],
+        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+        ['count', { name: 'count', type: 'double', userDefined: false }],
       ]),
     };
 
@@ -150,8 +155,8 @@ describe('STATS', () => {
     );
 
     expect(result).toEqual([
-      { name: 'AVG(field2)', type: 'double' },
-      { name: 'buckets', type: 'unknown' },
+      { name: 'AVG(field2)', type: 'double', userDefined: false },
+      { name: 'buckets', type: 'unknown', userDefined: false },
     ]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
index edefc48b218fe..ab84eab6667df 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts
@@ -101,7 +101,7 @@ export type GetColumnsByTypeFn = (
 export interface ESQLFieldWithMetadata {
   name: string;
   type: FieldType;
-  userDefined?: false; // TODO consider making required property
+  userDefined: false;
   isEcs?: boolean;
   hasConflict?: boolean;
   metadata?: {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
index 0f921740518ae..3966ee1eee529 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
@@ -6,36 +6,15 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { Visitor } from '../../../visitor';
 import type {
   ESQLFieldWithMetadata,
   ESQLUserDefinedColumn,
   ICommandContext,
 } from '../../../commands_registry/types';
-import type {
-  ESQLFunction,
-  ESQLAst,
-  ESQLAstItem,
-  ESQLCommand,
-  ESQLColumn,
-  ESQLIdentifier,
-} from '../../../types';
-import { isColumn, isFunctionExpression } from '../../../ast/is';
-import { getExpressionType } from '../expressions';
-import { EDITOR_MARKER } from '../../constants';
+import type { ESQLColumn, ESQLCommand, ESQLIdentifier } from '../../../types';
 import { fuzzySearch } from '../shared';
 
-function addToUserDefinedColumnOccurrences(
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>,
-  instance: ESQLUserDefinedColumn
-) {
-  if (!userDefinedColumns.has(instance.name)) {
-    userDefinedColumns.set(instance.name, []);
-  }
-  const userDefinedColumnsOccurrencies = userDefinedColumns.get(instance.name)!;
-  userDefinedColumnsOccurrencies.push(instance);
-}
-
+// FIXME
 export function excludeUserDefinedColumnsFromCurrentCommand(
   commands: ESQLCommand[],
   currentCommand: ESQLCommand,
@@ -57,147 +36,17 @@ export function excludeUserDefinedColumnsFromCurrentCommand(
   return resultUserDefinedColumns;
 }
 
-function addUserDefinedColumnFromAssignment(
-  assignOperation: ESQLFunction,
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>,
-  fields: Map<string, ESQLFieldWithMetadata>
-) {
-  if (isColumn(assignOperation.args[0])) {
-    const rightHandSideArgType = getExpressionType(assignOperation.args[1], fields);
-    addToUserDefinedColumnOccurrences(userDefinedColumns, {
-      name: assignOperation.args[0].parts.join('.'),
-      type: rightHandSideArgType /* fallback to number */,
-      location: assignOperation.args[0].location,
-    });
-  }
-}
-
-function addUserDefinedColumnFromExpression(
-  expressionOperation: ESQLFunction,
-  queryString: string,
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>,
-  fields: Map<string, ESQLFieldWithMetadata>
-) {
-  if (!expressionOperation.text.includes(EDITOR_MARKER)) {
-    const expressionText = queryString.substring(
-      expressionOperation.location.min,
-      expressionOperation.location.max + 1
-    );
-    const expressionType = getExpressionType(expressionOperation, fields);
-    addToUserDefinedColumnOccurrences(userDefinedColumns, {
-      name: expressionText,
-      type: expressionType,
-      location: expressionOperation.location,
-    });
-  }
-}
-
-function addToUserDefinedColumns(
-  oldArg: ESQLAstItem,
-  newArg: ESQLAstItem,
-  fields: Map<string, ESQLFieldWithMetadata>,
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>
-) {
-  if (isColumn(oldArg) && isColumn(newArg)) {
-    const newUserDefinedColumn: ESQLUserDefinedColumn = {
-      name: newArg.parts.join('.'),
-      type: 'double' /* fallback to number */,
-      location: newArg.location,
-    };
-    // Now workout the exact type
-    // it can be a rename of another userDefinedColumn as well
-    const oldRef =
-      fields.get(oldArg.parts.join('.')) || userDefinedColumns.get(oldArg.parts.join('.'));
-    if (oldRef) {
-      addToUserDefinedColumnOccurrences(userDefinedColumns, newUserDefinedColumn);
-      newUserDefinedColumn.type = Array.isArray(oldRef) ? oldRef[0].type : oldRef.type;
-    }
-  }
-}
-
-export function collectUserDefinedColumns(
-  ast: ESQLAst,
-  fields: Map<string, ESQLFieldWithMetadata>,
-  queryString: string
-): Map<string, ESQLUserDefinedColumn[]> {
-  const userDefinedColumns = new Map<string, ESQLUserDefinedColumn[]>();
-
-  const visitor = new Visitor()
-    .on('visitLiteralExpression', (ctx) => {
-      // TODO - add these as userDefinedColumns
-    })
-    .on('visitExpression', (_ctx) => {}) // required for the types :shrug:
-    .on('visitFunctionCallExpression', (ctx) => {
-      const node = ctx.node;
-
-      if (node.subtype === 'binary-expression' && node.name === 'where') {
-        ctx.visitArgument(0, undefined);
-        return;
-      }
-
-      if (node.name === 'as') {
-        const [oldArg, newArg] = ctx.node.args;
-        addToUserDefinedColumns(oldArg, newArg, fields, userDefinedColumns);
-      } else if (node.name === '=') {
-        addUserDefinedColumnFromAssignment(node, userDefinedColumns, fields);
-      } else {
-        addUserDefinedColumnFromExpression(node, queryString, userDefinedColumns, fields);
-      }
-    })
-    .on('visitCommandOption', (ctx) => {
-      if (ctx.node.name === 'by') {
-        return [...ctx.visitArguments()];
-      } else if (ctx.node.name === 'with') {
-        for (const assignFn of ctx.node.args) {
-          if (isFunctionExpression(assignFn)) {
-            const [newArg, oldArg] = assignFn?.args || [];
-            // TODO why is oldArg an array?
-            if (Array.isArray(oldArg)) {
-              addToUserDefinedColumns(oldArg[0], newArg, fields, userDefinedColumns);
-            }
-          }
-        }
-      }
-    })
-    .on('visitCommand', (ctx) => {
-      const ret = [];
-      if (['row', 'eval', 'stats', 'inlinestats', 'ts', 'rename'].includes(ctx.node.name)) {
-        ret.push(...ctx.visitArgs());
-      }
-      if (['stats', 'inlinestats', 'enrich'].includes(ctx.node.name)) {
-        // BY and WITH can contain userDefinedColumns
-        ret.push(...ctx.visitOptions());
-      }
-      if (ctx.node.name === 'fork') {
-        ret.push(...ctx.visitSubQueries());
-      }
-      return ret;
-    })
-    .on('visitQuery', (ctx) => [...ctx.visitCommands()]);
-
-  visitor.visitQuery(ast);
-
-  return userDefinedColumns;
-}
-
 /**
  * TODO - consider calling lookupColumn under the hood of this function. Seems like they should really do the same thing.
  */
-export function getColumnExists(
-  node: ESQLColumn | ESQLIdentifier,
-  { fields, userDefinedColumns }: ICommandContext
-) {
+export function getColumnExists(node: ESQLColumn | ESQLIdentifier, { columns }: ICommandContext) {
   const columnName = node.type === 'identifier' ? node.name : node.parts.join('.');
-  if (fields.has(columnName) || userDefinedColumns.has(columnName)) {
+  if (columns.has(columnName)) {
     return true;
   }
 
   // TODO — I don't see this fuzzy searching in lookupColumn... should it be there?
-  if (
-    Boolean(
-      fuzzySearch(columnName, fields.keys()) || fuzzySearch(columnName, userDefinedColumns.keys())
-    )
-  ) {
+  if (Boolean(fuzzySearch(columnName, columns.keys()))) {
     return true;
   }
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
index 3b5da9c6397bf..cca9bb69ef75f 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
@@ -24,7 +24,7 @@ import {
   listCompleteItem,
 } from '../../../commands_registry/complete_items';
 import type {
-  ESQLFieldWithMetadata,
+  ESQLColumnData,
   GetColumnsByTypeFn,
   ICommandCallbacks,
   ICommandContext,
@@ -51,7 +51,7 @@ import {
 import { getCompatibleLiterals, getDateLiterals } from '../literals';
 import { getSuggestionsToRightOfOperatorExpression } from '../operators';
 import { buildValueDefinitions } from '../values';
-import { collectUserDefinedColumns, excludeUserDefinedColumnsFromCurrentCommand } from './columns';
+import { excludeUserDefinedColumnsFromCurrentCommand } from './columns';
 import {
   extractTypeFromASTArg,
   getFieldsOrFunctionsSuggestions,
@@ -167,17 +167,15 @@ export async function getFunctionArgsSuggestions(
     signatures: filterFunctionSignatures(fnDefinition.signatures, hasMinimumLicenseRequired),
   };
 
-  const fieldsMap: Map<string, ESQLFieldWithMetadata> = context?.fields || new Map();
-  const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, innerText);
+  const columnMap: Map<string, ESQLColumnData> = context?.columns || new Map();
 
   const references = {
-    fields: fieldsMap,
-    userDefinedColumns: anyUserDefinedColumns,
+    columns: columnMap,
   };
   const userDefinedColumnsExcludingCurrentCommandOnes = excludeUserDefinedColumnsFromCurrentCommand(
     commands,
     command,
-    fieldsMap,
+    columnMap,
     innerText
   );
 
@@ -232,8 +230,7 @@ export async function getFunctionArgsSuggestions(
     arg &&
     isColumn(arg) &&
     !getColumnExists(arg, {
-      fields: fieldsMap,
-      userDefinedColumns: userDefinedColumnsExcludingCurrentCommandOnes,
+      columns: columnMap,
     });
   if (noArgDefined || isUnknownColumn) {
     // ... | EVAL fn( <suggest>)
@@ -437,7 +434,7 @@ async function getListArgsSuggestions(
   innerText: string,
   commands: ESQLCommand[],
   getFieldsByType: GetColumnsByTypeFn,
-  fieldsMap: Map<string, ESQLFieldWithMetadata>,
+  columnMap: Map<string, ESQLColumnData>,
   offset: number,
   hasMinimumLicenseRequired?: (minimumLicenseRequired: LicenseType) => boolean,
   activeProduct?: PricingProduct
@@ -460,18 +457,18 @@ async function getListArgsSuggestions(
       }
     }
 
-    const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, innerText);
-    // extract the current node from the userDefinedColumns inferred
-    anyUserDefinedColumns.forEach((values, key) => {
-      if (values.some((v) => v.location === node.location)) {
-        anyUserDefinedColumns.delete(key);
-      }
-    });
+    // FIXME
+    // const anyUserDefinedColumns = collectUserDefinedColumns(commands, columnMap, innerText);
+    // // extract the current node from the userDefinedColumns inferred
+    // anyUserDefinedColumns.forEach((values, key) => {
+    //   if (values.some((v) => v.location === node.location)) {
+    //     anyUserDefinedColumns.delete(key);
+    //   }
+    // });
     const [firstArg] = node.args;
     if (isColumn(firstArg)) {
       const argType = extractTypeFromASTArg(firstArg, {
-        fields: fieldsMap,
-        userDefinedColumns: anyUserDefinedColumns,
+        columns: columnMap,
       });
       if (argType) {
         // do not propose existing columns again
@@ -485,8 +482,7 @@ async function getListArgsSuggestions(
             getFieldsByType,
             {
               functions: true,
-              fields: true,
-              userDefinedColumns: anyUserDefinedColumns,
+              columns: true,
             },
             { ignoreColumns: [firstArg.name, ...otherArgs.map(({ name }) => name)] },
             hasMinimumLicenseRequired,
@@ -536,8 +532,7 @@ export const getInsideFunctionsSuggestions = async (
         queryText: innerText,
         location: getLocationFromCommandOrOptionName(command.name),
         rootOperator: node,
-        getExpressionType: (expression) =>
-          getExpressionType(expression, context?.fields, context?.userDefinedColumns),
+        getExpressionType: (expression) => getExpressionType(expression, context?.columns),
         getColumnsByType: callbacks?.getByType ?? (() => Promise.resolve([])),
         hasMinimumLicenseRequired: callbacks?.hasMinimumLicenseRequired,
         activeProduct: context?.activeProduct,
@@ -550,7 +545,7 @@ export const getInsideFunctionsSuggestions = async (
         innerText,
         ast,
         callbacks?.getByType ?? (() => Promise.resolve([])),
-        context?.fields ?? new Map(),
+        context?.columns ?? new Map(),
         cursorPosition ?? 0,
         callbacks?.hasMinimumLicenseRequired,
         context?.activeProduct
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
index 5537f876bb4bd..629e9bdae5957 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
@@ -22,7 +22,6 @@ import type {
 import type {
   ISuggestionItem,
   GetColumnsByTypeFn,
-  ESQLUserDefinedColumn,
   ICommandContext,
 } from '../../../commands_registry/types';
 import { Location } from '../../../commands_registry/types';
@@ -169,14 +168,12 @@ export async function getFieldsOrFunctionsSuggestions(
   getFieldsByType: GetColumnsByTypeFn,
   {
     functions,
-    fields,
-    userDefinedColumns,
+    columns: fields,
     values = false,
     literals = false,
   }: {
     functions: boolean;
-    fields: boolean;
-    userDefinedColumns?: Map<string, ESQLUserDefinedColumn[]>;
+    columns: boolean;
     literals?: boolean;
     values?: boolean;
   },
@@ -201,34 +198,6 @@ export async function getFieldsOrFunctionsSuggestions(
     functions
   );
 
-  const filteredColumnByType: string[] = [];
-  if (userDefinedColumns) {
-    for (const userDefinedColumn of userDefinedColumns.values()) {
-      if (
-        (types.includes('any') || types.includes(userDefinedColumn[0].type)) &&
-        !ignoreColumns.includes(userDefinedColumn[0].name)
-      ) {
-        filteredColumnByType.push(userDefinedColumn[0].name);
-      }
-    }
-    // due to a bug on the ES|QL table side, filter out fields list with underscored userDefinedColumns names (??)
-    // avg( numberField ) => avg_numberField_
-    const ALPHANUMERIC_REGEXP = /[^a-zA-Z\d]/g;
-    if (
-      filteredColumnByType.length &&
-      filteredColumnByType.some((v) => ALPHANUMERIC_REGEXP.test(v))
-    ) {
-      for (const userDefinedColumn of filteredColumnByType) {
-        const underscoredName = userDefinedColumn.replace(ALPHANUMERIC_REGEXP, '_');
-        const index = filteredFieldsByType.findIndex(
-          ({ label }) => underscoredName === label || `_${underscoredName}_` === label
-        );
-        if (index >= 0) {
-          filteredFieldsByType.splice(index);
-        }
-      }
-    }
-  }
   // could also be in stats (bucket) but our autocomplete is not great yet
   const displayDateSuggestions =
     types.includes('date') && [Location.WHERE, Location.EVAL].includes(location);
@@ -246,9 +215,10 @@ export async function getFieldsOrFunctionsSuggestions(
           activeProduct
         )
       : [],
-    userDefinedColumns
-      ? pushItUpInTheList(buildUserDefinedColumnsDefinitions(filteredColumnByType), functions)
-      : [],
+    // FIXME
+    // userDefinedColumns
+    //   ? pushItUpInTheList(buildUserDefinedColumnsDefinitions(filteredColumnByType), functions)
+    //   : [],
     literals ? getCompatibleLiterals(types) : []
   );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/columns.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/columns.ts
index c6661f38f173d..5e59f24313cd3 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/columns.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/columns.ts
@@ -16,12 +16,18 @@ import { fuzzySearch } from './shared';
  */
 export function getColumnExists(
   node: ESQLColumn | ESQLIdentifier,
-  { fields, userDefinedColumns }: Pick<ICommandContext, 'fields' | 'userDefinedColumns'>,
+  { columns }: Pick<ICommandContext, 'columns'>,
   excludeFields = false
 ) {
   const columnName = node.type === 'identifier' ? node.name : node.parts.join('.');
 
-  const set = new Set([...(!excludeFields ? fields.keys() : []), ...userDefinedColumns.keys()]);
+  const set = new Set(
+    !excludeFields
+      ? columns.keys()
+      : Array.from(columns.values())
+          .filter((col) => col.userDefined)
+          .map((col) => col.name)
+  );
 
   if (set.has(columnName)) {
     return true;
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.test.ts
index 31e68a6f444d2..777783f941db6 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.test.ts
@@ -9,6 +9,7 @@
 import { Parser } from '../../parser';
 import type { SupportedDataType } from '../types';
 import { FunctionDefinitionTypes } from '../types';
+import type { ESQLColumnData } from '../../commands_registry/types';
 import { Location } from '../../commands_registry/types';
 import { buildPartialMatcher, getExpressionType } from './expressions';
 import { setTestFunctions } from './test_functions';
@@ -127,27 +128,25 @@ describe('getExpressionType', () => {
               {
                 name: 'fieldName',
                 type: 'geo_shape',
+                userDefined: false,
               },
             ],
-          ]),
-          new Map()
+          ])
         )
       ).toBe('geo_shape');
 
       expect(
         getExpressionType(
           getASTForExpression('col0'),
-          new Map(),
           new Map([
             [
               'col0',
-              [
-                {
-                  name: 'col0',
-                  type: 'long',
-                  location: { min: 0, max: 0 },
-                },
-              ],
+              {
+                name: 'col0',
+                type: 'long',
+                location: { min: 0, max: 0 },
+                userDefined: true,
+              },
             ],
           ])
         )
@@ -155,9 +154,7 @@ describe('getExpressionType', () => {
     });
 
     it('handles fields and userDefinedColumns which do not exist', () => {
-      expect(getExpressionType(getASTForExpression('fieldName'), new Map(), new Map())).toBe(
-        'unknown'
-      );
+      expect(getExpressionType(getASTForExpression('fieldName'), new Map())).toBe('unknown');
     });
 
     // this is here to fix https://github.com/elastic/kibana/issues/215157
@@ -172,16 +169,16 @@ describe('getExpressionType', () => {
                 name: 'fieldName',
                 type: 'unsupported',
                 hasConflict: true,
+                userDefined: false,
               },
             ],
-          ]),
-          new Map()
+          ])
         )
       ).toBe('unknown');
     });
 
     it('handles fields defined by a named param', () => {
-      expect(getExpressionType(getASTForExpression('??field'), new Map(), new Map())).toBe('param');
+      expect(getExpressionType(getASTForExpression('??field'), new Map())).toBe('param');
     });
   });
 
@@ -279,17 +276,23 @@ describe('getExpressionType', () => {
       expect(
         getExpressionType(
           getASTForExpression('CASE(true, "", true, "", keywordField)'),
-          new Map([[`keywordField`, { name: 'keywordField', type: 'keyword' }]]),
-          new Map()
+          new Map([[`keywordField`, { name: 'keywordField', type: 'keyword', userDefined: false }]])
         )
       ).toBe('keyword');
 
       expect(
         getExpressionType(
           getASTForExpression('CASE(true, "", true, "", keywordVar)'),
-          new Map(),
-          new Map([
-            [`keywordVar`, [{ name: 'keywordVar', type: 'keyword', location: { min: 0, max: 0 } }]],
+          new Map<string, ESQLColumnData>([
+            [
+              `keywordVar`,
+              {
+                name: 'keywordVar',
+                type: 'keyword',
+                location: { min: 0, max: 0 },
+                userDefined: true,
+              },
+            ],
           ])
         )
       ).toBe('keyword');
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/operators.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/operators.ts
index c857cf1bc6972..fededccf342f1 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/operators.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/operators.ts
@@ -192,7 +192,7 @@ export async function getSuggestionsToRightOfOperatorExpression({
             getColumnsByType,
             {
               functions: true,
-              fields: true,
+              columns: true,
               values: Boolean(operator.subtype === 'binary-expression'),
             },
             {},
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/commands.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/commands.ts
index 514c98158d5c2..4c104cba12bed 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/commands.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/commands.ts
@@ -18,8 +18,7 @@ export const validateCommandArguments = (
   command: ESQLCommand,
   ast: ESQLAst,
   context: ICommandContext = {
-    userDefinedColumns: new Map(), // Ensure context is always defined
-    fields: new Map(),
+    columns: new Map(), // Ensure context is always defined
   },
   callbacks: ICommandCallbacks = {}
 ) => {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
index 0904f03358853..0e4edac4929c5 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
@@ -74,9 +74,7 @@ class FunctionValidator {
   ) {
     this.definition = getFunctionDefinition(fn.name);
     for (const arg of this.fn.args) {
-      this.argTypes.push(
-        getExpressionType(arg, this.context.fields, this.context.userDefinedColumns)
-      );
+      this.argTypes.push(getExpressionType(arg, this.context.columns));
       this.argLiteralsMask.push(isLiteral(arg));
     }
   }
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
deleted file mode 100644
index c6414c158df32..0000000000000
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts
+++ /dev/null
@@ -1,114 +0,0 @@
-/*
- * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-
-import type { ESQLAst, ESQLAstItem, ESQLCommand } from '@kbn/esql-ast';
-import { isColumn, isFunctionExpression } from '@kbn/esql-ast';
-import { Visitor } from '@kbn/esql-ast/src/visitor';
-import type {
-  ESQLUserDefinedColumn,
-  ESQLFieldWithMetadata,
-} from '@kbn/esql-ast/src/commands_registry/types';
-
-function addToUserDefinedColumnOccurrences(
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>,
-  instance: ESQLUserDefinedColumn
-) {
-  if (!userDefinedColumns.has(instance.name)) {
-    userDefinedColumns.set(instance.name, []);
-  }
-  const userDefinedColumnsOccurrencies = userDefinedColumns.get(instance.name)!;
-  userDefinedColumnsOccurrencies.push(instance);
-}
-
-function addToUserDefinedColumns(
-  oldArg: ESQLAstItem,
-  newArg: ESQLAstItem,
-  fields: Map<string, ESQLFieldWithMetadata>,
-  userDefinedColumns: Map<string, ESQLUserDefinedColumn[]>
-) {
-  if (isColumn(oldArg) && isColumn(newArg)) {
-    const newUserDefinedColumn: ESQLUserDefinedColumn = {
-      name: newArg.parts.join('.'),
-      type: 'double' /* fallback to number */,
-      location: newArg.location,
-    };
-    // Now workout the exact type
-    // it can be a rename of another userDefinedColumn as well
-    const oldRef =
-      fields.get(oldArg.parts.join('.')) || userDefinedColumns.get(oldArg.parts.join('.'));
-    if (oldRef) {
-      addToUserDefinedColumnOccurrences(userDefinedColumns, newUserDefinedColumn);
-      newUserDefinedColumn.type = Array.isArray(oldRef) ? oldRef[0].type : oldRef.type;
-    }
-  }
-}
-
-export function excludeUserDefinedColumnsFromCurrentCommand(
-  commands: ESQLCommand[],
-  currentCommand: ESQLCommand,
-  fieldsMap: Map<string, ESQLFieldWithMetadata>,
-  queryString: string
-) {
-  const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, queryString);
-  const currentCommandUserDefinedColumns = collectUserDefinedColumns(
-    [currentCommand],
-    fieldsMap,
-    queryString
-  );
-  const resultUserDefinedColumns = new Map<string, ESQLUserDefinedColumn[]>();
-  anyUserDefinedColumns.forEach((value, key) => {
-    if (!currentCommandUserDefinedColumns.has(key)) {
-      resultUserDefinedColumns.set(key, value);
-    }
-  });
-  return resultUserDefinedColumns;
-}
-
-export function collectUserDefinedColumns(
-  ast: ESQLAst,
-  fields: Map<string, ESQLFieldWithMetadata>,
-  queryString: string
-): Map<string, ESQLUserDefinedColumn[]> {
-  const userDefinedColumns = new Map<string, ESQLUserDefinedColumn[]>();
-
-  const visitor = new Visitor()
-    .on('visitLiteralExpression', (ctx) => {
-      // TODO - add these as userDefinedColumns
-    })
-    .on('visitExpression', (_ctx) => {}) // required for the types :shrug:
-    .on('visitCommandOption', (ctx) => {
-      if (ctx.node.name === 'with') {
-        for (const assignFn of ctx.node.args) {
-          if (isFunctionExpression(assignFn)) {
-            const [newArg, oldArg] = assignFn?.args || [];
-            // TODO why is oldArg an array?
-            if (Array.isArray(oldArg)) {
-              addToUserDefinedColumns(oldArg[0], newArg, fields, userDefinedColumns);
-            }
-          }
-        }
-      }
-    })
-    .on('visitCommand', (ctx) => {
-      const ret = [];
-      if (['enrich'].includes(ctx.node.name)) {
-        // BY and WITH can contain userDefinedColumns
-        ret.push(...ctx.visitOptions());
-      }
-      if (ctx.node.name === 'fork') {
-        ret.push(...ctx.visitSubQueries());
-      }
-      return ret;
-    })
-    .on('visitQuery', (ctx) => [...ctx.visitCommands()]);
-
-  visitor.visitQuery(ast);
-
-  return userDefinedColumns;
-}

From f868d8bdb55d27f1be48255381a8d2ea080f57f7 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 12:56:24 -0600
Subject: [PATCH 11/54] fix change_point

---
 .../change_point/columns_after.test.ts        | 50 ++++++++++---------
 .../commands/change_point/validate.test.ts    |  6 +--
 2 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
index beb59f176b7ea..0da11ee82a699 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
@@ -7,22 +7,21 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('CHANGE_POINT > columnsAfter', () => {
   const context = {
-    userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
-    fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword' }],
-      ['count', { name: 'count', type: 'double' }],
+    columns: new Map<string, ESQLColumnData>([
+      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+      ['count', { name: 'count', type: 'double', userDefined: false }],
     ]),
   };
   it('adds "type" and "pvalue" fields, when AS option not specified', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'count', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'count', type: 'double', userDefined: false },
+    ];
 
     const result = columnsAfter(
       synth.cmd`CHANGE_POINT count ON field1`,
@@ -30,19 +29,19 @@ describe('CHANGE_POINT > columnsAfter', () => {
       context
     );
 
-    expect(result).toEqual([
-      { name: 'field1', type: 'keyword' },
-      { name: 'count', type: 'double' },
-      { name: 'type', type: 'keyword' },
-      { name: 'pvalue', type: 'double' },
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'count', type: 'double', userDefined: false },
+      { name: 'type', type: 'keyword', userDefined: false },
+      { name: 'pvalue', type: 'double', userDefined: false },
     ]);
   });
 
   it('adds the given names as fields, when AS option is specified', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'count', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'count', type: 'double', userDefined: false },
+    ];
 
     const result = columnsAfter(
       synth.cmd`CHANGE_POINT count ON field1 AS changePointType, pValue`,
@@ -50,11 +49,16 @@ describe('CHANGE_POINT > columnsAfter', () => {
       context
     );
 
-    expect(result).toEqual([
-      { name: 'field1', type: 'keyword' },
-      { name: 'count', type: 'double' },
-      { name: 'changePointType', type: 'keyword' },
-      { name: 'pValue', type: 'double' },
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'count', type: 'double', userDefined: false },
+      {
+        name: 'changePointType',
+        type: 'keyword',
+        userDefined: true,
+        location: { min: 0, max: 0 }, // TODO — location is being assigned incorrectly in the AST
+      },
+      { name: 'pValue', type: 'double', userDefined: true, location: { min: 0, max: 0 } },
     ]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.test.ts
index af5d8857495a0..fecfc2760b627 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.test.ts
@@ -50,11 +50,11 @@ describe('CHANGE_POINT Validation', () => {
 
       test('raises error when the default @timestamp field is missing', () => {
         // make sure that @timestamp field is not present
-        const newFields = new Map(mockContext.fields);
-        newFields.delete('@timestamp');
+        const newColumns = new Map(mockContext.columns);
+        newColumns.delete('@timestamp');
         const context = {
           ...mockContext,
-          fields: newFields,
+          columns: newColumns,
         };
         changePointExpectErrors(
           'FROM a_index | CHANGE_POINT doubleField',

From e7cc9897bcc797af54e398c51e6a156550e09b8b Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 12:59:36 -0600
Subject: [PATCH 12/54] fix completion

---
 .../commands/completion/columns_after.test.ts | 41 +++++++++----------
 1 file changed, 20 insertions(+), 21 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
index 4020af765f063..f7fd48ef58354 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
@@ -7,52 +7,51 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('COMPLETION', () => {
   const context = {
-    userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
-    fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword' }],
-      ['count', { name: 'count', type: 'double' }],
+    columns: new Map<string, ESQLColumnData>([
+      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+      ['count', { name: 'count', type: 'double', userDefined: false }],
     ]),
   };
   it('adds "completion" field, when targetField is not specified', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'count', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandColumns: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'count', type: 'double', userDefined: false },
+    ];
 
     const result = columnsAfter(
       synth.cmd`COMPLETION "prompt" WITH {"inference_id": "my-inference-id"}`,
-      previousCommandFields,
+      previousCommandColumns,
       context
     );
 
     expect(result).toEqual([
-      { name: 'field1', type: 'keyword' },
-      { name: 'count', type: 'double' },
-      { name: 'completion', type: 'keyword' },
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'count', type: 'double', userDefined: false },
+      { name: 'completion', type: 'keyword', userDefined: false },
     ]);
   });
 
   it('adds the given targetField as field, when targetField is specified', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'count', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandColumns: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'count', type: 'double', userDefined: false },
+    ];
 
     const result = columnsAfter(
       synth.cmd`COMPLETION customField = "prompt" WITH {"inference_id": "my-inference-id"}`,
-      previousCommandFields,
+      previousCommandColumns,
       context
     );
 
     expect(result).toEqual([
-      { name: 'field1', type: 'keyword' },
-      { name: 'count', type: 'double' },
-      { name: 'customField', type: 'keyword' },
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'count', type: 'double', userDefined: false },
+      { name: 'customField', type: 'keyword', userDefined: true, location: { min: 0, max: 0 } }, // TODO location is being computed incorrectly in AST
     ]);
   });
 });

From 3401a13486c390c62018bbdf6f63d850c95bb203 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 13:01:15 -0600
Subject: [PATCH 13/54] fix dissect

---
 .../commands/dissect/columns_after.test.ts    | 25 +++++++++----------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
index c154efe00f4bf..ed76d849a4660 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import { columnsAfter, extractDissectColumnNames } from './columns_after';
 
 describe('DISSECT', () => {
@@ -69,17 +69,16 @@ describe('DISSECT', () => {
   });
   describe('columnsAfter', () => {
     const context = {
-      userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
-      fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword' }],
-        ['count', { name: 'count', type: 'double' }],
+      columns: new Map<string, ESQLColumnData>([
+        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+        ['count', { name: 'count', type: 'double', userDefined: false }],
       ]),
     };
     it('adds the DISSECT pattern columns as fields', () => {
-      const previousColumns = [
-        { name: 'field1', type: 'keyword' },
-        { name: 'field2', type: 'double' },
-      ] as ESQLFieldWithMetadata[];
+      const previousColumns: ESQLColumnData[] = [
+        { name: 'field1', type: 'keyword', userDefined: false },
+        { name: 'field2', type: 'double', userDefined: false },
+      ];
 
       const result = columnsAfter(
         synth.cmd`DISSECT agent "%{firstWord}"`,
@@ -87,10 +86,10 @@ describe('DISSECT', () => {
         context
       );
 
-      expect(result).toEqual([
-        { name: 'field1', type: 'keyword' },
-        { name: 'field2', type: 'double' },
-        { name: 'firstWord', type: 'keyword' },
+      expect(result).toEqual<ESQLColumnData[]>([
+        { name: 'field1', type: 'keyword', userDefined: false },
+        { name: 'field2', type: 'double', userDefined: false },
+        { name: 'firstWord', type: 'keyword', userDefined: false },
       ]);
     });
   });

From 4ddff6198851f3966fe87381b2f9975f89acc41b Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 13:03:49 -0600
Subject: [PATCH 14/54] fix drop

---
 .../commands/drop/columns_after.test.ts       | 19 ++++++------
 .../commands/drop/validate.test.ts            | 30 +++++++++----------
 2 files changed, 23 insertions(+), 26 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
index d0428da653a62..77a758516ebcc 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
@@ -7,25 +7,24 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('DROP', () => {
   const context = {
-    userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
-    fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword' }],
-      ['count', { name: 'count', type: 'double' }],
+    columns: new Map<string, ESQLColumnData>([
+      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+      ['count', { name: 'count', type: 'double', userDefined: false }],
     ]),
   };
   it('removes the columns defined in the command', () => {
-    const previousColumns = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+    ];
 
     const result = columnsAfter(synth.cmd`DROP field1`, previousColumns, context);
 
-    expect(result).toEqual([{ name: 'field2', type: 'double' }]);
+    expect(result).toEqual([{ name: 'field2', type: 'double', userDefined: false }]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/validate.test.ts
index 34e3d6fa91cc1..bab0ef07c51bd 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/validate.test.ts
@@ -21,24 +21,22 @@ describe('DROP Validation', () => {
 
   test('validates the most basic query', () => {
     // make sure that @timestamp field is not present
-    const newUserDefinedColumns = new Map(mockContext.userDefinedColumns);
-    newUserDefinedColumns.set('MIN(doubleField * 10)', [
-      {
-        name: 'MIN(doubleField * 10)',
-        type: 'double',
-        location: { min: 0, max: 10 },
-      },
-    ]);
-    newUserDefinedColumns.set('COUNT(*)', [
-      {
-        name: 'COUNT(*)',
-        type: 'integer',
-        location: { min: 0, max: 10 },
-      },
-    ]);
+    const newColumns = new Map(mockContext.columns);
+    newColumns.set('MIN(doubleField * 10)', {
+      name: 'MIN(doubleField * 10)',
+      type: 'double',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
+    newColumns.set('COUNT(*)', {
+      name: 'COUNT(*)',
+      type: 'integer',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
     const context = {
       ...mockContext,
-      userDefinedColumns: newUserDefinedColumns,
+      columns: newColumns,
     };
     dropExpectErrors('from index | drop textField, doubleField, dateField', []);
     dropExpectErrors('from index | drop `any#Char$Field`', []);

From 162ff70bd986754c5d41722fd2987ab81205a78a Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 13:38:50 -0600
Subject: [PATCH 15/54] make eval work

---
 .../change_point/columns_after.test.ts        |   2 +-
 .../commands/completion/columns_after.test.ts |   2 +-
 .../commands/eval/columns_after.test.ts       | 116 ++++++++++++++++++
 .../commands/eval/columns_after.ts            |   3 +-
 .../src/commands_registry/registry.ts         |   1 +
 .../definitions/utils/autocomplete/columns.ts |  30 +----
 .../utils/autocomplete/functions.ts           |   7 --
 .../kbn-esql-validation-autocomplete/index.ts |   2 -
 .../src/shared/helpers.ts                     |  54 +-------
 .../src/languages/esql/lib/hover/hover.ts     |   4 +-
 10 files changed, 129 insertions(+), 92 deletions(-)
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
index 0da11ee82a699..568371d494220 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
@@ -56,7 +56,7 @@ describe('CHANGE_POINT > columnsAfter', () => {
         name: 'changePointType',
         type: 'keyword',
         userDefined: true,
-        location: { min: 0, max: 0 }, // TODO — location is being assigned incorrectly in the AST
+        location: { min: 0, max: 0 },
       },
       { name: 'pValue', type: 'double', userDefined: true, location: { min: 0, max: 0 } },
     ]);
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
index f7fd48ef58354..627c9466847bb 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
@@ -51,7 +51,7 @@ describe('COMPLETION', () => {
     expect(result).toEqual([
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'count', type: 'double', userDefined: false },
-      { name: 'customField', type: 'keyword', userDefined: true, location: { min: 0, max: 0 } }, // TODO location is being computed incorrectly in AST
+      { name: 'customField', type: 'keyword', userDefined: true, location: { min: 0, max: 0 } },
     ]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts
new file mode 100644
index 0000000000000..4aa4edf901ca0
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts
@@ -0,0 +1,116 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { Parser, synth } from '../../../..';
+import type { ESQLColumnData } from '../../types';
+import { columnsAfter } from './columns_after';
+
+describe('EVAL > columnsAfter', () => {
+  const baseColumns: ESQLColumnData[] = [
+    { name: 'foo', type: 'integer', userDefined: false },
+    { name: 'bar', type: 'keyword', userDefined: false },
+  ];
+
+  it('adds a new column for a simple assignment', () => {
+    const command = synth.cmd`EVAL baz = foo + 1`;
+    const result = columnsAfter(command, baseColumns, '');
+
+    expect(result).toEqual([
+      ...baseColumns,
+      {
+        name: 'baz',
+        type: 'integer',
+        location: { min: 0, max: 0 },
+        userDefined: true,
+      },
+    ]);
+  });
+
+  it('adds multiple new columns for multiple assignments', () => {
+    const command = synth.cmd`EVAL baz = foo + 1, qux = bar`;
+    const result = columnsAfter(command, baseColumns, '');
+
+    expect(result).toEqual([
+      ...baseColumns,
+      {
+        name: 'baz',
+        type: 'integer',
+        location: { min: 0, max: 0 },
+        userDefined: true,
+      },
+      {
+        name: 'qux',
+        type: 'keyword',
+        location: { min: 0, max: 0 },
+        userDefined: true,
+      },
+    ]);
+  });
+
+  it('adds a column for a single expression (not assignment)', () => {
+    const queryString = `FROM index | EVAL foo + 1`;
+
+    // Can't use synth because it steps on the location information
+    // which is used to determine the name of the new column
+    const {
+      root: {
+        commands: [, command],
+      },
+    } = Parser.parseQuery(queryString);
+
+    const result = columnsAfter(command, baseColumns, queryString);
+
+    expect(result).toEqual([
+      ...baseColumns,
+      {
+        name: 'foo + 1',
+        type: 'integer',
+        location: { min: 18, max: 24 },
+        userDefined: true,
+      },
+    ]);
+  });
+
+  it('handles mix of assignments and expressions', () => {
+    const queryString = `FROM index | EVAL baz = foo + 1, TRIM(bar)`;
+
+    // Can't use synth because it steps on the location information
+    // which is used to determine the name of the new column
+    const {
+      root: {
+        commands: [, command],
+      },
+    } = Parser.parseQuery(queryString);
+
+    const result = columnsAfter(command, baseColumns, queryString);
+
+    expect(result).toEqual([
+      ...baseColumns,
+      {
+        name: 'baz',
+        type: 'integer',
+        location: { min: 18, max: 20 },
+        userDefined: true,
+      },
+      {
+        name: 'TRIM(bar)',
+        type: 'keyword',
+        location: { min: 33, max: 41 },
+        userDefined: true,
+      },
+    ]);
+  });
+
+  it('returns previous columns if no args', () => {
+    const command = { args: [] } as any;
+    const result = columnsAfter(command, baseColumns, '');
+
+    expect(result).toEqual(baseColumns);
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
index 1585c075c194c..13aa312ccf811 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
@@ -15,6 +15,7 @@ import type { ESQLColumnData, ESQLUserDefinedColumn, ICommandContext } from '../
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const columnMap = new Map<string, ESQLColumnData>();
@@ -36,7 +37,7 @@ export const columnsAfter = (
       newColumns.push(newColumn);
     } else if (!Array.isArray(expression)) {
       const newColumn: ESQLUserDefinedColumn = {
-        name: expression.text,
+        name: query.substring(expression.location.min, expression.location.max + 1),
         type: typeOf(expression),
         location: expression.location,
         userDefined: true,
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
index 9fbf19f0c604c..cdcf71961f13a 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
@@ -60,6 +60,7 @@ export interface ICommandMethods<TContext = any> {
   columnsAfter?: (
     command: ESQLCommand,
     previousColumns: ESQLColumnData[],
+    query: string,
     context?: TContext
   ) => ESQLColumnData[];
 }
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
index 3966ee1eee529..6f6c1972c0204 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts
@@ -6,36 +6,10 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import type {
-  ESQLFieldWithMetadata,
-  ESQLUserDefinedColumn,
-  ICommandContext,
-} from '../../../commands_registry/types';
-import type { ESQLColumn, ESQLCommand, ESQLIdentifier } from '../../../types';
+import type { ICommandContext } from '../../../commands_registry/types';
+import type { ESQLColumn, ESQLIdentifier } from '../../../types';
 import { fuzzySearch } from '../shared';
 
-// FIXME
-export function excludeUserDefinedColumnsFromCurrentCommand(
-  commands: ESQLCommand[],
-  currentCommand: ESQLCommand,
-  fieldsMap: Map<string, ESQLFieldWithMetadata>,
-  queryString: string
-) {
-  const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, queryString);
-  const currentCommandUserDefinedColumns = collectUserDefinedColumns(
-    [currentCommand],
-    fieldsMap,
-    queryString
-  );
-  const resultUserDefinedColumns = new Map<string, ESQLUserDefinedColumn[]>();
-  anyUserDefinedColumns.forEach((value, key) => {
-    if (!currentCommandUserDefinedColumns.has(key)) {
-      resultUserDefinedColumns.set(key, value);
-    }
-  });
-  return resultUserDefinedColumns;
-}
-
 /**
  * TODO - consider calling lookupColumn under the hood of this function. Seems like they should really do the same thing.
  */
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
index cca9bb69ef75f..06f97b81dd942 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
@@ -51,7 +51,6 @@ import {
 import { getCompatibleLiterals, getDateLiterals } from '../literals';
 import { getSuggestionsToRightOfOperatorExpression } from '../operators';
 import { buildValueDefinitions } from '../values';
-import { excludeUserDefinedColumnsFromCurrentCommand } from './columns';
 import {
   extractTypeFromASTArg,
   getFieldsOrFunctionsSuggestions,
@@ -172,12 +171,6 @@ export async function getFunctionArgsSuggestions(
   const references = {
     columns: columnMap,
   };
-  const userDefinedColumnsExcludingCurrentCommandOnes = excludeUserDefinedColumnsFromCurrentCommand(
-    commands,
-    command,
-    columnMap,
-    innerText
-  );
 
   const { typesToSuggestNext, hasMoreMandatoryArgs, enrichedArgs, argIndex } =
     getValidSignaturesAndTypesToSuggestNext(
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts
index 9459f53edd5fe..973717d840d86 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts
@@ -21,8 +21,6 @@ export { suggest } from './src/autocomplete/autocomplete';
  * Some utility functions that can be useful to build more feature
  * for the ES|QL language
  */
-export { collectUserDefinedColumns } from './src/shared/user_defined_columns';
-
 export {
   getColumnsByTypeHelper as getFieldsByTypeHelper,
   getPolicyHelper,
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index 72dddc1abfdca..edefc49925963 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -6,23 +6,12 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import {
-  type ESQLAstCommand,
-  esqlCommandRegistry,
-  type FieldType,
-  type FunctionDefinition,
-} from '@kbn/esql-ast';
-import type {
-  ESQLColumnData,
-  ESQLFieldWithMetadata,
-  ESQLUserDefinedColumn,
-} from '@kbn/esql-ast/src/commands_registry/types';
+import { esqlCommandRegistry, type ESQLAstCommand, type FunctionDefinition } from '@kbn/esql-ast';
+import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types';
 import type { ESQLParamLiteral } from '@kbn/esql-ast/src/types';
-import { uniqBy } from 'lodash';
 
 import { enrichFieldsWithECSInfo } from '../autocomplete/utils/ecs_metadata_helper';
 import type { ESQLCallbacks } from './types';
-import { collectUserDefinedColumns } from './user_defined_columns';
 
 export function nonNullable<T>(v: T): v is NonNullable<T> {
   return v != null;
@@ -78,28 +67,6 @@ export function getParamAtPosition(
   return params.length > position ? params[position] : minParams ? params[params.length - 1] : null;
 }
 
-// --- Fields helpers ---
-
-export function transformMapToESQLFields(
-  inputMap: Map<string, ESQLUserDefinedColumn[]>
-): ESQLFieldWithMetadata[] {
-  const esqlFields: ESQLFieldWithMetadata[] = [];
-
-  for (const [, userDefinedColumns] of inputMap) {
-    for (const userDefinedColumn of userDefinedColumns) {
-      // Only include userDefinedColumns that have a known type
-      if (userDefinedColumn.type) {
-        esqlFields.push({
-          name: userDefinedColumn.name,
-          type: userDefinedColumn.type as FieldType,
-        });
-      }
-    }
-  }
-
-  return esqlFields;
-}
-
 async function getEcsMetadata(resourceRetriever?: ESQLCallbacks) {
   if (!resourceRetriever?.getFieldsMetadata) {
     return undefined;
@@ -136,22 +103,11 @@ export async function getCurrentQueryAvailableColumns(
   const lastCommand = commands[commands.length - 1];
   const commandDefinition = esqlCommandRegistry.getCommandByName(lastCommand.name);
 
-  // @TODO — all logic in collectUserDefinedColumns
-  // should be will be moved to columnsAfter methods;
-  // though it may still be useful to delineate between
-  // user-defined columns and other fields... need to consider this
-  // If the command has a columnsAfter function, use it to get the fields
   if (commandDefinition?.methods.columnsAfter) {
-    return commandDefinition.methods.columnsAfter(lastCommand, previousPipeFields, {
-      userDefinedColumns: new Map(),
+    return commandDefinition.methods.columnsAfter(lastCommand, previousPipeFields, query, {
+      columns: new Map(),
     });
   } else {
-    // If the command doesn't have a columnsAfter function, use the default behavior
-    const userDefinedColumns = collectUserDefinedColumns(commands, cacheCopy, query);
-    const arrayOfUserDefinedColumns: ESQLFieldWithMetadata[] = transformMapToESQLFields(
-      userDefinedColumns ?? new Map<string, ESQLUserDefinedColumn[]>()
-    );
-    const allFields = uniqBy([...(previousPipeFields ?? []), ...arrayOfUserDefinedColumns], 'name');
-    return allFields;
+    return previousPipeFields;
   }
 }
diff --git a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
index 4a5803646b828..e134f2db820e1 100644
--- a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
+++ b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
@@ -24,7 +24,7 @@ import {
   type ESQLSingleAstItem,
   type ESQLSource,
 } from '@kbn/esql-ast/src/types';
-import { collectUserDefinedColumns, type ESQLCallbacks } from '@kbn/esql-validation-autocomplete';
+import { type ESQLCallbacks } from '@kbn/esql-validation-autocomplete';
 import { getColumnsByTypeRetriever } from '@kbn/esql-validation-autocomplete/src/autocomplete/autocomplete';
 import { getQueryForFields } from '@kbn/esql-validation-autocomplete/src/autocomplete/helper';
 import { getPolicyHelper } from '@kbn/esql-validation-autocomplete/src/shared/resources_helpers';
@@ -195,11 +195,9 @@ async function getHintForFunctionArg(
     return [];
   }
   const columnsMap = await getColumnMap();
-  const anyUserDefinedColumns = collectUserDefinedColumns(root.commands, columnsMap, query);
 
   const references = {
     columns: columnsMap,
-    userDefinedColumns: anyUserDefinedColumns,
   };
 
   const { typesToSuggestNext, enrichedArgs } = getValidSignaturesAndTypesToSuggestNext(

From 15f51b97aaed104b81b1efe77f0e0938ea9d94d5 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 13:48:13 -0600
Subject: [PATCH 16/54] add query to the signature

---
 .../commands/change_point/columns_after.test.ts   |  2 ++
 .../commands/change_point/columns_after.ts        |  1 +
 .../commands/completion/columns_after.test.ts     |  2 ++
 .../commands/completion/columns_after.ts          |  1 +
 .../commands/dissect/columns_after.test.ts        |  1 +
 .../commands/dissect/columns_after.ts             |  1 +
 .../commands/drop/columns_after.test.ts           |  2 +-
 .../commands/drop/columns_after.ts                |  1 +
 .../commands/fork/columns_after.test.ts           |  1 +
 .../commands/fork/columns_after.ts                |  1 +
 .../commands/grok/columns_after.test.ts           |  1 +
 .../commands/grok/columns_after.ts                |  1 +
 .../commands/keep/columns_after.test.ts           |  2 +-
 .../commands/keep/columns_after.ts                |  1 +
 .../commands/rename/columns_after.test.ts         | 15 +++++++++++++--
 .../commands/rename/columns_after.ts              |  1 +
 .../commands/stats/columns_after.test.ts          | 11 +++++++++--
 .../commands/stats/columns_after.ts               |  5 +++--
 18 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
index 568371d494220..aef7bcbba1560 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
@@ -26,6 +26,7 @@ describe('CHANGE_POINT > columnsAfter', () => {
     const result = columnsAfter(
       synth.cmd`CHANGE_POINT count ON field1`,
       previousCommandFields,
+      '',
       context
     );
 
@@ -46,6 +47,7 @@ describe('CHANGE_POINT > columnsAfter', () => {
     const result = columnsAfter(
       synth.cmd`CHANGE_POINT count ON field1 AS changePointType, pValue`,
       previousCommandFields,
+      '',
       context
     );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
index 7b938b60ebf87..cf8dd86602202 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
@@ -14,6 +14,7 @@ import type { ESQLColumnData, ICommandContext } from '../../types';
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const { target } = command as ESQLAstChangePointCommand;
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
index 627c9466847bb..d086c515a8544 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
@@ -26,6 +26,7 @@ describe('COMPLETION', () => {
     const result = columnsAfter(
       synth.cmd`COMPLETION "prompt" WITH {"inference_id": "my-inference-id"}`,
       previousCommandColumns,
+      '',
       context
     );
 
@@ -45,6 +46,7 @@ describe('COMPLETION', () => {
     const result = columnsAfter(
       synth.cmd`COMPLETION customField = "prompt" WITH {"inference_id": "my-inference-id"}`,
       previousCommandColumns,
+      '',
       context
     );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
index 96c98baa09b0a..c57dbe0ffce5e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
@@ -19,6 +19,7 @@ import type {
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const { targetField } = command as ESQLAstCompletionCommand;
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
index ed76d849a4660..fe8f646f61ff3 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
@@ -83,6 +83,7 @@ describe('DISSECT', () => {
       const result = columnsAfter(
         synth.cmd`DISSECT agent "%{firstWord}"`,
         previousColumns,
+        '',
         context
       );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
index e7e4fd180ca3a..18fdeaa09d09c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
@@ -36,6 +36,7 @@ export function extractDissectColumnNames(pattern: string): string[] {
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const columns: string[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
index 77a758516ebcc..b977ca03e7537 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
@@ -23,7 +23,7 @@ describe('DROP', () => {
       { name: 'field2', type: 'double', userDefined: false },
     ];
 
-    const result = columnsAfter(synth.cmd`DROP field1`, previousColumns, context);
+    const result = columnsAfter(synth.cmd`DROP field1`, previousColumns, '', context);
 
     expect(result).toEqual([{ name: 'field2', type: 'double', userDefined: false }]);
   });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
index 11ff6dee5db15..b17c457435ebd 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
@@ -14,6 +14,7 @@ import type { ICommandContext } from '../../types';
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const columnsToDrop: string[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
index 850321b897fa1..483af1df09c7e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
@@ -27,6 +27,7 @@ describe('FORK', () => {
     const result = columnsAfter(
       synth.cmd`FORK (LIMIT 10 ) (LIMIT 1000 ) `,
       previousCommandFields,
+      '',
       context
     );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index 7cbc681b9eee0..053dbcb064606 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -14,6 +14,7 @@ import type { ICommandContext } from '../../types';
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   return uniqBy(
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
index ffa3b48d1f53d..d07060fc023f3 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
@@ -56,6 +56,7 @@ describe('GROK', () => {
       const result = columnsAfter(
         synth.cmd`GROK agent "%{WORD:firstWord}"`,
         previousCommandFields,
+        '',
         context
       );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
index d58af7985f99d..f400521d26824 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
@@ -51,6 +51,7 @@ export function extractSemanticsFromGrok(pattern: string): string[] {
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const columns: string[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
index b287b9a3c7541..3af346c585b19 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
@@ -24,7 +24,7 @@ describe('KEEP', () => {
       { name: 'field2', type: 'double', userDefined: false },
     ] as ESQLFieldWithMetadata[];
 
-    const result = columnsAfter(synth.cmd`KEEP field1`, previousCommandFields, context);
+    const result = columnsAfter(synth.cmd`KEEP field1`, previousCommandFields, '', context);
 
     expect(result).toEqual([{ name: 'field1', type: 'keyword' }]);
   });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
index 6e6abc12c046a..cf8bac3dbf7cc 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
@@ -14,6 +14,7 @@ import type { ICommandContext } from '../../types';
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const columnsToKeep: string[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
index 5a6977b13b595..131462fb31378 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
@@ -24,7 +24,12 @@ describe('RENAME', () => {
       { name: 'field2', type: 'double', userDefined: false },
     ] as ESQLFieldWithMetadata[];
 
-    const result = columnsAfter(synth.cmd`RENAME field1 as meow`, previousCommandFields, context);
+    const result = columnsAfter(
+      synth.cmd`RENAME field1 as meow`,
+      previousCommandFields,
+      '',
+      context
+    );
 
     expect(result).toEqual([
       { name: 'meow', type: 'keyword' },
@@ -38,7 +43,12 @@ describe('RENAME', () => {
       { name: 'field2', type: 'double' },
     ] as ESQLFieldWithMetadata[];
 
-    const result = columnsAfter(synth.cmd`RENAME meow = field1`, previousCommandFields, context);
+    const result = columnsAfter(
+      synth.cmd`RENAME meow = field1`,
+      previousCommandFields,
+      '',
+      context
+    );
 
     expect(result).toEqual([
       { name: 'meow', type: 'keyword', userDefined: false },
@@ -55,6 +65,7 @@ describe('RENAME', () => {
     const result = columnsAfter(
       synth.cmd`RENAME meow = field1, field2 as woof`,
       previousCommandFields,
+      '',
       context
     );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
index 07a4545835fa5..f569faa0ff6b7 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
@@ -14,6 +14,7 @@ import type { ESQLColumnData, ICommandContext } from '../../types';
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const asRenamePairs: ESQLFunction[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
index 46a26ceb754d9..005b2c1e35522 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
@@ -36,7 +36,12 @@ describe('STATS', () => {
       ]),
     };
 
-    const result = columnsAfter(synth.cmd`STATS var0=AVG(field2)`, previousCommandFields, context);
+    const result = columnsAfter(
+      synth.cmd`STATS var0=AVG(field2)`,
+      previousCommandFields,
+      '',
+      context
+    );
 
     expect(result).toEqual([{ name: 'var0', type: 'double', userDefined: false }]);
   });
@@ -67,7 +72,7 @@ describe('STATS', () => {
       ]),
     };
 
-    const result = columnsAfter(synth.cmd`STATS AVG(field2)`, previousCommandFields, context);
+    const result = columnsAfter(synth.cmd`STATS AVG(field2)`, previousCommandFields, '', context);
 
     expect(result).toEqual([{ name: 'AVG(field2)', type: 'double', userDefined: false }]);
   });
@@ -101,6 +106,7 @@ describe('STATS', () => {
     const result = columnsAfter(
       synth.cmd`STATS AVG(field2) BY field1`,
       previousCommandFields,
+      '',
       context
     );
 
@@ -151,6 +157,7 @@ describe('STATS', () => {
     const result = columnsAfter(
       synth.cmd`STATS AVG(field2) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`,
       previousCommandFields,
+      '',
       context
     );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
index 0fc1502f82084..eb0a5b7c5b478 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
@@ -53,8 +53,9 @@ const getUserDefinedColumns = (
 };
 
 export const columnsAfter = (
-  _command: ESQLCommand,
+  command: ESQLCommand,
   previousColumns: ESQLColumnData[],
+  query: string,
   context?: ICommandContext
 ) => {
   const columnMap = new Map<string, ESQLColumnData>();
@@ -63,5 +64,5 @@ export const columnsAfter = (
   const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, columnMap);
 
   // TODO - is this uniqby helpful? Does it do what we expect?
-  return uniqBy([...previousColumns, ...getUserDefinedColumns(_command, typeOf)], 'name');
+  return uniqBy([...previousColumns, ...getUserDefinedColumns(command, typeOf)], 'name');
 };

From 500079855dab393f81b6497ba3293ff99d18a835 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 14:02:03 -0600
Subject: [PATCH 17/54] fix rename

---
 .../commands/rename/columns_after.test.ts     | 50 +++++++++----------
 .../commands/rename/columns_after.ts          | 22 +++++---
 .../commands/rename/validate.test.ts          | 30 ++++++-----
 3 files changed, 54 insertions(+), 48 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
index 131462fb31378..05a1e1249a4ea 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
@@ -7,22 +7,20 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('RENAME', () => {
-  const context = {
-    userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
-    fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-      ['count', { name: 'count', type: 'double', userDefined: false }],
-    ]),
-  };
+  const columns = new Map<string, ESQLColumnData>([
+    ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+    ['count', { name: 'count', type: 'double', userDefined: false }],
+  ]);
+  const context = { columns };
   it('renames the given columns with the new names using AS', () => {
-    const previousCommandFields = [
+    const previousCommandFields: ESQLColumnData[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
-    ] as ESQLFieldWithMetadata[];
+    ];
 
     const result = columnsAfter(
       synth.cmd`RENAME field1 as meow`,
@@ -31,17 +29,17 @@ describe('RENAME', () => {
       context
     );
 
-    expect(result).toEqual([
-      { name: 'meow', type: 'keyword' },
-      { name: 'field2', type: 'double' },
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'meow', type: 'keyword', userDefined: true, location: { max: 0, min: 0 } },
+      { name: 'field2', type: 'double', userDefined: false },
     ]);
   });
 
   it('renames the given columns with the new names using ASSIGN', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+    ];
 
     const result = columnsAfter(
       synth.cmd`RENAME meow = field1`,
@@ -50,17 +48,17 @@ describe('RENAME', () => {
       context
     );
 
-    expect(result).toEqual([
-      { name: 'meow', type: 'keyword', userDefined: false },
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'meow', type: 'keyword', userDefined: true, location: { max: 0, min: 0 } },
       { name: 'field2', type: 'double', userDefined: false },
     ]);
   });
 
   it('renames the given columns with the new names using a mix of ASSIGN and =', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+    ];
 
     const result = columnsAfter(
       synth.cmd`RENAME meow = field1, field2 as woof`,
@@ -69,9 +67,9 @@ describe('RENAME', () => {
       context
     );
 
-    expect(result).toEqual([
-      { name: 'meow', type: 'keyword', userDefined: false },
-      { name: 'woof', type: 'double', userDefined: false },
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'meow', type: 'keyword', userDefined: true, location: { max: 0, min: 0 } },
+      { name: 'woof', type: 'double', userDefined: true, location: { max: 0, min: 0 } },
     ]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
index f569faa0ff6b7..6f867858dcd8d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import uniqBy from 'lodash/uniqBy';
-import { isFunctionExpression } from '../../../ast/is';
+import { isColumn, isFunctionExpression } from '../../../ast/is';
 import type { ESQLAstBaseItem, ESQLCommand, ESQLFunction } from '../../../types';
 import type { ESQLColumnData, ICommandContext } from '../../types';
 
@@ -31,21 +31,31 @@ export const columnsAfter = (
   }
 
   // rename the columns with the user defined name
-  const newFields = previousColumns.map((oldColumn) => {
+  const newFields = previousColumns.map<ESQLColumnData>((oldColumn) => {
     const asRenamePair = asRenamePairs.find(
       (pair) => (pair?.args?.[0] as ESQLAstBaseItem)?.name === oldColumn.name
     );
 
-    if (asRenamePair?.args?.[1]) {
-      return { name: (asRenamePair.args[1] as ESQLAstBaseItem).name, type: oldColumn.type };
+    if (isColumn(asRenamePair?.args?.[1])) {
+      return {
+        name: asRenamePair.args[1].name,
+        type: oldColumn.type,
+        location: asRenamePair.args[1].location,
+        userDefined: true,
+      };
     }
 
     const assignRenamePair = assignRenamePairs.find(
       (pair) => (pair?.args?.[1] as ESQLAstBaseItem)?.name === oldColumn.name
     );
 
-    if (assignRenamePair?.args?.[0]) {
-      return { name: (assignRenamePair.args[0] as ESQLAstBaseItem).name, type: oldColumn.type };
+    if (isColumn(assignRenamePair?.args?.[0])) {
+      return {
+        name: assignRenamePair.args[0].name,
+        type: oldColumn.type,
+        location: assignRenamePair.args[0].location,
+        userDefined: true,
+      };
     }
 
     return oldColumn; // No rename found, keep the old name
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts
index 4f59f186c9d11..e7df537c67600 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts
@@ -20,24 +20,22 @@ describe('RENAME Validation', () => {
   });
 
   test('validates the most basic query', () => {
-    const newUserDefinedColumns = new Map(mockContext.userDefinedColumns);
-    newUserDefinedColumns.set('doubleField + 1', [
-      {
-        name: 'doubleField + 1',
-        type: 'double',
-        location: { min: 0, max: 10 },
-      },
-    ]);
-    newUserDefinedColumns.set('avg(doubleField)', [
-      {
-        name: 'avg(doubleField)',
-        type: 'double',
-        location: { min: 0, max: 10 },
-      },
-    ]);
+    const newColumns = new Map(mockContext.columns);
+    newColumns.set('doubleField + 1', {
+      name: 'doubleField + 1',
+      type: 'double',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
+    newColumns.set('avg(doubleField)', {
+      name: 'avg(doubleField)',
+      type: 'double',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
     const context = {
       ...mockContext,
-      userDefinedColumns: newUserDefinedColumns,
+      columns: newColumns,
     };
     renameExpectErrors('from a_index | rename textField as', [
       'AS expected 2 arguments, but got 1.',

From 0c0b6fca83e7b563103b49f058ccbfff49871ab0 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 29 Aug 2025 16:05:06 -0600
Subject: [PATCH 18/54] fork plus an eval fix

---
 .../commands/eval/columns_after.test.ts       |  28 +++-
 .../commands/eval/columns_after.ts            |   3 +-
 .../commands/fork/columns_after.test.ts       | 147 ++++++++++++++++--
 .../commands/fork/columns_after.ts            |  70 ++++++++-
 .../commands/keep/columns_after.test.ts       |  14 +-
 5 files changed, 236 insertions(+), 26 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts
index 4aa4edf901ca0..fdeffc0e157e2 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts
@@ -22,13 +22,13 @@ describe('EVAL > columnsAfter', () => {
     const result = columnsAfter(command, baseColumns, '');
 
     expect(result).toEqual([
-      ...baseColumns,
       {
         name: 'baz',
         type: 'integer',
         location: { min: 0, max: 0 },
         userDefined: true,
       },
+      ...baseColumns,
     ]);
   });
 
@@ -37,7 +37,6 @@ describe('EVAL > columnsAfter', () => {
     const result = columnsAfter(command, baseColumns, '');
 
     expect(result).toEqual([
-      ...baseColumns,
       {
         name: 'baz',
         type: 'integer',
@@ -50,6 +49,7 @@ describe('EVAL > columnsAfter', () => {
         location: { min: 0, max: 0 },
         userDefined: true,
       },
+      ...baseColumns,
     ]);
   });
 
@@ -67,13 +67,13 @@ describe('EVAL > columnsAfter', () => {
     const result = columnsAfter(command, baseColumns, queryString);
 
     expect(result).toEqual([
-      ...baseColumns,
       {
         name: 'foo + 1',
         type: 'integer',
         location: { min: 18, max: 24 },
         userDefined: true,
       },
+      ...baseColumns,
     ]);
   });
 
@@ -91,7 +91,6 @@ describe('EVAL > columnsAfter', () => {
     const result = columnsAfter(command, baseColumns, queryString);
 
     expect(result).toEqual([
-      ...baseColumns,
       {
         name: 'baz',
         type: 'integer',
@@ -104,6 +103,7 @@ describe('EVAL > columnsAfter', () => {
         location: { min: 33, max: 41 },
         userDefined: true,
       },
+      ...baseColumns,
     ]);
   });
 
@@ -113,4 +113,24 @@ describe('EVAL > columnsAfter', () => {
 
     expect(result).toEqual(baseColumns);
   });
+
+  it('handles overwriting columns', () => {
+    const command = synth.cmd`EVAL foo = "", bar = 23`;
+    const result = columnsAfter(command, baseColumns, '');
+
+    expect(result).toEqual([
+      {
+        name: 'foo',
+        type: 'keyword',
+        location: { min: 0, max: 0 },
+        userDefined: true,
+      },
+      {
+        name: 'bar',
+        type: 'integer',
+        location: { min: 0, max: 0 },
+        userDefined: true,
+      },
+    ]);
+  });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
index 13aa312ccf811..483cd61aa550a 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
@@ -7,6 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+import { uniqBy } from 'lodash';
 import { getExpressionType } from '../../../definitions/utils';
 import { isAssignment, isColumn } from '../../../ast/is';
 import type { ESQLAstItem, ESQLCommand } from '../../../types';
@@ -46,5 +47,5 @@ export const columnsAfter = (
     }
   }
 
-  return [...previousColumns, ...newColumns];
+  return uniqBy([...newColumns, ...previousColumns], 'name');
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
index 483af1df09c7e..7f5dba7bb939e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
@@ -7,22 +7,21 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('FORK', () => {
   const context = {
-    userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
-    fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword' }],
-      ['count', { name: 'count', type: 'double' }],
+    columns: new Map<string, ESQLColumnData>([
+      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+      ['count', { name: 'count', type: 'double', userDefined: false }],
     ]),
   };
   it('adds the _fork in the list of fields', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+    ];
 
     const result = columnsAfter(
       synth.cmd`FORK (LIMIT 10 ) (LIMIT 1000 ) `,
@@ -31,12 +30,136 @@ describe('FORK', () => {
       context
     );
 
-    expect(result).toEqual([
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
       {
         name: '_fork',
         type: 'keyword',
+        userDefined: false,
+      },
+    ]);
+  });
+
+  it('collects columns from branches', () => {
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+    ];
+
+    const result = columnsAfter(
+      synth.cmd`FORK (EVAL foo = 1 | RENAME foo AS bar) (EVAL lolz = 2 + 3 | EVAL field1 = 2.) `,
+      previousCommandFields,
+      '',
+      context
+    );
+
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'bar', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      { name: 'field1', type: 'double', userDefined: true, location: { min: 0, max: 0 } },
+      { name: 'field2', type: 'double', userDefined: false },
+      { name: 'lolz', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      {
+        name: '_fork',
+        type: 'keyword',
+        userDefined: false,
+      },
+    ]);
+  });
+
+  it('prefers userDefined columns over fields with the same name', () => {
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'foo', type: 'keyword', userDefined: false },
+      { name: 'bar', type: 'double', userDefined: false },
+    ];
+
+    // Branch 1: foo is userDefined, bar is not
+    // Branch 2: foo is not userDefined, bar is userDefined
+    const result = columnsAfter(
+      synth.cmd`FORK (EVAL foo = 1) (EVAL bar = 2)`,
+      previousCommandFields,
+      '',
+      context
+    );
+
+    // foo from branch 1 is userDefined, bar from branch 2 is userDefined
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      { name: 'bar', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      {
+        name: '_fork',
+        type: 'keyword',
+        userDefined: false,
+      },
+    ]);
+  });
+
+  it('keeps the first userDefined column if both branches define userDefined columns with the same name', () => {
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'foo', type: 'keyword', userDefined: false },
+    ];
+
+    // Both branches define foo as userDefined, but with different types
+    const result = columnsAfter(
+      synth.cmd`FORK (EVAL foo = 1) (EVAL foo = 2.5)`,
+      previousCommandFields,
+      '',
+      context
+    );
+
+    // The first userDefined column wins (type: integer)
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      {
+        name: '_fork',
+        type: 'keyword',
+        userDefined: false,
+      },
+    ]);
+  });
+
+  it("doesn't duplicate fields from branches", () => {
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'foo', type: 'keyword', userDefined: false },
+    ];
+
+    // All branches keep foo as a field, but with different types
+    const result = columnsAfter(
+      synth.cmd`FORK (KEEP foo) (KEEP foo) (KEEP foo)`,
+      previousCommandFields,
+      '',
+      context
+    );
+
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'foo', type: 'keyword', userDefined: false },
+      {
+        name: '_fork',
+        type: 'keyword',
+        userDefined: false,
+      },
+    ]);
+  });
+
+  it('prefers userDefined column if one branch overwrites a field', () => {
+    const previousCommandFields: ESQLColumnData[] = [
+      { name: 'foo', type: 'keyword', userDefined: false },
+    ];
+
+    // Branch 1: foo is userDefined, Branch 2: foo is not userDefined
+    const result = columnsAfter(
+      synth.cmd`FORK (EVAL foo = 1) (LIMIT 1)`,
+      previousCommandFields,
+      '',
+      context
+    );
+
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      {
+        name: '_fork',
+        type: 'keyword',
+        userDefined: false,
       },
     ]);
   });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index 053dbcb064606..9d20def537a45 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -7,9 +7,12 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { uniqBy } from 'lodash';
+import { col } from '../../../synth';
+import type { ESQLAstQueryExpression } from '../../../types';
 import { type ESQLCommand } from '../../../types';
 import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
 import type { ICommandContext } from '../../types';
+import { esqlCommandRegistry } from '../../../..';
 
 export const columnsAfter = (
   command: ESQLCommand,
@@ -17,9 +20,74 @@ export const columnsAfter = (
   query: string,
   context?: ICommandContext
 ) => {
+  const branches = command.args as ESQLAstQueryExpression[];
+
+  const columnsFromBranches = [];
+
+  for (const branch of branches) {
+    // start with columns from before FORK
+    let columnsFromBranch = [...previousColumns];
+    for (const branchCommand of branch.commands) {
+      const commandDef = esqlCommandRegistry.getCommandByName(branchCommand.name);
+      if (commandDef?.methods?.columnsAfter) {
+        columnsFromBranch = commandDef.methods?.columnsAfter?.(
+          branchCommand,
+          columnsFromBranch,
+          query,
+          context
+        );
+      }
+    }
+
+    columnsFromBranches.push(columnsFromBranch);
+  }
+
+  /**
+   * One of the branches may have overwritten
+   */
+
+  const maps = columnsFromBranches.map((cols) => new Map(cols.map((_col) => [_col.name, _col])));
+
+  const merged = new Map<string, ESQLColumnData>();
+
+  // O(b * n), where b is the branches and n is the number of columns
+  for (const map of maps) {
+    for (const [name, colData] of map) {
+      /**
+       * Check for conflicts...
+       *
+       * The branches often produce large numbers of columns with the same names
+       * because they will often include most or all fields in their results.
+       *
+       * Sometimes, user-defined columns have been created that overwrite
+       * fields with the same name, in which case user-defined columns should win.
+       *
+       * In other cases, the branches may define columns of the same name but with different
+       * data types. That is not allowed but, to keep things simple, we don't touch that case...
+       * we leave it to Elasticsearch validation.
+       */
+      if (merged.has(name)) {
+        const existingColData = merged.get(name) as ESQLColumnData;
+        if (existingColData.userDefined) {
+          // If the existing column is user-defined and the new one is not, keep the existing one
+          continue;
+        } else if (!existingColData.userDefined && colData.userDefined) {
+          // If the existing column is not user-defined and the new one is, use the new one
+          merged.set(name, colData);
+          continue;
+        } else {
+          // If both columns are user-defined or both are not, keep the existing one
+          continue;
+        }
+      }
+
+      merged.set(name, colData);
+    }
+  }
+
   return uniqBy(
     [
-      ...previousColumns,
+      ...merged.values(),
       {
         name: '_fork',
         type: 'keyword' as const,
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
index 3af346c585b19..e941131208280 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
@@ -7,17 +7,15 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('KEEP', () => {
-  const context = {
-    userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
-    fields: new Map<string, ESQLFieldWithMetadata>([
-      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-      ['count', { name: 'count', type: 'double', userDefined: false }],
-    ]),
-  };
+  const columns = new Map<string, ESQLColumnData>([
+    ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
+    ['count', { name: 'count', type: 'double', userDefined: false }],
+  ]);
+  const context = { columns };
   it('should return the correct fields after the command', () => {
     const previousCommandFields = [
       { name: 'field1', type: 'keyword', userDefined: false },

From 81a240c822c86195382234e1ce0b11c7cb23aee4 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 2 Sep 2025 15:09:12 -0600
Subject: [PATCH 19/54] grok

---
 .../commands/fork/columns_after.ts            |  1 -
 .../commands/grok/columns_after.test.ts       | 22 ++++++++-----------
 .../src/commands_registry/registry.ts         |  1 +
 3 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index 9d20def537a45..2af89aaead4b6 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -7,7 +7,6 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { uniqBy } from 'lodash';
-import { col } from '../../../synth';
 import type { ESQLAstQueryExpression } from '../../../types';
 import { type ESQLCommand } from '../../../types';
 import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
index d07060fc023f3..07bb630babcac 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import type { ESQLFieldWithMetadata } from '../../types';
 import { columnsAfter, extractSemanticsFromGrok } from './columns_after';
 
 describe('GROK', () => {
@@ -41,17 +41,13 @@ describe('GROK', () => {
   });
   describe('columnsAfter', () => {
     const context = {
-      userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([]),
-      fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword' }],
-        ['count', { name: 'count', type: 'double' }],
-      ]),
+      columns: new Map(),
     };
     it('adds the GROK columns from the pattern in the list', () => {
-      const previousCommandFields = [
-        { name: 'field1', type: 'keyword' },
-        { name: 'field2', type: 'double' },
-      ] as ESQLFieldWithMetadata[];
+      const previousCommandFields: ESQLFieldWithMetadata[] = [
+        { name: 'field1', type: 'keyword', userDefined: false },
+        { name: 'field2', type: 'double', userDefined: false },
+      ];
 
       const result = columnsAfter(
         synth.cmd`GROK agent "%{WORD:firstWord}"`,
@@ -61,9 +57,9 @@ describe('GROK', () => {
       );
 
       expect(result).toEqual([
-        { name: 'field1', type: 'keyword' },
-        { name: 'field2', type: 'double' },
-        { name: 'firstWord', type: 'keyword' },
+        { name: 'field1', type: 'keyword', userDefined: false },
+        { name: 'field2', type: 'double', userDefined: false },
+        { name: 'firstWord', type: 'keyword', userDefined: false },
       ]);
     });
   });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
index cdcf71961f13a..c876bb3826d6f 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
@@ -61,6 +61,7 @@ export interface ICommandMethods<TContext = any> {
     command: ESQLCommand,
     previousColumns: ESQLColumnData[],
     query: string,
+    // TODO is there any reason for this context to be here?
     context?: TContext
   ) => ESQLColumnData[];
 }

From 6e428be156c89bdc3d11ff8f707b6b2a9a385e34 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 2 Sep 2025 15:12:58 -0600
Subject: [PATCH 20/54] keep

---
 .../commands_registry/commands/keep/columns_after.test.ts | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
index e941131208280..0698135e6ee0a 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLColumnData } from '../../types';
+import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('KEEP', () => {
@@ -17,13 +17,13 @@ describe('KEEP', () => {
   ]);
   const context = { columns };
   it('should return the correct fields after the command', () => {
-    const previousCommandFields = [
+    const previousCommandFields: ESQLFieldWithMetadata[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
-    ] as ESQLFieldWithMetadata[];
+    ];
 
     const result = columnsAfter(synth.cmd`KEEP field1`, previousCommandFields, '', context);
 
-    expect(result).toEqual([{ name: 'field1', type: 'keyword' }]);
+    expect(result).toEqual([{ name: 'field1', type: 'keyword', userDefined: false }]);
   });
 });

From f2afdeeb70b28fcf3d86762cf467b05ad90fecf4 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 2 Sep 2025 15:37:02 -0600
Subject: [PATCH 21/54] stats working

---
 .../commands/stats/columns_after.test.ts      | 187 +++++++-----------
 .../commands/stats/columns_after.ts           |   9 +-
 2 files changed, 73 insertions(+), 123 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
index 005b2c1e35522..b5b197820c1a1 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
@@ -6,31 +6,19 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { synth } from '../../../..';
-import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
+import { Parser, synth } from '../../../..';
+import type { ESQLColumnData } from '../../types';
+import { type ESQLFieldWithMetadata } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('STATS', () => {
   it('adds the user defined column, when no grouping is given', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLFieldWithMetadata[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+    ];
     const context = {
-      userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([
-        [
-          'var0',
-          [
-            {
-              name: 'var0',
-              type: 'double',
-              location: { min: 0, max: 10 },
-              userDefined: true,
-            },
-          ],
-        ],
-      ]),
-      fields: new Map<string, ESQLFieldWithMetadata>([
+      columns: new Map<string, ESQLFieldWithMetadata>([
         ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
         ['count', { name: 'count', type: 'double', userDefined: false }],
       ]),
@@ -43,127 +31,88 @@ describe('STATS', () => {
       context
     );
 
-    expect(result).toEqual([{ name: 'var0', type: 'double', userDefined: false }]);
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'var0', type: 'double', userDefined: true, location: { min: 0, max: 0 } },
+    ]);
   });
 
   it('adds the escaped column, when no grouping is given', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLFieldWithMetadata[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+    ];
 
-    const context = {
-      userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([
-        [
-          'AVG(field2)',
-          [
-            {
-              name: 'AVG(field2)',
-              type: 'double',
-              location: { min: 0, max: 10 },
-              userDefined: true,
-            },
-          ],
-        ],
-      ]),
-      fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-        ['count', { name: 'count', type: 'double', userDefined: false }],
-      ]),
-    };
+    const queryString = `FROM index | STATS AVG(field2)`;
 
-    const result = columnsAfter(synth.cmd`STATS AVG(field2)`, previousCommandFields, '', context);
+    // Can't use synth because it steps on the location information
+    // which is used to determine the name of the new column
+    const {
+      root: {
+        commands: [, command],
+      },
+    } = Parser.parseQuery(queryString);
 
-    expect(result).toEqual([{ name: 'AVG(field2)', type: 'double', userDefined: false }]);
+    const result = columnsAfter(command, previousCommandFields, queryString, {
+      columns: new Map(),
+    });
+
+    expect(result).toEqual([
+      { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 19, max: 29 } },
+    ]);
   });
 
   it('adds the escaped and grouping columns', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLFieldWithMetadata[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+    ];
 
-    const context = {
-      userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([
-        [
-          'AVG(field2)',
-          [
-            {
-              name: 'AVG(field2)',
-              type: 'double',
-              location: { min: 0, max: 10 },
-              userDefined: true,
-            },
-          ],
-        ],
-      ]),
-      fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-        ['count', { name: 'count', type: 'double', userDefined: false }],
-      ]),
-    };
+    const queryString = `FROM a | STATS AVG(field2) BY field1`;
 
-    const result = columnsAfter(
-      synth.cmd`STATS AVG(field2) BY field1`,
-      previousCommandFields,
-      '',
-      context
-    );
+    // Can't use synth because it steps on the location information
+    // which is used to determine the name of the new column
+    const {
+      root: {
+        commands: [, command],
+      },
+    } = Parser.parseQuery(queryString);
+
+    const result = columnsAfter(command, previousCommandFields, queryString, {
+      columns: new Map(),
+    });
 
     expect(result).toEqual([
-      { name: 'field1', type: 'keyword', userDefined: false },
-      { name: 'AVG(field2)', type: 'double', userDefined: false },
+      { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 15, max: 25 } },
+      { name: 'field1', type: 'keyword', userDefined: true, location: { min: 30, max: 35 } },
     ]);
   });
 
   it('adds the user defined and grouping columns', () => {
-    const previousCommandFields = [
-      { name: 'field1', type: 'keyword' },
-      { name: 'field2', type: 'double' },
-      { name: '@timestamp', type: 'date' },
-    ] as ESQLFieldWithMetadata[];
+    const previousCommandFields: ESQLFieldWithMetadata[] = [
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+      { name: '@timestamp', type: 'date', userDefined: false },
+    ];
 
     const context = {
-      userDefinedColumns: new Map<string, ESQLUserDefinedColumn[]>([
-        [
-          'AVG(field2)',
-          [
-            {
-              name: 'AVG(field2)',
-              type: 'double',
-              location: { min: 0, max: 10 },
-              userDefined: true,
-            },
-          ],
-        ],
-        [
-          'buckets',
-          [
-            {
-              name: 'buckets',
-              type: 'unknown',
-              location: { min: 0, max: 10 },
-              userDefined: true,
-            },
-          ],
-        ],
-      ]),
-      fields: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-        ['count', { name: 'count', type: 'double', userDefined: false }],
-      ]),
+      columns: new Map(),
     };
 
-    const result = columnsAfter(
-      synth.cmd`STATS AVG(field2) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`,
-      previousCommandFields,
-      '',
-      context
-    );
+    const queryString = `FROM a | STATS AVG(field2) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`;
 
-    expect(result).toEqual([
-      { name: 'AVG(field2)', type: 'double', userDefined: false },
-      { name: 'buckets', type: 'unknown', userDefined: false },
+    // Can't use synth because it steps on the location information
+    // which is used to determine the name of the new column
+    const {
+      root: {
+        commands: [, command],
+      },
+    } = Parser.parseQuery(queryString);
+
+    const result = columnsAfter(command, previousCommandFields, queryString, context);
+
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 15, max: 25 } },
+      { name: 'buckets', type: 'date', userDefined: true, location: { min: 30, max: 36 } },
     ]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
index eb0a5b7c5b478..7c618ba88469c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
@@ -15,7 +15,8 @@ import type { ESQLColumnData, ESQLUserDefinedColumn, ICommandContext } from '../
 
 const getUserDefinedColumns = (
   command: ESQLCommand | ESQLCommandOption,
-  typeOf: (thing: ESQLAstItem) => SupportedDataType | 'unknown'
+  typeOf: (thing: ESQLAstItem) => SupportedDataType | 'unknown',
+  query: string
 ): ESQLUserDefinedColumn[] => {
   const columns = [];
 
@@ -33,13 +34,13 @@ const getUserDefinedColumns = (
     }
 
     if (isOptionNode(expression) && expression.name === 'by') {
-      columns.push(...getUserDefinedColumns(expression, typeOf));
+      columns.push(...getUserDefinedColumns(expression, typeOf, query));
       continue;
     }
 
     if (!isOptionNode(expression) && !Array.isArray(expression)) {
       const newColumn: ESQLUserDefinedColumn = {
-        name: expression.text,
+        name: query.substring(expression.location.min, expression.location.max + 1),
         type: typeOf(expression),
         location: expression.location,
         userDefined: true,
@@ -64,5 +65,5 @@ export const columnsAfter = (
   const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, columnMap);
 
   // TODO - is this uniqby helpful? Does it do what we expect?
-  return uniqBy([...previousColumns, ...getUserDefinedColumns(command, typeOf)], 'name');
+  return uniqBy([...getUserDefinedColumns(command, typeOf, query)], 'name');
 };

From 1942b3b9a1ba7cb7f865200cc56497f6cbb27562 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 3 Sep 2025 16:32:37 -0600
Subject: [PATCH 22/54] remove context from columns-after call signature

---
 .../change_point/columns_after.test.ts        | 16 ++--------
 .../commands/change_point/columns_after.ts    |  5 ++--
 .../commands/completion/columns_after.test.ts | 12 ++------
 .../commands/completion/columns_after.ts      | 10 ++-----
 .../commands/dissect/columns_after.test.ts    | 13 +--------
 .../commands/dissect/columns_after.ts         |  4 +--
 .../commands/drop/columns_after.test.ts       |  8 +----
 .../commands/drop/columns_after.ts            |  4 +--
 .../commands/eval/columns_after.ts            |  7 ++---
 .../commands/fork/columns_after.test.ts       | 24 ++++-----------
 .../commands/fork/columns_after.ts            |  9 ++----
 .../commands/grok/columns_after.test.ts       |  6 +---
 .../commands/grok/columns_after.ts            |  6 ++--
 .../commands/keep/columns_after.test.ts       |  9 ++----
 .../commands/keep/columns_after.ts            |  6 ++--
 .../commands/rename/columns_after.test.ts     | 22 ++------------
 .../commands/rename/columns_after.ts          |  5 ++--
 .../commands/stats/columns_after.test.ts      | 29 ++++---------------
 .../commands/stats/columns_after.ts           |  5 ++--
 .../src/commands_registry/registry.ts         |  4 +--
 .../src/shared/helpers.ts                     |  4 +--
 21 files changed, 45 insertions(+), 163 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
index aef7bcbba1560..56a086e22c668 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts
@@ -11,24 +11,13 @@ import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('CHANGE_POINT > columnsAfter', () => {
-  const context = {
-    columns: new Map<string, ESQLColumnData>([
-      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-      ['count', { name: 'count', type: 'double', userDefined: false }],
-    ]),
-  };
   it('adds "type" and "pvalue" fields, when AS option not specified', () => {
     const previousCommandFields: ESQLColumnData[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'count', type: 'double', userDefined: false },
     ];
 
-    const result = columnsAfter(
-      synth.cmd`CHANGE_POINT count ON field1`,
-      previousCommandFields,
-      '',
-      context
-    );
+    const result = columnsAfter(synth.cmd`CHANGE_POINT count ON field1`, previousCommandFields, '');
 
     expect(result).toEqual<ESQLColumnData[]>([
       { name: 'field1', type: 'keyword', userDefined: false },
@@ -47,8 +36,7 @@ describe('CHANGE_POINT > columnsAfter', () => {
     const result = columnsAfter(
       synth.cmd`CHANGE_POINT count ON field1 AS changePointType, pValue`,
       previousCommandFields,
-      '',
-      context
+      ''
     );
 
     expect(result).toEqual<ESQLColumnData[]>([
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
index cf8dd86602202..45f7cbedbfd0f 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts
@@ -9,13 +9,12 @@
 import uniqBy from 'lodash/uniqBy';
 import { LeafPrinter } from '../../../pretty_print/leaf_printer';
 import { type ESQLAstChangePointCommand, type ESQLCommand } from '../../../types';
-import type { ESQLColumnData, ICommandContext } from '../../types';
+import type { ESQLColumnData } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const { target } = command as ESQLAstChangePointCommand;
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
index d086c515a8544..c08c11e2d2498 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts
@@ -11,12 +11,6 @@ import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('COMPLETION', () => {
-  const context = {
-    columns: new Map<string, ESQLColumnData>([
-      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-      ['count', { name: 'count', type: 'double', userDefined: false }],
-    ]),
-  };
   it('adds "completion" field, when targetField is not specified', () => {
     const previousCommandColumns: ESQLColumnData[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
@@ -26,8 +20,7 @@ describe('COMPLETION', () => {
     const result = columnsAfter(
       synth.cmd`COMPLETION "prompt" WITH {"inference_id": "my-inference-id"}`,
       previousCommandColumns,
-      '',
-      context
+      ''
     );
 
     expect(result).toEqual([
@@ -46,8 +39,7 @@ describe('COMPLETION', () => {
     const result = columnsAfter(
       synth.cmd`COMPLETION customField = "prompt" WITH {"inference_id": "my-inference-id"}`,
       previousCommandColumns,
-      '',
-      context
+      ''
     );
 
     expect(result).toEqual([
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
index c57dbe0ffce5e..e128a20492fca 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts
@@ -9,18 +9,12 @@
 import uniqBy from 'lodash/uniqBy';
 import { LeafPrinter } from '../../../pretty_print/leaf_printer';
 import { type ESQLAstCompletionCommand, type ESQLCommand } from '../../../types';
-import type {
-  ESQLColumnData,
-  ESQLFieldWithMetadata,
-  ESQLUserDefinedColumn,
-  ICommandContext,
-} from '../../types';
+import type { ESQLColumnData, ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const { targetField } = command as ESQLAstCompletionCommand;
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
index fe8f646f61ff3..18d155383b8ba 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts
@@ -68,24 +68,13 @@ describe('DISSECT', () => {
     });
   });
   describe('columnsAfter', () => {
-    const context = {
-      columns: new Map<string, ESQLColumnData>([
-        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-        ['count', { name: 'count', type: 'double', userDefined: false }],
-      ]),
-    };
     it('adds the DISSECT pattern columns as fields', () => {
       const previousColumns: ESQLColumnData[] = [
         { name: 'field1', type: 'keyword', userDefined: false },
         { name: 'field2', type: 'double', userDefined: false },
       ];
 
-      const result = columnsAfter(
-        synth.cmd`DISSECT agent "%{firstWord}"`,
-        previousColumns,
-        '',
-        context
-      );
+      const result = columnsAfter(synth.cmd`DISSECT agent "%{firstWord}"`, previousColumns, '');
 
       expect(result).toEqual<ESQLColumnData[]>([
         { name: 'field1', type: 'keyword', userDefined: false },
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
index 18fdeaa09d09c..c8d7cbdf3c38c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts
@@ -9,7 +9,6 @@
 import { walk } from '../../../walker';
 import { type ESQLCommand } from '../../../types';
 import type { ESQLColumnData } from '../../types';
-import type { ICommandContext } from '../../types';
 
 function unquoteTemplate(inputString: string): string {
   if (inputString.startsWith('"') && inputString.endsWith('"') && inputString.length >= 2) {
@@ -36,8 +35,7 @@ export function extractDissectColumnNames(pattern: string): string[] {
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const columns: string[] = [];
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
index b977ca03e7537..32bcb28fb34c1 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts
@@ -11,19 +11,13 @@ import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('DROP', () => {
-  const context = {
-    columns: new Map<string, ESQLColumnData>([
-      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-      ['count', { name: 'count', type: 'double', userDefined: false }],
-    ]),
-  };
   it('removes the columns defined in the command', () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
     ];
 
-    const result = columnsAfter(synth.cmd`DROP field1`, previousColumns, '', context);
+    const result = columnsAfter(synth.cmd`DROP field1`, previousColumns, '');
 
     expect(result).toEqual([{ name: 'field2', type: 'double', userDefined: false }]);
   });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
index b17c457435ebd..491095041f95c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts
@@ -9,13 +9,11 @@
 import { walk } from '../../../walker';
 import { type ESQLCommand } from '../../../types';
 import type { ESQLColumnData } from '../../types';
-import type { ICommandContext } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const columnsToDrop: string[] = [];
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
index 483cd61aa550a..0e500606a63f6 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts
@@ -8,16 +8,15 @@
  */
 
 import { uniqBy } from 'lodash';
-import { getExpressionType } from '../../../definitions/utils';
 import { isAssignment, isColumn } from '../../../ast/is';
+import { getExpressionType } from '../../../definitions/utils';
 import type { ESQLAstItem, ESQLCommand } from '../../../types';
-import type { ESQLColumnData, ESQLUserDefinedColumn, ICommandContext } from '../../types';
+import type { ESQLColumnData, ESQLUserDefinedColumn } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const columnMap = new Map<string, ESQLColumnData>();
   previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
index 7f5dba7bb939e..2dec08c24aec5 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
@@ -11,12 +11,6 @@ import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('FORK', () => {
-  const context = {
-    columns: new Map<string, ESQLColumnData>([
-      ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-      ['count', { name: 'count', type: 'double', userDefined: false }],
-    ]),
-  };
   it('adds the _fork in the list of fields', () => {
     const previousCommandFields: ESQLColumnData[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
@@ -26,8 +20,7 @@ describe('FORK', () => {
     const result = columnsAfter(
       synth.cmd`FORK (LIMIT 10 ) (LIMIT 1000 ) `,
       previousCommandFields,
-      '',
-      context
+      ''
     );
 
     expect(result).toEqual<ESQLColumnData[]>([
@@ -50,8 +43,7 @@ describe('FORK', () => {
     const result = columnsAfter(
       synth.cmd`FORK (EVAL foo = 1 | RENAME foo AS bar) (EVAL lolz = 2 + 3 | EVAL field1 = 2.) `,
       previousCommandFields,
-      '',
-      context
+      ''
     );
 
     expect(result).toEqual<ESQLColumnData[]>([
@@ -78,8 +70,7 @@ describe('FORK', () => {
     const result = columnsAfter(
       synth.cmd`FORK (EVAL foo = 1) (EVAL bar = 2)`,
       previousCommandFields,
-      '',
-      context
+      ''
     );
 
     // foo from branch 1 is userDefined, bar from branch 2 is userDefined
@@ -103,8 +94,7 @@ describe('FORK', () => {
     const result = columnsAfter(
       synth.cmd`FORK (EVAL foo = 1) (EVAL foo = 2.5)`,
       previousCommandFields,
-      '',
-      context
+      ''
     );
 
     // The first userDefined column wins (type: integer)
@@ -127,8 +117,7 @@ describe('FORK', () => {
     const result = columnsAfter(
       synth.cmd`FORK (KEEP foo) (KEEP foo) (KEEP foo)`,
       previousCommandFields,
-      '',
-      context
+      ''
     );
 
     expect(result).toEqual<ESQLColumnData[]>([
@@ -150,8 +139,7 @@ describe('FORK', () => {
     const result = columnsAfter(
       synth.cmd`FORK (EVAL foo = 1) (LIMIT 1)`,
       previousCommandFields,
-      '',
-      context
+      ''
     );
 
     expect(result).toEqual<ESQLColumnData[]>([
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index 2af89aaead4b6..231dfd623d7ca 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -7,17 +7,15 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { uniqBy } from 'lodash';
+import { esqlCommandRegistry } from '../../../..';
 import type { ESQLAstQueryExpression } from '../../../types';
 import { type ESQLCommand } from '../../../types';
 import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
-import type { ICommandContext } from '../../types';
-import { esqlCommandRegistry } from '../../../..';
 
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const branches = command.args as ESQLAstQueryExpression[];
 
@@ -32,8 +30,7 @@ export const columnsAfter = (
         columnsFromBranch = commandDef.methods?.columnsAfter?.(
           branchCommand,
           columnsFromBranch,
-          query,
-          context
+          query
         );
       }
     }
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
index 07bb630babcac..1770473508517 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts
@@ -40,9 +40,6 @@ describe('GROK', () => {
     });
   });
   describe('columnsAfter', () => {
-    const context = {
-      columns: new Map(),
-    };
     it('adds the GROK columns from the pattern in the list', () => {
       const previousCommandFields: ESQLFieldWithMetadata[] = [
         { name: 'field1', type: 'keyword', userDefined: false },
@@ -52,8 +49,7 @@ describe('GROK', () => {
       const result = columnsAfter(
         synth.cmd`GROK agent "%{WORD:firstWord}"`,
         previousCommandFields,
-        '',
-        context
+        ''
       );
 
       expect(result).toEqual([
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
index f400521d26824..70cd7d18af47e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts
@@ -6,10 +6,9 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { walk } from '../../../walker';
 import { type ESQLCommand } from '../../../types';
+import { walk } from '../../../walker';
 import type { ESQLColumnData } from '../../types';
-import type { ICommandContext } from '../../types';
 
 function unquoteTemplate(inputString: string): string {
   if (inputString.startsWith('"') && inputString.endsWith('"') && inputString.length >= 2) {
@@ -51,8 +50,7 @@ export function extractSemanticsFromGrok(pattern: string): string[] {
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const columns: string[] = [];
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
index 0698135e6ee0a..b94cd73b215cc 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts
@@ -7,22 +7,17 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { synth } from '../../../..';
-import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
+import type { ESQLFieldWithMetadata } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('KEEP', () => {
-  const columns = new Map<string, ESQLColumnData>([
-    ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-    ['count', { name: 'count', type: 'double', userDefined: false }],
-  ]);
-  const context = { columns };
   it('should return the correct fields after the command', () => {
     const previousCommandFields: ESQLFieldWithMetadata[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
     ];
 
-    const result = columnsAfter(synth.cmd`KEEP field1`, previousCommandFields, '', context);
+    const result = columnsAfter(synth.cmd`KEEP field1`, previousCommandFields, '');
 
     expect(result).toEqual([{ name: 'field1', type: 'keyword', userDefined: false }]);
   });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
index cf8bac3dbf7cc..72235cd4a8f3c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts
@@ -6,16 +6,14 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { walk } from '../../../walker';
 import { type ESQLCommand } from '../../../types';
+import { walk } from '../../../walker';
 import type { ESQLColumnData } from '../../types';
-import type { ICommandContext } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const columnsToKeep: string[] = [];
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
index 05a1e1249a4ea..bf63511267fb6 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts
@@ -11,23 +11,13 @@ import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('RENAME', () => {
-  const columns = new Map<string, ESQLColumnData>([
-    ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-    ['count', { name: 'count', type: 'double', userDefined: false }],
-  ]);
-  const context = { columns };
   it('renames the given columns with the new names using AS', () => {
     const previousCommandFields: ESQLColumnData[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
     ];
 
-    const result = columnsAfter(
-      synth.cmd`RENAME field1 as meow`,
-      previousCommandFields,
-      '',
-      context
-    );
+    const result = columnsAfter(synth.cmd`RENAME field1 as meow`, previousCommandFields, '');
 
     expect(result).toEqual<ESQLColumnData[]>([
       { name: 'meow', type: 'keyword', userDefined: true, location: { max: 0, min: 0 } },
@@ -41,12 +31,7 @@ describe('RENAME', () => {
       { name: 'field2', type: 'double', userDefined: false },
     ];
 
-    const result = columnsAfter(
-      synth.cmd`RENAME meow = field1`,
-      previousCommandFields,
-      '',
-      context
-    );
+    const result = columnsAfter(synth.cmd`RENAME meow = field1`, previousCommandFields, '');
 
     expect(result).toEqual<ESQLColumnData[]>([
       { name: 'meow', type: 'keyword', userDefined: true, location: { max: 0, min: 0 } },
@@ -63,8 +48,7 @@ describe('RENAME', () => {
     const result = columnsAfter(
       synth.cmd`RENAME meow = field1, field2 as woof`,
       previousCommandFields,
-      '',
-      context
+      ''
     );
 
     expect(result).toEqual<ESQLColumnData[]>([
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
index 6f867858dcd8d..f5f5b236eda0e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
@@ -9,13 +9,12 @@
 import uniqBy from 'lodash/uniqBy';
 import { isColumn, isFunctionExpression } from '../../../ast/is';
 import type { ESQLAstBaseItem, ESQLCommand, ESQLFunction } from '../../../types';
-import type { ESQLColumnData, ICommandContext } from '../../types';
+import type { ESQLColumnData } from '../../types';
 
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const asRenamePairs: ESQLFunction[] = [];
   const assignRenamePairs: ESQLFunction[] = [];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
index b5b197820c1a1..847dcca5f8dd0 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts
@@ -17,19 +17,8 @@ describe('STATS', () => {
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
     ];
-    const context = {
-      columns: new Map<string, ESQLFieldWithMetadata>([
-        ['field1', { name: 'field1', type: 'keyword', userDefined: false }],
-        ['count', { name: 'count', type: 'double', userDefined: false }],
-      ]),
-    };
-
-    const result = columnsAfter(
-      synth.cmd`STATS var0=AVG(field2)`,
-      previousCommandFields,
-      '',
-      context
-    );
+
+    const result = columnsAfter(synth.cmd`STATS var0=AVG(field2)`, previousCommandFields, '');
 
     expect(result).toEqual<ESQLColumnData[]>([
       { name: 'var0', type: 'double', userDefined: true, location: { min: 0, max: 0 } },
@@ -52,9 +41,7 @@ describe('STATS', () => {
       },
     } = Parser.parseQuery(queryString);
 
-    const result = columnsAfter(command, previousCommandFields, queryString, {
-      columns: new Map(),
-    });
+    const result = columnsAfter(command, previousCommandFields, queryString);
 
     expect(result).toEqual([
       { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 19, max: 29 } },
@@ -77,9 +64,7 @@ describe('STATS', () => {
       },
     } = Parser.parseQuery(queryString);
 
-    const result = columnsAfter(command, previousCommandFields, queryString, {
-      columns: new Map(),
-    });
+    const result = columnsAfter(command, previousCommandFields, queryString);
 
     expect(result).toEqual([
       { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 15, max: 25 } },
@@ -94,10 +79,6 @@ describe('STATS', () => {
       { name: '@timestamp', type: 'date', userDefined: false },
     ];
 
-    const context = {
-      columns: new Map(),
-    };
-
     const queryString = `FROM a | STATS AVG(field2) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`;
 
     // Can't use synth because it steps on the location information
@@ -108,7 +89,7 @@ describe('STATS', () => {
       },
     } = Parser.parseQuery(queryString);
 
-    const result = columnsAfter(command, previousCommandFields, queryString, context);
+    const result = columnsAfter(command, previousCommandFields, queryString);
 
     expect(result).toEqual<ESQLColumnData[]>([
       { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 15, max: 25 } },
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
index 7c618ba88469c..8ebf756af6647 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
@@ -11,7 +11,7 @@ import { isAssignment, isColumn, isOptionNode } from '../../../ast/is';
 import type { SupportedDataType } from '../../../definitions/types';
 import { getExpressionType } from '../../../definitions/utils';
 import type { ESQLAstItem, ESQLCommand, ESQLCommandOption } from '../../../types';
-import type { ESQLColumnData, ESQLUserDefinedColumn, ICommandContext } from '../../types';
+import type { ESQLColumnData, ESQLUserDefinedColumn } from '../../types';
 
 const getUserDefinedColumns = (
   command: ESQLCommand | ESQLCommandOption,
@@ -56,8 +56,7 @@ const getUserDefinedColumns = (
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string,
-  context?: ICommandContext
+  query: string
 ) => {
   const columnMap = new Map<string, ESQLColumnData>();
   previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
index c876bb3826d6f..1d41f09a12974 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
@@ -60,9 +60,7 @@ export interface ICommandMethods<TContext = any> {
   columnsAfter?: (
     command: ESQLCommand,
     previousColumns: ESQLColumnData[],
-    query: string,
-    // TODO is there any reason for this context to be here?
-    context?: TContext
+    query: string
   ) => ESQLColumnData[];
 }
 
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index edefc49925963..a768ba4418302 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -104,9 +104,7 @@ export async function getCurrentQueryAvailableColumns(
   const commandDefinition = esqlCommandRegistry.getCommandByName(lastCommand.name);
 
   if (commandDefinition?.methods.columnsAfter) {
-    return commandDefinition.methods.columnsAfter(lastCommand, previousPipeFields, query, {
-      columns: new Map(),
-    });
+    return commandDefinition.methods.columnsAfter(lastCommand, previousPipeFields, query);
   } else {
     return previousPipeFields;
   }

From 696c468ca0c3e7631fbb6976e746d66fa7e374fa Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 3 Sep 2025 16:52:47 -0600
Subject: [PATCH 23/54] FIXME for EVAL validation

---
 .../src/commands_registry/commands/eval/columns_after.test.ts | 4 ++--
 .../src/validation/validation.ts                              | 2 ++
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts
index fdeffc0e157e2..0126c48f87211 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts
@@ -121,13 +121,13 @@ describe('EVAL > columnsAfter', () => {
     expect(result).toEqual([
       {
         name: 'foo',
-        type: 'keyword',
+        type: 'keyword', // originally integer
         location: { min: 0, max: 0 },
         userDefined: true,
       },
       {
         name: 'bar',
-        type: 'integer',
+        type: 'integer', // originally keyword
         location: { min: 0, max: 0 },
         userDefined: true,
       },
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index 63ce09106597e..474356ea93061 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -142,6 +142,8 @@ async function validateAst(
   const license = await callbacks?.getLicense?.();
   const hasMinimumLicenseRequired = license?.hasAtLeast;
   for (let i = 0; i < rootCommands.length; i++) {
+    // FIXME we need to expand FORK and EVAL like we do in getQueryForFields
+    // _before_ slicing off the later commands
     const partialQuery = queryString.slice(0, rootCommands[i].location.max + 1);
 
     const previousCommands = rootCommands.slice(0, i + 1);

From 6eb82c44054857a6e729f4685c3404bf0ee9eff1 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 4 Sep 2025 10:28:21 -0600
Subject: [PATCH 24/54] skip validation for new columns

---
 .../definitions/utils/validation/function.ts   | 18 ++++++++++++++----
 .../src/validation/__tests__/functions.test.ts | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
index 0e4edac4929c5..57ad7ec883501 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
@@ -153,11 +153,21 @@ class FunctionValidator {
     }
 
     // Validate column arguments
-    const columnMessages = this.fn.args.flat().flatMap((arg) => {
-      if (isColumn(arg) || isIdentifier(arg)) {
-        return new ColumnValidator(arg, this.context, this.parentCommand.name).validate();
+    const columnsToValidate = [];
+    const flatArgs = this.fn.args.flat();
+    for (let i = 0; i < flatArgs.length; i++) {
+      const arg = flatArgs[i];
+      if (
+        (isColumn(arg) || isIdentifier(arg)) &&
+        !(this.definition.name === '=' && i === 0) && // don't validate left-hand side of assignment
+        !(this.definition.name === 'as' && i === 1) // don't validate right-hand side of AS
+      ) {
+        columnsToValidate.push(arg);
       }
-      return [];
+    }
+
+    const columnMessages = columnsToValidate.flatMap((arg) => {
+      return new ColumnValidator(arg, this.context, this.parentCommand.name).validate();
     });
 
     // uniqBy is used to cover a special case in ENRICH where an implicit assignment is possible
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/functions.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/functions.test.ts
index 87b2f7e8971fe..4da35fc502e4b 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/functions.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/functions.test.ts
@@ -391,6 +391,24 @@ describe('function validation', () => {
             ]);
           });
         });
+
+        it('skips column validation for left assignment arg', async () => {
+          const { expectErrors } = await setup();
+
+          await expectErrors('FROM a_index | EVAL lolz = 2', []);
+          await expectErrors('FROM a_index | EVAL lolz = nonexistent', [
+            'Unknown column "nonexistent"',
+          ]);
+        });
+
+        it('skips column validation for right arg to AS', async () => {
+          const { expectErrors } = await setup();
+
+          await expectErrors('FROM a_index | RENAME keywordField AS lolz', []);
+          await expectErrors('FROM a_index | RENAME nonexistent AS lolz', [
+            'Unknown column "nonexistent"',
+          ]);
+        });
       });
     });
 

From 8ac4f193519ab91a3fc414f9513dd429015a9d5f Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 4 Sep 2025 11:47:20 -0600
Subject: [PATCH 25/54] validate by subquery

---
 .../src/autocomplete/helper.ts                |  12 +-
 .../__tests__/column_existence.test.ts        |  10 ++
 .../src/validation/validation.ts              | 127 +++++++++++++++---
 3 files changed, 125 insertions(+), 24 deletions(-)
 create mode 100644 src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_existence.test.ts

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
index 36954fa514f92..f3bebe1c5d11f 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
@@ -69,11 +69,13 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
   // fetch fields, hence return an empty string.
   return commands.length === 1 && ['row', 'show'].includes(commands[0].name)
     ? ''
-    : buildQueryUntilPreviousCommand(queryString, commands);
+    : buildQueryUntilPreviousCommand(root);
 }
 
-// TODO consider replacing this with a pretty printer-based solution
-function buildQueryUntilPreviousCommand(queryString: string, commands: ESQLCommand[]) {
-  const prevCommand = commands[Math.max(commands.length - 2, 0)];
-  return prevCommand ? queryString.substring(0, prevCommand.location.max + 1) : queryString;
+function buildQueryUntilPreviousCommand(root: ESQLAstQueryExpression) {
+  if (root.commands.length === 1) {
+    return BasicPrettyPrinter.print({ ...root.commands[0] });
+  } else {
+    return BasicPrettyPrinter.print({ ...root, commands: root.commands.slice(0, -1) });
+  }
 }
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_existence.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_existence.test.ts
new file mode 100644
index 0000000000000..c42459d577729
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_existence.test.ts
@@ -0,0 +1,10 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+// TODO
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index 474356ea93061..ff23d3c19b0b8 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -7,8 +7,14 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import type { ESQLAst, ESQLCommand, ESQLMessage, ErrorTypes } from '@kbn/esql-ast';
-import { EsqlQuery, walk, esqlCommandRegistry, Builder } from '@kbn/esql-ast';
+import type {
+  ESQLAst,
+  ESQLAstQueryExpression,
+  ESQLCommand,
+  ESQLMessage,
+  ErrorTypes,
+} from '@kbn/esql-ast';
+import { EsqlQuery, walk, esqlCommandRegistry, Builder, BasicPrettyPrinter } from '@kbn/esql-ast';
 import { getMessageFromId } from '@kbn/esql-ast/src/definitions/utils';
 import type {
   ESQLFieldWithMetadata,
@@ -141,32 +147,36 @@ async function validateAst(
 
   const license = await callbacks?.getLicense?.();
   const hasMinimumLicenseRequired = license?.hasAtLeast;
-  for (let i = 0; i < rootCommands.length; i++) {
-    // FIXME we need to expand FORK and EVAL like we do in getQueryForFields
-    // _before_ slicing off the later commands
-    const partialQuery = queryString.slice(0, rootCommands[i].location.max + 1);
-
-    const previousCommands = rootCommands.slice(0, i + 1);
-    const queryForFields = getQueryForFields(
-      partialQuery,
-      Builder.expression.query(previousCommands)
-    );
 
-    const { getColumnMap: getFieldsMap } = getColumnsByTypeHelper(queryForFields, callbacks);
-    const availableFields = await getFieldsMap();
+  /**
+   * Even though we are validating single commands, we work with subqueries.
+   *
+   * The reason is that building the list of columns available in each command requires
+   * the full command subsequence that precedes that command.
+   */
+  const subqueries = getSubqueriesToValidate(rootCommands);
+  for (const subquery of subqueries) {
+    const queryForFields = getQueryForFields(BasicPrettyPrinter.print(subquery), subquery);
+    const { getColumnMap } = getColumnsByTypeHelper(queryForFields, callbacks);
+    const availableColumns = await getColumnMap();
 
     const references: ReferenceMaps = {
       sources,
-      columns: availableFields,
+      columns: availableColumns,
       policies: availablePolicies,
       query: queryString,
       joinIndices: joinIndices?.indices || [],
     };
 
-    const commandMessages = validateCommand(rootCommands[i], references, rootCommands, {
-      ...callbacks,
-      hasMinimumLicenseRequired,
-    });
+    const commandMessages = validateCommand(
+      subquery.commands[subquery.commands.length - 1],
+      references,
+      rootCommands,
+      {
+        ...callbacks,
+        hasMinimumLicenseRequired,
+      }
+    );
     messages.push(...commandMessages);
   }
 
@@ -266,3 +276,82 @@ function validateUnsupportedTypeFields(fields: Map<string, ESQLFieldWithMetadata
   }
   return messages;
 }
+
+/**
+ * Returns a list of subqueries to validate
+ * @param rootCommands
+ *
+ * TODO - extract to another module?
+ */
+function getSubqueriesToValidate(rootCommands: ESQLCommand[]) {
+  const subsequences = [];
+  const expandedCommands = expandEvals(rootCommands);
+  for (let i = 0; i < expandedCommands.length; i++) {
+    const command = expandedCommands[i];
+
+    // every command within FORK's branches is its own subquery to be validated
+    if (command.name.toLowerCase() === 'fork') {
+      const branchSubqueries = getForkBranchSubqueries(command as ESQLCommand<'fork'>);
+      for (const subquery of branchSubqueries) {
+        subsequences.push([...expandedCommands.slice(0, i), ...subquery]);
+      }
+    }
+
+    subsequences.push(expandedCommands.slice(0, i + 1));
+  }
+
+  return subsequences.map((subsequence) => Builder.expression.query(subsequence));
+}
+
+/**
+ * Expands EVAL commands into separate commands for each expression.
+ *
+ * E.g. `EVAL 1 + 2, 3 + 4` becomes `EVAL 1 + 2 | EVAL 3 + 4`
+ *
+ * This is logically equivalent and makes validation and field existence detection much easier.
+ *
+ * @param commands The list of commands to expand.
+ * @returns The expanded list of commands.
+ */
+function expandEvals(commands: ESQLCommand[]): ESQLCommand[] {
+  const expanded: ESQLCommand[] = [];
+  for (const command of commands) {
+    if (command.name.toLowerCase() === 'eval') {
+      // treat each expression within EVAL as a separate EVAL command
+      for (const arg of command.args) {
+        expanded.push(
+          Builder.command({
+            name: 'eval',
+            args: [arg],
+            location: command.location,
+          })
+        );
+      }
+    } else {
+      expanded.push(command);
+    }
+  }
+  return expanded;
+}
+
+/**
+ * Expands a FORK command into queries for each command in each branch.
+ *
+ * E.g. FORK (EVAL 1 | LIMIT 10) (RENAME foo AS bar | DROP lolz)
+ *
+ * becomes [`EVAL 1`, `EVAL 1 | LIMIT 10`, `RENAME foo AS bar`, `RENAME foo AS bar | DROP lolz`]
+ *
+ * @param command a FORK command
+ * @returns an array of expanded subqueries
+ */
+function getForkBranchSubqueries(command: ESQLCommand<'fork'>): ESQLCommand[][] {
+  const expanded: ESQLCommand[][] = [];
+  const branches = command.args as ESQLAstQueryExpression[];
+  for (let j = 0; j < branches.length; j++) {
+    for (let k = 0; k < branches[j].commands.length; k++) {
+      const partialQuery = branches[j].commands.slice(0, k + 1);
+      expanded.push(partialQuery);
+    }
+  }
+  return expanded;
+}

From addabd170546222c7a626a190417dea3b3635a87 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 4 Sep 2025 13:20:36 -0600
Subject: [PATCH 26/54] support join fields

---
 .../commands/join/columns_after.test.ts       | 52 +++++++++++++++++++
 .../commands/join/columns_after.ts            | 20 +++++++
 .../commands_registry/commands/join/index.ts  |  2 +
 .../src/shared/helpers.ts                     | 39 ++++++++++++--
 .../src/shared/resources_helpers.ts           | 33 ++++++++----
 5 files changed, 130 insertions(+), 16 deletions(-)
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
new file mode 100644
index 0000000000000..0c48f3fa57658
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
@@ -0,0 +1,52 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
+import { columnsAfter } from './columns_after';
+
+describe('JOIN columnsAfter', () => {
+  it('returns previousColumns when joinColumns is undefined', () => {
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'fieldA', type: 'keyword', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ];
+    const result = columnsAfter({} as any, previousColumns, 'JOIN');
+    expect(result).toEqual(previousColumns);
+  });
+
+  it('prepends joinColumns to previousColumns when joinColumns is provided', () => {
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'fieldA', type: 'keyword', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ];
+    const joinColumns: ESQLFieldWithMetadata[] = [
+      { name: 'joinField1', type: 'keyword', userDefined: false },
+      { name: 'joinField2', type: 'double', userDefined: false },
+    ];
+    const result = columnsAfter({} as any, previousColumns, 'JOIN', joinColumns);
+    expect(result).toEqual([...joinColumns, ...previousColumns]);
+  });
+
+  it('overwrites previous columns with the same name', () => {
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'fieldA', type: 'keyword', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ];
+    const joinColumns: ESQLFieldWithMetadata[] = [
+      { name: 'fieldA', type: 'text', userDefined: false },
+      { name: 'fieldC', type: 'double', userDefined: false },
+    ];
+    const result = columnsAfter({} as any, previousColumns, 'JOIN', joinColumns);
+    expect(result).toEqual([
+      { name: 'fieldA', type: 'text', userDefined: false },
+      { name: 'fieldC', type: 'double', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ]);
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts
new file mode 100644
index 0000000000000..371f095e5403c
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts
@@ -0,0 +1,20 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+import { uniqBy } from 'lodash';
+import { type ESQLCommand } from '../../../types';
+import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
+
+export const columnsAfter = (
+  command: ESQLCommand,
+  previousColumns: ESQLColumnData[],
+  query: string,
+  joinColumns?: ESQLFieldWithMetadata[]
+) => {
+  return uniqBy([...(joinColumns ?? []), ...previousColumns], 'name');
+};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/index.ts
index db2fac4bc5556..44d6674dc7b07 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/index.ts
@@ -11,10 +11,12 @@ import type { ICommandMethods } from '../../registry';
 import { autocomplete } from './autocomplete';
 import { validate } from './validate';
 import type { ICommandContext } from '../../types';
+import { columnsAfter } from './columns_after';
 
 const joinCommandMethods: ICommandMethods<ICommandContext> = {
   validate,
   autocomplete,
+  columnsAfter,
 };
 
 export const joinCommand = {
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index a768ba4418302..c7f834068a307 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -6,9 +6,18 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { esqlCommandRegistry, type ESQLAstCommand, type FunctionDefinition } from '@kbn/esql-ast';
-import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types';
-import type { ESQLParamLiteral } from '@kbn/esql-ast/src/types';
+import {
+  esqlCommandRegistry,
+  mutate,
+  synth,
+  type ESQLAstCommand,
+  type FunctionDefinition,
+} from '@kbn/esql-ast';
+import type {
+  ESQLColumnData,
+  ESQLFieldWithMetadata,
+} from '@kbn/esql-ast/src/commands_registry/types';
+import type { ESQLAstQueryExpression, ESQLParamLiteral } from '@kbn/esql-ast/src/types';
 
 import { enrichFieldsWithECSInfo } from '../autocomplete/utils/ecs_metadata_helper';
 import type { ESQLCallbacks } from './types';
@@ -96,15 +105,35 @@ export async function getFieldsFromES(query: string, resourceRetriever?: ESQLCal
 export async function getCurrentQueryAvailableColumns(
   query: string,
   commands: ESQLAstCommand[],
-  previousPipeFields: ESQLColumnData[]
+  previousPipeFields: ESQLColumnData[],
+  fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>
 ) {
   const cacheCopy = new Map<string, ESQLColumnData>();
   previousPipeFields.forEach((field) => cacheCopy.set(field.name, field));
   const lastCommand = commands[commands.length - 1];
   const commandDefinition = esqlCommandRegistry.getCommandByName(lastCommand.name);
 
+  let joinFields;
+  // fetch fields for JOIN here
+  if (lastCommand.name === 'join') {
+    const joinSummary = mutate.commands.join.summarize({
+      type: 'query',
+      commands,
+    } as ESQLAstQueryExpression);
+    const joinIndices = joinSummary.map(({ target: { index } }) => index);
+    if (joinIndices.length) {
+      const joinFieldQuery = synth.cmd`FROM ${joinIndices}`.toString();
+      joinFields = await fetchFields(joinFieldQuery);
+    }
+  }
+
   if (commandDefinition?.methods.columnsAfter) {
-    return commandDefinition.methods.columnsAfter(lastCommand, previousPipeFields, query);
+    return commandDefinition.methods.columnsAfter(
+      lastCommand,
+      previousPipeFields,
+      query,
+      joinFields
+    );
   } else {
     return previousPipeFields;
   }
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index 5ad8ea05f4060..4c91bf3e146e5 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -8,7 +8,10 @@
  */
 
 import { parse } from '@kbn/esql-ast';
-import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types';
+import type {
+  ESQLColumnData,
+  ESQLFieldWithMetadata,
+} from '@kbn/esql-ast/src/commands_registry/types';
 import type { ESQLCallbacks } from './types';
 import { getFieldsFromES, getCurrentQueryAvailableColumns } from './helpers';
 import { removeLastPipe, processPipes, toSingleLine } from './query_string_utils';
@@ -42,7 +45,10 @@ function getValueInsensitive(keyToCheck: string) {
  * for the next time the same query is used.
  * @param queryText
  */
-async function cacheColumnsForQuery(queryText: string) {
+async function cacheColumnsForQuery(
+  queryText: string,
+  fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>
+) {
   const existsInCache = checkCacheInsensitive(queryText);
   if (existsInCache) {
     // this is already in the cache
@@ -56,7 +62,8 @@ async function cacheColumnsForQuery(queryText: string) {
     const availableFields = await getCurrentQueryAvailableColumns(
       queryText,
       root.commands,
-      fieldsAvailableAfterPreviousCommand
+      fieldsAvailableAfterPreviousCommand,
+      fetchFields
     );
     cache.set(queryText, availableFields);
   }
@@ -69,18 +76,22 @@ export function getColumnsByTypeHelper(queryText: string, resourceRetriever?: ES
       return;
     }
 
-    const [sourceCommand, ...partialQueries] = processPipes(queryText);
+    const getFields = async (query: string) => {
+      const cached = getValueInsensitive(query);
+      if (cached) {
+        return cached as ESQLFieldWithMetadata[];
+      }
+      const fields = await getFieldsFromES(query, resourceRetriever);
+      cache.set(query, fields);
+      return fields;
+    };
 
-    // retrieve the index fields from ES ONLY if the source command is not in the cache
-    const existsInCache = getValueInsensitive(sourceCommand);
-    if (!existsInCache) {
-      const fieldsWithMetadata = await getFieldsFromES(sourceCommand, resourceRetriever);
-      cache.set(sourceCommand, fieldsWithMetadata);
-    }
+    const [sourceCommand, ...partialQueries] = processPipes(queryText);
+    getFields(sourceCommand);
 
     // build fields cache for every partial query
     for (const query of partialQueries) {
-      await cacheColumnsForQuery(query);
+      await cacheColumnsForQuery(query, getFields);
     }
   };
 

From 21e1ec2b6f468bfac1e5a03f4caedd3c2824bc7f Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 4 Sep 2025 14:57:25 -0600
Subject: [PATCH 27/54] enrich mostly working

---
 .../kbn-esql-editor/src/esql_editor.tsx       |   1 +
 .../commands/enrich/columns_after.test.ts     | 104 ++++++++++++++++++
 .../commands/enrich/columns_after.ts          |  37 +++++++
 .../commands/enrich/index.ts                  |   2 +
 .../commands/join/columns_after.test.ts       |   8 +-
 .../commands/join/columns_after.ts            |   7 +-
 .../src/commands_registry/registry.ts         |  15 ++-
 .../src/shared/helpers.ts                     |  42 ++++---
 .../src/shared/resources_helpers.ts           |  12 +-
 .../src/validation/helpers.ts                 |   9 +-
 .../src/validation/resources.ts               |  39 +------
 .../src/validation/validation.ts              |  15 +--
 12 files changed, 206 insertions(+), 85 deletions(-)
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts

diff --git a/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx b/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
index 1d8fc029f003e..303ec4d409831 100644
--- a/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
+++ b/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
@@ -505,6 +505,7 @@ const ESQLEditorInternal = function ESQLEditor({
       },
       getPolicies: async () => {
         try {
+          // TODO cache?
           const policies = (await core.http.get(
             `/internal/index_management/enrich_policies`
           )) as SerializedEnrichPolicy[];
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
new file mode 100644
index 0000000000000..9c224f29704d8
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
@@ -0,0 +1,104 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { synth } from '../../../..';
+import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
+import { columnsAfter } from './columns_after';
+
+describe('ENRICH columnsAfter', () => {
+  it('returns previousColumns when no enrich columns', () => {
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'fieldA', type: 'keyword', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ];
+    const result = columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
+      fromJoin: undefined,
+    });
+    expect(result).toEqual(previousColumns);
+  });
+
+  it('adds all enrich columns to when no WITH clause', () => {
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'fieldA', type: 'keyword', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ];
+    const enrichColumns: ESQLFieldWithMetadata[] = [
+      { name: 'enrichField1', type: 'keyword', userDefined: false },
+      { name: 'enrichField2', type: 'double', userDefined: false },
+    ];
+    const result = columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
+      fromEnrich: enrichColumns,
+    });
+    expect(result).toEqual([...enrichColumns, ...previousColumns]);
+  });
+
+  it('adds only declared columns when WITH clause is present', () => {
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'fieldA', type: 'keyword', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ];
+    const enrichColumns: ESQLFieldWithMetadata[] = [
+      { name: 'enrichField1', type: 'keyword', userDefined: false },
+      { name: 'enrichField2', type: 'double', userDefined: false },
+    ];
+    const result = columnsAfter(
+      synth.cmd`ENRICH policy ON matchfield WITH enrichField2`,
+      previousColumns,
+      '',
+      {
+        fromEnrich: enrichColumns,
+      }
+    );
+    expect(result).toEqual([enrichColumns[1], ...previousColumns]);
+  });
+
+  it('renames enrichment fields using WITH', () => {
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'fieldA', type: 'keyword', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ];
+    const enrichColumns: ESQLFieldWithMetadata[] = [
+      { name: 'enrichField1', type: 'keyword', userDefined: false },
+      { name: 'enrichField2', type: 'double', userDefined: false },
+    ];
+    const result = columnsAfter(
+      synth.cmd`ENRICH policy ON matchfield WITH foo = enrichField1, bar = enrichField2`,
+      previousColumns,
+      '',
+      {
+        fromEnrich: enrichColumns,
+      }
+    );
+    const expected = [
+      { name: 'foo', type: 'keyword', userDefined: false },
+      { name: 'bar', type: 'double', userDefined: false },
+      ...previousColumns,
+    ];
+    expect(result).toEqual(expected);
+  });
+
+  it('overwrites previous columns with the same name', () => {
+    const previousColumns: ESQLColumnData[] = [
+      { name: 'fieldA', type: 'keyword', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ];
+    const enrichFields: ESQLFieldWithMetadata[] = [
+      { name: 'fieldA', type: 'text', userDefined: false },
+      { name: 'fieldC', type: 'double', userDefined: false },
+    ];
+    const result = columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
+      fromEnrich: enrichFields,
+    });
+    expect(result).toEqual([
+      { name: 'fieldA', type: 'text', userDefined: false },
+      { name: 'fieldC', type: 'double', userDefined: false },
+      { name: 'fieldB', type: 'long', userDefined: false },
+    ]);
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts
new file mode 100644
index 0000000000000..5c9dbe51adf11
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts
@@ -0,0 +1,37 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+import { uniqBy } from 'lodash';
+import { isAssignment, isColumn, isOptionNode } from '../../../ast/is';
+import type { ESQLCommandOption } from '../../../types';
+import { type ESQLCommand } from '../../../types';
+import type { ESQLColumnData } from '../../types';
+import type { IAdditionalFields } from '../../registry';
+
+export const columnsAfter = (
+  command: ESQLCommand,
+  previousColumns: ESQLColumnData[],
+  query: string,
+  newFields: IAdditionalFields
+) => {
+  const enrichFields = newFields.fromEnrich ?? [];
+  let fieldsToAdd = enrichFields;
+
+  // the with option scopes down the fields that are added
+  const withOption = command.args.find((arg) => isOptionNode(arg) && arg.name === 'with') as
+    | ESQLCommandOption
+    | undefined;
+  if (withOption) {
+    const declaredFields = withOption.args
+      .map((arg) => isAssignment(arg) && isColumn(arg.args[0]) && arg.args[0].parts.join('.'))
+      .filter(Boolean) as string[];
+    fieldsToAdd = enrichFields.filter((field) => declaredFields.includes(field.name));
+  }
+
+  return uniqBy([...fieldsToAdd, ...previousColumns], 'name');
+};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/index.ts
index 79303aa63ee77..230d81d28f831 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/index.ts
@@ -10,11 +10,13 @@ import { i18n } from '@kbn/i18n';
 import type { ICommandMethods } from '../../registry';
 import { autocomplete } from './autocomplete';
 import { validate } from './validate';
+import { columnsAfter } from './columns_after';
 import type { ICommandContext } from '../../types';
 
 const enrichCommandMethods: ICommandMethods<ICommandContext> = {
   validate,
   autocomplete,
+  columnsAfter,
 };
 
 export const enrichCommand = {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
index 0c48f3fa57658..e3d121fd37963 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
@@ -11,12 +11,12 @@ import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('JOIN columnsAfter', () => {
-  it('returns previousColumns when joinColumns is undefined', () => {
+  it('returns previousColumns when no join columns', () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
     ];
-    const result = columnsAfter({} as any, previousColumns, 'JOIN');
+    const result = columnsAfter({} as any, previousColumns, '', { fromJoin: undefined });
     expect(result).toEqual(previousColumns);
   });
 
@@ -29,7 +29,7 @@ describe('JOIN columnsAfter', () => {
       { name: 'joinField1', type: 'keyword', userDefined: false },
       { name: 'joinField2', type: 'double', userDefined: false },
     ];
-    const result = columnsAfter({} as any, previousColumns, 'JOIN', joinColumns);
+    const result = columnsAfter({} as any, previousColumns, '', { fromJoin: joinColumns });
     expect(result).toEqual([...joinColumns, ...previousColumns]);
   });
 
@@ -42,7 +42,7 @@ describe('JOIN columnsAfter', () => {
       { name: 'fieldA', type: 'text', userDefined: false },
       { name: 'fieldC', type: 'double', userDefined: false },
     ];
-    const result = columnsAfter({} as any, previousColumns, 'JOIN', joinColumns);
+    const result = columnsAfter({} as any, previousColumns, '', { fromJoin: joinColumns });
     expect(result).toEqual([
       { name: 'fieldA', type: 'text', userDefined: false },
       { name: 'fieldC', type: 'double', userDefined: false },
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts
index 371f095e5403c..25362d68488f7 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts
@@ -8,13 +8,14 @@
  */
 import { uniqBy } from 'lodash';
 import { type ESQLCommand } from '../../../types';
-import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
+import type { ESQLColumnData } from '../../types';
+import type { IAdditionalFields } from '../../registry';
 
 export const columnsAfter = (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
   query: string,
-  joinColumns?: ESQLFieldWithMetadata[]
+  newFields: IAdditionalFields
 ) => {
-  return uniqBy([...(joinColumns ?? []), ...previousColumns], 'name');
+  return uniqBy([...(newFields.fromJoin ?? []), ...previousColumns], 'name');
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
index 1d41f09a12974..966e2d06e9625 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
@@ -8,7 +8,12 @@
  */
 import type { LicenseType } from '@kbn/licensing-types';
 import type { ESQLMessage, ESQLCommand, ESQLAst } from '../types';
-import type { ISuggestionItem, ICommandCallbacks, ESQLColumnData } from './types';
+import type {
+  ISuggestionItem,
+  ICommandCallbacks,
+  ESQLColumnData,
+  ESQLFieldWithMetadata,
+} from './types';
 
 /**
  * Interface defining the methods that each ES|QL command should register.
@@ -60,7 +65,8 @@ export interface ICommandMethods<TContext = any> {
   columnsAfter?: (
     command: ESQLCommand,
     previousColumns: ESQLColumnData[],
-    query: string
+    query: string,
+    newFields: IAdditionalFields
   ) => ESQLColumnData[];
 }
 
@@ -114,6 +120,11 @@ export interface ICommandRegistry {
   getCommandByName(commandName: string): ICommand | undefined;
 }
 
+export interface IAdditionalFields {
+  fromJoin?: ESQLFieldWithMetadata[];
+  fromEnrich?: ESQLFieldWithMetadata[];
+}
+
 /**
  * Implementation of the ESQL Command Registry.
  * This class manages the registration, storage, and retrieval of ESQL command methods.
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index c7f834068a307..ec2e3eadfaf5b 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -8,6 +8,7 @@
  */
 import {
   esqlCommandRegistry,
+  isSource,
   mutate,
   synth,
   type ESQLAstCommand,
@@ -16,9 +17,11 @@ import {
 import type {
   ESQLColumnData,
   ESQLFieldWithMetadata,
+  ESQLPolicy,
 } from '@kbn/esql-ast/src/commands_registry/types';
 import type { ESQLAstQueryExpression, ESQLParamLiteral } from '@kbn/esql-ast/src/types';
 
+import type { IAdditionalFields } from '@kbn/esql-ast/src/commands_registry/registry';
 import { enrichFieldsWithECSInfo } from '../autocomplete/utils/ecs_metadata_helper';
 import type { ESQLCallbacks } from './types';
 
@@ -106,35 +109,40 @@ export async function getCurrentQueryAvailableColumns(
   query: string,
   commands: ESQLAstCommand[],
   previousPipeFields: ESQLColumnData[],
-  fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>
+  fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>,
+  policies: Map<string, ESQLPolicy>
 ) {
-  const cacheCopy = new Map<string, ESQLColumnData>();
-  previousPipeFields.forEach((field) => cacheCopy.set(field.name, field));
   const lastCommand = commands[commands.length - 1];
-  const commandDefinition = esqlCommandRegistry.getCommandByName(lastCommand.name);
+  const commandDef = esqlCommandRegistry.getCommandByName(lastCommand.name);
+  const extraFields: IAdditionalFields = {};
 
-  let joinFields;
-  // fetch fields for JOIN here
+  // Handle JOIN command: fetch fields from joined indices
   if (lastCommand.name === 'join') {
     const joinSummary = mutate.commands.join.summarize({
       type: 'query',
       commands,
     } as ESQLAstQueryExpression);
     const joinIndices = joinSummary.map(({ target: { index } }) => index);
-    if (joinIndices.length) {
+    if (joinIndices.length > 0) {
       const joinFieldQuery = synth.cmd`FROM ${joinIndices}`.toString();
-      joinFields = await fetchFields(joinFieldQuery);
+      extraFields.fromJoin = await fetchFields(joinFieldQuery);
     }
   }
 
-  if (commandDefinition?.methods.columnsAfter) {
-    return commandDefinition.methods.columnsAfter(
-      lastCommand,
-      previousPipeFields,
-      query,
-      joinFields
-    );
-  } else {
-    return previousPipeFields;
+  // Handle ENRICH command: fetch fields from enrich policy
+  if (lastCommand.name === 'enrich' && isSource(lastCommand.args[0])) {
+    const policyName = lastCommand.args[0].name;
+    const policy = policies.get(policyName);
+    if (policy) {
+      const fieldsQuery = `FROM ${policy.sourceIndices.join(
+        ', '
+      )} | KEEP ${policy.enrichFields.join(', ')}`;
+      extraFields.fromEnrich = await fetchFields(fieldsQuery);
+    }
+  }
+
+  if (commandDef?.methods.columnsAfter) {
+    return commandDef.methods.columnsAfter(lastCommand, previousPipeFields, query, extraFields);
   }
+  return previousPipeFields;
 }
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index 4c91bf3e146e5..3d16a5ede5245 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -11,6 +11,7 @@ import { parse } from '@kbn/esql-ast';
 import type {
   ESQLColumnData,
   ESQLFieldWithMetadata,
+  ESQLPolicy,
 } from '@kbn/esql-ast/src/commands_registry/types';
 import type { ESQLCallbacks } from './types';
 import { getFieldsFromES, getCurrentQueryAvailableColumns } from './helpers';
@@ -47,7 +48,8 @@ function getValueInsensitive(keyToCheck: string) {
  */
 async function cacheColumnsForQuery(
   queryText: string,
-  fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>
+  fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>,
+  policies: Map<string, ESQLPolicy>
 ) {
   const existsInCache = checkCacheInsensitive(queryText);
   if (existsInCache) {
@@ -63,7 +65,8 @@ async function cacheColumnsForQuery(
       queryText,
       root.commands,
       fieldsAvailableAfterPreviousCommand,
-      fetchFields
+      fetchFields,
+      policies
     );
     cache.set(queryText, availableFields);
   }
@@ -89,9 +92,12 @@ export function getColumnsByTypeHelper(queryText: string, resourceRetriever?: ES
     const [sourceCommand, ...partialQueries] = processPipes(queryText);
     getFields(sourceCommand);
 
+    const policies = (await resourceRetriever?.getPolicies?.()) ?? [];
+    const policyMap = new Map(policies.map((p) => [p.name, p]));
+
     // build fields cache for every partial query
     for (const query of partialQueries) {
-      await cacheColumnsForQuery(query, getFields);
+      await cacheColumnsForQuery(query, getFields, policyMap);
     }
   };
 
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
index 61842c3148a7f..a8d745c264a39 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
@@ -7,14 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { type ESQLAst, type ESQLCommand, type FunctionDefinition, Walker } from '@kbn/esql-ast';
-import type { ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types';
-
-export function buildQueryForFieldsInPolicies(policies: ESQLPolicy[]) {
-  return `from ${policies
-    .flatMap(({ sourceIndices }) => sourceIndices)
-    .join(', ')} | keep ${policies.flatMap(({ enrichFields }) => enrichFields).join(', ')}`;
-}
+import { type ESQLCommand, type FunctionDefinition, Walker } from '@kbn/esql-ast';
 
 /**
  * Returns the maximum and minimum number of parameters allowed by a function
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
index 9a0b08ad354b6..777984ca23a3f 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts
@@ -7,16 +7,12 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { isSource, type ESQLCommand } from '@kbn/esql-ast';
-import type { ESQLFieldWithMetadata, ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types';
-import { createMapFromList, nonNullable } from '../shared/helpers';
-import {
-  getColumnsByTypeHelper,
-  getPolicyHelper,
-  getSourcesHelper,
-} from '../shared/resources_helpers';
+import { type ESQLCommand } from '@kbn/esql-ast';
+import type { ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types';
+import { createMapFromList } from '../shared/helpers';
+import { getPolicyHelper, getSourcesHelper } from '../shared/resources_helpers';
 import type { ESQLCallbacks } from '../shared/types';
-import { buildQueryForFieldsInPolicies, getEnrichCommands } from './helpers';
+import { getEnrichCommands } from './helpers';
 
 export async function retrievePolicies(
   commands: ESQLCommand[],
@@ -44,28 +40,3 @@ export async function retrieveSources(
   const sources = await getSourcesHelper(callbacks)();
   return new Set(sources.map(({ name }) => name));
 }
-
-export async function retrievePoliciesFields(
-  commands: ESQLCommand[],
-  policies: Map<string, ESQLPolicy>,
-  callbacks?: ESQLCallbacks
-): Promise<Map<string, ESQLFieldWithMetadata>> {
-  if (!callbacks) {
-    return new Map();
-  }
-  const enrichCommands = getEnrichCommands(commands);
-  if (!enrichCommands.length) {
-    return new Map();
-  }
-  const policyNames = enrichCommands
-    .map(({ args }) => (isSource(args[0]) ? args[0].name : undefined))
-    .filter(nonNullable);
-  if (!policyNames.every((name) => policies.has(name))) {
-    return new Map();
-  }
-
-  const customQuery = buildQueryForFieldsInPolicies(
-    policyNames.map((name) => policies.get(name)) as ESQLPolicy[]
-  );
-  return await getColumnsByTypeHelper(customQuery, callbacks).getColumnMap();
-}
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index ff23d3c19b0b8..5c2d2d6081fc5 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -23,11 +23,7 @@ import type {
 import type { LicenseType } from '@kbn/licensing-types';
 
 import type { ESQLCallbacks } from '../shared/types';
-import {
-  retrievePolicies,
-  // retrievePoliciesFields,
-  retrieveSources,
-} from './resources';
+import { retrievePolicies, retrieveSources } from './resources';
 import type { ReferenceMaps, ValidationOptions, ValidationResult } from './types';
 import { getQueryForFields } from '../autocomplete/helper';
 import { getColumnsByTypeHelper } from '../shared/resources_helpers';
@@ -124,15 +120,6 @@ async function validateAst(
     callbacks?.getJoinIndices?.(),
   ]);
 
-  // if (availablePolicies.size) {
-  //   const fieldsFromPoliciesMap = await retrievePoliciesFields(
-  //     rootCommands,
-  //     availablePolicies,
-  //     callbacks
-  //   );
-  //   fieldsFromPoliciesMap.forEach((value, key) => availableFields.set(key, value));
-  // }
-
   const sourceFields = await getColumnsByTypeHelper(
     queryString.split('|')[0],
     callbacks

From 5c569459c4469493920107dfc4fc5ad848f0a07c Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 4 Sep 2025 16:13:11 -0600
Subject: [PATCH 28/54] support renaming in WITH

---
 .../commands/enrich/columns_after.ts          | 27 ++++++++++++++++---
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts
index 5c9dbe51adf11..1e1cb4bc649c6 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts
@@ -23,14 +23,33 @@ export const columnsAfter = (
   let fieldsToAdd = enrichFields;
 
   // the with option scopes down the fields that are added
+  // and potentially renames things
   const withOption = command.args.find((arg) => isOptionNode(arg) && arg.name === 'with') as
     | ESQLCommandOption
     | undefined;
   if (withOption) {
-    const declaredFields = withOption.args
-      .map((arg) => isAssignment(arg) && isColumn(arg.args[0]) && arg.args[0].parts.join('.'))
-      .filter(Boolean) as string[];
-    fieldsToAdd = enrichFields.filter((field) => declaredFields.includes(field.name));
+    const declaredFieldEntries = withOption.args
+      .map((arg) => {
+        if (
+          isAssignment(arg) &&
+          isColumn(arg.args[0]) &&
+          Array.isArray(arg.args[1]) &&
+          isColumn(arg.args[1][0])
+        ) {
+          return [arg.args[1][0].parts.join('.'), arg.args[0].parts.join('.')];
+        }
+        return undefined;
+      })
+      .filter(Boolean) as [string, string][];
+
+    const declaredFields = new Map<string, string>(declaredFieldEntries);
+
+    fieldsToAdd = enrichFields
+      .filter((field) => declaredFields.has(field.name))
+      .map((field) => {
+        const newName = declaredFields.get(field.name)!;
+        return { ...field, name: newName };
+      });
   }
 
   return uniqBy([...fieldsToAdd, ...previousColumns], 'name');

From 54e813da1cbc83402d06bf22d013e3375ba2fbc8 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Thu, 4 Sep 2025 16:57:15 -0600
Subject: [PATCH 29/54] resolve type issues

---
 .../kbn-esql-editor/src/esql_editor.tsx       |  1 +
 .../commands/fork/columns_after.ts            |  3 +-
 .../commands/fuse/validate.test.ts            |  6 ++--
 .../commands/inlinestats/validate.test.ts     | 28 ++++++++--------
 .../commands/keep/autocomplete.test.ts        | 10 ++----
 .../commands/keep/validate.test.ts            | 30 ++++++++---------
 .../commands/rerank/validate.ts               | 19 ++++-------
 .../commands/sort/validate.test.ts            | 20 +++++------
 .../commands/stats/validate.test.ts           | 33 +++++++++----------
 .../src/__tests__/helpers.ts                  | 20 +++++++----
 .../autocomplete.command.variables.test.ts    | 30 +++++++++++------
 .../src/autocomplete/__tests__/helpers.ts     | 22 ++++++++++---
 .../src/autocomplete/autocomplete.test.ts     | 27 +++++++++++----
 .../src/autocomplete/autocomplete.ts          |  2 +-
 .../utils/ecs_metadata_helper.test.ts         | 18 +++++-----
 .../src/validation/validation.test.ts         |  6 ++--
 .../esql/lib/hover/__tests__/fixtures.ts      |  5 +--
 .../containers/editor/monaco_editor.tsx       |  1 +
 18 files changed, 158 insertions(+), 123 deletions(-)

diff --git a/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx b/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
index 303ec4d409831..8e133f3113716 100644
--- a/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
+++ b/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
@@ -493,6 +493,7 @@ const ESQLEditorInternal = function ESQLEditor({
                   name: c.name,
                   type: c.meta.esType as FieldType,
                   hasConflict: c.meta.type === KBN_FIELD_TYPES.CONFLICT,
+                  userDefined: false,
                 };
               }) || [];
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index 231dfd623d7ca..b8890bac601d0 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -30,7 +30,8 @@ export const columnsAfter = (
         columnsFromBranch = commandDef.methods?.columnsAfter?.(
           branchCommand,
           columnsFromBranch,
-          query
+          query,
+          {} // TODO support nested JOIN commands :scared-hedgie:
         );
       }
     }
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts
index 0480c05cb0aa5..c9fcd81ba3c52 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts
@@ -22,13 +22,13 @@ describe('FUSE Validation', () => {
 
   describe('FUSE', () => {
     test('no errors for valid command', () => {
-      const newFields = new Map(mockContext.fields);
+      const newColumns = new Map(mockContext.columns);
       METADATA_FIELDS.forEach((fieldName) => {
-        newFields.set(fieldName, { name: fieldName, type: 'keyword' });
+        newColumns.set(fieldName, { name: fieldName, type: 'keyword', userDefined: false });
       });
       const context = {
         ...mockContext,
-        fields: newFields,
+        fields: newColumns,
       };
       fuseExpectErrors(
         `FROM index METADATA _id, _score, _index
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/validate.test.ts
index 4b47b7a2f19e1..de9cf10e3d76c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/validate.test.ts
@@ -26,21 +26,19 @@ describe('INLINESTATS Validation', () => {
   });
 
   describe('INLINESTATS <aggregates> [ BY <grouping> ]', () => {
-    const newUserDefinedColumns = new Map(mockContext.userDefinedColumns);
-    newUserDefinedColumns.set('doubleField * 3.281', [
-      {
-        name: 'doubleField * 3.281',
-        type: 'double',
-        location: { min: 0, max: 10 },
-      },
-    ]);
-    newUserDefinedColumns.set('avg_doubleField', [
-      {
-        name: 'avg_doubleField',
-        type: 'double',
-        location: { min: 0, max: 10 },
-      },
-    ]);
+    const newColumns = new Map(mockContext.columns);
+    newColumns.set('doubleField * 3.281', {
+      name: 'doubleField * 3.281',
+      type: 'double',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
+    newColumns.set('avg_doubleField', {
+      name: 'avg_doubleField',
+      type: 'double',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
 
     describe('... <aggregates> ...', () => {
       test('no errors on correct usage', () => {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/autocomplete.test.ts
index 1e71d0718b78d..3427910c9c9b6 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/autocomplete.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/autocomplete.test.ts
@@ -33,15 +33,11 @@ describe('KEEP Autocomplete', () => {
   });
 
   it('suggests available fields after KEEP', async () => {
-    const fieldsMap = mockContext.fields;
-    const userDefinedColumns = mockContext.userDefinedColumns;
-    const allFields = [
-      ...Array.from(fieldsMap.values()),
-      ...Array.from(userDefinedColumns.values()).flat(),
-    ];
+    const columns = mockContext.columns;
+
     keepExpectSuggestions(
       'FROM a | KEEP ',
-      allFields.map((field) => field.name)
+      Array.from(columns.values()).map((column) => column.name)
     );
   });
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/validate.test.ts
index 54f89a7341a80..b9ad1da948ff3 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/validate.test.ts
@@ -20,24 +20,22 @@ describe('KEEP Validation', () => {
   });
 
   test('validates the most basic query', () => {
-    const newUserDefinedColumns = new Map(mockContext.userDefinedColumns);
-    newUserDefinedColumns.set('MIN(doubleField * 10)', [
-      {
-        name: 'MIN(doubleField * 10)',
-        type: 'double',
-        location: { min: 0, max: 10 },
-      },
-    ]);
-    newUserDefinedColumns.set('COUNT(*)', [
-      {
-        name: 'COUNT(*)',
-        type: 'integer',
-        location: { min: 0, max: 10 },
-      },
-    ]);
+    const newColumns = new Map(mockContext.columns);
+    newColumns.set('MIN(doubleField * 10)', {
+      name: 'MIN(doubleField * 10)',
+      type: 'double',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
+    newColumns.set('COUNT(*)', {
+      name: 'COUNT(*)',
+      type: 'integer',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
     const context = {
       ...mockContext,
-      userDefinedColumns: newUserDefinedColumns,
+      columns: newColumns,
     };
     keepExpectErrors('from index | keep keywordField, doubleField, integerField, dateField', []);
     keepExpectErrors(
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts
index 0cbe21feee4dd..a51cbd8ecdbd1 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts
@@ -24,11 +24,7 @@ export const validate = (
   const messages: ESQLMessage[] = [];
 
   const { query, targetField, location, inferenceId } = command as ESQLAstRerankCommand;
-  const rerankExpressionType = getExpressionType(
-    query,
-    context?.fields,
-    context?.userDefinedColumns
-  );
+  const rerankExpressionType = getExpressionType(query, context?.columns);
 
   // check for supported query types
   if (!supportedQueryTypes.includes(rerankExpressionType)) {
@@ -47,13 +43,12 @@ export const validate = (
   const targetName = targetField?.name || 'rerank';
 
   // Sets the target field so the column is recognized after the command is applied
-  context?.userDefinedColumns.set(targetName, [
-    {
-      name: targetName,
-      location: targetField?.location || location,
-      type: 'keyword',
-    },
-  ]);
+  context?.columns.set(targetName, {
+    name: targetName,
+    location: targetField?.location || location,
+    type: 'keyword',
+    userDefined: true,
+  });
 
   messages.push(...validateCommandArguments(command, ast, context, callbacks));
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/validate.test.ts
index 60b0263d96ea3..d41ffc2e48593 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/validate.test.ts
@@ -10,6 +10,7 @@ import { mockContext } from '../../../__tests__/context_fixtures';
 import { validate } from './validate';
 import { expectErrors } from '../../../__tests__/validation';
 import { getNoValidCallSignatureError } from '../../../definitions/utils/validation/utils';
+import type { ICommandContext } from '../../types';
 
 const sortExpectErrors = (query: string, expectedErrors: string[], context = mockContext) => {
   return expectErrors(query, expectedErrors, context, 'sort', validate);
@@ -21,18 +22,17 @@ describe('SORT Validation', () => {
   });
 
   test('validates the most basic query', () => {
-    const newUserDefinedColumns = new Map(mockContext.userDefinedColumns);
+    const newColumns = new Map(mockContext.columns);
 
-    newUserDefinedColumns.set('COUNT(*)', [
-      {
-        name: 'COUNT(*)',
-        type: 'integer',
-        location: { min: 0, max: 10 },
-      },
-    ]);
-    const context = {
+    newColumns.set('COUNT(*)', {
+      name: 'COUNT(*)',
+      type: 'integer',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
+    const context: ICommandContext = {
       ...mockContext,
-      userDefinedColumns: newUserDefinedColumns,
+      columns: newColumns,
     };
     sortExpectErrors('from a_index | sort "field" ', []);
     sortExpectErrors('from a_index | sort wrongField ', ['Unknown column "wrongField"']);
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/validate.test.ts
index 4ff1a9ce9cf1b..75f66b4776f9d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/validate.test.ts
@@ -9,6 +9,7 @@
 import { mockContext } from '../../../__tests__/context_fixtures';
 import { validate } from './validate';
 import { expectErrors } from '../../../__tests__/validation';
+import type { ICommandContext } from '../../types';
 
 const statsExpectErrors = (query: string, expectedErrors: string[], context = mockContext) => {
   return expectErrors(query, expectedErrors, context, 'stats', validate);
@@ -20,24 +21,22 @@ describe('STATS Validation', () => {
   });
 
   describe('STATS <aggregates> [ BY <grouping> ]', () => {
-    const newUserDefinedColumns = new Map(mockContext.userDefinedColumns);
-    newUserDefinedColumns.set('doubleField * 3.281', [
-      {
-        name: 'doubleField * 3.281',
-        type: 'double',
-        location: { min: 0, max: 10 },
-      },
-    ]);
-    newUserDefinedColumns.set('avg_doubleField', [
-      {
-        name: 'avg_doubleField',
-        type: 'double',
-        location: { min: 0, max: 10 },
-      },
-    ]);
-    const context = {
+    const newColumns = new Map(mockContext.columns);
+    newColumns.set('doubleField * 3.281', {
+      name: 'doubleField * 3.281',
+      type: 'double',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
+    newColumns.set('avg_doubleField', {
+      name: 'avg_doubleField',
+      type: 'double',
+      location: { min: 0, max: 10 },
+      userDefined: true,
+    });
+    const context: ICommandContext = {
       ...mockContext,
-      userDefinedColumns: newUserDefinedColumns,
+      columns: newColumns,
     };
     test('no errors on correct usage', () => {
       statsExpectErrors('from a_index | stats by textField', []);
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts
index a18dfbcb50824..0202eb6489f55 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts
@@ -17,23 +17,28 @@ import type { ESQLCallbacks } from '../shared/types';
 export const metadataFields: ESQLFieldWithMetadata[] = METADATA_FIELDS.map((field) => ({
   name: field,
   type: 'keyword',
+  userDefined: false as false,
 }));
 
 export const fields: ESQLFieldWithMetadata[] = [
-  ...fieldTypes.map((type) => ({ name: `${camelCase(type)}Field`, type })),
-  { name: 'any#Char$Field', type: 'double' },
-  { name: 'kubernetes.something.something', type: 'double' },
-  { name: '@timestamp', type: 'date' },
+  ...fieldTypes.map((type) => ({
+    name: `${camelCase(type)}Field`,
+    type,
+    userDefined: false as false,
+  })),
+  { name: 'any#Char$Field', type: 'double', userDefined: false },
+  { name: 'kubernetes.something.something', type: 'double', userDefined: false },
+  { name: '@timestamp', type: 'date', userDefined: false },
 ];
 
 export const enrichFields: ESQLFieldWithMetadata[] = [
-  { name: 'otherField', type: 'text' },
-  { name: 'yetAnotherField', type: 'double' },
+  { name: 'otherField', type: 'text', userDefined: false },
+  { name: 'yetAnotherField', type: 'double', userDefined: false },
 ];
 
 // eslint-disable-next-line @typescript-eslint/naming-convention
 export const unsupported_field: ESQLFieldWithMetadata[] = [
-  { name: 'unsupported_field', type: 'unsupported' },
+  { name: 'unsupported_field', type: 'unsupported', userDefined: false },
 ];
 
 export const indexes = [
@@ -140,6 +145,7 @@ export function getCallbackMocks(): ESQLCallbacks {
           name: 'keywordField',
           type: 'unsupported',
           hasConflict: true,
+          userDefined: false,
         };
         return [field];
       }
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.variables.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.variables.test.ts
index 4fb4366355bb3..dc510dc939336 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.variables.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.variables.test.ts
@@ -24,7 +24,8 @@ describe('autocomplete.suggest', () => {
               type: ESQLVariableType.VALUES,
             },
           ],
-          getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'agent.name', type: 'keyword', userDefined: false }]),
         },
       });
 
@@ -54,7 +55,8 @@ describe('autocomplete.suggest', () => {
         callbacks: {
           canSuggestVariables: () => true,
           getVariables: () => [],
-          getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'clientip', type: 'ip', userDefined: false }]),
         },
       });
 
@@ -81,7 +83,8 @@ describe('autocomplete.suggest', () => {
               type: ESQLVariableType.FUNCTIONS,
             },
           ],
-          getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'clientip', type: 'ip', userDefined: false }]),
         },
       });
 
@@ -102,7 +105,8 @@ describe('autocomplete.suggest', () => {
         callbacks: {
           canSuggestVariables: () => true,
           getVariables: () => [],
-          getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'clientip', type: 'ip', userDefined: false }]),
         },
       });
 
@@ -129,7 +133,8 @@ describe('autocomplete.suggest', () => {
               type: ESQLVariableType.FIELDS,
             },
           ],
-          getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'clientip', type: 'ip', userDefined: false }]),
         },
       });
 
@@ -156,7 +161,8 @@ describe('autocomplete.suggest', () => {
               type: ESQLVariableType.TIME_LITERAL,
             },
           ],
-          getColumnsFor: () => Promise.resolve([{ name: '@timestamp', type: 'date' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: '@timestamp', type: 'date', userDefined: false }]),
         },
       });
 
@@ -177,7 +183,8 @@ describe('autocomplete.suggest', () => {
         callbacks: {
           canSuggestVariables: () => true,
           getVariables: () => [],
-          getColumnsFor: () => Promise.resolve([{ name: 'bytes', type: 'double' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'bytes', type: 'double', userDefined: false }]),
         },
       });
 
@@ -198,7 +205,8 @@ describe('autocomplete.suggest', () => {
         callbacks: {
           canSuggestVariables: () => true,
           getVariables: () => [],
-          getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'agent.name', type: 'keyword', userDefined: false }]),
         },
       });
 
@@ -225,7 +233,8 @@ describe('autocomplete.suggest', () => {
               type: ESQLVariableType.VALUES,
             },
           ],
-          getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'agent.name', type: 'keyword', userDefined: false }]),
         },
       });
 
@@ -246,7 +255,8 @@ describe('autocomplete.suggest', () => {
         callbacks: {
           canSuggestVariables: () => true,
           getVariables: () => [],
-          getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]),
+          getColumnsFor: () =>
+            Promise.resolve([{ name: 'agent.name', type: 'keyword', userDefined: false }]),
         },
       });
 
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts
index c18d4b92e9629..68f91fe9f3b64 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts
@@ -46,15 +46,27 @@ export const fields: TestField[] = [
   ...fieldTypes.map((type) => ({
     name: `${camelCase(type)}Field`,
     type,
+    userDefined: false as false,
+    // suggestedAs is optional and omitted here
   })),
-  { name: 'any#Char$Field', type: 'double', suggestedAs: '`any#Char$Field`' },
-  { name: 'kubernetes.something.something', type: 'double' },
+  {
+    name: 'any#Char$Field',
+    type: 'double',
+    suggestedAs: '`any#Char$Field`',
+    userDefined: false as false,
+  },
+  {
+    name: 'kubernetes.something.something',
+    type: 'double',
+    suggestedAs: undefined,
+    userDefined: false as false,
+  },
 ];
 
 export const lookupIndexFields: TestField[] = [
-  { name: 'booleanField', type: 'boolean' },
-  { name: 'dateField', type: 'date' },
-  { name: 'joinIndexOnlyField', type: 'text' },
+  { name: 'booleanField', type: 'boolean', userDefined: false },
+  { name: 'dateField', type: 'date', userDefined: false },
+  { name: 'joinIndexOnlyField', type: 'text', userDefined: false },
 ];
 
 export const indexes = (
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts
index 343d5ce672cc2..5c68093f0079b 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts
@@ -156,22 +156,27 @@ describe('autocomplete', () => {
             {
               name: 'round(doubleField) + 1',
               type: 'double',
+              userDefined: false,
             },
             {
               name: '`round(doubleField) + 1` + 1',
               type: 'double',
+              userDefined: false,
             },
             {
               name: '```round(doubleField) + 1`` + 1` + 1',
               type: 'double',
+              userDefined: false,
             },
             {
               name: '```````round(doubleField) + 1```` + 1`` + 1` + 1',
               type: 'double',
+              userDefined: false,
             },
             {
               name: '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1` + 1',
               type: 'double',
+              userDefined: false,
             },
           ],
         ]
@@ -184,12 +189,22 @@ describe('autocomplete', () => {
 
         expect(
           await getSuggestions('from a_index | EVAL foo = 1 | KEEP /', {
-            callbacks: { getColumnsFor: () => [...fields, { name: 'foo', type: 'integer' }] },
+            callbacks: {
+              getColumnsFor: () => [
+                ...fields,
+                { name: 'foo', type: 'integer', userDefined: false },
+              ],
+            },
           })
         ).toContain('foo');
         expect(
           await getSuggestions('from a_index | EVAL foo = 1 | KEEP foo, /', {
-            callbacks: { getColumnsFor: () => [...fields, { name: 'foo', type: 'integer' }] },
+            callbacks: {
+              getColumnsFor: () => [
+                ...fields,
+                { name: 'foo', type: 'integer', userDefined: false },
+              ],
+            },
           })
         ).not.toContain('foo');
 
@@ -1091,21 +1106,21 @@ describe('autocomplete', () => {
         'FROM index_a | KEEP field.nam/',
         [{ text: 'field.name', rangeToReplace: { start: 20, end: 29 } }],
         undefined,
-        [[{ name: 'field.name', type: 'double' }]]
+        [[{ name: 'field.name', type: 'double', userDefined: false }]]
       );
       // multi-line
       testSuggestions(
         'FROM index_a\n| KEEP field.nam/',
         [{ text: 'field.name', rangeToReplace: { start: 20, end: 29 } }],
         undefined,
-        [[{ name: 'field.name', type: 'double' }]]
+        [[{ name: 'field.name', type: 'double', userDefined: false }]]
       );
       // triple separator
       testSuggestions(
         'FROM index_c\n| KEEP field.name.f/',
         [{ text: 'field.name.foo', rangeToReplace: { start: 20, end: 32 } }],
         undefined,
-        [[{ name: 'field.name.foo', type: 'double' }]]
+        [[{ name: 'field.name.foo', type: 'double', userDefined: false }]]
       );
       // whitespace — we can't support this case yet because
       // we are relying on string checking instead of the AST :(
@@ -1113,7 +1128,7 @@ describe('autocomplete', () => {
         'FROM index_a | KEEP field . n/',
         [{ text: 'field . name', rangeToReplace: { start: 14, end: 22 } }],
         undefined,
-        [[{ name: 'field.name', type: 'double' }]]
+        [[{ name: 'field.name', type: 'double', userDefined: false }]]
       );
     });
   });
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
index 101f3ac131f45..79a1de540b271 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
@@ -251,7 +251,7 @@ async function getSuggestionsWithinCommandExpression(
 
     const augmentedColumnsMap = new Map(columnMap);
     extraColumnNames.forEach((name) => {
-      augmentedColumnsMap.set(name, { name, type: 'double' });
+      augmentedColumnsMap.set(name, { name, type: 'double', userDefined: false });
     });
     return findNewUserDefinedColumn(augmentedColumnsMap);
   };
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
index bfffbfef723d7..3ded634f8810e 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
@@ -13,9 +13,9 @@ import { type ECSMetadata, enrichFieldsWithECSInfo } from './ecs_metadata_helper
 describe('enrichFieldsWithECSInfo', () => {
   it('should return original columns if fieldsMetadata is not provided', async () => {
     const columns: ESQLFieldWithMetadata[] = [
-      { name: 'ecs.version', type: 'keyword' },
-      { name: 'field1', type: 'text' },
-      { name: 'field2', type: 'double' },
+      { name: 'ecs.version', type: 'keyword', userDefined: false },
+      { name: 'field1', type: 'text', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
     ];
 
     const result = await enrichFieldsWithECSInfo(columns);
@@ -24,9 +24,9 @@ describe('enrichFieldsWithECSInfo', () => {
 
   it('should return columns with metadata if both name and type match with ECS fields', async () => {
     const columns: ESQLFieldWithMetadata[] = [
-      { name: 'ecs.field', type: 'text' },
-      { name: 'ecs.fakeBooleanField', type: 'boolean' },
-      { name: 'field2', type: 'double' },
+      { name: 'ecs.field', type: 'text', userDefined: false },
+      { name: 'ecs.fakeBooleanField', type: 'boolean', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
     ];
     const fieldsMetadata = {
       getClient: jest.fn().mockResolvedValue({
@@ -66,9 +66,9 @@ describe('enrichFieldsWithECSInfo', () => {
 
   it('should handle keyword suffix correctly', async () => {
     const columns: ESQLFieldWithMetadata[] = [
-      { name: 'ecs.version', type: 'keyword' },
-      { name: 'ecs.version.keyword', type: 'keyword' },
-      { name: 'field2', type: 'double' },
+      { name: 'ecs.version', type: 'keyword', userDefined: false },
+      { name: 'ecs.version.keyword', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
     ];
     const fieldsMetadata = {
       getClient: jest.fn().mockResolvedValue({
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts
index 3f581967a9e4f..f26b20eaaf3a4 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts
@@ -178,9 +178,11 @@ describe('validation logic', () => {
           JSON.stringify(
             {
               indexes,
-              fields: fields.concat([{ name: policies[0].matchField, type: 'keyword' }]),
+              fields: fields.concat([
+                { name: policies[0].matchField, type: 'keyword', userDefined: false },
+              ]),
               enrichFields: enrichFields.concat([
-                { name: policies[0].matchField, type: 'keyword' },
+                { name: policies[0].matchField, type: 'keyword', userDefined: false },
               ]),
               policies,
               unsupported_field,
diff --git a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/__tests__/fixtures.ts b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/__tests__/fixtures.ts
index 52ce9bfe8d22a..b9e8cd967c68a 100644
--- a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/__tests__/fixtures.ts
+++ b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/__tests__/fixtures.ts
@@ -18,9 +18,10 @@ const fields: Array<ESQLFieldWithMetadata & { suggestedAs?: string }> = [
   ...types.map((type) => ({
     name: `${type}Field`,
     type,
+    userDefined: false as false,
   })),
-  { name: 'any#Char$Field', type: 'double', suggestedAs: '`any#Char$Field`' },
-  { name: 'kubernetes.something.something', type: 'double' },
+  { name: 'any#Char$Field', type: 'double', suggestedAs: '`any#Char$Field`', userDefined: false },
+  { name: 'kubernetes.something.something', type: 'double', userDefined: false },
 ];
 
 const indexes = (
diff --git a/src/platform/plugins/shared/console/public/application/containers/editor/monaco_editor.tsx b/src/platform/plugins/shared/console/public/application/containers/editor/monaco_editor.tsx
index 2cebacd4dc669..7470122bf9757 100644
--- a/src/platform/plugins/shared/console/public/application/containers/editor/monaco_editor.tsx
+++ b/src/platform/plugins/shared/console/public/application/containers/editor/monaco_editor.tsx
@@ -158,6 +158,7 @@ export const MonacoEditor = ({ localStorageValue, value, setValue }: EditorProps
                   name: c.name,
                   type: c.meta.esType as FieldType,
                   hasConflict: c.meta.type === KBN_FIELD_TYPES.CONFLICT,
+                  userDefined: false,
                 };
               }) || []
             );

From 7bd868c5293d2a8d6a4af3a80291d2c83f8d3b84 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 09:50:02 -0600
Subject: [PATCH 30/54] JOIN and ENRICH within FORK

---
 .../commands/enrich/columns_after.test.ts     |  35 +--
 .../commands/enrich/columns_after.ts          |   6 +-
 .../commands/fork/columns_after.test.ts       | 199 ++++++++++++------
 .../commands/fork/columns_after.ts            |  10 +-
 .../commands/join/columns_after.test.ts       |  27 ++-
 .../commands/join/columns_after.ts            |   7 +-
 .../src/commands_registry/registry.ts         |   6 +-
 .../src/shared/helpers.ts                     |  38 ++--
 8 files changed, 217 insertions(+), 111 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
index 9c224f29704d8..8f2d1d7fa4bd3 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
@@ -12,18 +12,19 @@ import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('ENRICH columnsAfter', () => {
-  it('returns previousColumns when no enrich columns', () => {
+  it('returns previousColumns when no enrich columns', async () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
     ];
-    const result = columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
-      fromJoin: undefined,
+    const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
+      fromJoin: () => Promise.resolve([]),
+      fromEnrich: () => Promise.resolve([]),
     });
     expect(result).toEqual(previousColumns);
   });
 
-  it('adds all enrich columns to when no WITH clause', () => {
+  it('adds all enrich columns to when no WITH clause', async () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
@@ -32,13 +33,14 @@ describe('ENRICH columnsAfter', () => {
       { name: 'enrichField1', type: 'keyword', userDefined: false },
       { name: 'enrichField2', type: 'double', userDefined: false },
     ];
-    const result = columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
-      fromEnrich: enrichColumns,
+    const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
+      fromJoin: () => Promise.resolve([]),
+      fromEnrich: () => Promise.resolve(enrichColumns),
     });
     expect(result).toEqual([...enrichColumns, ...previousColumns]);
   });
 
-  it('adds only declared columns when WITH clause is present', () => {
+  it('adds only declared columns when WITH clause is present', async () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
@@ -47,18 +49,19 @@ describe('ENRICH columnsAfter', () => {
       { name: 'enrichField1', type: 'keyword', userDefined: false },
       { name: 'enrichField2', type: 'double', userDefined: false },
     ];
-    const result = columnsAfter(
+    const result = await columnsAfter(
       synth.cmd`ENRICH policy ON matchfield WITH enrichField2`,
       previousColumns,
       '',
       {
-        fromEnrich: enrichColumns,
+        fromJoin: () => Promise.resolve([]),
+        fromEnrich: () => Promise.resolve(enrichColumns),
       }
     );
     expect(result).toEqual([enrichColumns[1], ...previousColumns]);
   });
 
-  it('renames enrichment fields using WITH', () => {
+  it('renames enrichment fields using WITH', async () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
@@ -67,12 +70,13 @@ describe('ENRICH columnsAfter', () => {
       { name: 'enrichField1', type: 'keyword', userDefined: false },
       { name: 'enrichField2', type: 'double', userDefined: false },
     ];
-    const result = columnsAfter(
+    const result = await columnsAfter(
       synth.cmd`ENRICH policy ON matchfield WITH foo = enrichField1, bar = enrichField2`,
       previousColumns,
       '',
       {
-        fromEnrich: enrichColumns,
+        fromEnrich: () => Promise.resolve(enrichColumns),
+        fromJoin: () => Promise.resolve([]),
       }
     );
     const expected = [
@@ -83,7 +87,7 @@ describe('ENRICH columnsAfter', () => {
     expect(result).toEqual(expected);
   });
 
-  it('overwrites previous columns with the same name', () => {
+  it('overwrites previous columns with the same name', async () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
@@ -92,8 +96,9 @@ describe('ENRICH columnsAfter', () => {
       { name: 'fieldA', type: 'text', userDefined: false },
       { name: 'fieldC', type: 'double', userDefined: false },
     ];
-    const result = columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
-      fromEnrich: enrichFields,
+    const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
+      fromEnrich: () => Promise.resolve(enrichFields),
+      fromJoin: () => Promise.resolve([]),
     });
     expect(result).toEqual([
       { name: 'fieldA', type: 'text', userDefined: false },
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts
index 1e1cb4bc649c6..9a8898a16369b 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts
@@ -13,13 +13,13 @@ import { type ESQLCommand } from '../../../types';
 import type { ESQLColumnData } from '../../types';
 import type { IAdditionalFields } from '../../registry';
 
-export const columnsAfter = (
+export const columnsAfter = async (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
   query: string,
-  newFields: IAdditionalFields
+  additionalFields: IAdditionalFields
 ) => {
-  const enrichFields = newFields.fromEnrich ?? [];
+  const enrichFields = await additionalFields.fromEnrich(command);
   let fieldsToAdd = enrichFields;
 
   // the with option scopes down the fields that are added
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
index 2dec08c24aec5..b790f90aa0b94 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
@@ -11,16 +11,20 @@ import type { ESQLColumnData } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('FORK', () => {
-  it('adds the _fork in the list of fields', () => {
+  it('adds the _fork in the list of fields', async () => {
     const previousCommandFields: ESQLColumnData[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
     ];
 
-    const result = columnsAfter(
+    const result = await columnsAfter(
       synth.cmd`FORK (LIMIT 10 ) (LIMIT 1000 ) `,
       previousCommandFields,
-      ''
+      '',
+      {
+        fromEnrich: () => Promise.resolve([]),
+        fromJoin: () => Promise.resolve([]),
+      }
     );
 
     expect(result).toEqual<ESQLColumnData[]>([
@@ -34,16 +38,20 @@ describe('FORK', () => {
     ]);
   });
 
-  it('collects columns from branches', () => {
+  it('collects columns from branches', async () => {
     const previousCommandFields: ESQLColumnData[] = [
       { name: 'field1', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
     ];
 
-    const result = columnsAfter(
+    const result = await columnsAfter(
       synth.cmd`FORK (EVAL foo = 1 | RENAME foo AS bar) (EVAL lolz = 2 + 3 | EVAL field1 = 2.) `,
       previousCommandFields,
-      ''
+      '',
+      {
+        fromEnrich: () => Promise.resolve([]),
+        fromJoin: () => Promise.resolve([]),
+      }
     );
 
     expect(result).toEqual<ESQLColumnData[]>([
@@ -59,47 +67,41 @@ describe('FORK', () => {
     ]);
   });
 
-  it('prefers userDefined columns over fields with the same name', () => {
+  it('supports JOIN and ENRICH', async () => {
     const previousCommandFields: ESQLColumnData[] = [
-      { name: 'foo', type: 'keyword', userDefined: false },
-      { name: 'bar', type: 'double', userDefined: false },
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
     ];
 
-    // Branch 1: foo is userDefined, bar is not
-    // Branch 2: foo is not userDefined, bar is userDefined
-    const result = columnsAfter(
-      synth.cmd`FORK (EVAL foo = 1) (EVAL bar = 2)`,
+    const result = await columnsAfter(
+      synth.cmd`FORK (LOOKUP JOIN lookup-index ON joinField) (ENRICH policy ON enrichField)`,
       previousCommandFields,
-      ''
-    );
-
-    // foo from branch 1 is userDefined, bar from branch 2 is userDefined
-    expect(result).toEqual<ESQLColumnData[]>([
-      { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
-      { name: 'bar', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      '',
       {
-        name: '_fork',
-        type: 'keyword',
-        userDefined: false,
-      },
-    ]);
-  });
-
-  it('keeps the first userDefined column if both branches define userDefined columns with the same name', () => {
-    const previousCommandFields: ESQLColumnData[] = [
-      { name: 'foo', type: 'keyword', userDefined: false },
-    ];
-
-    // Both branches define foo as userDefined, but with different types
-    const result = columnsAfter(
-      synth.cmd`FORK (EVAL foo = 1) (EVAL foo = 2.5)`,
-      previousCommandFields,
-      ''
+        fromEnrich: () =>
+          Promise.resolve([
+            {
+              name: 'from-enrich',
+              type: 'keyword',
+              userDefined: false,
+            },
+          ]),
+        fromJoin: () =>
+          Promise.resolve([
+            {
+              name: 'from-join',
+              type: 'keyword',
+              userDefined: false,
+            },
+          ]),
+      }
     );
 
-    // The first userDefined column wins (type: integer)
     expect(result).toEqual<ESQLColumnData[]>([
-      { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      { name: 'from-join', type: 'keyword', userDefined: false },
+      { name: 'field1', type: 'keyword', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
+      { name: 'from-enrich', type: 'keyword', userDefined: false },
       {
         name: '_fork',
         type: 'keyword',
@@ -108,42 +110,28 @@ describe('FORK', () => {
     ]);
   });
 
-  it("doesn't duplicate fields from branches", () => {
+  it('prefers userDefined columns over fields with the same name', async () => {
     const previousCommandFields: ESQLColumnData[] = [
       { name: 'foo', type: 'keyword', userDefined: false },
+      { name: 'bar', type: 'double', userDefined: false },
     ];
 
-    // All branches keep foo as a field, but with different types
-    const result = columnsAfter(
-      synth.cmd`FORK (KEEP foo) (KEEP foo) (KEEP foo)`,
+    // Branch 1: foo is userDefined, bar is not
+    // Branch 2: foo is not userDefined, bar is userDefined
+    const result = await columnsAfter(
+      synth.cmd`FORK (EVAL foo = 1) (EVAL bar = 2)`,
       previousCommandFields,
-      ''
-    );
-
-    expect(result).toEqual<ESQLColumnData[]>([
-      { name: 'foo', type: 'keyword', userDefined: false },
+      '',
       {
-        name: '_fork',
-        type: 'keyword',
-        userDefined: false,
-      },
-    ]);
-  });
-
-  it('prefers userDefined column if one branch overwrites a field', () => {
-    const previousCommandFields: ESQLColumnData[] = [
-      { name: 'foo', type: 'keyword', userDefined: false },
-    ];
-
-    // Branch 1: foo is userDefined, Branch 2: foo is not userDefined
-    const result = columnsAfter(
-      synth.cmd`FORK (EVAL foo = 1) (LIMIT 1)`,
-      previousCommandFields,
-      ''
+        fromEnrich: () => Promise.resolve([]),
+        fromJoin: () => Promise.resolve([]),
+      }
     );
 
+    // foo from branch 1 is userDefined, bar from branch 2 is userDefined
     expect(result).toEqual<ESQLColumnData[]>([
       { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+      { name: 'bar', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
       {
         name: '_fork',
         type: 'keyword',
@@ -151,4 +139,85 @@ describe('FORK', () => {
       },
     ]);
   });
+
+  describe('conflicts between branches', () => {
+    it('keeps the first userDefined column if both branches define userDefined columns with the same name', async () => {
+      const previousCommandFields: ESQLColumnData[] = [
+        { name: 'foo', type: 'keyword', userDefined: false },
+      ];
+
+      // Both branches define foo as userDefined, but with different types
+      const result = await columnsAfter(
+        synth.cmd`FORK (EVAL foo = 1) (EVAL foo = 2.5)`,
+        previousCommandFields,
+        '',
+        {
+          fromEnrich: () => Promise.resolve([]),
+          fromJoin: () => Promise.resolve([]),
+        }
+      );
+
+      // The first userDefined column wins (type: integer)
+      expect(result).toEqual<ESQLColumnData[]>([
+        { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+        {
+          name: '_fork',
+          type: 'keyword',
+          userDefined: false,
+        },
+      ]);
+    });
+
+    it("doesn't duplicate fields from branches", async () => {
+      const previousCommandFields: ESQLColumnData[] = [
+        { name: 'foo', type: 'keyword', userDefined: false },
+      ];
+
+      // All branches keep foo as a field, but with different types
+      const result = await columnsAfter(
+        synth.cmd`FORK (KEEP foo) (KEEP foo) (KEEP foo)`,
+        previousCommandFields,
+        '',
+        {
+          fromEnrich: () => Promise.resolve([]),
+          fromJoin: () => Promise.resolve([]),
+        }
+      );
+
+      expect(result).toEqual<ESQLColumnData[]>([
+        { name: 'foo', type: 'keyword', userDefined: false },
+        {
+          name: '_fork',
+          type: 'keyword',
+          userDefined: false,
+        },
+      ]);
+    });
+
+    it('prefers userDefined column if one branch overwrites a field', async () => {
+      const previousCommandFields: ESQLColumnData[] = [
+        { name: 'foo', type: 'keyword', userDefined: false },
+      ];
+
+      // Branch 1: foo is userDefined, Branch 2: foo is not userDefined
+      const result = await columnsAfter(
+        synth.cmd`FORK (EVAL foo = 1) (LIMIT 1)`,
+        previousCommandFields,
+        '',
+        {
+          fromEnrich: () => Promise.resolve([]),
+          fromJoin: () => Promise.resolve([]),
+        }
+      );
+
+      expect(result).toEqual<ESQLColumnData[]>([
+        { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } },
+        {
+          name: '_fork',
+          type: 'keyword',
+          userDefined: false,
+        },
+      ]);
+    });
+  });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index b8890bac601d0..0f22d746d23fd 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -11,11 +11,13 @@ import { esqlCommandRegistry } from '../../../..';
 import type { ESQLAstQueryExpression } from '../../../types';
 import { type ESQLCommand } from '../../../types';
 import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
+import type { IAdditionalFields } from '../../registry';
 
-export const columnsAfter = (
+export const columnsAfter = async (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
-  query: string
+  query: string,
+  additionalFields: IAdditionalFields
 ) => {
   const branches = command.args as ESQLAstQueryExpression[];
 
@@ -27,11 +29,11 @@ export const columnsAfter = (
     for (const branchCommand of branch.commands) {
       const commandDef = esqlCommandRegistry.getCommandByName(branchCommand.name);
       if (commandDef?.methods?.columnsAfter) {
-        columnsFromBranch = commandDef.methods?.columnsAfter?.(
+        columnsFromBranch = await commandDef.methods?.columnsAfter?.(
           branchCommand,
           columnsFromBranch,
           query,
-          {} // TODO support nested JOIN commands :scared-hedgie:
+          additionalFields
         );
       }
     }
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
index e3d121fd37963..1aeeef826a723 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
@@ -11,16 +11,21 @@ import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types';
 import { columnsAfter } from './columns_after';
 
 describe('JOIN columnsAfter', () => {
-  it('returns previousColumns when no join columns', () => {
+  it('returns previousColumns when no join columns', async () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
     ];
-    const result = columnsAfter({} as any, previousColumns, '', { fromJoin: undefined });
+
+    const result = await columnsAfter({} as any, previousColumns, '', {
+      fromJoin: () => Promise.resolve([]),
+      fromEnrich: () => Promise.resolve([]),
+    });
+
     expect(result).toEqual(previousColumns);
   });
 
-  it('prepends joinColumns to previousColumns when joinColumns is provided', () => {
+  it('prepends joinColumns to previousColumns when joinColumns is provided', async () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
@@ -29,11 +34,16 @@ describe('JOIN columnsAfter', () => {
       { name: 'joinField1', type: 'keyword', userDefined: false },
       { name: 'joinField2', type: 'double', userDefined: false },
     ];
-    const result = columnsAfter({} as any, previousColumns, '', { fromJoin: joinColumns });
+
+    const result = await columnsAfter({} as any, previousColumns, '', {
+      fromJoin: () => Promise.resolve(joinColumns),
+      fromEnrich: () => Promise.resolve([]),
+    });
+
     expect(result).toEqual([...joinColumns, ...previousColumns]);
   });
 
-  it('overwrites previous columns with the same name', () => {
+  it('overwrites previous columns with the same name', async () => {
     const previousColumns: ESQLColumnData[] = [
       { name: 'fieldA', type: 'keyword', userDefined: false },
       { name: 'fieldB', type: 'long', userDefined: false },
@@ -42,7 +52,12 @@ describe('JOIN columnsAfter', () => {
       { name: 'fieldA', type: 'text', userDefined: false },
       { name: 'fieldC', type: 'double', userDefined: false },
     ];
-    const result = columnsAfter({} as any, previousColumns, '', { fromJoin: joinColumns });
+
+    const result = await columnsAfter({} as any, previousColumns, '', {
+      fromJoin: () => Promise.resolve(joinColumns),
+      fromEnrich: () => Promise.resolve([]),
+    });
+
     expect(result).toEqual([
       { name: 'fieldA', type: 'text', userDefined: false },
       { name: 'fieldC', type: 'double', userDefined: false },
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts
index 25362d68488f7..b437293906aae 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts
@@ -11,11 +11,12 @@ import { type ESQLCommand } from '../../../types';
 import type { ESQLColumnData } from '../../types';
 import type { IAdditionalFields } from '../../registry';
 
-export const columnsAfter = (
+export const columnsAfter = async (
   command: ESQLCommand,
   previousColumns: ESQLColumnData[],
   query: string,
-  newFields: IAdditionalFields
+  additionalFields: IAdditionalFields
 ) => {
-  return uniqBy([...(newFields.fromJoin ?? []), ...previousColumns], 'name');
+  const joinFields = await additionalFields.fromJoin(command);
+  return uniqBy([...(joinFields ?? []), ...previousColumns], 'name');
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
index 966e2d06e9625..6329572156d50 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
@@ -67,7 +67,7 @@ export interface ICommandMethods<TContext = any> {
     previousColumns: ESQLColumnData[],
     query: string,
     newFields: IAdditionalFields
-  ) => ESQLColumnData[];
+  ) => Promise<ESQLColumnData[]> | ESQLColumnData[];
 }
 
 export interface ICommandMetadata {
@@ -121,8 +121,8 @@ export interface ICommandRegistry {
 }
 
 export interface IAdditionalFields {
-  fromJoin?: ESQLFieldWithMetadata[];
-  fromEnrich?: ESQLFieldWithMetadata[];
+  fromJoin: (cmd: ESQLCommand) => Promise<ESQLFieldWithMetadata[]>;
+  fromEnrich: (cmd: ESQLCommand) => Promise<ESQLFieldWithMetadata[]>;
 }
 
 /**
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index ec2e3eadfaf5b..8bb49d1e386f6 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -114,35 +114,49 @@ export async function getCurrentQueryAvailableColumns(
 ) {
   const lastCommand = commands[commands.length - 1];
   const commandDef = esqlCommandRegistry.getCommandByName(lastCommand.name);
-  const extraFields: IAdditionalFields = {};
 
-  // Handle JOIN command: fetch fields from joined indices
-  if (lastCommand.name === 'join') {
+  const getJoinFields = (command: ESQLAstCommand): Promise<ESQLFieldWithMetadata[]> => {
     const joinSummary = mutate.commands.join.summarize({
       type: 'query',
-      commands,
+      commands: [command],
     } as ESQLAstQueryExpression);
     const joinIndices = joinSummary.map(({ target: { index } }) => index);
     if (joinIndices.length > 0) {
       const joinFieldQuery = synth.cmd`FROM ${joinIndices}`.toString();
-      extraFields.fromJoin = await fetchFields(joinFieldQuery);
+      return fetchFields(joinFieldQuery);
+    }
+    return Promise.resolve([]);
+  };
+
+  const getEnrichFields = (command: ESQLAstCommand): Promise<ESQLFieldWithMetadata[]> => {
+    if (!isSource(command.args[0])) {
+      return Promise.resolve([]);
     }
-  }
 
-  // Handle ENRICH command: fetch fields from enrich policy
-  if (lastCommand.name === 'enrich' && isSource(lastCommand.args[0])) {
-    const policyName = lastCommand.args[0].name;
+    const policyName = command.args[0].name;
     const policy = policies.get(policyName);
     if (policy) {
       const fieldsQuery = `FROM ${policy.sourceIndices.join(
         ', '
       )} | KEEP ${policy.enrichFields.join(', ')}`;
-      extraFields.fromEnrich = await fetchFields(fieldsQuery);
+      return fetchFields(fieldsQuery);
     }
-  }
+
+    return Promise.resolve([]);
+  };
+
+  const additionalFields: IAdditionalFields = {
+    fromJoin: getJoinFields,
+    fromEnrich: getEnrichFields,
+  };
 
   if (commandDef?.methods.columnsAfter) {
-    return commandDef.methods.columnsAfter(lastCommand, previousPipeFields, query, extraFields);
+    return commandDef.methods.columnsAfter(
+      lastCommand,
+      previousPipeFields,
+      query,
+      additionalFields
+    );
   }
   return previousPipeFields;
 }

From 4b36c064ebbaefe41f468097c133049918675e89 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 10:19:51 -0600
Subject: [PATCH 31/54] Fix enrich validation

---
 .../commands/enrich/validate.test.ts          | 89 ++++++++-----------
 .../commands/enrich/validate.ts               | 44 +++++++--
 .../commands/fork/columns_after.test.ts       |  2 +-
 .../definitions/utils/validation/function.ts  | 11 +--
 4 files changed, 81 insertions(+), 65 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.test.ts
index c402abea064f7..cdf34554d19f6 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.test.ts
@@ -6,6 +6,7 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
+
 import { mockContext } from '../../../__tests__/context_fixtures';
 import { validate } from './validate';
 
@@ -22,79 +23,67 @@ describe('ENRICH Validation', () => {
   });
 
   test('validates the most basic query', () => {
-    enrichExpectErrors(`from a_index | enrich policy `, []);
-    enrichExpectErrors(`from a_index | enrich policy on textField with var0 = doubleField `, []);
-    enrichExpectErrors(
-      `from a_index | enrich policy on textField with var0 = doubleField, textField `,
-      []
-    );
-    enrichExpectErrors(
-      `from a_index | enrich policy on textField with col0 = doubleField, var0 = textField`,
-      []
-    );
-    enrichExpectErrors(`from a_index | enrich policy with doubleField`, []);
-    enrichExpectErrors(`from a_index | enrich policy | eval doubleField`, []);
-    enrichExpectErrors(`from a_index | enrich policy with col0 = doubleField | eval col0`, []);
+    enrichExpectErrors(`FROM a_index | ENRICH policy `, []);
   });
+
   test('validates the coordinators', () => {
     for (const value of ['any', 'coordinator', 'remote']) {
-      enrichExpectErrors(`from a_index | enrich _${value}:policy `, []);
-      enrichExpectErrors(`from a_index | enrich _${value} :  policy `, [
+      enrichExpectErrors(`FROM a_index | enrich _${value}:policy `, []);
+      enrichExpectErrors(`FROM a_index | enrich _${value} :  policy `, [
         `Unknown policy "_${value}"`,
       ]);
-      enrichExpectErrors(`from a_index | enrich _${value}:  policy `, [
+      enrichExpectErrors(`FROM a_index | enrich _${value}:  policy `, [
         `Unknown policy "_${value}"`,
       ]);
-      enrichExpectErrors(`from a_index | enrich _${camelCase(value)}:policy `, []);
-      enrichExpectErrors(`from a_index | enrich _${value.toUpperCase()}:policy `, []);
+      enrichExpectErrors(`FROM a_index | enrich _${camelCase(value)}:policy `, []);
+      enrichExpectErrors(`FROM a_index | enrich _${value.toUpperCase()}:policy `, []);
     }
 
-    enrichExpectErrors(`from a_index | enrich _unknown:policy`, [
+    enrichExpectErrors(`FROM a_index | enrich _unknown:policy`, [
       'Unrecognized value "_unknown" for ENRICH, mode needs to be one of [_any, _coordinator, _remote]',
     ]);
-    enrichExpectErrors(`from a_index | enrich any:policy`, [
+    enrichExpectErrors(`FROM a_index | enrich any:policy`, [
       'Unrecognized value "any" for ENRICH, mode needs to be one of [_any, _coordinator, _remote]',
     ]);
   });
   test('raises error on unknown policy', () => {
-    enrichExpectErrors(`from a_index | enrich _`, ['Unknown policy "_"']);
-    enrichExpectErrors(`from a_index | enrich _:`, ['Unknown policy "_"']);
-    enrichExpectErrors(`from a_index | enrich any:`, ['Unknown policy "any"']);
-    enrichExpectErrors(`from a_index | enrich _:policy`, [
+    enrichExpectErrors(`FROM a_index | enrich _`, ['Unknown policy "_"']);
+    enrichExpectErrors(`FROM a_index | enrich _:`, ['Unknown policy "_"']);
+    enrichExpectErrors(`FROM a_index | enrich any:`, ['Unknown policy "any"']);
+    enrichExpectErrors(`FROM a_index | enrich _:policy`, [
       'Unrecognized value "_" for ENRICH, mode needs to be one of [_any, _coordinator, _remote]',
     ]);
-    enrichExpectErrors(`from a_index | enrich _any:`, ['Unknown policy "_any"']);
-    enrichExpectErrors('from a_index | enrich `this``is fine`', ['Unknown policy "`this``is"']);
-    enrichExpectErrors('from a_index | enrich this is fine', ['Unknown policy "this"']);
-    enrichExpectErrors(`from a_index |enrich missing-policy `, ['Unknown policy "missing-policy"']);
+    enrichExpectErrors(`FROM a_index | enrich _any:`, ['Unknown policy "_any"']);
+    enrichExpectErrors('FROM a_index | enrich `this``is fine`', ['Unknown policy "`this``is"']);
+    enrichExpectErrors('FROM a_index | enrich this is fine', ['Unknown policy "this"']);
+    enrichExpectErrors(`FROM a_index |enrich missing-policy `, ['Unknown policy "missing-policy"']);
   });
 
-  test('validates the columns', () => {
-    enrichExpectErrors(`from a_index | enrich policy on b `, ['Unknown column "b"']);
+  test('validates match field', () => {
+    enrichExpectErrors(`FROM a_index | ENRICH policy ON b `, ['Unknown column "b"']);
 
-    enrichExpectErrors('from a_index | enrich policy on `this``is fine`', [
+    enrichExpectErrors('FROM a_index | ENRICH policy ON `this``is fine`', [
       'Unknown column "this`is fine"',
     ]);
-    enrichExpectErrors('from a_index | enrich policy on this is fine', ['Unknown column "this"']);
-    enrichExpectErrors(`from a_index | enrich policy on textField with col1 `, [
-      'Unknown column "col1"',
-    ]);
-    enrichExpectErrors(`from a_index |enrich policy on doubleField with col1 = `, [
-      'Unknown column "col1"',
-    ]);
-    enrichExpectErrors(`from a_index | enrich policy on textField with col1 = c `, [
-      'Unknown column "col1"',
-      'Unknown column "c"',
-    ]);
-    enrichExpectErrors(`from a_index |enrich policy on doubleField with col1 = , `, [
-      'Unknown column "col1"',
-    ]);
-    enrichExpectErrors(`from a_index | enrich policy on textField with var0 = doubleField, col1 `, [
-      'Unknown column "col1"',
-    ]);
+    enrichExpectErrors('FROM a_index | ENRICH policy ON this is fine', ['Unknown column "this"']);
+  });
+
+  test('validates enrich fields against policy', () => {
+    // otherField and yetAnotherField are enrich fields in the policy
+    enrichExpectErrors(`FROM a_index | ENRICH policy ON textField WITH otherField `, []);
     enrichExpectErrors(
-      `from a_index |enrich policy on doubleField with col0 = doubleField, col1 = `,
-      ['Unknown column "col1"']
+      `FROM a_index | ENRICH policy ON textField WITH otherField, newname = yetAnotherField `,
+      []
     );
+
+    // keywordField is a field in the index, but not an enrich field
+    enrichExpectErrors(`FROM a_index | ENRICH policy ON textField WITH keywordField `, [
+      'Unknown column "keywordField"',
+    ]);
+
+    // col1 shouldn't be validated, only foo
+    enrichExpectErrors(`FROM a_index |ENRICH policy ON doubleField WITH col1 = foo`, [
+      'Unknown column "foo"',
+    ]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.ts
index 2fc375fe8bdb1..a0fe9df7c9d2c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.ts
@@ -6,11 +6,18 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { getMessageFromId } from '../../../definitions/utils/errors';
-import type { ESQLSource, ESQLCommand, ESQLMessage, ESQLAst } from '../../../types';
-import { ENRICH_MODES } from './util';
-import type { ESQLPolicy, ICommandContext, ICommandCallbacks } from '../../types';
+import { isAssignment, isColumn, isOptionNode } from '../../../ast/is';
+import { errors, getMessageFromId } from '../../../definitions/utils/errors';
 import { validateCommandArguments } from '../../../definitions/utils/validation';
+import type {
+  ESQLAst,
+  ESQLCommand,
+  ESQLCommandOption,
+  ESQLMessage,
+  ESQLSource,
+} from '../../../types';
+import type { ESQLPolicy, ICommandCallbacks, ICommandContext } from '../../types';
+import { ENRICH_MODES } from './util';
 
 export const validate = (
   command: ESQLCommand,
@@ -53,7 +60,34 @@ export const validate = (
     }
   }
 
-  messages.push(...validateCommandArguments(command, ast, context, callbacks));
+  const policy = index && policies.get(index.valueUnquoted);
+  const withOption = command.args.find(
+    (arg) => isOptionNode(arg) && arg.name === 'with'
+  ) as ESQLCommandOption;
+
+  if (withOption && policy) {
+    withOption.args.forEach((arg) => {
+      if (isAssignment(arg) && Array.isArray(arg.args[1]) && isColumn(arg.args[1][0])) {
+        const column = arg.args[1][0];
+        if (!policy.enrichFields.includes(column.parts.join('.'))) {
+          messages.push(errors.unknownColumn(column));
+        }
+      }
+    });
+  }
+
+  messages.push(
+    ...validateCommandArguments(
+      {
+        ...command,
+        // exclude WITH from generic validation since it shouldn't be compared against the generic column list
+        args: command.args.filter((arg) => arg !== withOption),
+      },
+      ast,
+      context,
+      callbacks
+    )
+  );
 
   return messages;
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
index b790f90aa0b94..c705e048ae65f 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
@@ -74,7 +74,7 @@ describe('FORK', () => {
     ];
 
     const result = await columnsAfter(
-      synth.cmd`FORK (LOOKUP JOIN lookup-index ON joinField) (ENRICH policy ON enrichField)`,
+      synth.cmd`FORK (LOOKUP JOIN lookup-index ON joinField) (ENRICH policy ON matchField)`,
       previousCommandFields,
       '',
       {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
index 57ad7ec883501..787edbe01228f 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts
@@ -7,7 +7,6 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import type { LicenseType } from '@kbn/licensing-types';
-import { uniqBy } from 'lodash';
 import { errors, getFunctionDefinition } from '..';
 import { FunctionDefinitionTypes, within } from '../../../..';
 import {
@@ -35,8 +34,8 @@ import {
   getSignaturesWithMatchingArity,
   matchesArity,
 } from '../expressions';
-import { ColumnValidator } from './column';
 import { isArrayType } from '../operators';
+import { ColumnValidator } from './column';
 
 export function validateFunction({
   fn,
@@ -170,13 +169,7 @@ class FunctionValidator {
       return new ColumnValidator(arg, this.context, this.parentCommand.name).validate();
     });
 
-    // uniqBy is used to cover a special case in ENRICH where an implicit assignment is possible
-    // so the AST actually stores an explicit "columnX = columnX" which duplicates the message
-    //
-    // @TODO - we will no longer need to store an assignment in the AST approach when we
-    // align field availability detection with the system used by autocomplete
-    // (start using columnsAfter instead of collectUserDefinedColumns)
-    this.report(...uniqBy(columnMessages, ({ location }) => `${location.min}-${location.max}`));
+    this.report(...columnMessages);
   }
 
   /**

From 3172e2fc97866e0760ab4f20a5610781e7b64363 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 10:23:50 -0600
Subject: [PATCH 32/54] fix change_point validation

---
 .../commands/change_point/validate.ts         | 27 +++++++------------
 1 file changed, 9 insertions(+), 18 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
index dd2dd9870637f..3d378c7ff780d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts
@@ -54,24 +54,15 @@ export const validate = (
     });
   }
 
-  // validate AS
-  const asArg = command.args.find((arg) => isOptionNode(arg) && arg.name === 'as');
-  if (asArg && isOptionNode(asArg)) {
-    // populate userDefinedColumns references to prevent the common check from failing with unknown column
-    asArg.args.forEach((arg, index) => {
-      if (isColumn(arg)) {
-        // TODO - can we remove this?
-        context?.columns.set(arg.name, {
-          name: arg.name,
-          location: arg.location,
-          type: index === 0 ? 'keyword' : 'long',
-          userDefined: true,
-        });
-      }
-    });
-  }
-
-  messages.push(...validateCommandArguments(command, ast, context, callbacks));
+  messages.push(
+    ...validateCommandArguments(
+      // exclude AS option from generic validation
+      { ...command, args: command.args.filter((arg) => !(isOptionNode(arg) && arg.name === 'as')) },
+      ast,
+      context,
+      callbacks
+    )
+  );
 
   return messages;
 };

From 21e54530068a73286a019f0b530d16d47bc6116c Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 10:43:11 -0600
Subject: [PATCH 33/54] dedupe eval expansion

---
 .../src/autocomplete/helper.ts                | 16 +++-----
 .../src/shared/expand_evals.ts                | 41 +++++++++++++++++++
 .../src/validation/validation.ts              | 32 +--------------
 3 files changed, 48 insertions(+), 41 deletions(-)
 create mode 100644 src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/expand_evals.ts

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
index f3bebe1c5d11f..91f289a88db19 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
@@ -8,7 +8,8 @@
  */
 
 import type { ESQLAstQueryExpression } from '@kbn/esql-ast';
-import { type ESQLCommand, BasicPrettyPrinter, Builder, EDITOR_MARKER } from '@kbn/esql-ast';
+import { BasicPrettyPrinter, EDITOR_MARKER } from '@kbn/esql-ast';
+import { expandEvals } from '../shared/expand_evals';
 
 /**
  * This function is used to build the query that will be used to compute the
@@ -25,11 +26,11 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
   const lastCommand = commands[commands.length - 1];
   if (lastCommand && lastCommand.name === 'fork' && lastCommand.args.length > 0) {
     /**
-     * This translates the current fork command branch into a simpler but equivalent
+     * This flattens the current fork branch into a simpler but equivalent
      * query that is compatible with the existing field computation/caching strategy.
      *
      * The intuition here is that if the cursor is within a fork branch, the
-     * previous context is equivalent to a query without the FORK command.:
+     * previous context is equivalent to a query without the FORK command:
      *
      * Original query: FROM lolz | EVAL foo = 1 | FORK (EVAL bar = 2) (EVAL baz = 3 | WHERE /)
      * Simplified:     FROM lolz | EVAL foo = 1 | EVAL baz = 3
@@ -54,13 +55,8 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
        * Original query: FROM lolz | EVAL foo = 1, bar = foo + 1, baz = bar + 1, /
        * Simplified:     FROM lolz | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1
        */
-      const individualEvalCommands: ESQLCommand[] = [];
-      for (const expression of lastCommand.args) {
-        individualEvalCommands.push(Builder.command({ name: 'eval', args: [expression] }));
-      }
-      const newCommands = commands
-        .slice(0, -1)
-        .concat(endsWithComma ? individualEvalCommands : individualEvalCommands.slice(0, -1));
+      const expanded = expandEvals(commands);
+      const newCommands = expanded.slice(0, endsWithComma ? undefined : -1);
       return BasicPrettyPrinter.print({ ...root, commands: newCommands });
     }
   }
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/expand_evals.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/expand_evals.ts
new file mode 100644
index 0000000000000..07338095d16ef
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/expand_evals.ts
@@ -0,0 +1,41 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type { ESQLCommand } from '@kbn/esql-ast';
+import { Builder } from '@kbn/esql-ast';
+
+/**
+ * Expands EVAL commands into separate single-expression EVAL commands.
+ *
+ * E.g. EVAL foo = 1, bar = 2 => [EVAL foo = 1, EVAL bar = 2]
+ *
+ * This is logically equivalent and makes validation and field existence detection much easier.
+ *
+ * @param commands The list of commands to expand.
+ * @returns The expanded list of commands.
+ */
+export function expandEvals(commands: ESQLCommand[]): ESQLCommand[] {
+  const expanded: ESQLCommand[] = [];
+  for (const command of commands) {
+    if (command.name.toLowerCase() === 'eval') {
+      for (const arg of command.args) {
+        expanded.push(
+          Builder.command({
+            name: 'eval',
+            args: [arg],
+            location: command.location,
+          })
+        );
+      }
+    } else {
+      expanded.push(command);
+    }
+  }
+  return expanded;
+}
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index 5c2d2d6081fc5..2412b850cd99c 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -27,6 +27,7 @@ import { retrievePolicies, retrieveSources } from './resources';
 import type { ReferenceMaps, ValidationOptions, ValidationResult } from './types';
 import { getQueryForFields } from '../autocomplete/helper';
 import { getColumnsByTypeHelper } from '../shared/resources_helpers';
+import { expandEvals } from '../shared/expand_evals';
 
 /**
  * ES|QL validation public API
@@ -290,37 +291,6 @@ function getSubqueriesToValidate(rootCommands: ESQLCommand[]) {
   return subsequences.map((subsequence) => Builder.expression.query(subsequence));
 }
 
-/**
- * Expands EVAL commands into separate commands for each expression.
- *
- * E.g. `EVAL 1 + 2, 3 + 4` becomes `EVAL 1 + 2 | EVAL 3 + 4`
- *
- * This is logically equivalent and makes validation and field existence detection much easier.
- *
- * @param commands The list of commands to expand.
- * @returns The expanded list of commands.
- */
-function expandEvals(commands: ESQLCommand[]): ESQLCommand[] {
-  const expanded: ESQLCommand[] = [];
-  for (const command of commands) {
-    if (command.name.toLowerCase() === 'eval') {
-      // treat each expression within EVAL as a separate EVAL command
-      for (const arg of command.args) {
-        expanded.push(
-          Builder.command({
-            name: 'eval',
-            args: [arg],
-            location: command.location,
-          })
-        );
-      }
-    } else {
-      expanded.push(command);
-    }
-  }
-  return expanded;
-}
-
 /**
  * Expands a FORK command into queries for each command in each branch.
  *

From 0412eb586a4ca68ff64b87f0436847d4a0207def Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 10:46:55 -0600
Subject: [PATCH 34/54] move some utils

---
 .../src/validation/helpers.ts                 | 50 +++++++++++++-
 .../src/validation/validation.ts              | 66 ++-----------------
 2 files changed, 55 insertions(+), 61 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
index a8d745c264a39..7296cfbeb6aa9 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts
@@ -7,7 +7,9 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { type ESQLCommand, type FunctionDefinition, Walker } from '@kbn/esql-ast';
+import type { ESQLAstQueryExpression } from '@kbn/esql-ast';
+import { type ESQLCommand, type FunctionDefinition, Walker, Builder } from '@kbn/esql-ast';
+import { expandEvals } from '../shared/expand_evals';
 
 /**
  * Returns the maximum and minimum number of parameters allowed by a function
@@ -36,3 +38,49 @@ export function getMaxMinNumberOfParams(definition: FunctionDefinition) {
  */
 export const getEnrichCommands = (commands: ESQLCommand[]): ESQLCommand[] =>
   Walker.matchAll(commands, { type: 'command', name: 'enrich' }) as ESQLCommand[];
+
+/**
+ * Returns a list of subqueries to validate
+ * @param rootCommands
+ */
+export function getSubqueriesToValidate(rootCommands: ESQLCommand[]) {
+  const subsequences = [];
+  const expandedCommands = expandEvals(rootCommands);
+  for (let i = 0; i < expandedCommands.length; i++) {
+    const command = expandedCommands[i];
+
+    // every command within FORK's branches is its own subquery to be validated
+    if (command.name.toLowerCase() === 'fork') {
+      const branchSubqueries = getForkBranchSubqueries(command as ESQLCommand<'fork'>);
+      for (const subquery of branchSubqueries) {
+        subsequences.push([...expandedCommands.slice(0, i), ...subquery]);
+      }
+    }
+
+    subsequences.push(expandedCommands.slice(0, i + 1));
+  }
+
+  return subsequences.map((subsequence) => Builder.expression.query(subsequence));
+}
+
+/**
+ * Expands a FORK command into flat subqueries for each command in each branch.
+ *
+ * E.g. FORK (EVAL 1 | LIMIT 10) (RENAME foo AS bar | DROP lolz)
+ *
+ * becomes [`EVAL 1`, `EVAL 1 | LIMIT 10`, `RENAME foo AS bar`, `RENAME foo AS bar | DROP lolz`]
+ *
+ * @param command a FORK command
+ * @returns an array of expanded subqueries
+ */
+function getForkBranchSubqueries(command: ESQLCommand<'fork'>): ESQLCommand[][] {
+  const expanded: ESQLCommand[][] = [];
+  const branches = command.args as ESQLAstQueryExpression[];
+  for (let j = 0; j < branches.length; j++) {
+    for (let k = 0; k < branches[j].commands.length; k++) {
+      const partialQuery = branches[j].commands.slice(0, k + 1);
+      expanded.push(partialQuery);
+    }
+  }
+  return expanded;
+}
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index 2412b850cd99c..7fe5d6ef79bf1 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -7,27 +7,21 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import type {
-  ESQLAst,
-  ESQLAstQueryExpression,
-  ESQLCommand,
-  ESQLMessage,
-  ErrorTypes,
-} from '@kbn/esql-ast';
-import { EsqlQuery, walk, esqlCommandRegistry, Builder, BasicPrettyPrinter } from '@kbn/esql-ast';
-import { getMessageFromId } from '@kbn/esql-ast/src/definitions/utils';
+import type { ESQLAst, ESQLCommand, ESQLMessage, ErrorTypes } from '@kbn/esql-ast';
+import { BasicPrettyPrinter, EsqlQuery, esqlCommandRegistry, walk } from '@kbn/esql-ast';
 import type {
   ESQLFieldWithMetadata,
   ICommandCallbacks,
 } from '@kbn/esql-ast/src/commands_registry/types';
+import { getMessageFromId } from '@kbn/esql-ast/src/definitions/utils';
 import type { LicenseType } from '@kbn/licensing-types';
 
+import { getQueryForFields } from '../autocomplete/helper';
+import { getColumnsByTypeHelper } from '../shared/resources_helpers';
 import type { ESQLCallbacks } from '../shared/types';
 import { retrievePolicies, retrieveSources } from './resources';
 import type { ReferenceMaps, ValidationOptions, ValidationResult } from './types';
-import { getQueryForFields } from '../autocomplete/helper';
-import { getColumnsByTypeHelper } from '../shared/resources_helpers';
-import { expandEvals } from '../shared/expand_evals';
+import { getSubqueriesToValidate } from './helpers';
 
 /**
  * ES|QL validation public API
@@ -264,51 +258,3 @@ function validateUnsupportedTypeFields(fields: Map<string, ESQLFieldWithMetadata
   }
   return messages;
 }
-
-/**
- * Returns a list of subqueries to validate
- * @param rootCommands
- *
- * TODO - extract to another module?
- */
-function getSubqueriesToValidate(rootCommands: ESQLCommand[]) {
-  const subsequences = [];
-  const expandedCommands = expandEvals(rootCommands);
-  for (let i = 0; i < expandedCommands.length; i++) {
-    const command = expandedCommands[i];
-
-    // every command within FORK's branches is its own subquery to be validated
-    if (command.name.toLowerCase() === 'fork') {
-      const branchSubqueries = getForkBranchSubqueries(command as ESQLCommand<'fork'>);
-      for (const subquery of branchSubqueries) {
-        subsequences.push([...expandedCommands.slice(0, i), ...subquery]);
-      }
-    }
-
-    subsequences.push(expandedCommands.slice(0, i + 1));
-  }
-
-  return subsequences.map((subsequence) => Builder.expression.query(subsequence));
-}
-
-/**
- * Expands a FORK command into queries for each command in each branch.
- *
- * E.g. FORK (EVAL 1 | LIMIT 10) (RENAME foo AS bar | DROP lolz)
- *
- * becomes [`EVAL 1`, `EVAL 1 | LIMIT 10`, `RENAME foo AS bar`, `RENAME foo AS bar | DROP lolz`]
- *
- * @param command a FORK command
- * @returns an array of expanded subqueries
- */
-function getForkBranchSubqueries(command: ESQLCommand<'fork'>): ESQLCommand[][] {
-  const expanded: ESQLCommand[][] = [];
-  const branches = command.args as ESQLAstQueryExpression[];
-  for (let j = 0; j < branches.length; j++) {
-    for (let k = 0; k < branches[j].commands.length; k++) {
-      const partialQuery = branches[j].commands.slice(0, k + 1);
-      expanded.push(partialQuery);
-    }
-  }
-  return expanded;
-}

From e997976d068a6a4fc7753a4ca0404c6b9f7da987 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 11:17:04 -0600
Subject: [PATCH 35/54] comment cleanup

---
 .../src/commands_registry/commands/fork/columns_after.ts    | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
index 0f22d746d23fd..36cecba91e721 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts
@@ -41,15 +41,11 @@ export const columnsAfter = async (
     columnsFromBranches.push(columnsFromBranch);
   }
 
-  /**
-   * One of the branches may have overwritten
-   */
-
   const maps = columnsFromBranches.map((cols) => new Map(cols.map((_col) => [_col.name, _col])));
 
   const merged = new Map<string, ESQLColumnData>();
 
-  // O(b * n), where b is the branches and n is the number of columns
+  // O(b * n), where b is the branches and n is the max number of columns in a branch
   for (const map of maps) {
     for (const [name, colData] of map) {
       /**

From 8a40764a97428f7d198209a8a2eef35af6224a41 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 15:09:06 -0600
Subject: [PATCH 36/54] infer expression column names from the original query
 text, preserving formatting

also, restructure code to simplify API for column helper
---
 .../src/autocomplete/autocomplete.test.ts     | 138 +++++---------
 .../src/autocomplete/autocomplete.ts          |  52 +++---
 .../src/autocomplete/helper.ts                |  77 --------
 .../src/shared/helpers.ts                     |   6 +-
 .../src/shared/query_string_utils.test.ts     | 168 ------------------
 .../src/shared/query_string_utils.ts          |  51 ------
 .../resources_helpers.test.ts}                |   4 +-
 .../src/shared/resources_helpers.ts           | 150 ++++++++++++----
 .../src/validation/validation.ts              |  10 +-
 9 files changed, 204 insertions(+), 452 deletions(-)
 delete mode 100644 src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
 delete mode 100644 src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.test.ts
 delete mode 100644 src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.ts
 rename src/platform/packages/shared/kbn-esql-validation-autocomplete/src/{autocomplete/helper.test.ts => shared/resources_helpers.test.ts} (95%)

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts
index 5c68093f0079b..9f7220c810144 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts
@@ -124,97 +124,57 @@ describe('autocomplete', () => {
     testSuggestions('from a metadata _id | eval col0 = a | /', commands);
   });
 
-  for (const command of ['keep', 'drop']) {
-    describe(command, () => {
-      testSuggestions(`from a | ${command} /`, getFieldNamesByType('any'));
-      testSuggestions(
-        `from a | ${command} keywordField, /`,
-        getFieldNamesByType('any').filter((name) => name !== 'keywordField')
-      );
+  describe.each(['keep', 'drop'])('%s', (command) => {
+    testSuggestions(`from a | ${command} /`, getFieldNamesByType('any'));
+    testSuggestions(
+      `from a | ${command} keywordField, /`,
+      getFieldNamesByType('any').filter((name) => name !== 'keywordField')
+    );
 
-      testSuggestions(
-        `from a | ${command} keywordField,/`,
-        getFieldNamesByType('any').filter((name) => name !== 'keywordField'),
-        ','
-      );
+    testSuggestions(
+      `from a | ${command} keywordField,/`,
+      getFieldNamesByType('any').filter((name) => name !== 'keywordField'),
+      ','
+    );
 
-      testSuggestions(
-        `from a_index | eval round(doubleField) + 1 | eval \`round(doubleField) + 1\` + 1 | eval \`\`\`round(doubleField) + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`round(doubleField) + 1\`\`\`\` + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`\`\`\`\`\`\`\`\`round(doubleField) + 1\`\`\`\`\`\`\`\` + 1\`\`\`\` + 1\`\` + 1\` + 1 | ${command} /`,
-        [
-          ...getFieldNamesByType('any'),
-          '`round(doubleField) + 1`',
-          '```round(doubleField) + 1`` + 1`',
-          '```````round(doubleField) + 1```` + 1`` + 1`',
-          '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1`',
-          '```````````````````````````````round(doubleField) + 1```````````````` + 1```````` + 1```` + 1`` + 1`',
-        ],
-        undefined,
-        [
-          [
-            ...fields,
-            // the following non-field columns will come over the wire as part of the response
-            {
-              name: 'round(doubleField) + 1',
-              type: 'double',
-              userDefined: false,
-            },
-            {
-              name: '`round(doubleField) + 1` + 1',
-              type: 'double',
-              userDefined: false,
-            },
-            {
-              name: '```round(doubleField) + 1`` + 1` + 1',
-              type: 'double',
-              userDefined: false,
-            },
-            {
-              name: '```````round(doubleField) + 1```` + 1`` + 1` + 1',
-              type: 'double',
-              userDefined: false,
-            },
-            {
-              name: '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1` + 1',
-              type: 'double',
-              userDefined: false,
-            },
-          ],
-        ]
-      );
+    testSuggestions(
+      `from a_index | eval round(doubleField) + 1 | eval \`round(doubleField) + 1\` + 1 | eval \`\`\`round(doubleField) + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`round(doubleField) + 1\`\`\`\` + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`\`\`\`\`\`\`\`\`round(doubleField) + 1\`\`\`\`\`\`\`\` + 1\`\`\`\` + 1\`\` + 1\` + 1 | ${command} /`,
+      [
+        ...getFieldNamesByType('any'),
+        '`round(doubleField) + 1`',
+        '```round(doubleField) + 1`` + 1`',
+        '```````round(doubleField) + 1```` + 1`` + 1`',
+        '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1`',
+        '```````````````````````````````round(doubleField) + 1```````````````` + 1```````` + 1```` + 1`` + 1`',
+      ]
+    );
 
-      it('should not suggest already-used fields and user defined columns', async () => {
-        const { suggest: suggestTest } = await setup();
-        const getSuggestions = async (query: string, opts?: SuggestOptions) =>
-          (await suggestTest(query, opts)).map((value) => value.text);
-
-        expect(
-          await getSuggestions('from a_index | EVAL foo = 1 | KEEP /', {
-            callbacks: {
-              getColumnsFor: () => [
-                ...fields,
-                { name: 'foo', type: 'integer', userDefined: false },
-              ],
-            },
-          })
-        ).toContain('foo');
-        expect(
-          await getSuggestions('from a_index | EVAL foo = 1 | KEEP foo, /', {
-            callbacks: {
-              getColumnsFor: () => [
-                ...fields,
-                { name: 'foo', type: 'integer', userDefined: false },
-              ],
-            },
-          })
-        ).not.toContain('foo');
-
-        expect(await getSuggestions('from a_index | KEEP /')).toContain('doubleField');
-        expect(await getSuggestions('from a_index | KEEP doubleField, /')).not.toContain(
-          'doubleField'
-        );
-      });
+    it('should not suggest already-used fields and user defined columns', async () => {
+      const { suggest: suggestTest } = await setup();
+      const getSuggestions = async (query: string, opts?: SuggestOptions) =>
+        (await suggestTest(query, opts)).map((value) => value.text);
+
+      expect(
+        await getSuggestions('from a_index | EVAL foo = 1 | KEEP /', {
+          callbacks: {
+            getColumnsFor: () => [...fields, { name: 'foo', type: 'integer', userDefined: false }],
+          },
+        })
+      ).toContain('foo');
+      expect(
+        await getSuggestions('from a_index | EVAL foo = 1 | KEEP foo, /', {
+          callbacks: {
+            getColumnsFor: () => [...fields, { name: 'foo', type: 'integer', userDefined: false }],
+          },
+        })
+      ).not.toContain('foo');
+
+      expect(await getSuggestions('from a_index | KEEP /')).toContain('doubleField');
+      expect(await getSuggestions('from a_index | KEEP doubleField, /')).not.toContain(
+        'doubleField'
+      );
     });
-  }
+  });
 
   // @TODO: get updated eval block from main
   describe('values suggestions', () => {
@@ -241,7 +201,7 @@ describe('autocomplete', () => {
       const triggerOffset = statement.lastIndexOf(' ');
       await suggest(statement, triggerOffset + 1, callbackMocks);
       expect(callbackMocks.getColumnsFor).toHaveBeenCalledWith({
-        query: 'from index_b',
+        query: 'FROM index_b',
       });
     });
     it('should send the fields query aware of the location', async () => {
@@ -249,7 +209,7 @@ describe('autocomplete', () => {
       const statement = 'from index_d | drop | eval col0 = abs(doubleField) ';
       const triggerOffset = statement.lastIndexOf('p') + 1; // drop <here>
       await suggest(statement, triggerOffset + 1, callbackMocks);
-      expect(callbackMocks.getColumnsFor).toHaveBeenCalledWith({ query: 'from index_d' });
+      expect(callbackMocks.getColumnsFor).toHaveBeenCalledWith({ query: 'FROM index_d' });
     });
   });
 
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
index 79a1de540b271..0b6b41b5a187f 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts
@@ -7,37 +7,37 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+import type { ESQLAstQueryExpression } from '@kbn/esql-ast';
 import {
+  ESQL_VARIABLES_PREFIX,
+  EsqlQuery,
+  esqlCommandRegistry,
+  getCommandAutocompleteDefinitions,
   parse,
   type ESQLAstItem,
   type ESQLCommand,
   type ESQLCommandOption,
   type ESQLFunction,
-  esqlCommandRegistry,
-  getCommandAutocompleteDefinitions,
-  ESQL_VARIABLES_PREFIX,
 } from '@kbn/esql-ast';
-import { EDITOR_MARKER } from '@kbn/esql-ast/src/definitions/constants';
-import { correctQuerySyntax } from '@kbn/esql-ast/src/definitions/utils/ast';
-import {
-  getControlSuggestionIfSupported,
-  buildFieldsDefinitionsWithMetadata,
-} from '@kbn/esql-ast/src/definitions/utils';
 import { getRecommendedQueriesSuggestionsFromStaticTemplates } from '@kbn/esql-ast/src/commands_registry/options/recommended_queries';
 import type {
   ESQLColumnData,
   GetColumnsByTypeFn,
   ISuggestionItem,
 } from '@kbn/esql-ast/src/commands_registry/types';
+import {
+  buildFieldsDefinitionsWithMetadata,
+  getControlSuggestionIfSupported,
+} from '@kbn/esql-ast/src/definitions/utils';
+import { correctQuerySyntax } from '@kbn/esql-ast/src/definitions/utils/ast';
 import { ESQLVariableType } from '@kbn/esql-types';
 import type { LicenseType } from '@kbn/licensing-types';
-import { isSourceCommand } from '../shared/helpers';
 import { getAstContext } from '../shared/context';
+import { isSourceCommand } from '../shared/helpers';
 import { getColumnsByTypeHelper, getSourcesHelper } from '../shared/resources_helpers';
 import type { ESQLCallbacks } from '../shared/types';
-import { getQueryForFields } from './helper';
-import { mapRecommendedQueriesFromExtensions } from './utils/recommended_queries_helpers';
 import { getCommandContext } from './get_command_context';
+import { mapRecommendedQueriesFromExtensions } from './utils/recommended_queries_helpers';
 
 type GetColumnMapFn = () => Promise<Map<string, ESQLColumnData>>;
 
@@ -56,14 +56,12 @@ export async function suggest(
     return [];
   }
 
-  // build the correct query to fetch the list of fields
-  const queryForFields = getQueryForFields(correctedQuery, root);
-
   const { getColumnsByType, getColumnMap } = getColumnsByTypeRetriever(
-    queryForFields.replace(EDITOR_MARKER, ''),
-    resourceRetriever,
-    innerText
+    root,
+    innerText,
+    resourceRetriever
   );
+
   const supportsControls = resourceRetriever?.canSuggestVariables?.() ?? false;
   const getVariables = resourceRetriever?.getVariables;
   const getSources = getSourcesHelper(resourceRetriever);
@@ -112,9 +110,9 @@ export async function suggest(
         }
 
         const { getColumnsByType: getColumnsByTypeEmptyState } = getColumnsByTypeRetriever(
-          fromCommand,
-          resourceRetriever,
-          innerText
+          EsqlQuery.fromSrc(fromCommand).ast,
+          innerText,
+          resourceRetriever
         );
         const editorExtensions = (await resourceRetriever?.getEditorExtensions?.(fromCommand)) ?? {
           recommendedQueries: [],
@@ -171,15 +169,15 @@ export async function suggest(
 }
 
 export function getColumnsByTypeRetriever(
-  queryForFields: string,
-  resourceRetriever?: ESQLCallbacks,
-  fullQuery?: string
+  query: ESQLAstQueryExpression,
+  queryText: string,
+  resourceRetriever?: ESQLCallbacks
 ): { getColumnsByType: GetColumnsByTypeFn; getColumnMap: GetColumnMapFn } {
-  const helpers = getColumnsByTypeHelper(queryForFields, resourceRetriever);
+  const helpers = getColumnsByTypeHelper(query, queryText, resourceRetriever);
   const getVariables = resourceRetriever?.getVariables;
   const canSuggestVariables = resourceRetriever?.canSuggestVariables?.() ?? false;
 
-  const queryString = fullQuery ?? queryForFields;
+  const queryString = queryText;
   const lastCharacterTyped = queryString[queryString.length - 1];
   const lastCharIsQuestionMark = lastCharacterTyped === ESQL_VARIABLES_PREFIX;
   return {
@@ -192,7 +190,7 @@ export function getColumnsByTypeRetriever(
         ...options,
         supportsControls: canSuggestVariables && !lastCharIsQuestionMark,
       };
-      const editorExtensions = (await resourceRetriever?.getEditorExtensions?.(queryForFields)) ?? {
+      const editorExtensions = (await resourceRetriever?.getEditorExtensions?.(queryText)) ?? {
         recommendedQueries: [],
         recommendedFields: [],
       };
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
deleted file mode 100644
index 91f289a88db19..0000000000000
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-
-import type { ESQLAstQueryExpression } from '@kbn/esql-ast';
-import { BasicPrettyPrinter, EDITOR_MARKER } from '@kbn/esql-ast';
-import { expandEvals } from '../shared/expand_evals';
-
-/**
- * This function is used to build the query that will be used to compute the
- * available fields for the current cursor location.
- *
- * Generally, this is the user's query up to the end of the previous command.
- *
- * @param queryString The original query string
- * @param commands
- * @returns
- */
-export function getQueryForFields(queryString: string, root: ESQLAstQueryExpression): string {
-  const commands = root.commands;
-  const lastCommand = commands[commands.length - 1];
-  if (lastCommand && lastCommand.name === 'fork' && lastCommand.args.length > 0) {
-    /**
-     * This flattens the current fork branch into a simpler but equivalent
-     * query that is compatible with the existing field computation/caching strategy.
-     *
-     * The intuition here is that if the cursor is within a fork branch, the
-     * previous context is equivalent to a query without the FORK command:
-     *
-     * Original query: FROM lolz | EVAL foo = 1 | FORK (EVAL bar = 2) (EVAL baz = 3 | WHERE /)
-     * Simplified:     FROM lolz | EVAL foo = 1 | EVAL baz = 3
-     */
-    const currentBranch = lastCommand.args[lastCommand.args.length - 1] as ESQLAstQueryExpression;
-    const newCommands = commands.slice(0, -1).concat(currentBranch.commands.slice(0, -1));
-    return BasicPrettyPrinter.print({ ...root, commands: newCommands });
-  }
-
-  if (lastCommand && lastCommand.name === 'eval') {
-    const endsWithComma = queryString.replace(EDITOR_MARKER, '').trim().endsWith(',');
-    if (lastCommand.args.length > 1 || endsWithComma) {
-      /**
-       * If we get here, we know that we have a multi-expression EVAL statement.
-       *
-       * e.g. EVAL foo = 1, bar = foo + 1, baz = bar + 1
-       *
-       * In order for this to work with the caching system which expects field availability to be
-       * delineated by pipes, we need to split the current EVAL command into an equivalent
-       * set of single-expression EVAL commands.
-       *
-       * Original query: FROM lolz | EVAL foo = 1, bar = foo + 1, baz = bar + 1, /
-       * Simplified:     FROM lolz | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1
-       */
-      const expanded = expandEvals(commands);
-      const newCommands = expanded.slice(0, endsWithComma ? undefined : -1);
-      return BasicPrettyPrinter.print({ ...root, commands: newCommands });
-    }
-  }
-
-  // If there is only one source command and it does not require fields, do not
-  // fetch fields, hence return an empty string.
-  return commands.length === 1 && ['row', 'show'].includes(commands[0].name)
-    ? ''
-    : buildQueryUntilPreviousCommand(root);
-}
-
-function buildQueryUntilPreviousCommand(root: ESQLAstQueryExpression) {
-  if (root.commands.length === 1) {
-    return BasicPrettyPrinter.print({ ...root.commands[0] });
-  } else {
-    return BasicPrettyPrinter.print({ ...root, commands: root.commands.slice(0, -1) });
-  }
-}
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index 8bb49d1e386f6..9115c0b84bb8a 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -106,11 +106,11 @@ export async function getFieldsFromES(query: string, resourceRetriever?: ESQLCal
  * @returns a list of fields that are available for the current pipe
  */
 export async function getCurrentQueryAvailableColumns(
-  query: string,
   commands: ESQLAstCommand[],
   previousPipeFields: ESQLColumnData[],
   fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>,
-  policies: Map<string, ESQLPolicy>
+  policies: Map<string, ESQLPolicy>,
+  originalQueryText: string
 ) {
   const lastCommand = commands[commands.length - 1];
   const commandDef = esqlCommandRegistry.getCommandByName(lastCommand.name);
@@ -154,7 +154,7 @@ export async function getCurrentQueryAvailableColumns(
     return commandDef.methods.columnsAfter(
       lastCommand,
       previousPipeFields,
-      query,
+      originalQueryText,
       additionalFields
     );
   }
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.test.ts
deleted file mode 100644
index c69767d4a3cb1..0000000000000
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.test.ts
+++ /dev/null
@@ -1,168 +0,0 @@
-/*
- * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-import { removeLastPipe, processPipes, toSingleLine, removeComments } from './query_string_utils';
-
-describe('query_string_utils', () => {
-  describe('removeComments', () => {
-    it('should remove single-line comments', () => {
-      const text = `
-        FROM users // This is a comment
-        | WHERE age > 18
-      `;
-      const expected = `FROM users         | WHERE age > 18`;
-      expect(removeComments(text)).toBe(expected);
-    });
-
-    it('should remove multi-line comments', () => {
-      const text = `
-        FROM /* This is a
-        multi-line
-        comment */
-        products
-      `;
-      const expected = `FROM 
-        products`;
-      expect(removeComments(text)).toBe(expected);
-    });
-
-    it('should remove both single-line and multi-line comments', () => {
-      const text = `
-        FROM items // Get the name
-        /* The price of the
-        product */
-        | KEEP name, price
-      `;
-      const expected = `FROM items         
-        | KEEP name, price`;
-      expect(removeComments(text)).toBe(expected);
-    });
-
-    it('should handle text with no comments', () => {
-      const text = `
-        FROM orders
-        | KEEP order_id, status
-        | WHERE status = 'pending';
-      `;
-      expect(removeComments(text)).toBe(text.trim());
-    });
-
-    it('should handle comments at the beginning and end of the text', () => {
-      const text = `
-        // Initial comment
-        FROM logs
-        /* Final
-        comment */
-      `;
-      const expected = `FROM logs`;
-      expect(removeComments(text)).toBe(expected);
-    });
-
-    it('should handle consecutive single-line comments', () => {
-      const text = `
-        // Comment line 1
-        // Comment line 2
-        FROM events | STATS COUNT(*)
-      `;
-      const expected = `
-        FROM events | STATS COUNT(*)
-      `;
-      expect(removeComments(text)).toBe(expected.trim());
-    });
-  });
-
-  describe('removeLastPipe', () => {
-    it('should remove the last pipe and any trailing whitespace', () => {
-      expect(removeLastPipe('value1|value2|')).toBe('value1|value2');
-      expect(removeLastPipe('value1|value2|  ')).toBe('value1|value2');
-    });
-
-    it('should return the original string if there is no pipe', () => {
-      expect(removeLastPipe('FROM index')).toBe('FROM index');
-      expect(removeLastPipe('FROM index  ')).toBe('FROM index');
-    });
-
-    it('should handle strings with multiple pipes correctly', () => {
-      expect(removeLastPipe('FROM index | STATS count() | DROP field1  ')).toBe(
-        'FROM index | STATS count()'
-      );
-    });
-
-    it('should handle an empty string', () => {
-      expect(removeLastPipe('')).toBe('');
-    });
-
-    it('should handle a string with only whitespace', () => {
-      expect(removeLastPipe('  ')).toBe('');
-    });
-  });
-
-  describe('processPipes', () => {
-    it('should return an array of strings, each progressively including parts separated by " | "', () => {
-      const input = 'FROM index|EVAL col = ABS(numeric) | KEEP col';
-      const expected = [
-        'FROM index',
-        'FROM index | EVAL col = ABS(numeric)',
-        'FROM index | EVAL col = ABS(numeric) | KEEP col',
-      ];
-      expect(processPipes(input)).toEqual(expected);
-    });
-
-    it('should handle leading and trailing whitespace in parts', () => {
-      const input = ' FROM index |  EVAL col = ABS(numeric)  | KEEP col ';
-      const expected = [
-        'FROM index',
-        'FROM index | EVAL col = ABS(numeric)',
-        'FROM index | EVAL col = ABS(numeric) | KEEP col',
-      ];
-      expect(processPipes(input)).toEqual(expected);
-    });
-
-    it('should return an array with the trimmed input if there are no pipes', () => {
-      const input = 'FROM index';
-      const expected = ['FROM index'];
-      expect(processPipes(input)).toEqual(expected);
-
-      const inputWithWhitespace = '  FROM index  ';
-      const expectedWithWhitespace = ['FROM index'];
-      expect(processPipes(inputWithWhitespace)).toEqual(expectedWithWhitespace);
-    });
-
-    it('should ignore comments', () => {
-      const input = '// This is an ES|QL query \n FROM index';
-      const expected = ['FROM index'];
-      expect(processPipes(input)).toEqual(expected);
-    });
-  });
-
-  describe('toSingleLine', () => {
-    it('should convert a multi-line pipe-separated string to a single line with " | " as separator', () => {
-      const input = 'FROM index \n|EVAL col = ABS(numeric)\n|KEEP col';
-      const expected = 'FROM index | EVAL col = ABS(numeric) | KEEP col';
-      expect(toSingleLine(input)).toBe(expected);
-    });
-
-    it('should trim whitespace from each part', () => {
-      const input = ' FROM index |  EVAL col = ABS(numeric)  | KEEP col ';
-      const expected = 'FROM index | EVAL col = ABS(numeric) | KEEP col';
-      expect(toSingleLine(input)).toBe(expected);
-    });
-
-    it('should trim whitespace from each part for multi-line strings', () => {
-      const input = ' FROM index \n|  EVAL col = ABS(numeric)  \n| KEEP col ';
-      const expected = 'FROM index | EVAL col = ABS(numeric) | KEEP col';
-      expect(toSingleLine(input)).toBe(expected);
-    });
-
-    it('should ignore comments', () => {
-      const input = '// This is an ES|QL query \n FROM index';
-      const expected = 'FROM index';
-      expect(toSingleLine(input)).toEqual(expected);
-    });
-  });
-});
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.ts
deleted file mode 100644
index 43229074c02df..0000000000000
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.ts
+++ /dev/null
@@ -1,51 +0,0 @@
-/*
- * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-export function removeComments(text: string): string {
-  // Remove single-line comments
-  const withoutSingleLineComments = text.replace(/\/\/.*?(?:\r\n|\r|\n|$)/gm, '');
-  // Remove multi-line comments
-  const withoutMultiLineComments = withoutSingleLineComments.replace(/\/\*[\s\S]*?\*\//g, '');
-  return withoutMultiLineComments.trim();
-}
-
-export function removeLastPipe(inputString: string): string {
-  const queryNoComments = removeComments(inputString);
-  const lastPipeIndex = queryNoComments.lastIndexOf('|');
-  if (lastPipeIndex !== -1) {
-    return queryNoComments.substring(0, lastPipeIndex).trimEnd();
-  }
-  return queryNoComments.trimEnd();
-}
-
-export function processPipes(inputString: string) {
-  const queryNoComments = removeComments(inputString);
-  const parts = queryNoComments.split('|');
-  const results = [];
-  let currentString = '';
-
-  for (let i = 0; i < parts.length; i++) {
-    const part = parts[i].trim();
-    if (i === 0) {
-      currentString = part;
-    } else {
-      currentString += ' | ' + part;
-    }
-    results.push(currentString.trim());
-  }
-  return results;
-}
-
-export function toSingleLine(inputString: string): string {
-  const queryNoComments = removeComments(inputString);
-  return queryNoComments
-    .split('|')
-    .map((line) => line.trim())
-    .filter((line) => line !== '')
-    .join(' | ');
-}
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts
similarity index 95%
rename from src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts
rename to src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts
index bdaa985966c16..742572c1c0e53 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts
@@ -7,8 +7,8 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { EDITOR_MARKER, parse } from '@kbn/esql-ast';
-import { getQueryForFields } from './helper';
+import { parse, EDITOR_MARKER } from '@kbn/esql-ast';
+import { getQueryForFields } from './resources_helpers';
 
 describe('getQueryForFields', () => {
   const assert = (query: string, expected: string) => {
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index 3d16a5ede5245..a1ff3bee33166 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -7,7 +7,8 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { parse } from '@kbn/esql-ast';
+import type { ESQLAstQueryExpression } from '@kbn/esql-ast';
+import { BasicPrettyPrinter, Builder, EDITOR_MARKER, EsqlQuery } from '@kbn/esql-ast';
 import type {
   ESQLColumnData,
   ESQLFieldWithMetadata,
@@ -15,7 +16,7 @@ import type {
 } from '@kbn/esql-ast/src/commands_registry/types';
 import type { ESQLCallbacks } from './types';
 import { getFieldsFromES, getCurrentQueryAvailableColumns } from './helpers';
-import { removeLastPipe, processPipes, toSingleLine } from './query_string_utils';
+import { expandEvals } from './expand_evals';
 
 export const NOT_SUGGESTED_TYPES = ['unsupported'];
 
@@ -47,57 +48,81 @@ function getValueInsensitive(keyToCheck: string) {
  * @param queryText
  */
 async function cacheColumnsForQuery(
-  queryText: string,
+  query: ESQLAstQueryExpression,
   fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>,
-  policies: Map<string, ESQLPolicy>
+  policies: Map<string, ESQLPolicy>,
+  originalQueryText: string
 ) {
-  const existsInCache = checkCacheInsensitive(queryText);
+  const cacheKey = BasicPrettyPrinter.print(query);
+  const existsInCache = checkCacheInsensitive(cacheKey);
   if (existsInCache) {
     // this is already in the cache
     return;
   }
-  const queryTextWithoutLastPipe = removeLastPipe(queryText);
-  // retrieve the user defined fields from the query without an extra call
-  const fieldsAvailableAfterPreviousCommand = getValueInsensitive(queryTextWithoutLastPipe);
+
+  const fieldsAvailableAfterPreviousCommand = getValueInsensitive(
+    BasicPrettyPrinter.print({ ...query, commands: query.commands.slice(0, -1) })
+  );
+
   if (fieldsAvailableAfterPreviousCommand && fieldsAvailableAfterPreviousCommand?.length) {
-    const { root } = parse(queryText);
     const availableFields = await getCurrentQueryAvailableColumns(
-      queryText,
-      root.commands,
+      query.commands,
       fieldsAvailableAfterPreviousCommand,
       fetchFields,
-      policies
+      policies,
+      originalQueryText
     );
-    cache.set(queryText, availableFields);
+    cache.set(cacheKey, availableFields);
   }
 }
 
-export function getColumnsByTypeHelper(queryText: string, resourceRetriever?: ESQLCallbacks) {
-  const getColumns = async () => {
+/**
+ * Efficiently computes the list of columns available after the given query
+ * and returns column querying utilities.
+ *
+ * @param root
+ * @param originalQueryText the text of the original query, used to infer column names from expressions
+ * @param resourceRetriever
+ * @returns
+ */
+export function getColumnsByTypeHelper(
+  query: ESQLAstQueryExpression,
+  originalQueryText: string,
+  resourceRetriever?: ESQLCallbacks
+) {
+  const queryForFields = getQueryForFields(originalQueryText, query);
+  const root = EsqlQuery.fromSrc(queryForFields).ast;
+
+  const cacheColumns = async () => {
     // in some cases (as in the case of ROW or SHOW) the query is not set
-    if (!queryText) {
+    if (!originalQueryText) {
       return;
     }
 
-    const getFields = async (query: string) => {
-      const cached = getValueInsensitive(query);
+    const getFields = async (queryToES: string) => {
+      const cached = getValueInsensitive(queryToES);
       if (cached) {
         return cached as ESQLFieldWithMetadata[];
       }
-      const fields = await getFieldsFromES(query, resourceRetriever);
-      cache.set(query, fields);
+      const fields = await getFieldsFromES(queryToES, resourceRetriever);
+      cache.set(queryToES, fields);
       return fields;
     };
 
-    const [sourceCommand, ...partialQueries] = processPipes(queryText);
-    getFields(sourceCommand);
+    const subqueries = [];
+    for (let i = 0; i < root.commands.length; i++) {
+      subqueries.push(Builder.expression.query(root.commands.slice(0, i + 1)));
+    }
+
+    // source command
+    getFields(BasicPrettyPrinter.print(subqueries[0]));
 
     const policies = (await resourceRetriever?.getPolicies?.()) ?? [];
     const policyMap = new Map(policies.map((p) => [p.name, p]));
 
     // build fields cache for every partial query
-    for (const query of partialQueries) {
-      await cacheColumnsForQuery(query, getFields, policyMap);
+    for (const subquery of subqueries.slice(1)) {
+      await cacheColumnsForQuery(subquery, getFields, policyMap, originalQueryText);
     }
   };
 
@@ -107,9 +132,8 @@ export function getColumnsByTypeHelper(queryText: string, resourceRetriever?: ES
       ignored: string[] = []
     ): Promise<ESQLColumnData[]> => {
       const types = Array.isArray(expectedType) ? expectedType : [expectedType];
-      await getColumns();
-      const queryTextForCacheSearch = toSingleLine(queryText);
-      const cachedFields = getValueInsensitive(queryTextForCacheSearch);
+      await cacheColumns();
+      const cachedFields = getValueInsensitive(queryForFields);
       return (
         cachedFields?.filter(({ name, type }) => {
           const ts = Array.isArray(type) ? type : [type];
@@ -123,9 +147,8 @@ export function getColumnsByTypeHelper(queryText: string, resourceRetriever?: ES
       );
     },
     getColumnMap: async (): Promise<Map<string, ESQLColumnData>> => {
-      await getColumns();
-      const queryTextForCacheSearch = toSingleLine(queryText);
-      const cachedFields = getValueInsensitive(queryTextForCacheSearch);
+      await cacheColumns();
+      const cachedFields = getValueInsensitive(queryForFields);
       const cacheCopy = new Map<string, ESQLColumnData>();
       cachedFields?.forEach((field) => cacheCopy.set(field.name, field));
       return cacheCopy;
@@ -154,3 +177,70 @@ export function getSourcesHelper(resourceRetriever?: ESQLCallbacks) {
     return (await resourceRetriever?.getSources?.()) || [];
   };
 }
+
+/**
+ * This function is used to build the query that will be used to compute the
+ * fields available at the final position. It is robust to final partial commands
+ * e.g. "FROM logs* | EVAL foo = 1 | EVAL "
+ *
+ * Generally, this is the user's query up to the end of the previous command, but there
+ * are special cases for multi-expression EVAL and FORK branches.
+ *
+ * @param queryString The original query string
+ * @param commands
+ * @returns
+ */
+export function getQueryForFields(queryString: string, root: ESQLAstQueryExpression): string {
+  const commands = root.commands;
+  const lastCommand = commands[commands.length - 1];
+  if (lastCommand && lastCommand.name === 'fork' && lastCommand.args.length > 0) {
+    /**
+     * This flattens the current fork branch into a simpler but equivalent
+     * query that is compatible with the existing field computation/caching strategy.
+     *
+     * The intuition here is that if the cursor is within a fork branch, the
+     * previous context is equivalent to a query without the FORK command:
+     *
+     * Original query: FROM lolz | EVAL foo = 1 | FORK (EVAL bar = 2) (EVAL baz = 3 | WHERE /)
+     * Simplified:     FROM lolz | EVAL foo = 1 | EVAL baz = 3
+     */
+    const currentBranch = lastCommand.args[lastCommand.args.length - 1] as ESQLAstQueryExpression;
+    const newCommands = commands.slice(0, -1).concat(currentBranch.commands.slice(0, -1));
+    return BasicPrettyPrinter.print({ ...root, commands: newCommands });
+  }
+
+  if (lastCommand && lastCommand.name === 'eval') {
+    const endsWithComma = queryString.replace(EDITOR_MARKER, '').trim().endsWith(',');
+    if (lastCommand.args.length > 1 || endsWithComma) {
+      /**
+       * If we get here, we know that we have a multi-expression EVAL statement.
+       *
+       * e.g. EVAL foo = 1, bar = foo + 1, baz = bar + 1
+       *
+       * In order for this to work with the caching system which expects field availability to be
+       * delineated by pipes, we need to split the current EVAL command into an equivalent
+       * set of single-expression EVAL commands.
+       *
+       * Original query: FROM lolz | EVAL foo = 1, bar = foo + 1, baz = bar + 1, /
+       * Simplified:     FROM lolz | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1
+       */
+      const expanded = expandEvals(commands);
+      const newCommands = expanded.slice(0, endsWithComma ? undefined : -1);
+      return BasicPrettyPrinter.print({ ...root, commands: newCommands });
+    }
+  }
+
+  // If there is only one source command and it does not require fields, do not
+  // fetch fields, hence return an empty string.
+  return commands.length === 1 && ['row', 'show'].includes(commands[0].name)
+    ? ''
+    : buildQueryUntilPreviousCommand(root);
+}
+
+function buildQueryUntilPreviousCommand(root: ESQLAstQueryExpression) {
+  if (root.commands.length === 1) {
+    return BasicPrettyPrinter.print({ ...root.commands[0] });
+  } else {
+    return BasicPrettyPrinter.print({ ...root, commands: root.commands.slice(0, -1) });
+  }
+}
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index 7fe5d6ef79bf1..7846f8c71d7a8 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -8,7 +8,7 @@
  */
 
 import type { ESQLAst, ESQLCommand, ESQLMessage, ErrorTypes } from '@kbn/esql-ast';
-import { BasicPrettyPrinter, EsqlQuery, esqlCommandRegistry, walk } from '@kbn/esql-ast';
+import { EsqlQuery, esqlCommandRegistry, walk } from '@kbn/esql-ast';
 import type {
   ESQLFieldWithMetadata,
   ICommandCallbacks,
@@ -16,7 +16,6 @@ import type {
 import { getMessageFromId } from '@kbn/esql-ast/src/definitions/utils';
 import type { LicenseType } from '@kbn/licensing-types';
 
-import { getQueryForFields } from '../autocomplete/helper';
 import { getColumnsByTypeHelper } from '../shared/resources_helpers';
 import type { ESQLCallbacks } from '../shared/types';
 import { retrievePolicies, retrieveSources } from './resources';
@@ -115,8 +114,10 @@ async function validateAst(
     callbacks?.getJoinIndices?.(),
   ]);
 
+  const sourceQuery = queryString.split('|')[0];
   const sourceFields = await getColumnsByTypeHelper(
-    queryString.split('|')[0],
+    EsqlQuery.fromSrc(sourceQuery).ast,
+    sourceQuery,
     callbacks
   ).getColumnMap();
 
@@ -138,8 +139,7 @@ async function validateAst(
    */
   const subqueries = getSubqueriesToValidate(rootCommands);
   for (const subquery of subqueries) {
-    const queryForFields = getQueryForFields(BasicPrettyPrinter.print(subquery), subquery);
-    const { getColumnMap } = getColumnsByTypeHelper(queryForFields, callbacks);
+    const { getColumnMap } = getColumnsByTypeHelper(subquery, queryString, callbacks);
     const availableColumns = await getColumnMap();
 
     const references: ReferenceMaps = {

From cd00f669316b32561447aae3cb01785fe665631b Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 15:41:42 -0600
Subject: [PATCH 37/54] make more tests pass

---
 .../utils/ecs_metadata_helper.test.ts         | 15 ++--
 .../src/shared/resources_helpers.ts           |  4 +
 .../validation/__tests__/callbacks.test.ts    | 64 --------------
 .../esql_validation_meta_tests.json           | 84 ++++++++++++-------
 4 files changed, 69 insertions(+), 98 deletions(-)
 delete mode 100644 src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
index 3ded634f8810e..37e706e0aca0f 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
@@ -7,7 +7,8 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import type { ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types';
+import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types';
+import { type ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types';
 import { type ECSMetadata, enrichFieldsWithECSInfo } from './ecs_metadata_helper';
 
 describe('enrichFieldsWithECSInfo', () => {
@@ -58,9 +59,10 @@ describe('enrichFieldsWithECSInfo', () => {
         name: 'ecs.field',
         type: 'text',
         isEcs: true,
+        userDefined: false,
       },
-      { name: 'ecs.fakeBooleanField', type: 'boolean' },
-      { name: 'field2', type: 'double' },
+      { name: 'ecs.fakeBooleanField', type: 'boolean', isEcs: true, userDefined: false },
+      { name: 'field2', type: 'double', isEcs: false, userDefined: false },
     ]);
   });
 
@@ -90,14 +92,15 @@ describe('enrichFieldsWithECSInfo', () => {
       fieldsMetadataCache.fields as ECSMetadata
     );
 
-    expect(result).toEqual([
-      { name: 'ecs.version', type: 'keyword', isEcs: true },
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'ecs.version', type: 'keyword', isEcs: true, userDefined: false },
       {
         name: 'ecs.version.keyword',
         type: 'keyword',
         isEcs: true,
+        userDefined: false,
       },
-      { name: 'field2', type: 'double' },
+      { name: 'field2', type: 'double', userDefined: false },
     ]);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index a1ff3bee33166..6b74cb147f65a 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -114,6 +114,10 @@ export function getColumnsByTypeHelper(
       subqueries.push(Builder.expression.query(root.commands.slice(0, i + 1)));
     }
 
+    if (!subqueries.length) {
+      return;
+    }
+
     // source command
     getFields(BasicPrettyPrinter.print(subqueries[0]));
 
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts
deleted file mode 100644
index 16898cc737e58..0000000000000
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts
+++ /dev/null
@@ -1,64 +0,0 @@
-/*
- * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-
-import { setup } from './helpers';
-
-describe('FROM', () => {
-  test('does not load fields when validating only a single FROM, SHOW, ROW command', async () => {
-    const { validate, callbacks } = await setup();
-
-    await validate('FROM kib');
-    await validate('FROM kibana_ecommerce METADATA _i');
-    await validate('FROM kibana_ecommerce METADATA _id | ');
-    await validate('SHOW');
-    await validate('ROW \t');
-
-    expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(0);
-  });
-
-  test('loads fields with FROM source when commands after pipe present', async () => {
-    const { validate, callbacks } = await setup();
-
-    await validate('FROM kibana_ecommerce METADATA _id | eval');
-
-    expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1);
-  });
-
-  test('loads fields from JOIN index', async () => {
-    const { validate, callbacks } = await setup();
-
-    await validate('FROM index1 | LOOKUP JOIN index2 ON field1 | LIMIT 123');
-
-    expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1);
-
-    const query = (callbacks.getColumnsFor as any).mock.calls[0][0].query as string;
-
-    expect(query.includes('index1')).toBe(true);
-    expect(query.includes('index2')).toBe(true);
-  });
-
-  test('includes all "from" and "join" index for loading fields', async () => {
-    const { validate, callbacks } = await setup();
-
-    await validate(
-      'FROM index1, index2, index3 | LOOKUP JOIN index4 ON field1 | KEEP abc | LOOKUP JOIN index5 ON field2 | LIMIT 123'
-    );
-
-    expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1);
-
-    const query = (callbacks.getColumnsFor as any).mock.calls[0][0].query as string;
-
-    expect(query.includes('index1')).toBe(true);
-    expect(query.includes('index2')).toBe(true);
-    expect(query.includes('index3')).toBe(true);
-    expect(query.includes('index4')).toBe(true);
-    expect(query.includes('index5')).toBe(true);
-    expect(query.includes('index6')).toBe(false);
-  });
-});
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json
index 74209aa860b54..a21932393cefd 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json
@@ -10,113 +10,140 @@
   "fields": [
     {
       "name": "booleanField",
-      "type": "boolean"
+      "type": "boolean",
+      "userDefined": false
     },
     {
       "name": "dateField",
-      "type": "date"
+      "type": "date",
+      "userDefined": false
     },
     {
       "name": "doubleField",
-      "type": "double"
+      "type": "double",
+      "userDefined": false
     },
     {
       "name": "ipField",
-      "type": "ip"
+      "type": "ip",
+      "userDefined": false
     },
     {
       "name": "keywordField",
-      "type": "keyword"
+      "type": "keyword",
+      "userDefined": false
     },
     {
       "name": "integerField",
-      "type": "integer"
+      "type": "integer",
+      "userDefined": false
     },
     {
       "name": "longField",
-      "type": "long"
+      "type": "long",
+      "userDefined": false
     },
     {
       "name": "textField",
-      "type": "text"
+      "type": "text",
+      "userDefined": false
     },
     {
       "name": "unsignedLongField",
-      "type": "unsigned_long"
+      "type": "unsigned_long",
+      "userDefined": false
     },
     {
       "name": "versionField",
-      "type": "version"
+      "type": "version",
+      "userDefined": false
     },
     {
       "name": "cartesianPointField",
-      "type": "cartesian_point"
+      "type": "cartesian_point",
+      "userDefined": false
     },
     {
       "name": "cartesianShapeField",
-      "type": "cartesian_shape"
+      "type": "cartesian_shape",
+      "userDefined": false
     },
     {
       "name": "geoPointField",
-      "type": "geo_point"
+      "type": "geo_point",
+      "userDefined": false
     },
     {
       "name": "geoShapeField",
-      "type": "geo_shape"
+      "type": "geo_shape",
+      "userDefined": false
     },
     {
       "name": "counterIntegerField",
-      "type": "counter_integer"
+      "type": "counter_integer",
+      "userDefined": false
     },
     {
       "name": "counterLongField",
-      "type": "counter_long"
+      "type": "counter_long",
+      "userDefined": false
     },
     {
       "name": "counterDoubleField",
-      "type": "counter_double"
+      "type": "counter_double",
+      "userDefined": false
     },
     {
       "name": "unsupportedField",
-      "type": "unsupported"
+      "type": "unsupported",
+      "userDefined": false
     },
     {
       "name": "dateNanosField",
-      "type": "date_nanos"
+      "type": "date_nanos",
+      "userDefined": false
     },
     {
       "name": "functionNamedParametersField",
-      "type": "function_named_parameters"
+      "type": "function_named_parameters",
+      "userDefined": false
     },
     {
       "name": "any#Char$Field",
-      "type": "double"
+      "type": "double",
+      "userDefined": false
     },
     {
       "name": "kubernetes.something.something",
-      "type": "double"
+      "type": "double",
+      "userDefined": false
     },
     {
       "name": "@timestamp",
-      "type": "date"
+      "type": "date",
+      "userDefined": false
     },
     {
       "name": "otherStringField",
-      "type": "keyword"
+      "type": "keyword",
+      "userDefined": false
     }
   ],
   "enrichFields": [
     {
       "name": "otherField",
-      "type": "text"
+      "type": "text",
+      "userDefined": false
     },
     {
       "name": "yetAnotherField",
-      "type": "double"
+      "type": "double",
+      "userDefined": false
     },
     {
       "name": "otherStringField",
-      "type": "keyword"
+      "type": "keyword",
+      "userDefined": false
     }
   ],
   "policies": [
@@ -146,7 +173,8 @@
   "unsupported_field": [
     {
       "name": "unsupported_field",
-      "type": "unsupported"
+      "type": "unsupported",
+      "userDefined": false
     }
   ],
   "testCases": [

From d7f56e4c65c1dbd09c013ed1a20f863d768dc328 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Fri, 5 Sep 2025 15:49:08 -0600
Subject: [PATCH 38/54] restore a small check

---
 .../src/shared/resources_helpers.ts                        | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index 6b74cb147f65a..7cd4d0925ef9b 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -95,7 +95,8 @@ export function getColumnsByTypeHelper(
 
   const cacheColumns = async () => {
     // in some cases (as in the case of ROW or SHOW) the query is not set
-    if (!originalQueryText) {
+    // TODO consider having source commands including ROW use columnsAfter methods like any other command
+    if (!queryForFields) {
       return;
     }
 
@@ -114,10 +115,6 @@ export function getColumnsByTypeHelper(
       subqueries.push(Builder.expression.query(root.commands.slice(0, i + 1)));
     }
 
-    if (!subqueries.length) {
-      return;
-    }
-
     // source command
     getFields(BasicPrettyPrinter.print(subqueries[0]));
 

From b6182a18701e6a8268babd3c0575e279aedd336e Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Mon, 8 Sep 2025 11:31:40 -0600
Subject: [PATCH 39/54] Fix ast package tests

---
 .../commands/enrich/autocomplete.ts           | 32 +++++++++++--------
 .../commands/from/autocomplete.test.ts        | 10 +++---
 .../commands/fuse/validate.test.ts            |  6 ++--
 .../commands/rename/validate.test.ts          |  4 +--
 4 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts
index 29ab1d1a69a96..1f1c937c3c774 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts
@@ -6,30 +6,30 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
+import { findFinalWord, findPreviousWord } from '../../../definitions/utils/autocomplete/helpers';
+import { buildFieldsDefinitions } from '../../../definitions/utils/functions';
+import { getOperatorSuggestions } from '../../../definitions/utils/operators';
+import { unescapeColumnName } from '../../../definitions/utils/shared';
 import type { ESQLCommand, ESQLSource } from '../../../types';
 import {
-  pipeCompleteItem,
   commaCompleteItem,
   getNewUserDefinedColumnSuggestion,
+  pipeCompleteItem,
 } from '../../complete_items';
-import { findFinalWord, findPreviousWord } from '../../../definitions/utils/autocomplete/helpers';
-import { unescapeColumnName } from '../../../definitions/utils/shared';
-import type { ESQLPolicy, ICommandCallbacks } from '../../types';
-import { type ISuggestionItem, type ICommandContext, Location } from '../../types';
+import { TRIGGER_SUGGESTION_COMMAND } from '../../constants';
+import type { ESQLColumnData, ESQLPolicy, ICommandCallbacks } from '../../types';
+import { Location, type ICommandContext, type ISuggestionItem } from '../../types';
 import {
   Position,
   buildMatchingFieldsDefinition,
+  buildPoliciesDefinitions,
+  getPolicyMetadata,
   getPosition,
   modeSuggestions,
   noPoliciesAvailableSuggestion,
   onSuggestion,
   withSuggestion,
-  getPolicyMetadata,
-  buildPoliciesDefinitions,
 } from './util';
-import { TRIGGER_SUGGESTION_COMMAND } from '../../constants';
-import { getOperatorSuggestions } from '../../../definitions/utils/operators';
-import { buildFieldsDefinitions } from '../../../definitions/utils/functions';
 
 export async function autocomplete(
   query: string,
@@ -41,8 +41,14 @@ export async function autocomplete(
   const innerText = query.substring(0, cursorPosition);
   const pos = getPosition(innerText, command);
   const policies = context?.policies ?? new Map<string, ESQLPolicy>();
-  const columnMap = context?.columns ?? new Map<string, string>();
-  const allColumnNames = Array.from(columnMap.keys());
+  const columnMap = context?.columns ?? new Map<string, ESQLColumnData>();
+
+  const fieldNames: string[] = [];
+  for (const name of columnMap.keys()) {
+    const col = columnMap.get(name);
+    if (col && !col.userDefined) fieldNames.push(name);
+  }
+
   const policyName = (
     command.args.find((arg) => !Array.isArray(arg) && arg.type === 'source') as
       | ESQLSource
@@ -107,7 +113,7 @@ export async function autocomplete(
         return [];
       }
 
-      return buildMatchingFieldsDefinition(policyMetadata.matchField, allColumnNames);
+      return buildMatchingFieldsDefinition(policyMetadata.matchField, fieldNames);
     }
 
     case Position.AFTER_ON_CLAUSE:
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts
index a1f708fc05d87..ff0b42f940e13 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts
@@ -6,14 +6,14 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { mockContext, integrations } from '../../../__tests__/context_fixtures';
-import { autocomplete } from './autocomplete';
 import { expectSuggestions, getFieldNamesByType } from '../../../__tests__/autocomplete';
-import type { ICommandCallbacks } from '../../types';
+import { integrations, mockContext } from '../../../__tests__/context_fixtures';
 import { correctQuerySyntax, findAstPosition } from '../../../definitions/utils/ast';
 import { parse } from '../../../parser';
-import { getRecommendedQueriesTemplates } from '../../options/recommended_queries';
 import { METADATA_FIELDS } from '../../options/metadata';
+import { getRecommendedQueriesTemplates } from '../../options/recommended_queries';
+import type { ICommandCallbacks } from '../../types';
+import { autocomplete } from './autocomplete';
 
 const metadataFields = [...METADATA_FIELDS].sort();
 
@@ -57,7 +57,7 @@ describe('FROM Autocomplete', () => {
       getByType: jest.fn(),
     };
 
-    const expectedFields = getFieldNamesByType('any');
+    const expectedFields = getFieldNamesByType('any', true);
     (mockCallbacks.getByType as jest.Mock).mockResolvedValue(
       expectedFields.map((name) => ({ label: name, text: name }))
     );
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts
index c9fcd81ba3c52..d2905ab15923a 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts
@@ -6,10 +6,10 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
+import { METADATA_FIELDS } from '../../../..';
 import { mockContext } from '../../../__tests__/context_fixtures';
-import { validate } from './validate';
 import { expectErrors } from '../../../__tests__/validation';
-import { METADATA_FIELDS } from '../../../..';
+import { validate } from './validate';
 
 const fuseExpectErrors = (query: string, expectedErrors: string[], context = mockContext) => {
   return expectErrors(query, expectedErrors, context, 'fuse', validate);
@@ -28,7 +28,7 @@ describe('FUSE Validation', () => {
       });
       const context = {
         ...mockContext,
-        fields: newColumns,
+        columns: newColumns,
       };
       fuseExpectErrors(
         `FROM index METADATA _id, _score, _index
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts
index e7df537c67600..f1ec956163339 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts
@@ -7,8 +7,8 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { mockContext } from '../../../__tests__/context_fixtures';
-import { validate } from './validate';
 import { expectErrors } from '../../../__tests__/validation';
+import { validate } from './validate';
 
 const renameExpectErrors = (query: string, expectedErrors: string[], context = mockContext) => {
   return expectErrors(query, expectedErrors, context, 'rename', validate);
@@ -63,6 +63,6 @@ describe('RENAME Validation', () => {
     renameExpectErrors('from a_index |eval doubleField + 1 | rename `doubleField + 1` as ', [
       'AS expected 2 arguments, but got 1.',
     ]);
-    renameExpectErrors('from a_index | rename key* as keywords', ['Unknown column "keywords"']);
+    renameExpectErrors('from a_index | rename key* as keywords', []);
   });
 });

From 38a8fc01e87f8406d5c5fe0c7cdc4b9ded0fae1c Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Mon, 8 Sep 2025 12:21:57 -0600
Subject: [PATCH 40/54] make source commands non-exceptional

---
 .../commands/from/columns_after.test.ts       | 28 +++++++++++
 .../commands/from/columns_after.ts            | 32 ++++++++++++
 .../commands_registry/commands/from/index.ts  |  2 +
 .../commands/row/columns_after.test.ts        | 50 +++++++++++++++++++
 .../commands/row/columns_after.ts             | 46 +++++++++++++++++
 .../commands_registry/commands/row/index.ts   |  2 +
 .../commands/show/columns_after.test.ts       | 28 +++++++++++
 .../commands/show/columns_after.ts            | 20 ++++++++
 .../commands_registry/commands/show/index.ts  |  2 +
 .../src/commands_registry/registry.ts         |  1 +
 .../src/shared/helpers.ts                     |  6 +++
 .../src/shared/resources_helpers.ts           | 38 ++++++--------
 .../__tests__/fields_and_variables.test.ts    | 32 ++++++------
 13 files changed, 246 insertions(+), 41 deletions(-)
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.test.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.test.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.test.ts
new file mode 100644
index 0000000000000..6cda1b953ea20
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.test.ts
@@ -0,0 +1,28 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type { ESQLFieldWithMetadata } from '../../types';
+import { columnsAfter } from './columns_after';
+
+describe('FROM columnsAfter', () => {
+  it('returns fields from the source', async () => {
+    const sourceFields: ESQLFieldWithMetadata[] = [
+      { name: 'field1', type: 'double', userDefined: false },
+      { name: 'field2', type: 'keyword', userDefined: false },
+    ];
+
+    const result = await columnsAfter({} as any, [], '', {
+      fromFrom: () => Promise.resolve(sourceFields),
+      fromJoin: () => Promise.resolve([]),
+      fromEnrich: () => Promise.resolve([]),
+    });
+
+    expect(result).toEqual(sourceFields);
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts
new file mode 100644
index 0000000000000..ccec97ff9d890
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts
@@ -0,0 +1,32 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+import type { ESQLFieldWithMetadata } from '../../types';
+
+// from https://github.com/elastic/elasticsearch/blob/da50c723a43262a9f46228ca36099647706c90dd/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/show/ShowInfo.java#L56-L58
+const SHOW_INFO_FIELDS: ESQLFieldWithMetadata[] = [
+  {
+    name: 'version',
+    type: 'keyword',
+    userDefined: false,
+  },
+  {
+    name: 'date',
+    type: 'keyword',
+    userDefined: false,
+  },
+  {
+    name: 'hash',
+    type: 'keyword',
+    userDefined: false,
+  },
+];
+
+export const columnsAfter = () => {
+  return [...SHOW_INFO_FIELDS];
+};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/index.ts
index b5551e073892d..16fd887e4a34c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/index.ts
@@ -11,10 +11,12 @@ import type { ICommandMethods } from '../../registry';
 import { autocomplete } from './autocomplete';
 import { validate } from './validate';
 import type { ICommandContext } from '../../types';
+import { columnsAfter } from './columns_after';
 
 const fromCommandMethods: ICommandMethods<ICommandContext> = {
   autocomplete,
   validate,
+  columnsAfter,
 };
 
 export const fromCommand = {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.test.ts
new file mode 100644
index 0000000000000..2f7d253913d9c
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.test.ts
@@ -0,0 +1,50 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { Parser } from '../../../..';
+import type { ESQLUserDefinedColumn } from '../../types';
+
+import { columnsAfter } from './columns_after';
+
+describe('ROW > columnsAfter', () => {
+  it('adds new user-defined columns', () => {
+    const queryString = `ROW baz = 1, foo = FLOOR(2. + 2.), TO_LONG(23 * 4)`;
+
+    // Can't use synth because it steps on the location information
+    // which is used to determine the name of the new column
+    const {
+      root: {
+        commands: [command],
+      },
+    } = Parser.parseQuery(queryString);
+
+    const result = columnsAfter(command, [], queryString);
+
+    expect(result).toEqual<ESQLUserDefinedColumn[]>([
+      {
+        name: 'baz',
+        type: 'integer',
+        location: { min: 4, max: 6 },
+        userDefined: true,
+      },
+      {
+        name: 'foo',
+        type: 'double',
+        location: { min: 13, max: 15 },
+        userDefined: true,
+      },
+      {
+        name: 'TO_LONG(23 * 4)',
+        type: 'long',
+        location: { min: 35, max: 49 },
+        userDefined: true,
+      },
+    ]);
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.ts
new file mode 100644
index 0000000000000..2ff89393eba11
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.ts
@@ -0,0 +1,46 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { isAssignment, isColumn } from '../../../ast/is';
+import { getExpressionType } from '../../../definitions/utils';
+import type { ESQLAstItem, ESQLCommand } from '../../../types';
+import type { ESQLColumnData, ESQLUserDefinedColumn } from '../../types';
+
+export const columnsAfter = (
+  command: ESQLCommand,
+  _previousColumns: ESQLColumnData[], // will always be empty for ROW
+  query: string
+) => {
+  const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, new Map());
+
+  const columns = [];
+
+  for (const expression of command.args) {
+    if (isAssignment(expression) && isColumn(expression.args[0])) {
+      const name = expression.args[0].parts.join('.');
+      const newColumn: ESQLUserDefinedColumn = {
+        name,
+        type: typeOf(expression.args[1]),
+        location: expression.args[0].location,
+        userDefined: true,
+      };
+      columns.push(newColumn);
+    } else if (!Array.isArray(expression)) {
+      const newColumn: ESQLUserDefinedColumn = {
+        name: query.substring(expression.location.min, expression.location.max + 1),
+        type: typeOf(expression),
+        location: expression.location,
+        userDefined: true,
+      };
+      columns.push(newColumn);
+    }
+  }
+
+  return columns;
+};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/index.ts
index 986c5932dbd57..9a454af9e2ac9 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/index.ts
@@ -11,10 +11,12 @@ import type { ICommandMethods } from '../../registry';
 import { autocomplete } from './autocomplete';
 import type { ICommandContext } from '../../types';
 import { validate } from './validate';
+import { columnsAfter } from './columns_after';
 
 const rowCommandMethods: ICommandMethods<ICommandContext> = {
   autocomplete,
   validate,
+  columnsAfter,
 };
 
 export const rowCommand = {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts
new file mode 100644
index 0000000000000..6cda1b953ea20
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts
@@ -0,0 +1,28 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type { ESQLFieldWithMetadata } from '../../types';
+import { columnsAfter } from './columns_after';
+
+describe('FROM columnsAfter', () => {
+  it('returns fields from the source', async () => {
+    const sourceFields: ESQLFieldWithMetadata[] = [
+      { name: 'field1', type: 'double', userDefined: false },
+      { name: 'field2', type: 'keyword', userDefined: false },
+    ];
+
+    const result = await columnsAfter({} as any, [], '', {
+      fromFrom: () => Promise.resolve(sourceFields),
+      fromJoin: () => Promise.resolve([]),
+      fromEnrich: () => Promise.resolve([]),
+    });
+
+    expect(result).toEqual(sourceFields);
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts
new file mode 100644
index 0000000000000..6e5e8502964fa
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts
@@ -0,0 +1,20 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+import { type ESQLCommand } from '../../../types';
+import type { ESQLColumnData } from '../../types';
+import type { IAdditionalFields } from '../../registry';
+
+export const columnsAfter = (
+  command: ESQLCommand,
+  _previousColumns: ESQLColumnData[], // will always be empty for FROM
+  _query: string,
+  additionalFields: IAdditionalFields
+) => {
+  return additionalFields.fromFrom(command);
+};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/index.ts
index dc8e0b16fe35a..0f930f7fc9064 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/index.ts
@@ -10,9 +10,11 @@ import { i18n } from '@kbn/i18n';
 import type { ICommandMethods } from '../../registry';
 import { autocomplete } from './autocomplete';
 import type { ICommandContext } from '../../types';
+import { columnsAfter } from './columns_after';
 
 const showCommandMethods: ICommandMethods<ICommandContext> = {
   autocomplete,
+  columnsAfter,
 };
 
 export const showCommand = {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
index 6329572156d50..591bb3b484ff1 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts
@@ -123,6 +123,7 @@ export interface ICommandRegistry {
 export interface IAdditionalFields {
   fromJoin: (cmd: ESQLCommand) => Promise<ESQLFieldWithMetadata[]>;
   fromEnrich: (cmd: ESQLCommand) => Promise<ESQLFieldWithMetadata[]>;
+  fromFrom: (cmd: ESQLCommand) => Promise<ESQLFieldWithMetadata[]>;
 }
 
 /**
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index 9115c0b84bb8a..21663bd88141d 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -7,6 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import {
+  BasicPrettyPrinter,
   esqlCommandRegistry,
   isSource,
   mutate,
@@ -145,9 +146,14 @@ export async function getCurrentQueryAvailableColumns(
     return Promise.resolve([]);
   };
 
+  const getFromFields = (command: ESQLAstCommand): Promise<ESQLFieldWithMetadata[]> => {
+    return fetchFields(BasicPrettyPrinter.command(command));
+  };
+
   const additionalFields: IAdditionalFields = {
     fromJoin: getJoinFields,
     fromEnrich: getEnrichFields,
+    fromFrom: getFromFields,
   };
 
   if (commandDef?.methods.columnsAfter) {
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index 7cd4d0925ef9b..162fcc7370fb0 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -60,20 +60,21 @@ async function cacheColumnsForQuery(
     return;
   }
 
-  const fieldsAvailableAfterPreviousCommand = getValueInsensitive(
-    BasicPrettyPrinter.print({ ...query, commands: query.commands.slice(0, -1) })
+  const queryBeforeCurrentCommand = BasicPrettyPrinter.print({
+    ...query,
+    commands: query.commands.slice(0, -1),
+  });
+  const fieldsAvailableAfterPreviousCommand = getValueInsensitive(queryBeforeCurrentCommand) ?? [];
+
+  const availableFields = await getCurrentQueryAvailableColumns(
+    query.commands,
+    fieldsAvailableAfterPreviousCommand,
+    fetchFields,
+    policies,
+    originalQueryText
   );
 
-  if (fieldsAvailableAfterPreviousCommand && fieldsAvailableAfterPreviousCommand?.length) {
-    const availableFields = await getCurrentQueryAvailableColumns(
-      query.commands,
-      fieldsAvailableAfterPreviousCommand,
-      fetchFields,
-      policies,
-      originalQueryText
-    );
-    cache.set(cacheKey, availableFields);
-  }
+  cache.set(cacheKey, availableFields);
 }
 
 /**
@@ -94,8 +95,6 @@ export function getColumnsByTypeHelper(
   const root = EsqlQuery.fromSrc(queryForFields).ast;
 
   const cacheColumns = async () => {
-    // in some cases (as in the case of ROW or SHOW) the query is not set
-    // TODO consider having source commands including ROW use columnsAfter methods like any other command
     if (!queryForFields) {
       return;
     }
@@ -115,14 +114,11 @@ export function getColumnsByTypeHelper(
       subqueries.push(Builder.expression.query(root.commands.slice(0, i + 1)));
     }
 
-    // source command
-    getFields(BasicPrettyPrinter.print(subqueries[0]));
-
     const policies = (await resourceRetriever?.getPolicies?.()) ?? [];
     const policyMap = new Map(policies.map((p) => [p.name, p]));
 
     // build fields cache for every partial query
-    for (const subquery of subqueries.slice(1)) {
+    for (const subquery of subqueries) {
       await cacheColumnsForQuery(subquery, getFields, policyMap, originalQueryText);
     }
   };
@@ -231,11 +227,7 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
     }
   }
 
-  // If there is only one source command and it does not require fields, do not
-  // fetch fields, hence return an empty string.
-  return commands.length === 1 && ['row', 'show'].includes(commands[0].name)
-    ? ''
-    : buildQueryUntilPreviousCommand(root);
+  return buildQueryUntilPreviousCommand(root);
 }
 
 function buildQueryUntilPreviousCommand(root: ESQLAstQueryExpression) {
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/fields_and_variables.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/fields_and_variables.test.ts
index 9d62dcb95da83..4d99dc0d172b7 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/fields_and_variables.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/fields_and_variables.test.ts
@@ -16,7 +16,7 @@ import { Location } from '@kbn/esql-ast/src/commands_registry/types';
 import { setTestFunctions } from '@kbn/esql-ast/src/definitions/utils/test_functions';
 import { setup } from './helpers';
 
-describe('field and userDefinedColumn escaping', () => {
+describe('column escaping', () => {
   it('recognizes escaped fields', async () => {
     const { expectErrors } = await setup();
     // command level
@@ -39,7 +39,7 @@ describe('field and userDefinedColumn escaping', () => {
     );
   });
 
-  it('recognizes escaped userDefinedColumns', async () => {
+  it('recognizes escaped user-defined columns', async () => {
     const { expectErrors } = await setup();
     // command level
     await expectErrors('ROW `var$iable` = 1 | EVAL `var$iable`', []);
@@ -56,12 +56,12 @@ describe('field and userDefinedColumn escaping', () => {
       []
     );
 
-    // expression userDefinedColumn
+    // expression user-defined column
     await expectErrors('FROM index | EVAL doubleField + 20 | EVAL `doubleField + 20`', []);
     await expectErrors('ROW 21 + 20 | STATS AVG(`21 + 20`)', []);
   });
 
-  it('recognizes userDefinedColumns with spaces and comments', async () => {
+  it('recognizes user-defined columns with spaces and comments', async () => {
     const { expectErrors } = await setup();
     // command level
     await expectErrors(
@@ -114,13 +114,13 @@ describe('field and userDefinedColumn escaping', () => {
   });
 });
 
-describe('userDefinedColumn support', () => {
-  describe('userDefinedColumn data type detection', () => {
+describe('user-defined column support', () => {
+  describe('user-defined column data type detection', () => {
     // most of these tests are aspirational (and skipped) because we don't have
     // a good way to compute the type of an expression yet.
     beforeAll(() => {
       setTestFunctions([
-        // this test function is just used to test the type of the userDefinedColumn
+        // this test function is just used to test the type of the user-defined column
         {
           type: FunctionDefinitionTypes.SCALAR,
           description: 'Test function',
@@ -131,7 +131,7 @@ describe('userDefinedColumn support', () => {
           ],
         },
         // this test function is used to check that the correct return type is used
-        // when determining userDefinedColumn types
+        // when determining user-defined column types
         {
           type: FunctionDefinitionTypes.SCALAR,
           description: 'Test function',
@@ -157,10 +157,9 @@ describe('userDefinedColumn support', () => {
     });
 
     const expectType = (type: FunctionParameterType) =>
-      `Argument of [test] must be [cartesian_point], found value [var] type [${type}]`;
+      getNoValidCallSignatureError('test', [type]);
 
-    // @todo unskip after https://github.com/elastic/kibana/issues/195682
-    test.skip('literals', async () => {
+    test('literals', async () => {
       const { expectErrors } = await setup();
       // literal assignment
       await expectErrors('FROM index | EVAL var = 1, TEST(var)', [expectType('integer')]);
@@ -176,16 +175,14 @@ describe('userDefinedColumn support', () => {
       ]);
     });
 
-    // @todo unskip after https://github.com/elastic/kibana/issues/195682
-    test.skip('userDefinedColumns', async () => {
+    test('user-defined columns', async () => {
       const { expectErrors } = await setup();
       await expectErrors('FROM index | EVAL var = textField, col2 = var, TEST(col2)', [
-        `Argument of [test] must be [cartesian_point], found value [col2] type [text]`,
+        getNoValidCallSignatureError('test', ['text']),
       ]);
     });
 
-    // @todo unskip after https://github.com/elastic/kibana/issues/195682
-    test.skip('inline casting', async () => {
+    test('inline casting', async () => {
       const { expectErrors } = await setup();
       // inline cast assignment
       await expectErrors('FROM index | EVAL var = doubleField::long, TEST(var)', [
@@ -197,8 +194,7 @@ describe('userDefinedColumn support', () => {
       ]);
     });
 
-    // @todo unskip after https://github.com/elastic/kibana/issues/195682
-    test.skip('function results', async () => {
+    test('function results', async () => {
       const { expectErrors } = await setup();
       // function assignment
       await expectErrors('FROM index | EVAL var = RETURN_VALUE(doubleField), TEST(var)', [

From b1325c3183e80d19b49ff8073fdac82302450f3f Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Mon, 8 Sep 2025 13:42:20 -0600
Subject: [PATCH 41/54] swap from and show

---
 .../commands/from/columns_after.ts            | 32 ++++++-------------
 .../commands/show/columns_after.test.ts       | 20 +++---------
 .../commands/show/columns_after.ts            | 32 +++++++++++++------
 3 files changed, 37 insertions(+), 47 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts
index ccec97ff9d890..6e5e8502964fa 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts
@@ -6,27 +6,15 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import type { ESQLFieldWithMetadata } from '../../types';
+import { type ESQLCommand } from '../../../types';
+import type { ESQLColumnData } from '../../types';
+import type { IAdditionalFields } from '../../registry';
 
-// from https://github.com/elastic/elasticsearch/blob/da50c723a43262a9f46228ca36099647706c90dd/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/show/ShowInfo.java#L56-L58
-const SHOW_INFO_FIELDS: ESQLFieldWithMetadata[] = [
-  {
-    name: 'version',
-    type: 'keyword',
-    userDefined: false,
-  },
-  {
-    name: 'date',
-    type: 'keyword',
-    userDefined: false,
-  },
-  {
-    name: 'hash',
-    type: 'keyword',
-    userDefined: false,
-  },
-];
-
-export const columnsAfter = () => {
-  return [...SHOW_INFO_FIELDS];
+export const columnsAfter = (
+  command: ESQLCommand,
+  _previousColumns: ESQLColumnData[], // will always be empty for FROM
+  _query: string,
+  additionalFields: IAdditionalFields
+) => {
+  return additionalFields.fromFrom(command);
 };
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts
index 6cda1b953ea20..779cf69419ccc 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts
@@ -7,22 +7,12 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import type { ESQLFieldWithMetadata } from '../../types';
-import { columnsAfter } from './columns_after';
+import { SHOW_INFO_FIELDS, columnsAfter } from './columns_after';
 
-describe('FROM columnsAfter', () => {
-  it('returns fields from the source', async () => {
-    const sourceFields: ESQLFieldWithMetadata[] = [
-      { name: 'field1', type: 'double', userDefined: false },
-      { name: 'field2', type: 'keyword', userDefined: false },
-    ];
+describe('SHOW columnsAfter', () => {
+  it('returns info', async () => {
+    const result = await columnsAfter();
 
-    const result = await columnsAfter({} as any, [], '', {
-      fromFrom: () => Promise.resolve(sourceFields),
-      fromJoin: () => Promise.resolve([]),
-      fromEnrich: () => Promise.resolve([]),
-    });
-
-    expect(result).toEqual(sourceFields);
+    expect(result).toEqual(SHOW_INFO_FIELDS);
   });
 });
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts
index 6e5e8502964fa..22e129fbdca58 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts
@@ -6,15 +6,27 @@
  * your election, the "Elastic License 2.0", the "GNU Affero General Public
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
-import { type ESQLCommand } from '../../../types';
-import type { ESQLColumnData } from '../../types';
-import type { IAdditionalFields } from '../../registry';
+import type { ESQLFieldWithMetadata } from '../../types';
 
-export const columnsAfter = (
-  command: ESQLCommand,
-  _previousColumns: ESQLColumnData[], // will always be empty for FROM
-  _query: string,
-  additionalFields: IAdditionalFields
-) => {
-  return additionalFields.fromFrom(command);
+// from https://github.com/elastic/elasticsearch/blob/da50c723a43262a9f46228ca36099647706c90dd/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/show/ShowInfo.java#L56-L58
+export const SHOW_INFO_FIELDS: ESQLFieldWithMetadata[] = [
+  {
+    name: 'version',
+    type: 'keyword',
+    userDefined: false,
+  },
+  {
+    name: 'date',
+    type: 'keyword',
+    userDefined: false,
+  },
+  {
+    name: 'hash',
+    type: 'keyword',
+    userDefined: false,
+  },
+];
+
+export const columnsAfter = () => {
+  return [...SHOW_INFO_FIELDS];
 };

From e37f27e5afabe2a4ddf5fe0daf768e032da9c1be Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Mon, 8 Sep 2025 14:56:46 -0600
Subject: [PATCH 42/54] fix mo tests

---
 .../utils/ecs_metadata_helper.test.ts         | 56 +++++--------------
 .../src/shared/helpers.ts                     | 11 ++--
 .../src/shared/resources_helpers.test.ts      |  5 --
 .../src/shared/resources_helpers.ts           | 24 ++++++--
 .../__tests__/column_existence.test.ts        | 10 ----
 .../src/validation/validation.test.ts         | 11 ----
 6 files changed, 40 insertions(+), 77 deletions(-)
 delete mode 100644 src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_existence.test.ts

diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
index 37e706e0aca0f..47493ec786800 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts
@@ -29,30 +29,17 @@ describe('enrichFieldsWithECSInfo', () => {
       { name: 'ecs.fakeBooleanField', type: 'boolean', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
     ];
-    const fieldsMetadata = {
-      getClient: jest.fn().mockResolvedValue({
-        find: jest.fn().mockResolvedValue({
-          fields: {
-            'ecs.version': { description: 'ECS version field', type: 'keyword' },
-            'ecs.field': { description: 'ECS field description', type: 'text' },
-            'ecs.fakeBooleanField': {
-              description: 'ECS fake boolean field description',
-              type: 'keyword',
-            },
-          },
-        }),
-      }),
+
+    const fieldsMetadataCache: ECSMetadata = {
+      'ecs.version': { description: 'ECS version field', type: 'keyword' },
+      'ecs.field': { description: 'ECS field description', type: 'text' },
+      'ecs.fakeBooleanField': {
+        description: 'ECS fake boolean field description',
+        type: 'keyword',
+      },
     };
-    const fieldsMetadataCache = await (
-      await fieldsMetadata.getClient()
-    ).find({
-      attributes: ['type'],
-    });
 
-    const result = await enrichFieldsWithECSInfo(
-      columns,
-      fieldsMetadataCache.fields as ECSMetadata
-    );
+    const result = enrichFieldsWithECSInfo(columns, fieldsMetadataCache);
 
     expect(result).toEqual([
       {
@@ -61,8 +48,8 @@ describe('enrichFieldsWithECSInfo', () => {
         isEcs: true,
         userDefined: false,
       },
-      { name: 'ecs.fakeBooleanField', type: 'boolean', isEcs: true, userDefined: false },
-      { name: 'field2', type: 'double', isEcs: false, userDefined: false },
+      { name: 'ecs.fakeBooleanField', type: 'boolean', userDefined: false },
+      { name: 'field2', type: 'double', userDefined: false },
     ]);
   });
 
@@ -72,25 +59,12 @@ describe('enrichFieldsWithECSInfo', () => {
       { name: 'ecs.version.keyword', type: 'keyword', userDefined: false },
       { name: 'field2', type: 'double', userDefined: false },
     ];
-    const fieldsMetadata = {
-      getClient: jest.fn().mockResolvedValue({
-        find: jest.fn().mockResolvedValue({
-          fields: {
-            'ecs.version': { description: 'ECS version field', type: 'keyword' },
-          } as unknown as ECSMetadata,
-        }),
-      }),
+
+    const fieldsMetadataCache: ECSMetadata = {
+      'ecs.version': { description: 'ECS version field', type: 'keyword' },
     };
 
-    const fieldsMetadataCache = await (
-      await fieldsMetadata.getClient()
-    ).find({
-      attributes: ['type'],
-    });
-    const result = await enrichFieldsWithECSInfo(
-      columns,
-      fieldsMetadataCache.fields as ECSMetadata
-    );
+    const result = enrichFieldsWithECSInfo(columns, fieldsMetadataCache);
 
     expect(result).toEqual<ESQLColumnData[]>([
       { name: 'ecs.version', type: 'keyword', isEcs: true, userDefined: false },
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
index 21663bd88141d..8d77fb6432707 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts
@@ -110,7 +110,7 @@ export async function getCurrentQueryAvailableColumns(
   commands: ESQLAstCommand[],
   previousPipeFields: ESQLColumnData[],
   fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>,
-  policies: Map<string, ESQLPolicy>,
+  getPolicies: () => Promise<Map<string, ESQLPolicy>>,
   originalQueryText: string
 ) {
   const lastCommand = commands[commands.length - 1];
@@ -129,13 +129,16 @@ export async function getCurrentQueryAvailableColumns(
     return Promise.resolve([]);
   };
 
-  const getEnrichFields = (command: ESQLAstCommand): Promise<ESQLFieldWithMetadata[]> => {
+  const getEnrichFields = async (command: ESQLAstCommand): Promise<ESQLFieldWithMetadata[]> => {
     if (!isSource(command.args[0])) {
-      return Promise.resolve([]);
+      return [];
     }
 
     const policyName = command.args[0].name;
+
+    const policies = await getPolicies();
     const policy = policies.get(policyName);
+
     if (policy) {
       const fieldsQuery = `FROM ${policy.sourceIndices.join(
         ', '
@@ -143,7 +146,7 @@ export async function getCurrentQueryAvailableColumns(
       return fetchFields(fieldsQuery);
     }
 
-    return Promise.resolve([]);
+    return [];
   };
 
   const getFromFields = (command: ESQLAstCommand): Promise<ESQLFieldWithMetadata[]> => {
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts
index 742572c1c0e53..b623e75087d0e 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts
@@ -59,9 +59,4 @@ describe('getQueryForFields', () => {
     | EVAL foo = 1, ${EDITOR_MARKER}`;
     assert(query3, 'FROM index | EVAL foo = 1');
   });
-
-  it('should return empty string if non-FROM source command', () => {
-    assert('ROW field1 = 1', '');
-    assert('SHOW INFO', '');
-  });
 });
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
index 162fcc7370fb0..01151fec9ac14 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts
@@ -50,10 +50,20 @@ function getValueInsensitive(keyToCheck: string) {
 async function cacheColumnsForQuery(
   query: ESQLAstQueryExpression,
   fetchFields: (query: string) => Promise<ESQLFieldWithMetadata[]>,
-  policies: Map<string, ESQLPolicy>,
+  getPolicies: () => Promise<Map<string, ESQLPolicy>>,
   originalQueryText: string
 ) {
-  const cacheKey = BasicPrettyPrinter.print(query);
+  let cacheKey: string;
+  try {
+    cacheKey = BasicPrettyPrinter.print(query);
+  } catch {
+    // for some syntactically incorrect queries
+    // the printer will throw. They're incorrect
+    // anyways, so just move on — ANTLR errors will
+    // be reported.
+    return;
+  }
+
   const existsInCache = checkCacheInsensitive(cacheKey);
   if (existsInCache) {
     // this is already in the cache
@@ -70,7 +80,7 @@ async function cacheColumnsForQuery(
     query.commands,
     fieldsAvailableAfterPreviousCommand,
     fetchFields,
-    policies,
+    getPolicies,
     originalQueryText
   );
 
@@ -114,12 +124,14 @@ export function getColumnsByTypeHelper(
       subqueries.push(Builder.expression.query(root.commands.slice(0, i + 1)));
     }
 
-    const policies = (await resourceRetriever?.getPolicies?.()) ?? [];
-    const policyMap = new Map(policies.map((p) => [p.name, p]));
+    const getPolicies = async () => {
+      const policies = (await resourceRetriever?.getPolicies?.()) ?? [];
+      return new Map(policies.map((p) => [p.name, p]));
+    };
 
     // build fields cache for every partial query
     for (const subquery of subqueries) {
-      await cacheColumnsForQuery(subquery, getFields, policyMap, originalQueryText);
+      await cacheColumnsForQuery(subquery, getFields, getPolicies, originalQueryText);
     }
   };
 
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_existence.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_existence.test.ts
deleted file mode 100644
index c42459d577729..0000000000000
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_existence.test.ts
+++ /dev/null
@@ -1,10 +0,0 @@
-/*
- * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
- * Public License v 1"; you may not use this file except in compliance with, at
- * your election, the "Elastic License 2.0", the "GNU Affero General Public
- * License v3.0 only", or the "Server Side Public License, v 1".
- */
-
-// TODO
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts
index f26b20eaaf3a4..a3ceb6df4fda2 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts
@@ -601,17 +601,6 @@ describe('validation logic', () => {
         expect(callbackMocks.getColumnsFor).toHaveBeenCalledTimes(0);
       });
 
-      it('should call fields callbacks also for show command', async () => {
-        const callbackMocks = getCallbackMocks();
-        await validateQuery(`show info | keep name`, undefined, callbackMocks);
-        expect(callbackMocks.getSources).not.toHaveBeenCalled();
-        expect(callbackMocks.getPolicies).not.toHaveBeenCalled();
-        expect(callbackMocks.getColumnsFor).toHaveBeenCalledTimes(1);
-        expect(callbackMocks.getColumnsFor).toHaveBeenLastCalledWith({
-          query: 'show info',
-        });
-      });
-
       it(`should fetch additional fields if an enrich command is found`, async () => {
         const callbackMocks = getCallbackMocks();
         await validateQuery(

From cd77e169856b0a16ac77add293a2c5a1807f8ea0 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Mon, 8 Sep 2025 14:57:59 -0600
Subject: [PATCH 43/54] remove comment

---
 .../packages/private/kbn-esql-editor/src/esql_editor.tsx         | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx b/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
index 8e133f3113716..bda5dc608ebbd 100644
--- a/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
+++ b/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx
@@ -506,7 +506,6 @@ const ESQLEditorInternal = function ESQLEditor({
       },
       getPolicies: async () => {
         try {
-          // TODO cache?
           const policies = (await core.http.get(
             `/internal/index_management/enrich_policies`
           )) as SerializedEnrichPolicy[];

From dc7698c45943dd14954319b058fe527cfae51b0a Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Mon, 8 Sep 2025 15:25:38 -0600
Subject: [PATCH 44/54] fix type issue

---
 .../commands/enrich/columns_after.test.ts          |  5 +++++
 .../commands/fork/columns_after.test.ts            | 14 ++++++++++++++
 .../commands/join/columns_after.test.ts            |  3 +++
 .../src/languages/esql/lib/hover/hover.ts          |  4 +---
 4 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
index 8f2d1d7fa4bd3..0838236692969 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts
@@ -20,6 +20,7 @@ describe('ENRICH columnsAfter', () => {
     const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
       fromJoin: () => Promise.resolve([]),
       fromEnrich: () => Promise.resolve([]),
+      fromFrom: () => Promise.resolve([]),
     });
     expect(result).toEqual(previousColumns);
   });
@@ -36,6 +37,7 @@ describe('ENRICH columnsAfter', () => {
     const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
       fromJoin: () => Promise.resolve([]),
       fromEnrich: () => Promise.resolve(enrichColumns),
+      fromFrom: () => Promise.resolve([]),
     });
     expect(result).toEqual([...enrichColumns, ...previousColumns]);
   });
@@ -56,6 +58,7 @@ describe('ENRICH columnsAfter', () => {
       {
         fromJoin: () => Promise.resolve([]),
         fromEnrich: () => Promise.resolve(enrichColumns),
+        fromFrom: () => Promise.resolve([]),
       }
     );
     expect(result).toEqual([enrichColumns[1], ...previousColumns]);
@@ -77,6 +80,7 @@ describe('ENRICH columnsAfter', () => {
       {
         fromEnrich: () => Promise.resolve(enrichColumns),
         fromJoin: () => Promise.resolve([]),
+        fromFrom: () => Promise.resolve([]),
       }
     );
     const expected = [
@@ -99,6 +103,7 @@ describe('ENRICH columnsAfter', () => {
     const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', {
       fromEnrich: () => Promise.resolve(enrichFields),
       fromJoin: () => Promise.resolve([]),
+      fromFrom: () => Promise.resolve([]),
     });
     expect(result).toEqual([
       { name: 'fieldA', type: 'text', userDefined: false },
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
index c705e048ae65f..5d0c8803d1ed8 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts
@@ -24,6 +24,7 @@ describe('FORK', () => {
       {
         fromEnrich: () => Promise.resolve([]),
         fromJoin: () => Promise.resolve([]),
+        fromFrom: () => Promise.resolve([]),
       }
     );
 
@@ -51,6 +52,7 @@ describe('FORK', () => {
       {
         fromEnrich: () => Promise.resolve([]),
         fromJoin: () => Promise.resolve([]),
+        fromFrom: () => Promise.resolve([]),
       }
     );
 
@@ -94,6 +96,14 @@ describe('FORK', () => {
               userDefined: false,
             },
           ]),
+        fromFrom: () =>
+          Promise.resolve([
+            {
+              name: 'from-from',
+              type: 'keyword',
+              userDefined: false,
+            },
+          ]),
       }
     );
 
@@ -125,6 +135,7 @@ describe('FORK', () => {
       {
         fromEnrich: () => Promise.resolve([]),
         fromJoin: () => Promise.resolve([]),
+        fromFrom: () => Promise.resolve([]),
       }
     );
 
@@ -154,6 +165,7 @@ describe('FORK', () => {
         {
           fromEnrich: () => Promise.resolve([]),
           fromJoin: () => Promise.resolve([]),
+          fromFrom: () => Promise.resolve([]),
         }
       );
 
@@ -181,6 +193,7 @@ describe('FORK', () => {
         {
           fromEnrich: () => Promise.resolve([]),
           fromJoin: () => Promise.resolve([]),
+          fromFrom: () => Promise.resolve([]),
         }
       );
 
@@ -207,6 +220,7 @@ describe('FORK', () => {
         {
           fromEnrich: () => Promise.resolve([]),
           fromJoin: () => Promise.resolve([]),
+          fromFrom: () => Promise.resolve([]),
         }
       );
 
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
index 1aeeef826a723..3b2d34c78b92c 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts
@@ -20,6 +20,7 @@ describe('JOIN columnsAfter', () => {
     const result = await columnsAfter({} as any, previousColumns, '', {
       fromJoin: () => Promise.resolve([]),
       fromEnrich: () => Promise.resolve([]),
+      fromFrom: () => Promise.resolve([]),
     });
 
     expect(result).toEqual(previousColumns);
@@ -38,6 +39,7 @@ describe('JOIN columnsAfter', () => {
     const result = await columnsAfter({} as any, previousColumns, '', {
       fromJoin: () => Promise.resolve(joinColumns),
       fromEnrich: () => Promise.resolve([]),
+      fromFrom: () => Promise.resolve([]),
     });
 
     expect(result).toEqual([...joinColumns, ...previousColumns]);
@@ -56,6 +58,7 @@ describe('JOIN columnsAfter', () => {
     const result = await columnsAfter({} as any, previousColumns, '', {
       fromJoin: () => Promise.resolve(joinColumns),
       fromEnrich: () => Promise.resolve([]),
+      fromFrom: () => Promise.resolve([]),
     });
 
     expect(result).toEqual([
diff --git a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
index e134f2db820e1..fa81edfe2df7b 100644
--- a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
+++ b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts
@@ -26,7 +26,6 @@ import {
 } from '@kbn/esql-ast/src/types';
 import { type ESQLCallbacks } from '@kbn/esql-validation-autocomplete';
 import { getColumnsByTypeRetriever } from '@kbn/esql-validation-autocomplete/src/autocomplete/autocomplete';
-import { getQueryForFields } from '@kbn/esql-validation-autocomplete/src/autocomplete/helper';
 import { getPolicyHelper } from '@kbn/esql-validation-autocomplete/src/shared/resources_helpers';
 import { i18n } from '@kbn/i18n';
 import type { monaco } from '../../../../monaco_imports';
@@ -186,8 +185,7 @@ async function getHintForFunctionArg(
   offset: number,
   resourceRetriever?: ESQLCallbacks
 ) {
-  const queryForColumns = getQueryForFields(query, root);
-  const { getColumnMap } = getColumnsByTypeRetriever(queryForColumns, resourceRetriever);
+  const { getColumnMap } = getColumnsByTypeRetriever(root, query, resourceRetriever);
 
   const fnDefinition = getFunctionDefinition(fnNode.name);
   // early exit on no hit

From f5437e937cf110bca99f1dba74ba056ee2b3baa3 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 9 Sep 2025 11:51:00 -0600
Subject: [PATCH 45/54] fix inlinestats

---
 .../inlinestats/columns_after.test.ts         | 40 +++++++++++++++++++
 .../commands/inlinestats/columns_after.ts     | 22 ++++++++++
 .../commands/inlinestats/index.ts             |  2 +-
 .../commands/stats/columns_after.ts           |  3 +-
 4 files changed, 64 insertions(+), 3 deletions(-)
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
new file mode 100644
index 0000000000000..8029d1d7fed2e
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
@@ -0,0 +1,40 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+import { Parser } from '../../../..';
+import type { ESQLColumnData } from '../../types';
+import { type ESQLFieldWithMetadata } from '../../types';
+import { columnsAfter } from './columns_after';
+
+describe('INLINESTATS', () => {
+  it('gets the columns after the query', () => {
+    const previousCommandFields: ESQLFieldWithMetadata[] = [
+      { name: 'field1', type: 'double', userDefined: false },
+      { name: 'buckets', type: 'double', userDefined: false }, // should be overwritten
+      { name: '@timestamp', type: 'date', userDefined: false },
+    ];
+
+    const queryString = `FROM a | STATS AVG(field1) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`;
+
+    // Can't use synth because it steps on the location information
+    // which is used to determine the name of the new column
+    const {
+      root: {
+        commands: [, command],
+      },
+    } = Parser.parseQuery(queryString);
+
+    const result = columnsAfter(command, previousCommandFields, queryString);
+
+    expect(result).toEqual<ESQLColumnData[]>([
+      { name: 'AVG(field1)', type: 'double', userDefined: true, location: { min: 15, max: 25 } },
+      { name: 'buckets', type: 'date', userDefined: true, location: { min: 30, max: 36 } },
+      ...previousCommandFields.filter(({ name }) => name !== 'buckets'),
+    ]);
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
new file mode 100644
index 0000000000000..b599f132cbb68
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts
@@ -0,0 +1,22 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+import uniqBy from 'lodash/uniqBy';
+import type { ESQLCommand } from '../../../types';
+import { columnsAfter as columnsAfterStats } from '../stats/columns_after';
+import type { ESQLColumnData } from '../../types';
+
+export const columnsAfter = (
+  command: ESQLCommand,
+  previousColumns: ESQLColumnData[],
+  query: string
+) => {
+  const newColumns = columnsAfterStats(command, previousColumns, query);
+
+  return uniqBy([...newColumns, ...previousColumns], 'name');
+};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/index.ts
index 23ad88f522354..473779526c9aa 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/index.ts
@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import { i18n } from '@kbn/i18n';
-import { columnsAfter } from '../stats/columns_after';
+import { columnsAfter } from './columns_after';
 import { autocomplete } from '../stats/autocomplete';
 import { validate } from '../stats/validate';
 import type { ICommandContext } from '../../types';
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
index 8ebf756af6647..902b97406aa3d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts
@@ -18,7 +18,7 @@ const getUserDefinedColumns = (
   typeOf: (thing: ESQLAstItem) => SupportedDataType | 'unknown',
   query: string
 ): ESQLUserDefinedColumn[] => {
-  const columns = [];
+  const columns: ESQLUserDefinedColumn[] = [];
 
   for (const expression of command.args) {
     if (isAssignment(expression) && isColumn(expression.args[0])) {
@@ -63,6 +63,5 @@ export const columnsAfter = (
 
   const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, columnMap);
 
-  // TODO - is this uniqby helpful? Does it do what we expect?
   return uniqBy([...getUserDefinedColumns(command, typeOf, query)], 'name');
 };

From c838166a0c80a54c70d249785db89353af76d521 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 9 Sep 2025 12:02:46 -0600
Subject: [PATCH 46/54] remove superfluous validation

---
 .../commands_registry/commands/completion/validate.ts | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
index 307046501185a..19408c4857dd5 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
@@ -39,17 +39,6 @@ export const validate = (
     messages.push(errors.byId('inferenceIdRequired', command.location, { command: 'COMPLETION' }));
   }
 
-  const targetName = targetField?.name || 'completion';
-
-  // Sets the target field so the column is recognized after the command is applied
-  // @TODO can we remove now?
-  context?.columns.set(targetName, {
-    name: targetName,
-    location: targetField?.location || command.location,
-    type: 'keyword',
-    userDefined: true,
-  });
-
   messages.push(...validateCommandArguments(command, ast, context, callbacks));
 
   return messages;

From a6857e6a7f6abb73a76c47b2b4bfec62794bd31a Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 9 Sep 2025 12:14:35 -0600
Subject: [PATCH 47/54] PR feedback

---
 .../src/commands_registry/commands/rename/columns_after.ts  | 2 +-
 .../src/autocomplete/__tests__/helpers.ts                   | 6 +++---
 .../src/validation/validation.ts                            | 2 --
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
index f5f5b236eda0e..31b3915a54744 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts
@@ -60,5 +60,5 @@ export const columnsAfter = (
     return oldColumn; // No rename found, keep the old name
   });
 
-  return uniqBy(newFields, 'name') as ESQLColumnData[];
+  return uniqBy(newFields, 'name');
 };
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts
index 68f91fe9f3b64..74ad4ea8e47c2 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts
@@ -46,20 +46,20 @@ export const fields: TestField[] = [
   ...fieldTypes.map((type) => ({
     name: `${camelCase(type)}Field`,
     type,
-    userDefined: false as false,
+    userDefined: false as const,
     // suggestedAs is optional and omitted here
   })),
   {
     name: 'any#Char$Field',
     type: 'double',
     suggestedAs: '`any#Char$Field`',
-    userDefined: false as false,
+    userDefined: false as const,
   },
   {
     name: 'kubernetes.something.something',
     type: 'double',
     suggestedAs: undefined,
-    userDefined: false as false,
+    userDefined: false as const,
   },
 ];
 
diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
index 7846f8c71d7a8..38d78fc552cd6 100644
--- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
+++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts
@@ -165,8 +165,6 @@ async function validateAst(
   const parserErrors = parsingResult.errors;
 
   /**
-   * @TODO - move deeper
-   *
    * Some changes to the grammar deleted the literal names for some tokens.
    * This is a workaround to restore the literals that were lost.
    *

From 0b9e6ff0cad3c30194920d374f5a0522de6cf814 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 9 Sep 2025 12:22:29 -0600
Subject: [PATCH 48/54] fix TS

---
 .../commands/timeseries/columns_after.test.ts | 28 +++++++++++++++++++
 .../commands/timeseries/columns_after.ts      | 20 +++++++++++++
 .../commands/timeseries/index.ts              |  2 ++
 3 files changed, 50 insertions(+)
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.test.ts
 create mode 100644 src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.ts

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.test.ts
new file mode 100644
index 0000000000000..4e166a8759076
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.test.ts
@@ -0,0 +1,28 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import type { ESQLFieldWithMetadata } from '../../types';
+import { columnsAfter } from './columns_after';
+
+describe('TS columnsAfter', () => {
+  it('returns fields from the source', async () => {
+    const sourceFields: ESQLFieldWithMetadata[] = [
+      { name: 'field1', type: 'double', userDefined: false },
+      { name: 'field2', type: 'keyword', userDefined: false },
+    ];
+
+    const result = await columnsAfter({} as any, [], '', {
+      fromFrom: () => Promise.resolve(sourceFields),
+      fromJoin: () => Promise.resolve([]),
+      fromEnrich: () => Promise.resolve([]),
+    });
+
+    expect(result).toEqual(sourceFields);
+  });
+});
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.ts
new file mode 100644
index 0000000000000..24a08d97529e1
--- /dev/null
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.ts
@@ -0,0 +1,20 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+import { type ESQLCommand } from '../../../types';
+import type { ESQLColumnData } from '../../types';
+import type { IAdditionalFields } from '../../registry';
+
+export const columnsAfter = (
+  command: ESQLCommand,
+  _previousColumns: ESQLColumnData[], // will always be empty for TS
+  _query: string,
+  additionalFields: IAdditionalFields
+) => {
+  return additionalFields.fromFrom(command);
+};
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/index.ts
index c8b106a8654f8..2270c6579dc4e 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/index.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/index.ts
@@ -10,11 +10,13 @@ import { i18n } from '@kbn/i18n';
 import type { ICommandMethods } from '../../registry';
 import { autocomplete } from './autocomplete';
 import { validate } from './validate';
+import { columnsAfter } from './columns_after';
 import type { ICommandContext } from '../../types';
 
 const timeseriesCommandMethods: ICommandMethods<ICommandContext> = {
   autocomplete,
   validate,
+  columnsAfter,
 };
 
 export const timeseriesCommand = {

From c7a8428adf8f53e422a687789acdebd60b8a6ec7 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 9 Sep 2025 13:33:23 -0600
Subject: [PATCH 49/54] remove fixmes

---
 .../src/definitions/utils/autocomplete/functions.ts       | 8 --------
 .../src/definitions/utils/autocomplete/helpers.ts         | 4 ----
 2 files changed, 12 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
index 06f97b81dd942..e82882eb0cb1d 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts
@@ -450,14 +450,6 @@ async function getListArgsSuggestions(
       }
     }
 
-    // FIXME
-    // const anyUserDefinedColumns = collectUserDefinedColumns(commands, columnMap, innerText);
-    // // extract the current node from the userDefinedColumns inferred
-    // anyUserDefinedColumns.forEach((values, key) => {
-    //   if (values.some((v) => v.location === node.location)) {
-    //     anyUserDefinedColumns.delete(key);
-    //   }
-    // });
     const [firstArg] = node.args;
     if (isColumn(firstArg)) {
       const argType = extractTypeFromASTArg(firstArg, {
diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
index 629e9bdae5957..741af5ac94ac7 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts
@@ -215,10 +215,6 @@ export async function getFieldsOrFunctionsSuggestions(
           activeProduct
         )
       : [],
-    // FIXME
-    // userDefinedColumns
-    //   ? pushItUpInTheList(buildUserDefinedColumnsDefinitions(filteredColumnByType), functions)
-    //   : [],
     literals ? getCompatibleLiterals(types) : []
   );
 

From 0b1d93fd45e5d4164d51ae035eca43621ccceb4f Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Tue, 9 Sep 2025 13:36:02 -0600
Subject: [PATCH 50/54] clean unused var

---
 .../src/commands_registry/commands/completion/validate.ts       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
index 19408c4857dd5..2daf08a996565 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts
@@ -22,7 +22,7 @@ export const validate = (
 ): ESQLMessage[] => {
   const messages: ESQLMessage[] = [];
 
-  const { prompt, location, targetField, inferenceId } = command as ESQLAstCompletionCommand;
+  const { prompt, location, inferenceId } = command as ESQLAstCompletionCommand;
 
   const promptExpressionType = getExpressionType(prompt, context?.columns);
 

From 74432cbf04ae562f1cc3331714d336732f4cebed Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 10 Sep 2025 10:23:38 -0600
Subject: [PATCH 51/54] update stats to inline-stats

---
 .../commands/inlinestats/columns_after.test.ts                  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
index 8029d1d7fed2e..eaebad38ef1fe 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
@@ -19,7 +19,7 @@ describe('INLINESTATS', () => {
       { name: '@timestamp', type: 'date', userDefined: false },
     ];
 
-    const queryString = `FROM a | STATS AVG(field1) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`;
+    const queryString = `FROM a | INLINESTATS AVG(field1) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`;
 
     // Can't use synth because it steps on the location information
     // which is used to determine the name of the new column

From f64c58e21f3f932eb9017341fd5991a90bea8fe1 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 10 Sep 2025 10:33:03 -0600
Subject: [PATCH 52/54] remove target column from rerank validation list

---
 .../commands_registry/commands/rerank/validate.ts    | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts
index a51cbd8ecdbd1..956ffade901ce 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts
@@ -23,7 +23,7 @@ export const validate = (
 ): ESQLMessage[] => {
   const messages: ESQLMessage[] = [];
 
-  const { query, targetField, location, inferenceId } = command as ESQLAstRerankCommand;
+  const { query, location, inferenceId } = command as ESQLAstRerankCommand;
   const rerankExpressionType = getExpressionType(query, context?.columns);
 
   // check for supported query types
@@ -40,16 +40,6 @@ export const validate = (
     messages.push(errors.byId('inferenceIdRequired', command.location, { command: 'RERANK' }));
   }
 
-  const targetName = targetField?.name || 'rerank';
-
-  // Sets the target field so the column is recognized after the command is applied
-  context?.columns.set(targetName, {
-    name: targetName,
-    location: targetField?.location || location,
-    type: 'keyword',
-    userDefined: true,
-  });
-
   messages.push(...validateCommandArguments(command, ast, context, callbacks));
 
   return messages;

From 383cdbd51c1978856406ea0f637426273087a2c4 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 10 Sep 2025 11:00:59 -0600
Subject: [PATCH 53/54] remove _fork column setting

---
 .../src/commands_registry/commands/fork/validate.ts         | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
index 677eba017ce7c..7f26d94fc916b 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts
@@ -68,11 +68,5 @@ export const validate = (
     messages.push(errors.tooManyForks(forks[1]));
   }
 
-  context?.columns.set('_fork', {
-    name: '_fork',
-    type: 'keyword',
-    userDefined: false,
-  });
-
   return messages;
 };

From f025cd1314c9bb56864f51128a7c9b7716f93a35 Mon Sep 17 00:00:00 2001
From: Drew Tate <drew.tate@elastic.co>
Date: Wed, 10 Sep 2025 13:58:44 -0600
Subject: [PATCH 54/54] =?UTF-8?q?update=20location=20=F0=9F=99=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../commands/inlinestats/columns_after.test.ts                | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
index eaebad38ef1fe..875695d5cd378 100644
--- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
+++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts
@@ -32,8 +32,8 @@ describe('INLINESTATS', () => {
     const result = columnsAfter(command, previousCommandFields, queryString);
 
     expect(result).toEqual<ESQLColumnData[]>([
-      { name: 'AVG(field1)', type: 'double', userDefined: true, location: { min: 15, max: 25 } },
-      { name: 'buckets', type: 'date', userDefined: true, location: { min: 30, max: 36 } },
+      { name: 'AVG(field1)', type: 'double', userDefined: true, location: { min: 21, max: 31 } },
+      { name: 'buckets', type: 'date', userDefined: true, location: { min: 36, max: 42 } },
       ...previousCommandFields.filter(({ name }) => name !== 'buckets'),
     ]);
   });