diff --git a/.playground/index.html b/.playground/index.html
index 6d829f699f..3973665311 100644
--- a/.playground/index.html
+++ b/.playground/index.html
@@ -29,8 +29,8 @@
         /*display: inline-block;
         position: relative;
         */
-        width: 800px;
-        height: 400px;
+        width: 300px;
+        height: 300px;
         margin: 20px;
       }
     </style>
diff --git a/.playground/playground.tsx b/.playground/playground.tsx
index 7d1e4f864f..3a170dc677 100644
--- a/.playground/playground.tsx
+++ b/.playground/playground.tsx
@@ -17,10 +17,43 @@
  * under the License. */
 
 import React from 'react';
-import { example } from '../stories/sunburst/12_very_small';
+import { Chart, Partition, Settings, PartitionLayout, XYChartElementEvent, PartitionElementEvent } from '../src';
 
 export class Playground extends React.Component {
+  onElementClick = (elements: (XYChartElementEvent | PartitionElementEvent)[]) => {
+    // eslint-disable-next-line no-console
+    console.log(elements);
+  };
   render() {
-    return <div className="chart">{example()}</div>;
+    return (
+      <div className="chart">
+        <Chart>
+          <Settings onElementClick={this.onElementClick} />
+          <Partition
+            id="111"
+            config={{
+              partitionLayout: PartitionLayout.treemap,
+            }}
+            valueAccessor={(d: { v: number }) => {
+              return d.v;
+            }}
+            data={[
+              { g1: 'a', g2: 'a', v: 1 },
+              { g1: 'a', g2: 'b', v: 1 },
+              { g1: 'b', g2: 'a', v: 1 },
+              { g1: 'b', g2: 'b', v: 1 },
+            ]}
+            layers={[
+              {
+                groupByRollup: (datum: { g1: string }) => datum.g1,
+              },
+              {
+                groupByRollup: (datum: { g2: string }) => datum.g2,
+              },
+            ]}
+          />
+        </Chart>
+      </div>
+    );
   }
 }
