diff --git a/plugins/beta/datasets/src/DatasetsInit.jsx b/plugins/beta/datasets/src/DatasetsInit.jsx
index e2a8a4cc..1a15f4d8 100755
--- a/plugins/beta/datasets/src/DatasetsInit.jsx
+++ b/plugins/beta/datasets/src/DatasetsInit.jsx
@@ -59,6 +59,10 @@ export function DatasetsInit ({ pluginConfig, pluginState, appState, mapState, m
     initDatasets()
   }, [isMapStyleReady, appState.mode])
 
+  useEffect(() => {
+    dispatch({ type: 'BUILD_KEY_GROUPS', payload: null })
+  }, [pluginState.datasets])
+
   // Cleanup only on unmount
   useEffect(() => {
     return () => {
diff --git a/plugins/beta/datasets/src/panels/Key.jsx b/plugins/beta/datasets/src/panels/Key.jsx
index 26bcf9de..f450bd33 100755
--- a/plugins/beta/datasets/src/panels/Key.jsx
+++ b/plugins/beta/datasets/src/panels/Key.jsx
@@ -9,31 +9,6 @@ const SVG_SYMBOL_SIZE = 38
 const SVG_CENTER = SVG_SIZE / 2
 const PATTERN_INSET = 2
 
-const buildKeyGroups = (datasets) => {
-  const seenGroups = new Set()
-  const items = []
-  datasets.forEach(dataset => {
-    if (dataset.sublayers?.length) {
-      items.push({ type: 'sublayers', dataset })
-      return
-    }
-    if (dataset.groupLabel) {
-      if (seenGroups.has(dataset.groupLabel)) {
-        return
-      }
-      seenGroups.add(dataset.groupLabel)
-      items.push({
-        type: 'group',
-        groupLabel: dataset.groupLabel,
-        datasets: datasets.filter(d => !d.sublayers?.length && d.groupLabel === dataset.groupLabel)
-      })
-      return
-    }
-    items.push({ type: 'flat', dataset })
-  })
-  return items
-}
-
 export const Key = ({ mapState, pluginState, services }) => {
   const { mapStyle } = mapState
   const { symbolRegistry, patternRegistry } = services
@@ -123,11 +98,7 @@ export const Key = ({ mapState, pluginState, services }) => {
     </dl>
   )
 
-  const visibleDatasets = (pluginState.datasets || [])
-    .filter(dataset => dataset.showInKey && dataset.visibility !== 'hidden')
-
-  const keyGroups = buildKeyGroups(visibleDatasets)
-  const hasGroups = keyGroups.some(item => item.type === 'sublayers' || item.type === 'group')
+  const { items: keyGroups, hasGroups } = pluginState.key
   const containerClass = `im-c-datasets-key${hasGroups ? ' im-c-datasets-key--has-groups' : ''}`
 
   return (
diff --git a/plugins/beta/datasets/src/reducer.js b/plugins/beta/datasets/src/reducer.js
index abe59162..1409e86f 100755
--- a/plugins/beta/datasets/src/reducer.js
+++ b/plugins/beta/datasets/src/reducer.js
@@ -1,7 +1,12 @@
 import { applyDatasetDefaults } from './defaults.js'
+import { buildKeyState } from './utils/buildKeyState.js'
 
 const initialState = {
   datasets: null,
+  key: {
+    items: [],
+    hasGroups: false
+  },
   hiddenFeatures: {}, // { [layerId]: { idProperty: string, ids: string[] } }
   layerAdapter: null
 }
@@ -190,7 +195,17 @@ const setSublayerOpacity = (state, payload) => {
 
 const setLayerAdapter = (state, payload) => ({ ...state, layerAdapter: payload })
 
+const buildKeyGroups = (state) => {
+  const visibleDatasets = state.datasets
+    ? state.datasets.filter(dataset => dataset.showInKey && dataset.visibility !== 'hidden')
+    : []
+  const items = buildKeyState(visibleDatasets)
+  const hasGroups = items.some(item => item.type === 'sublayers' || item.type === 'group')
+  return { ...state, key: { items, hasGroups } }
+}
+
 const actions = {
+  BUILD_KEY_GROUPS: buildKeyGroups,
   SET_DATASETS: setDatasets,
   ADD_DATASET: addDataset,
   REMOVE_DATASET: removeDataset,
diff --git a/plugins/beta/datasets/src/utils/buildKeyState.js b/plugins/beta/datasets/src/utils/buildKeyState.js
new file mode 100644
index 00000000..5acef328
--- /dev/null
+++ b/plugins/beta/datasets/src/utils/buildKeyState.js
@@ -0,0 +1,24 @@
+export const buildKeyState = (datasets) => {
+  const seenGroups = new Set()
+  const items = []
+  datasets.forEach(dataset => {
+    if (dataset.sublayers?.length) {
+      items.push({ type: 'sublayers', dataset })
+      return
+    }
+    if (dataset.groupLabel) {
+      if (seenGroups.has(dataset.groupLabel)) {
+        return
+      }
+      seenGroups.add(dataset.groupLabel)
+      items.push({
+        type: 'group',
+        groupLabel: dataset.groupLabel,
+        datasets: datasets.filter(d => !d.sublayers?.length && d.groupLabel === dataset.groupLabel)
+      })
+      return
+    }
+    items.push({ type: 'flat', dataset })
+  })
+  return items
+}