[PoC] Add ES|QL highlighting to field formatters

src/platform/packages/shared/kbn-discover-utils/src/types.ts

@@ -8,6 +8,7 @@
  */
 
 import type { SearchHit } from '@elastic/elasticsearch/lib/api/types';
+import type { EsqlColumnHighlightMap } from '@kbn/esql-utils';
 import type { DatatableColumnMeta } from '@kbn/expressions-plugin/common';
 
 export type { IgnoredReason, ShouldShowFieldInTableHandler } from './utils';
@@ -26,6 +27,7 @@ export interface EsHitRecord extends Omit<DiscoverSearchHit, '_index' | '_id' |
   _index?: DiscoverSearchHit['_index'];
   _id?: DiscoverSearchHit['_id'];
   _source?: DiscoverSearchHit['_source'];
+  esql_highlight?: EsqlColumnHighlightMap;
 }
 
 /**

src/platform/packages/shared/kbn-esql-utils/index.ts

@@ -78,6 +78,12 @@ export {
   convertFiltersToESQLExpression,
   convertQueryToESQLExpression,
   injectWhereClauseAfterSourceCommand,
+  getColumnsWithHighlights,
+  getColumnsWithHighlightsMap,
+  DEFAULT_HIGHLIGHT_PRE_TAG,
+  DEFAULT_HIGHLIGHT_POST_TAG,
+  type EsqlColumnHighlight,
+  type EsqlColumnHighlightMap,
   type ESQLStatsQueryMeta,
 } from './src';
 

src/platform/packages/shared/kbn-esql-utils/src/index.ts

@@ -79,6 +79,14 @@ export { getAllEsqlControls, getEsqlControls } from './utils/get_esql_controls';
 export { convertFiltersToESQLExpression } from './utils/convert_filters_to_esql';
 export { convertQueryToESQLExpression } from './utils/convert_query_to_esql';
 export { injectWhereClauseAfterSourceCommand } from './utils/inject_where_after_source';
+export {
+  getColumnsWithHighlights,
+  getColumnsWithHighlightsMap,
+  DEFAULT_HIGHLIGHT_PRE_TAG,
+  DEFAULT_HIGHLIGHT_POST_TAG,
+  type EsqlColumnHighlight,
+  type EsqlColumnHighlightMap,
+} from './utils/get_columns_with_highlights';
 
 // Callback functions
 export * from './utils/callbacks';

src/platform/packages/shared/kbn-esql-utils/src/utils/get_columns_with_highlights.test.ts

@@ -0,0 +1,121 @@
+/*
+ * 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 {
+  DEFAULT_HIGHLIGHT_POST_TAG,
+  DEFAULT_HIGHLIGHT_PRE_TAG,
+  getColumnsWithHighlights,
+} from './get_columns_with_highlights';
+
+describe('getColumnsWithHighlights', () => {
+  it('returns column with default em tags when highlight is enabled', () => {
+    const query =
+      'FROM books | EVAL snippets = TOP_SNIPPETS(description, "Tolkien", { "highlight": true })';
+    expect(getColumnsWithHighlights(query)).toEqual([
+      {
+        column: 'snippets',
+        preTag: DEFAULT_HIGHLIGHT_PRE_TAG,
+        postTag: DEFAULT_HIGHLIGHT_POST_TAG,
+      },
+    ]);
+  });
+
+  it('returns custom pre_tag and post_tag from TOP_SNIPPETS options', () => {
+    const query =
+      'FROM books | EVAL snippets = TOP_SNIPPETS(description, "Tolkien", { "highlight": true, "pre_tag": "<mark>", "post_tag": "</mark>" })';
+    expect(getColumnsWithHighlights(query)).toEqual([
+      {
+        column: 'snippets',
+        preTag: '<mark>',
+        postTag: '</mark>',
+      },
+    ]);
+  });
+
+  it('ignores TOP_SNIPPETS without highlight option', () => {
+    const query = 'FROM books | EVAL snippets = TOP_SNIPPETS(description, "Tolkien")';
+    expect(getColumnsWithHighlights(query)).toEqual([]);
+  });
+
+  it('returns multiple columns with their respective tags', () => {
+    const query =
+      'FROM books | EVAL a = TOP_SNIPPETS(description, "one", { "highlight": true }) | EVAL b = TOP_SNIPPETS(title, "two", { "highlight": true, "pre_tag": "<mark>", "post_tag": "</mark>" })';
+    expect(getColumnsWithHighlights(query)).toEqual([
+      {
+        column: 'a',
+        preTag: DEFAULT_HIGHLIGHT_PRE_TAG,
+        postTag: DEFAULT_HIGHLIGHT_POST_TAG,
+      },
+      {
+        column: 'b',
+        preTag: '<mark>',
+        postTag: '</mark>',
+      },
+    ]);
+  });
+
+  it('handles EVAL unamed columns scenarios', () => {
+    const query = 'FROM books | EVAL TOP_SNIPPETS(description, "one", { "highlight": true })';
+    expect(getColumnsWithHighlights(query)).toEqual([
+      {
+        column: 'TOP_SNIPPETS(description, "one", { "highlight": true })',
+        preTag: DEFAULT_HIGHLIGHT_PRE_TAG,
+        postTag: DEFAULT_HIGHLIGHT_POST_TAG,
+      },
+    ]);
+  });
+
+  it('handles STATS user defined columns', () => {
+    const query =
+      'FROM books | STATS count(*) BY col0 = TOP_SNIPPETS(description, "one", { "highlight": true })';
+    expect(getColumnsWithHighlights(query)).toEqual([
+      {
+        column: 'col0',
+        preTag: DEFAULT_HIGHLIGHT_PRE_TAG,
+        postTag: DEFAULT_HIGHLIGHT_POST_TAG,
+      },
+    ]);
+  });
+
+  it('handles STATS unamed user defined columns', () => {
+    const query =
+      'FROM books | STATS count(*) BY TOP_SNIPPETS(description, "one", { "highlight": true })';
+    expect(getColumnsWithHighlights(query)).toEqual([
+      {
+        column: 'TOP_SNIPPETS(description, "one", { "highlight": true })',
+        preTag: DEFAULT_HIGHLIGHT_PRE_TAG,
+        postTag: DEFAULT_HIGHLIGHT_POST_TAG,
+      },
+    ]);
+  });
+
+  it('applies RENAME to resolved highlight column names', () => {
+    const query =
+      'FROM books | EVAL col0 = TOP_SNIPPETS(description, "one", { "highlight": true }) | RENAME col0 AS renamed';
+    expect(getColumnsWithHighlights(query)).toEqual([
+      {
+        column: 'renamed',
+        preTag: DEFAULT_HIGHLIGHT_PRE_TAG,
+        postTag: DEFAULT_HIGHLIGHT_POST_TAG,
+      },
+    ]);
+  });
+
+  it('can handle columns defined within quotes', () => {
+    const query =
+      'FROM books | EVAL `col0` = TOP_SNIPPETS(description, "one", { "highlight": true })';
+    expect(getColumnsWithHighlights(query)).toEqual([
+      {
+        column: 'col0',
+        preTag: DEFAULT_HIGHLIGHT_PRE_TAG,
+        postTag: DEFAULT_HIGHLIGHT_POST_TAG,
+      },
+    ]);
+  });
+});

src/platform/packages/shared/kbn-esql-utils/src/utils/get_columns_with_highlights.ts

@@ -0,0 +1,166 @@
+/*
+ * 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,
+  isBooleanLiteral,
+  isColumn,
+  isMap,
+  isStringLiteral,
+  LeafPrinter,
+  Parser,
+  Walker,
+} from '@elastic/esql';
+
+import type { ESQLAstQueryExpression, ESQLFunction, ESQLMap } from '@elastic/esql/types';
+import { replaceColumnNamesIfRenamed } from './query_parsing_helpers';
+
+export const DEFAULT_HIGHLIGHT_PRE_TAG = '<em>';
+export const DEFAULT_HIGHLIGHT_POST_TAG = '</em>';
+
+const HIGHLIGHT_OPTION_NAME = 'highlight';
+const PRE_TAG_OPTION_NAME = 'pre_tag';
+const POST_TAG_OPTION_NAME = 'post_tag';
+
+/**
+ * ES|QL functions that can produce highlight markup in output columns when
+ * called with `{ "highlight": true }`.
+ */
+export const FUNCTIONS_WITH_HIGHLIGHT_SUPPORT = ['top_snippets'];
+
+export interface EsqlColumnHighlight {
+  column: string;
+  preTag: string;
+  postTag: string;
+}
+
+export type EsqlColumnHighlightMap = Record<string, EsqlColumnHighlight>;
+
+export function getColumnsWithHighlightsMap(query: string): EsqlColumnHighlightMap {
+  const columnsWithHighlights: EsqlColumnHighlightMap = {};
+  const { root } = Parser.parse(query);
+
+  const highlightFunctionsCandidates = Walker.findAll(
+    root,
+    (node) => node.type === 'function' && FUNCTIONS_WITH_HIGHLIGHT_SUPPORT.includes(node.name)
+  ) as ESQLFunction[];
+
+  for (const fn of highlightFunctionsCandidates) {
+    const optionsMap = fn.args.find(isMap);
+
+    if (!optionsMap || !isHighlightEnabled(optionsMap)) {
+      continue;
+    }
+
+    const columnName = getHighlightedColumnName(root, fn, query);
+    if (!columnName) {
+      continue;
+    }
+
+    const preTag =
+      getHighlightTagName(optionsMap, PRE_TAG_OPTION_NAME) ?? DEFAULT_HIGHLIGHT_PRE_TAG;
+    const postTag =
+      getHighlightTagName(optionsMap, POST_TAG_OPTION_NAME) ?? DEFAULT_HIGHLIGHT_POST_TAG;
+
+    // Check if the column name has been renamed in the query
+    const [resolvedColumnName] = replaceColumnNamesIfRenamed(root, [columnName]);
+
+    columnsWithHighlights[resolvedColumnName] = {
+      column: resolvedColumnName,
+      preTag,
+      postTag,
+    };
+  }
+
+  return columnsWithHighlights;
+}
+
+/**
+ * Returns columns built using a highlighting algorithm,
+ * including the opening and closing markup tags configured for each column.
+ *
+ * Example:
+ * ```
+ * FROM books
+ *  | EVAL snippets = TOP_SNIPPETS(description, "Tolkien", { "highlight": true })
+ *  | EVAL titles = TOP_SNIPPETS(title, "Tolkien", { "highlight": true, "pre_tag": "<mark>", "post_tag": "</mark>" })
+ * ```
+ * Will return the following column:
+ * ```
+ * {
+ *   column: 'snippets',
+ *   preTag: '<em>',
+ *   postTag: '</em>',
+ * }
+ * {
+ *   column: 'titles',
+ *   preTag: '<mark>',
+ *   postTag: '</mark>',
+ * }
+ */
+export function getColumnsWithHighlights(query: string): EsqlColumnHighlight[] {
+  return Object.values(getColumnsWithHighlightsMap(query));
+}
+
+/**
+ * Given a map of options, returns true if the `highlight` option is set to `true`.
+ */
+const isHighlightEnabled = (optionsMap: ESQLMap): boolean => {
+  const highlightEntry = optionsMap.entries.find(
+    (entry) => isStringLiteral(entry.key) && entry.key.valueUnquoted === HIGHLIGHT_OPTION_NAME
+  );
+  if (!highlightEntry?.value) {
+    return false;
+  }
+
+  return (
+    isBooleanLiteral(highlightEntry.value) && highlightEntry.value.value.toLowerCase() === 'true'
+  );
+};
+
+/**
+ * Returns the tag name defined in the map options if it exists.
+ */
+const getHighlightTagName = (optionsMap: ESQLMap, optionName: string): string | undefined => {
+  const tagEntry = optionsMap.entries.find(
+    (entry) => isStringLiteral(entry.key) && entry.key.valueUnquoted === optionName
+  );
+
+  if (!tagEntry?.value) {
+    return undefined;
+  }
+  if (!isStringLiteral(tagEntry.value)) {
+    return undefined;
+  }
+  return tagEntry.value.valueUnquoted;
+};
+
+/**
+ * Returns the name of the column that was created using the highlight function.
+ *
+ * This function has an heuristic part, some combination of function could remove the highlighting tokens from the result.
+ * But it assumes that if the user used highlight:true, it's not interested in removing them.
+ * Doing a 100% accurate check would involve knowing the semantics of every invoked function.
+ * In the worst case of having a false positive,
+ */
+const getHighlightedColumnName = (
+  root: ESQLAstQueryExpression,
+  highlightFunction: ESQLFunction,
+  query: string
+): string | undefined => {
+  // Created using an assignment | EVAL col = TOP_SNIPPETS( ...
+  for (const parent of Walker.parents(root, highlightFunction)) {
+    if (isAssignment(parent) && isColumn(parent.args[0])) {
+      return LeafPrinter.column(parent.args[0]);
+    }
+  }
+
+  // Created using an expression text | EVAL TOP_SNIPPETS( ... or STATS count(*) BY TOP_SNIPPETS( ...
+  return query.substring(highlightFunction.location.min, highlightFunction.location.max + 1);
+};

