perf: Improve getKeyValues query performance for JSON keys

Author: pulpdrew

.changeset/sweet-vans-pump.md

@@ -0,0 +1,5 @@
+---
+"@hyperdx/common-utils": patch
+---
+
+perf: Improve getKeyValues query performance for JSON keys

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

@@ -251,22 +251,6 @@ describe('Metadata', () => {
             ],
           }),
       });
-
-      // Mock getColumns to return columns including materialized ones
-      mockCache.getOrFetch.mockImplementation((key, queryFn) => {
-        if (key.includes('.columns')) {
-          return Promise.resolve([
-            { name: 'regular_column', type: 'String', default_type: '' },
-            {
-              name: 'materialized_column',
-              type: 'String',
-              default_type: 'MATERIALIZED',
-            },
-            { name: 'default_column', type: 'String', default_type: 'DEFAULT' },
-          ]);
-        }
-        return queryFn();
-      });
     });
 
     it('should apply row limit when disableRowLimit is false', async () => {
@@ -355,54 +339,15 @@ describe('Metadata', () => {
       expect(result).toEqual([{ key: 'column1', value: ['value1', 'value2'] }]);
     });
 
-    it('should include materialized fields when selecting all columns', async () => {
+    it('should fallback to * when no keys are provided', async () => {
       const renderChartConfigSpy = jest.spyOn(
         renderChartConfigModule,
         'renderChartConfig',
       );
 
       await metadata.getKeyValues({
         chartConfig: mockChartConfig,
-        keys: ['column1'],
-        limit: 10,
-      });
-
-      // Verify that renderChartConfig was called with the expanded select list
-      // that includes all columns by name (including materialized ones)
-      expect(renderChartConfigSpy).toHaveBeenCalledWith(
-        expect.objectContaining({
-          with: [
-            expect.objectContaining({
-              name: 'sampledData',
-              chartConfig: expect.objectContaining({
-                // Should expand to all column names instead of using '*'
-                select:
-                  '`regular_column`, `materialized_column`, `default_column`',
-              }),
-            }),
-          ],
-        }),
-        metadata,
-      );
-    });
-
-    it('should fallback to * when no columns are found', async () => {
-      // Mock getColumns to return empty array
-      mockCache.getOrFetch.mockImplementation((key, queryFn) => {
-        if (key.includes('.columns')) {
-          return Promise.resolve([]);
-        }
-        return queryFn();
-      });
-
-      const renderChartConfigSpy = jest.spyOn(
-        renderChartConfigModule,
-        'renderChartConfig',
-      );
-
-      await metadata.getKeyValues({
-        chartConfig: mockChartConfig,
-        keys: ['column1'],
+        keys: [],
         limit: 10,
       });
 

packages/common-utils/src/metadata.ts

@@ -668,29 +668,20 @@ export class Metadata {
     return this.cache.getOrFetch(
       `${chartConfig.connection}.${chartConfig.from.databaseName}.${chartConfig.from.tableName}.${keys.join(',')}.${chartConfig.dateRange.toString()}.${disableRowLimit}.values`,
       async () => {
-        const selectClause = keys
-          .map((k, i) => `groupUniqArray(${limit})(${k}) AS param${i}`)
-          .join(', ');
-
         // When disableRowLimit is true, query directly without CTE
         // Otherwise, use CTE with row limits for sampling
         const sqlConfig = disableRowLimit
           ? {
               ...chartConfig,
-              select: selectClause,
+              select: keys
+                .map((k, i) => `groupUniqArray(${limit})(${k}) AS param${i}`)
+                .join(', '),
             }
           : await (async () => {
-              // Get all columns including materialized ones
-              const columns = await this.getColumns({
-                databaseName: chartConfig.from.databaseName,
-                tableName: chartConfig.from.tableName,
-                connectionId: chartConfig.connection,
-              });
-
               // Build select expression that includes all columns by name
               // This ensures materialized columns are included
               const selectExpr =
-                columns.map(col => `\`${col.name}\``).join(', ') || '*';
+                keys.map((k, i) => `${k} as param${i}`).join(', ') || '*';
 
               return {
                 with: [
@@ -710,7 +701,12 @@ export class Metadata {
                     isSubquery: true,
                   },
                 ],
-                select: selectClause,
+                select: keys
+                  .map(
+                    (_, i) =>
+                      `groupUniqArray(${limit})(param${i}) AS param${i}`,
+                  )
+                  .join(', '),
                 connection: chartConfig.connection,
                 from: { databaseName: '', tableName: 'sampledData' },
                 where: '',