perf: find the most frequent Map keys with topKWeighted

.changeset/dedupe-map-keys-reads.md

@@ -0,0 +1,10 @@
+---
+'@hyperdx/common-utils': patch
+---
+
+perf: share one Map key read between concurrent callers
+
+The search page's filter panel, search bar and group-by each load the same Map
+keys when the page opens. Each sent its own text index read, which on large
+tables can use several GiB of memory. Concurrent requests for the same keys now
+share one query.

.changeset/topk-map-keys.md

@@ -0,0 +1,12 @@
+---
+'@hyperdx/common-utils': patch
+---
+
+perf: find the most frequent Map keys with topKWeighted instead of grouping
+every text index token
+
+On tables whose Map keys include unique values (IDs, hashes), the keys text
+index holds millions of distinct tokens, and grouping them all took seconds and
+many GiB of memory. Keys now come from `topKWeighted`, which uses a fixed amount
+of memory. When the keys fit in its table (3 × the key limit), it picks the
+same keys as before.

packages/common-utils/src/__tests__/metadata.test.ts

@@ -3189,7 +3189,7 @@ describe('Metadata', () => {
     });
 
     it.each(textIndexLookups)(
-      'ranks keys by row count before the LIMIT on the %s text index path',
+      'keeps the maxKeys most frequent keys with topKWeighted on the %s text index path',
       async (_, textIndexLookup) => {
         const md = new Metadata(mockClickhouseClient, new MetadataCache());
         jest.spyOn(md, 'getServerVersion').mockResolvedValue([26, 3, 0, 0]);
@@ -3220,8 +3220,42 @@ describe('Metadata', () => {
           .calls[0][0];
         expect(query).toContain('FROM mergeTreeTextIndex(');
         expect(query).toMatch(
-          /GROUP BY key HAVING key != ''\s+ORDER BY sum\(cardinality\) DESC, key\s+LIMIT/,
+          /arrayJoin\(topKWeighted\(2\)\(\w+, cardinality\)\)/,
         );
+        expect(query).not.toContain('GROUP BY');
+      },
+    );
+
+    it.each(textIndexLookups)(
+      'shares one read between concurrent callers on the %s text index path',
+      async (_, textIndexLookup) => {
+        const md = new Metadata(mockClickhouseClient, new MetadataCache());
+        jest.spyOn(md, 'getServerVersion').mockResolvedValue([26, 3, 0, 0]);
+        jest
+          .spyOn(md, 'getMapColumnTextIndexes')
+          .mockResolvedValue(textIndexLookup);
+        (mockClickhouseClient.query as jest.Mock).mockResolvedValue({
+          json: () => Promise.resolve({ data: [{ key: 'http.method' }] }),
+        });
+        const args = {
+          databaseName: 'default',
+          tableName: 'otel_logs',
+          column: 'LogAttributes',
+          connectionId: 'test_connection',
+          dateRange: [
+            new Date('2024-01-01T00:00:00Z'),
+            new Date('2024-01-01T01:00:00Z'),
+          ] as [Date, Date],
+          timestampValueExpression: 'Timestamp',
+        };
+
+        const results = await Promise.all([
+          md.getMapKeys(args),
+          md.getMapKeys(args),
+        ]);
+
+        expect(results).toEqual([['http.method'], ['http.method']]);
+        expect(mockClickhouseClient.query).toHaveBeenCalledTimes(1);
       },
     );
   });

packages/common-utils/src/core/metadata.ts

@@ -1117,12 +1117,62 @@ export class Metadata {
     const cacheKey = conditionSql
       ? `${unscopedCacheKey}.${objectHash(conditionSql)}.condition`
       : unscopedCacheKey;
-    const cachedKeys = this.cache.get<string[]>(cacheKey);
 
-    if (cachedKeys != null) {
-      return cachedKeys;
-    }
+    // getOrFetch shares one in-flight read between concurrent callers (the
+    // filter panel, search bar and group-by all load fields at once); a text
+    // index read can take GiBs of memory. A condition changes with variable
+    // selections, so its entries expire instead of piling up.
+    return this.cache.getOrFetch<string[]>(
+      cacheKey,
+      () =>
+        this.queryMapKeys({
+          databaseName,
+          tableName,
+          column,
+          maxKeys,
+          connectionId,
+          metricName,
+          metadataMVs,
+          dateRange,
+          alignedDateRange,
+          timestampValueExpression,
+          conditionSql,
+          cacheKey,
+          signal,
+        }),
+      { ttlMs: conditionSql ? CONDITION_KEYS_CACHE_TTL_MS : undefined },
+    );
+  }
 