src/platform/packages/shared/kbn-esql-utils/src/utils/query_parsing_helpers.test.ts

@@ -25,6 +25,7 @@ import {
   fixESQLQueryWithVariables,
   getCategorizeColumns,
   getArgsFromRenameFunction,
+  replaceColumnNamesIfRenamed,
   getCategorizeField,
   findClosestColumn,
   getKqlSearchQueries,
@@ -922,6 +923,23 @@ describe('esql query helpers', () => {
     });
   });
 
+  describe('replaceColumnNameIfRenamed', () => {
+    it('returns column names unchanged when there is no RENAME', () => {
+      const { root } = Parser.parse('FROM index | KEEP col');
+      expect(replaceColumnNamesIfRenamed(root, ['col'])).toEqual(['col']);
+    });
+
+    it('replaces matching column names using RENAME', () => {
+      const { root } = Parser.parse('FROM index | RENAME old AS new');
+      expect(replaceColumnNamesIfRenamed(root, ['old', 'other'])).toEqual(['new', 'other']);
+    });
+
+    it('applies multiple RENAME commands in order', () => {
+      const { root } = Parser.parse('FROM index | RENAME a AS b | RENAME b AS c');
+      expect(replaceColumnNamesIfRenamed(root, ['a'])).toEqual(['c']);
+    });
+  });
+
   describe('getArgsFromRenameFunction', () => {
     it('should return the args from an = rename function', () => {
       const esql = 'FROM index | RENAME renamed = original';