exp: Add aggregation node
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/explore_page.ts b/ui/src/plugins/dev.perfetto.ExplorePage/explore_page.ts
index 0bbc872..a2fa23d 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/explore_page.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/explore_page.ts
@@ -28,6 +28,7 @@
 } from './query_builder/sources/slices_source';
 import {SqlSourceNode} from './query_builder/sources/sql_source';
 import {SubQueryNode} from './query_builder/sub_query_node';
+import {AggregationNode} from './query_builder/aggregation_node';
 import {Trace} from '../../public/trace';
 import {VisViewSource} from './data_visualiser/view_source';
 
@@ -50,8 +51,16 @@
 }
 
 export class ExplorePage implements m.ClassComponent<ExplorePageAttrs> {
-  private addNode(state: ExplorePageState, newNode: QueryNode) {
-    state.rootNodes.push(newNode);
+  private addNode(
+    state: ExplorePageState,
+    newNode: QueryNode,
+    prevNode?: QueryNode,
+  ) {
+    if (prevNode) {
+      prevNode.nextNodes.push(newNode);
+    } else {
+      state.rootNodes.push(newNode);
+    }
     this.selectNode(state, newNode);
   }
 
@@ -81,21 +90,28 @@
           sqlTable: selection.sqlTable,
           sourceCols: selection.sourceCols,
           filters: [],
-          groupByColumns: selection.groupByColumns,
-          aggregations: [],
         }),
       );
     }
   }
 
+  handleAddAggregation(state: ExplorePageState, node: QueryNode) {
+    const newNode = new AggregationNode({
+      prevNode: node,
+      sourceCols: node.finalCols,
+      groupByColumns: [],
+      aggregations: [],
+      filters: [],
+    });
+    this.addNode(state, newNode, node);
+  }
+
   handleAddSlicesSource(state: ExplorePageState) {
     this.addNode(
       state,
       new SlicesSourceNode({
         sourceCols: slicesSourceNodeColumns(true),
         filters: [],
-        groupByColumns: slicesSourceNodeColumns(false),
-        aggregations: [],
       }),
     );
   }
@@ -107,8 +123,6 @@
         trace: attrs.trace,
         sourceCols: [],
         filters: [],
-        groupByColumns: [],
-        aggregations: [],
       }),
     );
   }