diff --git a/src/chart_types/partition_chart/layout/viewmodel/viewmodel.ts b/src/chart_types/partition_chart/layout/viewmodel/viewmodel.ts
index 80655316ae..9f48c27c48 100644
--- a/src/chart_types/partition_chart/layout/viewmodel/viewmodel.ts
+++ b/src/chart_types/partition_chart/layout/viewmodel/viewmodel.ts
@@ -202,27 +202,9 @@ export function shapeViewModel(
   const innerRadius: Radius = outerRadius - (1 - emptySizeRatio) * outerRadius;
   const treeHeight = shownChildNodes.reduce((p: number, n: any) => Math.max(p, entryValue(n.node).depth), 0); // 1: pie, 2: two-ring donut etc.
   const ringThickness = (outerRadius - innerRadius) / treeHeight;
-
+  const partToShapeFn = partToShapeTreeNode(treemapLayout, innerRadius, ringThickness);
   const quadViewModel = makeQuadViewModel(
-    shownChildNodes.slice(1).map(
-      (n: Part): ShapeTreeNode => {
-        const node: ArrayEntry = n.node;
-        return {
-          dataName: entryKey(node),
-          depth: depthAccessor(node),
-          value: aggregateAccessor(node),
-          parent: parentAccessor(node),
-          sortIndex: sortIndexAccessor(node),
-          x0: n.x0,
-          x1: n.x1,
-          y0: n.y0,
-          y1: n.y1,
-          y0px: treemapLayout ? n.y0 : innerRadius + n.y0 * ringThickness,
-          y1px: treemapLayout ? n.y1 : innerRadius + n.y1 * ringThickness,
-          yMidPx: treemapLayout ? (n.y0 + n.y1) / 2 : innerRadius + ((n.y0 + n.y1) / 2) * ringThickness,
-        };
-      },
-    ),
+    shownChildNodes.slice(1).map(partToShapeFn),
     layers,
     config.sectorLineWidth,
     config.sectorLineStroke,
@@ -308,3 +290,23 @@ export function shapeViewModel(
     outerRadius,
   };
 }
+
+function partToShapeTreeNode(treemapLayout: boolean, innerRadius: Radius, ringThickness: number) {
+  return (n: Part): ShapeTreeNode => {
+    const node: ArrayEntry = n.node;
+    return {
+      dataName: entryKey(node),
+      depth: depthAccessor(node),
+      value: aggregateAccessor(node),
+      parent: parentAccessor(node),
+      sortIndex: sortIndexAccessor(node),
+      x0: n.x0,
+      x1: n.x1,
+      y0: n.y0,
+      y1: n.y1,
+      y0px: treemapLayout ? n.y0 : innerRadius + n.y0 * ringThickness,
+      y1px: treemapLayout ? n.y1 : innerRadius + n.y1 * ringThickness,
+      yMidPx: treemapLayout ? (n.y0 + n.y1) / 2 : innerRadius + ((n.y0 + n.y1) / 2) * ringThickness,
+    };
+  };
+}
diff --git a/src/chart_types/partition_chart/state/selectors/picked_shapes.test.ts b/src/chart_types/partition_chart/state/selectors/picked_shapes.test.ts
new file mode 100644
index 0000000000..25bb500631
--- /dev/null
+++ b/src/chart_types/partition_chart/state/selectors/picked_shapes.test.ts
@@ -0,0 +1,140 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you 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 { chartStoreReducer, GlobalChartState } from '../../../../state/chart_state';
+import { createStore, Store } from 'redux';
+import { PartitionSpec } from '../../specs';
+import { upsertSpec, specParsed, specParsing } from '../../../../state/actions/specs';
+import { MockGlobalSpec, MockSeriesSpec } from '../../../../mocks/specs';
+import { updateParentDimensions } from '../../../../state/actions/chart_settings';
+import { partitionGeometries } from './geometries';
+import { onMouseDown, onMouseUp, onPointerMove } from '../../../../state/actions/mouse';
+import { createOnElementClickCaller } from './on_element_click_caller';
+import { SettingsSpec, XYChartElementEvent, PartitionElementEvent } from '../../../../specs';
+
+describe('Picked shapes selector', () => {
+  function initStore() {
+    const storeReducer = chartStoreReducer('chartId');
+    return createStore(storeReducer);
+  }
+  function addSeries(store: Store<GlobalChartState>, spec: PartitionSpec, settings?: Partial<SettingsSpec>) {
+    store.dispatch(specParsing());
+    store.dispatch(upsertSpec(MockGlobalSpec.settings(settings)));
+    store.dispatch(upsertSpec(spec));
+    store.dispatch(specParsed());
+    store.dispatch(updateParentDimensions({ width: 300, height: 300, top: 0, left: 0 }));
+  }
+  let store: Store<GlobalChartState>;
+  let treemapSpec: PartitionSpec;
+  let sunburstSpec: PartitionSpec;
+  beforeEach(() => {
+    store = initStore();
+    const common = {
+      valueAccessor: (d: { v: number }) => {
+        return d.v;
+      },
+      data: [
+        { g1: 'a', g2: 'a', v: 1 },
+        { g1: 'a', g2: 'b', v: 1 },
+        { g1: 'b', g2: 'a', v: 1 },
+        { g1: 'b', g2: 'b', v: 1 },
+      ],
+      layers: [
+        {
+          groupByRollup: (datum: { g1: string }) => datum.g1,
+        },
+        {
+          groupByRollup: (datum: { g2: string }) => datum.g2,
+        },
+      ],
+    };
+    treemapSpec = MockSeriesSpec.treemap(common);
+    sunburstSpec = MockSeriesSpec.sunburst(common);
+  });
+  test('check initial geoms', () => {
+    addSeries(store, treemapSpec);
+    const treemapGeometries = partitionGeometries(store.getState());
+    expect(treemapGeometries.quadViewModel).toHaveLength(6);
+
+    addSeries(store, sunburstSpec);
+    const sunburstGeometries = partitionGeometries(store.getState());
+    expect(sunburstGeometries.quadViewModel).toHaveLength(6);
+  });
+  test('treemap check picked geometries', () => {
+    const onClickListener = jest.fn<undefined, Array<(XYChartElementEvent | PartitionElementEvent)[]>>(
+      (): undefined => undefined,
+    );
+    addSeries(store, treemapSpec, {
+      onElementClick: onClickListener,
+    });
+    const geometries = partitionGeometries(store.getState());
+    expect(geometries.quadViewModel).toHaveLength(6);
+
+    const onElementClickCaller = createOnElementClickCaller();
+    store.subscribe(() => {
+      onElementClickCaller(store.getState());
+    });
+    store.dispatch(onPointerMove({ x: 200, y: 200 }, 0));
+    store.dispatch(onMouseDown({ x: 200, y: 200 }, 1));
+    store.dispatch(onMouseUp({ x: 200, y: 200 }, 2));
+    expect(onClickListener).toBeCalled();
+    expect(onClickListener.mock.calls[0][0]).toEqual([
+      [
+        [
+          { groupByRollup: 'b', value: 2 },
+          { groupByRollup: 'b', value: 1 },
+        ],
+        {
+          specId: treemapSpec.id,
+          key: `spec{${treemapSpec.id}}`,
+        },
+      ],
+    ]);
+  });
+  test('sunburst check picked geometries', () => {
+    const onClickListener = jest.fn<undefined, Array<(XYChartElementEvent | PartitionElementEvent)[]>>(
+      (): undefined => undefined,
+    );
+    addSeries(store, sunburstSpec, {
+      onElementClick: onClickListener,
+    });
+    const geometries = partitionGeometries(store.getState());
+    expect(geometries.quadViewModel).toHaveLength(6);
+
+    const onElementClickCaller = createOnElementClickCaller();
+    store.subscribe(() => {
+      onElementClickCaller(store.getState());
+    });
+    store.dispatch(onPointerMove({ x: 200, y: 200 }, 0));
+    store.dispatch(onMouseDown({ x: 200, y: 200 }, 1));
+    store.dispatch(onMouseUp({ x: 200, y: 200 }, 2));
+    expect(onClickListener).toBeCalled();
+    expect(onClickListener.mock.calls[0][0]).toEqual([
+      [
+        [
+          { groupByRollup: 'b', value: 2 },
+          { groupByRollup: 'b', value: 1 },
+        ],
+        {
+          specId: sunburstSpec.id,
+          key: `spec{${sunburstSpec.id}}`,
+        },
+      ],
+    ]);
+  });
+});
diff --git a/src/chart_types/partition_chart/state/selectors/picked_shapes.ts b/src/chart_types/partition_chart/state/selectors/picked_shapes.ts
index ae370676a7..97667f9f07 100644
--- a/src/chart_types/partition_chart/state/selectors/picked_shapes.ts
+++ b/src/chart_types/partition_chart/state/selectors/picked_shapes.ts
@@ -42,8 +42,17 @@ export const getPickedShapes = createCachedSelector(
 /** @internal */
 export const getPickedShapesLayerValues = createCachedSelector(
   [getPickedShapes],
-  (pickedShapes): Array<Array<LayerValue>> => {
-    const elements = pickedShapes.map<Array<LayerValue>>((model) => {
+  pickShapesLayerValues,
+)((state) => state.chartId);
+
+/** @internal */
+export function pickShapesLayerValues(pickedShapes: QuadViewModel[]): Array<Array<LayerValue>> {
+  const maxDepth = pickedShapes.reduce((acc, curr) => {
+    return Math.max(acc, curr.depth);
+  }, 0);
+  const elements = pickedShapes
+    .filter(({ depth }) => depth === maxDepth)
+    .map<Array<LayerValue>>((model) => {
       const values: Array<LayerValue> = [];
       values.push({
         groupByRollup: model.dataName,
@@ -61,6 +70,5 @@ export const getPickedShapesLayerValues = createCachedSelector(
       }
       return values.reverse();
     });
-    return elements;
-  },
-)((state) => state.chartId);
+  return elements;
+}
diff --git a/src/mocks/specs/specs.ts b/src/mocks/specs/specs.ts
index 7e67a32f3b..61c9ce515a 100644
--- a/src/mocks/specs/specs.ts
+++ b/src/mocks/specs/specs.ts
@@ -32,6 +32,12 @@ import { ScaleType } from '../../scales';
 import { ChartTypes } from '../../chart_types';
 import { SettingsSpec, SpecTypes, TooltipType } from '../../specs';
 import { LIGHT_THEME } from '../../utils/themes/light_theme';
+import { PartitionSpec } from '../../chart_types/partition_chart/specs';
+import { config, percentFormatter } from '../../chart_types/partition_chart/layout/config/config';
+import { ShapeTreeNode } from '../../chart_types/partition_chart/layout/types/viewmodel_types';
+import { Datum } from '../../utils/commons';
+import { AGGREGATE_KEY, PrimitiveValue } from '../../chart_types/partition_chart/layout/utils/group_by_rollup';
+import { PartitionLayout } from '../../chart_types/partition_chart/layout/types/config_types';
 
 /** @internal */
 export class MockSeriesSpec {
@@ -100,6 +106,52 @@ export class MockSeriesSpec {
     data: [],
   };
 
+  private static readonly sunburstBase: PartitionSpec = {
+    chartType: ChartTypes.Partition,
+    specType: SpecTypes.Series,
+    id: 'spec1',
+    config: {
+      ...config,
+      partitionLayout: PartitionLayout.sunburst,
+    },
+    valueAccessor: (d: Datum) => (typeof d === 'number' ? d : 0),
+    valueGetter: (n: ShapeTreeNode): number => n[AGGREGATE_KEY],
+    valueFormatter: (d: number): string => String(d),
+    percentFormatter,
+    layers: [
+      {
+        groupByRollup: (d: Datum, i: number) => i,
+        nodeLabel: (d: PrimitiveValue) => String(d),
+        showAccessor: () => true,
+        fillLabel: {},
+      },
+    ],
+    data: [],
+  };
+
+  private static readonly treemapBase: PartitionSpec = {
+    chartType: ChartTypes.Partition,
+    specType: SpecTypes.Series,
+    id: 'spec1',
+    config: {
+      ...config,
+      partitionLayout: PartitionLayout.treemap,
+    },
+    valueAccessor: (d: Datum) => (typeof d === 'number' ? d : 0),
+    valueGetter: (n: ShapeTreeNode): number => n[AGGREGATE_KEY],
+    valueFormatter: (d: number): string => String(d),
+    percentFormatter,
+    layers: [
+      {
+        groupByRollup: (d: Datum, i: number) => i,
+        nodeLabel: (d: PrimitiveValue) => String(d),
+        showAccessor: () => true,
+        fillLabel: {},
+      },
+    ],
+    data: [],
+  };
+
   static bar(partial?: Partial<BarSeriesSpec>): BarSeriesSpec {
     return mergePartial<BarSeriesSpec>(MockSeriesSpec.barBase, partial as RecursivePartial<BarSeriesSpec>, {
       mergeOptionalPartialValues: true,
@@ -128,14 +180,25 @@ export class MockSeriesSpec {
     });
   }
 
+  static sunburst(partial?: Partial<PartitionSpec>): PartitionSpec {
+    return mergePartial<PartitionSpec>(MockSeriesSpec.sunburstBase, partial as RecursivePartial<PartitionSpec>, {
+      mergeOptionalPartialValues: true,
+    });
+  }
+
+  static treemap(partial?: Partial<PartitionSpec>): PartitionSpec {
+    return mergePartial<PartitionSpec>(MockSeriesSpec.treemapBase, partial as RecursivePartial<PartitionSpec>, {
+      mergeOptionalPartialValues: true,
+    });
+  }
+
   static byType(type?: 'line' | 'bar' | 'area'): BasicSeriesSpec {
     switch (type) {
       case 'line':
         return MockSeriesSpec.lineBase;
-      case 'bar':
-        return MockSeriesSpec.barBase;
       case 'area':
         return MockSeriesSpec.areaBase;
+      case 'bar':
       default:
         return MockSeriesSpec.barBase;
     }
diff --git a/stories/interactions/4_sunburst_slice_clicks.tsx b/stories/interactions/4_sunburst_slice_clicks.tsx
index 85f759eefe..01a6b084b5 100644
--- a/stories/interactions/4_sunburst_slice_clicks.tsx
+++ b/stories/interactions/4_sunburst_slice_clicks.tsx
@@ -18,8 +18,14 @@
 
 import { action } from '@storybook/addon-actions';
 import React from 'react';
-import { Chart, Position, Settings, Partition } from '../../src';
-import { indexInterpolatedFillColor, interpolatorCET2s } from '../utils/utils';
+import { Chart, Position, Settings, Partition, PartitionLayout } from '../../src';
+import {
+  indexInterpolatedFillColor,
+  interpolatorCET2s,
+  categoricalFillColor,
+  colorBrewerCategoricalPastel12,
+} from '../utils/utils';
+import { select } from '@storybook/addon-knobs';
 
 const onElementListeners = {
   onElementClick: action('onElementClick'),
@@ -47,12 +53,20 @@ const pieData: Array<PieDatum> = [
 ];
 
 export const example = () => {
+  const partitionLayout = select(
+    'layout',
+    { sunburst: PartitionLayout.sunburst, treemap: PartitionLayout.treemap },
+    'sunburst',
+  );
   return (
     <Chart className="story-chart">
       <Settings showLegend showLegendExtra legendPosition={Position.Right} {...onElementListeners} />
       <Partition
         id="pie"
         data={pieData}
+        config={{
+          partitionLayout,
+        }}
         valueAccessor={(d) => {
           return d[3];
         }}
@@ -66,8 +80,12 @@ export const example = () => {
             },
             shape: {
               fillColor: (d) => {
-                // pick color from color palette based on mean angle - rather distinct colors in the inner ring
-                return indexInterpolatedFillColor(interpolatorCET2s)(d, (d.x0 + d.x1) / 2 / (2 * Math.PI), []);
+                if (partitionLayout === 'sunburst') {
+                  // pick color from color palette based on mean angle - rather distinct colors in the inner ring
+                  return indexInterpolatedFillColor(interpolatorCET2s)(d, (d.x0 + d.x1) / 2 / (2 * Math.PI), []);
+                } else {
+                  return categoricalFillColor(colorBrewerCategoricalPastel12)(d.sortIndex);
+                }
               },
             },
           },
@@ -80,8 +98,12 @@ export const example = () => {
             },
             shape: {
               fillColor: (d) => {
-                // pick color from color palette based on mean angle - rather distinct colors in the inner ring
-                return indexInterpolatedFillColor(interpolatorCET2s)(d, (d.x0 + d.x1) / 2 / (2 * Math.PI), []);
+                if (partitionLayout === 'sunburst') {
+                  // pick color from color palette based on mean angle - rather distinct colors in the inner ring
+                  return indexInterpolatedFillColor(interpolatorCET2s)(d, (d.x0 + d.x1) / 2 / (2 * Math.PI), []);
+                } else {
+                  return categoricalFillColor(colorBrewerCategoricalPastel12)(d.sortIndex);
+                }
               },
             },
           },