diff --git a/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx b/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
index d0a17c75e1b51..a013d9b1bceae 100644
--- a/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
+++ b/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/data_panel_wrapper.tsx
@@ -11,6 +11,7 @@ import { DatasourceDataPanelProps, Datasource } from '../../../public';
 import { NativeRenderer } from '../../native_renderer';
 import { Action } from './state_management';
 import { DragContext } from '../../drag_drop';
+import { StateSetter } from '../../types';
 
 interface DataPanelWrapperProps {
   datasourceState: unknown;
@@ -21,11 +22,11 @@ interface DataPanelWrapperProps {
 }
 
 export const DataPanelWrapper = memo((props: DataPanelWrapperProps) => {
-  const setDatasourceState = useMemo(
-    () => (newState: unknown) => {
+  const setDatasourceState: StateSetter<unknown> = useMemo(
+    () => updater => {
       props.dispatch({
         type: 'UPDATE_DATASOURCE_STATE',
-        newState,
+        updater,
         datasourceId: props.activeDatasource!,
       });
     },
diff --git a/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx b/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
index ff393eabf8ca5..71dca959a0afc 100644
--- a/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
+++ b/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/editor_frame.tsx
@@ -59,7 +59,7 @@ export function EditorFrame(props: EditorFrameProps) {
             .then(datasourceState => {
               dispatch({
                 type: 'UPDATE_DATASOURCE_STATE',
-                newState: datasourceState,
+                updater: datasourceState,
                 datasourceId,
               });
             })
@@ -84,7 +84,7 @@ export function EditorFrame(props: EditorFrameProps) {
             dispatch({
               type: 'UPDATE_DATASOURCE_STATE',
               datasourceId: id,
-              newState,
+              updater: newState,
             });
           },
           layer
diff --git a/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts b/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
index 22ff323507487..8fdfe6c466828 100644
--- a/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
+++ b/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.test.ts
@@ -124,8 +124,8 @@ describe('editor_frame state management', () => {
       expect(newState.visualization.state).toBe(newVisState);
     });
 
-    it('should update the datasource state on update', () => {
-      const newDatasourceState = {};
+    it('should update the datasource state with passed in reducer', () => {
+      const datasourceReducer = jest.fn(() => ({ changed: true }));
       const newState = reducer(
         {
           datasourceStates: {
@@ -143,16 +143,17 @@ describe('editor_frame state management', () => {
         },
         {
           type: 'UPDATE_DATASOURCE_STATE',
-          newState: newDatasourceState,
+          updater: datasourceReducer,
           datasourceId: 'testDatasource',
         }
       );
 
-      expect(newState.datasourceStates.testDatasource.state).toBe(newDatasourceState);
+      expect(newState.datasourceStates.testDatasource.state).toEqual({ changed: true });
+      expect(datasourceReducer).toHaveBeenCalledTimes(1);
     });
 
-    it('should update the datasource state with passed in reducer', () => {
-      const layerReducer = jest.fn((_state, layerId) => ({ inserted: layerId }));
+    it('should update the layer state with passed in reducer', () => {
+      const newDatasourceState = {};
       const newState = reducer(
         {
           datasourceStates: {
@@ -169,15 +170,13 @@ describe('editor_frame state management', () => {
           },
         },
         {
-          type: 'UPDATE_LAYER',
-          layerId: 'abc',
-          updater: layerReducer,
+          type: 'UPDATE_DATASOURCE_STATE',
+          updater: newDatasourceState,
           datasourceId: 'testDatasource',
         }
       );
 
-      expect(newState.datasourceStates.testDatasource.state).toEqual({ inserted: 'abc' });
-      expect(layerReducer).toHaveBeenCalledTimes(1);
+      expect(newState.datasourceStates.testDatasource.state).toBe(newDatasourceState);
     });
 
     it('should should switch active visualization', () => {
diff --git a/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts b/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
index 6f76cc15dafed..fff3ae3d1f45b 100644
--- a/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
+++ b/x-pack/legacy/plugins/lens/public/editor_frame_plugin/editor_frame/state_management.ts
@@ -30,7 +30,7 @@ export type Action =
     }
   | {
       type: 'UPDATE_DATASOURCE_STATE';
-      newState: unknown;
+      updater: unknown | ((prevState: unknown) => unknown);
       datasourceId: string;
     }
   | {
@@ -168,7 +168,10 @@ export const reducer = (state: EditorFrameState, action: Action): EditorFrameSta
         datasourceStates: {
           ...state.datasourceStates,
           [action.datasourceId]: {
-            state: action.newState,
+            state:
+              typeof action.updater === 'function'
+                ? action.updater(state.datasourceStates[action.datasourceId].state)
+                : action.updater,
             isLoading: false,
           },
         },
diff --git a/x-pack/legacy/plugins/lens/public/indexpattern_plugin/dimension_panel/dimension_panel.tsx b/x-pack/legacy/plugins/lens/public/indexpattern_plugin/dimension_panel/dimension_panel.tsx
index 689e5a0f52c26..f6d80b38c6ab9 100644
--- a/x-pack/legacy/plugins/lens/public/indexpattern_plugin/dimension_panel/dimension_panel.tsx
+++ b/x-pack/legacy/plugins/lens/public/indexpattern_plugin/dimension_panel/dimension_panel.tsx
@@ -10,7 +10,7 @@ import { EuiButtonIcon } from '@elastic/eui';
 import { Storage } from 'ui/storage';
 import { i18n } from '@kbn/i18n';
 import { UiSettingsClientContract } from 'src/core/public';
-import { DatasourceDimensionPanelProps } from '../../types';
+import { DatasourceDimensionPanelProps, StateSetter } from '../../types';
 import {
   IndexPatternColumn,
   IndexPatternPrivateState,
@@ -26,7 +26,7 @@ import { isDraggedField } from '../utils';
 
 export type IndexPatternDimensionPanelProps = DatasourceDimensionPanelProps & {
   state: IndexPatternPrivateState;
-  setState: (newState: IndexPatternPrivateState) => void;
+  setState: StateSetter<IndexPatternPrivateState>;
   dragDropContext: DragContextState;
   uiSettings: UiSettingsClientContract;
   storage: Storage;
diff --git a/x-pack/legacy/plugins/lens/public/indexpattern_plugin/indexpattern.tsx b/x-pack/legacy/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
index b81ff9e553502..d674a96f7bda5 100644
--- a/x-pack/legacy/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
+++ b/x-pack/legacy/plugins/lens/public/indexpattern_plugin/indexpattern.tsx
@@ -305,7 +305,7 @@ export function getIndexPatternDatasource({
             <I18nProvider>
               <IndexPatternDimensionPanel
                 state={state}
-                setState={newState => setState(newState)}
+                setState={setState}
                 uiSettings={uiSettings}
                 storage={storage}
                 layerId={props.layerId}
@@ -317,10 +317,7 @@ export function getIndexPatternDatasource({
         },
 
         renderLayerPanel: (domElement: Element, props: DatasourceLayerPanelProps) => {
-          render(
-            <LayerPanel state={state} setState={newState => setState(newState)} {...props} />,
-            domElement
-          );
+          render(<LayerPanel state={state} setState={setState} {...props} />, domElement);
         },
 
         removeColumnInTableSpec: (columnId: string) => {
diff --git a/x-pack/legacy/plugins/lens/public/indexpattern_plugin/layerpanel.tsx b/x-pack/legacy/plugins/lens/public/indexpattern_plugin/layerpanel.tsx
index 27b128eb35ef2..c221371dd6aa2 100644
--- a/x-pack/legacy/plugins/lens/public/indexpattern_plugin/layerpanel.tsx
+++ b/x-pack/legacy/plugins/lens/public/indexpattern_plugin/layerpanel.tsx
@@ -18,13 +18,13 @@ import {
 } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 import { I18nProvider } from '@kbn/i18n/react';
-import { DatasourceLayerPanelProps } from '../types';
+import { DatasourceLayerPanelProps, StateSetter } from '../types';
 import { IndexPatternPrivateState, IndexPatternLayer } from './indexpattern';
 import { isLayerTransferable, updateLayerIndexPattern } from './state_helpers';
 
 export interface IndexPatternLayerPanelProps extends DatasourceLayerPanelProps {
   state: IndexPatternPrivateState;
-  setState: (newState: IndexPatternPrivateState) => void;
+  setState: StateSetter<IndexPatternPrivateState>;
 }
 
 function LayerPanelChooser({
diff --git a/x-pack/legacy/plugins/lens/public/indexpattern_plugin/operations.ts b/x-pack/legacy/plugins/lens/public/indexpattern_plugin/operations.ts
index 25821f505fe08..d2a0b16d647f2 100644
--- a/x-pack/legacy/plugins/lens/public/indexpattern_plugin/operations.ts
+++ b/x-pack/legacy/plugins/lens/public/indexpattern_plugin/operations.ts
@@ -7,7 +7,7 @@
 import _ from 'lodash';
 import { Storage } from 'ui/storage';
 import { UiSettingsClientContract } from 'src/core/public';
-import { DimensionPriority, OperationMetadata } from '../types';
+import { DimensionPriority, OperationMetadata, StateSetter } from '../types';
 import {
   IndexPatternColumn,
   IndexPatternField,
@@ -60,7 +60,7 @@ export function getOperations(): OperationType[] {
 
 export interface ParamEditorProps {
   state: IndexPatternPrivateState;
-  setState: (newState: IndexPatternPrivateState) => void;
+  setState: StateSetter<IndexPatternPrivateState>;
   columnId: string;
   layerId: string;
   uiSettings: UiSettingsClientContract;
diff --git a/x-pack/legacy/plugins/lens/public/types.ts b/x-pack/legacy/plugins/lens/public/types.ts
index 6deab82c5658e..3f2ee5b8067c9 100644
--- a/x-pack/legacy/plugins/lens/public/types.ts
+++ b/x-pack/legacy/plugins/lens/public/types.ts
@@ -63,6 +63,8 @@ export interface DatasourceMetaData {
   filterableIndexPatterns: Array<{ id: string; title: string }>;
 }
 
+export type StateSetter<T> = (newState: T | ((prevState: T) => T)) => void;
+
 /**
  * Interface for the datasource registry
  */
@@ -88,7 +90,7 @@ export interface Datasource<T = unknown, P = unknown> {
   getDatasourceSuggestionsForField: (state: T, field: unknown) => Array<DatasourceSuggestion<T>>;
   getDatasourceSuggestionsFromCurrentState: (state: T) => Array<DatasourceSuggestion<T>>;
 
-  getPublicAPI: (state: T, setState: (newState: T) => void, layerId: string) => DatasourcePublicAPI;
+  getPublicAPI: (state: T, setState: StateSetter<T>, layerId: string) => DatasourcePublicAPI;
 }
 
 /**
@@ -118,7 +120,7 @@ export type TableSpec = TableSpecColumn[];
 export interface DatasourceDataPanelProps<T = unknown> {
   state: T;
   dragDropContext: DragContextState;
-  setState: (newState: T) => void;
+  setState: StateSetter<T>;
 }
 
 // The only way a visualization has to restrict the query building