@@ -123,13 +137,26 @@
   }
 
   handleDeleteNode(state: ExplorePageState, node: QueryNode) {
-    const idx = state.rootNodes.indexOf(node);
-    if (idx !== -1) {
-      state.rootNodes.splice(idx, 1);
-      if (state.selectedNode === node) {
-        this.deselectNode(state);
+    // If the node is a root node, remove it from the root nodes array.
+    const rootIdx = state.rootNodes.indexOf(node);
+    if (rootIdx !== -1) {
+      state.rootNodes.splice(rootIdx, 1);
+    }
+
+    // If the node is a child of another node, remove it from the parent's
+    // nextNodes array.
+    if (node.prevNode) {
+      const prevNode = node.prevNode;
+      const childIdx = prevNode.nextNodes.indexOf(node);
+      if (childIdx !== -1) {
+        prevNode.nextNodes.splice(childIdx, 1);
       }
     }
+
+    // If the deleted node was selected, deselect it.
+    if (state.selectedNode === node) {
+      this.deselectNode(state);
+    }
   }
 
   handleAddSubQuery(state: ExplorePageState, node: QueryNode) {
@@ -137,10 +164,8 @@
       prevNode: node,
       sourceCols: node.finalCols,
       filters: [],
-      groupByColumns: [],
-      aggregations: [],
     });
-    this.addNode(state, newNode);
+    this.addNode(state, newNode, node);
   }
 
   private handleKeyDown(event: KeyboardEvent, attrs: ExplorePageAttrs) {
@@ -208,6 +233,8 @@
           onDuplicateNode: (node) => this.handleDuplicateNode(state, node),
           onDeleteNode: (node) => this.handleDeleteNode(state, node),
           onAddSubQueryNode: (node) => this.handleAddSubQuery(state, node),
+          onAddAggregationNode: (node) =>
+            this.handleAddAggregation(state, node),
         }),
       state.mode === ExplorePageModes.DATA_VISUALISER &&
         state.rootNodes.length !== 0 &&
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/aggregation_node.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/aggregation_node.ts
new file mode 100644
index 0000000..b201c1d
--- /dev/null
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/aggregation_node.ts
@@ -0,0 +1,403 @@
+// Copyright (C) 2025 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 m from 'mithril';
+import {
+  QueryNode,
+  QueryNodeState,
+  nextNodeId,
+  NodeType,
+  createSelectColumnsProto,
+} from '../query_node';
+import protos from '../../../protos';
+import {ColumnInfo, columnInfoFromName, newColumnInfoList} from './column_info';
+import {createFiltersProto} from './operations/operation_component';
+import {MultiselectInput} from '../../../widgets/multiselect_input';
+import {Select} from '../../../widgets/select';
+import {TextInput} from '../../../widgets/text_input';
+import {Button} from '../../../widgets/button';
+
+export interface AggregationNodeState extends QueryNodeState {
+  readonly prevNode: QueryNode;
+  readonly sourceCols: ColumnInfo[];
+  readonly groupByColumns: ColumnInfo[];
+  readonly aggregations: Aggregation[];
+}
+
+export interface Aggregation {
+  column?: ColumnInfo;
+  aggregationOp?: string;
+  newColumnName?: string;
+  isValid?: boolean;
+  isEditing?: boolean;
+}
+
+export class AggregationNode implements QueryNode {
+  readonly nodeId: string;
+  readonly type = NodeType.kAggregation;
+  readonly prevNode?: QueryNode;
+  nextNodes: QueryNode[];
+  readonly sourceCols: ColumnInfo[];
+  readonly state: AggregationNodeState;
+
+  get finalCols(): ColumnInfo[] {
+    const selected = this.state.groupByColumns.filter((c) => c.checked);
+    for (const agg of this.state.aggregations) {
+      selected.push(
+        columnInfoFromName(agg.newColumnName ?? placeholderNewColumnName(agg)),
+      );
+    }
+    return newColumnInfoList(selected, true);
+  }
+
+  constructor(state: AggregationNodeState) {
+    this.nodeId = nextNodeId();
+    this.state = {
+      ...state,
+      groupByColumns: newColumnInfoList(state.sourceCols, false),
+    };
+    this.prevNode = state.prevNode;
+    this.sourceCols = state.sourceCols;
+    this.nextNodes = [];
+  }
+
+  validate(): boolean {
+    return this.prevNode !== undefined;
+  }
+
+  getTitle(): string {
+    return this.state.customTitle ?? 'Aggregation';
+  }
+
+  nodeSpecificModify(): m.Child {
+    return m(AggregationOperationComponent, {
+      groupByColumns: this.state.groupByColumns,
+      aggregations: this.state.aggregations,
+      onchange: this.state.onchange,
+    });
+  }
+
+  clone(): QueryNode {
+    const stateCopy: AggregationNodeState = {
+      prevNode: this.state.prevNode,
+      sourceCols: newColumnInfoList(this.sourceCols),
+      groupByColumns: newColumnInfoList(this.state.groupByColumns),
+      aggregations: this.state.aggregations.map((a) => ({...a})),
+      filters: [],
+      customTitle: this.state.customTitle,
+      onchange: this.state.onchange,
+    };
+    return new AggregationNode(stateCopy);
+  }
+
+  getStructuredQuery(): protos.PerfettoSqlStructuredQuery | undefined {
+    if (!this.validate()) return;
+
+    const prevSq = this.prevNode!.getStructuredQuery();
+    if (!prevSq) return undefined;
+
+    const groupByProto = createGroupByProto(
+      this.state.groupByColumns,
+      this.state.aggregations,
+    );
+    const filtersProto = createFiltersProto(
+      this.state.filters,
+      this.sourceCols,
+    );
+
+    // If the previous node already has an aggregation, we need to create a
+    // subquery.
+    if (prevSq.groupBy) {
+      const sq = new protos.PerfettoSqlStructuredQuery();
+      sq.id = this.nodeId;
+      sq.innerQuery = prevSq;
+      if (filtersProto) sq.filters = filtersProto;
+      if (groupByProto) sq.groupBy = groupByProto;
+      const selectedColumns = createSelectColumnsProto(this);
+      if (selectedColumns) sq.selectColumns = selectedColumns;
+      return sq;
+    }
+
+    // Otherwise, we can just add the aggregation and filters to the previous
+    // query.
+    if (filtersProto) {
+      if (prevSq.filters) {
+        prevSq.filters.push(...filtersProto);
+      } else {
+        prevSq.filters = filtersProto;
+      }
+    }
+    if (groupByProto) {
+      prevSq.groupBy = groupByProto;
+    }
+    const selectedColumns = createSelectColumnsProto(this);
+    if (selectedColumns) {
+      prevSq.selectColumns = selectedColumns;
+    }
+    return prevSq;
+  }
+}
+
+export function createGroupByProto(
+  groupByColumns: ColumnInfo[],
+  aggregations: Aggregation[],
+): protos.PerfettoSqlStructuredQuery.GroupBy | undefined {
+  if (!groupByColumns.find((c) => c.checked)) return;
+
+  const groupByProto = new protos.PerfettoSqlStructuredQuery.GroupBy();
+  groupByProto.columnNames = groupByColumns
+    .filter((c) => c.checked)
+    .map((c) => c.column.name);
+
+  for (const agg of aggregations) {
+    agg.isValid = validateAggregation(agg);
+  }
+  groupByProto.aggregates = aggregations
+    .filter((agg) => agg.isValid)
+    .map(GroupByAggregationAttrsToProto);
+  return groupByProto;
+}
+
+function validateAggregation(aggregation: Aggregation): boolean {
+  if (!aggregation.column || !aggregation.aggregationOp) return false;
+  return true;
+}
+
+export function GroupByAggregationAttrsToProto(
+  agg: Aggregation,
+): protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate {
+  const newAgg = new protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate();
+  newAgg.columnName = agg.column!.column.name;
+  newAgg.op = stringToAggregateOp(agg.aggregationOp!);
+  newAgg.resultColumnName = agg.newColumnName ?? placeholderNewColumnName(agg);
+  return newAgg;
+}
+
+export function placeholderNewColumnName(agg: Aggregation) {
+  return agg.column && agg.aggregationOp
+    ? `${agg.column.name}_${agg.aggregationOp}`
+    : `agg_${agg.aggregationOp ?? ''}`;
+}
+
+function stringToAggregateOp(
+  s: string,
+): protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate.Op {
+  if (AGGREGATION_OPS.includes(s as (typeof AGGREGATION_OPS)[number])) {
+    return protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate.Op[
+      s as keyof typeof protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate.Op
+    ];
+  }
+  throw new Error(`Invalid AggregateOp '${s}'`);
+}
+
+const AGGREGATION_OPS = [
+  'COUNT',
+  'SUM',
+  'MIN',
+  'MAX',
+  'MEAN',
+  'DURATION_WEIGHTED_MEAN',
+] as const;
+
+interface AggregationOperationComponentAttrs {
+  groupByColumns: ColumnInfo[];
+  aggregations: Aggregation[];
+  onchange?: () => void;
+}
+
+class AggregationOperationComponent
+  implements m.ClassComponent<AggregationOperationComponentAttrs>
+{
+  view({attrs}: m.CVnode<AggregationOperationComponentAttrs>) {
+    const hasGroupByColumns = attrs.groupByColumns.some((c) => c.checked);
+
+    if (hasGroupByColumns && attrs.aggregations.length === 0) {
+      attrs.aggregations.push({isEditing: true});
+    }
+
+    if (!hasGroupByColumns && attrs.aggregations.length > 0) {
+      // Clear aggregations if no group by columns are selected
+      attrs.aggregations.length = 0;
+    }
+
+    const selectGroupByColumns = (): m.Child => {
+      return m(
+        '.pf-exp-multi-select-container',
+        m('label', 'GROUP BY columns'),
+        m(MultiselectInput, {
+          options: attrs.groupByColumns.map((col) => ({
+            key: col.name,
+            label: col.name,
+          })),
+          selectedOptions: attrs.groupByColumns
+            .filter((c) => c.checked)
+            .map((c) => c.name),
+          onOptionAdd: (key: string) => {
+            const column = attrs.groupByColumns.find((c) => c.name === key);
+            if (column) {
+              column.checked = true;
+              attrs.onchange?.();
+              m.redraw();
+            }
+          },
+          onOptionRemove: (key: string) => {
+            const column = attrs.groupByColumns.find((c) => c.name === key);
+            if (column) {
+              column.checked = false;
+              attrs.onchange?.();
+              m.redraw();
+            }
+          },
+        }),
+      );
+    };
+
+    const aggregationEditor = (agg: Aggregation, index: number): m.Child => {
+      const columnOptions = attrs.groupByColumns.map((col) =>
+        m(
+          'option',
+          {
+            value: col.name,
+            selected: agg.column?.name === col.name,
+          },
+          col.name,
+        ),
+      );
+
+      return m(
+        '.pf-exp-aggregation-editor',
+        m(
+          Select,
+          {
+            onchange: (e: Event) => {
+              agg.aggregationOp = (e.target as HTMLSelectElement).value;
+              m.redraw();
+            },
+          },
+          m(
+            'option',
+            {disabled: true, selected: !agg.aggregationOp},
+            'Operation',
+          ),
+          AGGREGATION_OPS.map((op) =>
+            m(
+              'option',
+              {
+                value: op,
+                selected: op === agg.aggregationOp,
+              },
+              op,
+            ),
+          ),
+        ),
+        m(
+          Select,
+          {
+            onchange: (e: Event) => {
+              const target = e.target as HTMLSelectElement;
+              agg.column = attrs.groupByColumns.find(
+                (c) => c.name === target.value,
+              );
+              attrs.onchange?.();
+              m.redraw();
+            },
+          },
+          m('option', {disabled: true, selected: !agg.column}, 'Column'),
+          columnOptions,
+        ),
+        'AS',
+        m(TextInput, {
+          placeholder: placeholderNewColumnName(agg),
+          oninput: (e: Event) => {
+            agg.newColumnName = (e.target as HTMLInputElement).value.trim();
+          },
+          value: agg.newColumnName,
+        }),
+        m(Button, {
+          className: 'delete-button',
+          icon: 'delete',
+          onclick: () => {
+            attrs.aggregations.splice(index, 1);
+            attrs.onchange?.();
+          },
+        }),
+        m(Button, {
+          label: 'Done',
+          className: 'is-primary',
+          disabled: !agg.isValid,
+          onclick: () => {
+            if (!agg.newColumnName) {
+              agg.newColumnName = placeholderNewColumnName(agg);
+            }
+            agg.isEditing = false;
+            attrs.onchange?.();
+          },
+        }),
+      );
+    };
+
+    const aggregationViewer = (agg: Aggregation, index: number): m.Child => {
+      return m(
+        '.pf-exp-aggregation-viewer',
+        {
+          onclick: () => {
+            attrs.aggregations.forEach((a, i) => {
+              a.isEditing = i === index;
+            });
+            m.redraw();
+          },
+        },
+        `${agg.aggregationOp}(${agg.column?.name}) AS ${agg.newColumnName}`,
+      );
+    };
+
+    const aggregationsList = (): m.Children => {
+      if (!hasGroupByColumns) {
+        return null;
+      }
+
+      const lastAgg = attrs.aggregations[attrs.aggregations.length - 1];
+      const showAddButton = lastAgg.isValid && !lastAgg.isEditing;
+
+      return [
+        ...attrs.aggregations.map((agg, index) => {
+          if (agg.isEditing) {
+            return aggregationEditor(agg, index);
+          } else {
+            return aggregationViewer(agg, index);
+          }
+        }),
+        showAddButton &&
+          m(Button, {
+            label: 'Add more aggregations',
+            onclick: () => {
+              attrs.aggregations.push({isEditing: true});
+              attrs.onchange?.();
+            },
+          }),
+      ];
+    };
+
+    return m('.pf-exp-query-operations', [
+      m(
+        '.pf-exp-section',
+        m(
+          '.pf-exp-operations-container',
+          selectGroupByColumns(),
+          m('.pf-exp-aggregations-list', aggregationsList()),
+        ),
+      ),
+    ]);
+  }
+}
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/arrow.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/arrow.ts
index 5c73a40..b000133 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/arrow.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/arrow.ts
@@ -18,12 +18,15 @@
 export interface ArrowAttrs {
   from: NodeBoxLayout;
   to: NodeBoxLayout;
+  portIndex: number;
+  portCount: number;
 }
 
 export const Arrow: m.Component<ArrowAttrs> = {
   view({attrs}) {
-    const {from, to} = attrs;
-    const x1 = from.x + (from.width ?? 0) / 2;
+    const {from, to, portIndex, portCount} = attrs;
+    const fromWidth = from.width ?? 0;
+    const x1 = from.x + (fromWidth * (portIndex + 1)) / (portCount + 1);
     const y1 = from.y + (from.height ?? 0);
     const x2 = to.x + (to.width ?? 0) / 2;
     const y2 = to.y;
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/builder.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/builder.ts
index e79a481..eab7c35 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/builder.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/builder.ts
@@ -28,7 +28,7 @@
 } from '../../../components/widgets/data_grid/common';
 import {InMemoryDataSource} from '../../../components/widgets/data_grid/in_memory_data_source';
 import {QueryResponse} from '../../../components/query_table/queries';
-import {columnInfoFromSqlColumn, newColumnInfoList} from './column_info';
+import {columnInfoFromSqlColumn} from './column_info';
 import {TableSourceNode} from './sources/table_source';
 import {SqlSourceNode} from './sources/sql_source';
 import {QueryService} from './query_service';
@@ -52,6 +52,7 @@
 
   // Add derived nodes.
   readonly onAddSubQueryNode: (node: QueryNode) => void;
+  readonly onAddAggregationNode: (node: QueryNode) => void;
 
   readonly onClearAllNodes: () => void;
   readonly onDuplicateNode: (node: QueryNode) => void;
@@ -139,17 +140,13 @@
             const sourceCols = sqlTable.columns.map((c) =>
               columnInfoFromSqlColumn(c, true),
             );
-            const groupByColumns = newColumnInfoList(sourceCols, false);
-
             onRootNodeCreated(
               new TableSourceNode({
                 trace,
                 sqlModules,
                 sqlTable,
                 sourceCols,
-                groupByColumns,
                 filters: [],
-                aggregations: [],
               }),
             );
           },
