fix(app): let the Database query drawer use the span attribute skip index

.changeset/db-statement-index-prefilter.md

@@ -0,0 +1,5 @@
+---
+'@hyperdx/app': patch
+---
+
+fix: Let the Services > Database query drawer use the span attribute skip index instead of scanning the whole trace window

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

@@ -1,4 +1,5 @@
 import type { ColumnMeta } from '@hyperdx/common-utils/dist/clickhouse';
+import { parseQuery } from '@hyperdx/common-utils/dist/filters';
 import type { TTraceSource } from '@hyperdx/common-utils/dist/types';
 import { SourceKind } from '@hyperdx/common-utils/dist/types';
 import { renderHook } from '@testing-library/react';
@@ -7,6 +8,7 @@ import * as metadataModule from '@/hooks/useMetadata';
 import {
   getExpressions,
   makeCoalescedFieldsAccessQuery,
+  makeDbStatementCondition,
   useServiceDashboardExpressions,
 } from '@/serviceDashboard';
 
@@ -146,6 +148,145 @@ describe('Service Dashboard', () => {
     });
   });
 
+  describe('dbStatementIndexHint', () => {
+    const attributesColumn = {
+      name: 'SpanAttributes',
+      type: 'Map(LowCardinality(String), String)',
+    } as ColumnMeta;
+    const itemsColumn = { name: 'SpanAttributeItems' } as ColumnMeta;
+
+    it('should prefer the attribute items column when the table has one', () => {
+      expect(
+        getExpressions(mockSource, [attributesColumn, itemsColumn], [])
+          .dbStatementIndexHint,
+      ).toEqual({ kind: 'items', column: 'SpanAttributeItems' });
+    });
+
+    it('should fall back to the attribute map, which older schemas index by value', () => {
+      expect(
+        getExpressions(mockSource, [attributesColumn], []).dbStatementIndexHint,
+      ).toEqual({ kind: 'mapValues', column: 'SpanAttributes' });
+    });
+
+    it('should be undefined when the attribute field is not a map', () => {
+      expect(
+        getExpressions(
+          mockSource,
+          [{ name: 'SpanAttributes', type: 'String' } as ColumnMeta],
+          [],
+        ).dbStatementIndexHint,
+      ).toBeUndefined();
+    });
+
+    it('should be undefined for JSON attribute columns', () => {
+      expect(
+        getExpressions(
+          mockSource,
+          [attributesColumn, itemsColumn],
+          ['SpanAttributes'],
+        ).dbStatementIndexHint,
+      ).toBeUndefined();
+    });
+
+    it('should not mistake the attribute field itself for an items column', () => {
+      expect(
+        getExpressions(
+          { ...mockSource, eventAttributesExpression: 'Attrs' },
+          [{ name: 'Attrs', type: 'Map(String, String)' } as ColumnMeta],
+          [],
+        ).dbStatementIndexHint,
+      ).toEqual({ kind: 'mapValues', column: 'Attrs' });
+    });
+  });
+
+  describe('makeDbStatementCondition', () => {
+    const dbStatement =
+      "coalesce(nullif(SpanAttributes['db.query.text'], ''), nullif(SpanAttributes['db.statement'], ''))";
+    const equality = `${dbStatement} IN ('SELECT 1')`;
+
+    it('should only compare the statement when nothing is indexed', () => {
+      expect(
+        makeDbStatementCondition({
+          expressions: { dbStatement, dbStatementIndexHint: undefined },
+          statement: 'SELECT 1',
+        }),
+      ).toBe(equality);
+    });
+
+    it('should prefilter on the items column, which newer schemas index', () => {
+      expect(
+        makeDbStatementCondition({
+          expressions: {
+            dbStatement,
+            dbStatementIndexHint: {
+              kind: 'items',
+              column: 'SpanAttributeItems',
+            },
+          },
+          statement: 'SELECT 1',
+        }),
+      ).toBe(
+        "(has(SpanAttributeItems, 'db.query.text=SELECT 1') OR " +
+          "has(SpanAttributeItems, 'db.statement=SELECT 1')) " +
+          `AND ${equality}`,
+      );
+    });
+
+    it('should prefilter on the attribute values, which older schemas index', () => {
+      expect(
+        makeDbStatementCondition({
+          expressions: {
+            dbStatement,
+            dbStatementIndexHint: {
+              kind: 'mapValues',
+              column: 'SpanAttributes',
+            },
+          },
+          statement: 'SELECT 1',
+        }),
+      ).toBe(
+        `(has(mapValues(SpanAttributes), 'SELECT 1')) AND ${equality}`,
+      );
+    });
+
+    it.each([
+      ['items' as const, 'SpanAttributeItems'],
+      ['mapValues' as const, 'SpanAttributes'],
+    ])('should escape quotes in the %s prefilter', (kind, column) => {
+      const condition = makeDbStatementCondition({
+        expressions: { dbStatement, dbStatementIndexHint: { kind, column } },
+        statement: "SELECT 'a'",
+      });
+
+      expect(condition).toContain("SELECT ''a''");
+      expect(condition).not.toContain("SELECT 'a'");
+    });
+
+    it.each([
+      [undefined],
+      [{ kind: 'items' as const, column: 'SpanAttributeItems' }],
+      [{ kind: 'mapValues' as const, column: 'SpanAttributes' }],
+    ])(
+      'should be read back by the search page as the same filter (%p)',
+      dbStatementIndexHint => {
+        const statement = "SELECT 'a' FROM t WHERE b IN (1) AND c = ?";
+
+        const { filters } = parseQuery([
+          {
+            type: 'sql',
+            condition: makeDbStatementCondition({
+              expressions: { dbStatement, dbStatementIndexHint },
+              statement,
+            }),
+          },
+        ]);
+
+        expect(Object.keys(filters)).toEqual([dbStatement]);
+        expect([...filters[dbStatement].included]).toEqual([statement]);
+      },
+    );
+  });
+
   describe('useServiceDashboardExpressions', () => {
     const mockColumns: ColumnMeta[] = [
       { name: 'Duration', type: 'UInt64' },

packages/app/src/components/ServiceDashboardDbQuerySidePanel.tsx

@@ -16,7 +16,10 @@ import { ChartCard } from '@/components/charts/ChartCard';
 import { DBTimeChart } from '@/components/DBTimeChart';
 import { DrawerBody, DrawerHeader } from '@/components/DrawerUtils';
 import SlowestEventsTile from '@/components/ServiceDashboardSlowestEventsTile';
-import { useServiceDashboardExpressions } from '@/serviceDashboard';
+import {
+  makeDbStatementCondition,
+  useServiceDashboardExpressions,
+} from '@/serviceDashboard';
 import { useSource } from '@/source';
 import { useZIndex, ZIndexContext } from '@/zIndex';
 
@@ -46,16 +49,23 @@ export default function ServiceDashboardDbQuerySidePanel({
   const drawerZIndex = contextZIndex + 10;
 
   const dbQueryFilters = useMemo(() => {
+    if (!expressions || !dbQuery) {
+      return [];
+    }
+
     const filters: Filter[] = [
       {
         type: 'sql',
-        condition: `${expressions?.dbStatement} IN ('${dbQuery}')`,
+        condition: makeDbStatementCondition({
+          expressions,
+          statement: dbQuery,
+        }),
       },
     ];
     if (service) {
       filters.push({
         type: 'sql',
-        condition: `${expressions?.service} IN ('${service}')`,
+        condition: `${expressions.service} IN ('${service}')`,
       });
     }
     return filters;

packages/app/src/serviceDashboard.ts

@@ -1,12 +1,27 @@
 import { useMemo } from 'react';
 import { ColumnMeta } from '@hyperdx/common-utils/dist/clickhouse';
 import { tcFromSource } from '@hyperdx/common-utils/dist/core/metadata';
+import { escapeSqlString } from '@hyperdx/common-utils/dist/core/utils';
 import { SourceKind, TTraceSource } from '@hyperdx/common-utils/dist/types';
 
 import { useColumns, useJsonColumns } from './hooks/useMetadata';
 
 const COALESCE_FIELDS_LIMIT = 100;
 
+// Ordered from highest to lowest precedence
+const DB_STATEMENT_KEYS = ['db.query.text', 'db.statement'];
+
+/**
+ * Where a skip index can be made to match a database statement. The OTel
+ * schemas index span attributes two ways depending on the ClickHouse version:
+ * newer ones an `<Name>AttributeItems` column of `key=value` strings, older
+ * ones `mapValues(<Name>Attributes)`.
+ */
+type DbStatementIndexHint = {
+  kind: 'items' | 'mapValues';
+  column: string;
+};
+
 // Helper function to format field access based on column type
 function formatFieldAccess(
   field: string,
@@ -89,18 +104,9 @@ function getDefaults({
   isAttributeFieldJSON?: boolean;
 } = {}) {
   const dbStatement = makeCoalescedFieldsAccessQuery(
-    [
-      formatFieldAccess(
-        spanAttributeField,
-        'db.query.text',
-        isAttributeFieldJSON,
-      ),
-      formatFieldAccess(
-        spanAttributeField,
-        'db.statement',
-        isAttributeFieldJSON,
-      ),
-    ],
+    DB_STATEMENT_KEYS.map(key =>
+      formatFieldAccess(spanAttributeField, key, isAttributeFieldJSON),
+    ),
     isAttributeFieldJSON,
   );
 
@@ -138,6 +144,77 @@ function getDefaults({
 
 const ENDPOINT_MATERIALIZED_COLUMN_NAME = 'endpoint';
 
+function getAttributeItemsColumn(
+  attributesField: string,
+  columns: ColumnMeta[],
+): string | undefined {
+  const itemsColumn = attributesField.replace(/Attributes$/, 'AttributeItems');
+  const exists =
+    itemsColumn !== attributesField &&
+    columns.some(column => column.name === itemsColumn);
+  return exists ? itemsColumn : undefined;
+}
+
+function getDbStatementIndexHint(
+  attributesField: string,
+  isAttributeFieldJSON: boolean,
+  columns: ColumnMeta[],
+): DbStatementIndexHint | undefined {
+  if (isAttributeFieldJSON) {
+    return undefined;
+  }
+
+  const itemsColumn = getAttributeItemsColumn(attributesField, columns);
+  if (itemsColumn) {
+    return { kind: 'items', column: itemsColumn };
+  }
+
+  const isMap = columns.some(
+    column =>
+      column.name === attributesField && column.type.startsWith('Map('),
+  );
+  return isMap ? { kind: 'mapValues', column: attributesField } : undefined;
+}
+
+function makeIndexPrefilter(
+  hint: DbStatementIndexHint,
+  statement: string,
+): string {
+  if (hint.kind === 'mapValues') {
+    return `has(mapValues(${hint.column}), '${escapeSqlString(statement)}')`;
+  }
+  return DB_STATEMENT_KEYS.map(
+    key =>
+      `has(${hint.column}, '${escapeSqlString(`${key}=${statement}`)}')`,
+  ).join(' OR ');
+}
+
+/**
+ * Condition selecting the spans of a single database statement.
+ *
+ * No skip index covers the coalesced attribute lookup, so on its own it makes
+ * ClickHouse read every granule in the time range. The prefilter matches a
+ * superset of those spans in a form an index does cover; the coalesced
+ * comparison still decides the match, so results are unchanged. It is also
+ * cheaper where no index applies, since it short-circuits the coalesce.
+ */
+export function makeDbStatementCondition({
+  expressions,
+  statement,
+}: {
+  expressions: Pick<
+    ReturnType<typeof getExpressions>,
+    'dbStatement' | 'dbStatementIndexHint'
+  >;
+  statement: string;
+}): string {
+  const equality = `${expressions.dbStatement} IN ('${escapeSqlString(statement)}')`;
+  const hint = expressions.dbStatementIndexHint;
+  return hint
+    ? `(${makeIndexPrefilter(hint, statement)}) AND ${equality}`
+    : equality;
+}
+
 export function getExpressions(
   source: TTraceSource,
   columns: ColumnMeta[],
@@ -172,6 +249,11 @@ export function getExpressions(
 
     // Database
     dbStatement: defaults.dbStatement,
+    dbStatementIndexHint: getDbStatementIndexHint(
+      spanAttributeField,
+      isAttributeFieldJSON,
+      columns,
+    ),
   };
 
   const auxExpressions = {