fix: label Map attribute columns by their key instead of arrayElement(...)

.changeset/map-attribute-column-alias.md

@@ -0,0 +1,14 @@
+---
+'@hyperdx/app': patch
+---
+
+fix: label Map attribute columns added from the search sidebars by their key
+
+Adding a Map attribute such as `ResourceAttributes['service.name']` as a column
+from the filters sidebar or the row side panel now writes it into the SELECT as
+`ResourceAttributes['service.name'] AS "service.name"`, so the results column
+reads `service.name` instead of
+`arrayElement(ResourceAttributes, 'service.name')`. The alias is visible and
+editable in the SELECT, and queries typed by hand are not changed. A key named
+like a table column, or one already used as a name in the SELECT, is added
+without an alias.

packages/app/src/DBSearchPage.tsx

@@ -130,9 +130,12 @@ import {
 } from '@/timeQuery';
 import {
   formatDurationMs,
+  orderByAfterRemovingSelectItem,
   QUERY_LOCAL_STORAGE,
+  selectItemExpression,
   useLocalStorage,
   usePrevious,
+  withMapKeyAlias,
 } from '@/utils';
 
 import ChartSQLPreview, { SQLPreview } from './components/ChartSQLPreview';
@@ -1104,7 +1107,7 @@ export function DBSearchPage() {
     [sources, lastSelectedSourceId],
   );
 
