feat(heatmap): enable brushing on categorical charts

packages/charts/src/chart_types/heatmap/layout/types/viewmodel_types.ts

@@ -19,10 +19,12 @@
 
 import { ChartType } from '../../..';
 import { Pixels } from '../../../../common/geometry';
+import { Box } from '../../../../common/text_utils';
 import { Fill, Line, Stroke } from '../../../../geoms/types';
 import { Point } from '../../../../utils/point';
+import { PrimitiveValue } from '../../../partition_chart/layout/utils/group_by_rollup';
 import { config } from '../config/config';
-import { HeatmapCellDatum, TextBox } from '../viewmodel/viewmodel';
+import { HeatmapCellDatum } from '../viewmodel/viewmodel';
 import { Config, HeatmapBrushEvent } from './config_types';
 
 /** @internal */
@@ -47,6 +49,13 @@ export interface Cell {
   datum: HeatmapCellDatum;
 }
 
+/** @internal */
+export interface TextBox extends Box {
+  value: NonNullable<PrimitiveValue>;
+  x: number;
+  y: number;
+}
+
 /** @internal */
 export interface HeatmapViewModel {
   gridOrigin: {

packages/charts/src/chart_types/heatmap/layout/viewmodel/viewmodel.ts

@@ -27,7 +27,9 @@ import { ScaleContinuous } from '../../../../scales';
 import { ScaleType } from '../../../../scales/constants';
 import { SettingsSpec } from '../../../../specs';
 import { CanvasTextBBoxCalculator } from '../../../../utils/bbox/canvas_text_bbox_calculator';
+import { clamp } from '../../../../utils/common';
 import { Dimensions } from '../../../../utils/dimensions';
+import { PrimitiveValue } from '../../../partition_chart/layout/utils/group_by_rollup';
 import { HeatmapSpec } from '../../specs';
 import { HeatmapTable } from '../../state/selectors/compute_chart_dimensions';
 import { ColorScaleType } from '../../state/selectors/get_color_scale';
@@ -39,21 +41,25 @@ import {
   PickDragShapeFunction,
   PickHighlightedArea,
   ShapeViewModel,
+  TextBox,
 } from '../types/viewmodel_types';
 
 /** @public */
 export interface HeatmapCellDatum {
-  x: string | number;
-  y: string | number;
+  x: NonNullable<PrimitiveValue>;
+  y: NonNullable<PrimitiveValue>;
   value: number;
   originalIndex: number;
 }
 
-/** @internal */
-export interface TextBox extends Box {
-  value: string | number;
-  x: number;
-  y: number;
+function getValuesInRange(
+  values: NonNullable<PrimitiveValue>[],
+  startValue: NonNullable<PrimitiveValue>,
+  endValue: NonNullable<PrimitiveValue>,
+) {
+  const startIndex = values.indexOf(startValue);
+  const endIndex = Math.min(values.indexOf(endValue) + 1, values.length);
+  return values.slice(startIndex, endIndex);
 }
 
 /**
@@ -93,7 +99,7 @@ export function shapeViewModel(
   const { table, yValues, xDomain } = heatmapTable;
 
   // measure the text width of all rows values to get the grid area width
-  const boxedYValues = yValues.map<Box & { value: string | number }>((value) => {
+  const boxedYValues = yValues.map<Box & { value: NonNullable<PrimitiveValue> }>((value) => {
     return {
       text: config.yAxisLabel.formatter(value),
       value,
@@ -102,9 +108,9 @@ export function shapeViewModel(
   });
 
   // compute the scale for the rows positions
-  const yScale = scaleBand<string | number>().domain(yValues).range([0, height]);
+  const yScale = scaleBand<NonNullable<PrimitiveValue>>().domain(yValues).range([0, height]);
 
-  const yInvertedScale = scaleQuantize<string | number>().domain([0, height]).range(yValues);
+  const yInvertedScale = scaleQuantize<NonNullable<PrimitiveValue>>().domain([0, height]).range(yValues);
 
   // TODO: Fix domain type to be `Array<number | string>`
   let xValues = xDomain.domain as any[];
@@ -137,9 +143,9 @@ export function shapeViewModel(
   }
 
   // compute the scale for the columns positions
-  const xScale = scaleBand<string | number>().domain(xValues).range([0, chartDimensions.width]);
+  const xScale = scaleBand<NonNullable<PrimitiveValue>>().domain(xValues).range([0, chartDimensions.width]);
 
-  const xInvertedScale = scaleQuantize<string | number>().domain([0, chartDimensions.width]).range(xValues);
+  const xInvertedScale = scaleQuantize<NonNullable<PrimitiveValue>>().domain([0, chartDimensions.width]).range(xValues);
 
   // compute the cell width (can be smaller then the available size depending on config
   const cellWidth =
@@ -150,6 +156,8 @@ export function shapeViewModel(
   // compute the cell height (we already computed the max size for that)
   const cellHeight = yScale.bandwidth();
 
+  const currentGridHeight = cellHeight * pageSize;
+
   const getTextValue = (
     formatter: (v: any, options: any) => string,
     scaleCallback: (x: any) => number | undefined | null = xScale,
@@ -258,19 +266,20 @@ export function shapeViewModel(
   const pickDragArea: PickDragFunction = (bound) => {
     const [start, end] = bound;
 
-    const { left, top } = chartDimensions;
-    const invertedBounds = {
-      startX: xInvertedScale(Math.min(start.x, end.x) - left),
-      startY: yInvertedScale(Math.min(start.y, end.y) - top),
-      endX: xInvertedScale(Math.max(start.x, end.x) - left),
-      endY: yInvertedScale(Math.max(start.y, end.y) - top),
-    };
+    const { left, top, width } = chartDimensions;
+    const topLeft = [Math.min(start.x, end.x) - left, Math.min(start.y, end.y) - top];
+    const bottomRight = [Math.max(start.x, end.x) - left, Math.max(start.y, end.y) - top];
 
-    let allXValuesInRange = [];
-    const invertedXValues: Array<string | number> = [];
-    const { startX, endX, startY, endY } = invertedBounds;
-    invertedXValues.push(startX);
-    if (typeof endX === 'number') {
+    const startX = xInvertedScale(clamp(topLeft[0], 0, width));
+    const endX = xInvertedScale(clamp(bottomRight[0], 0, width));
+    const startY = yInvertedScale(clamp(topLeft[1], 0, currentGridHeight - 1));
+    const endY = yInvertedScale(clamp(bottomRight[1], 0, currentGridHeight - 1));
+
+    let allXValuesInRange: Array<NonNullable<PrimitiveValue>> = [];
+    const invertedXValues: Array<NonNullable<PrimitiveValue>> = [];
+
+    if (timeScale && typeof endX === 'number') {
+      invertedXValues.push(startX);
       invertedXValues.push(endX + xDomain.minInterval);
       let [startXValue] = invertedXValues;
       if (typeof startXValue === 'number') {
@@ -280,19 +289,11 @@ export function shapeViewModel(
         }
       }
     } else {
-      invertedXValues.push(endX);
-      const startXIndex = xValues.indexOf(startX);
-      const endXIndex = Math.min(xValues.indexOf(endX) + 1, xValues.length);
-      allXValuesInRange = xValues.slice(startXIndex, endXIndex);
+      allXValuesInRange = getValuesInRange(xValues, startX, endX);
       invertedXValues.push(...allXValuesInRange);
     }
 
-    const invertedYValues: Array<string | number> = [];
-
-    const startYIndex = yValues.indexOf(startY);
-    const endYIndex = Math.min(yValues.indexOf(endY) + 1, yValues.length);
-    const allYValuesInRange = yValues.slice(startYIndex, endYIndex);
-    invertedYValues.push(...allYValuesInRange);
+    const allYValuesInRange: Array<NonNullable<PrimitiveValue>> = getValuesInRange(yValues, startY, endY);
 
     const cells: Cell[] = [];
 
@@ -306,7 +307,7 @@ export function shapeViewModel(
     return {
       cells: cells.filter(Boolean),
       x: invertedXValues,
-      y: invertedYValues,
+      y: allYValuesInRange,
     };
   };
 
@@ -315,26 +316,22 @@ export function shapeViewModel(
    * @param x
    * @param y
    */
-  const pickHighlightedArea: PickHighlightedArea = (x: Array<string | number>, y: Array<string | number>) => {
-    if (xDomain.type !== ScaleType.Time) {
-      return null;
-    }
-    const [startValue, endValue] = x;
-
-    if (typeof startValue !== 'number' || typeof endValue !== 'number') {
-      return null;
-    }
-    const start = Math.min(startValue, endValue);
-    const end = Math.max(startValue, endValue);
+  const pickHighlightedArea: PickHighlightedArea = (
+    x: Array<NonNullable<PrimitiveValue>>,
+    y: Array<NonNullable<PrimitiveValue>>,
+  ) => {
+    const startValue = x[0];
+    const endValue = x[x.length - 1];
 
     // find X coordinated based on the time range
-    const leftIndex = bisectLeft(xValues, start);
-    const rightIndex = bisectLeft(xValues, end);
+    const leftIndex = typeof startValue === 'number' ? bisectLeft(xValues, startValue) : xValues.indexOf(startValue);
+    const rightIndex = typeof endValue === 'number' ? bisectLeft(xValues, endValue) : xValues.indexOf(endValue);
 
-    const isOutOfRange = rightIndex > xValues.length - 1;
+    const isRightOutOfRange = rightIndex > xValues.length - 1 || rightIndex < 0;
+    const isLeftOutOfRange = leftIndex > xValues.length - 1 || leftIndex < 0;
 
-    const startFromScale = xScale(xValues[leftIndex]);
-    const endFromScale = xScale(isOutOfRange ? xValues[xValues.length - 1] : xValues[rightIndex]);
+    const startFromScale = xScale(isLeftOutOfRange ? xValues[0] : xValues[leftIndex]);
+    const endFromScale = xScale(isRightOutOfRange ? xValues[xValues.length - 1] : xValues[rightIndex]);
 
     if (startFromScale === undefined || endFromScale === undefined) {
       return null;
@@ -343,7 +340,7 @@ export function shapeViewModel(
     const xStart = chartDimensions.left + startFromScale;
 
     // extend the range in case the right boundary has been selected
-    const width = endFromScale - startFromScale + (isOutOfRange ? cellWidth : 0);
+    const width = endFromScale - startFromScale + cellWidth; // (isRightOutOfRange || isLeftOutOfRange ? cellWidth : 0);
 
     // resolve Y coordinated making sure the order is correct
     const { y: yStart, totalHeight } = y
@@ -358,7 +355,6 @@ export function shapeViewModel(
         },
         { y: 0, totalHeight: 0 },
       );
-
     return {
       x: xStart,
       y: yStart,
@@ -417,7 +413,7 @@ export function shapeViewModel(
   };
 }
 
-function getCellKey(x: string | number, y: string | number) {
+function getCellKey(x: NonNullable<PrimitiveValue>, y: NonNullable<PrimitiveValue>) {
   return [String(x), String(y)].join('&_&');
 }
 

packages/charts/src/chart_types/heatmap/state/selectors/get_brushed_highlighted_shapes.test.ts

@@ -0,0 +1,89 @@
+/*
+ * 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 { Store } from 'redux';
+
+import { MockGlobalSpec, MockSeriesSpec } from '../../../../mocks/specs/specs';
+import { MockStore } from '../../../../mocks/store/store';
+import { ScaleType } from '../../../../scales/constants';
+import { onMouseDown, onMouseUp, onPointerMove } from '../../../../state/actions/mouse';
+import { GlobalChartState } from '../../../../state/chart_state';
+import { createOnBrushEndCaller } from './on_brush_end_caller';
+
+describe('Heatmap brush', () => {
+  let store: Store<GlobalChartState>;
+  let onBrushEndMock = jest.fn();
+
+  beforeEach(() => {
+    store = MockStore.default({ width: 300, height: 300, top: 0, left: 0 }, 'chartId');
+    onBrushEndMock = jest.fn();
+    MockStore.addSpecs(
+      [
+        MockGlobalSpec.settingsNoMargins(),
+        MockSeriesSpec.heatmap({
+          xScaleType: ScaleType.Ordinal,
+          data: [
+            { x: 'a', y: 'ya', value: 1 },
+            { x: 'b', y: 'ya', value: 2 },
+            { x: 'c', y: 'ya', value: 3 },
+            { x: 'a', y: 'yb', value: 4 },
+            { x: 'b', y: 'yb', value: 5 },
+            { x: 'c', y: 'yb', value: 6 },
+            { x: 'a', y: 'yc', value: 7 },
+            { x: 'b', y: 'yc', value: 8 },
+            { x: 'c', y: 'yc', value: 9 },
+          ],
+          config: {
+            grid: {
+              cellHeight: {
+                max: 'fill',
+              },
+              cellWidth: {
+                max: 'fill',
+              },
+            },
+            xAxisLabel: {
+              visible: false,
+            },
+            yAxisLabel: {
+              visible: false,
+            },
+            margin: { top: 0, bottom: 0, left: 0, right: 0 },
+            onBrushEnd: onBrushEndMock,
+          },
+        }),
+      ],
+      store,
+    );
+  });
+
+  it('should brush on categorical scale', () => {
+    const caller = createOnBrushEndCaller();
+    store.dispatch(onPointerMove({ x: 50, y: 50 }, 0));
+    store.dispatch(onMouseDown({ x: 50, y: 50 }, 100));
+    store.dispatch(onPointerMove({ x: 150, y: 250 }, 200));
+    store.dispatch(onMouseUp({ x: 150, y: 250 }, 300));
+    caller(store.getState());
+    expect(onBrushEndMock).toBeCalledTimes(1);
+    const brushEvent = onBrushEndMock.mock.calls[0][0];
+    expect(brushEvent.cells).toHaveLength(6);
+    expect(brushEvent.x).toEqual(['a', 'b']);
+    expect(brushEvent.y).toEqual(['ya', 'yb', 'yc']);
+  });
+});

packages/charts/src/chart_types/heatmap/state/selectors/get_brushed_highlighted_shapes.ts

@@ -29,24 +29,11 @@ function getCurrentPointerStates(state: GlobalChartState) {
 /** @internal */
 export const getBrushedHighlightedShapesSelector = createCustomCachedSelector(
   [geometries, getCurrentPointerStates],
-  (geoms, pointerStates): DragShape | null => {
+  (geoms, pointerStates): DragShape => {
     if (!pointerStates.dragging || !pointerStates.down) {
       return null;
     }
 
-    const {
-      down: {
-        position: { x: startX, y: startY },
-      },
-      current: {
-        position: { x: endX, y: endY },
-      },
-    } = pointerStates;
-
-    const shape = geoms.pickDragShape([
-      { x: startX, y: startY },
-      { x: endX, y: endY },
-    ]);
-    return shape;
+    return geoms.pickDragShape([pointerStates.down.position, pointerStates.current.position]);
   },
 );

packages/charts/src/chart_types/heatmap/state/selectors/picked_shapes.ts

@@ -19,8 +19,7 @@
 
 import { GlobalChartState } from '../../../../state/chart_state';
 import { createCustomCachedSelector } from '../../../../state/create_selector';
-import { Cell } from '../../layout/types/viewmodel_types';
-import { TextBox } from '../../layout/viewmodel/viewmodel';
+import { Cell, TextBox } from '../../layout/types/viewmodel_types';
 import { geometries } from './geometries';
 
 function getCurrentPointerPosition(state: GlobalChartState) {