@@ -170,6 +167,7 @@
           onClearAllNodes,
           onDuplicateNode: attrs.onDuplicateNode,
           onAddSubQuery: attrs.onAddSubQueryNode,
+          onAddAggregation: attrs.onAddAggregationNode,
           onDeleteNode: (node: QueryNode) => {
             if (
               node.state.isExecuted &&
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/data_explorer.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/data_explorer.ts
index a0cd1fd..4ca1d34 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/data_explorer.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/data_explorer.ts
@@ -25,14 +25,12 @@
   renderCell,
 } from '../../../components/widgets/data_grid/data_grid';
 import {SqlValue} from '../../../trace_processor/query_result';
-import {Button, ButtonVariant} from '../../../widgets/button';
+import {Button} from '../../../widgets/button';
 import {Callout} from '../../../widgets/callout';
 import {DetailsShell} from '../../../widgets/details_shell';
 import {MenuItem, PopupMenu} from '../../../widgets/menu';
 import {TextParagraph} from '../../../widgets/text_paragraph';
 import {Query, QueryNode} from '../query_node';
-import {Intent} from '../../../widgets/common';
-import {AggregationsOperator} from './operations/aggregations';
 import {QueryService} from './query_service';
 
 import {findErrors} from './query_builder_utils';
@@ -56,8 +54,6 @@
 }
 
 export class DataExplorer implements m.ClassComponent<DataExplorerAttrs> {
-  private showAggregationCard: boolean = false;
-
   view({attrs}: m.CVnode<DataExplorerAttrs>) {
     const errors = findErrors(attrs.query, attrs.response);
     const statusText = this.getStatusText(attrs.query, attrs.response);
@@ -142,32 +138,6 @@
             )
           : null;
 
-      const maybeAggregateButton =
-        attrs.isFullScreen &&
-        m(
-          '.pf-ndv-floating-button',
-          m(Button, {
-            intent: Intent.Primary,
-            variant: ButtonVariant.Filled,
-            label: 'Aggregate',
-            onclick: () => {
-              this.showAggregationCard = !this.showAggregationCard;
-            },
-          }),
-        );
-
-      const maybeAggregationCard =
-        this.showAggregationCard &&
-        m(
-          '.pf-ndv-floating-card',
-          m(AggregationsOperator, {
-            groupByColumns: attrs.node.state.groupByColumns,
-            aggregations: attrs.node.state.aggregations,
-          }),
-        );
-
-      const hasAggregations = (attrs.node.state.aggregations?.length ?? 0) > 0;
-
       return [
         warning,
         m(DataGrid, {
@@ -176,18 +146,14 @@
           data: attrs.dataSource,
           showFiltersInToolbar: true,
           filters: attrs.node.state.filters,
-          onFiltersChanged: hasAggregations
-            ? undefined
-            : (filters: ReadonlyArray<FilterDefinition>) => {
-                attrs.node.state.filters = [...filters];
-                attrs.onchange?.();
-              },
+          onFiltersChanged: (filters: ReadonlyArray<FilterDefinition>) => {
+            attrs.node.state.filters = [...filters];
+            attrs.onchange?.();
+          },
           cellRenderer: (value: SqlValue, name: string) => {
             return renderCell(value, name);
           },
         }),
-        maybeAggregateButton,
-        maybeAggregationCard,
       ];
     }
     return null;
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/graph.scss b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/graph.scss
index 37d33f9..f2f2d88 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/graph.scss
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/graph.scss
@@ -72,12 +72,16 @@
   height: 10px;
   background-color: var(--pf-color-border);
   border-radius: 50%;
