diff --git a/x-pack/legacy/plugins/lens/public/datatable_visualization/visualization.tsx b/x-pack/legacy/plugins/lens/public/datatable_visualization/visualization.tsx
index 4248d722d55409..359c06a6a9ebce 100644
--- a/x-pack/legacy/plugins/lens/public/datatable_visualization/visualization.tsx
+++ b/x-pack/legacy/plugins/lens/public/datatable_visualization/visualization.tsx
@@ -150,6 +150,7 @@ export const datatableVisualization: Visualization<
           accessors: sortedColumns,
           supportsMoreColumns: true,
           filterOperations: () => true,
+          dataTestSubj: 'lnsDatatable_column',
         },
       ],
     };
diff --git a/x-pack/legacy/plugins/lens/public/editor_frame_service/editor_frame/config_panel_wrapper.tsx b/x-pack/legacy/plugins/lens/public/editor_frame_service/editor_frame/config_panel_wrapper.tsx
index c2cd0485de67ee..da812e948b23fd 100644
--- a/x-pack/legacy/plugins/lens/public/editor_frame_service/editor_frame/config_panel_wrapper.tsx
+++ b/x-pack/legacy/plugins/lens/public/editor_frame_service/editor_frame/config_panel_wrapper.tsx
@@ -129,7 +129,7 @@ function LayerPanels(
               },
             },
             visualization: {
-              activeId: activeVisualization.id,
+              ...prevState.visualization,
               state: newVisualizationState,
             },
             stagedPreview: undefined,
diff --git a/x-pack/test/functional/apps/lens/smokescreen.ts b/x-pack/test/functional/apps/lens/smokescreen.ts
index 317bb0b27e9729..5768e51ae5f9ff 100644
--- a/x-pack/test/functional/apps/lens/smokescreen.ts
+++ b/x-pack/test/functional/apps/lens/smokescreen.ts
@@ -9,7 +9,7 @@ import expect from '@kbn/expect';
 import { FtrProviderContext } from '../../ftr_provider_context';
 
 // eslint-disable-next-line import/no-default-export
-export default function({ getService, getPageObjects, ...rest }: FtrProviderContext) {
+export default function({ getService, getPageObjects }: FtrProviderContext) {
   const PageObjects = getPageObjects([
     'header',
     'common',
@@ -89,6 +89,17 @@ export default function({ getService, getPageObjects, ...rest }: FtrProviderCont
         field: 'bytes',
       });
 
+      await PageObjects.lens.configureDimension({
+        dimension:
+          '[data-test-subj="lnsXY_splitDimensionPanel"] [data-test-subj="lns-empty-dimension"]',
+        operation: 'terms',
+        field: '@message.raw',
+      });
+
+      await PageObjects.lens.switchToVisualization('lnsChartSwitchPopover_lnsDatatable');
+      await PageObjects.lens.removeDimension('lnsDatatable_column');
+      await PageObjects.lens.switchToVisualization('lnsChartSwitchPopover_bar_stacked');
+
       await PageObjects.lens.configureDimension({
         dimension:
           '[data-test-subj="lnsXY_splitDimensionPanel"] [data-test-subj="lns-empty-dimension"]',
diff --git a/x-pack/test/functional/page_objects/lens_page.ts b/x-pack/test/functional/page_objects/lens_page.ts
index 480814cb027817..1bf637c50b0ba4 100644
--- a/x-pack/test/functional/page_objects/lens_page.ts
+++ b/x-pack/test/functional/page_objects/lens_page.ts
@@ -103,8 +103,8 @@ export function LensPageProvider({ getService, getPageObjects }: FtrProviderCont
     /**
      * Changes the specified dimension to the specified operation and (optinally) field.
      *
-     * @param opts.from - the text of the dimension being changed
-     * @param opts.to - the desired operation for the dimension
+     * @param opts.dimension - the selector of the dimension being changed
+     * @param opts.operation - the desired operation ID for the dimension
      * @param opts.field - the desired field for the dimension
      */
     async configureDimension(opts: { dimension: string; operation?: string; field?: string }) {
@@ -123,6 +123,15 @@ export function LensPageProvider({ getService, getPageObjects }: FtrProviderCont
       }
     },
 
+    /**
+     * Removes the dimension matching a specific test subject
+     */
+    async removeDimension(dimensionTestSubj: string) {
+      await find.clickByCssSelector(
+        `[data-test-subj="${dimensionTestSubj}"] [data-test-subj="indexPattern-dimensionPopover-remove"]`
+      );
+    },
+
     /**
      * Save the current Lens visualization.
      */