Add collapsible groups to pivot table and simplify the sql_data_source caching model

Change-Id: Ie7d25759b7fe983e06cbe24164f49dea3102d5a9
diff --git a/ui/src/assets/components/datagrid.scss b/ui/src/assets/components/datagrid.scss
index 5d71d7e..0f289e2 100644
--- a/ui/src/assets/components/datagrid.scss
+++ b/ui/src/assets/components/datagrid.scss
@@ -204,3 +204,8 @@
     flex-grow: 1;
   }
 }
+
+.pf-data-grid__muted-cell {
+  color: var(--pf-color-text-muted);
+  opacity: 0.6;
+}
diff --git a/ui/src/components/widgets/datagrid/add_column_menu.ts b/ui/src/components/widgets/datagrid/add_column_menu.ts
index e5b628b..36d29da 100644
--- a/ui/src/components/widgets/datagrid/add_column_menu.ts
+++ b/ui/src/components/widgets/datagrid/add_column_menu.ts
@@ -101,7 +101,8 @@
     } else if (isParameterizedColumnDef(entry)) {
       // Parameterized column - show available keys from datasource
       const title = typeof entry.title === 'string' ? entry.title : columnName;
-      const availableKeys = context.dataSource.parameterKeys?.get(fullPath);
+      const availableKeys =
+        context.dataSource.getParameterKeys(fullPath).result;
       menuItems.push(
         m(
           MenuItem,
diff --git a/ui/src/components/widgets/datagrid/data_source.ts b/ui/src/components/widgets/datagrid/data_source.ts
index 8b09e06..d7a7c2a 100644
--- a/ui/src/components/widgets/datagrid/data_source.ts
+++ b/ui/src/components/widgets/datagrid/data_source.ts
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import {UseQueryResult} from '../../../trace_processor/query_cache';
 import {Row, SqlValue} from '../../../trace_processor/query_result';
 import {Column, Filter, Pivot} from './model';
 
@@ -21,31 +22,36 @@
 }
 
 export interface DataSource {
-  // The row data for the current data grid state (filters, sorting, pagination,
-  // etc)
-  readonly rows?: DataSourceRows;
-
-  // Available distinct values for specified columns (for filter dropdowns)
-  readonly distinctValues?: ReadonlyMap<string, readonly SqlValue[]>;
-
-  // Available parameter keys for parameterized columns (e.g., for 'args' ->
-  // ['foo', 'bar'])
-  readonly parameterKeys?: ReadonlyMap<string, readonly string[]>;
-
-  // Computed aggregate totals for each aggregate column (grand total across all
-  // filtered rows)
-  readonly aggregateTotals?: ReadonlyMap<string, SqlValue>;
-
-  // Whether the data source is currently loading data/updating.
-  readonly isLoading?: boolean;
-
-  // Called when the data grid parameters change (sorting, filtering,
-  // pagination, etc), which might trigger a data reload.
-  notify(model: DataSourceModel): void;
-
   // Export all data with current filters/sorting applied. Returns a promise
   // that resolves to all filtered and sorted rows.
   exportData(): Promise<readonly Row[]>;
+
+  getRows(
+    model: DataSourceModel,
+  ): UseQueryResult<{totalRows: number; offset: number; rows: Row[]}>;
+
+  getDistinctValues(columnPath: string): UseQueryResult<readonly SqlValue[]>;
+
+  getAggregateTotals(
+    model: DataSourceModel,
+  ): UseQueryResult<ReadonlyMap<string, SqlValue>>;
+
+  getPivotRollups(model: DataSourceModel): UseQueryResult<PivotRollups>;
+
+  getParameterKeys(prefix: string): UseQueryResult<readonly string[]>;
+}
+
+/**
+ * Rollup data for hierarchical pivot tables. Contains aggregated rows at each
+ * level of the groupBy hierarchy.
+ */
+export interface PivotRollups {
+  // Rollup rows indexed by depth level.
+  // Level 0: rows grouped by first groupBy column only (top-level groups)
+  // Level 1: rows grouped by first two groupBy columns
+  // etc.
+  // The deepest level (all groupBy columns) comes from the main `rows` field.
+  readonly byLevel: ReadonlyMap<number, readonly Row[]>;
 }
 
 export interface DataSourceModel {
diff --git a/ui/src/components/widgets/datagrid/datagrid.ts b/ui/src/components/widgets/datagrid/datagrid.ts
index 1ccdcb5..bbbe3c2 100644
--- a/ui/src/components/widgets/datagrid/datagrid.ts
+++ b/ui/src/components/widgets/datagrid/datagrid.ts
@@ -25,6 +25,7 @@
 import {
   Grid,
   GridApi,
+  GridAttrs,
   GridCell,
   GridColumn,
   GridHeaderCell,
@@ -41,7 +42,7 @@
 import {CellFilterMenu} from './cell_filter_menu';
 import {FilterMenu} from './column_filter_menu';
 import {ColumnInfoMenu} from './column_info_menu';
-import {DataSource} from './data_source';
+import {DataSource, PivotRollups} from './data_source';
 import {
   SchemaRegistry,
   getColumnInfo,
@@ -283,10 +284,10 @@
 }
 
 function getOrCreateDataSource(data: DataSource | readonly Row[]): DataSource {
-  if ('notify' in data) {
-    return data;
-  } else {
+  if (Array.isArray(data)) {
     return new InMemoryDataSource(data);
+  } else {
+    return data as DataSource;
   }
 }
 
@@ -298,7 +299,8 @@
   readonly schema: SchemaRegistry;
   readonly rootSchema: string;
   readonly datasource: DataSource;
-  readonly result: DataSource['rows'];
+  readonly rowsResult?: {totalRows: number; offset: number; rows: Row[]};
+  readonly aggregateTotals?: ReadonlyMap<string, SqlValue>;
   readonly columnInfoCache: Map<string, ReturnType<typeof getColumnInfo>>;
   readonly structuredQueryCompatMode: boolean;
   readonly enablePivotControls: boolean;
@@ -312,7 +314,9 @@
   readonly schema: SchemaRegistry;
   readonly rootSchema: string;
   readonly datasource: DataSource;
-  readonly result: DataSource['rows'];
+  readonly rowsResult?: {totalRows: number; offset: number; rows: Row[]};
+  readonly aggregateTotals?: ReadonlyMap<string, SqlValue>;
+  readonly pivotRollups?: PivotRollups;
   readonly pivot: Pivot;
   readonly structuredQueryCompatMode: boolean;
   readonly enablePivotControls: boolean;
@@ -337,6 +341,10 @@
   // Track parameterized columns needing key discovery
   private parameterKeyColumns = new Set<string>();
 
+  // Track expanded groups for hierarchical pivot mode
+  // Group keys are formed by joining parent values with '|' separator
+  private expandedGroups = new Set<string>();
+
   oninit({attrs}: m.Vnode<DataGridAttrs>) {
     if (attrs.initialColumns) {
       this.columns = attrs.initialColumns;
@@ -426,17 +434,30 @@
 
     // Notify the data source of the current model state.
     const datasource = getOrCreateDataSource(data);
-    datasource.notify({
+
+    // Check if hierarchical pivot mode - if so, we need all rows (no pagination)
+    // because we need to match children to expanded parents
+    const isHierarchicalPivotMode = this.isHierarchicalPivotMode(this.pivot);
+
+    const model = {
       columns: columnsWithDeps,
       filters: this.filters,
-      pagination: {
-        offset: this.paginationOffset,
-        limit: this.paginationLimit,
-      },
+      // Don't paginate in hierarchical pivot mode - we need all rows to match
+      // children to expanded parents
+      pagination: isHierarchicalPivotMode
+        ? undefined
+        : {
+            offset: this.paginationOffset,
+            limit: this.paginationLimit,
+          },
       pivot: this.pivot,
       distinctValuesColumns: this.distinctValuesColumns,
       parameterKeyColumns: this.parameterKeyColumns,
-    });
+    };
+    const {result: rowsResult, isLoading: rowsLoading} =
+      datasource.getRows(model);
+    const {result: aggregateTotals} = datasource.getAggregateTotals(model);
+    const {result: pivotRollups} = datasource.getPivotRollups(model);
 
     // Expose the API
     attrs.onReady?.({
@@ -450,13 +471,10 @@
         );
       },
       getRowCount: () => {
-        return datasource?.rows?.totalRows;
+        return rowsResult?.totalRows;
       },
     });
 
-    // Extract the result from the datasource
-    const result = datasource.rows;
-
     // Determine if we're in pivot mode (has groupBy columns and not drilling down)
     const isPivotMode =
       this.pivot !== undefined &&
@@ -474,7 +492,9 @@
         schema,
         rootSchema,
         datasource,
-        result,
+        rowsResult,
+        aggregateTotals,
+        pivotRollups,
         pivot: this.pivot!,
         structuredQueryCompatMode,
         enablePivotControls,
@@ -497,7 +517,8 @@
         schema,
         rootSchema,
         datasource,
-        result,
+        rowsResult,
+        aggregateTotals,
         columnInfoCache,
         structuredQueryCompatMode,
         enablePivotControls,
@@ -538,34 +559,40 @@
           }),
         ),
         drillDown: this.pivot?.drillDown,
-        drillDownFields: this.pivot?.groupBy.map(({field}) => {
-          const colInfo = getColumnInfo(schema, rootSchema, field);
-          const titleParts = colInfo?.titleParts ?? field.split('.');
-          const rawValue = this.pivot?.drillDown?.[field];
-          return {
-            title: buildColumnTitle(titleParts),
-            value: formatChipValue(rawValue, colInfo?.cellFormatter),
-          };
-        }),
+        drillDownFields: this.pivot?.groupBy
+          .filter(
+            ({field}) => this.pivot?.drillDown && field in this.pivot.drillDown,
+          )
+          .map(({field}) => {
+            const colInfo = getColumnInfo(schema, rootSchema, field);
+            const titleParts = colInfo?.titleParts ?? field.split('.');
+            const rawValue = this.pivot?.drillDown?.[field];
+            return {
+              title: buildColumnTitle(titleParts),
+              value: formatChipValue(rawValue, colInfo?.cellFormatter),
+            };
+          }),
         onExitDrillDown: () => this.exitDrillDown(attrs),
       }),
       m(LinearProgress, {
         className: 'pf-data-grid__loading',
-        state: datasource.isLoading ? 'indeterminate' : 'none',
+        state: rowsLoading ? 'indeterminate' : 'none',
       }),
       m(Grid, {
         className: 'pf-data-grid__table',
         columns: gridColumns,
-        rowData: {
-          data: gridRows,
-          total: result?.totalRows ?? 0,
-          offset: Math.max(result?.rowOffset ?? 0, this.paginationOffset),
-          onLoadData: (offset, limit) => {
-            this.paginationOffset = offset;
-            this.paginationLimit = limit;
-            m.redraw();
-          },
-        },
+        rowData: isHierarchicalPivotMode
+          ? gridRows
+          : {
+              data: gridRows,
+              total: rowsResult?.totalRows ?? 0,
+              offset: Math.max(rowsResult?.offset ?? 0, this.paginationOffset),
+              onLoadData: (offset, limit) => {
+                this.paginationOffset = offset;
+                this.paginationLimit = limit;
+                m.redraw();
+              },
+            },
         virtualization: {
           rowHeightPx: 25,
         },
@@ -581,8 +608,8 @@
           this.gridApi = api;
         },
         emptyState:
-          result?.totalRows === 0 &&
-          !datasource.isLoading &&
+          rowsResult?.totalRows === 0 &&
+          !rowsLoading &&
           m(
             EmptyState,
             {
@@ -600,7 +627,7 @@
                 onclick: () => this.clearFilters(attrs),
               }),
           ),
-      }),
+      } satisfies GridAttrs),
     );
   }
 
@@ -779,10 +806,14 @@
     if (!this.pivot) return;
 
     // Build drill-down filter from groupBy column values
+    // Only include fields that exist in the row (rollup/summary rows don't have
+    // values for lower-level groupBy columns)
     const drillDownRow: Row = {};
     for (const groupByCol of this.pivot.groupBy) {
       const field = groupByCol.field;
-      drillDownRow[field] = row[field];
+      if (field in row) {
+        drillDownRow[field] = row[field];
+      }
     }
 
     const newPivot: Pivot = {...this.pivot, drillDown: drillDownRow};
@@ -798,6 +829,44 @@
     attrs.onPivotChanged?.(newPivot);
   }
 
