diff --git a/src/column_display.tsx b/src/column_display.tsx
new file mode 100644
index 0000000..d34ce43
--- /dev/null
+++ b/src/column_display.tsx
@@ -0,0 +1,90 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import defaults from 'lodash/defaults';
+import React, { PureComponent } from 'react';
+import { MultiSelect, InlineLabel } from '@grafana/ui';
+import { QueryEditorProps, SelectableValue } from '@grafana/data';
+import { defaultQuery, PixieDataSourceOptions, PixieDataQuery } from './types';
+import { DataSource } from './datasource';
+
+type Props = QueryEditorProps<DataSource, PixieDataQuery, PixieDataSourceOptions>;
+
+export function getColumnsScript(
+  chosenOptions: Array<SelectableValue<number>>,
+  allColumnOptions: Array<SelectableValue<number>>
+): string {
+  const options = chosenOptions?.length ? chosenOptions : allColumnOptions;
+
+  return options.map(({ label }) => `'${label}'`).join(',');
+}
+
+export class ColDisplayComponents extends PureComponent<Props> {
+  onColSelect(chosenOptions: Array<SelectableValue<number>>) {
+    if (chosenOptions === undefined) {
+      return;
+    }
+    const { onChange, query, onRunQuery } = this.props;
+
+    // Update columns to groupby/aggregate based on any updates to the columns being displayed
+    const colsToDisplay = new Set(chosenOptions.map(({ label }) => label));
+
+    const groupByArr = query.queryMeta?.selectedColGroupby?.filter(({ label }) => colsToDisplay.has(label)) ?? [];
+
+    const aggData = groupByArr?.length
+      ? query.queryMeta?.aggData?.filter(({ aggColumn }) => colsToDisplay.has(aggColumn)) ?? []
+      : [];
+
+    onChange({
+      ...query,
+      queryMeta: {
+        ...query.queryMeta,
+        selectedColDisplay: chosenOptions,
+        selectedColGroupby: groupByArr,
+        aggData,
+      },
+    });
+    onRunQuery();
+  }
+
+  render() {
+    const query = defaults(this.props.query, defaultQuery);
+
+    return (
+      <div style={{ margin: '10px', display: 'flex' }}>
+        {query.queryMeta?.isColDisplay && (
+          <>
+            <InlineLabel transparent={false} width="auto">
+              Columns Displayed
+            </InlineLabel>
+            <MultiSelect
+              placeholder="Select Columns to Display"
+              options={query.queryMeta.columnOptions}
+              onChange={this.onColSelect.bind(this)}
+              closeMenuOnSelect={false}
+              width={32}
+              isClearable={true}
+              inputId="column-selection"
+              value={query.queryMeta.selectedColDisplay ?? undefined}
+            />
+          </>
+        )}
+      </div>
+    );
+  }
+}
diff --git a/src/column_filtering.tsx b/src/column_filtering.tsx
deleted file mode 100644
index 98f818e..0000000
--- a/src/column_filtering.tsx
+++ /dev/null
@@ -1,33 +0,0 @@
-/*
- * Copyright 2018- The Pixie Authors.
- *
- * 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.
- *
- * SPDX-License-Identifier: Apache-2.0
- */
-import { SelectableValue } from '@grafana/data';
-
-export function getColumnsScript(
-  chosenOptions: Array<SelectableValue<number>>,
-  allColumnOptions: Array<SelectableValue<number>>
-): string {
-  // Setting script default to all the columns in the script
-  let script = allColumnOptions.map(({ label }) => `'${label}'`).join();
-
-  // Updating the script if the user selected columns to filter
-  if (chosenOptions?.length > 0) {
-    script = chosenOptions.map(({ label }) => `'${label}'`).join();
-  }
-
-  return script;
-}
diff --git a/src/datasource.ts b/src/datasource.ts
index af25eb2..1835511 100644
--- a/src/datasource.ts
+++ b/src/datasource.ts
@@ -32,7 +32,7 @@ import {
   CLUSTER_VARIABLE_NAME as CLUSTER_VARIABLE_NAME,
   QueryType,
 } from './types';