-  left: 50%;
-  transform: translateX(-50%);
 }
 
 .pf-node-box-port-top {
   top: -5px;
+  left: 50%;
+  margin-left: -5px;
+}
+
+.pf-node-box-port-bottom {
+  bottom: -5px;
 }
 
 .pf-node-box-port-bottom {
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/graph.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/graph.ts
index 2959d17..0c73109 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/graph.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/graph.ts
@@ -66,15 +66,26 @@
   readonly onAddSlicesSource: () => void;
   readonly onAddSqlSource: () => void;
   readonly onAddSubQuery: (node: QueryNode) => void;
+  readonly onAddAggregation: (node: QueryNode) => void;
   readonly onClearAllNodes: () => void;
   readonly onDuplicateNode: (node: QueryNode) => void;
   readonly onDeleteNode: (node: QueryNode) => void;
 }
 
 export class Graph implements m.ClassComponent<GraphAttrs> {
+  // The node currently being dragged. This is used to apply styles and
+  // transformations to the node while it is being moved.
   private dragNode?: QueryNode;
+  // A map from nodes to their layout information (position and size). This
+  // allows us to quickly look up the position of any node in the graph.
   private nodeLayouts: Map<QueryNode, NodeBoxLayout> = new Map();
+  // The width of the node graph area. This is used to constrain the nodes
+  // within the bounds of the graph.
   private nodeGraphWidth: number = 0;
+  // The offset of the mouse cursor from the top-left corner of the dragged
+  // node. This is used to prevent the node from jumping to the cursor's
+  // position when the drag starts.
+  private dragOffset?: {x: number; y: number};
 
   oncreate({dom}: m.VnodeDOM<GraphAttrs>) {
     const box = dom as HTMLElement;
@@ -82,6 +93,25 @@
 
     box.ondragover = (event) => {
       event.preventDefault(); // Allow dropping
+      if (this.dragNode) {
+        const dragNodeLayout = this.nodeLayouts.get(this.dragNode);
+        if (dragNodeLayout && this.dragOffset) {
+          const rect = box.getBoundingClientRect();
+          const w = dragNodeLayout.width ?? DEFAULT_NODE_WIDTH;
+          const h = dragNodeLayout.height ?? NODE_HEIGHT;
+          // To provide real-time feedback to the user, we continuously update
+          // the node's position during the drag operation. This allows the
+          // connecting arrows to follow the node smoothly.
+          const x = event.clientX - rect.left - this.dragOffset.x;
+          const y = event.clientY - rect.top - this.dragOffset.y;
+          this.nodeLayouts.set(this.dragNode, {
+            ...dragNodeLayout,
+            x: Math.max(0, Math.min(x, rect.width - w)),
+            y: Math.max(0, Math.min(y, rect.height - h)),
+          });
+          m.redraw();
+        }
+      }
     };
 
     box.ondrop = (event) => {
@@ -91,6 +121,7 @@
     box.ondragend = () => {
       if (this.dragNode) {
         this.dragNode = undefined;
+        this.dragOffset = undefined;
         m.redraw();
       }
     };
@@ -106,15 +137,10 @@
     const w = dragNodeLayout.width ?? DEFAULT_NODE_WIDTH;
     const h = dragNodeLayout.height ?? NODE_HEIGHT;
 
-    const x = event.clientX - rect.left - w / 2;
-    const y = event.clientY - rect.top - h / 2;
-
-    const initialLayout: NodeBoxLayout = {
-      ...dragNodeLayout,
-      x: Math.max(0, Math.min(x, rect.width - w)),
-      y: Math.max(0, Math.min(y, rect.height - h)),
-    };
-
+    // The "Add Node" and "Clear All Nodes" buttons occupy a fixed area in the
+    // top-right corner of the graph. To prevent nodes from being dropped on
+    // top of these buttons, we define a reserved area that is treated as an
+    // obstacle.
     const buttonsReservedArea: NodeBoxLayout = {
       x: this.nodeGraphWidth - BUTTONS_AREA_WIDTH - PADDING,
       y: PADDING,
@@ -128,8 +154,12 @@
 
     const allLayouts = [...otherLayouts, buttonsReservedArea];
 
+    // After the node is dropped, we need to find a final position for it that
+    // doesn't overlap with any other nodes. This is important because the
+    // user can drag the node over other nodes, and we want to ensure that
+    // the graph is still readable after the drag operation is complete.
     const newLayout = findNonOverlappingLayout(
-      initialLayout,
+      dragNodeLayout,
       allLayouts,
       w,
       h,
@@ -145,6 +175,10 @@
     if (layout) {
       const newWidth = element.offsetWidth;
       const newHeight = element.offsetHeight;
+      // The dimensions of a node can change after it is rendered, for example
+      // if the user edits the node's title. To ensure that the layout is
+      // always accurate, we update the node's dimensions in the layout map
+      // whenever they change.
       if (layout.width !== newWidth || layout.height !== newHeight) {
         this.nodeLayouts.set(node, {
           ...layout,
@@ -168,6 +202,16 @@
       height: nodeElem.offsetHeight,
     });
 
+    // To prevent the node from jumping to the cursor's position when a drag
+    // starts, we calculate the initial offset of the cursor from the
+    // top-left corner of the node. This offset is then used to maintain the
+    // node's position relative to the cursor throughout the drag operation.
+    const rect = nodeElem.getBoundingClientRect();
+    this.dragOffset = {
+      x: event.clientX - rect.left,
+      y: event.clientY - rect.top,
+    };
+
     if (event.dataTransfer) {
       event.dataTransfer.setData('text/plain', node.getTitle());
       event.dataTransfer.effectAllowed = 'move';
@@ -254,10 +298,13 @@
 
     const allNodes: QueryNode[] = [];
     for (const root of rootNodes) {
-      let curr: QueryNode | undefined = root;
-      while (curr) {
+      const queue: QueryNode[] = [root];
+      while (queue.length > 0) {
+        const curr = queue.shift()!;
         allNodes.push(curr);
-        curr = curr.nextNode;
+        for (const child of curr.nextNodes) {
+          queue.push(child);
+        }
       }
     }
 
@@ -277,18 +324,26 @@
       children.push(this.renderEmptyNodeGraph(attrs));
     } else {
       for (const node of allNodes) {
-        if (node.prevNode) {
-          const prevLayout = this.nodeLayouts.get(node.prevNode);
-          const layout = this.nodeLayouts.get(node);
-          if (prevLayout && layout) {
-            children.push(m(Arrow, {from: prevLayout, to: layout}));
+        node.nextNodes.forEach((child, i) => {
+          const from = this.nodeLayouts.get(node);
+          const to = this.nodeLayouts.get(child);
+          if (from && to) {
+            children.push(
+              m(Arrow, {
+                from,
+                to,
+                portIndex: i,
+                portCount: node.nextNodes.length,
+              }),
+            );
           }
-        }
+        });
         let layout = this.nodeLayouts.get(node);
         if (!layout) {
           layout = findNextAvailablePosition(
             node,
             Array.from(this.nodeLayouts.values()),
+            this.nodeLayouts,
             this.nodeGraphWidth,
           );
           this.nodeLayouts.set(node, layout);
@@ -304,6 +359,7 @@
             onDuplicateNode: attrs.onDuplicateNode,
             onDeleteNode: attrs.onDeleteNode,
             onAddSubQuery: attrs.onAddSubQuery,
+            onAddAggregation: attrs.onAddAggregation,
             onNodeRendered: this.onNodeRendered,
           }),
         );
@@ -326,6 +382,11 @@
   }
 }
 
+// When a node is dropped, it might overlap with other nodes. This function
+// resolves such overlaps by finding the nearest available position for the
+// node. It works by checking for collisions and then shifting the node just
+// enough to clear the obstacle. This process is repeated until no more
+// overlaps are detected.
 function findNonOverlappingLayout(
   initialLayout: NodeBoxLayout,
   otherLayouts: NodeBoxLayout[],
@@ -340,18 +401,29 @@
       const layoutW = layout.width ?? DEFAULT_NODE_WIDTH;
       const layoutH = layout.height ?? NODE_HEIGHT;
 
+      // To resolve an overlap, we can move the node in one of four
+      // directions: right, left, down, or up. We calculate the target
+      // position for each of these moves.
       const right = layout.x + layoutW + PADDING;
       const left = layout.x - w - PADDING;
       const bottom = layout.y + layoutH + PADDING;
       const top = layout.y - h - PADDING;
 
+      // We want to move the node by the smallest possible amount to resolve
+      // the overlap. To do this, we calculate the distance to each of the
+      // four possible positions.
       const distRight = Math.abs(newLayout.x - right);
       const distLeft = Math.abs(newLayout.x - left);
       const distBottom = Math.abs(newLayout.y - bottom);
       const distTop = Math.abs(newLayout.y - top);
 
+      // The shortest distance determines the direction in which the node will
+      // be moved.
       const minDist = Math.min(distRight, distLeft, distBottom, distTop);
 
+      // By moving the node to the closest non-overlapping position, we
+      // ensure that the layout remains as stable as possible after the drag
+      // operation is complete.
       if (minDist === distRight) {
         newLayout.x = right;
       } else if (minDist === distLeft) {
@@ -364,12 +436,17 @@
     }
   }
 
+  // Finally, we ensure that the new layout is still within the bounds of the
+  // graph. This prevents nodes from being moved outside of the visible area.
   newLayout.x = Math.max(0, Math.min(newLayout.x, rect.width - w));
   newLayout.y = Math.max(0, Math.min(newLayout.y, rect.height - h));
 
   return newLayout;
 }
 
+// This is a standard axis-aligned bounding box (AABB) collision detection
+// algorithm. It checks if two rectangles are overlapping by comparing their
+// positions and dimensions.
 function isOverlapping(
   layout1: NodeBoxLayout,
   layout2: NodeBoxLayout,
@@ -388,9 +465,15 @@
   );
 }
 
+// When a new node is added to the graph, we need to find a suitable position
+// for it. This function implements a simple grid-based placement algorithm. It
+// iterates through the graph from top to bottom, left to right, and places the
+// new node in the first available slot that doesn't overlap with any existing
+// nodes.
 function findNextAvailablePosition(
   node: QueryNode,
   layouts: NodeBoxLayout[],
+  nodeLayouts: Map<QueryNode, NodeBoxLayout>,
   nodeGraphWidth: number,
 ): NodeBoxLayout {
   const w = Math.max(DEFAULT_NODE_WIDTH, node.getTitle().length * 8 + 60);
@@ -405,6 +488,38 @@
 
   const allLayouts = [...layouts, buttonsReservedArea];
 
+  // If the node is a nextNode (e.g., an aggregation or sub-query), it should
+  // be added below the previous node.
+  if (node.prevNode) {
+    const prevLayout = nodeLayouts.get(node.prevNode);
+    if (prevLayout) {
+      let x = prevLayout.x;
+      let y = prevLayout.y + (prevLayout.height ?? h) + PADDING * 2;
+      // Try to place the new node below the previous node, shifted by the
+      // number of siblings.
+      if (node.prevNode.nextNodes.length > 1) {
+        x +=
+          (node.prevNode.nextNodes.indexOf(node) -
+            (node.prevNode.nextNodes.length - 1) / 2) *
+          (w + PADDING);
+      }
+      while (true) {
+        const candidateLayout = {x, y, width: w, height: h};
+        let isInvalid = false;
+        for (const layout of allLayouts) {
+          if (isOverlapping(candidateLayout, layout, PADDING)) {
+            isInvalid = true;
+            y = layout.y + (layout.height ?? h) + PADDING;
+            break;
+          }
+        }
+        if (!isInvalid) {
+          return candidateLayout;
+        }
+      }
+    }
+  }
+
   let x = PADDING;
   let y = PADDING;
 
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_box.scss b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_box.scss
index 46a4888..0638942 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_box.scss
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_box.scss
@@ -85,3 +85,26 @@
 .pf-node-box {
   position: absolute;
 }
+
+.pf-node-box-add-button {
+  position: absolute;
+  bottom: -0.5rem;
+  right: -0.5rem;
+  background-color: var(--pf-color-background);
+  width: 1.6rem;
+  height: 1.6rem;
+  border-radius: 50%;
+  border: 1px solid;
+  display: none;
+  align-items: center;
+  justify-content: center;
+
+  .pf-node-box:hover &,
+  .pf-node-box__selected & {
+    display: flex;
+  }
+}
+
+.pf-node-box-port {
+  z-index: 1;
+}
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_box.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_box.ts
index 556dba6..e988596 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_box.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_box.ts
@@ -42,6 +42,7 @@
   readonly onDuplicateNode: (node: QueryNode) => void;
   readonly onDeleteNode: (node: QueryNode) => void;
   readonly onAddSubQuery: (node: QueryNode) => void;
+  readonly onAddAggregation: (node: QueryNode) => void;
   readonly onNodeRendered: (node: QueryNode, element: HTMLElement) => void;
 }
 
@@ -60,7 +61,8 @@
 }
 
 function renderContextMenu(attrs: NodeBoxAttrs): m.Child {
-  const {node, onDuplicateNode, onDeleteNode, onAddSubQuery} = attrs;
+  const {node, onDuplicateNode, onDeleteNode, onAddSubQuery, onAddAggregation} =
+    attrs;
   return m(
     PopupMenu,
     {
@@ -69,11 +71,16 @@
         icon: Icons.ContextMenuAlt,
       }),
     },
-    node.type !== NodeType.kSqlSource &&
+    node.type !== NodeType.kSqlSource && [
       m(MenuItem, {
         label: 'Add sub-query',
         onclick: () => onAddSubQuery(node),
       }),
+      m(MenuItem, {
+        label: 'Add aggregation',
+        onclick: () => onAddAggregation(node),
+      }),
+    ],
     m(MenuItem, {
       label: 'Duplicate',
       onclick: () => onDuplicateNode(node),
@@ -85,6 +92,23 @@
   );
 }
 
+function renderAddButton(attrs: NodeBoxAttrs): m.Child {
+  const {node, onAddAggregation} = attrs;
+  return m(
+    PopupMenu,
+    {
+      trigger: m(Icon, {
+        className: 'pf-node-box-add-button',
+        icon: 'add',
+      }),
+    },
+    m(MenuItem, {
+      label: 'Aggregate',
+      onclick: () => onAddAggregation(node),
+    }),
+  );
+}
+
 export const NodeBox: m.Component<NodeBoxAttrs> = {
   oncreate({attrs, dom}) {
     attrs.onNodeRendered(attrs.node, dom as HTMLElement);
@@ -128,7 +152,14 @@
       renderWarningIcon(node),
       m('span.pf-node-box__title', node.getTitle()),
       renderContextMenu(attrs),
-      node.nextNode && m('.pf-node-box-port.pf-node-box-port-bottom'),
+      node.nextNodes.map((_, i) => {
+        const portCount = node.nextNodes.length;
+        const left = `calc(${((i + 1) * 100) / (portCount + 1)}% - 5px)`;
+        return m('.pf-node-box-port.pf-node-box-port-bottom', {
+          style: {left},
+        });
+      }),
+      node.type !== NodeType.kSqlSource && renderAddButton(attrs),
     );
   },
 };
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_explorer.scss b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_explorer.scss
index 9ffe2ac..7059795 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_explorer.scss
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_explorer.scss
@@ -53,3 +53,46 @@
     border: none;
   }
 }
+
+.pf-exp-query-operations {
+  display: flex;
+  flex-direction: column;
+  gap: 1rem;
+  margin-bottom: 1rem;
+}
+
+.pf-exp-section {
+  display: flex;
+  flex-direction: column;
+  gap: 0.5rem;
+  padding: 1rem;
+  border: 1px solid var(--pf-color-border);
+  border-radius: 8px;
+}
+
+.pf-exp-operations-container {
+  display: flex;
+  flex-direction: column;
+  gap: 1rem;
+}
+
+.pf-exp-aggregations-list {
+  display: flex;
+  flex-direction: column;
+  gap: 0.5rem;
+}
+
+.pf-exp-aggregation-editor {
+  display: flex;
+  gap: 0.5rem;
+  align-items: center;
+}
+
+.pf-exp-aggregation-viewer {
+  cursor: pointer;
+  padding: 0.5rem;
+  border-radius: 4px;
+  &:hover {
+    background-color: var(--pf-color-hover);
+  }
+}
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_explorer.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_explorer.ts
index 11240ff..2b82b30 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_explorer.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/node_explorer.ts
@@ -116,10 +116,6 @@
           onchange?.();
         },
       },