-  const { control, setValue, reset, handleSubmit, formState } =
+  const { control, setValue, getValues, reset, handleSubmit, formState } =
     useForm<SearchConfigFromSchema>({
       values: {
         select: searchedConfig.select || '',
@@ -1737,13 +1740,38 @@ export function DBSearchPage() {
 
   const toggleColumn = useCallback(
     (column: string) => {
-      const newSelectArray = displayedColumns.includes(column)
-        ? displayedColumns.filter(s => s !== column)
-        : [...displayedColumns, column];
+      // A column added from the UI can carry an alias, so match the expression too
+      const selected = displayedColumns.find(
+        s => s === column || selectItemExpression(s) === column,
+      );
+      const newSelectArray = selected
+        ? displayedColumns.filter(s => s !== selected)
+        : [
+            ...displayedColumns,
+            withMapKeyAlias(column, displayedColumns, knownColumns),
+          ];
       setValue('select', newSelectArray.join(', '));
+      if (selected) {
+        const orderBy = getValues('orderBy') ?? '';
+        const nextOrderBy = orderByAfterRemovingSelectItem(
+          selected,
+          orderBy,
+          defaultSearchConfig.orderBy ?? '',
+        );
+        if (nextOrderBy !== orderBy) {
+          setValue('orderBy', nextOrderBy);
+        }
+      }
       onSubmit();
     },
-    [displayedColumns, setValue, onSubmit],
+    [
+      displayedColumns,
+      knownColumns,
+      setValue,
+      getValues,
+      defaultSearchConfig.orderBy,
+      onSubmit,
+    ],
   );
 
   const generateSearchUrl = useCallback(

packages/app/src/__tests__/utils.test.ts

@@ -21,14 +21,17 @@ import {
   getAllMetricTables,
   getColorFromCSSToken,
   getMetricTableName,
+  isColumnInSelect,
   mapKeyBy,
   mergePath,
+  orderByAfterRemovingSelectItem,
   orderByStringToSortingState,
   parseTimestampToMs,
   resolveConditionalColor,
   sortingStateToOrderByString,
   stripTrailingSlash,
   useQueryHistory,
+  withMapKeyAlias,
 } from '@/utils';
 
 describe('formatAttributeClause', () => {
@@ -1420,6 +1423,118 @@ describe('mergePath', () => {
   });
 });
 
+describe('withMapKeyAlias', () => {
+  const noColumns = new Set(['Timestamp', 'Body']);
+
+  it('labels a Map subscript with its key', () => {
+    expect(
+      withMapKeyAlias("ResourceAttributes['service.name']", [], noColumns),
+    ).toBe(`ResourceAttributes['service.name'] AS "service.name"`);
+    expect(withMapKeyAlias("`LogAttributes`['k8s.pod']", [], noColumns)).toBe(
+      '`LogAttributes`[\'k8s.pod\'] AS "k8s.pod"',
+    );
+  });
+
+  it('unescapes the key the way mergePath escaped it', () => {
+    expect(withMapKeyAlias("LogAttributes['it\\'s']", [], noColumns)).toBe(
+      `LogAttributes['it\\'s'] AS "it's"`,
+    );
+  });
+
+  it('leaves anything that is not a bare Map subscript alone', () => {
+    [
+      'Body',
+      'lower(Body)',
+      "LogAttributes['a']['b']",
+      "JSONExtractString(Body, 'a')",
+      "LogAttributes['a'] AS a",
+    ].forEach(column => {
+      expect(withMapKeyAlias(column, [], noColumns)).toBe(column);
+    });
+  });
+
+  it('keeps the derived name when the alias could change the query', () => {
+    // An alias shadows a column of the same name in WHERE and ORDER BY
+    expect(
+      withMapKeyAlias(
+        "ResourceAttributes['ServiceName']",
+        [],
+        new Set(['ServiceName']),
+      ),
+    ).toBe("ResourceAttributes['ServiceName']");
+    // Two selected columns cannot share a name
+    expect(
+      withMapKeyAlias(
+        "LogAttributes['http.method']",
+        [`SpanAttributes['http.method'] AS "http.method"`],
+        noColumns,
+      ),
+    ).toBe("LogAttributes['http.method']");
+    // Keys that would need escaping inside the quoted alias
+    expect(withMapKeyAlias(`LogAttributes['a"b']`, [], noColumns)).toBe(
+      `LogAttributes['a"b']`,
+    );
+    // Table columns not loaded yet: a clash with one cannot be ruled out
+    expect(
+      withMapKeyAlias("ResourceAttributes['ServiceName']", [], new Set()),
+    ).toBe("ResourceAttributes['ServiceName']");
+  });
+});
+
+describe('isColumnInSelect', () => {
+  it('matches a column whether or not it was aliased', () => {
+    const select = [
+      'Timestamp',
+      `ResourceAttributes['service.name'] AS "service.name"`,
+    ];
+    expect(isColumnInSelect(select, 'Timestamp')).toBe(true);
+    expect(isColumnInSelect(select, "ResourceAttributes['service.name']")).toBe(
+      true,
+    );
+    expect(isColumnInSelect(select, 'Body')).toBe(false);
+    expect(isColumnInSelect(undefined, 'Body')).toBe(false);
+  });
+
+  it('matches a hand-typed alias too', () => {
+    expect(isColumnInSelect(['Body AS b'], 'Body')).toBe(true);
+    expect(isColumnInSelect(['lower(Body) AS `b`'], 'lower(Body)')).toBe(true);
+  });
+
+  it('does not mistake a cast for an alias', () => {
+    expect(isColumnInSelect(['CAST(x AS String)'], 'CAST(x')).toBe(false);
+  });
+});
+
+describe('orderByAfterRemovingSelectItem', () => {
+  const item = `ResourceAttributes['service.name'] AS "service.name"`;
+
+  it('drops a sort on the alias of the removed column', () => {
+    expect(
+      orderByAfterRemovingSelectItem(
+        item,
+        '"service.name" DESC',
+        'Timestamp DESC',
+      ),
+    ).toBe('Timestamp DESC');
+    expect(orderByAfterRemovingSelectItem(item, '"service.name" ASC', '')).toBe(
+      '',
+    );
+  });
+
+  it('keeps any other sort', () => {
+    expect(
+      orderByAfterRemovingSelectItem(item, 'Timestamp ASC', 'Timestamp DESC'),
+    ).toBe('Timestamp ASC');
+    expect(orderByAfterRemovingSelectItem(item, '"http.method" DESC', '')).toBe(
+      '"http.method" DESC',
+    );
+    // A column without an alias sorts by its own expression, which is not affected
+    expect(
+      orderByAfterRemovingSelectItem('Body', 'Body DESC', 'Timestamp DESC'),
+    ).toBe('Body DESC');
+  });
+});
+
 describe('getColorFromCSSToken', () => {
   afterEach(() => {
     jest.restoreAllMocks();

packages/app/src/components/DBRowJsonViewer.tsx

@@ -33,7 +33,7 @@ import HyperJson, {
   LineAction,
 } from '@/components/HyperJson';
 import { useFormatTime } from '@/useFormatTime';
-import { mergePath } from '@/utils';
+import { isColumnInSelect, mergePath } from '@/utils';
 import {
   CLIPBOARD_ERROR_MESSAGE,
   copyTextToClipboard,
@@ -584,7 +584,7 @@ export function DBRowJsonViewer({
           }
         }
 
-        const isIncluded = displayedColumns?.includes(columnFieldPath);
+        const isIncluded = isColumnInSelect(displayedColumns, columnFieldPath);
         actions.push({
           key: 'toggle-column',
           label: isIncluded ? <IconMinus size={14} /> : <IconPlus size={14} />,

packages/app/src/components/DBSearchPageFilters.tsx

@@ -63,7 +63,7 @@ import {
   usePinnedFilters,
 } from '@/searchFilters';
 import { useSource } from '@/source';
-import { useLocalStorage } from '@/utils';
+import { isColumnInSelect, useLocalStorage } from '@/utils';
 
 import { FilterSettingsPanel } from './DBSearchPageFilters/FilterSettingsPopover';
 import { useFetchFacets } from './DBSearchPageFilters/hooks';
@@ -1554,7 +1554,10 @@ const DBSearchPageFiltersComponent = ({
                 onColumnToggle={
                   onColumnToggle ? () => onColumnToggle(facetSqlKey) : undefined
                 }
-                isColumnDisplayed={displayedColumns?.includes(facetSqlKey)}
+                isColumnDisplayed={isColumnInSelect(
+                  displayedColumns,
+                  facetSqlKey,
+                )}
                 onLoadMore={loadMoreFacetsForKey}
                 loadMoreLoading={loadMoreLoadingKeys.has(facet.key)}
                 hasLoadedMore={extraFacetKeys.has(facet.key)}

packages/app/src/components/DBSearchPageFilters/NestedFilterGroup.tsx

@@ -4,6 +4,7 @@ import { Accordion, Group, Text, Tooltip, UnstyledButton } from '@mantine/core';
 import { useVirtualizer } from '@tanstack/react-virtual';
 
 import { FilterGroup } from '@/components/DBSearchPageFilters';
+import { isColumnInSelect } from '@/utils';
 
 import classes from '@styles/SearchPage.module.scss';
 
@@ -243,7 +244,8 @@ export const NestedFilterGroup = ({
                                     onColumnToggle(child.sqlKey ?? child.key)
                                 : undefined
                             }
-                            isColumnDisplayed={displayedColumns?.includes(
+                            isColumnDisplayed={isColumnInSelect(
+                              displayedColumns,
                               child.sqlKey ?? child.key,
                             )}
                             onLoadMore={() => onLoadMore(child.key)}

packages/app/src/utils.ts

@@ -4,7 +4,10 @@ import { formatDistanceToNowStrict } from 'date-fns';
 import numbro from 'numbro';
 import type { SetStateAction } from 'react';
 import TimestampNano from 'timestamp-nano';
-import { TableConnection } from '@hyperdx/common-utils/dist/core/metadata';
+import {
+  TableConnection,
+  unquoteIdentifier,
+} from '@hyperdx/common-utils/dist/core/metadata';
 import {
   CATEGORICAL_PALETTE_TOKENS,
   ChartPaletteToken,
@@ -1201,6 +1204,93 @@ export const mergePath = (
     .join('')}`;
 };
 
+// A bare Map subscript as the UI builds it: `LogAttributes['key']` or `` `LogAttributes`['key'] ``
+const MAP_SUBSCRIPT_REGEX =
+  /^(`[^`]+`|[A-Za-z_][A-Za-z0-9_]*)\['((?:[^'\\]|\\.)*)'\]$/;
+// A trailing `AS alias` on a SELECT item, with the alias quoted or bare
+const SELECT_ALIAS_REGEX =
+  /\s+AS\s+("(?:[^"\\]|\\.)*"|`[^`]*`|[A-Za-z_][A-Za-z0-9_]*)\s*$/i;
+
+/**
+ * The expression of a SELECT item, without its `AS alias`
+ */
+export const selectItemExpression = (item: string): string =>
+  item.replace(SELECT_ALIAS_REGEX, '').trim();
+
+/**
+ * The alias of a SELECT item, unquoted, or undefined if it has none
+ */
+const selectItemAlias = (item: string): string | undefined => {
+  const alias = SELECT_ALIAS_REGEX.exec(item)?.[1];
+  return alias == null ? undefined : unquoteIdentifier(alias);
+};
+
+/**
+ * Whether a column picked in the UI is in the SELECT list, aliased or not
+ */
+export const isColumnInSelect = (
+  displayedColumns: string[] | undefined,
+  column: string,
+): boolean =>
+  displayedColumns?.some(
+    item => item === column || selectItemExpression(item) === column,
+  ) ?? false;
+
+/**
+ * Label a Map subscript picked in the UI with its key, so the results column
+ * reads `service.name` instead of `arrayElement(LogAttributes, 'service.name')`.
+ * The alias is left off whenever it could change the query.
+ */
+export const withMapKeyAlias = (
+  column: string,
+  selectItems: string[],
+  tableColumns: Set<string>,
+): string => {
+  const match = MAP_SUBSCRIPT_REGEX.exec(column);
+  if (!match) {
+    return column;
+  }
+  // Inverse of escapeSqlSingleQuoted, which mergePath uses to build the subscript
+  const key = match[2].replace(/\\(.)/g, '$1');
+  const selectedNames = selectItems.map(
+    item => selectItemAlias(item) ?? item.trim(),
+  );
+  if (
+    // Nothing to label
+    key === '' ||
+    // Would need escaping inside the quoted alias
+    /["\\]/.test(key) ||
+    // Table columns not loaded yet, so a clash with one cannot be ruled out
+    tableColumns.size === 0 ||
+    // ClickHouse resolves an alias before a column of the same name
+    tableColumns.has(key) ||
+    // Two selected columns cannot share a name
+    selectedNames.includes(key)
+  ) {
+    return column;
+  }
+  return `${column} AS "${key}"`;
+};
+
+/**
+ * The sort to keep once a SELECT item is removed. Sorting on an aliased column
+ * from the table header orders by the alias itself (`"service.name" DESC`),
+ * which stops resolving once the item is gone, so that falls back to the
+ * default sort.
+ */
+export const orderByAfterRemovingSelectItem = (
+  item: string,
+  orderBy: string,
+  defaultOrderBy: string,
+): string => {
+  const alias = selectItemAlias(item);
+  if (alias == null) {
+    return orderBy;
+  }
+  const sortedBy = orderBy.trim().replace(/\s+(ASC|DESC)$/i, '');
+  return sortedBy === `"${alias}"` ? defaultOrderBy : orderBy;
+};
+
 const _useTry = <T>(fn: () => T): [null | Error | unknown, null | T] => {
   let output = null;
   let error = null;