+  private async queryMapKeys({
+    databaseName,
+    tableName,
+    column,
+    maxKeys,
+    connectionId,
+    metricName,
+    metadataMVs,
+    dateRange,
+    alignedDateRange,
+    timestampValueExpression,
+    conditionSql,
+    cacheKey,
+    signal,
+  }: {
+    databaseName: string;
+    tableName: string;
+    column: string;
+    maxKeys: number;
+    connectionId: string;
+    metricName?: string;
+    metadataMVs?: MetadataMaterializedViews;
+    dateRange: [Date, Date];
+    alignedDateRange?: [Date, Date];
+    timestampValueExpression?: string;
+    conditionSql?: string;
+    cacheKey: string;
+    signal?: AbortSignal;
+  }): Promise<string[]> {
     const textIndexInfoLookup = await this.getMapColumnTextIndexes({
       databaseName,
       tableName,
@@ -1158,13 +1208,14 @@ export class Metadata {
       });
       const index = textIndexInfo.key.indexName;
       // `cardinality` is the token's row count in the part, not a distinct count.
+      // topKWeighted keeps a fixed-size table instead of grouping every token,
+      // so millions of one-off keys cost MiBs, not GiBs. While the distinct
+      // keys fit in that table (3 * maxKeys) it picks the same keys as an exact
+      // count, apart from ties at the cutoff; the order can differ slightly.
       const sql = chSql`
-        SELECT token AS key
+        SELECT arrayJoin(topKWeighted(${{ UNSAFE_RAW_SQL: inlineNonNegativeInt(maxKeys, 'maxKeys') }})(token, cardinality)) AS key
         FROM mergeTreeTextIndex(${{ String: databaseName }}, ${{ String: tableName }}, ${{ String: index }})
-        WHERE ${partsFilter}
-        GROUP BY key HAVING key != ''
-        ORDER BY sum(cardinality) DESC, key
-        LIMIT ${{ Int32: maxKeys }}`;
+        WHERE ${partsFilter} AND token != ''`;
       try {
         const keys = await this.clickhouseClient
           .query<'JSON'>({
@@ -1176,10 +1227,7 @@ export class Metadata {
           })
           .then(r => r.json<{ key: string }>())
           .then(d => d.data.map(r => r.key).filter(Boolean));
-        if (keys.length > 0) {
-          this.cache.set(cacheKey, keys);
-          return keys;
-        }
+        if (keys.length > 0) return keys;
       } catch (e) {
         // The caller gave up; don't start the fallback with a dead signal and
         // cache its instant rejection under the rollup key.
@@ -1206,14 +1254,16 @@ export class Metadata {
       });
       const index = textIndexInfo.kv.indexName;
       const separator = textIndexInfo.kv.separator;
-      // `cardinality` is the token's row count in the part, not a distinct count.
+      // `cardinality` is the token's row count in the part, not a distinct
+      // count. See the key text index query for why topKWeighted.
       const sql = chSql`
-        SELECT splitByString(${{ String: separator }}, token)[1] AS key
-        FROM mergeTreeTextIndex(${{ String: databaseName }}, ${{ String: tableName }}, ${{ String: index }})
-        WHERE ${partsFilter}
-        GROUP BY key HAVING key != ''
-        ORDER BY sum(cardinality) DESC, key
-        LIMIT ${{ Int32: maxKeys }}`;
+        SELECT arrayJoin(topKWeighted(${{ UNSAFE_RAW_SQL: inlineNonNegativeInt(maxKeys, 'maxKeys') }})(mapKey, cardinality)) AS key
+        FROM (
+          SELECT splitByString(${{ String: separator }}, token)[1] AS mapKey, cardinality
+          FROM mergeTreeTextIndex(${{ String: databaseName }}, ${{ String: tableName }}, ${{ String: index }})
+          WHERE ${partsFilter}
+        )
+        WHERE mapKey != ''`;
       try {
         const keys = await this.clickhouseClient
           .query<'JSON'>({
@@ -1225,10 +1275,7 @@ export class Metadata {
           })
           .then(r => r.json<{ key: string }>())
           .then(d => d.data.map(r => r.key).filter(Boolean));
-        if (keys.length > 0) {
-          this.cache.set(cacheKey, keys);
-          return keys;
-        }
+        if (keys.length > 0) return keys;
       } catch (e) {
         // See the key text index catch above.
         if (signal?.aborted) throw e;
@@ -1403,42 +1450,34 @@ export class Metadata {
       `;
     }
 
-    const fetchKeys = async () => {
-      const keys = await this.clickhouseClient
-        .query<'JSON'>({
-          query: sql.sql,
-          query_params: sql.params,
-          connectionId,
-          clickhouse_settings: {
-            ...this.getClickHouseSettings(),
-            // Max 15 seconds to get keys
-            timeout_overflow_mode: 'break',
-            max_execution_time: 15,
-            // Set the value to 0 (unlimited) so that the LIMIT is used instead
-            max_rows_to_read: '0',
-          },
-          abort_signal: signal,
-        })
-        .then(res => res.json<{ keysArr?: string[]; key?: string }>())
-        .then(d => {
-          let output: string[];
-          if (strategy === 'groupUniqArrayArray') {
-            output = d.data[0].keysArr ?? [];
-          } else {
-            output = d.data
-              .map(row => row.key)
-              .filter((k): k is string => Boolean(k));
-          }
+    return this.clickhouseClient
+      .query<'JSON'>({
+        query: sql.sql,
+        query_params: sql.params,
+        connectionId,
+        clickhouse_settings: {
+          ...this.getClickHouseSettings(),
+          // Max 15 seconds to get keys
+          timeout_overflow_mode: 'break',
+          max_execution_time: 15,
+          // Set the value to 0 (unlimited) so that the LIMIT is used instead
+          max_rows_to_read: '0',
+        },
+        abort_signal: signal,
+      })
+      .then(res => res.json<{ keysArr?: string[]; key?: string }>())
+      .then(d => {
+        let output: string[];
+        if (strategy === 'groupUniqArrayArray') {
+          output = d.data[0].keysArr ?? [];
+        } else {
+          output = d.data
+            .map(row => row.key)
+            .filter((k): k is string => Boolean(k));
+        }
 
-          return output.filter(r => r);
-        });
-      return keys;
-    };
-    // A condition changes with variable selections; expire its entries so
-    // they don't pile up.
-    return this.cache.getOrFetch<string[]>(cacheKey, fetchKeys, {
-      ttlMs: conditionSql ? CONDITION_KEYS_CACHE_TTL_MS : undefined,
-    });
+        return output.filter(r => r);
+      });
   }
 
   async getJSONKeys({