+  /**
+   * Toggle expansion state of a group in hierarchical pivot mode.
+   * @param groupKey The key identifying the group (e.g., "Chrome" or "Chrome|Main")
+   */
+  private toggleGroupExpansion(groupKey: string): void {
+    if (this.expandedGroups.has(groupKey)) {
+      this.expandedGroups.delete(groupKey);
+    } else {
+      this.expandedGroups.add(groupKey);
+    }
+  }
+
+  /**
+   * Check if hierarchical pivot mode is active.
+   * Hierarchical mode is when we have 2+ groupBy columns and are not in drillDown.
+   */
+  private isHierarchicalPivotMode(pivot: Pivot | undefined): boolean {
+    return pivot !== undefined && !pivot.drillDown && pivot.groupBy.length >= 2;
+  }
+
+  /**
+   * Build a group key from a row and the groupBy columns up to a certain level.
+   * @param row The row containing the values
+   * @param groupByFields Array of field names to include
+   * @param upToLevel How many levels to include (0 = first field only)
+   */
+  private buildGroupKey(
+    row: Row,
+    groupByFields: readonly string[],
+    upToLevel: number,
+  ): string {
+    const parts: string[] = [];
+    for (let i = 0; i <= upToLevel && i < groupByFields.length; i++) {
+      parts.push(String(row[groupByFields[i]] ?? ''));
+    }
+    return parts.join('|');
+  }
+
   private removeAggregateColumn(index: number, attrs: DataGridAttrs): void {
     if (!this.pivot?.aggregates) return;
 
@@ -1051,6 +1120,7 @@
       schema,
       rootSchema,
       datasource,
+      aggregateTotals,
       columnInfoCache,
       structuredQueryCompatMode,
       enablePivotControls,
@@ -1080,7 +1150,7 @@
         m(FilterMenu, {
           columnType,
           structuredQueryCompatMode,
-          distinctValues: datasource.distinctValues?.get(field),
+          distinctValues: datasource.getDistinctValues(field).result,
           valueFormatter: (v) => colInfo?.cellFormatter?.(v, {}) ?? String(v),
           onFilterAdd: (filter) => this.addFilter({field, ...filter}, attrs),
           onRequestDistinctValues: () => this.distinctValuesColumns.add(field),
@@ -1143,7 +1213,7 @@
       // Build subContent showing grand total if column has an aggregate
       let subContent: m.Children;
       if (aggregate) {
-        const totalValue = datasource.aggregateTotals?.get(field);
+        const totalValue = aggregateTotals?.get(field);
         const isLoading = totalValue === undefined;
         // Don't show grand total for ANY aggregation (it's just an arbitrary value)
         let totalContent: m.Children;
@@ -1185,13 +1255,13 @@
    * Builds grid rows for flat (non-pivot) mode.
    */
   private buildFlatRows(ctx: FlatGridBuildContext): m.Children[][] {
-    const {attrs, result, columnInfoCache} = ctx;
+    const {attrs, rowsResult, columnInfoCache} = ctx;
 
-    if (result === undefined) return [];
+    if (rowsResult === undefined) return [];
 
     // Find the intersection of rows between what we have and what is required
     // and only render those.
-    const start = Math.max(result.rowOffset, this.paginationOffset);
+    const start = Math.max(rowsResult.offset, this.paginationOffset);
 
     const rowIndices = Array.from(
       {length: this.paginationLimit},
@@ -1200,7 +1270,7 @@
 
     return rowIndices
       .map((index) => {
-        const row = result.rows[index - result.rowOffset];
+        const row = rowsResult.rows[index - rowsResult.offset];
         if (row === undefined) return undefined;
 
         return this.columns.map((col) => {
@@ -1246,6 +1316,7 @@
       schema,
       rootSchema,
       datasource,
+      aggregateTotals,
       pivot,
       structuredQueryCompatMode,
       enablePivotControls,
@@ -1284,7 +1355,7 @@
         m(FilterMenu, {
           columnType,
           structuredQueryCompatMode,
-          distinctValues: datasource.distinctValues?.get(field),
+          distinctValues: datasource.getDistinctValues(field).result,
           valueFormatter: (v) => colInfo?.cellFormatter?.(v, {}) ?? String(v),
           onFilterAdd: (filter) => this.addFilter({field, ...filter}, attrs),
           onRequestDistinctValues: () => this.distinctValuesColumns.add(field),
@@ -1421,7 +1492,7 @@
 
       // Build subContent showing grand total with aggregate symbol
       // Don't show grand total for ANY aggregation (it's just an arbitrary value)
-      const aggregateTotalValue = datasource.aggregateTotals?.get(alias);
+      const aggregateTotalValue = aggregateTotals?.get(alias);
       const symbol = agg.function;
       const isLoading = aggregateTotalValue === undefined;
       let aggTotalContent: m.Children;
@@ -1462,26 +1533,34 @@
   /**
    * Builds grid rows for pivot mode.
    * Each row contains values for groupBy columns and aggregate columns.
+   * In hierarchical mode (2+ groupBy columns), builds a tree structure with
+   * expand/collapse capability.
    */
   private buildPivotRows(ctx: PivotGridBuildContext): m.Children[][] {
-    const {attrs, schema, rootSchema, result, pivot, enablePivotControls} = ctx;
+    const {attrs, schema, rootSchema, rowsResult, pivot, enablePivotControls} =
+      ctx;
 
-    if (result === undefined) return [];
+    if (rowsResult === undefined) return [];
 
+    // For hierarchical mode, build visible rows from rollups + expanded children
+    if (this.isHierarchicalPivotMode(pivot)) {
+      return this.buildHierarchicalPivotRows(ctx);
+    }
+
+    const aggregates = pivot.aggregates ?? [];
+
+    // Non-hierarchical pivot mode (single groupBy column)
     // Find the intersection of rows between what we have and what is required
-    // and only render those.
-    const start = Math.max(result.rowOffset, this.paginationOffset);
+    const start = Math.max(rowsResult.offset, this.paginationOffset);
 
     const rowIndices = Array.from(
       {length: this.paginationLimit},
       (_, i) => i + start,
     );
 
-    const aggregates = pivot.aggregates ?? [];
-
     return rowIndices
       .map((index) => {
-        const row = result.rows[index - result.rowOffset];
+        const row = rowsResult.rows[index - rowsResult.offset];
         if (row === undefined) return undefined;
 
         const cells: m.Children[] = [];
@@ -1585,6 +1664,212 @@
       .filter(exists);
   }
 
+  /**
+   * Builds hierarchical pivot rows with expand/collapse functionality.
+   * Shows rollup rows at each level, with children visible when expanded.
+   */
+  private buildHierarchicalPivotRows(
+    ctx: PivotGridBuildContext,
+  ): m.Children[][] {
+    const {
+      attrs,
+      schema,
+      rootSchema,
+      rowsResult,
+      pivotRollups,
+      pivot,
+      enablePivotControls,
+    } = ctx;
+
+    const aggregates = pivot.aggregates ?? [];
+    const groupByFields = pivot.groupBy.map(({field}) => field);
+    const numLevels = groupByFields.length;
+    const rollups = pivotRollups;
+
+    // Build the visible rows list based on expansion state
+    const visibleRows: Array<{row: Row; level: number; parentKey: string}> = [];
+
+    // Helper to add rows recursively based on expansion state
+    const addVisibleRows = (parentKey: string, level: number) => {
+      // Get rows at this level
+      let rowsAtLevel: readonly Row[];
+      if (level === numLevels - 1) {
+        // Deepest level comes from main result rows
+        rowsAtLevel = rowsResult?.rows ?? [];
+      } else {
+        // Rollup rows from datasource
+        rowsAtLevel = rollups?.byLevel.get(level) ?? [];
+      }
+
+      // Filter rows that belong to this parent
+      const matchingRows = rowsAtLevel.filter((row) => {
+        if (parentKey === '') return true; // Top level, all rows match
+        // Check if row's parent values match the parent key
+        const rowParentKey = this.buildGroupKey(row, groupByFields, level - 1);
+        return rowParentKey === parentKey;
+      });
+
+      for (const row of matchingRows) {
+        visibleRows.push({row, level, parentKey});
+
+        // If this level is not the deepest and is expanded, add children
+        if (level < numLevels - 1) {
+          const groupKey = this.buildGroupKey(row, groupByFields, level);
+          if (this.expandedGroups.has(groupKey)) {
+            addVisibleRows(groupKey, level + 1);
+          }
+        }
+      }
+    };
+
+    // Start from level 0 (top-level groups)
+    addVisibleRows('', 0);
+
+    // Build cells for each visible row
+    return visibleRows.map(({row, level}) => {
+      const cells: m.Children[] = [];
+      const groupKey = this.buildGroupKey(row, groupByFields, level);
+      const isExpandable = level < numLevels - 1;
+      const isExpanded = this.expandedGroups.has(groupKey);
+
+      // Render groupBy columns
+      for (let i = 0; i < numLevels; i++) {
+        const field = groupByFields[i];
+        const value = row[field];
+        const colInfo = getColumnInfo(schema, rootSchema, field);
+        const cellRenderer =
+          colInfo?.cellRenderer ?? ((v: SqlValue) => renderCell(v, field));
+
+        // For cells to the right of current level (i > level), show blank
+        // These are child columns that don't apply to parent/summary rows
+        if (i > level) {
+          cells.push(
+            m(GridCell, {
+              className: 'pf-data-grid__groupby-column',
+            }),
+          );
+        } else if (i < level) {
+          // For cells to the left (parent columns), show the inherited value
+          // with indentation to indicate hierarchy
+          const rendered = cellRenderer(value, row);
+          const isRich = isCellRenderResult(rendered);
+
+          cells.push(
+            m(
+              GridCell,
+              {
+                align: isRich ? rendered.align ?? 'left' : getAligment(value),
+                nullish: isRich
+                  ? rendered.nullish ?? value === null
+                  : value === null,
+                className: classNames(
+                  'pf-data-grid__groupby-column',
+                  'pf-data-grid__muted-cell',
+                ),
+                indent: 1,
+              },
+              isRich ? rendered.content : rendered,
+            ),
+          );
+        } else {
+          const rendered = cellRenderer(value, row);
+          const isRich = isCellRenderResult(rendered);
+
+          // Determine chevron state for GridCell's built-in chevron support
+          const chevronState = isExpandable
+            ? isExpanded
+              ? 'expanded'
+              : 'collapsed'
+            : undefined;
+
+          cells.push(
+            m(
+              GridCell,
+              {
+                align: isRich ? rendered.align ?? 'left' : getAligment(value),
+                nullish: isRich
+                  ? rendered.nullish ?? value === null
+                  : value === null,
+                className: 'pf-data-grid__groupby-column',
+                chevron: chevronState,
+                onChevronClick: isExpandable
+                  ? () => this.toggleGroupExpansion(groupKey)
+                  : undefined,
+                menuItems: [
+                  enablePivotControls && [
+                    m(MenuItem, {
+                      label: 'Drill down',
+                      icon: 'zoom_in',
+                      onclick: () => this.drillDown(row, attrs),
+                    }),
+                    m(MenuDivider),
+                  ],
+                  m(CellFilterMenu, {
+                    value,
+                    onFilterAdd: (filter) =>
+                      this.addFilter({field, ...filter}, attrs),
+                  }),
+                ],
+              },
+              isRich ? rendered.content : rendered,
+            ),
+          );
+        }
+      }
+
+      // Drill-down cell
+      if (enablePivotControls) {
+        cells.push(
+          m(
+            '.pf-datagrid__dd-cell',
+            {className: 'pf-datagrid__dd'},
+            m(Button, {
+              className:
+                'pf-visible-on-row-hover pf-datagrid__drilldown-button',
+              icon: Icons.GoTo,
+              rounded: true,
+              title: 'Drill down into this group',
+              fillWidth: true,
+              onclick: () => {
+                this.drillDown(row, attrs);
+              },
+            }),
+          ),
+        );
+      }
+
+      // Render aggregate columns
+      for (const agg of aggregates) {
+        const alias = getAggregateAlias(agg);
+        const value = row[alias];
+
+        let cellRenderer =
+          'field' in agg
+            ? getColumnInfo(schema, rootSchema, agg.field)?.cellRenderer
+            : undefined;
+        cellRenderer ??= (v: SqlValue) => renderCell(v, alias);
+
+        const rendered = cellRenderer(value, row);
+        const isRich = isCellRenderResult(rendered);
+
+        cells.push(
+          m(
+            GridCell,
+            {
+              align: isRich ? rendered.align ?? 'left' : getAligment(value),
+              nullish: isRich
+                ? rendered.nullish ?? value === null
+                : value === null,
+            },
+            isRich ? rendered.content : rendered,
+          ),
+        );
+      }
+
+      return cells;
+    });
+  }
+
   private async formatData(
     dataSource: DataSource,
     schema: SchemaRegistry | undefined,
diff --git a/ui/src/components/widgets/datagrid/in_memory_data_source.ts b/ui/src/components/widgets/datagrid/in_memory_data_source.ts
index c05f5ff..a9b6d87 100644
--- a/ui/src/components/widgets/datagrid/in_memory_data_source.ts
+++ b/ui/src/components/widgets/datagrid/in_memory_data_source.ts
@@ -12,219 +12,365 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import {stringifyJsonWithBigints} from '../../../base/json_utils';
 import {assertUnreachable} from '../../../base/logging';
+import {UseQueryResult} from '../../../trace_processor/query_cache';
 import {Row, SqlValue} from '../../../trace_processor/query_result';
-import {DataSource, DataSourceModel, DataSourceRows} from './data_source';
-import {Column, Filter, Pivot} from './model';
+import {DataSource, DataSourceModel, PivotRollups} from './data_source';
+import {AggregateColumn, Column, Filter, Pivot} from './model';
 
 export class InMemoryDataSource implements DataSource {
   private data: ReadonlyArray<Row> = [];
-  private filteredSortedData: ReadonlyArray<Row> = [];
   private distinctValuesCache = new Map<string, ReadonlyArray<SqlValue>>();
   private parameterKeysCache = new Map<string, ReadonlyArray<string>>();
-  private aggregateTotalsCache = new Map<string, SqlValue>();
 
-  // Cached state for diffing
-  private oldColumns?: readonly Column[];
-  private oldFilters: ReadonlyArray<Filter> = [];
-  private oldPivot?: Pivot;
+  // Last-result caches for expensive operations
+  private rowsCache?: {
+    key: string;
+    result: {totalRows: number; offset: number; rows: Row[]};
+  };
+  private aggregateTotalsCache?: {
+    key: string;
+    result: ReadonlyMap<string, SqlValue>;
+  };
+  private pivotRollupsCache?: {
+    key: string;
+    result: PivotRollups;
+  };
 
   constructor(data: ReadonlyArray<Row>) {
     this.data = data;
-    this.filteredSortedData = data;
   }
 
-  get rows(): DataSourceRows {
+  /**
+   * Creates a cache key from model parameters.
+   * Uses JSON.stringify for simplicity - works well for small objects.
+   */
+  private buildCacheKey(parts: Record<string, unknown>): string {
+    return JSON.stringify(parts);
+  }
+
+  getRows(
+    model: DataSourceModel,
+  ): UseQueryResult<{totalRows: number; offset: number; rows: Row[]}> {
+    const {columns, filters = [], pivot, pagination} = model;
+
+    // Check cache first
+    const cacheKey = this.buildCacheKey({columns, filters, pivot, pagination});
+    if (this.rowsCache?.key === cacheKey) {
+      return {result: this.rowsCache.result, isLoading: false};
+    }
+
+    // In pivot mode, separate filters into pre-pivot and post-pivot
+    const aggregates = pivot?.aggregates ?? [];
+    const aggregateFields = new Set(
+      aggregates.map((a) => ('field' in a ? a.field : '__count__')),
+    );
+    const prePivotFilters =
+      pivot && !pivot.drillDown
+        ? filters.filter((f) => !aggregateFields.has(f.field))
+        : filters;
+    const postPivotFilters =
+      pivot && !pivot.drillDown
+        ? filters.filter((f) => aggregateFields.has(f.field))
+        : [];
+
+    // Apply pre-pivot filters (on source data)
+    let result = this.applyFilters(this.data, prePivotFilters);
+
+    // Apply pivot (but not in drilldown mode - drilldown shows raw data)
+    if (pivot && !pivot.drillDown) {
+      result = this.applyPivoting(result, pivot);
+      // Apply post-pivot filters (on aggregate results)
+      if (postPivotFilters.length > 0) {
+        result = this.applyFilters(result, postPivotFilters);
+      }
+    } else if (pivot?.drillDown) {
+      // Drilldown mode: filter to show only rows matching the drillDown values
+      result = this.applyDrillDown(result, pivot);
+    }
+
+    // Apply sorting - find sorted column from columns or pivot
+    const sortedColumn = this.findSortedColumn(columns, pivot);
+    if (sortedColumn) {
+      result = this.applySorting(
+        result,
+        sortedColumn.field,
+        sortedColumn.direction,
+      );
+    }
+
+    const totalRows = result.length;
+    const offset = pagination?.offset ?? 0;
+    const limit = pagination?.limit ?? result.length;
+    const rows = [...result.slice(offset, offset + limit)];
+
+    const cachedResult = {totalRows, offset, rows};
+    this.rowsCache = {key: cacheKey, result: cachedResult};
+
     return {
-      rowOffset: 0,
-      rows: this.filteredSortedData,
-      totalRows: this.filteredSortedData.length,
+      result: cachedResult,
+      isLoading: false,
     };
   }
 
-  get distinctValues(): ReadonlyMap<string, readonly SqlValue[]> | undefined {
-    return this.distinctValuesCache.size > 0
-      ? this.distinctValuesCache
-      : undefined;
+  getDistinctValues(columnPath: string): UseQueryResult<readonly SqlValue[]> {
+    // Check cache first
+    const cached = this.distinctValuesCache.get(columnPath);
+    if (cached) {
+      return {result: cached, isLoading: false};
+    }
+
+    // Compute distinct values from base data (not filtered)
+    const uniqueValues = new Set<SqlValue>();
+    for (const row of this.data) {
+      uniqueValues.add(row[columnPath]);
+    }
+
+    // Sort with null-aware comparison
+    const sorted = Array.from(uniqueValues).sort((a, b) => {
+      // Nulls come first
+      if (a === null && b === null) return 0;
+      if (a === null) return -1;
+      if (b === null) return 1;
+
+      // Type-specific sorting
+      if (typeof a === 'number' && typeof b === 'number') {
+        return a - b;
+      }
+      if (typeof a === 'bigint' && typeof b === 'bigint') {
+        return Number(a - b);
+      }
+      if (typeof a === 'string' && typeof b === 'string') {
+        return a.localeCompare(b);
+      }
+
+      // Default: convert to string and compare
+      return String(a).localeCompare(String(b));
+    });
+
+    this.distinctValuesCache.set(columnPath, sorted);
+    return {result: sorted, isLoading: false};
   }
 
-  get parameterKeys(): ReadonlyMap<string, readonly string[]> | undefined {
-    return this.parameterKeysCache.size > 0
-      ? this.parameterKeysCache
-      : undefined;
+  getAggregateTotals(
+    model: DataSourceModel,
+  ): UseQueryResult<ReadonlyMap<string, SqlValue>> {
+    const {columns, filters = [], pivot} = model;
+
+    // Check cache first
+    const cacheKey = this.buildCacheKey({columns, filters, pivot});
+    if (this.aggregateTotalsCache?.key === cacheKey) {
+      return {result: this.aggregateTotalsCache.result, isLoading: false};
+    }
+
+    // First get the filtered/pivoted data
+    const aggregates = pivot?.aggregates ?? [];
+    const aggregateFields = new Set(
+      aggregates.map((a) => ('field' in a ? a.field : '__count__')),
+    );
+    const prePivotFilters =
+      pivot && !pivot.drillDown
+        ? filters.filter((f) => !aggregateFields.has(f.field))
+        : filters;
+    const postPivotFilters =
+      pivot && !pivot.drillDown
+        ? filters.filter((f) => aggregateFields.has(f.field))
+        : [];
+
+    let result = this.applyFilters(this.data, prePivotFilters);
+
+    const totals = new Map<string, SqlValue>();
+
+    if (pivot && !pivot.drillDown) {
+      result = this.applyPivoting(result, pivot);
+      if (postPivotFilters.length > 0) {
+        result = this.applyFilters(result, postPivotFilters);
+      }
+      // Compute aggregate totals across all filtered pivot rows
+      this.computeAggregateTotalsFromData(result, pivot, totals);
+    } else if (columns && !pivot) {
+      // Non-pivot mode: compute column-level aggregations
+      this.computeColumnAggregatesFromData(result, columns, totals);
+    }
+
+    this.aggregateTotalsCache = {key: cacheKey, result: totals};
+    return {result: totals, isLoading: false};
   }
 
-  get aggregateTotals(): ReadonlyMap<string, SqlValue> | undefined {
-    return this.aggregateTotalsCache.size > 0
-      ? this.aggregateTotalsCache
-      : undefined;
-  }
+  getPivotRollups(model: DataSourceModel): UseQueryResult<PivotRollups> {
+    const {filters = [], pivot} = model;
 
-  notify({
-    columns,
-    filters = [],
-    pivot,
-    distinctValuesColumns,
-    parameterKeyColumns,
-  }: DataSourceModel): void {
-    if (
-      !this.areColumnsEqual(columns, this.oldColumns) ||
-      !this.areFiltersEqual(filters, this.oldFilters) ||
-      !arePivotsEqual(pivot, this.oldPivot)
-    ) {
-      this.oldColumns = columns;
-      this.oldFilters = filters;
-      this.oldPivot = pivot;
+    // Only generate rollups for hierarchical pivot mode (2+ groupBy columns)
+    if (!pivot || pivot.drillDown || pivot.groupBy.length < 2) {
+      return {result: {byLevel: new Map()}, isLoading: false};
+    }
 
-      // Clear aggregate totals cache
-      this.aggregateTotalsCache.clear();
+    // Check cache
+    const cacheKey = this.buildCacheKey({filters, pivot, _type: 'rollups'});
+    if (this.pivotRollupsCache?.key === cacheKey) {
+      return {result: this.pivotRollupsCache.result, isLoading: false};
+    }
 
-      // In pivot mode, separate filters into pre-pivot and post-pivot
-      // Post-pivot filters apply to aggregate columns
-      const aggregates = pivot?.aggregates ?? [];
-      const aggregateFields = new Set(
-        aggregates.map((a) => ('field' in a ? a.field : '__count__')),
+    // Apply pre-pivot filters
+    const aggregates = pivot.aggregates ?? [];
+    const aggregateFields = new Set(
+      aggregates.map((a) => ('field' in a ? a.field : '__count__')),
+    );
+    const prePivotFilters = filters.filter(
+      (f) => !aggregateFields.has(f.field),
+    );
+    const filteredData = this.applyFilters(this.data, prePivotFilters);
+
+    // Build rollups for levels 0 to N-2
+    const byLevel = new Map<number, Row[]>();
+    const numLevels = pivot.groupBy.length;
+
+    for (let level = 0; level < numLevels - 1; level++) {
+      // Group by the first (level+1) columns
+      const groupByFieldsForLevel = pivot.groupBy
+        .slice(0, level + 1)
+        .map(({field}) => field);
+      const rollupRows = this.computeRollupForLevel(
+        filteredData,
+        groupByFieldsForLevel,
+        aggregates,
       );
-      const prePivotFilters =
-        pivot && !pivot.drillDown
-          ? filters.filter((f) => !aggregateFields.has(f.field))
-          : filters;
-      const postPivotFilters =
-        pivot && !pivot.drillDown
-          ? filters.filter((f) => aggregateFields.has(f.field))
-          : [];
-
-      // Apply pre-pivot filters (on source data)
-      let result = this.applyFilters(this.data, prePivotFilters);
-
-      // Apply pivot (but not in drilldown mode - drilldown shows raw data)
-      if (pivot && !pivot.drillDown) {
-        result = this.applyPivoting(result, pivot);
-        // Apply post-pivot filters (on aggregate results)
-        if (postPivotFilters.length > 0) {
-          result = this.applyFilters(result, postPivotFilters);
-        }
-        // Compute aggregate totals across all filtered pivot rows
-        this.computeAggregateTotals(result, pivot);
-      } else if (pivot?.drillDown) {
-        // Drilldown mode: filter to show only rows matching the drillDown values
-        result = this.applyDrillDown(result, pivot);
-      } else if (columns) {
-        // Non-pivot mode: compute column-level aggregations
-        this.computeColumnAggregates(result, columns);
-      }
-
-      // Apply sorting - find sorted column from columns or pivot
-      const sortedColumn = this.findSortedColumn(columns, pivot);
-      if (sortedColumn) {
-        result = this.applySorting(
-          result,
-          sortedColumn.field,
-          sortedColumn.direction,
-        );
-      }
-
-      // Store the filtered and sorted data
-      this.filteredSortedData = result;
+      byLevel.set(level, rollupRows);
     }
 
-    // Handle distinct values requests
-    if (distinctValuesColumns) {
-      for (const column of distinctValuesColumns) {
-        if (!this.distinctValuesCache.has(column)) {
-          // Compute distinct values from base data (not filtered)
-          const uniqueValues = new Set<SqlValue>();
-          for (const row of this.data) {
-            uniqueValues.add(row[column]);
+    const result: PivotRollups = {byLevel};
+    this.pivotRollupsCache = {key: cacheKey, result};
+    return {result, isLoading: false};
+  }
+
+  /**
+   * Compute rollup rows for a specific level by grouping on the given fields.
+   */
+  private computeRollupForLevel(
+    data: ReadonlyArray<Row>,
+    groupByFields: readonly string[],
+    aggregates: readonly AggregateColumn[],
+  ): Row[] {
+    const groups = new Map<string, Row[]>();
+
+    for (const row of data) {
+      const key = groupByFields.map((field) => row[field]).join('-');
+      if (!groups.has(key)) {
+        groups.set(key, []);
+      }
+      groups.get(key)!.push(row);
+    }
+
+    const result: Row[] = [];
+
+    for (const group of groups.values()) {
+      const newRow: Row = {};
+
+      // Copy the groupBy field values
+      for (const field of groupByFields) {
+        newRow[field] = group[0][field];
+      }
+
+      // Compute aggregates
+      for (const agg of aggregates) {
+        const alias =
+          agg.function === 'COUNT'
+            ? '__count__'
+            : 'field' in agg
+              ? agg.field
+              : '__unknown__';
+
+        if (agg.function === 'COUNT') {
+          newRow[alias] = group.length;
+          continue;
+        }
+
+        const aggField = 'field' in agg ? agg.field : null;
+        if (!aggField) {
+          newRow[alias] = null;
+          continue;
+        }
+
+        const values = group
+          .map((row) => row[aggField])
+          .filter((v) => v !== null);
+        if (values.length === 0) {
+          newRow[alias] = null;
+          continue;
+        }
+
+        switch (agg.function) {
+          case 'SUM':
+            newRow[alias] = values.reduce(
+              (acc: number, val) => acc + (Number(val) || 0),
+              0,
+            );
+            break;
+          case 'AVG':
+            newRow[alias] =
+              (values.reduce(
+                (acc: number, val) => acc + (Number(val) || 0),
+                0,
+              ) as number) / values.length;
+            break;
+          case 'MIN':
+            newRow[alias] = values.reduce(
+              (acc, val) => (val < acc ? val : acc),
+              values[0],
+            );
+            break;
+          case 'MAX':
+            newRow[alias] = values.reduce(
+              (acc, val) => (val > acc ? val : acc),
+              values[0],
+            );
+            break;
+          case 'ANY':
+            newRow[alias] = values[0];
+            break;
+        }
+      }
+
+      result.push(newRow);
+    }
+
+    return result;
+  }
+
+  getParameterKeys(prefix: string): UseQueryResult<readonly string[]> {
+    // Check cache first
+    const cached = this.parameterKeysCache.get(prefix);
+    if (cached) {
+      return {result: cached, isLoading: false};
+    }
+
+    // Find all keys that match the prefix pattern
+    const uniqueKeys = new Set<string>();
+    const prefixWithDot = prefix + '.';
+
+    for (const row of this.data) {
+      for (const key of Object.keys(row)) {
+        if (key.startsWith(prefixWithDot)) {
+          const paramKey = key.slice(prefixWithDot.length);
+          if (!paramKey.includes('.')) {
+            uniqueKeys.add(paramKey);
           }
-
-          // Sort with null-aware comparison
-          const sorted = Array.from(uniqueValues).sort((a, b) => {
-            // Nulls come first
-            if (a === null && b === null) return 0;
-            if (a === null) return -1;
-            if (b === null) return 1;
-
-            // Type-specific sorting
-            if (typeof a === 'number' && typeof b === 'number') {
-              return a - b;
-            }
-            if (typeof a === 'bigint' && typeof b === 'bigint') {
-              return Number(a - b);
-            }
-            if (typeof a === 'string' && typeof b === 'string') {
-              return a.localeCompare(b);
-            }
-
-            // Default: convert to string and compare
-            return String(a).localeCompare(String(b));
-          });
-
-          this.distinctValuesCache.set(column, sorted);
         }
       }
     }
 
-    // Handle parameter keys requests
-    if (parameterKeyColumns) {
-      for (const prefix of parameterKeyColumns) {
-        if (!this.parameterKeysCache.has(prefix)) {
-          // Find all keys that match the prefix pattern (e.g., "skills.typescript" for prefix "skills")
-          const uniqueKeys = new Set<string>();
-          const prefixWithDot = prefix + '.';
-
-          for (const row of this.data) {
-            for (const key of Object.keys(row)) {
-              if (key.startsWith(prefixWithDot)) {
-                // Extract the parameter key (everything after the prefix)
-                const paramKey = key.slice(prefixWithDot.length);
-                // Only add top-level keys (no further dots)
-                if (!paramKey.includes('.')) {
-                  uniqueKeys.add(paramKey);
-                }
-              }
-            }
-          }
-
-          // Sort alphabetically
-          const sorted = Array.from(uniqueKeys).sort((a, b) =>
-            a.localeCompare(b),
-          );
-
-          this.parameterKeysCache.set(prefix, sorted);
-        }
-      }
-    }
+    const sorted = Array.from(uniqueKeys).sort((a, b) => a.localeCompare(b));
+    this.parameterKeysCache.set(prefix, sorted);
+    return {result: sorted, isLoading: false};
   }
 
   /**
    * Export all data with current filters/sorting applied.
    */
   async exportData(): Promise<readonly Row[]> {
-    // Return all the filtered and sorted data
-    return this.filteredSortedData;
-  }
-
-  /**
-   * Compare columns for equality (including sort state).
-   */
-  private areColumnsEqual(
-    a: readonly Column[] | undefined,
-    b: readonly Column[] | undefined,
-  ): boolean {
-    if (a === b) return true;
-    if (!a || !b) return false;
-    if (a.length !== b.length) return false;
-
-    return a.every((colA, i) => {
-      const colB = b[i];
-      return (
-        colA.field === colB.field &&
-        colA.sort === colB.sort &&
-        colA.aggregate === colB.aggregate
-      );
-    });
+    // For export, return all data without pagination
+    const result = this.getRows({});
+    return result.result?.rows ?? [];
   }
 
   /**
@@ -262,22 +408,6 @@
     return undefined;
   }
 
-  // Helper function to compare arrays of filter definitions for equality.
-  private areFiltersEqual(
-    filtersA: ReadonlyArray<Filter>,
-    filtersB: ReadonlyArray<Filter>,
-  ): boolean {
-    if (filtersA.length !== filtersB.length) return false;
-
-    // Compare each filter
-    return filtersA.every((filterA, index) => {
-      const filterB = filtersB[index];
-      return (
-        stringifyJsonWithBigints(filterA) === stringifyJsonWithBigints(filterB)
-      );
-    });
-  }
-
   private applyFilters(
     data: ReadonlyArray<Row>,
     filters: ReadonlyArray<Filter>,
@@ -499,9 +629,10 @@
    * For AVG, we compute the average of averages (weighted by count would be better but we don't have that info).
    * For MIN/MAX, we find the min/max across all groups.
    */
-  private computeAggregateTotals(
+  private computeAggregateTotalsFromData(
     pivotedData: ReadonlyArray<Row>,
     pivot: Pivot,
+    totals: Map<string, SqlValue>,
   ): void {
     const aggregates = pivot.aggregates ?? [];
     for (const agg of aggregates) {
@@ -516,7 +647,7 @@
         .filter((v) => v !== null);
 
       if (values.length === 0) {
-        this.aggregateTotalsCache.set(alias, null);
+        totals.set(alias, null);
         continue;
       }
 
@@ -524,14 +655,14 @@
         case 'SUM':
         case 'COUNT':
           // Sum up all the sums/counts
-          this.aggregateTotalsCache.set(
+          totals.set(
             alias,
             values.reduce((acc: number, val) => acc + (Number(val) || 0), 0),
           );
           break;
         case 'AVG':
           // Average of averages (simple, not weighted)
-          this.aggregateTotalsCache.set(
+          totals.set(
             alias,
             (values.reduce(
               (acc: number, val) => acc + (Number(val) || 0),
@@ -540,20 +671,20 @@
           );
           break;
         case 'MIN':
-          this.aggregateTotalsCache.set(
+          totals.set(
             alias,
             values.reduce((acc, val) => (val < acc ? val : acc), values[0]),
           );
           break;
         case 'MAX':
-          this.aggregateTotalsCache.set(
+          totals.set(
             alias,
             values.reduce((acc, val) => (val > acc ? val : acc), values[0]),
           );
           break;
         case 'ANY':
           // For ANY, just take the first value
-          this.aggregateTotalsCache.set(alias, values[0]);
+          totals.set(alias, values[0]);
           break;
       }
     }
@@ -563,9 +694,10 @@
    * Compute aggregates for columns with aggregation functions defined.
    * This is used in non-pivot mode when columns have individual aggregations.
    */
-  private computeColumnAggregates(
+  private computeColumnAggregatesFromData(
     data: ReadonlyArray<Row>,
     columns: ReadonlyArray<Column>,
+    totals: Map<string, SqlValue>,
   ): void {
     for (const col of columns) {
       if (!col.aggregate) continue;
@@ -575,19 +707,19 @@
         .filter((v) => v !== null);
 
       if (values.length === 0) {
-        this.aggregateTotalsCache.set(col.field, null);
+        totals.set(col.field, null);
         continue;
       }
 
       switch (col.aggregate) {
         case 'SUM':
-          this.aggregateTotalsCache.set(
+          totals.set(
             col.field,
             values.reduce((acc: number, val) => acc + (Number(val) || 0), 0),
           );
           break;
         case 'AVG':
-          this.aggregateTotalsCache.set(
+          totals.set(
             col.field,
             (values.reduce(
               (acc: number, val) => acc + (Number(val) || 0),
@@ -596,19 +728,19 @@
           );
           break;
         case 'MIN':
-          this.aggregateTotalsCache.set(
+          totals.set(
             col.field,
             values.reduce((acc, val) => (val < acc ? val : acc), values[0]),
           );
           break;
         case 'MAX':
-          this.aggregateTotalsCache.set(
+          totals.set(
             col.field,
             values.reduce((acc, val) => (val > acc ? val : acc), values[0]),
           );
           break;
         case 'ANY':
-          this.aggregateTotalsCache.set(col.field, values[0]);
+          totals.set(col.field, values[0]);
           break;
       }
     }
@@ -668,19 +800,3 @@
     return Number(a) - Number(b);
   }
 }
-
-function arePivotsEqual(a?: Pivot, b?: Pivot): boolean {
-  if (a === b) return true;
-  if (a === undefined || b === undefined) return false;
-  // Compare groupBy fields (including sort)
-  if (JSON.stringify(a.groupBy) !== JSON.stringify(b.groupBy)) return false;
-  // Compare aggregates
-  if (JSON.stringify(a.aggregates) !== JSON.stringify(b.aggregates)) {
-    return false;
-  }
-  // Check drillDown equality
-  if (a.drillDown === b.drillDown) return true;
-  if (a.drillDown === undefined || b.drillDown === undefined) return false;
-  if (JSON.stringify(a.drillDown) !== JSON.stringify(b.drillDown)) return false;
-  return true;
-}
diff --git a/ui/src/components/widgets/datagrid/in_memory_data_source_unittest.ts b/ui/src/components/widgets/datagrid/in_memory_data_source_unittest.ts
index 3da25b6..5c59988 100644
--- a/ui/src/components/widgets/datagrid/in_memory_data_source_unittest.ts
+++ b/ui/src/components/widgets/datagrid/in_memory_data_source_unittest.ts
@@ -76,8 +76,10 @@
   });
 
   test('initialization', () => {
-    const result = dataSource.rows;
-    expect(result.rowOffset).toBe(0);
+    const queryResult = dataSource.getRows({});
+    expect(queryResult.isLoading).toBe(false);
+    const result = queryResult.result!;
+    expect(result.offset).toBe(0);
     expect(result.totalRows).toBe(sampleData.length);
     expect(result.rows).toEqual(sampleData);
   });
@@ -85,24 +87,21 @@
   describe('filtering', () => {
     test('equality filter', () => {
       const filters: Filter[] = [{field: 'name', op: '=', value: 'Alice'}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(1);
       expect(result.rows[0].name).toBe('Alice');
     });
 
     test('inequality filter', () => {
       const filters: Filter[] = [{field: 'active', op: '!=', value: 1}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(3); // Bob, David, Mallory
       result.rows.forEach((row) => expect(row.active).toBe(0));
     });
 
     test('less than filter', () => {
       const filters: Filter[] = [{field: 'value', op: '<', value: 150}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       // David (null), Alice (100), Eve (100)
       expect(result.totalRows).toBe(3);
       expect(result.rows.map((r) => r.id).sort()).toEqual([1, 4, 5]);
@@ -110,8 +109,7 @@
 
     test('less than or equal filter', () => {
       const filters: Filter[] = [{field: 'value', op: '<=', value: 150}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       // David (null), Alice (100), Charlie (150), Eve (100)
       expect(result.totalRows).toBe(4);
       expect(result.rows.map((r) => r.id).sort()).toEqual([1, 3, 4, 5]);
@@ -119,48 +117,42 @@
 
     test('greater than filter', () => {
       const filters: Filter[] = [{field: 'value', op: '>', value: 200}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(2); // Mallory (300n), Trent (250n)
       expect(result.rows.map((r) => r.id).sort()).toEqual([6, 7]);
     });
 
     test('greater than or equal filter with bigint', () => {
       const filters: Filter[] = [{field: 'value', op: '>=', value: 250n}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(2); // Mallory, Trent
       expect(result.rows.map((r) => r.id).sort()).toEqual([6, 7]);
     });
 
     test('is null filter', () => {
       const filters: Filter[] = [{field: 'value', op: 'is null'}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(1);
       expect(result.rows[0].id).toBe(4); // David
     });
 
     test('is not null filter', () => {
       const filters: Filter[] = [{field: 'blob', op: 'is not null'}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(6); // All except Charlie
       expect(result.rows.find((r) => r.id === 3)).toBeUndefined();
     });
 
     test('glob filter', () => {
       const filters: Filter[] = [{field: 'name', op: 'glob', value: 'A*e'}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(1);
       expect(result.rows[0].name).toBe('Alice');
     });
 
     test('glob filter with ?', () => {
       const filters: Filter[] = [{field: 'name', op: 'glob', value: 'B?b'}];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(1);
       expect(result.rows[0].name).toBe('Bob');
     });
@@ -170,8 +162,7 @@
         {field: 'active', op: '=', value: 1},
         {field: 'tag', op: '=', value: 'A'},
       ];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(3); // Alice, Charlie, Trent
       result.rows.forEach((row) => {
         expect(row.active).toBe(1);
@@ -183,18 +174,43 @@
       const filters: Filter[] = [
         {field: 'name', op: '=', value: 'NonExistent'},
       ];
-      dataSource.notify({filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({filters}).result!;
       expect(result.totalRows).toBe(0);
       expect(result.rows.length).toBe(0);
     });
+
+    test('in filter', () => {
+      const filters: Filter[] = [{field: 'tag', op: 'in', value: ['A', 'B']}];
+      const result = dataSource.getRows({filters}).result!;
+      expect(result.totalRows).toBe(5); // Alice, Bob, Charlie, Eve, Trent
+      result.rows.forEach((row) => {
+        expect(['A', 'B']).toContain(row.tag);
+      });
+    });
+
+    test('not in filter', () => {
+      const filters: Filter[] = [
+        {field: 'tag', op: 'not in', value: ['A', 'B']},
+      ];
+      const result = dataSource.getRows({filters}).result!;
+      expect(result.totalRows).toBe(2); // David, Mallory
+      result.rows.forEach((row) => {
+        expect(row.tag).toBe('C');
+      });
+    });
+
+    test('not glob filter', () => {
+      const filters: Filter[] = [{field: 'name', op: 'not glob', value: 'A*'}];
+      const result = dataSource.getRows({filters}).result!;
+      expect(result.totalRows).toBe(6); // All except Alice
+      expect(result.rows.find((r) => r.name === 'Alice')).toBeUndefined();
+    });
   });
 
   describe('sorting', () => {
     test('sort by string ascending', () => {
       const columns: Column[] = [{field: 'name', sort: 'ASC'}];
-      dataSource.notify({columns, filters: []});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({columns}).result!;
       expect(result.rows.map((r) => r.name)).toEqual([
         'Alice',
         'Bob',
@@ -208,8 +224,7 @@
 
     test('sort by string descending', () => {
       const columns: Column[] = [{field: 'name', sort: 'DESC'}];
-      dataSource.notify({columns, filters: []});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({columns}).result!;
       expect(result.rows.map((r) => r.name)).toEqual([
         'Trent',
         'Mallory',
@@ -223,31 +238,27 @@
 
     test('sort by number ascending (includes nulls)', () => {
       const columns: Column[] = [{field: 'value', sort: 'ASC'}];
-      dataSource.notify({columns, filters: []});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({columns}).result!;
       // Nulls first, then 100, 100, 150, 200, 250n, 300n
       expect(result.rows.map((r) => r.id)).toEqual([4, 1, 5, 3, 2, 7, 6]);
     });
 
     test('sort by number descending (includes nulls and bigint)', () => {
       const columns: Column[] = [{field: 'value', sort: 'DESC'}];
-      dataSource.notify({columns, filters: []});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({columns}).result!;
       // 300n, 250n, 200, 150, 100, 100, Nulls last
       expect(result.rows.map((r) => r.id)).toEqual([6, 7, 2, 3, 1, 5, 4]);
     });
 
     test('sort by boolean ascending', () => {
       const columns: Column[] = [{field: 'active', sort: 'ASC'}]; // 0 then 1
-      dataSource.notify({columns, filters: []});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({columns}).result!;
       expect(result.rows.map((r) => r.active)).toEqual([0, 0, 0, 1, 1, 1, 1]);
     });
 
     test('sort by Uint8Array ascending (by length)', () => {
       const columns: Column[] = [{field: 'blob', sort: 'ASC'}];
-      dataSource.notify({columns, filters: []});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({columns}).result!;
       // null (Charlie, id:3), len 1 (David id:4, Mallory id:6), len 2 (Alice id:1, Trent id:7), len 3 (Bob id:2), len 4 (Eve id:5)
       // Original order for same length: David before Mallory, Alice before Trent.
       expect(result.rows.map((r) => r.id)).toEqual([3, 4, 6, 1, 7, 2, 5]);
@@ -255,21 +266,14 @@
 
     test('sort by Uint8Array descending (by length)', () => {
       const columns: Column[] = [{field: 'blob', sort: 'DESC'}];
-      dataSource.notify({columns, filters: []});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({columns}).result!;
       // len 4, len 3, len 2, len 2, len 1, len 0, null
       expect(result.rows.map((r) => r.id)).toEqual([5, 2, 1, 7, 4, 6, 3]);
     });
 
-    test('unsorted', () => {
-      // Apply some sort first
-      dataSource.notify({
-        columns: [{field: 'name', sort: 'ASC'}],
-      });
-      // Then unsort
-      dataSource.notify({});
-      const result = dataSource.rows;
-      // Should revert to original order if no filters applied
+    test('unsorted returns original order', () => {
+      const result = dataSource.getRows({}).result!;
+      // Should be in original order
       expect(result.rows.map((r) => r.id)).toEqual(sampleData.map((r) => r.id));
     });
   });
@@ -278,8 +282,7 @@
     test('filter then sort', () => {
       const filters: Filter[] = [{field: 'active', op: '=', value: 1}];
       const columns: Column[] = [{field: 'value', sort: 'DESC'}];
-      dataSource.notify({columns, filters});
-      const result = dataSource.rows;
+      const result = dataSource.getRows({columns, filters}).result!;
       // Active: Alice (100), Charlie (150), Eve (100), Trent (250n)
       // Sorted by value desc: Trent, Charlie, Alice, Eve (Alice/Eve order by original due to stable sort on value)
       expect(result.rows.map((r) => r.id)).toEqual([7, 3, 1, 5]);
@@ -287,87 +290,89 @@
     });
   });
 
-  describe('caching behavior', () => {
-    test('data is not reprocessed if columns and filters are identical', () => {
-      const filters: Filter[] = [{field: 'tag', op: '=', value: 'A'}];
-      const columns: Column[] = [{field: 'name', sort: 'ASC'}];
-
-      dataSource.notify({columns, filters});
-      const result1 = dataSource.rows.rows; // Access internal array
-
-      // Spy on internal methods if possible, or check object identity
-      // For this test, we'll check if the returned array reference is the same
-      dataSource.notify({columns, filters}); // Identical call
-      const result2 = dataSource.rows.rows;
-
-      expect(result1).toBe(result2); // Should be the same array instance due to caching
+  describe('pagination', () => {
+    test('offset and limit', () => {
+      const columns: Column[] = [{field: 'id', sort: 'ASC'}];
+      const result = dataSource.getRows({
+        columns,
+        pagination: {offset: 2, limit: 3},
+      }).result!;
+      expect(result.totalRows).toBe(7); // Total is still 7
+      expect(result.offset).toBe(2);
+      expect(result.rows.length).toBe(3);
+      expect(result.rows.map((r) => r.id)).toEqual([3, 4, 5]);
     });
 
-    test('data is reprocessed if sorting changes', () => {
-      const filters: Filter[] = [{field: 'tag', op: '=', value: 'A'}];
-      const columns1: Column[] = [{field: 'name', sort: 'ASC'}];
-      const columns2: Column[] = [{field: 'name', sort: 'DESC'}];
+    test('offset beyond data', () => {
+      const result = dataSource.getRows({
+        pagination: {offset: 100, limit: 10},
+      }).result!;
+      expect(result.totalRows).toBe(7);
+      expect(result.rows.length).toBe(0);
+    });
+  });
 
-      dataSource.notify({columns: columns1, filters});
-      const result1 = dataSource.rows.rows;
-
-      dataSource.notify({columns: columns2, filters}); // Different sort
-      const result2 = dataSource.rows.rows;
-
-      expect(result1).not.toBe(result2);
-      expect(result1.map((r) => r.id)).not.toEqual(result2.map((r) => r.id));
+  describe('pivot mode', () => {
+    test('basic pivot with count', () => {
+      const result = dataSource.getRows({
+        pivot: {
+          groupBy: [{field: 'tag'}],
+          aggregates: [{function: 'COUNT'}],
+        },
+      }).result!;
+      // 3 unique tags: A, B, C
+      expect(result.totalRows).toBe(3);
+      const tagCounts = new Map(result.rows.map((r) => [r.tag, r.__count__]));
+      expect(tagCounts.get('A')).toBe(3); // Alice, Charlie, Trent
+      expect(tagCounts.get('B')).toBe(2); // Bob, Eve
+      expect(tagCounts.get('C')).toBe(2); // David, Mallory
     });
 
-    test('data is reprocessed if filters change', () => {
-      const filters1: Filter[] = [{field: 'tag', op: '=', value: 'A'}];
-      const filters2: Filter[] = [{field: 'tag', op: '=', value: 'B'}];
-      const columns: Column[] = [{field: 'name', sort: 'ASC'}];
+    test('pivot with sum aggregate', () => {
+      const result = dataSource.getRows({
+        pivot: {
+          groupBy: [{field: 'active'}],
+          aggregates: [{function: 'SUM', field: 'value'}],
+        },
+      }).result!;
+      expect(result.totalRows).toBe(2); // active: 0 and 1
+      const activeSums = new Map(result.rows.map((r) => [r.active, r.value]));
+      // active=0: Bob (200) + David (null) + Mallory (300n) = 500
+      // active=1: Alice (100) + Charlie (150) + Eve (100) + Trent (250n) = 600
+      expect(activeSums.get(0)).toBe(500);
+      expect(activeSums.get(1)).toBe(600);
+    });
+  });
 
-      dataSource.notify({columns, filters: filters1});
-      const result1 = dataSource.rows.rows;
-
-      dataSource.notify({columns, filters: filters2}); // Different filters
-      const result2 = dataSource.rows.rows;
-
-      expect(result1).not.toBe(result2);
-      expect(result1.map((r) => r.id)).not.toEqual(result2.map((r) => r.id));
+  describe('distinct values', () => {
+    test('returns distinct values sorted', () => {
+      const result = dataSource.getDistinctValues('tag');
+      expect(result.isLoading).toBe(false);
+      expect(result.result).toEqual(['A', 'B', 'C']);
     });
 
-    test('data is reprocessed if filter value changes (Uint8Array)', () => {
-      const filters1: Filter[] = [
-        {field: 'blob', op: '=', value: new Uint8Array([1, 2])},
-      ];
-      const filters2: Filter[] = [
-        {field: 'blob', op: '=', value: new Uint8Array([3, 4, 5])},
-      ];
-
-      dataSource.notify({filters: filters1});
-      const result1 = dataSource.rows.rows;
-      expect(result1.length).toBe(1);
-      expect(result1[0].id).toBe(1);
-
-      dataSource.notify({filters: filters2});
-      const result2 = dataSource.rows.rows;
-      expect(result2.length).toBe(1);
-      expect(result2[0].id).toBe(2);
-
-      expect(result1).not.toBe(result2);
+    test('handles null values', () => {
+      const result = dataSource.getDistinctValues('blob');
+      expect(result.isLoading).toBe(false);
+      // null should come first
+      expect(result.result![0]).toBe(null);
     });
   });
 
   test('empty data source', () => {
     const emptyDataSource = new InMemoryDataSource([]);
-    const result = emptyDataSource.rows;
-    expect(result.rowOffset).toBe(0);
+    const queryResult = emptyDataSource.getRows({});
+    expect(queryResult.isLoading).toBe(false);
+    const result = queryResult.result!;
+    expect(result.offset).toBe(0);
     expect(result.totalRows).toBe(0);
     expect(result.rows).toEqual([]);
 
-    emptyDataSource.notify({
+    const filteredResult = emptyDataSource.getRows({
       columns: [{field: 'id', sort: 'DESC'}],
       filters: [{field: 'name', op: '=', value: 'test'}],
-    });
-    const resultAfterUpdate = emptyDataSource.rows;
-    expect(resultAfterUpdate.totalRows).toBe(0);
-    expect(resultAfterUpdate.rows).toEqual([]);
+    }).result!;
+    expect(filteredResult.totalRows).toBe(0);
+    expect(filteredResult.rows).toEqual([]);
   });
 });
diff --git a/ui/src/components/widgets/datagrid/model.ts b/ui/src/components/widgets/datagrid/model.ts
index 2c0d2e4..257c9ab 100644
--- a/ui/src/components/widgets/datagrid/model.ts
+++ b/ui/src/components/widgets/datagrid/model.ts
@@ -64,6 +64,12 @@
   // This allows drilling down into a specific pivot group to see the
   // underlying data. The keys are the groupBy column names.
   readonly drillDown?: Row;
+
+  // Set of collapsed group keys. Groups not in this set are expanded.
+  // Keys are formed by joining parent group values with '|' separator.
+  // E.g., for groupBy [process, thread], a collapsed Chrome group would be "Chrome".
+  // Only applicable when there are 2+ groupBy columns (hierarchical mode).
+  readonly collapsed?: ReadonlySet<string>;
 }
 
 export interface Model {
diff --git a/ui/src/components/widgets/datagrid/sql_data_source.ts b/ui/src/components/widgets/datagrid/sql_data_source.ts
index 6c92bba..135a5bd 100644
--- a/ui/src/components/widgets/datagrid/sql_data_source.ts
+++ b/ui/src/components/widgets/datagrid/sql_data_source.ts
@@ -12,18 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import {AsyncLimiter} from '../../../base/async_limiter';
 import {assertUnreachable} from '../../../base/logging';
 import {maybeUndefined} from '../../../base/utils';
 import {Engine} from '../../../trace_processor/engine';
-import {NUM, Row, SqlValue} from '../../../trace_processor/query_result';
-import {runQueryForQueryTable} from '../../query_table/queries';
 import {
-  DataSource,
-  DataSourceModel,
-  DataSourceRows,
-  Pagination,
-} from './data_source';
+  createQueryCache,
+  UseQueryResult,
+} from '../../../trace_processor/query_cache';
+import {
+  NUM,
+  QueryResult,
+  Row,
+  SqlValue,
+} from '../../../trace_processor/query_result';
+import {runQueryForQueryTable} from '../../query_table/queries';
+import {DataSource, DataSourceModel, PivotRollups} from './data_source';
 import {Column, Filter, Pivot} from './model';
 import {
   isSQLExpressionDef,
@@ -61,32 +64,6 @@
   readonly preamble?: string;
 }
 
-// Cache entry for row count resolution
-interface RowCountCache {
-  query: string;
-  count: number;
-}
-
-// Cache entry for rows resolution
-interface RowsCache {
-  query: string;
-  pagination: Pagination | undefined;
-  offset: number;
-  rows: Row[];
-}
-
-// Cache entry for aggregate resolution
-interface AggregatesCache {
-  query: string;
-  totals: Map<string, SqlValue>;
-}
-
-// Cache entry for distinct values resolution
-interface DistinctValuesCache {
-  query: string;
-  values: ReadonlyArray<SqlValue>;
-}
-
 /**
  * SQL data source for DataGrid.
  *
@@ -127,343 +104,207 @@
  */
 export class SQLDataSource implements DataSource {
   private readonly engine: Engine;
-  private readonly limiter = new AsyncLimiter();
   private readonly sqlSchema: SQLSchemaRegistry;
   private readonly rootSchemaName: string;
   private readonly prelude?: string;
+  private readonly queryCache: ReturnType<typeof createQueryCache>;
 
-  // Cache for each resolution type
-  private rowCountCache?: RowCountCache;
-  private rowsCache?: RowsCache;
-  private aggregatesCache?: AggregatesCache;
-  private distinctValuesCache = new Map<string, DistinctValuesCache>();
-  private parameterKeysCache = new Map<string, ReadonlyArray<string>>();
-
-  // Current results
-  private cachedResult?: DataSourceRows;
-  private cachedDistinctValues?: ReadonlyMap<string, ReadonlyArray<SqlValue>>;
-  private cachedAggregateTotals?: ReadonlyMap<string, SqlValue>;
-  private isLoadingFlag = false;
+  // Track the last built query for exportData
+  private lastRowsQuery?: string;
 
   constructor(config: SQLDataSourceConfig) {
     this.engine = config.engine;
     this.sqlSchema = config.sqlSchema;
     this.rootSchemaName = config.rootSchemaName;
     this.prelude = config.preamble;
+    this.queryCache = createQueryCache(this.engine);
   }
 
-  /**
-   * Getter for the current rows result
-   */
-  get rows(): DataSourceRows | undefined {
-    return this.cachedResult;
-  }
-
-  get isLoading(): boolean {
-    return this.isLoadingFlag;
-  }
-
-  get distinctValues(): ReadonlyMap<string, readonly SqlValue[]> | undefined {
-    return this.cachedDistinctValues;
-  }
-
-  get parameterKeys(): ReadonlyMap<string, readonly string[]> | undefined {
-    return this.parameterKeysCache.size > 0
-      ? this.parameterKeysCache
-      : undefined;
-  }
-
-  get aggregateTotals(): ReadonlyMap<string, SqlValue> | undefined {
-    return this.cachedAggregateTotals;
-  }
-
-  /**
-   * Get the current working query for the datasource.
-   * Useful for debugging or creating debug tracks.
-   */
-  getCurrentQuery(): string {
-    return this.rowsCache?.query ?? '';
-  }
-
-  /**
-   * Notify of parameter changes and trigger data update
-   */
-  notify(model: DataSourceModel): void {
-    this.limiter.schedule(async () => {
-      // Defer setting loading flag to avoid setting it synchronously during the
-      // view() call that triggered notify(). This avoids the bug that the
-      // current frame always has isLoading = true.
-      await Promise.resolve();
-      this.isLoadingFlag = true;
-
-      try {
-        // Resolve row count
-        const rowCount = await this.resolveRowCount(model);
-
-        // Resolve aggregates
-        const aggregateTotals = await this.resolveAggregates(model);
-
-        // Resolve rows
-        const {offset, rows} = await this.resolveRows(model);
-
-        // Resolve distinct values
-        const distinctValues = await this.resolveDistinctValues(model);
-
-        // Resolve parameter keys
-        await this.resolveParameterKeys(model);
-
-        // Build final result
-        this.cachedResult = {
-          rowOffset: offset,
-          totalRows: rowCount,
-          rows,
-        };
-        this.cachedDistinctValues = distinctValues;
-        this.cachedAggregateTotals = aggregateTotals;
-      } finally {
-        this.isLoadingFlag = false;
-      }
-    });
-  }
-
-  /**
-   * Resolves the row count. Compares query against cache and reuses if unchanged.
-   */
-  private async resolveRowCount(model: DataSourceModel): Promise<number> {
-    // Build query without ORDER BY - ordering is irrelevant for counting
-    const countQuery = this.buildQuery(model, {includeOrderBy: false});
-
-    // Check cache
-    if (this.rowCountCache?.query === countQuery) {
-      return this.rowCountCache.count;
-    }
-
-    // Fetch new count
-    const result = await this.engine.query(
-      this.wrapQueryWithPrelude(`
-      WITH data AS (${countQuery})
-      SELECT COUNT(*) AS total_count
-      FROM data
-    `),
-    );
-    const count = result.firstRow({total_count: NUM}).total_count;
-
-    // Update cache
-    this.rowCountCache = {query: countQuery, count};
-
-    return count;
-  }
-
-  /**
-   * Resolves the rows for the current page. Compares query and pagination against cache.
-   */
-  private async resolveRows(
+  getRows(
     model: DataSourceModel,
-  ): Promise<{offset: number; rows: Row[]}> {
+  ): UseQueryResult<{totalRows: number; offset: number; rows: Row[]}> {
     const {pagination} = model;
 
-    // Build query with ORDER BY for proper pagination ordering
-    const rowsQuery = this.buildQuery(model, {includeOrderBy: true});
+    // Build base query (without ORDER BY for counting)
+    const baseQuery = this.buildQuery(model, {includeOrderBy: false});
 
-    // Check cache - both query and pagination must match
-    if (
-      this.rowsCache?.query === rowsQuery &&
-      comparePagination(this.rowsCache.pagination, pagination)
-    ) {
-      return {offset: this.rowsCache.offset, rows: this.rowsCache.rows};
+    // Count query
+    const countQuery = this.wrapQueryWithPrelude(`
+      WITH data AS (${baseQuery})
+      SELECT COUNT(*) AS total_count
+      FROM data
+    `);
+    const countResult = this.queryCache.useQuery(countQuery);
+    if (countResult.isLoading) {
+      return {isLoading: true};
     }
 
-    // Fetch new rows
+    const rowCount = countResult.result.firstRow({
+      total_count: NUM,
+    }).total_count;
+
+    // Rows query (with ORDER BY for proper pagination)
+    const rowsQuery = this.buildQuery(model, {includeOrderBy: true});
+    this.lastRowsQuery = rowsQuery; // Track for exportData
     let query = `
       WITH data AS (${rowsQuery})
       SELECT *
       FROM data
     `;
-
     if (pagination) {
       query += `LIMIT ${pagination.limit} OFFSET ${pagination.offset}`;
     }
 
-    const result = await runQueryForQueryTable(
+    const rowsResult = this.queryCache.useQuery(
       this.wrapQueryWithPrelude(query),
-      this.engine,
     );
+    if (rowsResult.isLoading) {
+      return {
+        result: {
+          totalRows: rowCount,
+          offset: 0,
+          rows: [],
+        },
+        isLoading: true,
+      };
+    }
 
     const offset = pagination?.offset ?? 0;
-    const rows = result.rows;
+    const rows = queryResultToRows(rowsResult.result);
 
-    // Update cache
-    this.rowsCache = {query: rowsQuery, pagination, offset, rows};
-
-    return {offset, rows};
+    return {
+      result: {totalRows: rowCount, offset, rows},
+      isLoading: countResult.isLoading || rowsResult.isLoading,
+    };
   }
 
-  /**
-   * Resolves aggregate totals. Handles both pivot aggregates and column aggregates.
-   */
-  private async resolveAggregates(
+  getDistinctValues(columnPath: string): UseQueryResult<readonly SqlValue[]> {
+    const query = this.buildDistinctValuesQuery(columnPath);
+    if (!query) {
+      return {result: [], isLoading: false};
+    }
+
+    const {result, isLoading} = this.queryCache.useQuery(
+      this.wrapQueryWithPrelude(query),
+    );
+
+    if (isLoading) {
+      return {isLoading: true};
+    }
+
+    const values = queryResultToRows(result).map((r) => r['value']);
+    return {result: values, isLoading: false};
+  }
+
+  getAggregateTotals(
     model: DataSourceModel,
-  ): Promise<Map<string, SqlValue> | undefined> {
+  ): UseQueryResult<ReadonlyMap<string, SqlValue>> {
     const {columns, filters = [], pivot} = model;
 
-    // Build a unique query string for the aggregates
-    const aggregateQuery = this.buildAggregateQuery(model);
-
-    // If no aggregates needed, return undefined
-    if (!aggregateQuery) {
-      this.aggregatesCache = undefined;
-      return undefined;
-    }
-
-    // Check cache
-    if (this.aggregatesCache?.query === aggregateQuery) {
-      return this.aggregatesCache.totals;
-    }
-
-    // Compute aggregates
-    const totals = new Map<string, SqlValue>();
-
     // Pivot aggregates (but not drill-down mode)
     const pivotAggregates = pivot?.aggregates ?? [];
     if (pivot && !pivot.drillDown && pivotAggregates.length > 0) {
-      const aggregates = await this.fetchPivotAggregates(filters, pivot);
-      for (const [key, value] of Object.entries(aggregates)) {
-        totals.set(key, value);
+      const query = this.buildPivotAggregateQuery(filters, pivot);
+      if (!query) {
+        return {result: new Map(), isLoading: false};
       }
+
+      const {result, isLoading} = this.queryCache.useQuery(
+        this.wrapQueryWithPrelude(query),
+      );
+
+      if (isLoading) {
+        return {isLoading: true};
+      }
+
+      const row = queryResultToRows(result)[0] ?? {};
+      const totals = new Map<string, SqlValue>(Object.entries(row));
+      return {result: totals, isLoading: false};
     }
 
     // Column-level aggregations (non-pivot mode)
     const columnsWithAggregation = columns?.filter((c) => c.aggregate);
     if (columnsWithAggregation && columnsWithAggregation.length > 0 && !pivot) {
-      const columnAggregates = await this.fetchColumnAggregates(
+      const query = this.buildColumnAggregateQuery(
         filters,
         columnsWithAggregation,
       );
-      for (const [key, value] of Object.entries(columnAggregates)) {
-        totals.set(key, value as SqlValue);
-      }
-    }
-
-    // Update cache
-    this.aggregatesCache = {query: aggregateQuery, totals};
-
-    return totals;
-  }
-
-  /**
-   * Builds a unique string representing the aggregate query for cache comparison.
-   */
-  private buildAggregateQuery(model: DataSourceModel): string | undefined {
-    const {columns, filters = [], pivot} = model;
-
-    const parts: string[] = [];
-
-    // Include pivot aggregates
-    if (pivot && !pivot.drillDown && Boolean(pivot.aggregates?.length)) {
-      parts.push(`pivot:${JSON.stringify(pivot.aggregates)}`);
-    }
-
-    // Include column aggregates
-    const columnsWithAggregation = columns?.filter((c) => c.aggregate);
-    if (columnsWithAggregation && columnsWithAggregation.length > 0 && !pivot) {
-      const colAggs = columnsWithAggregation.map(
-        (c) => `${c.field}:${c.aggregate}`,
-      );
-      parts.push(`columns:${colAggs.join(',')}`);
-    }
-
-    if (parts.length === 0) {
-      return undefined;
-    }
-
-    // Include filters in the cache key
-    const filterKey = filters.map((f) => {
-      const value = 'value' in f ? f.value : '';
-      return `${f.field}:${f.op}:${value}`;
-    });
-    parts.push(`filters:${filterKey.join(',')}`);
-
-    return parts.join('|');
-  }
-
-  /**
-   * Resolves distinct values for requested columns.
-   */
-  private async resolveDistinctValues(
-    model: DataSourceModel,
-  ): Promise<Map<string, ReadonlyArray<SqlValue>>> {
-    const {distinctValuesColumns} = model;
-
-    const result = new Map<string, ReadonlyArray<SqlValue>>();
-
-    if (!distinctValuesColumns) {
-      return result;
-    }
-
-    for (const columnPath of distinctValuesColumns) {
-      const query = this.buildDistinctValuesQuery(columnPath);
-      if (!query) continue;
-
-      // Check cache
-      const cached = this.distinctValuesCache.get(columnPath);
-      if (cached?.query === query) {
-        result.set(columnPath, cached.values);
-        continue;
+      if (!query) {
+        return {result: new Map(), isLoading: false};
       }
 
-      // Fetch new values
-      const queryResult = await runQueryForQueryTable(
+      const {result, isLoading} = this.queryCache.useQuery(
         this.wrapQueryWithPrelude(query),
-        this.engine,
       );
-      const values = queryResult.rows.map((r) => r['value']);
 
-      // Update cache
-      this.distinctValuesCache.set(columnPath, {query, values});
-      result.set(columnPath, values);
+      if (isLoading) {
+        return {isLoading: true};
+      }
+
+      const row = queryResultToRows(result)[0] ?? {};
+      const totals = new Map<string, SqlValue>(Object.entries(row));
+      return {result: totals, isLoading: false};
     }
 
-    return result;
+    return {result: new Map(), isLoading: false};
   }
 
-  /**
-   * Resolves parameter keys for parameterized columns.
-   */
-  private async resolveParameterKeys(model: DataSourceModel): Promise<void> {
-    const {parameterKeyColumns} = model;
+  getPivotRollups(model: DataSourceModel): UseQueryResult<PivotRollups> {
+    const {pivot, filters = []} = model;
 
-    if (!parameterKeyColumns) {
-      return;
+    // Only resolve rollups for hierarchical pivot mode (2+ groupBy columns)
+    if (!pivot || pivot.drillDown || pivot.groupBy.length < 2) {
+      return {result: {byLevel: new Map()}, isLoading: false};
     }
 
-    for (const prefix of parameterKeyColumns) {
-      // Already cached
-      if (this.parameterKeysCache.has(prefix)) {
-        continue;
+    const byLevel = new Map<number, Row[]>();
+    let anyLoading = false;
+
+    // For N groupBy columns, we need rollups for levels 0 to N-2
+    for (let level = 0; level < pivot.groupBy.length - 1; level++) {
+      const query = this.buildRollupQuery(filters, pivot, level);
+      const {result, isLoading} = this.queryCache.useQuery(
+        this.wrapQueryWithPrelude(query),
+      );
+
+      if (isLoading) {
+        anyLoading = true;
+      } else {
+        byLevel.set(level, queryResultToRows(result));
       }
+    }
 
-      const schema = this.sqlSchema[this.rootSchemaName];
-      const colDef = maybeUndefined(schema?.columns[prefix]);
+    if (anyLoading) {
+      return {result: {byLevel}, isLoading: true};
+    }
 
-      if (colDef && isSQLExpressionDef(colDef) && colDef.parameterized) {
-        if (colDef.parameterKeysQuery) {
-          const baseTable = schema.table;
-          const baseAlias = `${baseTable}_0`;
-          const query = colDef.parameterKeysQuery(baseTable, baseAlias);
+    return {result: {byLevel}, isLoading: false};
+  }
 
-          try {
-            const result = await runQueryForQueryTable(
-              this.wrapQueryWithPrelude(query),
-              this.engine,
-            );
-            const keys = result.rows.map((r) => String(r['key']));
-            this.parameterKeysCache.set(prefix, keys);
-          } catch {
-            this.parameterKeysCache.set(prefix, []);
-          }
+  getParameterKeys(prefix: string): UseQueryResult<readonly string[]> {
+    const schema = this.sqlSchema[this.rootSchemaName];
+    const colDef = maybeUndefined(schema?.columns[prefix]);
+
+    if (colDef && isSQLExpressionDef(colDef) && colDef.parameterized) {
+      if (colDef.parameterKeysQuery) {
+        const baseTable = schema.table;
+        const baseAlias = `${baseTable}_0`;
+        const query = colDef.parameterKeysQuery(baseTable, baseAlias);
+
+        const {result, isLoading} = this.queryCache.useQuery(
+          this.wrapQueryWithPrelude(query),
+        );
+
+        if (isLoading) {
+          return {isLoading: true};
         }
+
+        const keys = queryResultToRows(result).map((r) => String(r['key']));
+        return {result: keys, isLoading: false};
       }
     }
+
+    return {result: [], isLoading: false};
   }
 
   private wrapQueryWithPrelude(query: string): string {
@@ -477,12 +318,11 @@
    * Export all data with current filters/sorting applied.
    */
   async exportData(): Promise<Row[]> {
-    const workingQuery = this.rowsCache?.query;
-    if (!workingQuery) {
+    if (!this.lastRowsQuery) {
       return [];
     }
 
-    const query = `SELECT * FROM (${workingQuery})`;
+    const query = `SELECT * FROM (${this.lastRowsQuery})`;
     const result = await runQueryForQueryTable(
       this.wrapQueryWithPrelude(query),
       this.engine,
@@ -557,9 +397,10 @@
       query += `\nWHERE ${whereConditions.join(' AND ')}`;
     }
 
-    // Add drill-down conditions
+    // Add drill-down conditions (only for fields present in drillDown)
     if (pivot?.drillDown) {
       const drillDownConditions = pivot.groupBy
+        .filter((col) => col.field in pivot.drillDown!)
         .map((col) => {
           const field = col.field;
           const value = pivot.drillDown![field];
@@ -799,6 +640,220 @@
   }
 
   /**
+   * Builds a pivot aggregate query.
+   */
+  private buildPivotAggregateQuery(
+    filters: ReadonlyArray<Filter>,
+    pivot: Pivot,
+  ): string | undefined {
+    const resolver = new SQLSchemaResolver(this.sqlSchema, this.rootSchemaName);
+
+    const baseTable = resolver.getBaseTable();
+    const baseAlias = resolver.getBaseAlias();
+
+    // Resolve filter columns first to ensure JOINs are added
+    for (const filter of filters) {
+      resolver.resolveColumnPath(filter.field);
+    }
+
+    const aggregates = pivot.aggregates ?? [];
+    const selectClauses = aggregates
+      .map((agg) => {
+        if (agg.function === 'COUNT') {
+          return `COUNT(*) AS __count__`;
+        }
+        const field = 'field' in agg ? agg.field : null;
+        if (!field) return null;
+        const alias = this.pathToAlias(field);
+        if (agg.function === 'ANY') {
+          return `NULL AS ${alias}`;
+        }
+        const colExpr = resolver.resolveColumnPath(field);
+        if (!colExpr) {
+          return `NULL AS ${alias}`;
+        }
+        return `${agg.function}(${colExpr}) AS ${alias}`;
+      })
+      .filter(Boolean)
+      .join(', ');
+
+    if (!selectClauses) {
+      return undefined;
+    }
+
+    const joinClauses = resolver.buildJoinClauses();
+
+    let query = `
+SELECT ${selectClauses}
+FROM ${baseTable} AS ${baseAlias}
+${joinClauses}`;
+
+    if (filters.length > 0) {
+      const filterResolver = new SQLSchemaResolver(
+        this.sqlSchema,
+        this.rootSchemaName,
+      );
+      const whereConditions = filters.map((filter) => {
+        const sqlExpr = filterResolver.resolveColumnPath(filter.field);
+        return this.filterToSql(filter, sqlExpr ?? filter.field);
+      });
+      query += `\nWHERE ${whereConditions.join(' AND ')}`;
+    }
+
+    return query;
+  }
+
+  /**
+   * Builds a column aggregate query.
+   */
+  private buildColumnAggregateQuery(
+    filters: ReadonlyArray<Filter>,
+    columns: ReadonlyArray<Column>,
+  ): string | undefined {
+    const resolver = new SQLSchemaResolver(this.sqlSchema, this.rootSchemaName);
+
+    const baseTable = resolver.getBaseTable();
+    const baseAlias = resolver.getBaseAlias();
+
+    const selectClauses = columns
+      .filter((col) => col.aggregate)
+      .map((col) => {
+        const func = col.aggregate!;
+        const colExpr = resolver.resolveColumnPath(col.field);
+        const alias = this.pathToAlias(col.field);
+
+        if (!colExpr) {
+          return `NULL AS ${alias}`;
+        }
+        if (func === 'ANY') {
+          return `MIN(${colExpr}) AS ${alias}`;
+        }
+        return `${func}(${colExpr}) AS ${alias}`;
+      })
+      .join(', ');
+
+    if (!selectClauses) {
+      return undefined;
+    }
+
+    // Resolve filter column paths first to ensure JOINs are added
+    for (const filter of filters) {
+      resolver.resolveColumnPath(filter.field);
+    }
+
+    const joinClauses = resolver.buildJoinClauses();
+
+    let query = `
+SELECT ${selectClauses}
+FROM ${baseTable} AS ${baseAlias}
+${joinClauses}`;
+
+    if (filters.length > 0) {
+      const whereConditions = filters.map((filter) => {
+        const sqlExpr = resolver.resolveColumnPath(filter.field);
+        return this.filterToSql(filter, sqlExpr ?? filter.field);
+      });
+      query += `\nWHERE ${whereConditions.join(' AND ')}`;
+    }
+
+    return query;
+  }
+
+  /**
+   * Builds a rollup query for a specific grouping level.
+   */
+  private buildRollupQuery(
+    filters: ReadonlyArray<Filter>,
+    pivot: Pivot,
+    level: number,
+  ): string {
+    const resolver = new SQLSchemaResolver(this.sqlSchema, this.rootSchemaName);
+
+    const baseTable = resolver.getBaseTable();
+    const baseAlias = resolver.getBaseAlias();
+
+    // Get groupBy columns for this level (first N+1 columns for level N)
+    const groupByColumnsForLevel = pivot.groupBy.slice(0, level + 1);
+
+    // Build groupBy expressions
+    const groupByExprs: string[] = [];
+    const groupByFields: string[] = [];
+
+    for (const col of groupByColumnsForLevel) {
+      const field = col.field;
+      groupByFields.push(field);
+      const sqlExpr = resolver.resolveColumnPath(field);
+      if (sqlExpr) {
+        const alias = this.pathToAlias(field);
+        groupByExprs.push(`${sqlExpr} AS ${alias}`);
+      }
+    }
+
+    // Build aggregate expressions
+    const aggregates = pivot.aggregates ?? [];
+    const aggregateExprs = aggregates.map((agg) => {
+      if (agg.function === 'COUNT') {
+        return `COUNT(*) AS __count__`;
+      }
+      const field = 'field' in agg ? agg.field : null;
+      if (!field) {
+        return `NULL AS __unknown__`;
+      }
+      const alias = this.pathToAlias(field);
+      const colExpr = resolver.resolveColumnPath(field);
+      if (!colExpr) {
+        return `NULL AS ${alias}`;
+      }
+      if (agg.function === 'ANY') {
+        return `MIN(${colExpr}) AS ${alias}`;
+      }
+      return `${agg.function}(${colExpr}) AS ${alias}`;
+    });
+
+    // Resolve filter columns to ensure JOINs are added
+    for (const filter of filters) {
+      resolver.resolveColumnPath(filter.field);
+    }
+
+    const selectClauses = [...groupByExprs, ...aggregateExprs];
+    const joinClauses = resolver.buildJoinClauses();
+
+    let query = `
+SELECT ${selectClauses.join(',\n       ')}
+FROM ${baseTable} AS ${baseAlias}
+${joinClauses}`;
+
+    // Add WHERE clause for filters
+    if (filters.length > 0) {
+      const whereConditions = filters.map((filter) => {
+        const sqlExpr = resolver.resolveColumnPath(filter.field);
+        if (!sqlExpr) {
+          return this.filterToSql(filter, filter.field);
+        }
+        return this.filterToSql(filter, sqlExpr);
+      });
+      query += `\nWHERE ${whereConditions.join(' AND ')}`;
+    }
+
+    // Add GROUP BY
+    if (groupByFields.length > 0) {
+      const groupByOrigExprs = groupByFields.map(
+        (field) => resolver.resolveColumnPath(field) ?? field,
+      );
+      query += `\nGROUP BY ${groupByOrigExprs.join(', ')}`;
+    }
+
+    // Add ORDER BY based on the groupBy columns' sort settings
+    const sortedCol = groupByColumnsForLevel.find((col) => col.sort);
+    if (sortedCol) {
+      const alias = this.pathToAlias(sortedCol.field);
+      query += `\nORDER BY ${alias} ${sortedCol.sort}`;
+    }
+
+    return query;
+  }
+
+  /**
    * Converts a column path to a valid SQL alias.
    */
   private pathToAlias(path: string): string {
@@ -836,128 +891,6 @@
         assertUnreachable(filter);
     }
   }
-
-  private async fetchPivotAggregates(
-    filters: ReadonlyArray<Filter>,
-    pivot: Pivot,
-  ): Promise<Row> {
-    const resolver = new SQLSchemaResolver(this.sqlSchema, this.rootSchemaName);
-
-    const baseTable = resolver.getBaseTable();
-    const baseAlias = resolver.getBaseAlias();
-
-    // Resolve filter columns first to ensure JOINs are added
-    for (const filter of filters) {
-      resolver.resolveColumnPath(filter.field);
-    }
-
-    const aggregates = pivot.aggregates ?? [];
-    const selectClauses = aggregates
-      .map((agg) => {
-        if (agg.function === 'COUNT') {
-          return `COUNT(*) AS __count__`;
-        }
-        const field = 'field' in agg ? agg.field : null;
-        if (!field) return null;
-        const alias = this.pathToAlias(field);
-        if (agg.function === 'ANY') {
-          return `NULL AS ${alias}`;
-        }
-        const colExpr = resolver.resolveColumnPath(field);
-        if (!colExpr) {
-          return `NULL AS ${alias}`;
-        }
-        return `${agg.function}(${colExpr}) AS ${alias}`;
-      })
-      .filter(Boolean)
-      .join(', ');
-
-    if (!selectClauses) {
-      return {};
-    }
-
-    const joinClauses = resolver.buildJoinClauses();
-
-    let query = `
-SELECT ${selectClauses}
-FROM ${baseTable} AS ${baseAlias}
-${joinClauses}`;
-
-    if (filters.length > 0) {
-      const filterResolver = new SQLSchemaResolver(
-        this.sqlSchema,
-        this.rootSchemaName,
-      );
-      const whereConditions = filters.map((filter) => {
-        const sqlExpr = filterResolver.resolveColumnPath(filter.field);
-        return this.filterToSql(filter, sqlExpr ?? filter.field);
-      });
-      query += `\nWHERE ${whereConditions.join(' AND ')}`;
-    }
-
-    const result = await runQueryForQueryTable(
-      this.wrapQueryWithPrelude(query),
-      this.engine,
-    );
-    return result.rows[0] ?? {};
-  }
-
-  private async fetchColumnAggregates(
-    filters: ReadonlyArray<Filter>,
-    columns: ReadonlyArray<Column>,
-  ): Promise<Row> {
-    const resolver = new SQLSchemaResolver(this.sqlSchema, this.rootSchemaName);
-
-    const baseTable = resolver.getBaseTable();
-    const baseAlias = resolver.getBaseAlias();
-
-    const selectClauses = columns
-      .filter((col) => col.aggregate)
-      .map((col) => {
-        const func = col.aggregate!;
-        const colExpr = resolver.resolveColumnPath(col.field);
-        const alias = this.pathToAlias(col.field);
-
-        if (!colExpr) {
-          return `NULL AS ${alias}`;
-        }
-        if (func === 'ANY') {
-          return `MIN(${colExpr}) AS ${alias}`;
-        }
-        return `${func}(${colExpr}) AS ${alias}`;
-      })
-      .join(', ');
-
-    if (!selectClauses) {
-      return {};
-    }
-
-    // Resolve filter column paths first to ensure JOINs are added
-    for (const filter of filters) {
-      resolver.resolveColumnPath(filter.field);
-    }
-
-    const joinClauses = resolver.buildJoinClauses();
-
-    let query = `
-SELECT ${selectClauses}
-FROM ${baseTable} AS ${baseAlias}
-${joinClauses}`;
-
-    if (filters.length > 0) {
-      const whereConditions = filters.map((filter) => {
-        const sqlExpr = resolver.resolveColumnPath(filter.field);
-        return this.filterToSql(filter, sqlExpr ?? filter.field);
-      });
-      query += `\nWHERE ${whereConditions.join(' AND ')}`;
-    }
-
-    const result = await runQueryForQueryTable(
-      this.wrapQueryWithPrelude(query),
-      this.engine,
-    );
-    return result.rows[0] ?? {};
-  }
 }
 
 function sqlValue(value: SqlValue): string {
@@ -972,8 +905,15 @@
   }
 }
 
-function comparePagination(a?: Pagination, b?: Pagination): boolean {
-  if (!a && !b) return true;
-  if (!a || !b) return false;
-  return a.limit === b.limit && a.offset === b.offset;
+function queryResultToRows(queryResult: QueryResult): Row[] {
+  const rows: Row[] = [];
+  const columns = queryResult.columns();
+  for (const iter = queryResult.iter({}); iter.valid(); iter.next()) {
+    const row: Row = {};
+    for (const colName of columns) {
+      row[colName] = iter.get(colName);
+    }
+    rows.push(row);
+  }
+  return rows;
 }
diff --git a/ui/src/plugins/dev.perfetto.Sched/cpu_slice_selection_aggregator.ts b/ui/src/plugins/dev.perfetto.Sched/cpu_slice_selection_aggregator.ts
index b561943..a947ea4 100644
--- a/ui/src/plugins/dev.perfetto.Sched/cpu_slice_selection_aggregator.ts
+++ b/ui/src/plugins/dev.perfetto.Sched/cpu_slice_selection_aggregator.ts
@@ -39,6 +39,7 @@
   dur: LONG,
   ts: LONG,
   utid: NUM,
+  ucpu: NUM,
 };
 
 export class CpuSliceSelectionAggregator implements Aggregator {
@@ -110,6 +111,7 @@
             sched.dur,
             sched.dur * 1.0 / sum(sched.dur) OVER () as fraction_of_total,
             sched.dur * 1.0 / ${area.end - area.start} as fraction_of_selection,
+            ucpu,
             __groupid,
             __partition
           from ${iiTable.name} as sched
@@ -130,12 +132,7 @@
 
   getColumnDefinitions(): AggregatePivotModel {
     return {
-      groupBy: [
-        {field: 'pid'},
-        {field: 'process_name'},
-        {field: 'tid'},
-        {field: 'thread_name'},
-      ],
+      groupBy: [{field: 'process_name'}, {field: 'thread_name'}],
       aggregates: [
         {function: 'COUNT'},
         {field: 'dur', function: 'SUM', sort: 'DESC'},
@@ -191,7 +188,7 @@
           formatHint: 'NUMERIC',
         },
         {
-          title: 'Process Name',
+          title: 'Process',
           columnId: 'process_name',
           formatHint: 'STRING',
         },
@@ -201,7 +198,7 @@
           formatHint: 'NUMERIC',
         },
         {
-          title: 'Thread Name',
+          title: 'Thread',
           columnId: 'thread_name',
           formatHint: 'STRING',
         },
@@ -221,14 +218,9 @@
           formatHint: 'PERCENT',
         },
         {
-          title: 'Partition',
-          columnId: '__partition',
-          formatHint: 'ID',
-        },
-        {
-          title: 'GroupID',
-          columnId: '__groupid',
-          formatHint: 'ID',
+          title: 'CPU',
+          columnId: 'ucpu',
+          formatHint: 'NUM',
         },
       ],
     };
diff --git a/ui/src/plugins/dev.perfetto.Sched/thread_state_selection_aggregator.ts b/ui/src/plugins/dev.perfetto.Sched/thread_state_selection_aggregator.ts
index a79759a..0d9851c 100644
--- a/ui/src/plugins/dev.perfetto.Sched/thread_state_selection_aggregator.ts
+++ b/ui/src/plugins/dev.perfetto.Sched/thread_state_selection_aggregator.ts
@@ -170,7 +170,7 @@
 
   getColumnDefinitions(): AggregatePivotModel {
     return {
-      groupBy: [{field: 'utid'}, {field: 'state'}],
+      groupBy: [{field: 'thread_name'}, {field: 'state'}],
       aggregates: [
         {function: 'COUNT'},
         {field: 'process_name', function: 'ANY'},
diff --git a/ui/src/trace_processor/query_cache.ts b/ui/src/trace_processor/query_cache.ts
new file mode 100644
index 0000000..0f0fbe5
--- /dev/null
+++ b/ui/src/trace_processor/query_cache.ts
@@ -0,0 +1,48 @@
+// Copyright (C) 2026 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import {Engine} from './engine';
+import {QueryResult} from './query_result';
+
+export type UseQueryResult<T> =
+  | {
+      result: T;
+      isLoading: boolean;
+    }
+  | {
+      result?: undefined;
+      isLoading: true;
+    };
+
+export function createQueryCache(engine: Engine) {
+  const cache = new Map<string, QueryResult>();
+  // Fetches or triggers a cached query - returning undefined if not cached
+  return {
+    useQuery(query: string): UseQueryResult<QueryResult> {
+      const result = cache.get(query);
+      if (!result) {
+        // Kick off the query and store the result in the cache when done
+        engine.query(query).then((res) => {
+          cache.set(query, res);
+        });
+        return {isLoading: true};
+      }
+      return {result, isLoading: false};
+    },
+    get isLoading(): boolean {
+      // TODO: implement proper loading state
+      return false;
+    },
+  };
+}