[ES|QL][Discover] Highlight results when using ES|QL highlight

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 { EsqlColumnsWithHighlights } from '@kbn/esql-utils';
 import type { DatatableColumnMeta } from '@kbn/expressions-plugin/common';
 
 export type { IgnoredReason, ShouldShowFieldInTableHandler } from './utils';
@@ -26,6 +27,11 @@ export interface EsHitRecord extends Omit<DiscoverSearchHit, '_index' | '_id' |
   _index?: DiscoverSearchHit['_index'];
   _id?: DiscoverSearchHit['_id'];
   _source?: DiscoverSearchHit['_source'];
+  /**
+   * As oposed to DSL, ES|QL highlights are inlined in the hit value.
+   * This record holds which columns have highlights and what tag was used for it.
+   */
+  inline_highlights?: EsqlColumnsWithHighlights;
 }
 
 /**

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

@@ -58,7 +58,7 @@ export function formatHitReact(
     return cached.formattedHit;
   }
 
-  const highlights = hit.raw.highlight ?? {};
+  const highlights = hit.raw.highlight ?? hit.raw.inline_highlights ?? {};
   const flattened = hit.flattened;
   const renderedPairs: PartialHitReactPair[] = [];
   const otherPairs: PartialHitReactPair[] = [];
@@ -114,7 +114,6 @@ export function formatHitReact(
       fieldName: key,
       columnMeta: columnsMeta?.[key],
     });
-
     pair[1] = formatFieldValueReact({
       value: flattened[key],
       hit: hit.raw,

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

@@ -72,13 +72,16 @@ export {
   hasTimeseriesInfoCommand,
   isComputedColumn,
   getQuerySummary,
+  getColumnsWithHighlights,
   buildRenameSourceFieldMap,
   getEsqlControls,
   getAllEsqlControls,
   convertFiltersToESQLExpression,
   convertQueryToESQLExpression,
   injectWhereClauseAfterSourceCommand,
   type ESQLStatsQueryMeta,
+  type EsqlColumnsWithHighlights,
+  type ESQLHighlightTags,
 } from './src';
 
 export { ENABLE_ESQL, GROUP_NOT_SET_VALUE } from './constants';

src/platform/packages/shared/kbn-esql-utils/moon.yml

@@ -36,6 +36,7 @@ dependsOn:
   - '@kbn/controls-constants'
   - '@kbn/presentation-publishing'
   - '@kbn/controls-schemas'
+  - '@kbn/code-editor'
 tags:
   - shared-common
   - package

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

@@ -74,6 +74,11 @@ export {
 } from './utils/cascaded_documents_helpers/utils';
 export { getProjectRoutingFromEsqlQuery } from './utils/set_instructions_helpers';
 export { isComputedColumn, getQuerySummary } from './utils/get_query_summary';
+export {
+  getColumnsWithHighlights,
+  type EsqlColumnsWithHighlights,
+  type ESQLHighlightTags,
+} from './utils/get_columns_with_highlights';
 export { buildRenameSourceFieldMap } from './utils/build_rename_source_field_map';
 export { getAllEsqlControls, getEsqlControls } from './utils/get_esql_controls';
 export { convertFiltersToESQLExpression } from './utils/convert_filters_to_esql';

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

@@ -0,0 +1,112 @@
+/*
+ * 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({
+      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({
+      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({
+      a: {
+        preTag: DEFAULT_HIGHLIGHT_PRE_TAG,
+        postTag: DEFAULT_HIGHLIGHT_POST_TAG,
+      },
+      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({
+      '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({
+      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({
+      '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({
+      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({
+      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,160 @@
+/*
+ * 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 ESQLHighlightTags {
+  preTag: string;
+  postTag: string;
+}
+
+export type EsqlColumnsWithHighlights = Record<string, ESQLHighlightTags>;
+
+/**
+ * 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 map:
+ * ```
+ * {
+ *   snippets: {
+ *     preTag: '<em>',
+ *     postTag: '</em>',
+ *   },
+ *   titles: {
+ *     preTag: '<mark>',
+ *     postTag: '</mark>',
+ *   },
+ * }
+ */
+export function getColumnsWithHighlights(query: string): EsqlColumnsWithHighlights {
+  const columnsWithHighlights: EsqlColumnsWithHighlights = {};
+  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] = {
+      preTag,
+      postTag,
+    };
+  }
+
+  return columnsWithHighlights;
+}
+
+/**
+ * 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

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
-import type { monaco } from '@kbn/monaco';
+import type { monaco } from '@kbn/code-editor';
 import type { ESQLColumn } from '@elastic/esql/types';
 import { Parser, walk } from '@elastic/esql';
 import { ESQLVariableType, type ESQLControlVariable } from '@kbn/esql-types';
@@ -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';