perf: find the most frequent Map keys with topKWeighted

.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,9 @@ 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');
       },
     );
 

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

@@ -1208,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'>({
@@ -1253,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'>({