feat: Add disableQueryChunking option

Author: pulpdrew

packages/app/src/hooks/__tests__/useChartConfig.test.tsx

@@ -74,15 +74,26 @@ describe('useChartConfig', () => {
       expect(
         getGranularityAlignedTimeWindows({
           granularity: '1 hour',
-        }),
+          timestampValueExpression: 'TimestampTime',
+        } as ChartConfigWithOptDateRange),
       ).toEqual([undefined]);
     });
 
     it('returns [undefined] if no granularity is provided', () => {
       expect(
         getGranularityAlignedTimeWindows({
           dateRange: [new Date('2023-01-01'), new Date('2023-01-02')],
-        }),
+          timestampValueExpression: 'TimestampTime',
+        } as ChartConfigWithOptDateRange),
+      ).toEqual([undefined]);
+    });
+
+    it('returns [undefined] if no timestampValueExpression is provided', () => {
+      expect(
+        getGranularityAlignedTimeWindows({
+          dateRange: [new Date('2023-01-01'), new Date('2023-01-02')],
+          granularity: '1 hour',
+        } as ChartConfigWithOptDateRange),
       ).toEqual([undefined]);
     });
 
@@ -95,7 +106,8 @@ describe('useChartConfig', () => {
               new Date('2023-01-10 00:10:00'),
             ],
             granularity: '1 minute',
-          },
+            timestampValueExpression: 'TimestampTime',
+          } as ChartConfigWithOptDateRange,
           [
             30, // 30s
             60, // 1m
@@ -143,7 +155,8 @@ describe('useChartConfig', () => {
               new Date('2023-01-10 00:10:00'),
             ],
             granularity: '1 minute',
-          },
+            timestampValueExpression: 'TimestampTime',
+          } as ChartConfigWithOptDateRange,
           [
             15, // 15s
           ],
@@ -175,8 +188,9 @@ describe('useChartConfig', () => {
               new Date('2023-01-10 00:00:30'),
             ],
             granularity: '1 minute',
+            timestampValueExpression: 'TimestampTime',
             dateRangeEndInclusive: true,
-          },
+          } as ChartConfigWithOptDateRange,
           [
             15 * 60, // 15m
             30 * 60, // 30m
@@ -230,7 +244,8 @@ describe('useChartConfig', () => {
               new Date('2023-01-10 00:02:00'),
             ],
             granularity: '1 minute',
-          },
+            timestampValueExpression: 'TimestampTime',
+          } as ChartConfigWithOptDateRange,
           [
             60, // 1m
           ],
@@ -373,6 +388,104 @@ describe('useChartConfig', () => {
       expect(result.current.isPending).toBe(false);
     });
 
