fix(treemap): align onElementClick parameters to sunburst

This commit align the shape of passed parameter of the onElementClick listener to the one passed by
a sunburst. For each single hovered shape (usually only one shape at time) it returns the values for
each layer of the treemap

fix elastic#624

.playground/index.html

@@ -29,8 +29,8 @@
         /*display: inline-block;
         position: relative;
         */
-        width: 800px;
-        height: 400px;
+        width: 300px;
+        height: 300px;
         margin: 20px;
       }
     </style>

.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>
+    );
   }
 }

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,
+    };
+  };
+}

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}}`,
+        },
+      ],
+    ]);
+  });
+});

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;
+}

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;
     }