-import { getColumnsScript } from './column_filtering';
+import { getColumnsScript } from './column_display';
 import { getGroupByScript } from './groupby';
 import { checkExhaustive, getClusterId } from 'utils';
 
diff --git a/src/groupby.tsx b/src/groupby.tsx
index b53da3d..2aa8d19 100644
--- a/src/groupby.tsx
+++ b/src/groupby.tsx
@@ -143,6 +143,7 @@ export class GroupbyComponents extends PureComponent<Props> {
             onChange={this.onGroupBySelect.bind(this)}
             closeMenuOnSelect={false}
             width={34}
+            isClearable={true}
             value={query.queryMeta?.selectedColGroupby ?? undefined}
           />
         </div>
diff --git a/src/query_editor.tsx b/src/query_editor.tsx
index 1fd798a..2f0c871 100644
--- a/src/query_editor.tsx
+++ b/src/query_editor.tsx
@@ -18,7 +18,7 @@
 
 import defaults from 'lodash/defaults';
 import React, { PureComponent } from 'react';
-import { Select, MultiSelect, Button, InlineLabel } from '@grafana/ui';
+import { Select, Button, InlineLabel } from '@grafana/ui';
 import { QueryEditorProps, SelectableValue } from '@grafana/data';
 import Editor from 'react-simple-code-editor';
 import { highlight, languages } from 'prismjs';
@@ -29,6 +29,7 @@ import { DataSource } from './datasource';
 import { scriptOptions, Script } from './pxl_scripts';
 import { defaultQuery, PixieDataSourceOptions, PixieDataQuery, QueryType } from './types';
 import { GroupbyComponents } from './groupby';
+import { ColDisplayComponents } from './column_display';
 
 type Props = QueryEditorProps<DataSource, PixieDataQuery, PixieDataSourceOptions>;
 
@@ -66,20 +67,11 @@ export class QueryEditor extends PureComponent<Props> {
           isGroupBy: option.value.isGroupBy || false,
           columnOptions: option.columnOptions,
           groupByColOptions: option.groupByColOptions,
-          selectedColDisplay: [],
+          selectedColDisplay: option.columnOptions,
           selectedColGroupby: [],
           aggData: [],
         },
       });
-
-      onRunQuery();
-    }
-  }
-
-  onColSelect(chosenOptions: Array<SelectableValue<number>>) {
-    if (chosenOptions !== undefined) {
-      const { onChange, query, onRunQuery } = this.props;
-      onChange({ ...query, queryMeta: { ...query.queryMeta, selectedColDisplay: chosenOptions } });
       onRunQuery();
     }
   }
@@ -104,24 +96,14 @@ export class QueryEditor extends PureComponent<Props> {
             />
           </div>
 
-          <div style={{ margin: '10px', display: 'flex' }}>
-            {query.queryMeta?.isColDisplay && (
-              <>
-                <InlineLabel transparent={false} width="auto">
-                  Columns Displayed
-                </InlineLabel>
-                <MultiSelect
-                  placeholder="Select Columns to Display"
-                  options={query.queryMeta.columnOptions}
-                  onChange={this.onColSelect.bind(this)}
-                  closeMenuOnSelect={false}
-                  width={32}
-                  inputId="column-selection"
-                  value={query.queryMeta.selectedColDisplay}
-                />
-              </>
-            )}
-          </div>
+          {query.queryMeta?.isColDisplay && (
+            <ColDisplayComponents
+              datasource={this.props.datasource}
+              query={query}
+              onRunQuery={onRunQuery}
+              onChange={onChange}
+            />
+          )}
 
           {query.queryMeta?.isGroupBy && (
             <GroupbyComponents