+    it('fetches data without chunking when no timestampValueExpression is provided', async () => {
+      const config = createMockChartConfig({
+        dateRange: [
+          new Date('2025-10-01 00:00:00Z'),
+          new Date('2025-10-02 00:00:00Z'),
+        ],
+        granularity: '1 hour',
+        timestampValueExpression: undefined,
+      });
+
+      const mockResponse = createMockQueryResponse([
+        {
+          'count()': '71',
+          SeverityText: 'info',
+          __hdx_time_bucket: '2025-10-01T00:00:00Z',
+        },
+        {
+          'count()': '73',
+          SeverityText: 'info',
+          __hdx_time_bucket: '2025-10-02T00:00:00Z',
+        },
+      ]);
+
+      mockClickhouseClient.queryChartConfig.mockResolvedValue(mockResponse);
+
+      const { result } = renderHook(() => useQueriedChartConfig(config), {
+        wrapper,
+      });
+
+      await waitFor(() => expect(result.current.isSuccess).toBe(true));
+      await waitFor(() => expect(result.current.isFetching).toBe(false));
+
+      // Should only be called once since chunking is disabled without timestampValueExpression
+      expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledTimes(1);
+      expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledWith({
+        config,
+        metadata: expect.any(Object),
+        opts: {
+          abort_signal: expect.any(AbortSignal),
+        },
+      });
+      expect(result.current.data).toEqual({
+        data: mockResponse.data,
+        meta: mockResponse.meta,
+        rows: mockResponse.rows,
+      });
+    });
+
+    it('fetches data without chunking when disableQueryChunking is true', async () => {
+      const config = createMockChartConfig({
+        dateRange: [
+          new Date('2025-10-01 00:00:00Z'),
+          new Date('2025-10-02 00:00:00Z'),
+        ],
+        granularity: '1 hour',
+      });
+
+      const mockResponse = createMockQueryResponse([
+        {
+          'count()': '71',
+          SeverityText: 'info',
+          __hdx_time_bucket: '2025-10-01T00:00:00Z',
+        },
+        {
+          'count()': '73',
+          SeverityText: 'info',
+          __hdx_time_bucket: '2025-10-02T00:00:00Z',
+        },
+      ]);
+
+      mockClickhouseClient.queryChartConfig.mockResolvedValue(mockResponse);
+
+      const { result } = renderHook(
+        () => useQueriedChartConfig(config, { disableQueryChunking: true }),
+        {
+          wrapper,
+        },
+      );
+
+      await waitFor(() => expect(result.current.isSuccess).toBe(true));
+      await waitFor(() => expect(result.current.isFetching).toBe(false));
+
+      // Should only be called once since chunking is explicitly disabled
+      expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledTimes(1);
+      expect(mockClickhouseClient.queryChartConfig).toHaveBeenCalledWith({
+        config,
+        metadata: expect.any(Object),
+        opts: {
+          abort_signal: expect.any(AbortSignal),
+        },
+      });
+      expect(result.current.data).toEqual({
+        data: mockResponse.data,
+        meta: mockResponse.meta,
+        rows: mockResponse.rows,
+      });
+    });
+
     it('fetches data with chunking when granularity and date range are provided', async () => {
       const config = createMockChartConfig({
         dateRange: [
@@ -521,6 +634,8 @@ describe('useChartConfig', () => {
         rows: 4,
       });
       expect(result.current.isFetching).toBe(true);
+      expect(result.current.isLoading).toBe(false); // isLoading is false because we have partial data
+      expect(result.current.isSuccess).toBe(true); // isSuccess is true because we have partial data
 
       // Resolve the final promise to simulate data arriving
       const mockResponse3Data = createMockQueryResponse([

packages/app/src/hooks/useChartConfig.tsx

@@ -7,6 +7,7 @@ import {
 import { ClickhouseClient } from '@hyperdx/common-utils/dist/clickhouse/browser';
 import {
   DEFAULT_AUTO_GRANULARITY_MAX_BUCKETS,
+  isUsingGranularity,
   renderChartConfig,
 } from '@hyperdx/common-utils/dist/renderChartConfig';
 import { format } from '@hyperdx/common-utils/dist/sqlFormatter';
@@ -29,6 +30,13 @@ import { generateTimeWindowsDescending } from '@/utils/searchWindows';
 
 interface AdditionalUseQueriedChartConfigOptions {
   onError?: (error: Error | ClickHouseQueryError) => void;
+  /**
+   * By default, queries with large date ranges are split into multiple smaller queries to
+   * avoid overloading the ClickHouse server and running into timeouts. In some cases, such
+   * as when data is being sampled across the entire range, this chunking is not desirable
+   * and can be disabled.
+   */
+  disableQueryChunking?: boolean;
 }
 
 type TimeWindow = {
@@ -39,15 +47,14 @@ type TimeWindow = {
 type TQueryFnData = Pick<ResponseJSON<any>, 'data' | 'meta' | 'rows'> & {};
 
 export const getGranularityAlignedTimeWindows = (
-  config: Pick<
-    ChartConfigWithOptDateRange,
-    'dateRange' | 'granularity' | 'dateRangeEndInclusive'
-  >,
+  config: ChartConfigWithOptDateRange,
   windowDurationsSeconds?: number[],
 ): TimeWindow[] | [undefined] => {
-  // Granularity is required for pagination, otherwise we could break other group-bys.
-  // Date range is required for pagination, otherwise we'd have infinite pages, or some unbounded page(s).
-  if (!config.dateRange || !config.granularity) return [undefined];
+  // Granularity is required for chunking, otherwise we could break other group-bys.
+  if (!isUsingGranularity(config)) return [undefined];
+
+  // Date range is required for chunking, otherwise we'd have infinite chunks, or some unbounded chunk(s).
+  if (!config.dateRange) return [undefined];
 
   const [startDate, endDate] = config.dateRange;
   const windowsUnaligned = generateTimeWindowsDescending(
@@ -95,8 +102,11 @@ async function* fetchDataInChunks(
   config: ChartConfigWithOptDateRange,
   clickhouseClient: ClickhouseClient,
   signal: AbortSignal,
+  disableQueryChunking: boolean = false,
 ) {
-  const windows = getGranularityAlignedTimeWindows(config);
+  const windows = disableQueryChunking
+    ? ([undefined] as const)
+    : getGranularityAlignedTimeWindows(config);
 
   let query = null;
   if (IS_MTVIEWS_ENABLED) {
@@ -128,6 +138,20 @@ async function* fetchDataInChunks(
   }
 }
 
+/**
+ * A hook providing data queried based on the provided chart config.
+ *
+ * If all of the following are true, the query will be chunked into multiple smaller queries:
+ * - The config includes a date range
+ * - The config includes a granularity
+ * - `options.disableQueryChunking` is falsy
+ *
+ * For chunked queries, note the following:
+ * - The config's limit, if provided, is applied to each chunk, so the total number
+ *   of rows returned may be up to `limit * number_of_chunks`.
+ * - The returned data will be ordered within each chunk, and chunks will
+ *   be ordered oldest-first, by the timestampValueExpression.
+ */
 export function useQueriedChartConfig(
   config: ChartConfigWithOptDateRange,
   options?: Partial<UseQueryOptions<ResponseJSON<any>>> &
@@ -139,7 +163,12 @@ export function useQueriedChartConfig(
     queryKey: [config],
     queryFn: streamedQuery({
       streamFn: context =>
-        fetchDataInChunks(config, clickhouseClient, context.signal),
+        fetchDataInChunks(
+          config,
+          clickhouseClient,
+          context.signal,
+          options?.disableQueryChunking,
+        ),
       /**
        * This mode ensures that data remains in the cache until the next full streamed result is available.
        * By default, the cache would be cleared before new data starts arriving, which results in the query briefly

packages/app/src/hooks/usePatterns.tsx

@@ -143,7 +143,11 @@ function usePatterns({
 
   const { data: sampleRows } = useQueriedChartConfig(
     configWithPrimaryAndPartitionKey ?? config, // `config` satisfying type, never used due to `enabled` check
-    { enabled: configWithPrimaryAndPartitionKey != null && enabled },
+    {
+      enabled: configWithPrimaryAndPartitionKey != null && enabled,
+      // Disable chunking so that we get a uniform sample across the entire date range
+      disableQueryChunking: true,
+    },
   );
 
   const { data: pyodide, isLoading: isLoadingPyodide } = usePyodide({

packages/common-utils/src/renderChartConfig.ts

@@ -58,7 +58,7 @@ export function isUsingGroupBy(
   return chartConfig.groupBy != null && chartConfig.groupBy.length > 0;
 }
 
-function isUsingGranularity(
+export function isUsingGranularity(
   chartConfig: ChartConfigWithOptDateRange,
 ): chartConfig is Omit<
   Omit<Omit<ChartConfigWithDateRange, 'granularity'>, 'dateRange'>,