[Pivot Table] Improve UI design of pivot tables

Improved the UI flow of pivot tables, added a pop up window to edit
pivot and aggregation selection, drag and drop to reorder columns in the
pop up window and table, and an option to sort aggregations from the
table.
Next Step: Modify query generation and response to support expandable
tables.

Bug: 149193812
Change-Id: Ifeef60cfbe61effa8534b321cc6b02e841867a45

15 files changed, 1219 insertions, 403 deletions

diff --git a/ui/src/assets/common.scss b/ui/src/assets/common.scss
index 0a12f9c69..e99453ad4 100644
--- a/ui/src/assets/common.scss
+++ b/ui/src/assets/common.scss
@@ -207,15 +207,24 @@ button.close {
 
 .query-table-container {
   width: 100%;
-  overflow-x: auto;
 }
 
 .query-table {
     width: 100%;
-    border-collapse: collapse;
     font-size: 14px;
     border: 0;
+    thead.pivot-table-header {
+      cursor: pointer;
+      td.drop-location {
+        background-color: hsla(210, 38%, 95%, 1);
+      }
+      .disabled {
+        cursor: default;
+      }
+    }
     thead td {
+        position: sticky;
+        top: 0;
         background-color: hsl(214, 22%, 90%);
         color: #262f3c;
         text-align: center;
@@ -627,3 +636,110 @@ button.query-ctrl {
 .disallow-selection {
   user-select: none;
 }
+
+.pivot-table-editor-container {
+  font: inherit;
+  width: 670px;
+  height: 420px;
+  h2 {
+    font-weight: bold;
+    text-align: left;
+  }
+  label {
+    cursor: pointer;
+  }
+  select {
+    font-weight: 100;
+    margin: 3px;
+    color: #333;
+    font-size: 15px;
+    align-items: center;
+    cursor: pointer;
+  }
+  span:nth-of-type(2) {
+    margin-left: 1rem;
+  }
+  section.table-group {
+    display: table-row;
+    table {
+      margin: 15px;
+      td {
+        width: 300px;
+        font-size: 17px;
+        cursor: pointer;
+        &.drop-location {
+          background-color: hsla(0, 0%, 85%, 1);
+        }
+      }
+      th {
+        text-align: center;
+        width: 300px;
+        border-bottom: 1px solid rgba(60, 76, 92, 0.4);
+      }
+      &:first-child {
+        float: left;
+      }
+      &:last-child {
+        float: right;
+      }
+    }
+  }
+  .scroll {
+    height: 150px;
+    overflow: auto;
+  }
+  section.button-group {
+    text-align: center;
+    button {
+      background-color: #262f3c;
+      color: #fff;
+      border-radius: 10px;
+      padding: 2px 5px;
+      font-weight: bold;
+      font-size: 13px;
+      min-width: 7em;
+      margin-right: 1rem;
+    }
+  }
+  section {
+    margin: 1rem;
+  }
+}
+
+.pivot-table-tab {
+  button {
+    background: #262f3c;
+    color: white;
+    border-radius: 10px;
+    font-size: 12px;
+    height: 20px;
+    line-height: 18px;
+    min-width: 7em;
+    margin: 0.2rem;
+    &:disabled {
+      opacity: 0.75;
+      cursor: default;
+    }
+  }
+  span {
+    user-select: text;
+    flex-grow: 1;
+  }
+}
+
+.spinner {
+  display: inline-block;
+  vertical-align: middle;
+  box-sizing: border-box;
+  width: 18px;
+  height: 18px;
+  margin-left: 10px;
+  border-radius: 50%;
+  border: 2px solid #408ee0;
+  border-color: #408ee0 transparent #408ee0 transparent;
+  animation: spinner 1.25s linear infinite;
+  @keyframes spinner {
+    0% { transform: rotate(0deg); }
+    100% { transform: rotate(360deg); }
+  }
+}
diff --git a/ui/src/common/actions.ts b/ui/src/common/actions.ts
index 5d5efd486..a7000820b 100644
--- a/ui/src/common/actions.ts
+++ b/ui/src/common/actions.ts
@@ -30,7 +30,7 @@ import {
 import {PROCESS_SUMMARY_TRACK} from '../tracks/process_summary/common';
 
 import {DEFAULT_VIEWING_OPTION} from './flamegraph_util';
-import {AggregationAttrs, PivotAttrs} from './pivot_table_query_generator';
+import {AggregationAttrs, PivotAttrs, TableAttrs} from './pivot_table_data';
 import {
   AdbRecordingTarget,
   Area,
@@ -118,72 +118,6 @@ function rankIndex<T>(element: T, array: T[]): number {
   return index;
 }
 
-function getSelectedPivotTableColumnAttrs(
-    state: StateDraft, args: {pivotTableId: string}):
-    {aggregation?: string, tableName: string, columnName: string} {
-  const availableColumns = state.pivotTableConfig.availableColumns;
-  const columnCount = state.pivotTableConfig.totalColumnsCount;
-  let columnIdx = state.pivotTable[args.pivotTableId].selectedColumnIndex;
-  if (!availableColumns || !columnCount) {
-    throw Error('No columns available');
-  }
-  if (columnIdx === undefined) {
-    throw Error('No column selected');
-  }
-  let tableName, columnName;
-  // Finds column index relative to its table.
-  for (let i = 0; i < availableColumns.length; ++i) {
-    if (columnIdx >= availableColumns[i].columns.length) {
-      columnIdx -= availableColumns[i].columns.length;
-      continue;
-    }
-    tableName = availableColumns[i].tableName;
-    columnName = availableColumns[i].columns[columnIdx];
-    break;
-  }
-  if (tableName === undefined || columnName === undefined) {
-    throw Error('Pivot table requested column does not exist');
-  }
-
-  // Get aggregation if selected column is not a pivot, undefined otherwise.
-  let aggregation;
-  if (state.pivotTable[args.pivotTableId].isPivot === false) {
-    const aggregations = state.pivotTableConfig.availableAggregations;
-    const aggregationIdx =
-        state.pivotTable[args.pivotTableId].selectedAggregationIndex;
-    if (!aggregations) {
-      throw Error('No aggregations available');
-    }
-    if (aggregationIdx === undefined) {
-      throw Error('No aggregation selected');
-    }
-    if (aggregationIdx >= aggregations.length) {
-      throw Error('Pivot table requested aggregation out of bounds');
-    }
-    aggregation = aggregations[aggregationIdx];
-  }
-
-  return {aggregation, tableName, columnName};
-}
-
-function equalTableAttrs(
-    left: PivotAttrs|AggregationAttrs, right: PivotAttrs|AggregationAttrs) {
-  if (left.columnName !== right.columnName) {
-    return false;
-  }
-
-  if (left.tableName !== right.tableName) {
-    return false;
-  }
-
-  if ('aggregation' in left && 'aggregation' in right) {
-    if (left.aggregation !== right.aggregation) {
-      return false;
-    }
-  }
-  return true;
-}
-
 export const StateActions = {
 
   openTraceFromFile(state: StateDraft, args: {file: File}): void {
@@ -915,7 +849,7 @@ export const StateActions = {
       name: args.name,
       selectedPivots: args.selectedPivots,
       selectedAggregations: args.selectedAggregations,
-      isPivot: true,
+      isLoadingQuery: false,
     };
   },
 
@@ -923,40 +857,6 @@ export const StateActions = {
     delete state.pivotTable[args.pivotTableId];
   },
 
-  // Adds column to selectedPivots or selectedAggregations if it doesn't
-  // already exist, remove otherwise.
-  toggleRequestedPivotTablePivot(
-      state: StateDraft, args: {pivotTableId: string}): void {
-    const {aggregation, tableName, columnName} =
-        getSelectedPivotTableColumnAttrs(
-            state, {pivotTableId: args.pivotTableId});
-    let storage: Array<PivotAttrs|AggregationAttrs>;
-    let attrs: PivotAttrs|AggregationAttrs;
-    if (state.pivotTable[args.pivotTableId].isPivot) {
-      storage = state.pivotTable[args.pivotTableId].selectedPivots;
-      attrs = {tableName, columnName};
-    } else {
-      storage = state.pivotTable[args.pivotTableId].selectedAggregations;
-      attrs = {tableName, columnName, aggregation, order: 'DESC'};
-    }
-
-    // Gets index of requested column in selectedPivots or selectedAggregations
-    // if exists.
-    const index = storage.findIndex(element => equalTableAttrs(element, attrs));
-
-    if (index === -1) {
-      storage.push(attrs);
-    } else {
-      storage.splice(index, 1);
-    }
-  },
-
-  clearPivotTableColumns(state: StateDraft, args: {pivotTableId: string}):
-      void {
-        state.pivotTable[args.pivotTableId].selectedPivots = [];
-        state.pivotTable[args.pivotTableId].selectedAggregations = [];
-      },
-
   resetPivotTableRequest(state: StateDraft, args: {pivotTableId: string}):
       void {
         state.pivotTable[args.pivotTableId].requestedAction = undefined;
@@ -967,39 +867,30 @@ export const StateActions = {
     state.pivotTable[args.pivotTableId].requestedAction = args.action;
   },
 
-  setAvailablePivotTableColumns(state: StateDraft, args: {
-    availableColumns: Array<{tableName: string, columns: string[]}>,
-    totalColumnsCount: number,
-    availableAggregations: string[]
-  }): void {
-    state.pivotTableConfig.availableColumns = args.availableColumns;
-    state.pivotTableConfig.totalColumnsCount = args.totalColumnsCount;
-    state.pivotTableConfig.availableAggregations = args.availableAggregations;
-  },
-
-  // Dictates if the selected indexes refer to a pivot or aggregation.
-  togglePivotSelection(state: StateDraft, args: {pivotTableId: string}): void {
-    state.pivotTable[args.pivotTableId].isPivot =
-        !state.pivotTable[args.pivotTableId].isPivot;
-  },
+  setAvailablePivotTableColumns(
+      state: StateDraft,
+      args: {availableColumns: TableAttrs[], availableAggregations: string[]}):
+      void {
+        state.pivotTableConfig.availableColumns = args.availableColumns;
+        state.pivotTableConfig.availableAggregations =
+            args.availableAggregations;
+      },
 
-  setSelectedPivotTableColumnIndex(
-      state: StateDraft, args: {pivotTableId: string, index: number}): void {
-    if (!state.pivotTableConfig.availableColumns ||
-        !state.pivotTableConfig.totalColumnsCount ||
-        args.index >= state.pivotTableConfig.totalColumnsCount) {
-      throw Error('Column selection out of bounds');
-    }
-    state.pivotTable[args.pivotTableId].selectedColumnIndex = args.index;
+  toggleQueryLoading(state: StateDraft, args: {pivotTableId: string}): void {
+    state.pivotTable[args.pivotTableId].isLoadingQuery =
+        !state.pivotTable[args.pivotTableId].isLoadingQuery;
   },
 
-  setSelectedPivotTableAggregationIndex(
-      state: StateDraft, args: {pivotTableId: string, index: number}): void {
-    if (!state.pivotTableConfig.availableAggregations ||
-        args.index >= state.pivotTableConfig.availableAggregations.length) {
-      throw Error('Aggregation selection out of bounds');
-    }
-    state.pivotTable[args.pivotTableId].selectedAggregationIndex = args.index;
+  setSelectedPivotsAndAggregations(state: StateDraft, args: {
+    pivotTableId: string,
+    selectedPivots: PivotAttrs[],
+    selectedAggregations: AggregationAttrs[]
+  }) {
+    state.pivotTable[args.pivotTableId].selectedPivots =
+        args.selectedPivots.map(pivot => Object.assign({}, pivot));
+    state.pivotTable[args.pivotTableId].selectedAggregations =
+        args.selectedAggregations.map(
+            aggregation => Object.assign({}, aggregation));
   },
 };
 
diff --git a/ui/src/common/pivot_table_data.ts b/ui/src/common/pivot_table_data.ts
new file mode 100644
index 000000000..981012857
--- /dev/null
+++ b/ui/src/common/pivot_table_data.ts
@@ -0,0 +1,51 @@
+// Copyright (C) 2021 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 {Row} from './query_result';
+
+export const AVAILABLE_TABLES = ['slice'];
+export const AVAILABLE_AGGREGATIONS = ['COUNT', 'SUM', 'AVG', 'MIN', 'MAX'];
+
+export interface AggregationAttrs {
+  tableName: string;
+  columnName: string;
+  aggregation: string;
+  order: string;
+}
+
+export interface PivotAttrs {
+  tableName: string;
+  columnName: string;
+}
+
+export interface TableAttrs {
+  tableName: string;
+  columns: string[];
+}
+
+export interface ColumnAttrs {
+  name: string;
+  index: number;
+  tableName: string;
+  columnName: string;
+  aggregation?: string;
+  order?: string;
+}
+
+export interface PivotTableQueryResponse {
+  columns: ColumnAttrs[];
+  rows: Row[];
+  error?: string;
+  durationMs: number;
+}
\ No newline at end of file
diff --git a/ui/src/common/pivot_table_query_generator.ts b/ui/src/common/pivot_table_query_generator.ts
index ba9574527..16dfd4f46 100644
--- a/ui/src/common/pivot_table_query_generator.ts
+++ b/ui/src/common/pivot_table_query_generator.ts
@@ -12,37 +12,36 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-export interface AggregationAttrs {
-  tableName: string;
-  columnName: string;
-  aggregation: string;
-  order: string;
-}
-
-export interface PivotAttrs {
-  tableName: string;
-  columnName: string;
-}
+import {AggregationAttrs, PivotAttrs} from './pivot_table_data';
 
-function getPivotAlias(pivot: PivotAttrs): string {
-  return `${pivot.tableName}_${pivot.columnName}`;
+export function getPivotAlias(pivot: PivotAttrs): string {
+  return `${pivot.tableName} ${pivot.columnName}`;
 }
 
-function getAggregationAlias(
+export function getAggregationAlias(
     aggregation: AggregationAttrs, index?: number): string {
-  let alias = `${aggregation.aggregation}_${aggregation.tableName}_${
-      aggregation.columnName}`;
+  let alias = `${aggregation.tableName} ${aggregation.columnName} (${
+      aggregation.aggregation})`;
   if (index !== undefined) {
-    alias += `_${index}`;
+    alias += ` ${index}`;
   }
   return alias;
 }
 
+export function getSqlPivotAlias(pivot: PivotAttrs): string {
+  return `"${getPivotAlias(pivot)}"`;
+}
+
+export function getSqlAggregationAlias(
+    aggregation: AggregationAttrs, index?: number): string {
+  return `"${getAggregationAlias(aggregation, index)}"`;
+}
+
 function getAliasedPivotColumns(
     pivots: PivotAttrs[], lastIndex: number): string[] {
   const pivotCols = [];
   for (let i = 0; i < lastIndex; ++i) {
-    pivotCols.push(getPivotAlias(pivots[i]));
+    pivotCols.push(getSqlPivotAlias(pivots[i]));
   }
   return pivotCols;
 }
@@ -50,13 +49,13 @@ function getAliasedPivotColumns(
 function getAliasedAggregationColumns(
     pivots: PivotAttrs[], aggregations: AggregationAttrs[]): string[] {
   const aggCols = [];
-  for (let i = 0; i < aggregations.length; ++i) {
+  for (const aggregation of aggregations) {
     if (pivots.length === 0) {
-      aggCols.push(getAggregationAlias(aggregations[i]));
+      aggCols.push(getSqlAggregationAlias(aggregation));
       continue;
     }
     for (let j = 0; j < pivots.length; ++j) {
-      aggCols.push(getAggregationAlias(aggregations[i], j + 1));
+      aggCols.push(getSqlAggregationAlias(aggregation, j + 1));
     }
   }
   return aggCols;
@@ -65,25 +64,25 @@ function getAliasedAggregationColumns(
 export class PivotTableQueryGenerator {
   // Generates a query that selects all pivots and aggregations and joins any
   // tables needed by them together. All pivots are renamed into the format
-  // tableName_columnName and all aggregations are renamed into
-  // aggregation_tableName_columnName (see getPivotAlias or
+  // tableName columnName and all aggregations are renamed into
+  // tableName columnName (aggregation) (see getPivotAlias or
   // getAggregationAlias).
   private generateJoinQuery(
       pivots: PivotAttrs[], aggregations: AggregationAttrs[]): string {
     let joinQuery = 'SELECT\n';
 
     const pivotCols = [];
-    for (let i = 0; i < pivots.length; ++i) {
+    for (const pivot of pivots) {
       pivotCols.push(
-          `${pivots[i].tableName}.${pivots[i].columnName} AS ` +
-          `${getPivotAlias(pivots[i])}`);
+          `${pivot.tableName}.${pivot.columnName} AS ` +
+          `${getSqlPivotAlias(pivot)}`);
     }
 
     const aggCols = [];
-    for (let i = 0; i < aggregations.length; ++i) {
+    for (const aggregation of aggregations) {
       aggCols.push(
-          `${aggregations[i].tableName}.${aggregations[i].columnName} AS ` +
-          `${getAggregationAlias(aggregations[i])}`);
+          `${aggregation.tableName}.${aggregation.columnName} AS ` +
+          `${getSqlAggregationAlias(aggregation)}`);
     }
 
     joinQuery += pivotCols.concat(aggCols).join(',\n  ');
@@ -105,14 +104,14 @@ export class PivotTableQueryGenerator {
     const pivotCols = getAliasedPivotColumns(pivots, pivots.length);
 
     const aggCols = [];
-    for (let i = 0; i < aggregations.length; ++i) {
-      const aggColPrefix = `${aggregations[i].aggregation}(${
-          getAggregationAlias(aggregations[i])})`;
+    for (const aggregation of aggregations) {
+      const aggColPrefix =
+          `${aggregation.aggregation}(${getSqlAggregationAlias(aggregation)})`;
 
       if (pivots.length === 0) {
         // Don't partition over pivots if there are no pivots.
         aggCols.push(
-            `${aggColPrefix} AS ${getAggregationAlias(aggregations[i])}`);
+            `${aggColPrefix} AS ${getSqlAggregationAlias(aggregation)}`);
         continue;
       }
 
@@ -120,7 +119,7 @@ export class PivotTableQueryGenerator {
         aggCols.push(
             `${aggColPrefix} OVER (PARTITION BY ` +
             `${getAliasedPivotColumns(pivots, j + 1).join(',  ')}) AS ` +
-            `${getAggregationAlias(aggregations[i], j + 1)}`);
+            `${getSqlAggregationAlias(aggregation, j + 1)}`);
       }
     }
 
diff --git a/ui/src/common/pivot_table_query_generator_unittest.ts b/ui/src/common/pivot_table_query_generator_unittest.ts
index 7d23643ca..b3760238e 100644
--- a/ui/src/common/pivot_table_query_generator_unittest.ts
+++ b/ui/src/common/pivot_table_query_generator_unittest.ts
@@ -15,8 +15,8 @@
 import {
   AggregationAttrs,
   PivotAttrs,
-  PivotTableQueryGenerator
-} from './pivot_table_query_generator';
+} from './pivot_table_data';
+import {PivotTableQueryGenerator} from './pivot_table_query_generator';
 
 test('Generate query with pivots and aggregations', () => {
   const pivotTableQueryGenerator = new PivotTableQueryGenerator();
@@ -28,23 +28,23 @@ test('Generate query with pivots and aggregations', () => {
     {aggregation: 'SUM', tableName: 'slice', columnName: 'dur', order: 'DESC'}
   ];
   const expectedQuery = '\nSELECT\n' +
-      'slice_type,\n' +
-      '  slice_id,\n' +
-      '  SUM_slice_dur_1,\n' +
-      '  SUM_slice_dur_2\n' +
+      '"slice type",\n' +
+      '  "slice id",\n' +
+      '  "slice dur (SUM) 1",\n' +
+      '  "slice dur (SUM) 2"\n' +
       'FROM (\n' +
       'SELECT\n' +
-      'slice_type,\n' +
-      '  slice_id,\n' +
-      '  SUM(SUM_slice_dur) OVER (PARTITION BY slice_type) AS SUM_slice_dur_1' +
-      ',\n' +
-      '  SUM(SUM_slice_dur) OVER (PARTITION BY slice_type,  slice_id)' +
-      ' AS SUM_slice_dur_2\n' +
+      '"slice type",\n' +
+      '  "slice id",\n' +
+      '  SUM("slice dur (SUM)") OVER (PARTITION BY "slice type")' +
+      ' AS "slice dur (SUM) 1",\n' +
+      '  SUM("slice dur (SUM)") OVER (PARTITION BY "slice type",  "slice id")' +
+      ' AS "slice dur (SUM) 2"\n' +
       'FROM (\n' +
       'SELECT\n' +
-      'slice.type AS slice_type,\n' +
-      '  slice.id AS slice_id,\n' +
-      '  slice.dur AS SUM_slice_dur\n' +
+      'slice.type AS "slice type",\n' +
+      '  slice.id AS "slice id",\n' +
+      '  slice.dur AS "slice dur (SUM)"\n' +
       'FROM slice WHERE slice.dur != -1\n' +
       ')\n' +
       ')\n' +
@@ -63,12 +63,12 @@ test('Generate query with pivots', () => {
   ];
   const selectedAggregations: AggregationAttrs[] = [];
   const expectedQuery = '\nSELECT\n' +
-      'slice_type,\n' +
-      '  slice_id\n' +
+      '"slice type",\n' +
+      '  "slice id"\n' +
       'FROM (\n' +
       'SELECT\n' +
-      'slice.type AS slice_type,\n' +
-      '  slice.id AS slice_id\n' +
+      'slice.type AS "slice type",\n' +
+      '  slice.id AS "slice id"\n' +
       'FROM slice WHERE slice.dur != -1\n' +
       ')\n' +
       'GROUP BY 1,  2\n';
@@ -85,16 +85,16 @@ test('Generate query with aggregations', () => {
     {aggregation: 'MAX', tableName: 'slice', columnName: 'dur', order: 'ASC'}
   ];
   const expectedQuery = '\nSELECT\n' +
-      'SUM_slice_dur,\n' +
-      '  MAX_slice_dur\n' +
+      '"slice dur (SUM)",\n' +
+      '  "slice dur (MAX)"\n' +
       'FROM (\n' +
       'SELECT\n' +
-      'SUM(SUM_slice_dur) AS SUM_slice_dur,\n' +
-      '  MAX(MAX_slice_dur) AS MAX_slice_dur\n' +
+      'SUM("slice dur (SUM)") AS "slice dur (SUM)",\n' +
+      '  MAX("slice dur (MAX)") AS "slice dur (MAX)"\n' +
       'FROM (\n' +
       'SELECT\n' +
-      'slice.dur AS SUM_slice_dur,\n' +
-      '  slice.dur AS MAX_slice_dur\n' +
+      'slice.dur AS "slice dur (SUM)",\n' +
+      '  slice.dur AS "slice dur (MAX)"\n' +
       'FROM slice WHERE slice.dur != -1\n' +
       ')\n' +
       ')\n' +
diff --git a/ui/src/common/queries.ts b/ui/src/common/queries.ts
index cd7273777..1e7e67170 100644
--- a/ui/src/common/queries.ts
+++ b/ui/src/common/queries.ts
@@ -13,7 +13,6 @@
 // limitations under the License.
 
 import {Engine} from './engine';
-
 import {Row} from './query_result';
 
 const MAX_DISPLAY_ROWS = 10000;
diff --git a/ui/src/common/state.ts b/ui/src/common/state.ts
index 065ba2fc9..2935d2d6c 100644
--- a/ui/src/common/state.ts
+++ b/ui/src/common/state.ts
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import {AggregationAttrs, PivotAttrs} from './pivot_table_query_generator';
+import {AggregationAttrs, PivotAttrs, TableAttrs} from './pivot_table_data';
 
 /**
  * A plain js object, holding objects of type |Class| keyed by string id.
@@ -283,11 +283,7 @@ export interface MetricsState {
 }
 
 export interface PivotTableConfig {
-  availableColumns?: Array<{
-    tableName: string,
-    columns: string[]
-  }>;                                // Undefined until list is loaded.
-  totalColumnsCount?: number;        // Total columns in all tables.
+  availableColumns?: TableAttrs[];   // Undefined until list is loaded.
   availableAggregations?: string[];  // Undefined until list is loaded.
 }
 
@@ -296,11 +292,9 @@ export interface PivotTableState {
   name: string;
   selectedPivots: PivotAttrs[];
   selectedAggregations: AggregationAttrs[];
-  selectedColumnIndex?: number;
-  selectedAggregationIndex?: number;
-  isPivot: boolean;
   requestedAction?:
       string;  // Unset after pivot table column request is handled.
+  isLoadingQuery: boolean;
 }
 
 export interface State {
diff --git a/ui/src/controller/pivot_table_controller.ts b/ui/src/controller/pivot_table_controller.ts
index 8412f0ade..4643d0a43 100644
--- a/ui/src/controller/pivot_table_controller.ts
+++ b/ui/src/controller/pivot_table_controller.ts
@@ -14,26 +14,78 @@
 
 import {Actions} from '../common/actions';
 import {Engine} from '../common/engine';
-import {PivotTableQueryGenerator} from '../common/pivot_table_query_generator';
-import {runQuery} from '../common/queries';
-import {publishQueryResult} from '../frontend/publish';
+import {
+  AVAILABLE_AGGREGATIONS,
+  AVAILABLE_TABLES,
+  PivotTableQueryResponse,
+} from '../common/pivot_table_data';
+import {
+  getAggregationAlias,
+  getPivotAlias,
+  PivotTableQueryGenerator
+} from '../common/pivot_table_query_generator';
+import {
+  QueryResponse,
+  runQuery,
+} from '../common/queries';
+import {PivotTableHelper} from '../frontend/pivot_table_helper';
+import {publishPivotTableHelper, publishQueryResult} from '../frontend/publish';
 
 import {Controller} from './controller';
 import {globals} from './globals';
 
-const AVAILABLE_TABLES = ['slice'];
-const AVAILABLE_AGGREGATIONS = ['COUNT', 'SUM', 'AVG', 'MIN', 'MAX'];
-
 export interface PivotTableControllerArgs {
   pivotTableId: string;
   engine: Engine;
 }
 
+function getPivotTableQueryResponse(
+    pivotTableId: string, queryResp: QueryResponse): PivotTableQueryResponse {
+  const columns = [];
+  const pivotTable = globals.state.pivotTable[pivotTableId];
+  for (const column of queryResp.columns) {
+    let isPivot = false;
+    let index = pivotTable.selectedAggregations.findIndex(
+        element => column.startsWith(getAggregationAlias(element)));
+    if (index === -1) {
+      isPivot = true;
+      index = pivotTable.selectedPivots.findIndex(
+          element => column.startsWith(getPivotAlias(element)));
+    }
+    if (index === -1) {
+      throw Error(
+          'Column in query response not in selectedAggregations or ' +
+          'selectedPivots.');
+    }
+    let tableName;
+    let columnName;
+    let aggregation;
+    let order;
+    if (isPivot) {
+      tableName = pivotTable.selectedPivots[index].tableName;
+      columnName = pivotTable.selectedPivots[index].columnName;
+    } else {
+      tableName = pivotTable.selectedAggregations[index].tableName;
+      columnName = pivotTable.selectedAggregations[index].columnName;
+      aggregation = pivotTable.selectedAggregations[index].aggregation;
+      order = pivotTable.selectedAggregations[index].order;
+    }
+    columns.push(
+        {name: column, index, tableName, columnName, aggregation, order});
+  }
+
+  return {
+    columns,
+    rows: queryResp.rows,
+    error: queryResp.error,
+    durationMs: queryResp.durationMs
+  };
+}
+
 export class PivotTableController extends Controller<'main'> {
+  private pivotTableId: string;
   private pivotTableQueryGenerator = new PivotTableQueryGenerator();
   private engine: Engine;
-  private pivotTableId: string;
-  private availableColumns: Array<{tableName: string, columns: string[]}> = [];
   private previousQuery = '';
 
   constructor(args: PivotTableControllerArgs) {
@@ -50,13 +102,7 @@ export class PivotTableController extends Controller<'main'> {
     if (!requestedAction) return;
     globals.dispatch(
         Actions.resetPivotTableRequest({pivotTableId: this.pivotTableId}));
-
     switch (requestedAction) {
-      case 'UPDATE':
-        globals.dispatch(Actions.toggleRequestedPivotTablePivot(
-            {pivotTableId: this.pivotTableId}));
-        break;
-
       case 'QUERY':
         // Generates and executes new query based on selectedPivots and
         // selectedAggregations.
@@ -65,9 +111,16 @@ export class PivotTableController extends Controller<'main'> {
             pivotTable.selectedPivots, pivotTable.selectedAggregations);
         if (query === this.previousQuery) break;
         if (query !== '') {
+          globals.dispatch(
+              Actions.toggleQueryLoading({pivotTableId: this.pivotTableId}));
           runQuery(this.pivotTableId, query, this.engine).then(resp => {
             console.log(`Query ${query} took ${resp.durationMs} ms`);
-            publishQueryResult({id: this.pivotTableId, data: resp});
+            publishQueryResult({
+              id: this.pivotTableId,
+              data: getPivotTableQueryResponse(this.pivotTableId, resp)
+            });
+            globals.dispatch(
+                Actions.toggleQueryLoading({pivotTableId: this.pivotTableId}));
           });
         } else {
           publishQueryResult({id: this.pivotTableId, data: undefined});
@@ -76,27 +129,37 @@ export class PivotTableController extends Controller<'main'> {
         break;
 
       default:
-        throw new Error(`Unexpected state ${this.state}`);
+        throw new Error(`Unexpected requested action ${requestedAction}`);
     }
   }
 
   private async setup(): Promise<void> {
+    const pivotTable = globals.state.pivotTable[this.pivotTableId];
+    const selectedPivots = pivotTable.selectedPivots;
+    const selectedAggregations = pivotTable.selectedAggregations;
+    let availableColumns = globals.state.pivotTableConfig.availableColumns;
     // No need to retrieve table columns if they are already stored.
     // Only needed when first pivot table is created.
-    if (globals.state.pivotTableConfig.availableColumns !== undefined) return;
-    let totalColumnsCount = 0;
-    for (const table of AVAILABLE_TABLES) {
-      const columns = await this.getColumnsForTable(table);
-      totalColumnsCount += columns.length;
-      if (columns.length > 0) {
-        this.availableColumns.push({tableName: table, columns});
+    if (availableColumns === undefined) {
+      availableColumns = [];
+      for (const table of AVAILABLE_TABLES) {
+        const columns = await this.getColumnsForTable(table);
+        if (columns.length > 0) {
+          availableColumns.push({tableName: table, columns});
+        }
       }
+      globals.dispatch(Actions.setAvailablePivotTableColumns(
+          {availableColumns, availableAggregations: AVAILABLE_AGGREGATIONS}));
     }
-    globals.dispatch(Actions.setAvailablePivotTableColumns({
-      availableColumns: this.availableColumns,
-      totalColumnsCount,
-      availableAggregations: AVAILABLE_AGGREGATIONS
-    }));
+    publishPivotTableHelper({
+      id: this.pivotTableId,
+      data: new PivotTableHelper(
+          this.pivotTableId,
+          availableColumns,
+          AVAILABLE_AGGREGATIONS,
+          selectedPivots,
+          selectedAggregations)
+    });
   }
 
   private async getColumnsForTable(tableName: string): Promise<string[]> {
diff --git a/ui/src/frontend/details_panel.ts b/ui/src/frontend/details_panel.ts
index 75c873de8..4424d10d9 100644
--- a/ui/src/frontend/details_panel.ts
+++ b/ui/src/frontend/details_panel.ts
@@ -30,9 +30,11 @@ import {
 import {globals} from './globals';
 import {HeapProfileDetailsPanel} from './heap_profile_panel';
 import {LogPanel} from './logs_panel';
+import {showModal} from './modal';
 import {NotesEditorPanel} from './notes_panel';
 import {AnyAttrsVnode, PanelContainer} from './panel_container';
 import {PivotTable} from './pivot_table';
+import {ColumnDisplay, ColumnPicker} from './pivot_table_editor';
 import {QueryTable} from './query_table';
 import {SliceDetailsPanel} from './slice_panel';
 import {ThreadStatePanel} from './thread_state_panel';
@@ -284,12 +286,49 @@ export class DetailsPanel implements m.ClassComponent {
       });
     }
 
-    if (globals.frontendLocalState.showPivotTable) {
-      const pivotTableId = 'pivot-table';
+    const pivotTableId = 'pivot-table';
+    const pivotTable = globals.state.pivotTable[pivotTableId];
+    const helper = globals.pivotTableHelper.get(pivotTableId);
+
+    if (globals.frontendLocalState.showPivotTable && pivotTable !== undefined) {
+      if (helper !== undefined) {
+        helper.setSelectedPivotsAndAggregations(
+            pivotTable.selectedPivots, pivotTable.selectedAggregations);
+      }
       detailsPanels.push({
         key: pivotTableId,
-        name: globals.state.pivotTable[pivotTableId].name,
-        vnode: m(PivotTable, {key: pivotTableId, pivotTableId})
+        name: pivotTable.name,
+        vnode: m(PivotTable, {key: pivotTableId, pivotTableId, helper})
+      });
+    }
+
+    if (helper !== undefined && helper.editPivotTableModalOpen) {
+      let content;
+      if (helper.availableColumns.length === 0 ||
+          helper.availableAggregations.length === 0) {
+        content =
+            m('.pivot-table-editor-container',
+              helper.availableColumns.length === 0 ?
+                  m('div', 'No columns available.') :
+                  null,
+              helper.availableAggregations.length === 0 ?
+                  m('div', 'No aggregations available.') :
+                  null);
+      } else {
+        const attrs = {helper};
+        content =
+            m('.pivot-table-editor-container',
+              m(ColumnPicker, attrs),
+              m(ColumnDisplay, attrs));
+      }
+
+      showModal({
+        title: 'Edit Pivot Table',
+        content,
+        buttons: [],
+      }).finally(() => {
+        helper.toggleEditPivotTableModal();
+        globals.rafScheduler.scheduleFullRedraw();
       });
     }
 
diff --git a/ui/src/frontend/globals.ts b/ui/src/frontend/globals.ts
index bc0d8e927..e2967dd79 100644
--- a/ui/src/frontend/globals.ts
+++ b/ui/src/frontend/globals.ts
@@ -27,6 +27,7 @@ import {fromNs, toNs} from '../common/time';
 
 import {Analytics, initAnalytics} from './analytics';
 import {FrontendLocalState} from './frontend_local_state';
+import {PivotTableHelper} from './pivot_table_helper';
 import {RafScheduler} from './raf_scheduler';
 import {Router} from './router';
 import {ServiceWorkerController} from './service_worker_controller';
@@ -34,6 +35,7 @@ import {ServiceWorkerController} from './service_worker_controller';
 type Dispatch = (action: DeferredAction) => void;
 type TrackDataStore = Map<string, {}>;
 type QueryResultsStore = Map<string, {}|undefined>;
+type PivotTableHelperStore = Map<string, PivotTableHelper>;
 type AggregateDataStore = Map<string, AggregateData>;
 type Description = Map<string, string>;
 export interface SliceDetails {
@@ -164,6 +166,7 @@ class Globals {
   // TODO(hjd): Unify trackDataStore, queryResults, overviewStore, threads.
   private _trackDataStore?: TrackDataStore = undefined;
   private _queryResults?: QueryResultsStore = undefined;
+  private _pivotTableHelper?: PivotTableHelperStore = undefined;
   private _overviewStore?: OverviewStore = undefined;
   private _aggregateDataStore?: AggregateDataStore = undefined;
   private _threadMap?: ThreadMap = undefined;
@@ -216,6 +219,7 @@ class Globals {
     // TODO(hjd): Unify trackDataStore, queryResults, overviewStore, threads.
     this._trackDataStore = new Map<string, {}>();
     this._queryResults = new Map<string, {}>();
+    this._pivotTableHelper = new Map<string, PivotTableHelper>();
     this._overviewStore = new Map<string, QuantizedLoad[]>();
     this._aggregateDataStore = new Map<string, AggregateData>();
     this._threadMap = new Map<number, ThreadDesc>();
@@ -282,6 +286,10 @@ class Globals {
     return assertExists(this._queryResults);
   }
 
+  get pivotTableHelper(): PivotTableHelperStore {
+    return assertExists(this._pivotTableHelper);
+  }
+
   get threads() {
     return assertExists(this._threadMap);
   }
@@ -495,6 +503,7 @@ class Globals {
     // TODO(hjd): Unify trackDataStore, queryResults, overviewStore, threads.
     this._trackDataStore = undefined;
     this._queryResults = undefined;
+    this._pivotTableHelper = undefined;
     this._overviewStore = undefined;
     this._threadMap = undefined;
     this._sliceDetails = undefined;
diff --git a/ui/src/frontend/pivot_table.ts b/ui/src/frontend/pivot_table.ts
index 7d441d2f8..ec23d5a7a 100644
--- a/ui/src/frontend/pivot_table.ts
+++ b/ui/src/frontend/pivot_table.ts
@@ -12,224 +12,162 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-
 import * as m from 'mithril';
 
 import {Actions} from '../common/actions';
-import {QueryResponse} from '../common/queries';
+import {
+  ColumnAttrs,
+  PivotTableQueryResponse,
+} from '../common/pivot_table_data';
 import {Row} from '../common/query_result';
 
-import {queryResponseToClipboard} from './clipboard';
 import {globals} from './globals';
 import {Panel} from './panel';
+import {
+  PivotTableHelper,
+} from './pivot_table_helper';
 
 interface PivotTableRowAttrs {
   row: Row;
-  columns: string[];
+  columns: ColumnAttrs[];
 }
 
-class PivotTableRow implements m.ClassComponent<PivotTableRowAttrs> {
-  view(vnode: m.Vnode<PivotTableRowAttrs>) {
-    const cells = [];
-    const {row, columns} = vnode.attrs;
-    for (const col of columns) {
-      cells.push(m('td', row[col]));
-    }
-
-    return m('tr', cells);
-  }
+interface PivotTableHeaderAttrs {
+  helper: PivotTableHelper;
 }
 
 interface PivotTableAttrs {
   pivotTableId: string;
+  helper?: PivotTableHelper;
 }
 
-class ColumnPicker implements m.ClassComponent<PivotTableAttrs> {
-  view(vnode: m.Vnode<PivotTableAttrs>) {
-    const {pivotTableId} = vnode.attrs;
-    const availableColumns = globals.state.pivotTableConfig.availableColumns;
-    const availableColumnsCount =
-        globals.state.pivotTableConfig.totalColumnsCount;
-    const availableAggregations =
-        globals.state.pivotTableConfig.availableAggregations;
-    if (availableColumns === undefined || availableColumnsCount === undefined) {
-      return 'Loading columns...';
-    }
-    if (availableAggregations === undefined) {
-      return 'Loading aggregations...';
-    }
-    if (availableColumnsCount === 0) {
-      return 'No columns available';
-    }
-    if (availableAggregations.length === 0) {
-      return 'No aggregations available';
-    }
-
-    if (globals.state.pivotTable[pivotTableId].selectedColumnIndex ===
-        undefined) {
-      globals.state.pivotTable[pivotTableId].selectedColumnIndex = 0;
-    }
-    if (globals.state.pivotTable[pivotTableId].selectedAggregationIndex ===
-        undefined) {
-      globals.state.pivotTable[pivotTableId].selectedAggregationIndex = 0;
-    }
+class PivotTableHeader implements m.ClassComponent<PivotTableHeaderAttrs> {
+  view(vnode: m.Vnode<PivotTableHeaderAttrs>) {
+    const {helper} = vnode.attrs;
+    const pivotTableId = helper.pivotTableId;
+    const pivotTable = globals.state.pivotTable[pivotTableId];
+    const resp =
+        globals.queryResults.get(pivotTableId) as PivotTableQueryResponse;
 
-    // Fills available aggregations options in aggregation select.
-    const aggregationOptions = [];
-    for (let i = 0; i < availableAggregations.length; ++i) {
-      aggregationOptions.push(
-          m('option',
-            {value: availableAggregations[i], key: availableAggregations[i]},
-            availableAggregations[i]));
+    const cols = [];
+    for (const column of resp.columns) {
+      const isPivot = column.aggregation === undefined;
+      let sortIcon;
+      if (!isPivot) {
+        sortIcon =
+            column.order === 'DESC' ? 'arrow_drop_down' : 'arrow_drop_up';
+      }
+      cols.push(m(
+          'td',
+          {
+            class: pivotTable.isLoadingQuery ? 'disabled' : '',
+            draggable: pivotTable.isLoadingQuery ? false : true,
+            ondragstart: (e: DragEvent) => {
+              helper.selectedColumnOnDrag(e, isPivot, column.index);
+            },
+            ondrop: (e: DragEvent) => {
+              helper.removeHighlightFromDropLocation(e);
+              helper.selectedColumnOnDrop(e, isPivot, column.index);
+              helper.queryPivotTableChanges();
+            },
+            ondragenter: (e: DragEvent) => {
+              helper.highlightDropLocation(e, isPivot);
+            },
+            ondragleave: (e: DragEvent) => {
+              helper.removeHighlightFromDropLocation(e);
+            }
+          },
+          column.name,
+          (!isPivot && sortIcon !== undefined ?
+               m('i.material-icons',
+                 {
+                   onclick: () => {
+                     if (!pivotTable.isLoadingQuery) {
+                       helper.togglePivotTableAggregationSorting(column.index);
+                       helper.queryPivotTableChanges();
+                     }
+                   }
+                 },
+                 sortIcon) :
+               null)));
     }
+    return m('tr', cols);
+  }
+}
 
-    // Fills available columns options divided according to their table in
-    // column select.
-    const columnOptionGroup = [];
-    for (let i = 0; i < availableColumns.length; ++i) {
-      const options = [];
-      for (let j = 0; j < availableColumns[i].columns.length; ++j) {
-        options.push(
-            m('option',
-              {
-                value: availableColumns[i].columns[j],
-                key: availableColumns[i].columns[j]
-              },
-              availableColumns[i].columns[j]));
-      }
-      columnOptionGroup.push(
-          m('optgroup', {label: availableColumns[i].tableName}, options));
+class PivotTableRow implements m.ClassComponent<PivotTableRowAttrs> {
+  view(vnode: m.Vnode<PivotTableRowAttrs>) {
+    const cells = [];
+    const {row, columns} = vnode.attrs;
+    for (const col of columns) {
+      cells.push(m('td', row[col.name]));
     }
 
-    return m('div', [
-      'Select a column: ',
-      // Pivot radio button.
-      m(`input[type=radio][name=type][id=pivot]`, {
-        checked: globals.state.pivotTable[pivotTableId].isPivot,
-        onchange: () =>
-            globals.dispatch(Actions.togglePivotSelection({pivotTableId}))
-      }),
-      m(`label[for=pivot]`, 'Pivot'),
-      // Aggregation radio button.
-      m(`input[type=radio][name=type][id=aggregation]`, {
-        checked: !globals.state.pivotTable[pivotTableId].isPivot,
-        onchange: () =>
-            globals.dispatch(Actions.togglePivotSelection({pivotTableId}))
-      }),
-      m(`label[for=aggregation]`, 'Aggregation'),
-      ' ',
-      // Aggregation select.
-      m('select',
-        {
-          disabled: (globals.state.pivotTable[pivotTableId].isPivot === true),
-          selectedIndex:
-              globals.state.pivotTable[pivotTableId].selectedAggregationIndex,
-          onchange: (e: InputEvent) => {
-            globals.dispatch(Actions.setSelectedPivotTableAggregationIndex({
-              pivotTableId,
-              index: (e.target as HTMLSelectElement).selectedIndex
-            }));
-          }
-        },
-        aggregationOptions),
-      ' ',
-      // Column select.
-      m('select',
-        {
-          selectedIndex:
-              globals.state.pivotTable[pivotTableId].selectedColumnIndex,
-          onchange: (e: InputEvent) => {
-            globals.dispatch(Actions.setSelectedPivotTableColumnIndex({
-              pivotTableId,
-              index: (e.target as HTMLSelectElement).selectedIndex
-            }));
-          }
-        },
-        columnOptionGroup),
-      ' ',
-      // Button to toggle selected column.
-      m('button.query-ctrl',
-        {
-          onclick: () => {
-            globals.dispatch(
-                Actions.setPivotTableRequest({pivotTableId, action: 'UPDATE'}));
-          }
-        },
-        'Add/Remove'),
-      // Button to execute query based on added/removed columns.
-      m('button.query-ctrl',
-        {
-          onclick: () => {
-            globals.dispatch(
-                Actions.setPivotTableRequest({pivotTableId, action: 'QUERY'}));
-          }
-        },
-        'Query'),
-      // Button to clear table and all selected columns.
-      m('button.query-ctrl',
-        {
-          onclick: () => {
-            globals.dispatch(Actions.clearPivotTableColumns({pivotTableId}));
-            globals.dispatch(
-                Actions.setPivotTableRequest({pivotTableId, action: 'QUERY'}));
-          }
-        },
-        'Clear'),
-    ]);
+    return m('tr', cells);
   }
 }
 
 export class PivotTable extends Panel<PivotTableAttrs> {
   view(vnode: m.CVnode<PivotTableAttrs>) {
-    const {pivotTableId} = vnode.attrs;
-    const resp = globals.queryResults.get(pivotTableId) as QueryResponse;
+    const {pivotTableId, helper} = vnode.attrs;
+    const pivotTable = globals.state.pivotTable[pivotTableId];
+    const resp =
+        globals.queryResults.get(pivotTableId) as PivotTableQueryResponse;
+
     // Query resulting from query generator should always be valid.
     if (resp !== undefined && resp.error) {
       throw Error(`Pivot table query resulted in SQL error: ${resp.error}`);
     }
-    const cols = [];
+
     const rows = [];
     let header;
 
-    if (resp !== undefined) {
-      for (const col of resp.columns) {
-        cols.push(m('td', col));
-      }
-      header = m('tr', cols);
+    if (helper !== undefined && resp !== undefined) {
+      header = m(PivotTableHeader, {helper});
 
-      for (let i = 0; i < resp.rows.length; i++) {
-        rows.push(m(PivotTableRow, {row: resp.rows[i], columns: resp.columns}));
+      for (const row of resp.rows) {
+        rows.push(m(PivotTableRow, {row, columns: resp.columns}));
       }
     }
 
     return m(
-        'div',
+        'div.pivot-table-tab',
         m(
             'header.overview',
-            m(
-                'span.code',
-                m(ColumnPicker, {pivotTableId}),
-                ),
-            (resp === undefined || resp.error) ?
-                null :
-                m('button.query-ctrl',
-                  {
-                    onclick: () => {
-                      queryResponseToClipboard(resp);
-                    },
-                  },
-                  'Copy as .tsv'),
-            m('button.query-ctrl',
+            m('span',
+              m('button',
+                {
+                  disabled: helper === undefined || pivotTable.isLoadingQuery,
+                  onclick: () => {
+                    if (helper !== undefined) {
+                      helper.toggleEditPivotTableModal();
+                      globals.rafScheduler.scheduleFullRedraw();
+                    }
+                  }
+                },
+                'Edit'),
+              ' ',
+              (pivotTable.isLoadingQuery ? m('div.spinner') : null),
+              (resp !== undefined && !pivotTable.isLoadingQuery ?
+                   m('span.code',
+                     `Query took ${Math.round(resp.durationMs)} ms`) :
+                   null)),
+            m('button',
               {
+                disabled: helper === undefined || pivotTable.isLoadingQuery,
                 onclick: () => {
                   globals.frontendLocalState.togglePivotTable();
+                  globals.queryResults.delete(pivotTableId);
+                  globals.pivotTableHelper.delete(pivotTableId);
                   globals.dispatch(Actions.deletePivotTable({pivotTableId}));
                 }
               },
               'Close'),
             ),
-        m('query-table-container',
-          m('table.query-table', m('thead', header), m('tbody', rows))));
+        m('.query-table-container',
+          m('table.query-table',
+            m('thead.pivot-table-header', header),
+            m('tbody', rows))));
   }
 
   renderCanvas() {}
diff --git a/ui/src/frontend/pivot_table_editor.ts b/ui/src/frontend/pivot_table_editor.ts
new file mode 100644
index 000000000..37f281dcf
--- /dev/null
+++ b/ui/src/frontend/pivot_table_editor.ts
@@ -0,0 +1,248 @@
+// Copyright (C) 2021 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 * as m from 'mithril';
+
+import {globals} from './globals';
+import {hideModel} from './modal';
+import {
+  PivotTableHelper,
+} from './pivot_table_helper';
+
+interface PivotTableEditorAttrs {
+  helper: PivotTableHelper;
+}
+
+export class ColumnPicker implements m.ClassComponent<PivotTableEditorAttrs> {
+  view(vnode: m.Vnode<PivotTableEditorAttrs>) {
+    const {helper} = vnode.attrs;
+
+    // Fills available aggregations options in aggregation select.
+    const aggregationOptions = [];
+    for (const aggregation of helper.availableAggregations) {
+      aggregationOptions.push(
+          m('option', {value: aggregation, key: aggregation}, aggregation));
+    }
+
+    // Fills available columns options divided according to their table in
+    // column select.
+    const columnOptionGroup = [];
+    for (const {tableName, columns} of helper.availableColumns) {
+      const options = [];
+      for (const column of columns) {
+        options.push(m('option', {value: column, key: column}, column));
+      }
+      columnOptionGroup.push(m('optgroup', {label: tableName}, options));
+    }
+
+    return m(
+        'div',
+        m(
+            'section',
+            m('h2', 'Select column type: '),
+            // Pivot radio button.
+            m(
+                'span',
+                m(`input[type=radio][name=type][id=pivot]`, {
+                  checked: helper.isPivot,
+                  onchange: () => {
+                    helper.togglePivotSelection();
+                    globals.rafScheduler.scheduleFullRedraw();
+                  }
+                }),
+                m(`label[for=pivot]`, 'Pivot'),
+                ),
+            // Aggregation radio button.
+            m('span', m(`input[type=radio][name=type][id=aggregation]`, {
+                checked: !helper.isPivot,
+                onchange: () => {
+                  helper.togglePivotSelection();
+                  globals.rafScheduler.scheduleFullRedraw();
+                }
+              })),
+            m(`label[for=aggregation]`, 'Aggregation'),
+            ),
+        m(
+            'section',
+            m('h2', 'Select a column: '),
+            // Aggregation select.
+            m('select',
+              {
+                disabled: helper.isPivot,
+                selectedIndex: helper.selectedAggregationIndex,
+                onchange: (e: InputEvent) => {
+                  helper.setSelectedPivotTableAggregationIndex(
+                      (e.target as HTMLSelectElement).selectedIndex);
+                }
+              },
+              aggregationOptions),
+            ' ',
+            // Column select.
+            m('select',
+              {
+                selectedIndex: helper.selectedColumnIndex,
+                onchange: (e: InputEvent) => {
+                  helper.setSelectedPivotTableColumnIndex(
+                      (e.target as HTMLSelectElement).selectedIndex);
+                }
+              },
+              columnOptionGroup),
+            ),
+        m('section.button-group',
+          // Button to toggle selected column.
+          m('button',
+            {
+              onclick: () => {
+                helper.updatePivotTableColumnOnSelectedIndex();
+                globals.rafScheduler.scheduleFullRedraw();
+              }
+            },
+            'Add/Remove'),
+          // Button to clear table and all selected columns.
+          m('button',
+            {
+              onclick: () => {
+                helper.clearPivotTableColumns();
+                globals.rafScheduler.scheduleFullRedraw();
+              }
+            },
+            'Clear')));
+  }
+}
+
+export class ColumnDisplay implements m.ClassComponent<PivotTableEditorAttrs> {
+  view(vnode: m.Vnode<PivotTableEditorAttrs>) {
+    const {helper} = vnode.attrs;
+    const selectedPivotsDisplay = [];
+    const selectedAggregationsDisplay = [];
+
+    for (let i = 0; i < helper.selectedPivots.length; ++i) {
+      const columnAttrs = helper.selectedPivots[i];
+      selectedPivotsDisplay.push(m(
+          'tr',
+          m('td',
+            {
+              draggable: true,
+              ondragstart: (e: DragEvent) => {
+                helper.selectedColumnOnDrag(e, true, i);
+              },
+              ondrop: (e: DragEvent) => {
+                helper.removeHighlightFromDropLocation(e);
+                helper.selectedColumnOnDrop(e, true, i);
+                globals.rafScheduler.scheduleFullRedraw();
+              },
+              onclick: () => {
+                helper.selectPivotTableColumn(columnAttrs);
+                globals.rafScheduler.scheduleFullRedraw();
+              },
+              ondragenter: (e: DragEvent) => {
+                helper.highlightDropLocation(e, true);
+              },
+              ondragleave: (e: DragEvent) => {
+                helper.removeHighlightFromDropLocation(e);
+              }
+            },
+            m('i.material-icons',
+              {
+                onclick: () => {
+                  helper.updatePivotTableColumnOnColumnAttributes(columnAttrs);
+                  globals.rafScheduler.scheduleFullRedraw();
+                },
+              },
+              'remove'),
+            ' ',
+            `${columnAttrs.tableName} ${columnAttrs.columnName}`)));
+    }
+
+    for (let i = 0; i < helper.selectedAggregations.length; ++i) {
+      const columnAttrs = helper.selectedAggregations[i];
+      const sortIcon = helper.selectedAggregations[i].order === 'DESC' ?
+          'arrow_drop_down' :
+          'arrow_drop_up';
+      selectedAggregationsDisplay.push(m(
+          'tr',
+          m('td',
+            {
+              draggable: 'true',
+              ondragstart: (e: DragEvent) => {
+                helper.selectedColumnOnDrag(e, false, i);
+              },
+              ondrop: (e: DragEvent) => {
+                helper.removeHighlightFromDropLocation(e);
+                helper.selectedColumnOnDrop(e, false, i);
+                globals.rafScheduler.scheduleFullRedraw();
+              },
+              onclick: () => {
+                helper.selectPivotTableColumn(columnAttrs);
+                globals.rafScheduler.scheduleFullRedraw();
+              },
+              ondragenter: (e: DragEvent) => {
+                helper.highlightDropLocation(e, false);
+              },
+              ondragleave: (e: DragEvent) => {
+                helper.removeHighlightFromDropLocation(e);
+              }
+            },
+            m('i.material-icons',
+              {
+                onclick: () => {
+                  helper.updatePivotTableColumnOnColumnAttributes(columnAttrs);
+                  globals.rafScheduler.scheduleFullRedraw();
+                },
+              },
+              'remove'),
+            ' ',
+            `${columnAttrs.tableName} ${columnAttrs.columnName} (${
+                columnAttrs.aggregation})`,
+            m('i.material-icons',
+              {
+                onclick: () => {
+                  helper.togglePivotTableAggregationSorting(i);
+                  globals.rafScheduler.scheduleFullRedraw();
+                }
+              },
+              sortIcon))));
+    }
+
+    return m(
+        'div',
+        m('section.table-group',
+          // Table that displays selected pivots.
+          m('table',
+            m('thead', m('tr', m('th', 'Selected Pivots'))),
+            m('div.scroll', m('tbody', selectedPivotsDisplay))),
+          // Table that displays selected aggregations.
+          m('table',
+            m('thead', m('tr', m('th', 'Selected Aggregations'))),
+            m('div.scroll', m('tbody', selectedAggregationsDisplay)))),
+        m('section.button-group',
+          // Button to toggle selected column.
+          m('button',
+            {
+              onclick: () => {
+                helper.queryPivotTableChanges();
+                hideModel();
+              }
+            },
+            'Query'),
+          // Button to clear table and all selected columns.
+          m('button',
+            {
+              onclick: () => {
+                hideModel();
+              }
+            },
+            'Cancel')));
+  }
+}
diff --git a/ui/src/frontend/pivot_table_helper.ts b/ui/src/frontend/pivot_table_helper.ts
new file mode 100644
index 000000000..a7f440921
--- /dev/null
+++ b/ui/src/frontend/pivot_table_helper.ts
@@ -0,0 +1,337 @@
+// Copyright (C) 2021 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 {Actions} from '../common/actions';
+import {
+  AggregationAttrs,
+  PivotAttrs,
+  TableAttrs
+} from '../common/pivot_table_data';
+import {globals} from './globals';
+
+export function isAggregationAttrs(attrs: PivotAttrs|AggregationAttrs):
+    attrs is AggregationAttrs {
+  return (attrs as AggregationAttrs).aggregation !== undefined;
+}
+
+function equalTableAttrs(
+    left: PivotAttrs|AggregationAttrs, right: PivotAttrs|AggregationAttrs) {
+  if (left.columnName !== right.columnName) {
+    return false;
+  }
+
+  if (left.tableName !== right.tableName) {
+    return false;
+  }
+
+  if (isAggregationAttrs(left) && isAggregationAttrs(right)) {
+    if (left.aggregation !== right.aggregation) {
+      return false;
+    }
+  }
+  return true;
+}
+
+export function getDataTransferType(isPivot: boolean) {
+  if (isPivot) {
+    return 'perfetto/pivot-table-dragged-pivot';
+  }
+  return 'perfetto/pivot-table-dragged-aggregation';
+}
+
+export class PivotTableHelper {
+  readonly pivotTableId: string;
+  readonly availableColumns: TableAttrs[];
+  readonly availableAggregations: string[];
+  readonly totalColumnsCount = 0;
+
+  private _selectedPivots: PivotAttrs[] = [];
+  private _selectedAggregations: AggregationAttrs[] = [];
+  private _isPivot = true;
+  private _selectedColumnIndex = 0;
+  private _selectedAggregationIndex = 0;
+  private _editPivotTableModalOpen = false;
+
+  constructor(
+      pivotTableId: string, availableColumns: TableAttrs[],
+      availableAggregations: string[], selectedPivots: PivotAttrs[],
+      selectedAggregations: AggregationAttrs[]) {
+    this.pivotTableId = pivotTableId;
+    this.availableColumns = availableColumns;
+    for (const table of this.availableColumns) {
+      this.totalColumnsCount += table.columns.length;
+    }
+    this.availableAggregations = availableAggregations;
+    this.setSelectedPivotsAndAggregations(selectedPivots, selectedAggregations);
+  }
+
+  // Sets selected pivots and aggregations if the editor modal is not open.
+  setSelectedPivotsAndAggregations(
+      selectedPivots: PivotAttrs[], selectedAggregations: AggregationAttrs[]) {
+    if (!this.editPivotTableModalOpen) {
+      // Making a copy of selectedPivots and selectedAggregations to preserve
+      // the original state.
+      this._selectedPivots =
+          selectedPivots.map(pivot => Object.assign({}, pivot));
+      this._selectedAggregations = selectedAggregations.map(
+          aggregation => Object.assign({}, aggregation));
+    }
+  }
+
+  // Dictates if the selected indexes refer to a pivot or aggregation.
+  togglePivotSelection() {
+    this._isPivot = !this._isPivot;
+  }
+
+  setSelectedPivotTableColumnIndex(index: number) {
+    if (index < 0 && index >= this.totalColumnsCount) {
+      throw Error(`Selected column index "${index}" out of bounds.`);
+    }
+    this._selectedColumnIndex = index;
+  }
+
+  setSelectedPivotTableAggregationIndex(index: number) {
+    if (index < 0 && index >= this.availableAggregations.length) {
+      throw Error(`Selected aggregation index "${index}" out of bounds.`);
+    }
+    this._selectedAggregationIndex = index;
+  }
+
+  // Get column attributes on selectedColumnIndex and
+  // selectedAggregationIndex.
+  getSelectedPivotTableColumnAttrs(): PivotAttrs|AggregationAttrs {
+    let tableName, columnName;
+    // Finds column index relative to its table.
+    let colIdx = this._selectedColumnIndex;
+    for (const {tableName: tblName, columns} of this.availableColumns) {
+      if (colIdx < columns.length) {
+        tableName = tblName;
+        columnName = columns[colIdx];
+        break;
+      }
+      colIdx -= columns.length;
+    }
+    if (tableName === undefined || columnName === undefined) {
+      throw Error(
+          'Pivot table selected column does not exist in availableColumns.');
+    }
+
+    // Get aggregation if selected column is not a pivot, undefined otherwise.
+    if (!this._isPivot) {
+      const aggregation =
+          this.availableAggregations[this._selectedAggregationIndex];
+      return {tableName, columnName, aggregation};
+    }
+
+    return {tableName, columnName};
+  }
+
+  // Adds column based on selected index to selectedPivots or
+  // selectedAggregations if it doesn't already exist, remove otherwise.
+  updatePivotTableColumnOnSelectedIndex() {
+    const columnAttrs = this.getSelectedPivotTableColumnAttrs();
+    this.updatePivotTableColumnOnColumnAttributes(columnAttrs);
+  }
+
+  // Adds column based on column attributes to selectedPivots or
+  // selectedAggregations if it doesn't already exist, remove otherwise.
+  updatePivotTableColumnOnColumnAttributes(columnAttrs: PivotAttrs|
+                                           AggregationAttrs) {
+    let storage: Array<PivotAttrs|AggregationAttrs>|undefined;
+    let attrs: PivotAttrs|AggregationAttrs;
+    if (isAggregationAttrs(columnAttrs)) {
+      storage = this._selectedAggregations;
+      attrs = {
+        tableName: columnAttrs.tableName,
+        columnName: columnAttrs.columnName,
+        aggregation: columnAttrs.aggregation,
+        order: 'DESC'
+      };
+    } else {
+      storage = this._selectedPivots;
+      attrs = {
+        tableName: columnAttrs.tableName,
+        columnName: columnAttrs.columnName
+      };
+    }
+    const index =
+        storage.findIndex(element => equalTableAttrs(element, columnAttrs));
+
+    if (index === -1) {
+      storage.push(attrs);
+    } else {
+      storage.splice(index, 1);
+    }
+  }
+
+  clearPivotTableColumns() {
+    this._selectedPivots = [];
+    this._selectedAggregations = [];
+  }
+
+  // Changes aggregation sorting from DESC to ASC and vice versa.
+  togglePivotTableAggregationSorting(index: number) {
+    if (index < 0 || index >= this._selectedAggregations.length) {
+      throw Error(`Column index "${index}" is out of bounds.`);
+    }
+    this._selectedAggregations[index].order =
+        this._selectedAggregations[index].order === 'DESC' ? 'ASC' : 'DESC';
+  }
+
+  // Moves target column to the requested destination.
+  reorderPivotTableDraggedColumn(
+      isPivot: boolean, targetColumnIdx: number, destinationColumnIdx: number) {
+    let storage;
+    if (isPivot) {
+      storage = this._selectedPivots;
+    } else {
+      storage = this._selectedAggregations;
+    }
+
+    if (targetColumnIdx < 0 || targetColumnIdx >= storage.length) {
+      throw Error(`Target column index "${targetColumnIdx}" out of bounds.`);
+    }
+    if (destinationColumnIdx < 0 || destinationColumnIdx >= storage.length) {
+      throw Error(
+          `Destination column index "${destinationColumnIdx}" out of bounds.`);
+    }
+
+    const targetColumn = storage[targetColumnIdx];
+    storage.splice(targetColumnIdx, 1);
+    storage.splice(destinationColumnIdx, 0, targetColumn);
+  }
+
+  selectedColumnOnDrag(e: DragEvent, isPivot: boolean, targetIdx: number) {
+    const dataTransferType = getDataTransferType(isPivot);
+    if (e.dataTransfer === null) {
+      return;
+    }
+    e.dataTransfer.setData(dataTransferType, targetIdx.toString());
+  }
+
+  selectedColumnOnDrop(
+      e: DragEvent, isPivot: boolean, destinationColumnIdx: number) {
+    const dataTransferType = getDataTransferType(isPivot);
+    if (e.dataTransfer === null) {
+      return;
+    }
+    // Prevents dragging pivots to aggregations and vice versa.
+    if (!e.dataTransfer.types.includes(dataTransferType)) {
+      return;
+    }
+
+    const targetColumnIdxString = e.dataTransfer.getData(dataTransferType);
+    const targetColumnIdx = Number(targetColumnIdxString);
+    if (!Number.isInteger(targetColumnIdx)) {
+      throw Error(
+          `Target column index "${targetColumnIdxString}" is not valid.`);
+    }
+
+    this.reorderPivotTableDraggedColumn(
+        isPivot, targetColumnIdx, destinationColumnIdx);
+    e.dataTransfer.clearData(dataTransferType);
+  }
+
+
+  // Highlights valid drop locations when dragging over them.
+  highlightDropLocation(e: DragEvent, isPivot: boolean) {
+    if (e.dataTransfer === null) {
+      return;
+    }
+    // Prevents highlighting aggregations when dragging pivots over them
+    // and vice versa.
+    if (!e.dataTransfer.types.includes(getDataTransferType(isPivot))) {
+      return;
+    }
+    (e.target as HTMLTableDataCellElement).classList.add('drop-location');
+  }
+
+  removeHighlightFromDropLocation(e: DragEvent) {
+    (e.target as HTMLTableDataCellElement).classList.remove('drop-location');
+  }
+
+  // Gets column index in availableColumns based on its attributes.
+  getColumnIndex(columnAttrs: PivotAttrs|AggregationAttrs) {
+    let index = 0;
+    for (const {tableName, columns} of this.availableColumns) {
+      if (tableName === columnAttrs.tableName) {
+        const colIdx =
+            columns.findIndex(column => column === columnAttrs.columnName);
+        return colIdx === -1 ? -1 : index + colIdx;
+      }
+      index += columns.length;
+    }
+    return -1;
+  }
+
+  selectPivotTableColumn(columnAttrs: PivotAttrs|AggregationAttrs) {
+    this._isPivot = !isAggregationAttrs(columnAttrs);
+
+    const colIndex = this.getColumnIndex(columnAttrs);
+    if (colIndex === -1) {
+      throw Error(`Selected column "${columnAttrs.tableName} ${
+          columnAttrs.columnName}" not found in availableColumns.`);
+    }
+    this.setSelectedPivotTableColumnIndex(colIndex);
+
+    if (isAggregationAttrs(columnAttrs)) {
+      const aggIndex = this.availableAggregations.findIndex(
+          aggregation => aggregation === columnAttrs.aggregation);
+      if (aggIndex === -1) {
+        throw Error(`Selected aggregation "${
+            columnAttrs.aggregation}" not found in availableAggregations.`);
+      }
+      this.setSelectedPivotTableAggregationIndex(aggIndex);
+    }
+  }
+
+  queryPivotTableChanges() {
+    globals.dispatch(Actions.setSelectedPivotsAndAggregations({
+      pivotTableId: this.pivotTableId,
+      selectedPivots: this._selectedPivots,
+      selectedAggregations: this._selectedAggregations
+    }));
+    globals.dispatch(Actions.setPivotTableRequest(
+        {pivotTableId: this.pivotTableId, action: 'QUERY'}));
+  }
+
+  toggleEditPivotTableModal() {
+    this._editPivotTableModalOpen = !this._editPivotTableModalOpen;
+  }
+
+  get selectedPivots() {
+    return this._selectedPivots.map(pivot => Object.assign({}, pivot));
+  }
+
+  get selectedAggregations() {
+    return this._selectedAggregations.map(
+        aggregation => Object.assign({}, aggregation));
+  }
+
+  get isPivot() {
+    return this._isPivot;
+  }
+
+  get selectedColumnIndex() {
+    return this._selectedColumnIndex;
+  }
+
+  get selectedAggregationIndex() {
+    return this._selectedAggregationIndex;
+  }
+
+  get editPivotTableModalOpen() {
+    return this._editPivotTableModalOpen;
+  }
+}
diff --git a/ui/src/frontend/pivot_table_helper_unittest.ts b/ui/src/frontend/pivot_table_helper_unittest.ts
new file mode 100644
index 000000000..6c3ae29a8
--- /dev/null
+++ b/ui/src/frontend/pivot_table_helper_unittest.ts
@@ -0,0 +1,125 @@
+// Copyright (C) 2021 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 {TableAttrs} from '../common/pivot_table_data';
+import {PivotTableHelper} from './pivot_table_helper';
+
+const AVAILABLE_COLUMNS: TableAttrs[] =
+    [{tableName: 'slice', columns: ['id', 'type', 'dur']}];
+const ID_COL_IDX = 0;
+const TYPE_COL_IDX = 1;
+const DUR_COL_IDX = 2;
+
+const AVAILABLE_AGGREGATIONS = ['SUM', 'AVG'];
+const SUM_AGG_IDX = 0;
+const AVG_AGG_IDX = 1;
+
+function createNewHelper() {
+  return new PivotTableHelper(
+      'pivotTable', AVAILABLE_COLUMNS, AVAILABLE_AGGREGATIONS, [], []);
+}
+
+test('Update selected pivots based on selected indices', () => {
+  const helper = createNewHelper();
+  helper.setSelectedPivotTableColumnIndex(ID_COL_IDX);
+
+  helper.updatePivotTableColumnOnSelectedIndex();
+  expect(helper.selectedPivots).toEqual([{
+    tableName: 'slice',
+    columnName: 'id',
+  }]);
+
+  helper.updatePivotTableColumnOnSelectedIndex();
+  expect(helper.selectedPivots).toEqual([]);
+});
+
+test('Update selected aggregations based on selected indices', () => {
+  const helper = createNewHelper();
+  helper.togglePivotSelection();
+  helper.setSelectedPivotTableColumnIndex(DUR_COL_IDX);
+  helper.setSelectedPivotTableAggregationIndex(SUM_AGG_IDX);
+
+  helper.updatePivotTableColumnOnSelectedIndex();
+  expect(helper.selectedAggregations).toEqual([
+    {tableName: 'slice', columnName: 'dur', aggregation: 'SUM', order: 'DESC'}
+  ]);
+
+  helper.updatePivotTableColumnOnSelectedIndex();
+  expect(helper.selectedAggregations).toEqual([]);
+});
+
+test('Change aggregation sorting based on aggregation index', () => {
+  const helper = createNewHelper();
+  helper.togglePivotSelection();
+  helper.setSelectedPivotTableColumnIndex(DUR_COL_IDX);
+  helper.setSelectedPivotTableAggregationIndex(SUM_AGG_IDX);
+  helper.updatePivotTableColumnOnSelectedIndex();
+
+  expect(helper.selectedAggregations).toEqual([
+    {tableName: 'slice', columnName: 'dur', aggregation: 'SUM', order: 'DESC'}
+  ]);
+  helper.togglePivotTableAggregationSorting(0);
+  expect(helper.selectedAggregations).toEqual([
+    {tableName: 'slice', columnName: 'dur', aggregation: 'SUM', order: 'ASC'}
+  ]);
+});
+
+test(
+    'Changing aggregation sorting with invalid index results in an error',
+    () => {
+      const helper = createNewHelper();
+      expect(() => helper.togglePivotTableAggregationSorting(1))
+          .toThrow('Column index "1" is out of bounds.');
+    });
+
+test('Reorder columns based on target and destination indices', () => {
+  const helper = createNewHelper();
+  helper.setSelectedPivotTableColumnIndex(ID_COL_IDX);
+  helper.updatePivotTableColumnOnSelectedIndex();
+  helper.setSelectedPivotTableColumnIndex(TYPE_COL_IDX);
+  helper.updatePivotTableColumnOnSelectedIndex();
+
+  expect(helper.selectedPivots).toEqual([
+    {tableName: 'slice', columnName: 'id'},
+    {tableName: 'slice', columnName: 'type'}
+  ]);
+  helper.reorderPivotTableDraggedColumn(true, 0, 1);
+  expect(helper.selectedPivots).toEqual([
+    {tableName: 'slice', columnName: 'type'},
+    {tableName: 'slice', columnName: 'id'}
+  ]);
+});
+
+test('Reordering columns with invalid indices results in an error', () => {
+  const helper = createNewHelper();
+  expect(() => helper.reorderPivotTableDraggedColumn(true, 0, 1))
+      .toThrow('Target column index "0" out of bounds.');
+});
+
+test('Select column based on attributes', () => {
+  const helper = createNewHelper();
+  helper.selectPivotTableColumn(
+      {tableName: 'slice', columnName: 'dur', aggregation: 'AVG'});
+  expect(helper.isPivot).toEqual(false);
+  expect(helper.selectedColumnIndex).toEqual(DUR_COL_IDX);
+  expect(helper.selectedAggregationIndex).toEqual(AVG_AGG_IDX);
+});
+
+test('Selecting a column with invalid attributes results in an error', () => {
+  const helper = createNewHelper();
+  expect(() => helper.selectPivotTableColumn({
+    tableName: 'foo',
+    columnName: 'bar',
+  })).toThrow('Selected column "foo bar" not found in availableColumns.');
+});
diff --git a/ui/src/frontend/publish.ts b/ui/src/frontend/publish.ts
index 247bbbdb1..eb12d687a 100644
--- a/ui/src/frontend/publish.ts
+++ b/ui/src/frontend/publish.ts
@@ -35,6 +35,7 @@ import {
   ThreadDesc,
   ThreadStateDetails
 } from './globals';
+import {PivotTableHelper} from './pivot_table_helper';
 
 export function publishOverviewData(
     data: {[key: string]: QuantizedLoad|QuantizedLoad[]}) {
@@ -146,6 +147,12 @@ export function publishQueryResult(args: {id: string, data?: {}}) {
   globals.publishRedraw();
 }
 
+export function publishPivotTableHelper(
+    args: {id: string, data: PivotTableHelper}) {
+  globals.pivotTableHelper.set(args.id, args.data);
+  globals.publishRedraw();
+}
+
 export function publishThreads(data: ThreadDesc[]) {
   globals.threads.clear();
   data.forEach(thread => {