-      groupby: {
-        groupByColumns: node.state.groupByColumns,
-        aggregations: node.state.aggregations,
-      },
       onchange: () => {
         setOperationChanged(node);
         onchange?.();
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/operations/aggregations.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/operations/aggregations.ts
deleted file mode 100644
index 01e5b6f..0000000
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/operations/aggregations.ts
+++ /dev/null
@@ -1,259 +0,0 @@
-// Copyright (C) 2025 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 m from 'mithril';
-import {ColumnInfo} from '../column_info';
-import {MultiselectInput} from '../../../../widgets/multiselect_input';
-import {Select} from '../../../../widgets/select';
-import {TextInput} from '../../../../widgets/text_input';
-import {Button} from '../../../../widgets/button';
-import protos from '../../../../protos';
-
-export interface Aggregation {
-  column?: ColumnInfo;
-  aggregationOp?: string;
-  newColumnName?: string;
-  isValid?: boolean;
-  isEditing?: boolean;
-}
-
-export interface AggregationsOperatorAttrs {
-  groupByColumns: ColumnInfo[];
-  aggregations: Aggregation[];
-  onchange?: () => void;
-}
-
-const AGGREGATION_OPS = [
-  'COUNT',
-  'SUM',
-  'MIN',
-  'MAX',
-  'MEAN',
-  'DURATION_WEIGHTED_MEAN',
-] as const;
-
-export class AggregationsOperator
-  implements m.ClassComponent<AggregationsOperatorAttrs>
-{
-  view({attrs}: m.CVnode<AggregationsOperatorAttrs>) {
-    if (attrs.groupByColumns.length === 0) {
-      return;
-    }
-
-    const hasGroupByColumns = attrs.groupByColumns.some((c) => c.checked);
-
-    if (hasGroupByColumns && attrs.aggregations.length === 0) {
-      attrs.aggregations.push({isEditing: true});
-    }
-
-    if (!hasGroupByColumns && attrs.aggregations.length > 0) {
-      // Clear aggregations if no group by columns are selected
-      attrs.aggregations.length = 0;
-    }
-
-    const selectGroupByColumns = (): m.Child => {
-      return m(
-        '.pf-exp-multi-select-container',
-        m('label', 'GROUP BY columns'),
-        m(MultiselectInput, {
-          options: attrs.groupByColumns.map((col) => ({
-            key: col.name,
-            label: col.name,
-          })),
-          selectedOptions: attrs.groupByColumns
-            .filter((c) => c.checked)
-            .map((c) => c.name),
-          onOptionAdd: (key: string) => {
-            const column = attrs.groupByColumns.find((c) => c.name === key);
-            if (column) {
-              column.checked = true;
-              attrs.onchange?.();
-              m.redraw();
-            }
-          },
-          onOptionRemove: (key: string) => {
-            const column = attrs.groupByColumns.find((c) => c.name === key);
-            if (column) {
-              column.checked = false;
-              attrs.onchange?.();
-              m.redraw();
-            }
-          },
-        }),
-      );
-    };
-
-    const aggregationEditor = (agg: Aggregation, index: number): m.Child => {
-      const columnOptions = attrs.groupByColumns.map((col) =>
-        m(
-          'option',
-          {
-            value: col.name,
-            selected: agg.column?.name === col.name,
-          },
-          col.name,
-        ),
-      );
-
-      return m(
-        '.pf-exp-aggregation-editor',
-        m(
-          Select,
-          {
-            onchange: (e: Event) => {
-              agg.aggregationOp = (e.target as HTMLSelectElement).value;
-              m.redraw();
-            },
-          },
-          m(
-            'option',
-            {disabled: true, selected: !agg.aggregationOp},
-            'Operation',
-          ),
-          AGGREGATION_OPS.map((op) =>
-            m(
-              'option',
-              {
-                value: op,
-                selected: op === agg.aggregationOp,
-              },
-              op,
-            ),
-          ),
-        ),
-        m(
-          Select,
-          {
-            onchange: (e: Event) => {
-              const target = e.target as HTMLSelectElement;
-              agg.column = attrs.groupByColumns.find(
-                (c) => c.name === target.value,
-              );
-              attrs.onchange?.();
-              m.redraw();
-            },
-          },
-          m('option', {disabled: true, selected: !agg.column}, 'Column'),
-          columnOptions,
-        ),
-        'AS',
-        m(TextInput, {
-          placeholder: placeholderNewColumnName(agg),
-          oninput: (e: Event) => {
-            agg.newColumnName = (e.target as HTMLInputElement).value.trim();
-          },
-          value: agg.newColumnName,
-        }),
-        m(Button, {
-          className: 'delete-button',
-          icon: 'delete',
-          onclick: () => {
-            attrs.aggregations.splice(index, 1);
-            attrs.onchange?.();
-          },
-        }),
-        m(Button, {
-          label: 'Done',
-          className: 'is-primary',
-          disabled: !agg.isValid,
-          onclick: () => {
-            if (!agg.newColumnName) {
-              agg.newColumnName = placeholderNewColumnName(agg);
-            }
-            agg.isEditing = false;
-            attrs.onchange?.();
-          },
-        }),
-      );
-    };
-
-    const aggregationViewer = (agg: Aggregation, index: number): m.Child => {
-      return m(
-        '.pf-exp-aggregation-viewer',
-        {
-          onclick: () => {
-            attrs.aggregations.forEach((a, i) => {
-              a.isEditing = i === index;
-            });
-            m.redraw();
-          },
-        },
-        `${agg.aggregationOp}(${agg.column?.name}) AS ${agg.newColumnName}`,
-      );
-    };
-
-    const aggregationsList = (): m.Children => {
-      if (!hasGroupByColumns) {
-        return null;
-      }
-
-      const lastAgg = attrs.aggregations[attrs.aggregations.length - 1];
-      const showAddButton = lastAgg.isValid && !lastAgg.isEditing;
-
-      return [
-        ...attrs.aggregations.map((agg, index) => {
-          if (agg.isEditing) {
-            return aggregationEditor(agg, index);
-          } else {
-            return aggregationViewer(agg, index);
-          }
-        }),
-        showAddButton &&
-          m(Button, {
-            label: 'Add more aggregations',
-            onclick: () => {
-              attrs.aggregations.push({isEditing: true});
-              attrs.onchange?.();
-            },
-          }),
-      ];
-    };
-
-    return m(
-      '.pf-exp-section',
-      m(
-        '.pf-exp-operations-container',
-        selectGroupByColumns(),
-        m('.pf-exp-aggregations-list', aggregationsList()),
-      ),
-    );
-  }
-}
-
-function stringToAggregateOp(
-  s: string,
-): protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate.Op {
-  if (AGGREGATION_OPS.includes(s as (typeof AGGREGATION_OPS)[number])) {
-    return protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate.Op[
-      s as keyof typeof protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate.Op
-    ];
-  }
-  throw new Error(`Invalid AggregateOp '${s}'`);
-}
-
-export function GroupByAggregationAttrsToProto(
-  agg: Aggregation,
-): protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate {
-  const newAgg = new protos.PerfettoSqlStructuredQuery.GroupBy.Aggregate();
-  newAgg.columnName = agg.column!.column.name;
-  newAgg.op = stringToAggregateOp(agg.aggregationOp!);
-  newAgg.resultColumnName = agg.newColumnName ?? placeholderNewColumnName(agg);
-  return newAgg;
-}
-
-export function placeholderNewColumnName(agg: Aggregation) {
-  return agg.column && agg.aggregationOp
-    ? `${agg.column.name}_${agg.aggregationOp}`
-    : `agg_${agg.aggregationOp ?? ''}`;
-}
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/operations/operation_component.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/operations/operation_component.ts
index b424d12..c1d145c 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/operations/operation_component.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/operations/operation_component.ts
@@ -14,28 +14,19 @@
 
 import m from 'mithril';
 import {ALL_FILTER_OPS, FilterAttrs, FilterOperation} from './filter';
-import {
-  Aggregation,
-  GroupByAggregationAttrsToProto,
-  AggregationsOperatorAttrs,
-  AggregationsOperator,
-} from './aggregations';
 import {FilterDefinition} from '../../../../components/widgets/data_grid/common';
-import {Button, ButtonVariant} from '../../../../widgets/button';
 import protos from '../../../../protos';
 import {ColumnInfo} from '../column_info';
 
 export interface OperatorAttrs {
   filter: FilterAttrs;
-  groupby: AggregationsOperatorAttrs;
   onchange?: () => void;
 }
 
 export class Operator implements m.ClassComponent<OperatorAttrs> {
-  private showAggregations = false;
-
   view({attrs}: m.CVnode<OperatorAttrs>): m.Children {
-    return m('.pf-exp-query-operations', [
+    return m(
+      '.pf-exp-query-operations',
       m(FilterOperation, {
         ...attrs.filter,
         onFiltersChanged: (filters: ReadonlyArray<FilterDefinition>) => {
@@ -43,20 +34,7 @@
           attrs.onchange?.();
         },
       }),
-      this.showAggregations
-        ? m(AggregationsOperator, {
-            ...attrs.groupby,
-            onchange: attrs.onchange,
-          })
-        : m(Button, {
-            label: 'Aggregate data',
-            onclick: () => {
-              this.showAggregations = true;
-              attrs.onchange?.();
-            },
-            variant: ButtonVariant.Filled,
-          }),
-    ]);
+    );
   }
 }
 
@@ -99,30 +77,3 @@
   );
   return protoFilters;
 }
-
-export function createGroupByProto(
-  groupByColumns: ColumnInfo[],
-  aggregations: Aggregation[],
-): protos.PerfettoSqlStructuredQuery.GroupBy | undefined {
-  if (!groupByColumns.find((c) => c.checked)) return;
-
-  const groupByProto = new protos.PerfettoSqlStructuredQuery.GroupBy();
-  groupByProto.columnNames = groupByColumns
-    .filter((c) => c.checked)
-    .map((c) => c.column.name);
-
-  for (const agg of aggregations) {
-    agg.isValid = validateAggregation(agg);
-  }
-  groupByProto.aggregates = aggregations
-    .filter((agg) => agg.isValid)
-    .map(GroupByAggregationAttrsToProto);
-  return groupByProto;
-}
-
-// Both 'column' and 'aggregationOp' must be present for an aggregation to be considered valid.
-// This ensures that the aggregation operation is applied to a specific column.
-function validateAggregation(aggregation: Aggregation): boolean {
-  if (!aggregation.column || !aggregation.aggregationOp) return false;
-  return true;
-}
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/source_node.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/source_node.ts
index b51e715..2115667 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/source_node.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/source_node.ts
@@ -26,7 +26,7 @@
 export abstract class SourceNode implements QueryNode {
   readonly nodeId: string;
   readonly prevNode = undefined;
-  nextNode?: QueryNode;
+  nextNodes: QueryNode[];
 
   sourceCols: ColumnInfo[];
   finalCols: ColumnInfo[];
@@ -38,6 +38,7 @@
     this.state = state;
     this.sourceCols = state.sourceCols ?? [];
     this.finalCols = createFinalColumns(this);
+    this.nextNodes = [];
   }
 
   abstract get type(): NodeType;
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/slices_source.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/slices_source.ts
index d512434..58f8cc9 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/slices_source.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/slices_source.ts
@@ -28,10 +28,7 @@
 import {TextInput} from '../../../../widgets/text_input';
 import {SqlColumn} from '../../../dev.perfetto.SqlModules/sql_modules';
 import {TableAndColumnImpl} from '../../../dev.perfetto.SqlModules/sql_modules_impl';
-import {
-  createFiltersProto,
-  createGroupByProto,
-} from '../operations/operation_component';
+import {createFiltersProto} from '../operations/operation_component';
 import {SourceNode} from '../source_node';
 
 export interface SlicesSourceState extends QueryNodeState {
@@ -50,6 +47,7 @@
     this.state = attrs;
     this.state.onchange = attrs.onchange;
     this.sourceCols = slicesSourceNodeColumns(true);
+    this.nextNodes = [];
   }
 
   get type() {
@@ -63,9 +61,7 @@
       process_name: this.state.process_name?.slice(),
       track_name: this.state.track_name?.slice(),
       sourceCols: newColumnInfoList(this.sourceCols),
-      groupByColumns: newColumnInfoList(this.state.groupByColumns),
       filters: this.state.filters.map((f) => ({...f})),
-      aggregations: this.state.aggregations.map((a) => ({...a})),
       customTitle: this.state.customTitle,
     };
     return new SlicesSourceNode(stateCopy);
@@ -94,11 +90,6 @@
       this.sourceCols,
     );
     if (filtersProto) sq.filters = filtersProto;
-    const groupByProto = createGroupByProto(
-      this.state.groupByColumns,
-      this.state.aggregations,
-    );
-    if (groupByProto) sq.groupBy = groupByProto;
 
     const selectedColumns = createSelectColumnsProto(this);
     if (selectedColumns) sq.selectColumns = selectedColumns;
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/sql_source.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/sql_source.ts
index ea31cf9..4e25f2a 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/sql_source.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/sql_source.ts
@@ -24,10 +24,7 @@
 import {Editor} from '../../../../widgets/editor';
 import {Icon} from '../../../../widgets/icon';
 import {Icons} from '../../../../base/semantic_icons';
-import {
-  createFiltersProto,
-  createGroupByProto,
-} from '../operations/operation_component';
+import {createFiltersProto} from '../operations/operation_component';
 import {
   QueryHistoryComponent,
   queryHistoryStorage,
@@ -48,6 +45,7 @@
   constructor(attrs: SqlSourceState) {
     super(attrs);
     this.state = attrs;
+    this.nextNodes = [];
   }
 
   get type() {
@@ -68,9 +66,7 @@
       sql: this.state.sql,
       onExecute: this.state.onExecute,
       sourceCols: newColumnInfoList(this.sourceCols),
-      groupByColumns: [],
       filters: [],
-      aggregations: [],
       customTitle: this.state.customTitle,
       trace: this.state.trace,
     };
@@ -99,11 +95,6 @@
       this.sourceCols,
     );
     if (filtersProto) sq.filters = filtersProto;
-    const groupByProto = createGroupByProto(
-      this.state.groupByColumns,
-      this.state.aggregations,
-    );
-    if (groupByProto) sq.groupBy = groupByProto;
 
     const selectedColumns = createSelectColumnsProto(this);
     if (selectedColumns) sq.selectColumns = selectedColumns;
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/table_source.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/table_source.ts
index ff4a107..b10e018 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/table_source.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sources/table_source.ts
@@ -32,10 +32,7 @@
 import {TextParagraph} from '../../../../widgets/text_paragraph';
 import {Button} from '../../../../widgets/button';
 import {Trace} from '../../../../public/trace';
-import {
-  createFiltersProto,
-  createGroupByProto,
-} from '../operations/operation_component';
+import {createFiltersProto} from '../operations/operation_component';
 import {closeModal, showModal} from '../../../../widgets/modal';
 import {TableList} from '../table_list';
 import {redrawModal} from '../../../../widgets/modal';
@@ -52,7 +49,6 @@
 interface TableSelectionResult {
   sqlTable: SqlTable;
   sourceCols: ColumnInfo[];
-  groupByColumns: ColumnInfo[];
 }
 
 export function modalForTableSelection(
@@ -77,8 +73,7 @@
               const sourceCols = sqlTable.columns.map((c) =>
                 columnInfoFromSqlColumn(c, true),
               );
-              const groupByColumns = newColumnInfoList(sourceCols, false);
-              resolve({sqlTable, sourceCols, groupByColumns});
+              resolve({sqlTable, sourceCols});
               closeModal();
             },
             searchQuery,
@@ -105,9 +100,6 @@
     this.state.onchange = attrs.onchange;
 
     this.state.filters = attrs.filters ?? [];
-    this.state.groupByColumns =
-      attrs.groupByColumns ?? newColumnInfoList(this.sourceCols, false);
-    this.state.aggregations = attrs.aggregations ?? [];
   }
 
   get type() {
@@ -120,9 +112,7 @@
       sqlModules: this.state.sqlModules,
       sqlTable: this.state.sqlTable,
       sourceCols: newColumnInfoList(this.sourceCols),
-      groupByColumns: newColumnInfoList(this.state.groupByColumns),
       filters: this.state.filters.map((f) => ({...f})),
-      aggregations: this.state.aggregations.map((a) => ({...a})),
       customTitle: this.state.customTitle,
       onchange: this.state.onchange,
     };
@@ -201,11 +191,6 @@
       this.sourceCols,
     );
     if (filtersProto) sq.filters = filtersProto;
-    const groupByProto = createGroupByProto(
-      this.state.groupByColumns,
-      this.state.aggregations,
-    );
-    if (groupByProto) sq.groupBy = groupByProto;
 
     const selectedColumns = createSelectColumnsProto(this);
     if (selectedColumns) sq.selectColumns = selectedColumns;
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sub_query_node.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sub_query_node.ts
index 709c73b..af28701 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sub_query_node.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_builder/sub_query_node.ts
@@ -22,17 +22,14 @@
   QueryNode,
   QueryNodeState,
 } from '../query_node';
-import {
-  createFiltersProto,
-  createGroupByProto,
-} from './operations/operation_component';
+import {createFiltersProto} from './operations/operation_component';
 import {ColumnInfo, newColumnInfoList} from './column_info';
 
 export class SubQueryNode implements QueryNode {
   readonly nodeId: string;
   readonly type = NodeType.kSubQuery;
   readonly prevNode?: QueryNode;
-  readonly nextNode?: QueryNode;
+  nextNodes: QueryNode[];
   readonly sourceCols: ColumnInfo[];
   readonly finalCols: ColumnInfo[];
   readonly state: QueryNodeState;
@@ -43,6 +40,7 @@
     this.prevNode = state.prevNode;
     this.sourceCols = this.prevNode!.finalCols;
     this.finalCols = createFinalColumns(this);
+    this.nextNodes = [];
   }
 
   validate(): boolean {
@@ -61,9 +59,7 @@
     const stateCopy: QueryNodeState = {
       prevNode: this.state.prevNode,
       sourceCols: newColumnInfoList(this.sourceCols),
-      groupByColumns: newColumnInfoList(this.state.groupByColumns),
       filters: this.state.filters.map((f) => ({...f})),
-      aggregations: this.state.aggregations.map((a) => ({...a})),
       customTitle: this.state.customTitle,
       onchange: this.state.onchange,
     };
@@ -82,11 +78,6 @@
       this.sourceCols,
     );
     if (filtersProto) sq.filters = filtersProto;
-    const groupByProto = createGroupByProto(
-      this.state.groupByColumns,
-      this.state.aggregations,
-    );
-    if (groupByProto) sq.groupBy = groupByProto;
 
     const selectedColumns = createSelectColumnsProto(this);
     if (selectedColumns) sq.selectColumns = selectedColumns;
diff --git a/ui/src/plugins/dev.perfetto.ExplorePage/query_node.ts b/ui/src/plugins/dev.perfetto.ExplorePage/query_node.ts
index 768c359..939aeb3 100644
--- a/ui/src/plugins/dev.perfetto.ExplorePage/query_node.ts
+++ b/ui/src/plugins/dev.perfetto.ExplorePage/query_node.ts
@@ -14,15 +14,7 @@
 
 import protos from '../../protos';
 import m from 'mithril';
-import {
-  ColumnInfo,
-  columnInfoFromName,
-  newColumnInfoList,
-} from './query_builder/column_info';
-import {
-  Aggregation,
-  placeholderNewColumnName,
-} from './query_builder/operations/aggregations';
+import {ColumnInfo, newColumnInfoList} from './query_builder/column_info';
 import {FilterDefinition} from '../../components/widgets/data_grid/common';
 import {Engine} from '../../trace_processor/engine';
 
@@ -39,6 +31,7 @@
 
   // Single node operations
   kSubQuery,
+  kAggregation,
 }
 
 // All information required to create a new node.
@@ -49,8 +42,6 @@
 
   // Operations
   filters: FilterDefinition[];
-  groupByColumns: ColumnInfo[];
-  aggregations: Aggregation[];
 
   // Errors
   queryError?: Error;
@@ -69,7 +60,7 @@
   readonly graphTableName?: string;
   readonly type: NodeType;
   readonly prevNode?: QueryNode;
-  readonly nextNode?: QueryNode;
+  nextNodes: QueryNode[];
 
   // Columns that are available in the source data.
   readonly sourceCols: ColumnInfo[];
@@ -114,16 +105,6 @@
 }
 
 export function createFinalColumns(node: QueryNode) {
-  if (node.state.groupByColumns.find((c) => c.checked)) {
-    const selected = node.state.groupByColumns.filter((c) => c.checked);
-    for (const agg of node.state.aggregations) {
-      selected.push(
-        columnInfoFromName(agg.newColumnName ?? placeholderNewColumnName(agg)),
-      );
-    }
-    return newColumnInfoList(selected, true);
-  }
-
   return newColumnInfoList(node.sourceCols, true);
 }
 
@@ -221,7 +202,11 @@
       break;
     }
     curr.state.hasOperationChanged = true;
-    curr = curr.nextNode;
+    const queue: QueryNode[] = [];
+    curr.nextNodes.forEach((child) => {
+      queue.push(child);
+    });
+    curr = queue.shift();
   }
 }